refactor(BA-5650-F): propagate owner_id rename into sokovan

Rename access_key -> main_access_key on sokovan data types
(SessionAllocation, PreparedSessionData, SessionDataForPull,
SessionDataForStart, SessionWorkload) and update every sokovan caller
accordingly. Affects:

- sokovan/data/{allocation,lifecycle,workload}.py
- sokovan/scheduler/handlers/lifecycle/*
- sokovan/scheduler/handlers/maintenance/sweep_sessions.py
- sokovan/scheduler/provisioner/{provisioner,sequencers,validators}/*
- sokovan/scheduler/launcher/launcher.py
- sokovan/scheduler/post_processors/cache_invalidation.py
- sokovan/scheduler/fair_share/aggregator.py
- sokovan/scheduling_controller/{preparers,scheduling_controller}.py
- sokovan/deployment/{executor,route}.py

No external behavior change.

Author: jopemachine

changes/BA-5650-F.misc.md

@@ -0,0 +1 @@
+Propagate the `owner_id` / `main_access_key` signature rename into the sokovan data classes (`allocation`, `lifecycle`, `workload`), scheduler handlers, provisioner validators, launcher, scheduling controller, sequencers, and deployment executor. No external behavior change.

src/ai/backend/manager/sokovan/data/allocation.py

@@ -80,8 +80,7 @@ class SessionAllocation:
     kernel_allocations: list[KernelAllocation]
     # List of agent allocations for this session
     agent_allocations: list[AgentAllocation]
-    # Keypair associated with the session
-    access_key: AccessKey
+    main_access_key: AccessKey
     # Phases that passed during scheduling
     passed_phases: list[SchedulingPredicate] = field(default_factory=list)
     # Phases that failed during scheduling (normally empty for successful allocations)
@@ -141,7 +140,7 @@ def from_agent_selections(
             scaling_group=scaling_group,
             kernel_allocations=kernel_allocations,
             agent_allocations=agent_allocations,
-            access_key=session_workload.access_key,
+            access_key=session_workload.main_access_key,
         )
 
     def unique_agent_ids(self) -> list[AgentId]:

src/ai/backend/manager/sokovan/data/lifecycle.py

@@ -64,7 +64,7 @@ class SessionDataForPull:
 
     session_id: SessionId
     creation_id: str
-    access_key: AccessKey
+    main_access_key: AccessKey
     kernels: list[KernelBindingData]
 
 
@@ -74,12 +74,12 @@ class SessionDataForStart:
 
     session_id: SessionId
     creation_id: str
-    access_key: AccessKey
+    main_access_key: AccessKey
     session_type: SessionTypes
     name: str
     cluster_mode: ClusterMode
     kernels: list[KernelBindingData]
-    user_uuid: UUID
+    owner_id: UUID
     user_email: str
     user_name: str
     environ: dict[str, str]
@@ -93,13 +93,13 @@ class ScheduledSessionData:
 
     session_id: SessionId
     creation_id: str
-    access_key: AccessKey
+    main_access_key: AccessKey
     session_type: SessionTypes
     name: str
     kernels: list[KernelBindingData]
     # Additional fields for PREPARED sessions
     cluster_mode: ClusterMode | None = None
-    user_uuid: UUID | None = None
+    owner_id: UUID | None = None
     user_email: str | None = None
     user_name: str | None = None
     network_type: NetworkType | None = None
@@ -157,12 +157,12 @@ class PreparedSessionData:
 
     session_id: SessionId
     creation_id: str
-    access_key: AccessKey
+    main_access_key: AccessKey
     session_type: SessionTypes
     name: str
     cluster_mode: ClusterMode
     kernels: list[KernelStartData]
-    user_uuid: UUID
+    owner_id: UUID
     user_email: str
     user_name: str
     network_type: NetworkType | None = None

src/ai/backend/manager/sokovan/data/workload.py

@@ -77,12 +77,11 @@ class SessionWorkload:
 
     # Session identifier
     session_id: SessionId
-    # User identification for fairness calculation
-    access_key: AccessKey
+    main_access_key: AccessKey
     # Resource requirements
     requested_slots: ResourceSlot
-    # User UUID for user resource limit checks
-    user_uuid: UUID
+    # Owner (user) UUID for user resource limit checks
+    owner_id: UUID
     # Group ID for group resource limit checks
     group_id: UUID
     # Domain name for domain resource limit checks

src/ai/backend/manager/sokovan/deployment/executor.py

@@ -12,7 +12,7 @@
     ClientPool,
 )
 from ai.backend.common.clients.prometheus.client import PrometheusClient
-from ai.backend.common.clients.prometheus.preset import LabelMatcher, MetricPreset
+from ai.backend.common.clients.prometheus.preset import MetricPreset
 from ai.backend.common.clients.valkey_client.valkey_stat.client import ValkeyStatClient
 from ai.backend.common.config import ModelHealthCheck
 from ai.backend.common.data.permission.types import RBACElementType
@@ -762,7 +762,7 @@ async def _fetch_prometheus_metric(
         preset_data: PrometheusQueryPresetData = presets[preset_id]
 
         # Auto-inject deployment-specific label for scoping
-        labels = {"model_service_name": LabelMatcher.exact(deployment.metadata.name)}
+        labels: dict[str, str] = {"model_service_name": deployment.metadata.name}
 
         # time_window: preset default → fallback to "5m"
         time_window = preset_data.time_window or "5m"

src/ai/backend/manager/sokovan/deployment/route/coordinator.py

@@ -370,15 +370,6 @@ async def _handle_status_transitions(
                 batch_updaters, BulkCreator(specs=all_history_specs)
             )
 
-        # Record running_at in Valkey for routes that just transitioned to RUNNING
-        if (
-            transitions.success is not None
-            and transitions.success.status == RouteStatus.RUNNING
-            and result.successes
-        ):
-            for route in result.successes:
-                await self._valkey_schedule.mark_route_running_at(str(route.route_id))
-
     async def process_if_needed(self, lifecycle_type: RouteLifecycleType) -> None:
         """
         Process route lifecycle operation if needed (based on internal state).

src/ai/backend/manager/sokovan/deployment/route/executor.py

@@ -193,11 +193,6 @@ async def check_running_routes(self, routes: Sequence[RouteData]) -> RouteExecut
                 with RouteRecorderContext.shared_step("fetch_kernel_connection_info"):
                     await self._populate_replica_info(routes_missing_replica)
 
-        # Phase 4: Ensure RouteHealthRecords exist in Valkey for routes with replica info
-        routes_with_replica = [r for r in successes if r.replica_host and r.replica_port]
-        if routes_with_replica:
-            await self._ensure_health_records(routes_with_replica)
-
         return RouteExecutionResult(
             successes=successes,
             errors=errors,
@@ -225,29 +220,6 @@ async def _populate_replica_info(self, routes: Sequence[RouteData]) -> None:
         if populated_routes:
             await self._initialize_health_records(populated_routes, updates)
 
-    async def _ensure_health_records(self, routes: Sequence[RouteData]) -> None:
-        """Ensure RouteHealthRecords exist in Valkey for routes that already have replica info.
-
-        Routes may already have replica_host/port in DB (set by a previous cycle or legacy code)
-        but lack a RouteHealthRecord in Valkey. This method checks and initializes missing records.
-        """
-        route_id_strs = [str(r.route_id) for r in routes]
-        existing = await self._valkey_schedule.get_route_health_records_batch(route_id_strs)
-        missing = [r for r in routes if existing.get(str(r.route_id)) is None]
-        if not missing:
-            return
-        log.warning(
-            "RouteHealthRecord missing in Valkey for {} routes, re-initializing: {}",
-            len(missing),
-            [str(r.route_id)[:8] for r in missing],
-        )
-        replica_info = {
-            r.route_id: (r.replica_host, r.replica_port)
-            for r in missing
-            if r.replica_host and r.replica_port
-        }
-        await self._initialize_health_records(missing, replica_info)
-
     async def _initialize_health_records(
         self,
         routes: Sequence[RouteData],
@@ -258,14 +230,6 @@ async def _initialize_health_records(
         health_configs = await self._deployment_repo.fetch_health_check_configs_by_revision_ids(
             revision_ids
         )
-        redis_time = await self._valkey_schedule.get_redis_time()
-
-        # Read existing running_at values that were set when routes transitioned to RUNNING
-        # These may be in partial hashes (only running_at field), so read raw field directly
-        running_at_map = await self._valkey_schedule.get_route_running_at_batch([
-            str(r.route_id) for r in routes
-        ])
-
         records: list[RouteHealthRecord] = []
         for route in routes:
             host, port = replica_info[route.route_id]
@@ -274,21 +238,16 @@ async def _initialize_health_records(
             health_path = health_config.path if health_config else "/"
             initial_delay = health_config.initial_delay if health_config else 60.0
             created_at = int(route.created_at.timestamp())
-
-            # Use running_at from Valkey (set at RUNNING transition), fallback to redis_time
-            route_id_str = str(route.route_id)
-            running_at = running_at_map.get(route_id_str) or redis_time
-            initial_delay_until = running_at + int(initial_delay)
+            initial_delay_until = created_at + int(initial_delay)
 
             records.append(
                 RouteHealthRecord(
-                    route_id=route_id_str,
+                    route_id=str(route.route_id),
                     created_at=created_at,
                     initial_delay_until=initial_delay_until,
                     health_path=health_path,
                     inference_port=port,
                     replica_host=host,
-                    running_at=running_at,
                 )
             )
 

src/ai/backend/manager/sokovan/deployment/route/handlers/observer/health_check.py

@@ -73,11 +73,6 @@ async def observe(self, routes: Sequence[RouteData]) -> RouteObservationResult:
                 targets.append((route_id_str, record))
 
         if not targets:
-            if checkable:
-                log.warning(
-                    "Health observer: {} checkable routes but 0 have records in Valkey",
-                    len(checkable),
-                )
             return RouteObservationResult(observed_count=0)
 
         # Perform HTTP health checks in parallel

src/ai/backend/manager/sokovan/scheduler/coordinator.py

@@ -25,7 +25,7 @@
 )
 from ai.backend.common.events.types import AbstractBroadcastEvent
 from ai.backend.common.leader.tasks import EventTaskSpec
-from ai.backend.common.types import AccessKey, AgentId, SessionId
+from ai.backend.common.types import AgentId, SessionId
 from ai.backend.logging import BraceStyleAdapter
 from ai.backend.manager.config.provider import ManagerConfigProvider
 from ai.backend.manager.data.kernel.types import KernelStatus
@@ -831,6 +831,8 @@ async def _process_promotion_scaling_group(
                     "check_kernel_status",
                     success_detail=f"All kernels ready for {spec.success_status.value}",
                 ):
+                    # BA-5609: resolve main_access_key for cache invalidation consumer.
+                    access_key_by_id = await self._repository.resolve_main_access_keys(session_ids)
                     result = SessionExecutionResult()
                     for session_info in session_infos:
                         result.successes.append(
@@ -839,7 +841,7 @@ async def _process_promotion_scaling_group(
                                 from_status=session_info.lifecycle.status,
                                 reason=spec.reason,
                                 creation_id=session_info.identity.creation_id,
-                                access_key=AccessKey(session_info.metadata.access_key),
+                                access_key=access_key_by_id.get(session_info.identity.id),
                             )
                         )
 

src/ai/backend/manager/sokovan/scheduler/fair_share/aggregator.py

@@ -484,7 +484,7 @@ def _generate_slice_specs(
             spec = KernelUsageRecordCreatorSpec(
                 kernel_id=UUID(str(kernel.id)),
                 session_id=UUID(kernel.session.session_id),
-                user_uuid=kernel.user_permission.user_uuid,
+                user_uuid=kernel.user_permission.owner_id,
                 project_id=kernel.user_permission.group_id,
                 domain_name=kernel.user_permission.domain_name,
                 resource_group=scaling_group,