Add use_num_extras flag temporarily

lablup · Oct 26, 2023 · b228c3e · b228c3e
1 parent fec83ce
commit b228c3e
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 3 deletions.
diff --git a/changes/1655.fix.md b/changes/1655.fix.md
@@ -0,0 +1 @@
+Replace RoundRobin flag with AgentSelectionStrategy.RoundRobin
diff --git a/src/ai/backend/manager/scheduler/drf.py b/src/ai/backend/manager/scheduler/drf.py
@@ -113,6 +113,7 @@ async def _assign_agent(
         return await self.select_agent(
             possible_agents,
             pending_session_or_kernel,
+            False,
             roundrobin_context,
         )
 

diff --git a/src/ai/backend/manager/scheduler/fifo.py b/src/ai/backend/manager/scheduler/fifo.py
@@ -71,6 +71,7 @@ async def assign_agent_for_session(
         return await self.select_agent(
             agents,
             pending_session,
+            True,
             roundrobin_context,
         )
 
@@ -82,6 +83,7 @@ async def assign_agent_for_kernel(
         return await self.select_agent(
             agents,
             pending_kernel,
+            True,
         )
 
 
@@ -106,6 +108,7 @@ async def assign_agent_for_session(
         return await self.select_agent(
             agents,
             pending_session,
+            True,
             roundrobin_context,
         )
 
@@ -117,4 +120,5 @@ async def assign_agent_for_kernel(
         return await self.select_agent(
             agents,
             pending_kernel,
+            True,
         )
diff --git a/src/ai/backend/manager/scheduler/types.py b/src/ai/backend/manager/scheduler/types.py
@@ -472,6 +472,7 @@ async def select_agent(
         self,
         possible_agents: Sequence[AgentRow],
         pending_session_or_kernel: SessionRow | KernelInfo,
+        use_num_extras: bool,
         roundrobin_context: Optional[RoundRobinContext] = None,
     ) -> Optional[AgentId]:
         """
@@ -536,7 +537,7 @@ async def select_agent(
                 chosen_agent = max(
                     possible_agents,
                     key=lambda agent: [
-                        -get_num_extras(agent, requested_slots),
+                        -get_num_extras(agent, requested_slots) if use_num_extras else 0,
                         *[
                             agent.available_slots.get(key, -sys.maxsize)
                             for key in resource_priorities
@@ -547,7 +548,7 @@ async def select_agent(
                 chosen_agent = min(
                     possible_agents,
                     key=lambda agent: [
-                        get_num_extras(agent, requested_slots),
+                        get_num_extras(agent, requested_slots) if use_num_extras else 0,
                         *[
                             (agent.available_slots - agent.occupied_slots).get(key, sys.maxsize)
                             for key in resource_priorities
@@ -558,7 +559,7 @@ async def select_agent(
                 chosen_agent = max(
                     possible_agents,
                     key=lambda agent: [
-                        -get_num_extras(agent, requested_slots),
+                        -get_num_extras(agent, requested_slots) if use_num_extras else 0,
                         *[
                             (agent.available_slots - agent.occupied_slots).get(key, -sys.maxsize)
                             for key in resource_priorities
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1 @@
		Replace RoundRobin flag with AgentSelectionStrategy.RoundRobin