Skip to content

fix(BA-XXXX): claim PREPARED sessions instead of locking every START pass - #15978

Draft
hhoikoo wants to merge 4 commits into
mainfrom
fix/start-sessions-claim
Draft

hhoikoo wants to merge 4 commits into
mainfrom
fix/start-sessions-claim

Conversation

@hhoikoo

@hhoikoo hhoikoo commented Oct 22, 2026 •

Copy link
Copy Markdown
Member

resolves #NNN (BA-XXXX)

BA-XXXX and #NNN are placeholders until the ticket exists.

The sokovan START handler no longer holds the global LOCKID_SOKOVAN_TARGET_CREATING lock across create_kernels. Each START pass claims its PREPARED sessions in one transaction before the RPC and launches only the sessions it claimed, so two passes cannot start the same session and a pass does not wait for another pass's kernel creation.

Problem

  • StartSessionsLifecycleHandler.lock_id returns LOCKID_SOKOVAN_TARGET_CREATING. coordinator._process_lifecycle_handler_schedule takes it once per pass and holds it while asyncio.gather runs every resource group's handler.execute(), which awaits create_kernels on the agents (START_SESSION_TIMEOUT_SEC per session).
  • The lock is the only thing that keeps two passes from starting the same session: the session and its kernels stay PREPARED for the whole RPC (the coordinator writes CREATING only after execute() returns; update_kernel_status_preparing/_creating do not match a PREPARED kernel).
  • While the lock is held, a session that reaches PREPARED in a later pass, in any resource group, cannot start until the slowest create_kernels in flight returns. With Docker kernels this is seconds. Where the create RPC is slow (a large image pulled inside the create, a VM-based runtime), it is minutes, and waiters hit lock_timeout and retry every cycle.
  • When the lock is lost while the holder still runs, a second pass starts the session again. That happened on a CoCo test rig: PostgreSQL dropped the lock connection after 60 s, and the queued pass dispatched create_kernels a second time for the same PREPARED session. A network blip or a PostgreSQL failover can lose the lock the same way.

Change

  • ScheduleDBSource.claim_sessions_for_start(session_ids, reason) (repositories/scheduler/db_source/db_source.py): one READ COMMITTED transaction moves the sessions that are still PREPARED to CREATING (UPDATE sessions ... WHERE status = PREPARED RETURNING id) and their PREPARED kernels to CREATING. The conditional session update is the claim: a second pass blocks on the session row until the first commits, rechecks it, finds it CREATING and leaves it out. Kernels move only for the winner, in the same transaction, so a cluster session is claimed all or nothing. The kernel UPDATE is shared with the existing update_kernels_to_creating_for_sessions through a private _update_kernels_to_creating.
  • ScheduleDBSource.release_sessions_for_start(session_ids): moves sessions that are still CREATING, and their CREATING kernels, back to PREPARED. A session that moved on (for example, terminated) keeps its status.
  • StartSessionsLifecycleHandler.execute: claims first, drops the sessions another pass claimed (not reported, so the coordinator records nothing for them in this pass), queries and launches only the claimed ones, and reports them as successes as before. If the session query fails or the pass is cancelled before the launch, it releases the claim and re-raises.
  • lock_id returns None. LOCKID_SOKOVAN_TARGET_CREATING stays in defs.py unused, like the other retired lock ids there.
  • Coordinator guard (opt-in): TransitionStatus gains except_session_statuses (default empty, so every other transition is unchanged), passed to SessionStatusBatchUpdater.except_statuses. START's success transition uses frozenset(SessionStatus) - {CREATING}; its expired and give_up transitions use frozenset(SessionStatus) - {PREPARED}: a failed session has had its claim released to PREPARED, and if another pass claimed it again meanwhile it is CREATING and belongs to that pass. _apply_transition resets kernels to PENDING only for the sessions the update moved, so a skipped session's kernels are not reset and their agents not cleared while it is being created; the sessions a guarded transition left in their status are logged at info. Reason: the session is CREATING during the RPC now, so promote-to-running (CHECK_CREATING_PROGRESS) can promote it to RUNNING as soon as its kernels report started, while the pass still waits for other sessions. The unguarded write would put it back to CREATING; a session terminated during the RPC would also be written back (that part already happened on main from PREPARED).
  • update_with_history returns the IDs of the sessions it moved and records history only for those (update, then history, in one write_ops transaction, as mark_sessions_status does). _apply_transition returns the moved session infos, and _handle_result broadcasts the success event only for them, so a session promoted during the pass gets no late CREATING event (the event is cached per session and re-sent to new subscribers) and no CREATING history row after its RUNNING row.
  • SchedulerRepository gains the two pass-through methods; the scheduler README's start phase describes the claim.

Provisional (design choices of this PR, not decided elsewhere)

  • Session-level claim (session PREPARED to CREATING before the RPC) instead of a kernel-only claim. A kernel-only claim keeps the coordinator untouched, but a manager crash or cancel mid-RPC would then leave the session PREPARED with CREATING or RUNNING kernels, which no handler selects. With the session CREATING, promote-to-running picks it up once the kernels start.
  • The guard is opt-in per transition, not a general "only from the handler's target statuses" rule: schedule, reschedule and preempt mark statuses inside execute().

Not changed / known gaps

  • At most once instead of at least once: if the manager dies or is cancelled mid-RPC and the agent never creates the kernels, the session stays CREATING with CREATING kernels. That is the same shape as main's failed-RPC case (session CREATING). Neither has a timeout for CREATING sessions; the user can terminate it. A repository retry after a claim whose commit reached the server but whose acknowledgement was lost leaves the same shape.
  • The same CREATING shape results if start_sessions_for_handler itself raises or is cancelled after the claim (for example at shutdown): the claim is released only for a failure before the launch. There is no CREATING sweeper.
  • Failures inside the launcher before the RPC (network setup, missing image) are unchanged: the launcher records the error in status_data and the session stays CREATING as before.
  • A START handler that reports failures must release their claim first (release_sessions_for_start; the status_transitions docstring says so). Main reports none; a hook that rejects a kernel before the RPC would add them.
  • Accepted trade-off for such a hook path: if the release of hook-rejected sessions fails (logged), a session on its final attempt stays CREATING, because expired and give_up skip a CREATING session.
  • Follow-up (no code here): a claimed-but-never-dispatched session (manager crash or cancel mid-start) stays CREATING, and there is no sweeper for the age of a CREATING session. feat(BA-XXXX): make the manager session-create timeout configurable #15738's start timeout is an in-pass wait_for around the RPC, not a sweeper.
  • Without the lock, START passes from the short timer, long timer and events can each have a create_kernels batch in flight at once; the remaining bounds are the agents' creation limits and the agent client pool. No other handler used this lock.
  • Post-processors and metrics still count a session the guard skipped (one extra CHECK_CREATING_PROGRESS mark).

Test

  • tests/unit/manager/repositories/scheduler/test_start_claim.py (real PostgreSQL):
    • two concurrent claims of one 3-kernel session: exactly one wins, session and all kernels CREATING;
    • a claim queued behind another transaction holding the session row (which then sets the session CREATING and commits) returns nothing and leaves the kernels PREPARED (a kernel-only claim fails this test);
    • a claimed session is not claimed again; a CREATING or TERMINATING session is not claimed;
    • release returns session and kernels to PREPARED and the session can be claimed again; release leaves a terminated session alone;
    • START's success transition through update_with_history: a claimed CREATING session gets its history row; a session moved to RUNNING or TERMINATING keeps its status, [] is returned, no start-sessions history row.
    • START's expired and give_up transitions through update_with_history: a session still PREPARED moves (PENDING / TERMINATING) with one history row; a session claimed again (CREATING) is not updated and gets no history row.
  • test_lifecycle_handlers.py SC-ST-006: two passes racing for one PREPARED session, the launcher is awaited once. SC-ST-007 (error and CancelledError): a pass that fails before the launch releases the claim and does not launch.
  • test_coordinator.py SC-CO-019d: START's expired and give_up transitions exclude CREATING but not PREPARED; a skipped session is not returned and its kernels are not reset; a moved one is (kernel reset on expired only). SC-CO-019e: a guarded PENDING transition resets the kernels of the moved session only. SC-CO-026: a START pass whose launcher is blocked on an event for session A, then a second START pass for session B with a process-local lock factory; the second pass finishes within 1 s while the first is still blocked. SC-CO-019c: START's success updater excludes every status but CREATING, and a skipped session is not returned for the broadcast.
  • Before the round-2 commit (719ce6f): SC-CO-019d assert <SessionStatus.CREATING> in frozenset(), SC-CO-019e expected await not found (reset called with both sessions), DB failure tests assert [SessionID(...)] == [] (the CREATING session was moved). After: pass.
  • Before the change: SC-ST-006 Expected start_sessions_for_handler to have been awaited once. Awaited 2 times.; SC-ST-007 Expected release_sessions_for_start to have been awaited once. Awaited 0 times.; SC-CO-026 TimeoutError (the second pass waits on the lock); SC-CO-019c fails on the updater's empty except_statuses; the DB tests fail with AttributeError (no claim method). After: all pass.
  • pants test tests/unit/manager/sokovan:: tests/unit/manager/repositories/scheduler:: tests/unit/manager/repositories/scheduling_history:: passes (108 targets); pants lint check --changed-since=origin/main --changed-dependents=direct clean (ruff, visibility, mypy).
  • Not run on a live cluster.

Related

  • None. This PR does not depend on another PR.

Checklist: (if applicable)

  • Backport targets: a fix: PR goes to every version in .github/maintained-versions.yml automatically.
    A Backport: line here names the targets instead — some of them, other ones, or none.
  • Mention to the original issue
  • Installer updates including:
    • Fixtures for db schema changes
    • New mandatory config options
  • Update of end-to-end CLI integration tests in ai.backend.test
  • API server-client counterparts (e.g., manager API -> client SDK)
  • Test case(s) to:
    • Demonstrate the difference of before/after
    • Demonstrate the flow of abstract/conceptual models with a concrete implementation
  • Documentation
    • Contents in the docs directory
    • docstrings in public interfaces and type annotations

@github-actions github-actions Bot added size:XL 500~ LoC comp:manager Related to Manager component labels Oct 22, 2026
…pass

The START handler held the global LOCKID_SOKOVAN_TARGET_CREATING lock across
create_kernels, so a slow kernel creation made every later START pass wait,
and a lost lock let a second pass start the same session again. Each pass now
claims its sessions in one transaction (PREPARED sessions and their PREPARED
kernels to CREATING) and launches only what it claimed; the handler takes no
lock. START's success transition leaves a session that moved on during the
pass (promoted to RUNNING, terminated) alone. A pass that fails before the
launch moves the sessions back to PREPARED.
…ved on during START

START's success transition now skips a session promoted to RUNNING or
terminated while the pass waited for create_kernels. update_with_history
returns the ids it moved and records history only for them, and the
coordinator broadcasts only those.
…that pass

START's expired and give_up transitions now move a session only while it is
still PREPARED; a session another pass claimed again after its release is
CREATING and is skipped. The coordinator resets kernels to PENDING only for
the sessions a transition moved, so a skipped session's kernels are not reset
and their agents not cleared while it is being created, and it logs the
sessions a guarded transition left in their status.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

comp:manager Related to Manager component size:XL 500~ LoC

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant