Repository navigation
fix: preserve registered checkpoints when duplicate registration is rejected - #113
Open
JunkaiWang-TheoPhy wants to merge 1 commit into
Conversation
A registration precondition failure entered the allocation rollback handler and removed a checkpoint that belonged to an earlier successful call. Check duplicate names and shared-pool occupancy before entering that handler, so rejected attempts cannot destroy the existing registration. Confidence: high Scope-risk: narrow Tested: New CPU regressions fail on main and pass with the fix; full CPU suite Not-tested: Required pre-commit hooks could not initialize because the shared disk is full; GPU, RDMA and actual transfer
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Rejecting a repeated
register_checkpoint("weights", ...)call currently unregisters the checkpoint that was successfully registered by the first call. The duplicate-name assertion runs inside the sametryblock as allocation, so its exception reaches the rollback handler and removes the existing buffers. Repeating the name of the active shared-pool checkpoint also clears its owner marker, making those weights unavailable through that checkpoint name.Move the existing duplicate-name and shared-pool occupancy checks before the allocation/rollback block. Rejected calls then preserve the earlier registration; failures after registration starts still use the existing cleanup path. The public signature, assertion messages, and transfer behavior remain unchanged.
Regression evidence
The CPU regression performs an initial registration through the real
ParameterServer.register_checkpoint, then rejects a second call with the same name. It checks that the original pool remains available with unchanged tensor contents, the shared owner remains intact, and no second allocation or host-cache cleanup occurs. Hardware initialization and pinned allocation are replaced at their boundaries; no GPU or transfer result is claimed.On unmodified
main0ac607299eb711cdd6b164e9eba04a7ade9cf98b, retaining the new test:Both ordinary and shared-pool cases fail with
RuntimeError: checkpoint weights is not registeredafter the duplicate assertion is caught.With the fix:
Full repository pre-commit hooks passed, including repository-pinned Ruff 0.12.2 lint/format and codespell. Python compilation and whitespace checks passed. A separate-agent review reran the regression and verified that actual allocation and subsequent P2P-registration failures still enter the existing rollback handler.
Scope
Only
checkpoint_engine/ps.pyandtests/test_registration_rejection.pychange. No dependency, public API, CI, GPU, RDMA or model behavior changes. This addresses registration preconditions, independently of the P2P registration-state fixes proposed in #91.AI assistance: investigation, implementation, local validation and review used OpenAI Codex.