fix: Address CodeRabbitAI review feedback on WeightSynchronizer

saumishr · saumishr · commit 66232a653615 · 2026-05-19T22:53:37.000Z
- Validate refit_buffer_size_gb &gt; 0 in factory and _compute_buffer_size
- Safely parse NRL_REFIT_BUFFER_MEMORY_RATIO env var with error handling
- Wrap IPC and HTTP sync_weights transfer in try/finally so phase
  restoration (offload_after_refit, kv_cache rebuild) always runs
- Add 7 new tests covering validation and phase-restoration behavior

Signed-off-by: Saurabh Mishra &lt;sauramishra@nvidia.com&gt;
diff --git a/nemo_rl/weight_sync/factory.py b/nemo_rl/weight_sync/factory.py
@@ -74,6 +74,8 @@ def create_weight_synchronizer(
                 generation=generation,
             )
         elif generation_backend in (VLLM_BACKEND, MEGATRON_BACKEND):
+            if refit_buffer_size_gb is not None and refit_buffer_size_gb <= 0:
+                raise ValueError("refit_buffer_size_gb must be > 0")
             from nemo_rl.weight_sync.ipc_weight_synchronizer import (
                 IPCWeightSynchronizer,
             )
diff --git a/nemo_rl/weight_sync/http_weight_synchronizer.py b/nemo_rl/weight_sync/http_weight_synchronizer.py
@@ -61,30 +61,34 @@ def sync_weights(
         self._policy.offload_before_refit()
         self._generation.prepare_for_generation(tags=["weights"])
 
-        timer_context = (
-            timer.time("prepare_for_generation/transfer_and_update_weights")
-            if timer is not None
-            else nullcontext()
-        )
-        with timer_context:
-            sglang_url_to_gpu_uuids = (
-                self._generation.get_sglang_url_to_gpu_uuids()
+        sync_succeeded = False
+        try:
+            timer_context = (
+                timer.time("prepare_for_generation/transfer_and_update_weights")
+                if timer is not None
+                else nullcontext()
             )
-
-            flush_success = self._generation.invalidate_kv_cache()
-            if not flush_success:
-                print(
-                    "SGLang KV cache invalidation failed before weight update. "
+            with timer_context:
+                sglang_url_to_gpu_uuids = (
+                    self._generation.get_sglang_url_to_gpu_uuids()
                 )
 
-            futures_train = self._policy.stream_weights_via_http(
-                sglang_url_to_gpu_uuids=sglang_url_to_gpu_uuids,
-            )
-            ray.get(futures_train)
+                flush_success = self._generation.invalidate_kv_cache()
+                if not flush_success:
+                    print(
+                        "SGLang KV cache invalidation failed before weight update. "
+                    )
+
+                futures_train = self._policy.stream_weights_via_http(
+                    sglang_url_to_gpu_uuids=sglang_url_to_gpu_uuids,
+                )
+                ray.get(futures_train)
+            sync_succeeded = True
+        finally:
+            self._policy.offload_after_refit()
+            self._generation.prepare_for_generation(tags=["kv_cache"])
 
-        self._policy.offload_after_refit()
-        self._generation.prepare_for_generation(tags=["kv_cache"])
-        self._stale = False
+        self._stale = not sync_succeeded
 
     @property
     def is_stale(self) -> bool:
diff --git a/nemo_rl/weight_sync/ipc_weight_synchronizer.py b/nemo_rl/weight_sync/ipc_weight_synchronizer.py
@@ -72,34 +72,38 @@ def sync_weights(
         self._policy.offload_before_refit()
         self._generation.prepare_for_generation(tags=["weights"])
 
-        timer_context = (
-            timer.time("prepare_for_generation/transfer_and_update_weights")
-            if timer is not None
-            else nullcontext()
-        )
-        with timer_context:
-            buffer_size_bytes = self._compute_buffer_size()
-
-            futures_train = self._policy.stream_weights_via_ipc_zmq(
-                buffer_size_bytes=buffer_size_bytes
+        sync_succeeded = False
+        try:
+            timer_context = (
+                timer.time("prepare_for_generation/transfer_and_update_weights")
+                if timer is not None
+                else nullcontext()
             )
-            futures_inference = self._generation.update_weights_via_ipc_zmq()
+            with timer_context:
+                buffer_size_bytes = self._compute_buffer_size()
 
-            ray.get(futures_train)
-            results = ray.get(futures_inference)
-            update_success = all(
-                result for result in results if result is not None
-            )
+                futures_train = self._policy.stream_weights_via_ipc_zmq(
+                    buffer_size_bytes=buffer_size_bytes
+                )
+                futures_inference = self._generation.update_weights_via_ipc_zmq()
 
-            if not update_success:
-                raise RuntimeError(
-                    "Weight transfer failed during IPC/ZMQ sync. "
-                    "This often indicates an issue with cuda-ipc or the vLLM worker."
+                ray.get(futures_train)
+                results = ray.get(futures_inference)
+                update_success = all(
+                    result for result in results if result is not None
                 )
 
-        self._policy.offload_after_refit()
-        self._generation.prepare_for_generation(tags=["kv_cache"])
-        self._stale = False
+                if not update_success:
+                    raise RuntimeError(
+                        "Weight transfer failed during IPC/ZMQ sync. "
+                        "This often indicates an issue with cuda-ipc or the vLLM worker."
+                    )
+            sync_succeeded = True
+        finally:
+            self._policy.offload_after_refit()
+            self._generation.prepare_for_generation(tags=["kv_cache"])
+
+        self._stale = not sync_succeeded
 
     @property
     def is_stale(self) -> bool:
@@ -117,7 +121,19 @@ def shutdown(self) -> None:
 
     def _compute_buffer_size(self) -> int:
         if self._refit_buffer_size_gb is not None:
+            if self._refit_buffer_size_gb <= 0:
+                raise ValueError("refit_buffer_size_gb must be > 0")
             return self._refit_buffer_size_gb * (1024**3)
 
-        memory_ratio = os.getenv("NRL_REFIT_BUFFER_MEMORY_RATIO", "0.3")
-        return int(self._policy.get_free_memory_bytes() * float(memory_ratio))
+        memory_ratio_raw = os.getenv("NRL_REFIT_BUFFER_MEMORY_RATIO", "0.3")
+        try:
+            memory_ratio = float(memory_ratio_raw)
+        except ValueError as exc:
+            raise ValueError(
+                f"NRL_REFIT_BUFFER_MEMORY_RATIO must be a valid float, got {memory_ratio_raw!r}"
+            ) from exc
+        if memory_ratio <= 0:
+            raise ValueError(
+                f"NRL_REFIT_BUFFER_MEMORY_RATIO must be > 0, got {memory_ratio}"
+            )
+        return int(self._policy.get_free_memory_bytes() * memory_ratio)
diff --git a/tests/unit/weight_sync/test_weight_synchronizer.py b/tests/unit/weight_sync/test_weight_synchronizer.py
@@ -179,6 +179,46 @@ def test_init_communicator(self):
         policy.prepare_refit_info.assert_called_once()
         gen.prepare_refit_info.assert_called_once()
 
+    @patch("nemo_rl.weight_sync.ipc_weight_synchronizer.ray")
+    def test_phase_restoration_on_transfer_failure(self, mock_ray):
+        """offload_after_refit and kv_cache prep run even when transfer raises."""
+        mock_ray.get.side_effect = RuntimeError("IPC transfer exploded")
+        policy = _mock_policy()
+        gen = _mock_generation()
+        sync = IPCWeightSynchronizer(policy, gen)
+
+        with pytest.raises(RuntimeError, match="IPC transfer exploded"):
+            sync.sync_weights()
+
+        policy.offload_after_refit.assert_called_once()
+        gen.prepare_for_generation.assert_any_call(tags=["kv_cache"])
+        assert sync.is_stale
+
+    def test_negative_buffer_size_raises(self):
+        policy = _mock_policy()
+        gen = _mock_generation()
+        sync = IPCWeightSynchronizer(policy, gen, refit_buffer_size_gb=-1)
+        with pytest.raises(ValueError, match="refit_buffer_size_gb must be > 0"):
+            sync._compute_buffer_size()
+
+    @patch("nemo_rl.weight_sync.ipc_weight_synchronizer.ray")
+    def test_invalid_env_ratio_raises(self, mock_ray, monkeypatch):
+        monkeypatch.setenv("NRL_REFIT_BUFFER_MEMORY_RATIO", "not_a_number")
+        policy = _mock_policy()
+        gen = _mock_generation()
+        sync = IPCWeightSynchronizer(policy, gen)
+        with pytest.raises(ValueError, match="must be a valid float"):
+            sync._compute_buffer_size()
+
+    @patch("nemo_rl.weight_sync.ipc_weight_synchronizer.ray")
+    def test_zero_env_ratio_raises(self, mock_ray, monkeypatch):
+        monkeypatch.setenv("NRL_REFIT_BUFFER_MEMORY_RATIO", "0")
+        policy = _mock_policy()
+        gen = _mock_generation()
+        sync = IPCWeightSynchronizer(policy, gen)
+        with pytest.raises(ValueError, match="must be > 0"):
+            sync._compute_buffer_size()
+
 
 # ---------------------------------------------------------------------------
 # HTTPWeightSynchronizer
@@ -237,6 +277,21 @@ def test_init_communicator(self):
         policy.prepare_refit_info.assert_called_once()
         gen.prepare_refit_info.assert_called_once()
 
+    @patch("nemo_rl.weight_sync.http_weight_synchronizer.ray")
+    def test_phase_restoration_on_transfer_failure(self, mock_ray):
+        """offload_after_refit and kv_cache prep run even when transfer raises."""
+        mock_ray.get.side_effect = RuntimeError("HTTP transfer exploded")
+        policy = _mock_policy()
+        gen = _mock_generation()
+        sync = HTTPWeightSynchronizer(policy, gen)
+
+        with pytest.raises(RuntimeError, match="HTTP transfer exploded"):
+            sync.sync_weights()
+
+        policy.offload_after_refit.assert_called_once()
+        gen.prepare_for_generation.assert_any_call(tags=["kv_cache"])
+        assert sync.is_stale
+
 
 # ---------------------------------------------------------------------------
 # CollectiveWeightSynchronizer
@@ -399,3 +454,27 @@ def test_unknown_backend_raises(self):
                 generation_backend="vlllm",
                 colocated=True,
             )
+
+    def test_negative_refit_buffer_size_raises(self):
+        policy = _mock_policy()
+        gen = _mock_generation()
+        with pytest.raises(ValueError, match="refit_buffer_size_gb must be > 0"):
+            create_weight_synchronizer(
+                policy=policy,
+                generation=gen,
+                generation_backend=VLLM_BACKEND,
+                colocated=True,
+                refit_buffer_size_gb=-1,
+            )
+
+    def test_zero_refit_buffer_size_raises(self):
+        policy = _mock_policy()
+        gen = _mock_generation()
+        with pytest.raises(ValueError, match="refit_buffer_size_gb must be > 0"):
+            create_weight_synchronizer(
+                policy=policy,
+                generation=gen,
+                generation_backend=VLLM_BACKEND,
+                colocated=True,
+                refit_buffer_size_gb=0,
+            )

Original file line number	Diff line number	Diff line change
`@@ -74,6 +74,8 @@ def create_weight_synchronizer(`
`74`	`74`	`generation=generation,`
`75`	`75`	`)`
`76`	`76`	`elif generation_backend in (VLLM_BACKEND, MEGATRON_BACKEND):`
	`77`	`+ if refit_buffer_size_gb is not None and refit_buffer_size_gb <= 0:`
	`78`	`+ raise ValueError("refit_buffer_size_gb must be > 0")`
`77`	`79`	`from nemo_rl.weight_sync.ipc_weight_synchronizer import (`
`78`	`80`	`IPCWeightSynchronizer,`
`79`	`81`	`)`