fix(sglang): avoid duplicate model allocation on retry

modelexpress_client/python/modelexpress/adapter.py

@@ -34,6 +34,15 @@ def __init__(self, message: str, *, mutated: bool = False):
         self.mutated = mutated
 
 
+class StrategyRecoveryError(RuntimeError):
+    """Raised when a failed strategy cannot restore a safe model state.
+
+    The strategy chain must stop immediately: trying another loader with a
+    partially cleared or otherwise unrecoverable model would hide the original
+    recovery failure and may publish invalid weights.
+    """
+
+
 def gated_capability(method):
     """Create an optional adapter method that engines must override to support it.
 
@@ -154,7 +163,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,75 @@ 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(
-                    self.model_config,
-                    self.load_config,
-                    quant_config,
+        try:
+            with set_default_torch_dtype(self.model_config.dtype):
+                with self.target_device:
+                    fresh_model = _initialize_model(
+                        self.model_config,
+                        self.load_config,
+                        quant_config,
+                    )
+            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__}"
                 )
-        return LoadResult(value=model, model=model, publishable=result.publishable)
+        except BaseException:
+            # The old parameter graph was intentionally released before fresh
+            # initialization and cannot be restored without retaining the HBM
+            # that caused duplicate-model OOM. Restore the envelope to the
+            # engine-owned empty root so callers do not observe None, then let
+            # the original failure abort startup rather than attempting another
+            # strategy with an invalid model.
+            result.value = model
+            result.model = model
+            result.publishable = publishable
+            result.metadata = metadata
+            raise
+
+        # 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

@@ -16,12 +16,13 @@
 
 from modelexpress.tracing import tracer
 
-from ..adapter import StrategyFailed, UnsupportedCapability
+from ..adapter import StrategyFailed, StrategyRecoveryError, UnsupportedCapability
 from .base import (
     LoadContext,
     LoadResult,
     LoadStrategy,
     SourceTransferError,
+    clear_exception_tracebacks,
     publish_source_if_supported,
     register_tensors,
     publish_metadata,
@@ -92,13 +93,20 @@ def run(model: nn.Module, ctx: LoadContext) -> nn.Module:
                     publish_source_if_supported(result, ctx)
                     span.set_attribute("weight_loading_strategy", strategy.name)
                     return result.value
+                except StrategyRecoveryError:
+                    # Recovery already failed, so no later strategy can safely
+                    # use the current model. Fail closed and retain the original
+                    # recovery error as the exception cause.
+                    strategy.rollback(ctx)
+                    raise
                 except StrategyFailed as e:
                     logger.warning(
                         f"[Worker {ctx.global_rank}] Strategy {strategy.name} failed, "
                         f"trying next: {e}"
                     )
                     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

@@ -11,7 +11,7 @@
 import time
 
 from .. import envs, p2p_pb2
-from ..adapter import EngineAdapter, StrategyFailed
+from ..adapter import EngineAdapter, StrategyFailed, StrategyRecoveryError
 from ..metadata.payload import (
     accelerators_compatible,
     worker_tensor_count,
@@ -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,20 @@ 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:
+                            vars(result).update(vars(reinitialized))
                     except Exception as reinit_error:
-                        raise StrategyFailed(
+                        raise StrategyRecoveryError(
                             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 +263,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(