refactor: Split PendingSession Scheduler into PendingSession Scheduler and AgentSelector

changes/1655.fix.md

@@ -0,0 +1 @@
+Refactor `PendingSession` Scheduler into `PendingSession` scheduler and `AgentSelector`, and replace `roundrobin` flag with `AgentSelectionStrategy.RoundRobin` policy.

src/ai/backend/common/config.py

@@ -89,6 +89,14 @@ class BaseSchema(BaseModel):
     ),
 }).allow_extra("*")
 
+# Used in Etcd as a global config.
+# If `scalingGroup.scheduler_opts` contains an `agent_selector_config`, it will override this.
+agent_selector_globalconfig_iv = t.Dict({}).allow_extra("*")
+
+# Used in `scalingGroup.scheduler_opts` as a per scaling_group config.
+agent_selector_config_iv = t.Dict({}) | agent_selector_globalconfig_iv
+
+
 model_definition_iv = t.Dict({
     t.Key("models"): t.List(
         t.Dict({

src/ai/backend/common/types.py

@@ -27,6 +27,7 @@
     NewType,
     NotRequired,
     Optional,
+    Self,
     Sequence,
     Tuple,
     Type,
@@ -218,7 +219,9 @@ def check_typed_dict(value: Mapping[Any, Any], expected_type: Type[TD]) -> TD:
 SessionId = NewType("SessionId", uuid.UUID)
 KernelId = NewType("KernelId", uuid.UUID)
 ImageAlias = NewType("ImageAlias", str)
+ArchName = NewType("ArchName", str)
 
+ResourceGroupID = NewType("ResourceGroupID", str)
 AgentId = NewType("AgentId", str)
 DeviceName = NewType("DeviceName", str)
 DeviceId = NewType("DeviceId", str)
@@ -840,7 +843,7 @@ def to_json(self) -> dict[str, Any]:
         raise NotImplementedError
 
     @classmethod
-    def from_json(cls, obj: Mapping[str, Any]) -> JSONSerializableMixin:
+    def from_json(cls, obj: Mapping[str, Any]) -> Self:
         return cls(**cls.as_trafaret().check(obj))
 
     @classmethod
@@ -985,7 +988,7 @@ def to_json(self) -> dict[str, Any]:
         return {host: [perm.value for perm in perms] for host, perms in self.items()}
 
     @classmethod
-    def from_json(cls, obj: Mapping[str, Any]) -> JSONSerializableMixin:
+    def from_json(cls, obj: Mapping[str, Any]) -> Self:
         return cls(**cls.as_trafaret().check(obj))
 
     @classmethod
@@ -1197,6 +1200,7 @@ class AcceleratorMetadata(TypedDict):
 class AgentSelectionStrategy(enum.StrEnum):
     DISPERSED = "dispersed"
     CONCENTRATED = "concentrated"
+    ROUNDROBIN = "roundrobin"
     # LEGACY chooses the largest agent (the sort key is a tuple of resource slots).
     LEGACY = "legacy"
 
@@ -1215,29 +1219,6 @@ class VolumeMountableNodeType(enum.StrEnum):
     STORAGE_PROXY = enum.auto()
 
 
-@dataclass
-class RoundRobinState(JSONSerializableMixin):
-    schedulable_group_id: str
-    next_index: int
-
-    def to_json(self) -> dict[str, Any]:
-        return dataclasses.asdict(self)
-
-    @classmethod
-    def from_json(cls, obj: Mapping[str, Any]) -> RoundRobinState:
-        return cls(**cls.as_trafaret().check(obj))
-
-    @classmethod
-    def as_trafaret(cls) -> t.Trafaret:
-        return t.Dict({
-            t.Key("schedulable_group_id"): t.String,
-            t.Key("next_index"): t.Int,
-        })
-
-
-# States of the round-robin scheduler for each resource group and architecture.
-RoundRobinStates: TypeAlias = dict[str, dict[str, RoundRobinState]]
-
 SSLContextType: TypeAlias = bool | Fingerprint | SSLContext
 
 

src/ai/backend/common/validators.py

@@ -47,7 +47,6 @@
 from .types import BinarySize as _BinarySize
 from .types import HostPortPair as _HostPortPair
 from .types import QuotaScopeID as _QuotaScopeID
-from .types import RoundRobinState, RoundRobinStates
 from .types import VFolderID as _VFolderID
 
 __all__ = (
@@ -727,25 +726,3 @@ def check_and_return(self, value: Any) -> float:
                 return 0
             case _:
                 self._failure(f"Value must be (float, tuple of float or None), not {type(value)}.")
-
-
-class RoundRobinStatesJSONString(t.Trafaret):
-    def check_and_return(self, value: Any) -> RoundRobinStates:
-        try:
-            rr_states_dict: dict[str, dict[str, dict[str, Any]]] = json.loads(value)
-        except (KeyError, ValueError, json.decoder.JSONDecodeError):
-            self._failure(
-                f"Expected valid JSON string, got `{value}`. RoundRobinStatesJSONString should"
-                " be a valid JSON string",
-                value=value,
-            )
-
-        rr_states: RoundRobinStates = {}
-        for resource_group, arch_rr_states_dict in rr_states_dict.items():
-            rr_states[resource_group] = {}
-            for arch, rr_state_dict in arch_rr_states_dict.items():
-                if "next_index" not in rr_state_dict or "schedulable_group_id" not in rr_state_dict:
-                    self._failure("Invalid roundrobin states")
-                rr_states[resource_group][arch] = RoundRobinState.from_json(rr_state_dict)
-
-        return rr_states

src/ai/backend/manager/BUILD

@@ -43,7 +43,12 @@ python_distribution(
             "fifo": "ai.backend.manager.scheduler.fifo:FIFOSlotScheduler",
             "lifo": "ai.backend.manager.scheduler.fifo:LIFOSlotScheduler",
             "drf": "ai.backend.manager.scheduler.drf:DRFScheduler",
-            "mof": "ai.backend.manager.scheduler.mof:MOFScheduler",
+        },
+        "backendai_agentselector_v10": {
+            "legacy": "ai.backend.manager.scheduler.agent_selector:LegacyAgentSelector",
+            "roundrobin": "ai.backend.manager.scheduler.agent_selector:RoundRobinAgentSelector",
+            "concentrated": "ai.backend.manager.scheduler.agent_selector:ConcentratedAgentSelector",
+            "dispersed": "ai.backend.manager.scheduler.agent_selector:DispersedAgentSelector",
         },
         "backendai_error_monitor_v20": {
             "intrinsic": "ai.backend.manager.plugin.error_monitor:ErrorMonitor",

src/ai/backend/manager/config.py

@@ -174,7 +174,6 @@
 
 from __future__ import annotations
 
-import json
 import logging
 import os
 import secrets
@@ -212,7 +211,6 @@
 from ai.backend.common.lock import EtcdLock, FileLock, RedisLock
 from ai.backend.common.types import (
     HostPortPair,
-    RoundRobinState,
     SlotName,
     SlotTypes,
     current_resource_slots,
@@ -350,6 +348,7 @@
     "plugins": {
         "accelerator": {},
         "scheduler": {},
+        "agent_selector": {},
     },
     "watcher": {
         "token": None,
@@ -439,6 +438,9 @@ def container_registry_serialize(v: dict[str, Any]) -> dict[str, str]:
         t.Key("scheduler", default=_config_defaults["plugins"]["scheduler"]): t.Mapping(
             t.String, t.Mapping(t.String, t.Any)
         ),
+        t.Key("agent_selector", default=_config_defaults["plugins"]["agent_selector"]): t.Mapping(
+            t.String, config.agent_selector_globalconfig_iv
+        ),
     }).allow_extra("*"),
     t.Key("network", default=_config_defaults["network"]): t.Dict({
         t.Key("subnet", default=_config_defaults["network"]["subnet"]): t.Dict({
@@ -471,7 +473,6 @@ def container_registry_serialize(v: dict[str, Any]) -> dict[str, str]:
             ): session_hang_tolerance_iv,
         },
     ).allow_extra("*"),
-    t.Key("roundrobin_states", default=None): t.Null | tx.RoundRobinStatesJSONString,
 }).allow_extra("*")
 
 _volume_defaults: dict[str, Any] = {
@@ -847,38 +848,3 @@ def get_redis_url(self, db: int = 0) -> yarl.URL:
             self.data["redis"]["addr"][1]
         ).with_password(self.data["redis"]["password"]) / str(db)
         return url
-
-    async def get_roundrobin_state(
-        self, resource_group_name: str, architecture: str
-    ) -> RoundRobinState | None:
-        """
-        Return the roundrobin state for the given resource group and architecture.
-        If given resource group's roundrobin states or roundrobin state of the given architecture is not found, return None.
-        """
-        if (rr_state_str := await self.get_raw("roundrobin_states")) is not None:
-            rr_states_dict: dict[str, dict[str, Any]] = json.loads(rr_state_str)
-            resource_group_rr_states_dict = rr_states_dict.get(resource_group_name, None)
-
-            if resource_group_rr_states_dict is not None:
-                rr_state_dict = resource_group_rr_states_dict.get(architecture, None)
-
-                if rr_state_dict is not None:
-                    return RoundRobinState(
-                        schedulable_group_id=rr_state_dict["schedulable_group_id"],
-                        next_index=rr_state_dict["next_index"],
-                    )
-
-        return None
-
-    async def put_roundrobin_state(
-        self, resource_group_name: str, architecture: str, state: RoundRobinState
-    ) -> None:
-        """
-        Update the roundrobin states using the given resource group and architecture key.
-        """
-        rr_states_dict = json.loads(await self.get_raw("roundrobin_states") or "{}")
-        if resource_group_name not in rr_states_dict:
-            rr_states_dict[resource_group_name] = {}
-
-        rr_states_dict[resource_group_name][architecture] = state.to_json()
-        await self.etcd.put("roundrobin_states", json.dumps(rr_states_dict))

src/ai/backend/manager/models/alembic/versions/c4b7ec740b36_migrate_roundrobin_strategy_to_agent_.py

@@ -0,0 +1,59 @@
+"""migrate-roundrobin-strategy-to-agent-selector-config
+
+Revision ID: c4b7ec740b36
+Revises: 59a622c31820
+Create Date: 2024-09-17 00:31:31.379466
+
+"""
+
+from alembic import op
+from sqlalchemy.sql import text
+
+# revision identifiers, used by Alembic.
+revision = "c4b7ec740b36"
+down_revision = "59a622c31820"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    op.execute(
+        text("""
+    UPDATE scaling_groups
+    SET scheduler_opts =
+    CASE
+        WHEN scheduler_opts ? 'roundrobin' AND scheduler_opts->>'roundrobin' = 'true' THEN
+            jsonb_set(
+                jsonb_set(
+                    scheduler_opts - 'roundrobin',
+                    '{agent_selection_strategy}', '"roundrobin"'::jsonb
+                ),
+                '{agent_selector_config}', '{}'::jsonb
+            )
+        ELSE jsonb_set(
+            scheduler_opts - 'roundrobin',
+            '{agent_selector_config}', '{}'::jsonb
+        )
+    END;
+    """)
+    )
+
+
+def downgrade() -> None:
+    op.execute(
+        text("""
+    UPDATE scaling_groups
+    SET scheduler_opts =
+    CASE
+        WHEN scheduler_opts->>'agent_selection_strategy' = 'roundrobin' THEN
+            jsonb_set(
+                jsonb_set(
+                    scheduler_opts - 'agent_selector_config',
+                    '{agent_selection_strategy}', '"legacy"'::jsonb
+                ),
+                '{roundrobin}', 'true'::jsonb
+            )
+        ELSE scheduler_opts - 'agent_selector_config'
+    END;
+    """)
+    )

src/ai/backend/manager/models/scaling_group.py

@@ -29,6 +29,7 @@
 from sqlalchemy.sql.expression import true
 
 from ai.backend.common import validators as tx
+from ai.backend.common.config import agent_selector_config_iv
 from ai.backend.common.types import (
     AgentSelectionStrategy,
     JSONSerializableMixin,
@@ -98,17 +99,20 @@ class ScalingGroupOpts(JSONSerializableMixin):
         ],
     )
     pending_timeout: timedelta = timedelta(seconds=0)
-    config: Mapping[str, Any] = attr.Factory(dict)
+    config: Mapping[str, Any] = attr.field(factory=dict)
+
+    # Scheduler has a dedicated database column to store its name,
+    # but agent selector configuration is stored as a part of the scheduler_opts column.
     agent_selection_strategy: AgentSelectionStrategy = AgentSelectionStrategy.DISPERSED
-    roundrobin: bool = False
+    agent_selector_config: Mapping[str, Any] = attr.field(factory=dict)
 
     def to_json(self) -> dict[str, Any]:
         return {
             "allowed_session_types": [item.value for item in self.allowed_session_types],
             "pending_timeout": self.pending_timeout.total_seconds(),
             "config": self.config,
             "agent_selection_strategy": self.agent_selection_strategy,
-            "roundrobin": self.roundrobin,
+            "agent_selector_config": self.agent_selector_config,
         }
 
     @classmethod
@@ -127,7 +131,7 @@ def as_trafaret(cls) -> t.Trafaret:
             t.Key("agent_selection_strategy", default=AgentSelectionStrategy.DISPERSED): tx.Enum(
                 AgentSelectionStrategy
             ),
-            t.Key("roundrobin", default=False): t.Bool(),
+            t.Key("agent_selector_config", default={}): agent_selector_config_iv,
         }).allow_extra("*")