[JunkaiWang-TheoPhy]

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 same try block 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 main 0ac607299eb711cdd6b164e9eba04a7ade9cf98b, retaining the new test:

python -m pytest tests/test_registration_rejection.py -q --tb=short
2 failed in 10.36s

Both ordinary and shared-pool cases fail with RuntimeError: checkpoint weights is not registered after the duplicate assertion is caught.

With the fix:

python -m pytest tests/test_registration_rejection.py -q --tb=short
2 passed in 4.09s

python -m pytest tests/ -m "not gpu" -q --tb=short
109 passed, 1 skipped, 12 deselected, 1 warning in 1.71s

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.py and tests/test_registration_rejection.py change. 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.

[koriyoshi2041]

Reviewed exact head 3c8f4e33. Moving both duplicate-registration guards ahead of the cleanup try is the right state boundary: a rejected duplicate no longer reaches unregister_checkpoint(), so the existing dedicated pool or shared-pool user remains intact and host_empty_cache() is not invoked. The two-path regression test also proves allocation is attempted only once and the original buffer remains readable. Locally, the focused cases passed 2/2 and the complete non-GPU suite passed 109 tests (1 skipped, 12 GPU-deselected); git diff --check origin/main..HEAD is clean.

[sunruize93-cmyk]

Reviewed 3c8f4e330e8295f075fc49172384d701e5968e4a against official main/merge-base 0ac607299eb711cdd6b164e9eba04a7ade9cf98b.

Standards

No documented-standard findings. Changed files pass repository-pinned Ruff 0.12.2 lint/format and whitespace checks.

Spec

Moving the dedicated duplicate and shared-occupancy guards outside the rollback block preserves existing registrations on rejection. Independent controls cover repeated rejection, same/different shared-pool names, empty dedicated registrations, unchanged buffer contents, and P2P registration bookkeeping. No new blocking finding in this diff.

One independent CPU comparison on macOS arm64 / Python 3.12.13 / PyTorch 2.14.1:

• Committed non-GPU suite: 109 passed, 1 skipped, 12 GPU cases deselected. The PR's two regressions are included; both fail against main.
• Additional controls: 23 passed on head; 13 passed and 10 duplicate-protection variants failed on main. Success/unregister/re-register, allocation-error propagation without P2P, failed shared-pool reuse, and unrelated-registration preservation remain working.
• Three additional desired P2P failure contracts are unmet on both revisions, recorded as strict xfails and then expanded for diagnosis: dedicated allocation failure masks the original error with “checkpoint weights is not registered”; failed dedicated P2P registration reaches a second unregister and raises KeyError; failed shared P2P registration clears the owner but retains named_tensors entries. These are unchanged baseline cleanup limitations, not regressions introduced here, and the xfails are not passes. This review does not certify complete P2P rollback.

The controls execute production registration/unregistration, packing/copying and P2PStore methods. Hardware initialization, pinned allocation, device cache operations and the transfer engine are CPU doubles. No real GPU/pinned-memory/RDMA transfer, distributed concurrency or cross-platform validation was performed.