fix(sglang): avoid duplicate model allocation on retry

[daii-0818]

Summary

Prevent SGLang NIXL/RDMA retries from temporarily keeping two complete model allocations alive after a transfer mutates the target.

• preserve the engine-owned root model identity while replacing its complete internal state
• release the old module/parameter/buffer graph before calling SGLang _initialize_model()
• clear completed exception traceback frames that can retain transfer-local tensors
• keep the shared LoadResult envelope synchronized across RDMA source retries and native fallback
• preserve existing stale-source selection, alternate-source retry, and default loader behavior

Root cause

SGLang initializes the model before delegating to MxModelLoader. On a mutated RDMA failure, the previous retry path cleared LoadResult and constructed a fresh model, but the original root was still strongly referenced by:

• SGLang RemoteInstanceModelLoader.load_model()
• ModelExpress MxModelLoader and LoadStrategyChain.run() frames
• transfer exception tracebacks and NIXL tensor registration state

empty_cache() cannot release active parameter allocations. Models that already consume most of HBM therefore OOM while _initialize_model() attempts to allocate the second complete model.

Design

SGLang's loader API does not allow ModelExpress to delete the caller's root reference. The adapter now:

1. requires LoadResult.value and .model to reference the same root
2. clears the result envelope and the retained root's complete nn.Module state
3. runs GC and releases allocator cache after NIXL rollback/registration cleanup
4. initializes fresh SGLang model state
5. moves that fresh state onto the original root object
6. restores the same LoadResult envelope

The temporary fresh root and original root only reference the same new child graph; there is never a second parameter-storage set. This retains fresh-model semantics for post-load/quantization hooks and is safer than overwriting a partially transformed model in place.

VMM remains orthogonal: it reduces allocation and MR registration overhead, but does not remove engine-owned model references. The retry lifecycle works with and without a VMM arena.

Validation

Focused client tests on the latest main:

241 passed, 22 skipped

Development GPU host:

92 passed

A real CUDA allocation regression was run on one NVIDIA H20. The fake SGLang model held a 50 GiB CUDA buffer, more than half the memory required to demonstrate that a second live copy would not fit safely:

root_identity_preserved=True
old_storage_released_before_init=True
allocated_before_init_gib=0.0
allocated_after_retry_gib=50.0
max_allocated_gib=50.0

The regression tests also cover:

• first RDMA source fails after mutation, second source succeeds
• a later source miss remains eligible for native fallback without a second reinitialization
• native fallback uses the restored engine-owned root
• exception traceback locals release retained allocations
• VMM context present and absent
• normal RDMA success and existing stale-source paths remain unchanged

Summary by CodeRabbit

• Bug Fixes

  • Improved model retry and recovery behavior after loading failures.
  • Released retained resources and stale allocation state before retrying.
  • Preserved model identity, metadata, and data types during reinitialization.
  • Improved fallback loading reliability when the primary transfer method fails.
  • Corrected failure reporting after clean transfer failures.

• Tests

  • Added coverage for resource release, retry behavior, model preservation, and native loading fallback.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[coderabbitai]

Walkthrough

Changes

Retry reinitialization and cleanup

Layer / File(s)	Summary
Model-root retry reinitialization modelexpress_client/python/modelexpress/adapter.py, modelexpress_client/python/modelexpress/engines/sglang/adapter.py, modelexpress_client/python/tests/test_sglang_loader.py	The retry contract permits root preservation. SGLang clears the old model, initializes replacement state, restores metadata, and keeps the original root identity.
Exception cleanup and transfer recovery modelexpress_client/python/modelexpress/load_strategy/base.py, modelexpress_client/python/modelexpress/load_strategy/__init__.py, modelexpress_client/python/modelexpress/engines/sglang/loader.py, modelexpress_client/python/modelexpress/load_strategy/rdma_strategy.py	Retry paths clear chained exception tracebacks, reset transfer state, copy replacement result state, and report final non-mutating failures.
Retry regression coverage modelexpress_client/python/tests/test_source_selection.py	Tests cover traceback resource release, retry identity, VMM arena states, metadata failures, and clean transfer failures.

Estimated code review effort: 4 (Complex) | ~45 minutes

Poem

I’m a rabbit guarding the retry trail,
Old frames vanish, and roots prevail.
New weights hop into the same model shell,
Native paths recover when transfers fail.
Metadata stays tucked in tight—
The loader lands its state just right.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 9.52% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and concisely describes the main change: preventing duplicate model allocation during SGLang retries.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (2)

modelexpress_client/python/tests/test_source_selection.py (1)

694-717: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Add coverage for a mutated failure on the final candidate.

These tests cover the terminal mutated=False paths. The complementary path is not covered: when the last candidate fails with mutated=True, load() re-raises the original StrategyFailed at Line 218 of modelexpress_client/python/modelexpress/load_strategy/rdma_strategy.py without calling reinit_for_retry. The outer chain depends on that flag to reinitialize the model. A regression there would silently skip outer reinitialization.

💚 Proposed test

def test_load_last_candidate_mutated_failure_propagates_mutated():
    strat = RdmaStrategy()
    strat._find_source_instances = MagicMock(return_value=_sources(1))
    strat._fetch_worker_metadata = MagicMock(return_value=MagicMock())
    strat._load_as_target = MagicMock(
        side_effect=StrategyFailed("mutated failure", mutated=True)
    )
    ctx = MagicMock(global_rank=0)
    ctx.accelerator_backend.name = ""

    with pytest.raises(StrategyFailed, match="mutated failure") as exc:
        strat.load(MagicMock(), ctx)

    assert exc.value.mutated is True
    ctx.adapter.reinit_for_retry.assert_not_called()

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@modelexpress_client/python/tests/test_source_selection.py` around lines 694 -
717, Add a test alongside
test_load_internal_reinit_then_clean_failure_stays_clean that configures a
single source candidate and makes _load_as_target raise StrategyFailed with
mutated=True. Assert load() propagates the original mutated failure and that
ctx.adapter.reinit_for_retry is not called, covering the final-candidate path.

modelexpress_client/python/modelexpress/load_strategy/rdma_strategy.py (1)

241-245: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick win

Copy the complete LoadResult state.

LoadResult is a non-slotted, non-frozen dataclass. Replace the manual field assignments so future dataclass fields are preserved:

if reinitialized is not result:
    vars(result).update(vars(reinitialized))

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@modelexpress_client/python/modelexpress/load_strategy/rdma_strategy.py`
around lines 241 - 245, In the reinitialization handling around the LoadResult
objects, replace the manual assignments to value, model, publishable, and
metadata with vars(result).update(vars(reinitialized)) so the complete dataclass
state, including future fields, is copied whenever reinitialized is not result.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@modelexpress_client/python/modelexpress/engines/sglang/adapter.py`:
- Around line 178-213: Wrap the in-place reinitialization and type validation in
a failure-safe path that restores the original model envelope and LoadResult
state before propagating the original exception. Preserve the pre-clear model
root and restore its fields, along with result.value and result.model, if
_initialize_model or the fresh_model type check fails; successful initialization
should continue with the existing flow.

---

Nitpick comments:
In `@modelexpress_client/python/modelexpress/load_strategy/rdma_strategy.py`:
- Around line 241-245: In the reinitialization handling around the LoadResult
objects, replace the manual assignments to value, model, publishable, and
metadata with vars(result).update(vars(reinitialized)) so the complete dataclass
state, including future fields, is copied whenever reinitialized is not result.

In `@modelexpress_client/python/tests/test_source_selection.py`:
- Around line 694-717: Add a test alongside
test_load_internal_reinit_then_clean_failure_stays_clean that configures a
single source candidate and makes _load_as_target raise StrategyFailed with
mutated=True. Assert load() propagates the original mutated failure and that
ctx.adapter.reinit_for_retry is not called, covering the final-candidate path.

🪄 Autofix

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 79fd16db-7fe9-4cb2-a887-bf660e3417a0

📥 Commits

Reviewing files that changed from the base of the PR and between 7e0421e and b5d396f.

📒 Files selected for processing (8)

• modelexpress_client/python/modelexpress/adapter.py
• modelexpress_client/python/modelexpress/engines/sglang/adapter.py
• modelexpress_client/python/modelexpress/engines/sglang/loader.py
• modelexpress_client/python/modelexpress/load_strategy/__init__.py
• modelexpress_client/python/modelexpress/load_strategy/base.py
• modelexpress_client/python/modelexpress/load_strategy/rdma_strategy.py
• modelexpress_client/python/tests/test_sglang_loader.py
• modelexpress_client/python/tests/test_source_selection.py

[daii-0818]

Addressed the remaining review items in 7c41dcd:

• copy the complete LoadResult state with vars(...).update(...)
• cover a mutated failure on the final RDMA candidate
• restore the retry envelope when fresh initialization fails
• fail closed so an empty model root never reaches a later strategy

Validation:

• focused client tests: 244 passed, 22 skipped
• H20 retry test with a 50 GiB fake model: 50 GiB before, 50 GiB after, 50 GiB peak
• forced reinitialization failure: original exception preserved and active CUDA allocation returned to 0 GiB

[daii-0818]

Hi ! @zhengluo-nv @nv-hwoo Could you please review this SGLang retry lifecycle fix? It addresses the duplicate-model CUDA OOM after an RDMA/NIXL transfer failure while preserving stale-source retry and native fallback behavior. The retry path was validated on H20 with a 50 GiB allocation, and all CI checks are green. Feedback welcome!