fix(sglang): avoid duplicate model allocation on retry

modelexpress_client/python/modelexpress/adapter.py

@@ -154,7 +154,12 @@ def load_via_native(self, result: LoadResult) -> LoadResult:
 
     @gated_capability
     def reinit_for_retry(self, result: LoadResult) -> LoadResult:
-        """Replace a possibly-mutated model with a fresh engine model instance."""
+        """Restore a possibly-mutated model to freshly initialized state.
+
+        Adapters may return a different model object, or preserve the root
+        object's identity while replacing its complete internal state when an
+        engine-owned caller retains the original root reference.
+        """
         ...
 
     def get_unique_id(self) -> str:

modelexpress_client/python/modelexpress/engines/sglang/adapter.py

@@ -6,6 +6,7 @@
 from __future__ import annotations
 
 import copy
+import gc
 import logging
 import uuid
 from importlib.metadata import version as pkg_version
@@ -165,27 +166,62 @@ def reinit_for_retry(self, result: LoadResult) -> LoadResult:
         )
         from sglang.srt.model_loader.utils import set_default_torch_dtype
 
-        old_value = result.value
+        model = result.model
+        if model is None:
+            raise RuntimeError("SGLang retry reinitialization requires result.model")
+        if result.value is not model:
+            raise RuntimeError(
+                "SGLang retry reinitialization requires result.value and "
+                "result.model to reference the same model root"
+            )
+
+        publishable = result.publishable
+        metadata = result.metadata
         result.value = None
         result.model = None
-        del old_value
+
+        # SGLang's RemoteInstanceModelLoader and MxModelLoader both retain the
+        # root model object while this hook runs. Deleting LoadResult references
+        # cannot release its parameters, so constructing a replacement directly
+        # would temporarily allocate two full models. Preserve the externally
+        # owned root identity, but first turn it into an empty shell so the old
+        # CUDA allocations can be reclaimed before initialization starts.
+        model.__dict__.clear()
+        gc.collect()
         self.accelerator_backend.empty_cache()
 
         logger.info(
-            "[Worker %s] Re-initializing SGLang model after failed strategy",
+            "[Worker %s] Re-initializing SGLang model state in-place after "
+            "failed strategy",
             self.get_global_rank(),
         )
         quant_config = _get_quantization_config(self.model_config, self.load_config)
         # Match SGLang's initial load path so retry parameters use the model's
         # configured dtype instead of PyTorch's default float32.
         with set_default_torch_dtype(self.model_config.dtype):
             with self.target_device:
-                model = _initialize_model(
+                fresh_model = _initialize_model(
                     self.model_config,
                     self.load_config,
                     quant_config,
                 )
-        return LoadResult(value=model, model=model, publishable=result.publishable)
+        if type(fresh_model) is not type(model):
+            raise RuntimeError(
+                "SGLang retry initialization returned a different model type: "
+                f"expected {type(model).__qualname__}, "
+                f"got {type(fresh_model).__qualname__}"
+            )
+
+        # Both roots briefly reference the same new children, so there is still
+        # only one set of parameter storage. The externally owned root remains
+        # valid after the temporary fresh root is dropped.
+        model.__dict__.update(fresh_model.__dict__)
+        del fresh_model
+        result.value = model
+        result.model = model
+        result.publishable = publishable
+        result.metadata = metadata
+        return result
 
     def _process_weights_after_loading(self, result: LoadResult) -> LoadResult:
         if result.model is None:

modelexpress_client/python/modelexpress/engines/sglang/loader.py

@@ -14,6 +14,7 @@
 
 from ... import envs, p2p_pb2
 from ...load_strategy import LoadContext, LoadStrategyChain
+from ...load_strategy.base import clear_exception_tracebacks
 from ...load_strategy.context import LoadResult
 from ...metadata.publisher import PublisherThread
 from ...metadata.payload import tensor_source_metadata, worker_tensor_descriptors
@@ -197,6 +198,9 @@ def _load_model_via_transfer_engine(
                     exc,
                     exc_info=True,
                 )
+                registered_tensors = None
+                tensors = {}
+                clear_exception_tracebacks(exc)
                 result = ctx.adapter.reinit_for_retry(result)
                 result = ctx.adapter.load_via_native(result)
                 tensors = ctx.adapter.discover_tensors(result)

modelexpress_client/python/modelexpress/load_strategy/__init__.py

@@ -22,6 +22,7 @@
     LoadResult,
     LoadStrategy,
     SourceTransferError,
+    clear_exception_tracebacks,
     publish_source_if_supported,
     register_tensors,
     publish_metadata,
@@ -99,6 +100,7 @@ def run(model: nn.Module, ctx: LoadContext) -> nn.Module:
                     )
                     strategy.rollback(ctx)
                     if e.mutated:
+                        clear_exception_tracebacks(e)
                         result = LoadStrategyChain._reinit_for_retry(result, ctx, strategy)
                     continue
                 except Exception as e:

modelexpress_client/python/modelexpress/load_strategy/base.py

@@ -6,6 +6,7 @@
 from __future__ import annotations
 
 import logging
+import traceback
 import uuid
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, ClassVar
@@ -25,6 +26,30 @@
 logger = logging.getLogger("modelexpress.load_strategy")
 
 
+def clear_exception_tracebacks(exc: BaseException) -> None:
+    """Drop completed failure frames before releasing a mutated model.
+
+    Transfer failures commonly retain target tensors through traceback frame
+    locals (for example ``local_tensor`` in the NIXL matching loop). Clearing
+    only ``LoadResult`` and ``LoadContext`` therefore does not guarantee that
+    CUDA allocations become unreachable before retry initialization.
+    """
+    pending: list[BaseException] = [exc]
+    seen: set[int] = set()
+    while pending:
+        current = pending.pop()
+        if id(current) in seen:
+            continue
+        seen.add(id(current))
+        if current.__cause__ is not None:
+            pending.append(current.__cause__)
+        if current.__context__ is not None:
+            pending.append(current.__context__)
+        if current.__traceback__ is not None:
+            traceback.clear_frames(current.__traceback__)
+            current.__traceback__ = None
+
+
 class SourceTransferError(Exception):
     """Raised when a failure is demonstrably from the remote source side.
 

modelexpress_client/python/modelexpress/load_strategy/rdma_strategy.py

@@ -30,6 +30,7 @@
     LoadStrategy,
     SourceTransferError,
     _as_load_result,
+    clear_exception_tracebacks,
     register_tensors,
 )
 from .context import LoadResult
@@ -156,7 +157,6 @@ def load(self, result: LoadResult, ctx: LoadContext) -> LoadResult:
 
         attempts = candidates[:MAX_SOURCE_RETRIES]
         policy = configured_policy_label()
-        needs_outer_reinit = False
         for attempt_index, instance in enumerate(attempts):
             mx_source_id = instance.mx_source_id
             worker_id = instance.worker_id
@@ -215,8 +215,6 @@ def load(self, result: LoadResult, ctx: LoadContext) -> LoadResult:
                     "transfer_retry" if has_next_candidate else "transfer_fallback",
                 )
                 if not has_next_candidate:
-                    if needs_outer_reinit and not e.mutated:
-                        raise StrategyFailed(str(e), mutated=True) from e
                     raise
 
                 logger.warning(
@@ -233,14 +231,24 @@ def load(self, result: LoadResult, ctx: LoadContext) -> LoadResult:
                     ) from cleanup_error
                 if e.mutated:
                     try:
-                        result = ctx.adapter.reinit_for_retry(result)
+                        clear_exception_tracebacks(e)
+                        reinitialized = ctx.adapter.reinit_for_retry(result)
+                        # LoadResult is the stable envelope shared with the
+                        # outer strategy chain. Some adapters return a new
+                        # envelope, so copy its restored state back rather than
+                        # leaving the outer owner with the cleared pre-retry
+                        # object if all later candidates miss.
+                        if reinitialized is not result:
+                            result.value = reinitialized.value
+                            result.model = reinitialized.model
+                            result.publishable = reinitialized.publishable
+                            result.metadata = reinitialized.metadata
                     except Exception as reinit_error:
                         raise StrategyFailed(
                             f"Failed to reinitialize target after source worker "
                             f"{worker_id} failed: {reinit_error}",
                             mutated=True,
                         ) from reinit_error
-                    needs_outer_reinit = True
                 continue
             except BaseException:
                 selection_metrics.observe_transfer_seconds(
@@ -259,11 +267,9 @@ def load(self, result: LoadResult, ctx: LoadContext) -> LoadResult:
             f"[Worker {ctx.global_rank}] Tried {tried} of {len(candidates)} source workers "
             f"(max retries={MAX_SOURCE_RETRIES}), falling through"
         )
-        # An internal reinit returns a new result, but the outer strategy chain
-        # still owns the original result that the adapter cleared.
         raise StrategyFailed(
             "No RDMA source succeeded",
-            mutated=needs_outer_reinit,
+            mutated=False,
         )
 
     def _find_source_instances(

modelexpress_client/python/tests/test_sglang_loader.py

@@ -5,6 +5,7 @@
 
 import os
 import sys
+import weakref
 from contextlib import contextmanager
 from types import ModuleType
 from types import SimpleNamespace
@@ -20,6 +21,7 @@
     build_sglang_load_context,
 )
 from modelexpress.engines.sglang.loader import MxModelLoader
+from modelexpress.load_strategy.context import LoadResult
 
 
 def _load_config(**overrides):
@@ -381,6 +383,8 @@ def test_sglang_retry_initializes_model_with_configured_dtype(monkeypatch):
     loader_mod = ModuleType("sglang.srt.model_loader.loader")
     model_loader_utils_mod = ModuleType("sglang.srt.model_loader.utils")
     observed_dtypes = []
+    initial_model = nn.Linear(2, 2)
+    initial_weight_ref = weakref.ref(initial_model.weight)
 
     @contextmanager
     def set_default_torch_dtype(dtype):
@@ -394,6 +398,7 @@ def set_default_torch_dtype(dtype):
     loader_mod._get_quantization_config = lambda *_: None
 
     def initialize_model(*_):
+        assert initial_weight_ref() is None
         observed_dtypes.append(torch.get_default_dtype())
         return nn.Linear(2, 2)
 
@@ -411,16 +416,85 @@ def initialize_model(*_):
 
     model_config = _model_config(dtype=torch.bfloat16)
     adapter = SglangAdapter(_load_config(), model_config, _device_config())
-    result = SimpleNamespace(
-        value=nn.Linear(2, 2),
-        model=nn.Linear(2, 2),
+    result = LoadResult(
+        value=initial_model,
+        model=initial_model,
         publishable=True,
     )
 
-    adapter.reinit_for_retry(result)
+    retried = adapter.reinit_for_retry(result)
 
     assert observed_dtypes == [torch.bfloat16]
     assert torch.get_default_dtype() == original_dtype
+    assert retried.value is initial_model
+    assert retried.model is initial_model
+    assert list(initial_model.parameters())
+
+
+def test_sglang_retry_reuses_root_for_native_fallback(monkeypatch):
+    sglang_mod = ModuleType("sglang")
+    srt_mod = ModuleType("sglang.srt")
+    model_loader_mod = ModuleType("sglang.srt.model_loader")
+    loader_mod = ModuleType("sglang.srt.model_loader.loader")
+    model_loader_utils_mod = ModuleType("sglang.srt.model_loader.utils")
+    configs_mod = ModuleType("sglang.srt.configs")
+    load_config_mod = ModuleType("sglang.srt.configs.load_config")
+
+    @contextmanager
+    def set_default_torch_dtype(_dtype):
+        yield
+
+    initial_model = nn.Linear(2, 2)
+    initial_weight_ref = weakref.ref(initial_model.weight)
+    native_roots = []
+
+    loader_mod._get_quantization_config = lambda *_: None
+
+    def initialize_model(*_):
+        assert initial_weight_ref() is None
+        return nn.Linear(2, 2)
+
+    class DefaultModelLoader:
+        def __init__(self, _load_config):
+            pass
+
+        def _get_all_weights(self, _model_config, model):
+            native_roots.append(model)
+            return iter([])
+
+        @staticmethod
+        def load_weights_and_postprocess(model, _weights, _target_device):
+            model.weight.data.fill_(7)
+
+    loader_mod._initialize_model = initialize_model
+    loader_mod.DefaultModelLoader = DefaultModelLoader
+    model_loader_utils_mod.set_default_torch_dtype = set_default_torch_dtype
+    load_config_mod.LoadFormat = SimpleNamespace(AUTO="auto")
+    monkeypatch.setitem(sys.modules, "sglang", sglang_mod)
+    monkeypatch.setitem(sys.modules, "sglang.srt", srt_mod)
+    monkeypatch.setitem(sys.modules, "sglang.srt.model_loader", model_loader_mod)
+    monkeypatch.setitem(sys.modules, "sglang.srt.model_loader.loader", loader_mod)
+    monkeypatch.setitem(
+        sys.modules,
+        "sglang.srt.model_loader.utils",
+        model_loader_utils_mod,
+    )
+    monkeypatch.setitem(sys.modules, "sglang.srt.configs", configs_mod)
+    monkeypatch.setitem(
+        sys.modules,
+        "sglang.srt.configs.load_config",
+        load_config_mod,
+    )
+
+    adapter = SglangAdapter(_load_config(), _model_config(), _device_config())
+    result = LoadResult(value=initial_model, model=initial_model)
+
+    retried = adapter.reinit_for_retry(result)
+    loaded = adapter.load_via_native(retried)
+
+    assert loaded.model is initial_model
+    assert native_roots == [initial_model]
+    assert torch.all(initial_model.weight == 7)
 
 
 def test_mx_model_loader_delegates_to_shared_strategy_chain():