diff --git a/changes/13707.fix.md b/changes/13707.fix.md new file mode 100644 index 00000000000..f71cd976b8e --- /dev/null +++ b/changes/13707.fix.md @@ -0,0 +1 @@ +Stop scheduling sessions behind one whose resource group ran out of resources, so lower-priority sessions no longer take what a blocked higher-priority session is waiting for diff --git a/src/ai/backend/manager/models/scheduling_history/row.py b/src/ai/backend/manager/models/scheduling_history/row.py index 622a42e5fc5..bd113cfb8bc 100644 --- a/src/ai/backend/manager/models/scheduling_history/row.py +++ b/src/ai/backend/manager/models/scheduling_history/row.py @@ -78,17 +78,28 @@ class SessionSchedulingHistoryRow(Base): onupdate=sa.func.now(), ) + def records_an_attempt(self) -> bool: + """Whether this record describes an attempt rather than a skip.""" + return self.result != SchedulingResult.SKIPPED + def should_merge_with(self, new_row: SessionSchedulingHistoryRow) -> bool: """Check if a new entry should be merged with this one. Merge conditions: - Same phase, error_code, and to_status -> merge (increment attempts) - - from_status and result (success/failure) do not affect merge decision + - Both must describe an attempt, or both a skip + - from_status and which attempt result (success/failure/give-up) do + not affect the merge decision + + Skips are kept apart because ``attempts`` drives the give-up + (deprioritization) classification, which may only count attempts a + session really got. """ return ( self.phase == new_row.phase and self.error_code == new_row.error_code and self.to_status == new_row.to_status + and self.records_an_attempt() == new_row.records_an_attempt() ) def to_data(self) -> SessionSchedulingHistoryData: diff --git a/src/ai/backend/manager/sokovan/scheduler/coordinator.py b/src/ai/backend/manager/sokovan/scheduler/coordinator.py index f0fc0f191fa..1c4c2e632af 100644 --- a/src/ai/backend/manager/sokovan/scheduler/coordinator.py +++ b/src/ai/backend/manager/sokovan/scheduler/coordinator.py @@ -39,6 +39,7 @@ ) from ai.backend.manager.metrics.scheduler import SchedulerOperationMetricObserver from ai.backend.manager.models.kernel.conditions import KernelConditions +from ai.backend.manager.models.scheduling_history.row import SessionSchedulingHistoryRow from ai.backend.manager.models.session.conditions import SessionConditions from ai.backend.manager.models.specs.pagination import NoPagination, OffsetPagination from ai.backend.manager.repositories.base import BatchQuerier @@ -58,6 +59,7 @@ from ai.backend.manager.types import DistributedLockFactory from ai.backend.manager.views.sokovan.lifecycle import ( KernelCreationInfo, + LastPhase, SessionWithKernels, ) from ai.backend.manager.views.sokovan.result import PromotionSpec @@ -140,7 +142,8 @@ class FailureClassificationResult: Classification priority (first match wins): 1. give_up: per-handler ``max_retry_count`` is set AND - ``phase_attempts`` reached it. ``None`` means "no retry limit" + the last phase record's ``attempts`` reached it. ``None`` means + "no retry limit" — give-up never fires. 2. expired: per-handler ``timeout`` is set AND phase elapsed time exceeded it. ``None`` means "no timeout" — expired never fires. @@ -745,19 +748,11 @@ async def _process_resource_group( # Extract session IDs for recorder entity_ids session_ids = [s.session_info.identity.id for s in sessions] - # Populate phase_attempts and phase_started_at from scheduling history for failure classification + # Carry the last phase record onto the sessions for failure classification # Get last history records (regardless of phase), then compare phase at application level history_map = await self._repository.get_last_session_histories(session_ids) handler_name = handler.name() - for session in sessions: - history = history_map.get(session.session_info.identity.id) - # Only use history data if the last history is for the current phase - if history and history.phase == handler_name: - session.phase_attempts = history.attempts - session.phase_started_at = history.created_at - else: - session.phase_attempts = 0 - session.phase_started_at = None + self._populate_phase_history(sessions, history_map, handler_name) # Create recorder scoped to this resource group recorder_scope = f"{schedule_type.value}:{resource_group_id}" @@ -802,6 +797,30 @@ async def _process_resource_group( resource_group_id, ) + def _populate_phase_history( + self, + sessions: list[SessionWithKernels], + history_map: Mapping[SessionId, SessionSchedulingHistoryRow], + handler_name: str, + ) -> None: + """Carry this phase's last history record onto the sessions. + + Only the last record counts, and only while it is still this phase's; + anything else leaves the session without one. Skips are carried like + any other record — ``LastPhase.result`` says what was counted, and the + give-up rule in :meth:`_classify_failures` is what excludes skips. + """ + for session in sessions: + history = history_map.get(session.session_info.identity.id) + if history and history.phase == handler_name: + session.last_phase = LastPhase( + attempts=history.attempts, + started_at=history.created_at, + result=SchedulingResult(history.result), + ) + else: + session.last_phase = None + async def _process_promotion_resource_group( self, spec: PromotionSpec, @@ -1120,7 +1139,7 @@ async def _handle_result( handler: The lifecycle handler that produced the result result: Execution result containing successes, failures, and skipped records: Mapping of session IDs to their execution records for sub_steps - sessions: Original sessions with phase_attempts for failure classification + sessions: Original sessions with last_phase for failure classification Returns: FailureClassificationResult if there were failures, None otherwise. @@ -1230,13 +1249,15 @@ def _classify_failures( is ``None``; ``expired`` never fires when ``timeout`` is ``None``. Classification priority (first match wins): - 1. give_up: max_retry_count set AND phase_attempts >= max_retry_count + 1. give_up: max_retry_count set AND last_phase.attempts >= max_retry_count, + and those attempts are not skips (a session queued behind a blocked + one was never attempted, so it may not be given up on) 2. expired: timeout set AND phase elapsed > timeout 3. need_retry: default Args: failures: Failed session transition info - sessions: Original sessions with phase_attempts and phase_started_at populated + sessions: Original sessions with last_phase populated current_time: Current database time for timeout comparison handler_name: ``SessionLifecycleHandler.name()`` for resolving per-handler policy from ``SessionHandlerOptions`` @@ -1257,14 +1278,21 @@ def _classify_failures( continue policy = session.session_info.handler_options.resolve(handler_name) - - # 1. Check max retries exceeded → give_up - if policy.is_retry_exhausted(session.phase_attempts): + last_phase = session.last_phase + + # 1. Check max retries exceeded → give_up. + # A skip is recorded like any other result but is not an attempt: + # a session queued behind a blocked one may not be given up on for + # work it never got to do. + attempts = last_phase.attempts if last_phase is not None else 0 + was_skipped = last_phase is not None and last_phase.result == SchedulingResult.SKIPPED + if not was_skipped and policy.is_retry_exhausted(attempts): give_up_failures.append(failure) continue # 2. Check timeout exceeded → expired - if policy.is_timed_out(session.phase_started_at, current_time): + started_at = last_phase.started_at if last_phase is not None else None + if policy.is_timed_out(started_at, current_time): expired_failures.append(failure) continue diff --git a/src/ai/backend/manager/sokovan/scheduler/handlers/lifecycle/schedule_sessions.py b/src/ai/backend/manager/sokovan/scheduler/handlers/lifecycle/schedule_sessions.py index cb82cb88ec2..be373b1d083 100644 --- a/src/ai/backend/manager/sokovan/scheduler/handlers/lifecycle/schedule_sessions.py +++ b/src/ai/backend/manager/sokovan/scheduler/handlers/lifecycle/schedule_sessions.py @@ -116,6 +116,10 @@ async def execute( - successes: Sessions that were scheduled - failures: Sessions whose scheduling attempt failed in the Provisioner - skipped: Sessions that were not attempted (priority-based, resource constraints) + + A session left unattempted must stay out of `failures`: the retry + pressure that eventually deprioritizes a session may only be charged + to attempts that actually happened. """ result = SessionExecutionResult() @@ -141,6 +145,7 @@ async def execute( failure_map = { failure.session_id: failure for failure in schedule_result.scheduling_failures } + skip_map = {skip.session_id: skip for skip in schedule_result.scheduling_skips} # Reservation-backed placements: the sessions hold their resources # (kernels already RESERVED) and wait for their victims; the victims @@ -180,6 +185,9 @@ async def execute( elif session_id in failure_map: reason = failure_map[session_id].msg or "scheduling-failed" result.failures.append(self._to_transition_info(session, reason)) + elif session_id in skip_map: + reason = skip_map[session_id].msg or "not-attempted-this-cycle" + result.skipped.append(self._to_transition_info(session, reason)) else: result.skipped.append(self._to_transition_info(session, "not-scheduled-this-cycle")) diff --git a/src/ai/backend/manager/sokovan/scheduler/provisioner/provisioner.py b/src/ai/backend/manager/sokovan/scheduler/provisioner/provisioner.py index fc8023fb475..b7318a20874 100644 --- a/src/ai/backend/manager/sokovan/scheduler/provisioner/provisioner.py +++ b/src/ai/backend/manager/sokovan/scheduler/provisioner/provisioner.py @@ -19,7 +19,11 @@ from ai.backend.manager.sokovan.recorder import ( RecorderContext, ) -from ai.backend.manager.sokovan.scheduler.results import PreemptionPlanEntry, ScheduleResult +from ai.backend.manager.sokovan.scheduler.results import ( + PreemptionPlanEntry, + ScheduleResult, + SchedulingSkip, +) from ai.backend.manager.views.sokovan.allocation import ( KernelAllocation, SchedulingFailure, @@ -30,6 +34,7 @@ from ai.backend.manager.views.sokovan.snapshot import ScopeVictimCandidates, SystemSnapshot from ai.backend.manager.views.sokovan.workload import SessionWorkload +from .selectors.exceptions import BatchAgentSelectionFailedError, NoAvailableAgentError from .selectors.selector import ( AgentSelection, AgentSelectionCriteria, @@ -136,6 +141,10 @@ async def schedule_resource_group( 3. Agent selection: Select agents using configured strategy 4. Allocation: Persist allocations to database + Once a session fails because the group ran out of resources, the + sessions behind it are left unattempted and reported as skips, so + lower-priority sessions cannot take the resources it is waiting for. + Args: scheduling_data: Pre-fetched scheduling data from Handler. @@ -172,8 +181,9 @@ async def schedule_resource_group( reserved_allocations: list[SessionAllocation] = [] claimed_victim_ids: set[SessionId] = set() scheduling_failures: list[SchedulingFailure] = [] + scheduling_skips: list[SchedulingSkip] = [] - for session_workload in sequenced_workloads: + for index, session_workload in enumerate(sequenced_workloads): try: # Sequencing phase is automatically included via shared phases session_allocation = await self._schedule_workload( @@ -186,33 +196,43 @@ async def schedule_resource_group( claimed_victim_ids.update(session_allocation.preempting_session_ids) else: session_allocations.append(session_allocation) - except Exception as e: - log.debug( - "Scheduling failed for workload {}: {}", - session_workload.meta.session_id, - e, - ) - scheduling_failures.append( - SchedulingFailure( - session_id=session_workload.meta.session_id, - msg=str(e), + except (BatchAgentSelectionFailedError, NoAvailableAgentError) as e: + # The group ran out of resources. This is the one failure the + # sessions behind share: they would take exactly what this one + # is waiting for, so they are left unattempted. + self._record_failure(scheduling_failures, session_workload, e) + scheduling_skips.extend( + self._skips_behind( + sequenced_workloads[index + 1 :], session_workload.meta.session_id ) ) - continue + break + except Exception as e: + # Specific to the session that hit it (an architecture mismatch, + # an unsatisfied dependency, an exceeded quota, a group with no + # agents at all), so it must not stall the sessions behind it. + self._record_failure(scheduling_failures, session_workload, e) log.info( - "Processing {} allocations, {} reservations and {} failures in resource group {}", + "Processing {} allocations, {} reservations, {} failures" + " and {} skips in resource group {}", len(session_allocations), len(reserved_allocations), len(scheduling_failures), + len(scheduling_skips), resource_group_id, ) with self._phase_metrics.measure_phase("scheduler", resource_group_id, "allocation"): scheduled_session_ids = await self._repository.allocate_sessions(session_allocations) reserved_session_ids = await self._repository.reserve_sessions(reserved_allocations) - failure_ids = [f.session_id for f in scheduling_failures] - await self._valkey_schedule.set_pending_queue(state.resource_group.name, failure_ids) + # The pending queue is the whole still-waiting queue in sequencing + # order (the legacy GraphQL resolver re-sorts by this order), so the + # unattempted sessions belong in it just as much as the failed ones. + pending_ids = [f.session_id for f in scheduling_failures] + [ + s.session_id for s in scheduling_skips + ] + await self._valkey_schedule.set_pending_queue(state.resource_group.name, pending_ids) reserved_ids = set(reserved_session_ids) return ScheduleResult( scheduled_session_ids=scheduled_session_ids, @@ -226,8 +246,44 @@ async def schedule_resource_group( for allocation in reserved_allocations if allocation.session_id in reserved_ids ], + scheduling_skips=scheduling_skips, ) + def _record_failure( + self, + scheduling_failures: list[SchedulingFailure], + session_workload: SessionWorkload, + error: Exception, + ) -> None: + log.debug( + "Scheduling failed for workload {}: {}", + session_workload.meta.session_id, + error, + ) + scheduling_failures.append( + SchedulingFailure( + session_id=session_workload.meta.session_id, + msg=str(error), + ) + ) + + def _skips_behind( + self, + remaining_workloads: Sequence[SessionWorkload], + blocking_session_id: SessionId, + ) -> list[SchedulingSkip]: + """The unattempted tail of the queue, in sequencing order.""" + return [ + SchedulingSkip( + session_id=workload.meta.session_id, + msg=( + "Not attempted: the resource group ran out of resources at the" + f" earlier session {blocking_session_id}" + ), + ) + for workload in remaining_workloads + ] + async def _schedule_workload( self, state: SchedulingState, diff --git a/src/ai/backend/manager/sokovan/scheduler/results.py b/src/ai/backend/manager/sokovan/scheduler/results.py index 07b76d647b6..af979abf8a0 100644 --- a/src/ai/backend/manager/sokovan/scheduler/results.py +++ b/src/ai/backend/manager/sokovan/scheduler/results.py @@ -27,6 +27,19 @@ class PreemptionPlanEntry: victim_session_ids: tuple[SessionId, ...] +@dataclass +class SchedulingSkip: + """A session that was not attempted this pass. + + Distinct from a :class:`SchedulingFailure`: nothing was tried, so no + retry pressure (and no eventual deprioritization) may be charged to it. + """ + + session_id: SessionId + # Human-readable skip reason (transition reason / pending queue) + msg: str + + @dataclass class ScheduleResult: """Result of a scheduling operation.""" @@ -40,6 +53,9 @@ class ScheduleResult: reserved_session_ids: list[SessionId] # The preemption plans backing those reservations. preemption_plan: list[PreemptionPlanEntry] + # Sessions left unattempted because an earlier session exhausted the + # group's resources, in sequencing order. + scheduling_skips: list[SchedulingSkip] = field(default_factory=list) def success_count(self) -> int: """Get the count of successfully scheduled sessions.""" diff --git a/src/ai/backend/manager/views/sokovan/lifecycle.py b/src/ai/backend/manager/views/sokovan/lifecycle.py index d8096640884..5135706ebe6 100644 --- a/src/ai/backend/manager/views/sokovan/lifecycle.py +++ b/src/ai/backend/manager/views/sokovan/lifecycle.py @@ -21,7 +21,7 @@ ) from ai.backend.manager.data.kernel.types import KernelInfo, KernelStatus from ai.backend.manager.data.network.types import NetworkType -from ai.backend.manager.data.session.types import SessionInfo +from ai.backend.manager.data.session.types import SchedulingResult, SessionInfo from ai.backend.manager.defs import DEFAULT_ROLE from ai.backend.manager.errors.kernel import MainKernelNotFound, TooManyKernelsFound @@ -300,6 +300,26 @@ class SessionRunningData: occupying_slots: ResourceSlot +@dataclass(frozen=True) +class LastPhase: + """The session's last scheduling-history record of the phase in progress. + + Absent when the session has no record of that phase yet. Read by the + coordinator's failure classification: ``attempts`` against the retry + budget, ``started_at`` against the timeout, and ``result`` to tell an + attempt from a skip. + + Attributes: + attempts: How many times the phase was recorded, skips included + started_at: When the phase was first recorded + result: What the record ended in + """ + + attempts: int + started_at: datetime + result: SchedulingResult + + @dataclass class SessionWithKernels: """ @@ -311,16 +331,13 @@ class SessionWithKernels: Attributes: session_info: Session information including lifecycle data kernel_infos: List of kernels belonging to this session - phase_attempts: Number of attempts for current phase from scheduling history - (used for failure classification: give_up when >= max_retries) - phase_started_at: When the current phase started from scheduling history - (used for failure classification: expired when timeout exceeded) + last_phase: The session's last record of the phase being processed, + or None when it has none yet """ session_info: SessionInfo kernel_infos: list[KernelInfo] - phase_attempts: int = 0 - phase_started_at: datetime | None = None + last_phase: LastPhase | None = None @property def main_kernel(self) -> KernelInfo: diff --git a/tests/unit/manager/models/test_scheduling_history_row.py b/tests/unit/manager/models/test_scheduling_history_row.py new file mode 100644 index 00000000000..86cf97500e7 --- /dev/null +++ b/tests/unit/manager/models/test_scheduling_history_row.py @@ -0,0 +1,85 @@ +"""Tests for the scheduling-history merge rule.""" + +from __future__ import annotations + +import uuid + +import pytest + +from ai.backend.common.types import SessionId +from ai.backend.manager.data.session.types import SchedulingResult, SessionStatus +from ai.backend.manager.models.scheduling_history.row import SessionSchedulingHistoryRow + +SESSION_ID = SessionId(uuid.uuid4()) + + +def _make_history( + *, + result: SchedulingResult, + phase: str = "schedule-sessions", + to_status: SessionStatus = SessionStatus.PENDING, + error_code: str | None = None, +) -> SessionSchedulingHistoryRow: + return SessionSchedulingHistoryRow( + session_id=SESSION_ID, + phase=phase, + from_status=str(SessionStatus.PENDING), + to_status=str(to_status), + result=str(result), + error_code=error_code, + message="", + sub_steps=[], + attempts=1, + ) + + +class TestSessionSchedulingHistoryMerge: + """``attempts`` drives the give-up (deprioritization) classification, so + what merges into one record decides what counts as a retry.""" + + @pytest.fixture + def failure(self) -> SessionSchedulingHistoryRow: + return _make_history(result=SchedulingResult.NEED_RETRY) + + def test_same_result_merges(self, failure: SessionSchedulingHistoryRow) -> None: + assert failure.should_merge_with(_make_history(result=SchedulingResult.NEED_RETRY)) + + @pytest.mark.parametrize( + "later_result", + [SchedulingResult.GIVE_UP, SchedulingResult.EXPIRED, SchedulingResult.SUCCESS], + ) + def test_another_attempt_result_still_merges( + self, + failure: SessionSchedulingHistoryRow, + later_result: SchedulingResult, + ) -> None: + """Which attempt result was recorded is not part of the merge key. + + Handlers that declare no transition for give_up/expired record them + with an unchanged ``to_status``; splitting those off would reset the + retry budget every time give_up fires, so it could never stick. + """ + assert failure.should_merge_with(_make_history(result=later_result)) + + def test_skip_does_not_merge_into_an_attempt( + self, failure: SessionSchedulingHistoryRow + ) -> None: + """A session that was never tried may not inflate the retry counter.""" + assert not failure.should_merge_with(_make_history(result=SchedulingResult.SKIPPED)) + + def test_attempt_does_not_merge_into_a_skip(self) -> None: + skipped = _make_history(result=SchedulingResult.SKIPPED) + + assert not skipped.should_merge_with(_make_history(result=SchedulingResult.NEED_RETRY)) + + def test_different_phase_does_not_merge(self, failure: SessionSchedulingHistoryRow) -> None: + assert not failure.should_merge_with( + _make_history(result=SchedulingResult.NEED_RETRY, phase="start-sessions") + ) + + def test_different_error_code_does_not_merge( + self, failure: SessionSchedulingHistoryRow + ) -> None: + assert not failure.should_merge_with( + _make_history(result=SchedulingResult.NEED_RETRY, error_code="E-1") + ) diff --git a/tests/unit/manager/repositories/scheduler/test_update_with_history.py b/tests/unit/manager/repositories/scheduler/test_update_with_history.py index 4cb10d720fe..695f04e2421 100644 --- a/tests/unit/manager/repositories/scheduler/test_update_with_history.py +++ b/tests/unit/manager/repositories/scheduler/test_update_with_history.py @@ -645,7 +645,11 @@ async def test_update_with_history_merge_same_phase_error_to_status( db_with_cleanup: ExtendedAsyncSAEngine, test_session_id: SessionId, ) -> None: - """Test that repeated calls with same phase+error_code+to_status merge (increment attempts).""" + """Test that repeated calls with same phase+error_code+to_status merge (increment attempts). + + Neither ``from_status`` nor which attempt result was recorded is part + of the merge key — only attempt-vs-skip is (see the skip test below). + """ db_source = ScheduleDBSource(db_with_cleanup) # First call - creates history record @@ -755,6 +759,70 @@ async def test_update_with_history_merge_same_phase_error_to_status( assert len(records) == 1 assert records[0].attempts == 3 + async def test_update_with_history_no_merge_skipped_after_failure( + self, + db_with_cleanup: ExtendedAsyncSAEngine, + test_session_id: SessionId, + ) -> None: + """Skips are counted, but on their own record. + + Skips must be visible and countable, while ``attempts`` on the + attempt record stays the number the give-up (deprioritization) + classification is allowed to see. + """ + db_source = ScheduleDBSource(db_with_cleanup) + + updater = BatchUpdater( + spec=SessionStatusBatchUpdaterSpec( + to_status=SessionStatus.PENDING, + status_changed_at=datetime.now(tzutc()), + reason="attempted", + ), + conditions=[lambda: SessionRow.id.in_([test_session_id])], + ) + await db_source.update_with_history( + updater, + BulkCreator( + specs=[ + SessionSchedulingHistoryCreatorSpec( + session_id=test_session_id, + phase="schedule", + result=SchedulingResult.NEED_RETRY, + message="No resources available", + to_status=SessionStatus.PENDING, + ) + ] + ), + ) + + # Same phase and to_status, but nothing was attempted these two cycles + for _ in range(2): + await db_source.update_with_history( + updater, + BulkCreator( + specs=[ + SessionSchedulingHistoryCreatorSpec( + session_id=test_session_id, + phase="schedule", + result=SchedulingResult.SKIPPED, + message="Not attempted", + to_status=SessionStatus.PENDING, + ) + ] + ), + ) + + async with db_with_cleanup.begin_readonly_session() as db_sess: + history_stmt = sa.select(SessionSchedulingHistoryRow).where( + SessionSchedulingHistoryRow.session_id == test_session_id + ) + records = (await db_sess.execute(history_stmt)).scalars().all() + attempts_by_result = {r.result: r.attempts for r in records} + # The skips are counted on their own record... + assert attempts_by_result[str(SchedulingResult.SKIPPED)] == 2 + # ...and the attempt record keeps the count give-up may see + assert attempts_by_result[str(SchedulingResult.NEED_RETRY)] == 1 + async def test_update_with_history_no_merge_different_phase( self, db_with_cleanup: ExtendedAsyncSAEngine, diff --git a/tests/unit/manager/sokovan/scheduler/handlers/conftest.py b/tests/unit/manager/sokovan/scheduler/handlers/conftest.py index f8beb47586f..1bc5e06f6aa 100644 --- a/tests/unit/manager/sokovan/scheduler/handlers/conftest.py +++ b/tests/unit/manager/sokovan/scheduler/handlers/conftest.py @@ -87,8 +87,6 @@ def _create_session( kernel_status: KernelStatus = KernelStatus.PENDING, session_type: SessionTypes = SessionTypes.INTERACTIVE, cluster_mode: ClusterMode = ClusterMode.SINGLE_NODE, - phase_attempts: int = 0, - phase_started_at: datetime | None = None, ) -> SessionWithKernels: """Create SessionWithKernels with sensible defaults.""" sid = session_id or SessionId(uuid4()) @@ -238,8 +236,6 @@ def _create_session( return SessionWithKernels( session_info=session_info, kernel_infos=kernel_infos, - phase_attempts=phase_attempts, - phase_started_at=phase_started_at, ) diff --git a/tests/unit/manager/sokovan/scheduler/handlers/test_lifecycle_handlers.py b/tests/unit/manager/sokovan/scheduler/handlers/test_lifecycle_handlers.py index 10d1f085191..504731475e8 100644 --- a/tests/unit/manager/sokovan/scheduler/handlers/test_lifecycle_handlers.py +++ b/tests/unit/manager/sokovan/scheduler/handlers/test_lifecycle_handlers.py @@ -31,7 +31,7 @@ from ai.backend.manager.sokovan.scheduler.handlers.lifecycle.terminate_sessions import ( TerminateSessionsLifecycleHandler, ) -from ai.backend.manager.sokovan.scheduler.results import ScheduleResult +from ai.backend.manager.sokovan.scheduler.results import ScheduleResult, SchedulingSkip from ai.backend.manager.views.sokovan.allocation import SchedulingFailure from ai.backend.manager.views.sokovan.lifecycle import ( SessionsForPullWithImages, @@ -254,6 +254,55 @@ async def test_mixed_results_categorized_correctly( for skipped in result.skipped: assert skipped.reason == "not-scheduled-this-cycle" + async def test_unattempted_sessions_reported_as_skipped( + self, + handler: ScheduleSessionsLifecycleHandler, + mock_provisioner: AsyncMock, + mock_repository: AsyncMock, + pending_sessions_multiple: list[SessionWithKernels], + ) -> None: + """SC-SS-008: Sessions left unattempted behind a blocked one are skipped. + + Given: Multiple PENDING sessions in the scaling group + When: Provisioner fails the first on exhausted resources and reports + the rest as skips + Then: Only the attempted session is a failure; the rest are skipped + with the provisioner's reason (no retry pressure charged) + """ + # Arrange + blocked_session, *rest = pending_sessions_multiple + mock_repository.get_scheduling_data.return_value = MagicMock() + mock_provisioner.schedule_resource_group.return_value = ScheduleResult( + scheduled_session_ids=[], + scheduling_failures=[ + SchedulingFailure( + session_id=blocked_session.session_info.identity.id, + msg="no agents can be allocated", + ) + ], + reserved_session_ids=[], + preemption_plan=[], + scheduling_skips=[ + SchedulingSkip( + session_id=session.session_info.identity.id, + msg="not attempted: resources exhausted", + ) + for session in rest + ], + ) + + # Act + result = await handler.execute(ResourceGroupID(uuid.uuid4()), pending_sessions_multiple) + + # Assert + assert len(result.successes) == 0 + assert [f.session_id for f in result.failures] == [blocked_session.session_info.identity.id] + assert {s.session_id for s in result.skipped} == { + session.session_info.identity.id for session in rest + } + for skipped in result.skipped: + assert skipped.reason == "not attempted: resources exhausted" + async def test_empty_session_list_returns_empty_result( self, handler: ScheduleSessionsLifecycleHandler, diff --git a/tests/unit/manager/sokovan/scheduler/provisioner/conftest.py b/tests/unit/manager/sokovan/scheduler/provisioner/conftest.py index 3c0b20ee863..6ae2c20beae 100644 --- a/tests/unit/manager/sokovan/scheduler/provisioner/conftest.py +++ b/tests/unit/manager/sokovan/scheduler/provisioner/conftest.py @@ -51,6 +51,7 @@ from ai.backend.manager.views.sokovan.workload import ( KernelWorkload, ResourceRequest, + SessionDependencyInfo, SessionGroupPolicy, SessionPlacement, SessionWorkload, @@ -72,6 +73,9 @@ def _make_workload( priority: int = 0, cluster_mode: ClusterMode = ClusterMode.SINGLE_NODE, session_group: SessionGroupPolicy | None = None, + architecture: str = "x86_64", + session_type: SessionTypes = SessionTypes.INTERACTIVE, + requested_starts_at: datetime | None = None, ) -> SessionWorkload: if kernel_slots is None: kernel_slots = [{"cpu": "1", "mem": "1024"}] @@ -91,7 +95,7 @@ def _make_workload( kernels=[ KernelWorkload( kernel_id=KernelId(uuid.uuid4()), - architecture=ArchName("x86_64"), + architecture=ArchName(architecture), requested_slots=ResourceRequest( slots={ ResourceSlotName(name): Decimal(amount) @@ -107,8 +111,8 @@ def _make_workload( ), priority=priority, job_priority=0, - session_type=SessionTypes.INTERACTIVE, - requested_starts_at=None, + session_type=session_type, + requested_starts_at=requested_starts_at, is_preemptible=False, ) @@ -141,16 +145,23 @@ def _make_scheduling_data( agents: list[AgentMeta] | None = None, scheduler: str = "fifo", agent_selection_strategy: AgentSelectionStrategy = AgentSelectionStrategy.CONCENTRATED, + session_dependencies: Mapping[SessionId, list[SessionDependencyInfo]] | None = None, + resource_policy: ResourcePolicySnapshot | None = None, + max_container_count: int | None = None, ) -> SchedulingData: if agents is None: agents = [_make_agent_meta()] + if resource_policy is None: + resource_policy = ResourcePolicySnapshot(by_user={}, by_project={}, by_domain={}) return SchedulingData( resource_group=ResourceGroupMeta(id=RESOURCE_GROUP_ID, name=RESOURCE_GROUP_NAME), workloads=workloads, system_snapshot=SystemSnapshot( resource_group=ResourceGroupScopeSnapshot( resources=ResourceGroupResource(agents=agents), - session_dependencies=SessionDependencySnapshot(by_session={}), + session_dependencies=SessionDependencySnapshot( + by_session=session_dependencies or {} + ), policy=ResourceGroupSchedulingPolicy( scheduler=scheduler, agent_selection_strategy=agent_selection_strategy, @@ -159,8 +170,8 @@ def _make_scheduling_data( ), global_scope=GlobalScopeSnapshot( occupancy=ResourceOccupancySnapshot(by_user={}, by_project={}, by_domain={}), - resource_policy=ResourcePolicySnapshot(by_user={}, by_project={}, by_domain={}), - agent_limit=AgentLimit(max_container_count=None), + resource_policy=resource_policy, + agent_limit=AgentLimit(max_container_count=max_container_count), ), observed_at=datetime.now(UTC), ), diff --git a/tests/unit/manager/sokovan/scheduler/provisioner/test_provisioner.py b/tests/unit/manager/sokovan/scheduler/provisioner/test_provisioner.py index 895c2f2d61a..2e19858f7aa 100644 --- a/tests/unit/manager/sokovan/scheduler/provisioner/test_provisioner.py +++ b/tests/unit/manager/sokovan/scheduler/provisioner/test_provisioner.py @@ -3,11 +3,15 @@ from __future__ import annotations from collections.abc import Sequence +from datetime import UTC, datetime, timedelta +from decimal import Decimal from unittest.mock import AsyncMock, MagicMock import pytest -from ai.backend.common.types import AgentId, SessionId +from ai.backend.common.identifier.resource_slot import ResourceSlotName +from ai.backend.common.types import AgentId, SessionId, SessionResult, SessionTypes +from ai.backend.manager.data.session.types import SessionStatus from ai.backend.manager.sokovan.recorder import RecorderContext from ai.backend.manager.sokovan.scheduler.provisioner.provisioner import ( SchedulingState, @@ -19,13 +23,23 @@ from ai.backend.manager.sokovan.scheduler.provisioner.validators.dependencies import ( DependenciesValidator, ) +from ai.backend.manager.sokovan.scheduler.provisioner.validators.reserved_batch import ( + ReservedBatchSessionValidator, +) +from ai.backend.manager.sokovan.scheduler.provisioner.validators.resource_policy import ( + ResourcePolicyValidator, +) from ai.backend.manager.sokovan.scheduler.provisioner.validators.validator import ( SchedulingValidator, ) from ai.backend.manager.sokovan.scheduler.results import ScheduleResult from ai.backend.manager.views.sokovan.allocation import SessionAllocation from ai.backend.manager.views.sokovan.scheduling import SchedulingData -from ai.backend.manager.views.sokovan.workload import SessionWorkload +from ai.backend.manager.views.sokovan.snapshot import ( + ResourcePolicySnapshot, + UserResourceLimit, +) +from ai.backend.manager.views.sokovan.workload import SessionDependencyInfo, SessionWorkload from .conftest import ( RESOURCE_GROUP_NAME, @@ -43,7 +57,11 @@ def _make_provisioner( config_provider.config.manager.agent_selection_resource_priority = ["cpu", "mem"] return SessionProvisioner( SessionProvisionerArgs( - validator=SchedulingValidator([DependenciesValidator()]), + validator=SchedulingValidator([ + DependenciesValidator(), + ReservedBatchSessionValidator(), + ResourcePolicyValidator(), + ]), default_sequencer=FIFOSequencer(), agent_selector=create_agent_selector(["cpu", "mem"]), repository=repository, @@ -199,7 +217,7 @@ async def test_partial_failure_keeps_other_sessions( agent_meta_factory: AgentMetaFactory, scheduling_data_factory: SchedulingDataFactory, ) -> None: - """One session failing does not abort the other sessions of the pass.""" + """A failure does not undo the sessions already scheduled in the pass.""" fitting = workload_factory(kernel_slots=[{"cpu": "2", "mem": "2048"}]) too_big = workload_factory(kernel_slots=[{"cpu": "100", "mem": "999999"}]) data = scheduling_data_factory( @@ -231,3 +249,179 @@ async def test_in_batch_occupancy_blocks_later_sessions( assert result.scheduled_session_ids == [first.meta.session_id] assert [f.session_id for f in result.scheduling_failures] == [second.meta.session_id] + + +class TestQueueBlockingOnResourceExhaustion: + """Only a resource-exhausted session holds the queue back. + + Every other failure belongs to the session that hit it, so the sessions + behind it must still be attempted. + """ + + async def test_exhausted_resources_skip_lower_priority_sessions( + self, + provisioner: SessionProvisioner, + valkey_schedule: AsyncMock, + workload_factory: WorkloadFactory, + agent_meta_factory: AgentMetaFactory, + scheduling_data_factory: SchedulingDataFactory, + ) -> None: + """A small low-priority session may not take what a big high-priority + session is waiting for.""" + big = workload_factory(kernel_slots=[{"cpu": "8", "mem": "16384"}], priority=10) + small = workload_factory(kernel_slots=[{"cpu": "1", "mem": "1024"}], priority=1) + data = scheduling_data_factory( + workloads=[small, big], + agents=[agent_meta_factory("agent-1", {"cpu": "4", "mem": "8192"})], + ) + + result = await _schedule(provisioner, data, [small, big]) + + assert result.scheduled_session_ids == [] + assert [f.session_id for f in result.scheduling_failures] == [big.meta.session_id] + assert [s.session_id for s in result.scheduling_skips] == [small.meta.session_id] + assert result.scheduling_skips[0].msg + # Failed and skipped sessions alike stay in the pending queue, in + # sequencing order + valkey_schedule.set_pending_queue.assert_awaited_once_with( + RESOURCE_GROUP_NAME, [big.meta.session_id, small.meta.session_id] + ) + + async def test_container_limit_exhaustion_skips_later_sessions( + self, + provisioner: SessionProvisioner, + workload_factory: WorkloadFactory, + agent_meta_factory: AgentMetaFactory, + scheduling_data_factory: SchedulingDataFactory, + ) -> None: + """The per-agent container cap is a resource too: it blocks the queue.""" + first = workload_factory(priority=10) + second = workload_factory(priority=1) + data = scheduling_data_factory( + workloads=[first, second], + agents=[agent_meta_factory("agent-1")], + max_container_count=0, + ) + + result = await _schedule(provisioner, data, [first, second]) + + assert result.scheduled_session_ids == [] + assert [f.session_id for f in result.scheduling_failures] == [first.meta.session_id] + assert [s.session_id for s in result.scheduling_skips] == [second.meta.session_id] + + async def test_incompatible_architecture_does_not_block_later_sessions( + self, + provisioner: SessionProvisioner, + workload_factory: WorkloadFactory, + scheduling_data_factory: SchedulingDataFactory, + ) -> None: + """A session no agent can ever host would block the queue forever.""" + foreign = workload_factory(architecture="aarch64", priority=10) + fitting = workload_factory(priority=1) + data = scheduling_data_factory(workloads=[foreign, fitting]) + + result = await _schedule(provisioner, data, [foreign, fitting]) + + assert result.scheduled_session_ids == [fitting.meta.session_id] + assert [f.session_id for f in result.scheduling_failures] == [foreign.meta.session_id] + assert result.scheduling_skips == [] + + async def test_empty_resource_group_does_not_block_later_sessions( + self, + provisioner: SessionProvisioner, + workload_factory: WorkloadFactory, + scheduling_data_factory: SchedulingDataFactory, + ) -> None: + """With no agents at all every session fails on its own account.""" + first = workload_factory(priority=10) + second = workload_factory(priority=1) + data = scheduling_data_factory(workloads=[first, second], agents=[]) + + result = await _schedule(provisioner, data, [first, second]) + + assert result.scheduled_session_ids == [] + assert [f.session_id for f in result.scheduling_failures] == [ + first.meta.session_id, + second.meta.session_id, + ] + assert result.scheduling_skips == [] + + async def test_reserved_batch_session_does_not_block_later_sessions( + self, + provisioner: SessionProvisioner, + workload_factory: WorkloadFactory, + scheduling_data_factory: SchedulingDataFactory, + ) -> None: + """A batch session starting in an hour must not stall the group for an hour.""" + later = workload_factory( + priority=10, + session_type=SessionTypes.BATCH, + requested_starts_at=datetime.now(UTC) + timedelta(hours=1), + ) + fitting = workload_factory(priority=1) + data = scheduling_data_factory(workloads=[later, fitting]) + + result = await _schedule(provisioner, data, [later, fitting]) + + assert result.scheduled_session_ids == [fitting.meta.session_id] + assert [f.session_id for f in result.scheduling_failures] == [later.meta.session_id] + assert result.scheduling_skips == [] + + async def test_unsatisfied_dependency_does_not_block_later_sessions( + self, + provisioner: SessionProvisioner, + workload_factory: WorkloadFactory, + scheduling_data_factory: SchedulingDataFactory, + ) -> None: + """The session it depends on may be the one queued behind it.""" + waiting = workload_factory(priority=10) + fitting = workload_factory(priority=1) + data = scheduling_data_factory( + workloads=[waiting, fitting], + session_dependencies={ + waiting.meta.session_id: [ + SessionDependencyInfo( + depends_on=fitting.meta.session_id, + dependency_name="upstream", + dependency_status=SessionStatus.RUNNING, + dependency_result=SessionResult.UNDEFINED, + ) + ] + }, + ) + + result = await _schedule(provisioner, data, [waiting, fitting]) + + assert result.scheduled_session_ids == [fitting.meta.session_id] + assert [f.session_id for f in result.scheduling_failures] == [waiting.meta.session_id] + assert result.scheduling_skips == [] + + async def test_quota_exceeded_does_not_block_other_users( + self, + provisioner: SessionProvisioner, + workload_factory: WorkloadFactory, + scheduling_data_factory: SchedulingDataFactory, + ) -> None: + """One user over quota must not stall the whole resource group.""" + over_quota = workload_factory(kernel_slots=[{"cpu": "2", "mem": "2048"}], priority=10) + fitting = workload_factory(priority=1) + data = scheduling_data_factory( + workloads=[over_quota, fitting], + resource_policy=ResourcePolicySnapshot( + by_user={ + over_quota.meta.owner.user_uuid: UserResourceLimit( + slots={ResourceSlotName("cpu"): Decimal(1)}, + max_session_count=None, + max_sftp_session_count=None, + ) + }, + by_project={}, + by_domain={}, + ), + ) + + result = await _schedule(provisioner, data, [over_quota, fitting]) + + assert result.scheduled_session_ids == [fitting.meta.session_id] + assert [f.session_id for f in result.scheduling_failures] == [over_quota.meta.session_id] + assert result.scheduling_skips == [] diff --git a/tests/unit/manager/sokovan/scheduler/test_coordinator.py b/tests/unit/manager/sokovan/scheduler/test_coordinator.py index 88b804197b1..040a2bd1205 100644 --- a/tests/unit/manager/sokovan/scheduler/test_coordinator.py +++ b/tests/unit/manager/sokovan/scheduler/test_coordinator.py @@ -29,6 +29,7 @@ StatusTransitions, TransitionStatus, ) +from ai.backend.manager.models.scheduling_history.row import SessionSchedulingHistoryRow from ai.backend.manager.repositories.scheduler.updaters import SessionStatusBatchUpdaterSpec from ai.backend.manager.sokovan.scheduler.coordinator import ( FailureClassificationResult, @@ -43,6 +44,7 @@ SessionExecutionResult, SessionTransitionInfo, ) +from ai.backend.manager.views.sokovan.lifecycle import LastPhase # ============================================================================= # Test Fixtures @@ -52,13 +54,25 @@ _TEST_HANDLER_NAME = "test-handler" +def _last_phase( + attempts: int, + started_at: datetime | None = None, + result: SchedulingResult = SchedulingResult.NEED_RETRY, +) -> LastPhase: + """The phase record as the coordinator carries it onto a session.""" + return LastPhase( + attempts=attempts, + started_at=started_at if started_at is not None else datetime.now(tzutc()), + result=result, + ) + + def _create_session_with_kernels( session_id: SessionId, status: SessionStatus = SessionStatus.PREPARING, - phase_attempts: int = 0, - phase_started_at: datetime | None = None, timeout: int | None = None, max_retry_count: int | None = None, + last_phase: LastPhase | None = None, ) -> MagicMock: """Create a mock SessionWithKernels with a stub ``handler_options``. @@ -69,8 +83,7 @@ def _create_session_with_kernels( mock = MagicMock() mock.session_info.identity.id = session_id mock.session_info.lifecycle.status = status - mock.phase_attempts = phase_attempts - mock.phase_started_at = phase_started_at + mock.last_phase = last_phase mock.kernel_infos = [] mock.session_info.handler_options.resolve.return_value = HandlerOptions( timeout=timeout, @@ -79,6 +92,26 @@ def _create_session_with_kernels( return mock +def _create_history_row( + phase: str, + result: SchedulingResult, + attempts: int, +) -> SessionSchedulingHistoryRow: + """Create a last-history record as the repository hands it to the coordinator.""" + return SessionSchedulingHistoryRow( + session_id=uuid4(), + phase=phase, + from_status=str(SessionStatus.PENDING), + to_status=str(SessionStatus.PENDING), + result=str(result), + error_code=None, + message="", + sub_steps=[], + attempts=attempts, + created_at=datetime.now(tzutc()), + ) + + def _create_session_transition_info( session_id: SessionId | None = None, from_status: SessionStatus = SessionStatus.PREPARING, @@ -104,7 +137,7 @@ class TestScheduleCoordinatorFailureClassification: The coordinator classifies failures into: - give_up: per-handler max_retry_count is set AND - phase_attempts >= max_retry_count + the last phase record's attempts reached it - expired: per-handler timeout is set AND elapsed > timeout - need_retry: default (can be retried) @@ -121,7 +154,7 @@ def test_give_up_on_max_attempts_exceeded(self) -> None: session = _create_session_with_kernels( session_id=session_id, status=SessionStatus.PREPARING, - phase_attempts=5, + last_phase=_last_phase(attempts=5), max_retry_count=5, # Limit reached ) @@ -140,6 +173,40 @@ def test_give_up_on_max_attempts_exceeded(self) -> None: assert len(result.need_retry) == 0 assert result.give_up[0].session_id == session_id + def test_no_give_up_when_the_attempts_are_skips(self) -> None: + """SC-CO-001b: A skipped session is not given up on, however many skips. + + A session queued behind a blocked one accumulates skips without ever + being attempted; charging those to the retry budget would deprioritize + it for work it never got to do. + """ + # Arrange + session_id = SessionId(uuid4()) + failure = _create_session_transition_info(session_id=session_id) + + session = _create_session_with_kernels( + session_id=session_id, + status=SessionStatus.PENDING, + last_phase=_last_phase( + attempts=99, result=SchedulingResult.SKIPPED + ), # all of them skips + max_retry_count=5, + ) + + # Act + result = ScheduleCoordinator._classify_failures( + None, # type: ignore[arg-type] + failures=[failure], + sessions=[session], + current_time=datetime.now(tzutc()), + handler_name=_TEST_HANDLER_NAME, + ) + + # Assert + assert len(result.give_up) == 0 + assert len(result.need_retry) == 1 + assert result.need_retry[0].session_id == session_id + def test_expired_on_timeout_exceeded(self) -> None: """SC-CO-002: Expire when timeout is set and exceeded.""" # Arrange @@ -153,8 +220,7 @@ def test_expired_on_timeout_exceeded(self) -> None: session = _create_session_with_kernels( session_id=session_id, status=SessionStatus.PREPARING, - phase_attempts=0, - phase_started_at=past_time, + last_phase=_last_phase(attempts=0, started_at=past_time), timeout=900, # 15 minutes — past_time is 20 minutes ago ) @@ -183,8 +249,7 @@ def test_need_retry_on_retryable_failure(self) -> None: session = _create_session_with_kernels( session_id=session_id, status=SessionStatus.PREPARING, - phase_attempts=1, - phase_started_at=recent_time, + last_phase=_last_phase(attempts=1, started_at=recent_time), timeout=900, max_retry_count=5, ) @@ -214,8 +279,7 @@ def test_give_up_takes_priority_over_expired(self) -> None: session = _create_session_with_kernels( session_id=session_id, status=SessionStatus.PREPARING, - phase_attempts=5, - phase_started_at=past_time, + last_phase=_last_phase(attempts=5, started_at=past_time), timeout=900, max_retry_count=5, ) @@ -241,7 +305,7 @@ def test_mixed_classification_results(self) -> None: failure_1 = _create_session_transition_info(session_id=session_id_1) session_1 = _create_session_with_kernels( session_id=session_id_1, - phase_attempts=5, + last_phase=_last_phase(attempts=5), max_retry_count=5, ) @@ -252,8 +316,7 @@ def test_mixed_classification_results(self) -> None: session_2 = _create_session_with_kernels( session_id=session_id_2, status=SessionStatus.PREPARING, - phase_attempts=1, - phase_started_at=past_time, + last_phase=_last_phase(attempts=1, started_at=past_time), timeout=900, max_retry_count=5, ) @@ -264,8 +327,7 @@ def test_mixed_classification_results(self) -> None: recent_time = datetime.now(tzutc()) - timedelta(minutes=1) session_3 = _create_session_with_kernels( session_id=session_id_3, - phase_attempts=1, - phase_started_at=recent_time, + last_phase=_last_phase(attempts=1, started_at=recent_time), timeout=900, max_retry_count=5, ) @@ -324,7 +386,7 @@ def test_no_limits_means_always_need_retry(self) -> None: Given: Session whose handler_options.resolve() returns ``HandlerOptions(timeout=None, max_retry_count=None)``, - even with high phase_attempts and an old phase_started_at. + even with a high attempt count and an old phase start. Then: Always classified as need_retry. """ # Arrange @@ -338,8 +400,7 @@ def test_no_limits_means_always_need_retry(self) -> None: session = _create_session_with_kernels( session_id=session_id, status=SessionStatus.PENDING, - phase_attempts=999, - phase_started_at=past_time, + last_phase=_last_phase(attempts=999, started_at=past_time), timeout=None, max_retry_count=None, ) @@ -359,6 +420,96 @@ def test_no_limits_means_always_need_retry(self) -> None: assert len(result.need_retry) == 1 +# ============================================================================= +# TestScheduleCoordinatorPhaseHistory Tests (SC-CO-008b) +# ============================================================================= + + +class TestScheduleCoordinatorPhaseHistory: + """Tests for carrying the phase's retry pressure onto sessions. + + The record feeds the give-up and timeout classifications, so only the + current phase's may be carried over. + """ + + @pytest.fixture + def session_id(self) -> SessionId: + return SessionId(uuid4()) + + @pytest.fixture + def session(self, session_id: SessionId) -> MagicMock: + return _create_session_with_kernels(session_id=session_id) + + def test_attempts_carried_from_the_same_phase( + self, + session_id: SessionId, + session: MagicMock, + ) -> None: + """SC-CO-008b: A failed attempt of this phase carries its counter.""" + history = _create_history_row( + phase=_TEST_HANDLER_NAME, result=SchedulingResult.NEED_RETRY, attempts=3 + ) + + ScheduleCoordinator._populate_phase_history( + None, # type: ignore[arg-type] + sessions=[session], + history_map={session_id: history}, + handler_name=_TEST_HANDLER_NAME, + ) + + assert session.last_phase == LastPhase( + attempts=3, + started_at=history.created_at, + result=SchedulingResult.NEED_RETRY, + ) + + def test_skips_are_counted_and_marked( + self, + session_id: SessionId, + session: MagicMock, + ) -> None: + """SC-CO-008c: Skips are counted like any other record. + + ``LastPhase.result`` marks what was counted; excluding skips from + give_up is the classifier's job, not this one's. + """ + history = _create_history_row( + phase=_TEST_HANDLER_NAME, result=SchedulingResult.SKIPPED, attempts=9 + ) + + ScheduleCoordinator._populate_phase_history( + None, # type: ignore[arg-type] + sessions=[session], + history_map={session_id: history}, + handler_name=_TEST_HANDLER_NAME, + ) + + assert session.last_phase == LastPhase( + attempts=9, + started_at=history.created_at, + result=SchedulingResult.SKIPPED, + ) + + def test_other_phase_history_carries_no_attempts( + self, + session_id: SessionId, + session: MagicMock, + ) -> None: + """SC-CO-008d: Another phase's counter does not leak into this one.""" + history = _create_history_row( + phase="other-handler", result=SchedulingResult.NEED_RETRY, attempts=4 + ) + + ScheduleCoordinator._populate_phase_history( + None, # type: ignore[arg-type] + sessions=[session], + history_map={session_id: history}, + handler_name=_TEST_HANDLER_NAME, + ) + + assert session.last_phase is None + + # ============================================================================= # TestScheduleCoordinatorHookExecution Tests (SC-CO-009 ~ SC-CO-013) # ============================================================================= @@ -667,7 +818,7 @@ async def test_failure_classified_and_transitioned( session = _create_session_with_kernels( session_id=session_id, - phase_attempts=5, + last_phase=_last_phase(attempts=5), max_retry_count=5, # Will be classified as give_up )