Skip to content

fix(agent): save the kernel registry from a snapshot, not the live map - #14746

Open
daemyung-lablup wants to merge 6 commits into
mainfrom
fix/agent-registry-save-snapshot
Open

daemyung-lablup wants to merge 6 commits into
mainfrom
fix/agent-registry-save-snapshot

Conversation

@daemyung-lablup

Copy link
Copy Markdown
Contributor

save_kernel_registry iterates the live kernel_registry map and awaits once per kernel. A kernel created on the same node in between registers itself into that map, and the save raises RuntimeError: dictionary changed size during iteration inside the creating call -- which then fails the kernel it was creating. Measured with 27 two-node sessions started at once on three nodes: 1 of 27 failed on exactly this.

Change Why
Iterate a snapshot of the registry A kernel that arrives mid-save is written by the save its own create runs.
Commit a record only while the id still maps to the snapshot's object A restart re-creates a kernel under the same id and keeps its directory; the old object's REPL ports must not land there. The record is staged beside recovery.json and renamed into place in one step, with no await between the check and the rename -- which also means a crash mid-write cannot leave a torn record.
Never create the config directory from a save It is the kernel's own, removed with it; its absence means the kernel is gone, and re-making it left an empty directory behind. A failed or cancelled write removes its staged file.
Skip a PREPARING kernel quietly, keep the traceback for any other Registered a step before its REPL ports exist, a PREPARING kernel is ordinary under concurrent creates. A RUNNING kernel without them is a broken invariant.

Independent of #14551, which fixes the periodic reconciler destroying a just-created kernel; no file overlap.

Tests: tests/unit/agent/kernel_registry/test_writer_container.py (mid-save registration, same-id replacement, destroyed-while-saving, PREPARING vs RUNNING), tests/unit/agent/test_scratch_config.py (stage/commit/discard, no directory creation, cleanup on failed serialization and on cancellation)

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

The save awaits once per kernel while iterating the registry itself, and a
create on the same node registers its kernel in between -- `dictionary
changed size during iteration`, raised inside the creating call, which then
fails the kernel it was creating. Measured with 27 two-node sessions started
at once on three nodes: one of 27 failed on exactly this. Iterate over a
copy; a kernel that arrives mid-save is written by the save its own create
runs.

(cherry picked from commit 99b04f557afa987db00a66d9c3a66652fd4f4110)
A kernel is registered a step before its REPL ports exist, so a save that runs
in between cannot record it yet -- ordinary under concurrent creates, and its
own start writes it a moment later. It was logged as an exception with a full
traceback, which reads as a fault and buries the real ones: one line per
kernel per save, at warning, says what happened.
… a save

The other side of the same race: the save's snapshot still names a kernel
whose scratch directory a destroy has just removed, and writing its recovery
record re-created the directory -- an empty scratch left behind for a kernel
that is gone. The config directory is the kernel's own, made when it was
created and removed with it, so the save writes into it and treats its
absence as the kernel having left.
…ecord

The parse error covers every missing REPL key, and treating all of them as
the start-up race hid a broken invariant behind a warning: a kernel that is
PREPARING has no ports yet by design, a RUNNING one has them by contract.
Only the first is skipped quietly; the second keeps the exception log.
…that kernel

A restart destroys and re-creates a kernel under the same id and keeps its
scratch directory, so a save that snapshotted the old object could land the
old REPL ports in the new incarnation's directory, and the next agent restart
would recover it unreachable. The record is now staged beside `recovery.json`
and renamed into place only after the registry still maps the id to the
snapshot's object -- one atomic rename, no await between the check and it.
The rename also means a crash mid-write cannot leave a torn record.

Staging never creates the config directory, which replaces the
`create_config_dir` flag: the directory is the kernel's own, and its absence
means the kernel is gone.
Staging opened its temporary file before serializing, so a serialization or
write error -- or a cancelled save -- left `recovery.json.<uuid>.tmp` in the
kernel's directory with no handle to discard it. Serialize first; if opening
or writing fails for any reason, remove the file and re-raise.

The existing `ScratchConfig` tests still called the removed
`save_json_recovery_data`; they now cover the stage/commit contract, which
also takes in the scratch tests added under a separate directory.
Copilot AI balanced review requested due to automatic review settings September 18, 2026 00:54
@daemyung-lablup
daemyung-lablup requested a review from a team as a code owner September 18, 2026 00:54
@github-actions github-actions Bot added size:L 100~500 LoC comp:agent Related to Agent component labels Sep 18, 2026

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

@HyeockJinKim HyeockJinKim added the skip:changelog Make the action workflow to skip towncrier check label Sep 18, 2026

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:agent Related to Agent component size:L 100~500 LoC skip:changelog Make the action workflow to skip towncrier check

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants