fix(tinker): finalize failed saves for unloaded models

[LOGO127]

Summary

Move the unloaded-model checks into the existing checkpoint-status context for
training and sampler saves. The request's future and checkpoint then both reach
FAILED with the same error instead of leaving the checkpoint PENDING.

When a queued save is dispatched after its model is no longer loaded, the current
early ErrorResponse completes the future as FAILED but bypasses checkpoint
finalization. Deleting that checkpoint returns 425, and reusing its name returns
409. This change allows the failed checkpoint record to be deleted and retried.

The two-file patch retains successful-save behavior and ephemeral sampler
persist=False behavior. It does not repair historical orphan records, change
retention policy (#2116), or implement Megatron checkpoint publication (#2145).

Validation

• Added 12 CPU regression/control cases: success, backend failure and unloaded
  model across training, persistent sampler and ephemeral sampler saves, plus
  actual ASGI/SQLite delete-and-retry paths. Only the backend is mocked in these
  committed tests.
• Same tests on baseline 0b286bac: 6 failed / 6 passed. Failures are three
  PENDING-vs-FAILED statuses and three delete responses of 425 instead of 204.
• Signed candidate: all 12 pass. An earlier selected CPU suite including these
  tests passed 94 tests (not an additional 94).
• Supplementary, uncommitted real-JAX lifecycle probe: 6 pass, versus 3 failures
  and 3 passing controls on baseline. This deliberately stages unload before
  dispatch; it does not claim that ordinary FIFO requests reorder themselves.
• Supplementary real TCP smoke: the unmodified API launches its JAX engine child;
  Tinker SDK 0.24.0 creates a tiny-Qwen LoRA model and successfully saves a training
  checkpoint. One test passed in 44.90s. This is a successful-save control, not a
  cross-process reproduction of the unload condition.
• Repository Ruff, Black and secret-detection hooks passed on the reviewed patch;
  signed contents match that patch, and git diff --check passes.

Local validation used WSL Ubuntu, Python 3.12 and an independent locked CPU
environment (JAX 0.11.0, Flax 0.12.8, CPU Torch 2.11.0, Tinker 0.24.0), not the
full production dependency installation. Supplementary probes and environment
files are not included in this scoped patch. No GPU training, PostgreSQL,
SDK >=0.25 protobuf compatibility, model quality or performance is claimed.

To run the committed regression suite in a configured JAX development environment:

uv run --isolated --extra dev --extra jax pytest tests/tinker/test_checkpoint_status.py -q

Attribution

AI-assisted investigation, implementation and test automation with Codex. The
contributor reviewed the patch and personally signed commit 8671853.

---

Note

Medium Risk
Changes checkpoint failure finalization and error handling on save paths; behavior for successful saves is intended unchanged but affects DB state and delete/retry semantics for failed saves.

Overview
Fixes stale save requests when a model is no longer loaded: training and sampler weight saves now validate has_model inside _checkpoint_status_context instead of returning early, so checkpoint rows move from PENDING to FAILED alongside the request future.

Introduces _ModelNotLoadedError for that case; checkpoint and batch request handlers still return client ErrorResponses but skip full exception logging for this expected stale-request path. Save-weight processors no longer declare ErrorResponse in their return types because failures are raised through the context manager.

Adds ASGI/SQLite regression tests covering unloaded-model and backend-failure saves (training and sampler) and that a failed checkpoint can be deleted (204) instead of blocking on PENDING/425.

^{Reviewed by Cursor Bugbot for commit e9940cf. Bugbot is set up for automated code reviews on this repo. Configure here.}

[gemini-code-assist]

Code Review

This pull request moves the model existence check inside the checkpoint status context block in both process_save_weights and process_save_weights_for_sampler, raising a ValueError if the model is not found. It also introduces a new test suite to verify checkpoint status behaviors and retries. The feedback suggests using a custom exception instead of a generic ValueError to avoid generating noisy tracebacks in the logs for expected stale requests.

[gemini-code-assist]

Raising a generic ValueError here will propagate through _checkpoint_status_context and process_single_requests. Both of these catch Exception and log it using logger.exception, which will produce two noisy, full tracebacks in the logs for what is an expected and common scenario (a stale request for an unloaded model after a server restart).

To keep the logs clean, consider introducing a custom exception (e.g., ModelNotLoadedError) and updating _checkpoint_status_context and process_single_requests to handle it gracefully (e.g., logging a simple warning or info message without a traceback) while still marking the checkpoint and future as failed.

Suggested change

	if not self.backend.has_model(model_id):
	raise ValueError(_model_not_found_error(model_id).error)
	if not self.backend.has_model(model_id):
	raise ModelNotLoadedError(_model_not_found_error(model_id).error)

[LOGO127]

Thanks for catching the noisy stale-request logs. Addressed in e0f30b5d8676e3a69e1321ac71cc3c2444478574 with a dedicated internal _ModelNotLoadedError in both save paths. The two exception handlers omit tracebacks only for that subtype; the existing _model_not_found_error helper still emits one info message, without an additional warning. Checkpoint/future failure finalization and error text are preserved.

The regression matrix covers training, sampler and ephemeral sampler saves: unloaded requests produce no traceback and one info message, while ordinary backend ValueError and OSError retain both existing exception logs. Successful saves remain covered. On the exact signed commit, the checkpoint suite plus test_engine.py reran with 33 passed in 28.61s in the pinned local CPU environment. This does not claim GPU validation or resolve the previously reported upstream Anyscale credentials issue.

The contributor personally reviewed/signed the increment; implementation and verification used Codex assistance. The signed file blobs match the reviewed and previously hook-checked candidate; no history rewrite was performed.

[gemini-code-assist]

Raising a generic ValueError here will propagate through _checkpoint_status_context and process_single_requests. Both of these catch Exception and log it using logger.exception, which will produce two noisy, full tracebacks in the logs for what is an expected and common scenario (a stale request for an unloaded model after a server restart).

To keep the logs clean, consider introducing a custom exception (e.g., ModelNotLoadedError) and updating _checkpoint_status_context and process_single_requests to handle it gracefully (e.g., logging a simple warning or info message without a traceback) while still marking the checkpoint and future as failed.

Suggested change

	if not self.backend.has_model(model_id):
	raise ValueError(_model_not_found_error(model_id).error)
	if not self.backend.has_model(model_id):
	raise ModelNotLoadedError(_model_not_found_error(model_id).error)

[LOGO127]

Addressed in the same signed e0f30b5 increment as the training-save path. The sampler and persist=False ephemeral sampler cases are included in the logging/status regression matrix; full details and validation limits are in my reply on the first thread.

[greptile-apps]

Confidence Score: 5/5

The PR appears safe to merge with no actionable defects identified in the changed paths.

Diagram

sequenceDiagram
    participant API
    participant DB
    participant Engine
    participant Backend
    API->>DB: Create PENDING checkpoint and future
    Engine->>DB: Enter checkpoint-status context
    Engine->>Backend: Check whether model is loaded
    alt Model is loaded
        Engine->>Backend: Save checkpoint
        Engine->>DB: Mark checkpoint COMPLETED
        Engine->>DB: Mark future COMPLETED
    else Model is unloaded
        Engine->>DB: Mark checkpoint FAILED
        Engine->>Engine: Convert ValueError to ErrorResponse
        Engine->>DB: Mark future FAILED
    end

Loading

[LOGO127]

Thanks for taking a look! A brief CI note: the GPU Actions job stopped while submitting its Anyscale job, with Your user credentials are invalid, before the GPU tests ran. I am not counting this as GPU validation or attributing it to the patch. The CPU workflow is still running at this check.

Could a maintainer please advise the normal way to run this check for a fork PR when convenient? I have left the workflow and permissions unchanged; I am not asking for secrets to be shared or checks to be bypassed. Local CPU evidence and its limitations are in the PR description. Thank you! This CI investigation and follow-up were prepared with AI assistance.

[LOGO127]

Thanks for reviewing. I investigated the CPU failure at e0f30b5: the two failures are test_qwen3_moe_layer_lora[1-1] and [2-1], in the merged-weight numerical comparison. All 33 checkpoint/engine tests passed in that run.

The MoE code/tests, dependency files, and CPU workflow are unchanged by this PR. I ran the three MoE LoRA configurations on both the pre-patch baseline 0b286bac and PR head in the same pinned local CPU environment; both passed. This is not an exact reproduction of the GitHub runner, so I cannot yet establish the cause or rule out an interaction.

Could a maintainer please rerun the failed CPU job when convenient? My rerun attempt was denied for insufficient repository permissions. I have left the test tolerances and workflow unchanged. If it fails again, I can continue investigating with the runner environment and a reproducible input. Thank you!

This investigation and message were prepared with AI assistance.

[ktanishqk]

fix is right, move it in the with statement. however, testing is a bit too extensive and adds failure points that might hurt more than help. going to ask for some consolidation and reduction there - check reviews!

[ktanishqk]

the grid is easy to grow and hard to read. imo: unloaded × training/sampler + one ValueError (for the _ModelNotLoadedError check) is enough

[ktanishqk]

seems brittle. i wouldn't test logs here.
if someone later adds a logger.info on this path, this test fails even though the checkpoint status is still right.
i’d drop the counts and just assert the DB rows.

[ktanishqk]

can we not just check for checkpoint PENDING -> FAILED and then delete returns 204? the rest aren't this bug and give the test extra ways to fail.

[ktanishqk]

hey @LOGO127 - i had started working on this issue after creating it because i was waiting on scope confirmation before opening a pr.
however, thanks for starting work on this. going to turn into a reviewer instead so we can coordinate and land the best version!

[LOGO127]

Thanks @ktanishqk for coordinating and taking the reviewer role. The test-only follow-up has now been consolidated locally into three cases: unloaded training save, unloaded sampler save, and one ordinary backend ValueError. It checks the PENDING -> FAILED database transition and DELETE 204, without logger-count, duplicate/retry, or successful-save matrices.

The reduced tests still reproduce the original bug: unmodified pre-fix source gives 2 failed / 1 passed; the candidate plus test_engine.py gives 21 passed. Ruff, Black and secret-detection hooks pass. These are local CPU results in the previously documented independent pinned environment, not a resolution of the existing MoE/Anyscale CI failures.

The follow-up is local pending contributor review/sign-off; this PR's published head remains e0f30b5. Thanks for helping keep the regression focused.

[LOGO127]

Addressed the review in e9940cf by reducing the checkpoint regression coverage to the three requested cases: unloaded training save, unloaded sampler save, and one ordinary backend ValueError. The test now only checks PENDING -> FAILED plus DELETE 204; the logger-count, retry/duplicate, ephemeral, and success matrices were removed.

Validation: old production code + the reduced tests gives 2 failed / 1 passed (the two unloaded cases stay PENDING); current PR head + the reduced tests and test_engine.py gives 21 passed. Ruff, Black, and hardcoded-secret hooks pass as well.

Thanks for the direction — the production fix is unchanged.