【Task.026】Version-gated Mooncake loss guards

.gitignore

@@ -54,3 +54,7 @@ tensorboard_log
 # .github
 .github/copilot-instructions.md
 env.sh
+
+# Machine-local multimodal acceptance fixtures (generated by
+# scripts/benchmarks/make_multimodal_fixture.py; hundreds of MB, never commit)
+tests/fixtures/

relax/backends/megatron/actor.py

@@ -13,7 +13,6 @@
 import requests
 import torch
 import torch.distributed as dist
-import transfer_queue as tq
 from megatron.core import mpu
 
 
@@ -61,6 +60,7 @@
 from relax.utils.rotate_ckpt import rotate_ckpt
 from relax.utils.s3_model_loader import prepare_model_maybe_update_args
 from relax.utils.timer import Timer, inverse_timer, timer, with_defer
+from relax.utils.tq_lifecycle import attach_tq_client, detach_tq_client
 from relax.utils.tracking_utils import init_tracking
 from relax.utils.training import train_dump_utils
 from relax.utils.training.data_fields import build_data_fields
@@ -151,6 +151,20 @@ def _per_step_rollout(self) -> bool:
         periodic predict steps; Megatron stays awake between."""
         return not is_sft_mode(self.args)
 
+    def __del__(self) -> None:
+        # Best-effort detach on graceful teardown; ray.kill / fate-sharing
+        # kills skip destructors, in which case the Mooncake master TTL
+        # reclaims the segment.
+        generation = getattr(self, "_tq_client_generation", None)
+        if getattr(self, "data_system_client", None) is None or generation is None:
+            return
+        try:
+            detach_tq_client(generation)
+            self._tq_client_generation = None
+            self.data_system_client = None
+        except Exception:  # destructor must never raise (interpreter shutdown)
+            return
+
     def init(
         self,
         args: Namespace,
@@ -187,8 +201,12 @@ def _init(
         init(args)
         if repatch is not None:
             repatch(args)
-        tq.init(args.tq_config)
-        self.data_system_client = tq.get_client()
+        self.data_system_client = attach_tq_client(
+            args.tq_config,
+            requested_gdr=getattr(args, "tq_use_gdr", False),
+            role=role,
+            lease_owner=self,
+        )
         if is_megatron_main_rank():
             init_tracking(args, primary=False)
 

relax/components/actor.py

@@ -7,7 +7,6 @@
 from typing import Any, Dict, Optional
 
 import ray
-import transfer_queue as tq
 from fastapi import FastAPI
 from ray import serve
 
@@ -17,6 +16,7 @@
 from relax.engine.sft.runtime import is_sft_mode, sft_partition_id, sft_task_name
 from relax.utils.async_utils import run
 from relax.utils.opd.opd_utils import set_managed_opd_teacher_on_train_group
+from relax.utils.tq_lifecycle import attach_tq_client
 
 
 app = FastAPI()
@@ -71,8 +71,12 @@ def __init__(
 
         self.actor_model = allocate_train_group(args=config, num_gpus=num_gpus, pg=pgs, runtime_env=runtime_env)
 
-        tq.init(self.config.tq_config)
-        self.data_system_client = tq.get_client()
+        self.data_system_client = attach_tq_client(
+            self.config.tq_config,
+            requested_gdr=getattr(self.config, "tq_use_gdr", False),
+            role=self.role,
+            lease_owner=self,
+        )
 
         self.steps = ray.get(
             self.actor_model.async_init(

relax/components/actor_fwd.py

@@ -6,12 +6,12 @@
 from typing import Any, Optional
 
 import ray
-import transfer_queue as tq
 from fastapi import FastAPI
 from ray import serve
 
 from relax.components.base import Base
 from relax.distributed.ray.placement_group import allocate_train_group
+from relax.utils.tq_lifecycle import attach_tq_client
 
 
 app = FastAPI()
@@ -35,8 +35,12 @@ def __init__(
         self._run_thread = None
         self._done_event: Optional[asyncio.Event] = None
         self._thread_error: Optional[Exception] = None
-        tq.init(self.config.tq_config)
-        self.data_system_client = tq.get_client()
+        self.data_system_client = attach_tq_client(
+            self.config.tq_config,
+            requested_gdr=getattr(self.config, "tq_use_gdr", False),
+            role=self.role,
+            lease_owner=self,
+        )
         self.actor_model = allocate_train_group(args=config, num_gpus=num_gpus, pg=pgs, runtime_env=runtime_env)
         ray.get(self.actor_model.async_init(config, role=self.role, with_ref=False))
         self.step = 0

relax/components/advantages.py

@@ -6,7 +6,6 @@
 from typing import Any, Dict
 
 import torch
-import transfer_queue as tq
 from megatron.core import mpu
 from ray import serve
 from tensordict import TensorDict
@@ -17,6 +16,7 @@
     apply_opd_to_advantages,
     consume_opd_advantage_data,
 )
+from relax.utils.tq_lifecycle import attach_tq_client
 from relax.utils.training.ppo_utils import (
     compute_approx_kl,
     get_advantages_and_returns_batch,
@@ -39,8 +39,12 @@ def __init__(
         self._lock = threading.RLock()
         self.healthy = healthy
 
-        tq.init(self.config.tq_config)
-        self.data_system_client = tq.get_client()
+        self.data_system_client = attach_tq_client(
+            self.config.tq_config,
+            requested_gdr=getattr(self.config, "tq_use_gdr", False),
+            role="advantages",
+            lease_owner=self,
+        )
         self.step = 0
 
     async def run(self) -> None:

relax/components/base.py

@@ -89,6 +89,24 @@ def __init__(self) -> None:
         self.step = 0
         self._logger_instance = None
         self._lock = threading.Lock()
+        self._tq_client_generation: int | None = None
+
+    def __del__(self) -> None:
+        # Ray Serve calls the destructor on replica shutdown (normal stop,
+        # global restart, in-place restart). Components that attached a
+        # TransferQueue client must detach so a MooncakeStore segment
+        # deregisters before client_ttl instead of leaving a stale endpoint.
+        generation = getattr(self, "_tq_client_generation", None)
+        if getattr(self, "data_system_client", None) is None or generation is None:
+            return
+        try:
+            from relax.utils.tq_lifecycle import detach_tq_client
+
+            detach_tq_client(generation)
+            self._tq_client_generation = None
+            self.data_system_client = None
+        except Exception:  # destructor must never raise (interpreter shutdown)
+            return
 
     @property
     def _logger(self):

relax/components/critic.py

@@ -7,7 +7,6 @@
 from typing import Any, Optional
 
 import ray
-import transfer_queue as tq
 from ray import serve
 from ray.serve.schema import LoggingConfig
 
@@ -16,6 +15,7 @@
 from relax.distributed.ray.placement_group import allocate_train_group
 from relax.engine.sft.runtime import sft_partition_id
 from relax.utils.async_utils import run
+from relax.utils.tq_lifecycle import attach_tq_client
 
 
 @serve.deployment(
@@ -40,8 +40,12 @@ def __init__(
         self.healthy = healthy
         self.role = role
 
-        tq.init(self.config.tq_config)
-        self.data_system_client = tq.get_client()
+        self.data_system_client = attach_tq_client(
+            self.config.tq_config,
+            requested_gdr=getattr(self.config, "tq_use_gdr", False),
+            role=self.role,
+            lease_owner=self,
+        )
 
         self.critic_model = allocate_train_group(
             args=config, num_gpus=num_gpus, pg=pgs, role=self.role, runtime_env=runtime_env

relax/components/rollout.py

@@ -9,7 +9,6 @@
 
 import httpx
 import ray
-import transfer_queue as tq
 from fastapi import FastAPI, HTTPException, Request
 from fastapi.responses import StreamingResponse
 from pydantic import BaseModel, Field
@@ -20,6 +19,7 @@
 from relax.distributed.ray.placement_group import create_rollout_manager
 from relax.utils.env import Envs
 from relax.utils.http_utils import _wrap_ipv6
+from relax.utils.tq_lifecycle import attach_tq_client
 
 
 app = FastAPI()
@@ -333,8 +333,12 @@ def __init__(
         self.config = config
         self.healthy = healthy
 
-        tq.init(self.config.tq_config)
-        self.data_system_client = tq.get_client()
+        self.data_system_client = attach_tq_client(
+            self.config.tq_config,
+            requested_gdr=getattr(self.config, "tq_use_gdr", False),
+            role="rollout",
+            lease_owner=self,
+        )
         self.rollout_manager, self.num_rollout_per_epoch = create_rollout_manager(
             config, pg, data_source=data_source, runtime_env=runtime_env
         )

relax/components/sft.py

@@ -29,7 +29,6 @@
 import random
 from typing import Any
 
-import transfer_queue as tq
 from ray import serve
 from transformers import AutoConfig, AutoTokenizer
 
@@ -39,6 +38,7 @@
 from relax.utils.data.processor_pool import ProcessorPool
 from relax.utils.misc import load_function
 from relax.utils.s3_model_loader import prepare_model_maybe_update_args
+from relax.utils.tq_lifecycle import attach_tq_client
 from relax.utils.training.eval_config import build_named_prompt_data_configs
 from relax.utils.utils import dict_to_tensordict
 
@@ -81,8 +81,12 @@ def __init__(self, healthy, pgs, num_gpus, config, role, runtime_env=None):  # n
         self.healthy = healthy
         self.step = getattr(config, "start_rollout_id", 0)
 
-        tq.init(self.config.tq_config)
-        self.data_system_client = tq.get_client()
+        self.data_system_client = attach_tq_client(
+            self.config.tq_config,
+            requested_gdr=getattr(self.config, "tq_use_gdr", False),
+            role=self.role,
+            lease_owner=self,
+        )
 
         self._dataset: Any | None = None
         self._eval_dataset: Any | None = None