diff --git a/resources_servers/code_gen/README.md b/resources_servers/code_gen/README.md index fc0f0dff43..197f226316 100644 --- a/resources_servers/code_gen/README.md +++ b/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 diff --git a/resources_servers/code_gen/app.py b/resources_servers/code_gen/app.py index d6787203db..dbfa1e0352 100644 --- a/resources_servers/code_gen/app.py +++ b/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 diff --git a/resources_servers/code_gen/configs/code_gen.yaml b/resources_servers/code_gen/configs/code_gen.yaml index 0006cd667e..0e8ffe6c0d 100644 --- a/resources_servers/code_gen/configs/code_gen.yaml +++ b/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: diff --git a/resources_servers/code_gen/lcb_integration/compute_code_generation_metrics.py b/resources_servers/code_gen/lcb_integration/compute_code_generation_metrics.py index 5d9b09a179..f6420a2875 100644 --- a/resources_servers/code_gen/lcb_integration/compute_code_generation_metrics.py +++ b/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): diff --git a/resources_servers/code_gen/tests/test_app.py b/resources_servers/code_gen/tests/test_app.py index a0c82ca3bc..6f9f267f86 100644 --- a/resources_servers/code_gen/tests/test_app.py +++ b/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( diff --git a/resources_servers/code_gen/tests/test_compute_code_generation_metrics.py b/resources_servers/code_gen/tests/test_compute_code_generation_metrics.py index 93da0dbec1..0a677c5259 100644 --- a/resources_servers/code_gen/tests/test_compute_code_generation_metrics.py +++ b/resources_servers/code_gen/tests/test_compute_code_generation_metrics.py @@ -6,11 +6,15 @@ # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 -"""Tests for `check_correctness` worker + Manager lifecycle.""" +"""Tests for `check_correctness` worker and direct IPC lifecycle.""" from __future__ import annotations import json +import os +import pickle +import sys +import time from unittest.mock import MagicMock import pytest @@ -23,80 +27,177 @@ @pytest.fixture def patched_mp(monkeypatch): - """Replace multiprocessing.Manager + Process so we can assert lifecycle ordering.""" - manager: MagicMock = MagicMock(name="Manager") - manager.list.side_effect = lambda: [] - manager_factory: MagicMock = MagicMock(return_value=manager) + """Replace the pipe and process so lifecycle ordering is observable.""" + result_connection: MagicMock = MagicMock(name="result_connection") + child_connection: MagicMock = MagicMock(name="child_connection") + pipe_factory: MagicMock = MagicMock(return_value=(result_connection, child_connection)) process_instance: MagicMock = MagicMock(name="Process") process_factory: MagicMock = MagicMock(return_value=process_instance) - monkeypatch.setattr(compute_code_generation_metrics.multiprocessing, "Manager", manager_factory) + monkeypatch.setattr(compute_code_generation_metrics.multiprocessing, "Pipe", pipe_factory) monkeypatch.setattr(compute_code_generation_metrics.multiprocessing, "Process", process_factory) - return manager, manager_factory, process_instance, process_factory + return result_connection, child_connection, pipe_factory, process_instance, process_factory class TestCheckCorrectnessReap: - """check_correctness must reap its worker and shut down its Manager on every exit path.""" + """check_correctness must close IPC and reap its worker on every exit path.""" def test_kill_is_followed_by_reap_join(self, patched_mp): - manager, _manager_factory, process, _process_factory = patched_mp - process.is_alive.side_effect = [True, False] + result_connection, child_connection, _pipe_factory, process, _process_factory = patched_mp + result_connection.poll.return_value = False + process.is_alive.return_value = True result, metadata = check_correctness(_SAMPLE, generation="ignored", timeout=1, debug=False) process.start.assert_called_once() - assert process.join.call_count == 2 process.kill.assert_called_once() - first_join_timeout = process.join.call_args_list[0].kwargs.get("timeout") - second_join_timeout = process.join.call_args_list[1].kwargs.get("timeout") - assert first_join_timeout is not None - assert second_join_timeout is not None + process.join.assert_called_once_with(timeout=5) + result_connection.close.assert_called_once() + assert child_connection.close.call_count == 2 assert result == [-1, -1] assert metadata is None - manager.shutdown.assert_called_once() - def test_manager_shutdown_runs_when_process_start_raises(self, patched_mp): - manager, _manager_factory, process, _process_factory = patched_mp + def test_connections_close_when_process_start_raises(self, patched_mp): + result_connection, child_connection, _pipe_factory, process, _process_factory = patched_mp process.start.side_effect = RuntimeError("boom") + process.is_alive.return_value = False with pytest.raises(RuntimeError, match="boom"): check_correctness(_SAMPLE, generation="ignored", timeout=1, debug=False) - manager.shutdown.assert_called_once() - - def test_happy_path_drains_results_before_manager_shutdown(self, patched_mp): - manager, _manager_factory, process, _process_factory = patched_mp + result_connection.close.assert_called_once() + child_connection.close.assert_called_once() + + def test_happy_path_receives_result_before_reaping(self, patched_mp): + result_connection, child_connection, _pipe_factory, process, _process_factory = patched_mp + result_connection.poll.return_value = True + result_connection.recv_bytes.return_value = pickle.dumps( + { + "version": compute_code_generation_metrics._WORKER_RESULT_VERSION, + "result": [1, 1], + "metadata": {"ok": True}, + } + ) process.is_alive.return_value = False - captured_lists: list[list] = [] + result, metadata = check_correctness(_SAMPLE, generation="ignored", timeout=1, debug=False) + + assert result == [1, 1] + assert metadata == {"ok": True} + result_connection.recv_bytes.assert_called_once_with( + maxlength=compute_code_generation_metrics._DEFAULT_RESULT_MAX_BYTES + ) + process.join.assert_called_once_with(timeout=5) + assert child_connection.close.call_count == 2 + result_connection.close.assert_called_once() - def _make_list() -> list: - new_list: list = [] - captured_lists.append(new_list) - return new_list + @pytest.mark.parametrize("sample", [{}, {"input_output": "not-json"}, {"input_output": {"outputs": []}}]) + def test_invalid_input_output_short_circuits(self, patched_mp, sample): + _result_connection, _child_connection, pipe_factory, process, _process_factory = patched_mp - manager.list.side_effect = _make_list + result, metadata = check_correctness(sample, generation="g", timeout=1, debug=False) - def _start_side_effect() -> None: - assert len(captured_lists) >= 2 - captured_lists[0].append([1, 1]) - captured_lists[1].append({"ok": True}) + assert result == [-1] + assert metadata is None + pipe_factory.assert_not_called() + process.start.assert_not_called() - process.start.side_effect = _start_side_effect + @pytest.mark.parametrize( + "side_effect", + [EOFError(), OSError("oversized"), pickle.UnpicklingError("malformed")], + ids=["eof", "oversized", "malformed"], + ) + def test_invalid_child_payload_fails_closed(self, patched_mp, side_effect): + result_connection, _child_connection, _pipe_factory, process, _process_factory = patched_mp + result_connection.poll.return_value = True + if isinstance(side_effect, pickle.UnpicklingError): + result_connection.recv_bytes.return_value = b"not-pickle" + else: + result_connection.recv_bytes.side_effect = side_effect + process.is_alive.return_value = False result, metadata = check_correctness(_SAMPLE, generation="ignored", timeout=1, debug=False) - assert result == [1, 1] - assert metadata == {"ok": True} - manager.shutdown.assert_called_once() + assert result == [-1, -1] + assert metadata is None + process.join.assert_called_once_with(timeout=5) + + def test_global_timeout_caps_input_scaled_backstop(self, patched_mp): + result_connection, _child_connection, _pipe_factory, process, _process_factory = patched_mp + result_connection.poll.return_value = False + process.is_alive.return_value = True + sample = {"input_output": {"inputs": ["1"] * 50, "outputs": ["1"] * 50}} + + check_correctness( + sample, + generation="ignored", + timeout=10, + debug=False, + global_timeout_seconds=20, + ) + + result_connection.poll.assert_called_once_with(20) + + def test_worker_interruption_still_kills_and_reaps_child(self, patched_mp): + result_connection, _child_connection, _pipe_factory, process, _process_factory = patched_mp + result_connection.poll.side_effect = KeyboardInterrupt + process.is_alive.return_value = True + + with pytest.raises(KeyboardInterrupt): + check_correctness(_SAMPLE, generation="ignored", timeout=1, debug=False) - def test_invalid_input_output_short_circuits(self, patched_mp): - _manager, manager_factory, process, _process_factory = patched_mp + process.kill.assert_called_once() + process.join.assert_called_once_with(timeout=5) + result_connection.close.assert_called_once() - result, metadata = check_correctness({"input_output": "not-json"}, generation="g", timeout=1, debug=False) - assert result == [-1] +@pytest.mark.skipif(sys.platform != "linux", reason="requires fork and Linux procfs") +def test_direct_ipc_preserves_result_parity_and_leaves_no_descendants(monkeypatch): + children_path = f"/proc/{os.getpid()}/task/{os.getpid()}/children" + + def child_pids(): + with open(children_path) as stream: + return set(stream.read().split()) + + baseline_children = child_pids() + baseline_fd_count = len(os.listdir("/proc/self/fd")) + monkeypatch.setattr( + compute_code_generation_metrics, + "run_test", + lambda *args, **kwargs: ([True, False], {"error_code": 0, "detail": "complete"}), + ) + + for sample in (_SAMPLE, {"input_output": json.loads(_SAMPLE["input_output"])}): + result, metadata = check_correctness(sample, generation="ignored", timeout=1, debug=False) + assert result == [True, False] + assert metadata == {"error_code": 0, "detail": "complete"} + + time.sleep(0.05) + assert child_pids() == baseline_children + assert len(os.listdir("/proc/self/fd")) <= baseline_fd_count + + +@pytest.mark.skipif(sys.platform != "linux", reason="requires fork and Linux procfs") +def test_repeated_child_crashes_leave_no_descendants_or_file_descriptors(monkeypatch): + children_path = f"/proc/{os.getpid()}/task/{os.getpid()}/children" + + def child_pids(): + with open(children_path) as stream: + return set(stream.read().split()) + + def crash(*args, **kwargs): + os._exit(7) + + baseline_children = child_pids() + baseline_fd_count = len(os.listdir("/proc/self/fd")) + monkeypatch.setattr(compute_code_generation_metrics, "_temp_run", crash) + + for _ in range(16): + result, metadata = check_correctness(_SAMPLE, generation="ignored", timeout=1, debug=False) + assert result == [-1, -1] assert metadata is None - manager_factory.assert_not_called() - process.start.assert_not_called() + + time.sleep(0.05) + assert child_pids() == baseline_children + assert len(os.listdir("/proc/self/fd")) <= baseline_fd_count