fix(code-gen): replace Manager proxies with bounded IPC

resources_servers/code_gen/README.md

@@ -18,7 +18,7 @@ The dataset is in preparation and the example data can be found in `data/example
 - Failed test cases result in a reward of 0.0 with detailed error information
 
 ### Test execution (for now)
-We use the LiveCodeBench execution code.
+We use the LiveCodeBench execution code. Each verification uses one fresh child process because the reliability guard mutates process-global state. The child returns one versioned payload over a one-way pipe; `unit_test_result_max_bytes` bounds that payload, and `unit_test_global_timeout_secs` caps the input-scaled execution backstop. `num_processes` limits locally awaited Ray tasks. Cancelling an HTTP request asks Ray to interrupt its task; if cancellation delivery is delayed, that remote task remains bounded by the global timeout.
 
 ### Example of rollouts and usage
 

resources_servers/code_gen/app.py

@@ -13,13 +13,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from asyncio import Semaphore
+import asyncio
+import contextlib
 from time import time
 from typing import Any, Dict, List, Optional, Union
 
+import ray
 from lcb_integration.compute_code_generation_metrics import check_correctness_remote
 from lcb_integration.extraction_utils import LMStyle, extract_code
-from pydantic import BaseModel
+from pydantic import BaseModel, PositiveInt
 
 from nemo_gym.base_resources_server import (
     BaseResourcesServerConfig,
@@ -40,8 +42,10 @@
 # Config
 # ----------------------------
 class CompCodingResourcesServerConfig(BaseResourcesServerConfig):
-    num_processes: int
-    unit_test_timeout_secs: int
+    num_processes: PositiveInt
+    unit_test_timeout_secs: PositiveInt
+    unit_test_global_timeout_secs: PositiveInt = 600
+    unit_test_result_max_bytes: PositiveInt = 16 * 1024 * 1024
     debug: bool
     reasoning_format_penalty: float = 0.0
 
@@ -79,11 +83,21 @@ class CompCodingVerifyResponse(BaseVerifyResponse):
 # ----------------------------
 # Server
 # ----------------------------
+async def _await_remote_result(future: Any) -> Any:
+    try:
+        return await future
+    except asyncio.CancelledError:
+        with contextlib.suppress(Exception):
+            ray.cancel(future, force=False)
+        raise
+
+
 class CompCodingResourcesServer(SimpleResourcesServer):
     config: CompCodingResourcesServerConfig
 
     def model_post_init(self, context):
-        self._semaphore: Semaphore = Semaphore(value=self.config.num_processes)
+        super().model_post_init(context)
+        self._semaphore = asyncio.Semaphore(value=self.config.num_processes)
 
     @staticmethod
     def _has_reasoning_format_violation(response) -> bool:
@@ -197,14 +211,16 @@ async def verify(self, body: CompCodingVerifyRequest) -> CompCodingVerifyRespons
             start_time = time()
 
             task_args = (
-                {"input_output": tests.model_dump_json()},  # sample
+                {"input_output": tests.model_dump(mode="json")},  # sample
                 code,  # generation
                 self.config.unit_test_timeout_secs,  # timeout
                 self.config.debug,  # debug
+                self.config.unit_test_global_timeout_secs,
+                self.config.unit_test_result_max_bytes,
             )
 
             future = check_correctness_remote.remote(*task_args)
-            result, metadata = await future
+            result, metadata = await _await_remote_result(future)
 
             unit_tests_time_taken = time() - start_time
 

resources_servers/code_gen/configs/code_gen.yaml

@@ -9,6 +9,10 @@ code_gen:
       value: Improve competitive coding capabilities
       num_processes: 8
       unit_test_timeout_secs: 10
+      # Cap the full child lifetime independently of the number of submitted tests.
+      unit_test_global_timeout_secs: 600
+      # Maximum pickled result and metadata payload accepted from the isolated child.
+      unit_test_result_max_bytes: 16777216
       debug: false
 code_gen_simple_agent:
   responses_api_agents:

resources_servers/code_gen/lcb_integration/compute_code_generation_metrics.py

@@ -19,6 +19,7 @@
 import json
 import multiprocessing
 import os
+import pickle
 import sys
 from collections import defaultdict
 from concurrent.futures import ProcessPoolExecutor, as_completed
@@ -36,10 +37,33 @@
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
 
-def _temp_run(in_outs, generation, debug, result, metadata_list, timeout):
-    res, metadata = run_test(in_outs, test=generation, debug=debug, timeout=timeout)
-    result.append(res)
-    metadata_list.append(metadata)
+_WORKER_RESULT_VERSION = 1
+_DEFAULT_GLOBAL_TIMEOUT_SECONDS = 600
+_DEFAULT_RESULT_MAX_BYTES = 16 * 1024 * 1024
+
+
+def _temp_run(in_outs, generation, debug, result_connection, timeout, result_max_bytes):
+    try:
+        result, metadata = run_test(in_outs, test=generation, debug=debug, timeout=timeout)
+        payload = pickle.dumps(
+            {
+                "version": _WORKER_RESULT_VERSION,
+                "result": result,
+                "metadata": metadata,
+            },
+            protocol=pickle.HIGHEST_PROTOCOL,
+        )
+        if len(payload) > result_max_bytes:
+            payload = pickle.dumps(
+                {
+                    "version": _WORKER_RESULT_VERSION,
+                    "error": "result_too_large",
+                },
+                protocol=pickle.HIGHEST_PROTOCOL,
+            )
+        result_connection.send_bytes(payload)
+    finally:
+        result_connection.close()
 
 
 # Using SPREAD scheduling so that Ray assigns tasks to as many distinct nodes as possible.
@@ -58,56 +82,82 @@ def _temp_run(in_outs, generation, debug, result, metadata_list, timeout):
         "env_vars": {"PYTHONPATH": _CODE_GEN_DIR},
     },
 )
-def check_correctness_remote(sample, generation, timeout, debug=True):
+def check_correctness_remote(
+    sample,
+    generation,
+    timeout,
+    debug=True,
+    global_timeout_seconds=_DEFAULT_GLOBAL_TIMEOUT_SECONDS,
+    result_max_bytes=_DEFAULT_RESULT_MAX_BYTES,
+):
     """Ray wrapper of check_correctness for remote execution."""
-    return check_correctness(sample, generation, timeout, debug)
+    return check_correctness(
+        sample,
+        generation,
+        timeout,
+        debug,
+        global_timeout_seconds,
+        result_max_bytes,
+    )
 
 
-def check_correctness(sample, generation, timeout, debug=True):
+def check_correctness(
+    sample,
+    generation,
+    timeout,
+    debug=True,
+    global_timeout_seconds=_DEFAULT_GLOBAL_TIMEOUT_SECONDS,
+    result_max_bytes=_DEFAULT_RESULT_MAX_BYTES,
+):
     """Check correctness of code generation with a global timeout.
     The global timeout is to catch some extreme/rare cases not handled by the timeouts
     inside `run_test`"""
 
-    # Parse JSON once at the beginning to avoid multiple parsing
     try:
-        in_outs = json.loads(sample["input_output"])
-    except (ValueError, MemoryError):
+        input_output = sample["input_output"]
+        in_outs = json.loads(input_output) if isinstance(input_output, str) else input_output
+        num_inputs = len(in_outs["inputs"])
+    except (KeyError, TypeError, ValueError, MemoryError):
         return [-1], None
 
-    manager = multiprocessing.Manager()
+    result_connection, child_connection = multiprocessing.Pipe(duplex=False)
     p: multiprocessing.Process | None = None
     try:
-        result = manager.list()
-        metadata_list = manager.list()
         p = multiprocessing.Process(
             target=_temp_run,
-            args=(in_outs, generation, debug, result, metadata_list, timeout),
+            args=(in_outs, generation, debug, child_connection, timeout, result_max_bytes),
         )
         p.start()
-        p.join(timeout=(timeout + 1) * len(in_outs["inputs"]) + 5)
-        if p.is_alive():
-            p.kill()
-            # Reap the worker after SIGKILL to release joinable resources.
-            p.join(timeout=5)
-
-        # Drain ListProxy values into plain lists before Manager shutdown, since access
-        # raises once the Manager helper process exits.
-        if result:
-            result_local: list = list(result)
-            metadata_local: list = list(metadata_list)
-            return result_local[0], metadata_local[0]
+        child_connection.close()
+        join_backstop = min((timeout + 1) * num_inputs + 5, global_timeout_seconds)
+        if result_connection.poll(join_backstop):
+            try:
+                payload = result_connection.recv_bytes(maxlength=result_max_bytes)
+                message = pickle.loads(payload)
+                result = message.get("result") if isinstance(message, dict) else None
+                metadata = message.get("metadata") if isinstance(message, dict) else None
+                if (
+                    isinstance(message, dict)
+                    and message.get("version") == _WORKER_RESULT_VERSION
+                    and isinstance(result, list)
+                    and (metadata is None or isinstance(metadata, dict))
+                ):
+                    return result, metadata
+            except Exception:
+                pass
 
         if debug:
             print("global timeout")
         # consider that all tests failed
-        return [-1 for _ in range(len(in_outs["inputs"]))], None
+        return [-1 for _ in range(num_inputs)], None
     finally:
         if p is not None and p.is_alive():
-            # Defensive: reap the worker if an exception bypassed the join above.
             p.kill()
             p.join(timeout=5)
-        # Always shut down the Manager so its helper process doesn't leak under stress.
-        manager.shutdown()
+        elif p is not None:
+            p.join(timeout=5)
+        child_connection.close()
+        result_connection.close()
 
 
 def evaluate_generations_by_problem(args):

resources_servers/code_gen/tests/test_app.py

@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import asyncio
 from typing import Generator
 from unittest.mock import MagicMock
 
@@ -23,6 +24,7 @@
     CompCodingResourcesServerConfig,
     CompCodingVerifyRequest,
     CompCodingVerifyResponse,
+    _await_remote_result,
 )
 from fastapi.testclient import TestClient
 from lcb_integration.testing_util import MockStdinWithBuffer
@@ -272,6 +274,20 @@ async def test_verify_runtime_error(self, code_gen_resources_server_client: Test
         res = CompCodingVerifyResponse.model_validate(response.json())
         assert res.reward == 0.0 and res.metadata["error_message"] == "Runtime Error"
 
+    @pytest.mark.parametrize("cancel_error", [None, RuntimeError("Ray unavailable")])
+    async def test_cancelled_wait_requests_remote_ray_cancellation(self, monkeypatch, cancel_error) -> None:
+        future = asyncio.get_running_loop().create_future()
+        cancel = MagicMock(side_effect=cancel_error)
+        monkeypatch.setattr(ray, "cancel", cancel)
+        task = asyncio.create_task(_await_remote_result(future))
+        await asyncio.sleep(0)
+
+        task.cancel()
+        with pytest.raises(asyncio.CancelledError):
+            await task
+
+        cancel.assert_called_once_with(future, force=False)
+
 
 def _make_server():
     return CompCodingResourcesServer(