feat(BA-2791): Implement scaling group filtering with injectable rules (#6424)

Co-authored-by: Claude <[email protected]>

Author: HyeockJinKim

changes/6424.feature.md

@@ -0,0 +1 @@
+Implement scaling group filtering with injectable rules for session scheduling. The new ScalingGroupFilter applies configurable filter rules (public/private access, session type support) to determine eligible scaling groups before session creation, replacing the previous validation-only approach. This enables more flexible and extensible scaling group selection logic.

src/ai/backend/manager/api/auth.py

@@ -531,7 +531,6 @@ async def _authenticate_via_jwt(
 
     try:
         # 1. Decode token without verification to extract access_key
-        log.info("jwt token: {}", jwt_token)
         unverified_payload = pyjwt.decode(
             jwt_token,
             options={"verify_signature": False},
@@ -727,8 +726,6 @@ async def auth_middleware(request: web.Request, handler) -> web.StreamResponse:
     # Detect authentication method and route to appropriate flow
     jwt_token = request.headers.get("X-BackendAI-Token")
     auth_header = request.headers.get("Authorization")
-    log.info("AUTH.MIDDLEWARE(path:{}, headers:{})", request.path, dict(request.headers))
-
     if jwt_token:
         # JWT authentication flow (GraphQL Federation)
         await _authenticate_via_jwt(request, root_ctx, jwt_token)

src/ai/backend/manager/errors/resource.py

@@ -354,3 +354,16 @@ def error_code(cls) -> ErrorCode:
             operation=ErrorOperation.READ,
             error_detail=ErrorDetail.INVALID_PARAMETERS,
         )
+
+
+class NoAvailableScalingGroup(BackendAIError, web.HTTPBadRequest):
+    error_type = "https://api.backend.ai/probs/no-available-scaling-group"
+    error_title = "No scaling groups available for this session."
+
+    @classmethod
+    def error_code(cls) -> ErrorCode:
+        return ErrorCode(
+            domain=ErrorDomain.SCALING_GROUP,
+            operation=ErrorOperation.ACCESS,
+            error_detail=ErrorDetail.NOT_FOUND,
+        )

src/ai/backend/manager/sokovan/scheduling_controller/scheduling_controller.py

@@ -43,9 +43,10 @@
     ClusterValidationRule,
     ContainerLimitRule,
     MountNameValidationRule,
-    ScalingGroupAccessRule,
+    PublicPrivateFilterRule,
+    ScalingGroupFilter,
     ServicePortRule,
-    SessionTypeRule,
+    SessionTypeFilterRule,
     SessionValidator,
 )
 
@@ -76,6 +77,7 @@ class SchedulingController:
 
     # Services
     _scaling_group_resolver: ScalingGroupResolver
+    _scaling_group_filter: ScalingGroupFilter
     _validator: SessionValidator
     _preparer: SessionPreparer
     _resource_calculator: ResourceCalculator
@@ -98,11 +100,16 @@ def __init__(self, args: SchedulingControllerArgs) -> None:
         # Initialize services
         self._scaling_group_resolver = ScalingGroupResolver()
 
+        # Initialize scaling group filter with rules
+        filter_rules = [
+            PublicPrivateFilterRule(),
+            SessionTypeFilterRule(),
+        ]
+        self._scaling_group_filter = ScalingGroupFilter(filter_rules)
+
         # Initialize validator with rules
         validator_rules = [
             ContainerLimitRule(),
-            ScalingGroupAccessRule(),
-            SessionTypeRule(),
             ServicePortRule(),
             ClusterValidationRule(),
             MountNameValidationRule(),
@@ -196,10 +203,22 @@ async def enqueue_session(
                 allowed_vfolder_types,
             )
 
-        # Phase 3: Validate
+        # Phase 3: Filter and validate
         with self._metric_observer.measure_phase(
             "scheduling_controller", validated_scaling_group.name, "validation"
         ):
+            # Filter scaling groups based on session requirements
+            # This will raise NoAvailableScalingGroup if filtering fails
+            filter_result = self._scaling_group_filter.filter(
+                session_spec,
+                creation_context.allowed_scaling_groups,
+            )
+
+            # Update context with filtered scaling groups for remaining validation
+            creation_context.allowed_scaling_groups = filter_result.allowed_groups
+            session_spec.scaling_group = filter_result.selected_scaling_group
+
+            # Run remaining validation rules
             self._validator.validate(
                 session_spec,
                 creation_context,

src/ai/backend/manager/sokovan/scheduling_controller/validators/__init__.py

@@ -1,23 +1,35 @@
 """Validators for session creation."""
 
-from .base import SessionValidatorRule
+from .base import (
+    ScalingGroupFilterResult,
+    ScalingGroupFilterRule,
+    ScalingGroupFilterRuleResult,
+    SessionValidatorRule,
+)
 from .cluster import ClusterValidationRule
 from .mount import MountNameValidationRule
 from .rules import (
     ContainerLimitRule,
     ResourceLimitRule,
-    ScalingGroupAccessRule,
     ServicePortRule,
-    SessionTypeRule,
+)
+from .scaling_group_filter import (
+    PublicPrivateFilterRule,
+    ScalingGroupFilter,
+    SessionTypeFilterRule,
 )
 from .validator import SessionValidator
 
 __all__ = [
     "SessionValidator",
     "SessionValidatorRule",
     "ContainerLimitRule",
-    "ScalingGroupAccessRule",
-    "SessionTypeRule",
+    "ScalingGroupFilter",
+    "ScalingGroupFilterRule",
+    "ScalingGroupFilterResult",
+    "ScalingGroupFilterRuleResult",
+    "PublicPrivateFilterRule",
+    "SessionTypeFilterRule",
     "ServicePortRule",
     "ResourceLimitRule",
     "ClusterValidationRule",

src/ai/backend/manager/sokovan/scheduling_controller/validators/base.py

@@ -1,6 +1,9 @@
-"""Base classes for session validation rules."""
+"""Base classes for session validation and filtering rules."""
+
+from __future__ import annotations
 
 from abc import ABC, abstractmethod
+from dataclasses import dataclass
 
 from ai.backend.manager.repositories.scheduler.types.session_creation import (
     AllowedScalingGroup,
@@ -9,6 +12,58 @@
 )
 
 
+@dataclass
+class ScalingGroupFilterRuleResult:
+    """Result of a single scaling group filter rule."""
+
+    allowed_groups: list[AllowedScalingGroup]
+    """Scaling groups that passed this rule."""
+
+    rejected_groups: dict[str, str]
+    """Scaling groups that were rejected by this rule, mapped to rejection reason."""
+
+
+@dataclass
+class ScalingGroupFilterResult:
+    """Final result of scaling group filtering with selected group."""
+
+    allowed_groups: list[AllowedScalingGroup]
+    """Scaling groups that passed all filters."""
+
+    selected_scaling_group: str
+    """The selected scaling group name (either specified or auto-selected)."""
+
+
+class ScalingGroupFilterRule(ABC):
+    """
+    Abstract base class for scaling group filter rules.
+    Each rule filters scaling groups based on specific criteria.
+    """
+
+    @abstractmethod
+    def name(self) -> str:
+        """Return the filter rule name."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def filter(
+        self,
+        spec: SessionCreationSpec,
+        allowed_groups: list[AllowedScalingGroup],
+    ) -> ScalingGroupFilterRuleResult:
+        """
+        Filter scaling groups based on session creation specification.
+
+        Args:
+            spec: Session creation specification
+            allowed_groups: List of scaling groups to filter
+
+        Returns:
+            ScalingGroupFilterRuleResult containing allowed and rejected groups with reasons
+        """
+        raise NotImplementedError
+
+
 class SessionValidatorRule(ABC):
     """
     Abstract base class for session validator rules.
@@ -25,15 +80,13 @@ def validate(
         self,
         spec: SessionCreationSpec,
         context: SessionCreationContext,
-        allowed_groups: list[AllowedScalingGroup],
     ) -> None:
         """
         Validate a session creation specification.
 
         Args:
             spec: Session creation specification
             context: Pre-fetched context with all required data
-            allowed_groups: List of allowed scaling groups for the user
 
         Raises:
             InvalidAPIParameters or QuotaExceeded: If validation fails

src/ai/backend/manager/sokovan/scheduling_controller/validators/cluster.py

@@ -2,7 +2,6 @@
 
 from ai.backend.manager.errors.api import InvalidAPIParameters
 from ai.backend.manager.repositories.scheduler.types.session_creation import (
-    AllowedScalingGroup,
     SessionCreationContext,
     SessionCreationSpec,
 )
@@ -20,7 +19,6 @@ def validate(
         self,
         spec: SessionCreationSpec,
         context: SessionCreationContext,
-        allowed_groups: list[AllowedScalingGroup],
     ) -> None:
         """Validate cluster configuration and kernel specifications."""
         # Check if kernel_specs exists

src/ai/backend/manager/sokovan/scheduling_controller/validators/mount.py

@@ -5,7 +5,6 @@
 from ai.backend.manager.errors.api import InvalidAPIParameters
 from ai.backend.manager.models import verify_vfolder_name
 from ai.backend.manager.repositories.scheduler.types.session_creation import (
-    AllowedScalingGroup,
     SessionCreationContext,
     SessionCreationSpec,
 )
@@ -23,7 +22,6 @@ def validate(
         self,
         spec: SessionCreationSpec,
         context: SessionCreationContext,
-        allowed_groups: list[AllowedScalingGroup],
     ) -> None:
         """Validate mount names if mount map is provided."""
         mount_map = spec.creation_spec.get("mount_map") or {}

src/ai/backend/manager/sokovan/scheduling_controller/validators/rules.py

@@ -7,9 +7,7 @@
 from ai.backend.common.types import SlotName, SlotTypes
 from ai.backend.manager.errors.api import InvalidAPIParameters
 from ai.backend.manager.errors.kernel import QuotaExceeded
-from ai.backend.manager.models import PRIVATE_SESSION_TYPES
 from ai.backend.manager.repositories.scheduler.types.session_creation import (
-    AllowedScalingGroup,
     SessionCreationContext,
     SessionCreationSpec,
 )
@@ -29,7 +27,6 @@ def validate(
         self,
         spec: SessionCreationSpec,
         context: SessionCreationContext,
-        allowed_groups: list[AllowedScalingGroup],
     ) -> None:
         max_containers = spec.resource_policy.get("max_containers_per_session", 1)
         if spec.cluster_size > int(max_containers):
@@ -38,68 +35,6 @@ def validate(
             )
 
 
-class ScalingGroupAccessRule(SessionValidatorRule):
-    """Validates that the scaling group is accessible."""
-
-    @override
-    def name(self) -> str:
-        return "scaling_group_access"
-
-    @override
-    def validate(
-        self,
-        spec: SessionCreationSpec,
-        context: SessionCreationContext,
-        allowed_groups: list[AllowedScalingGroup],
-    ) -> None:
-        if not spec.scaling_group:
-            # Should have been resolved already
-            return
-
-        public_sgroup_only = spec.session_type not in PRIVATE_SESSION_TYPES
-
-        # Find the scaling group in allowed list
-        for sg in allowed_groups:
-            if sg.name == spec.scaling_group:
-                if public_sgroup_only and sg.is_private:
-                    raise InvalidAPIParameters(
-                        f"Scaling group {spec.scaling_group} is not allowed for {spec.session_type} sessions"
-                    )
-                return
-
-        raise InvalidAPIParameters(f"Scaling group {spec.scaling_group} is not accessible")
-
-
-class SessionTypeRule(SessionValidatorRule):
-    """Validates session type compatibility with scaling group."""
-
-    @override
-    def name(self) -> str:
-        return "session_type"
-
-    @override
-    def validate(
-        self,
-        spec: SessionCreationSpec,
-        context: SessionCreationContext,
-        allowed_groups: list[AllowedScalingGroup],
-    ) -> None:
-        if spec.scaling_group is None:
-            # Should have been resolved already
-            return
-
-        for sg in allowed_groups:
-            if sg.name == spec.scaling_group:
-                allowed_session_types = sg.scheduler_opts.allowed_session_types
-                if spec.session_type not in allowed_session_types:
-                    raise InvalidAPIParameters(
-                        f"Session type {spec.session_type} is not allowed in scaling group {sg.name}"
-                    )
-                return
-
-        raise InvalidAPIParameters(f"Scaling group {spec.scaling_group} is not accessible")
-
-
 class ServicePortRule(SessionValidatorRule):
     """Validates preopen ports against service ports."""
 
@@ -112,7 +47,6 @@ def validate(
         self,
         spec: SessionCreationSpec,
         context: SessionCreationContext,
-        allowed_groups: list[AllowedScalingGroup],
     ) -> None:
         # Check preopen_ports from creation_config (applies to all kernels)
         creation_preopen_ports = spec.creation_spec.get("preopen_ports")
@@ -186,7 +120,6 @@ def validate(
         self,
         spec: SessionCreationSpec,
         context: SessionCreationContext,
-        allowed_groups: list[AllowedScalingGroup],
     ) -> None:
         # Note: This validation should ideally be done after resource calculation
         # For now, we'll validate what we can from the spec