From 46609009885e42781382bda17f0e270997d0c80f Mon Sep 17 00:00:00 2001 From: Mohit Kulkarni Date: Thu, 22 Oct 2026 15:06:53 -0400 Subject: [PATCH 1/2] Own CPU vector resources and make close idempotent --- src/bindings_cpu.cpp | 77 +++++++++++++++++++-------- tests/test_pufferlib4_lifecycle.py | 83 ++++++++++++++++++++++++++++++ 2 files changed, 138 insertions(+), 22 deletions(-) create mode 100644 tests/test_pufferlib4_lifecycle.py diff --git a/src/bindings_cpu.cpp b/src/bindings_cpu.cpp index 5ba4dc81e5..e4ec733d03 100644 --- a/src/bindings_cpu.cpp +++ b/src/bindings_cpu.cpp @@ -6,6 +6,8 @@ #define _PUFFER_STRINGIFY(x) #x #define PUFFER_STRINGIFY(x) _PUFFER_STRINGIFY(x) #include +#include +#include // vecenv.h header section gives us StaticVec, Dict, cudaStream_t typedef #include "vecenv.h" @@ -75,11 +77,29 @@ static double get_config(py::dict& kwargs, const char* key) { return kwargs[key].cast(); } -static Dict* py_dict_to_c_dict(py::dict py_dict) { - Dict* c_dict = create_dict(py_dict.size()); +static std::atomic live_dicts{0}; +static std::atomic live_vectors{0}; +struct DictDeleter { + void operator()(Dict* dict) const noexcept { + if (!dict) return; + // Keys are borrowed from the Python dictionary, never owned here. + free(dict->items); + free(dict); + --live_dicts; + } +}; +using OwnedDict = std::unique_ptr; +static OwnedDict owned_dict(int capacity) { + OwnedDict result(create_dict(capacity)); + ++live_dicts; + return result; +} +static OwnedDict py_dict_to_c_dict(py::dict py_dict) { + auto c_dict = owned_dict(py_dict.size()); for (auto item : py_dict) { const char* key = PyUnicode_AsUTF8(item.first.ptr()); - try { dict_set(c_dict, key, item.second.cast()); } + if (!key) throw py::error_already_set(); + try { dict_set(c_dict.get(), key, item.second.cast()); } catch (const py::cast_error&) {} } return c_dict; @@ -90,7 +110,17 @@ static Dict* py_dict_to_c_dict(py::dict py_dict) { // ============================================================================ struct VecEnv { - StaticVec* vec; + StaticVec* vec = nullptr; + ~VecEnv() { close(); } + void close() noexcept { + if (!vec) return; + static_vec_close(vec); + vec = nullptr; + --live_vectors; + } + void require_open() const { + if (!vec) throw std::runtime_error("VecEnv is closed"); + } int total_agents; int obs_size; int num_atns; @@ -105,13 +135,14 @@ static std::unique_ptr create_vec(py::dict args, int gpu = 0) { py::dict env_kwargs = args["env"].cast(); int total_agents = (int)get_config(vec_kwargs, "total_agents"); int num_buffers = (int)get_config(vec_kwargs, "num_buffers"); - Dict* vec_dict = py_dict_to_c_dict(vec_kwargs); - Dict* env_dict = py_dict_to_c_dict(env_kwargs); + auto vec_dict = py_dict_to_c_dict(vec_kwargs); + auto env_dict = py_dict_to_c_dict(env_kwargs); auto ve = std::make_unique(); { py::gil_scoped_release no_gil; - ve->vec = create_static_vec(total_agents, num_buffers, 0, vec_dict, env_dict); + ve->vec = create_static_vec(total_agents, num_buffers, 0, vec_dict.get(), env_dict.get()); + if (ve->vec) ++live_vectors; } ve->total_agents = total_agents; ve->obs_size = get_obs_size(); @@ -127,33 +158,31 @@ static std::unique_ptr create_vec(py::dict args, int gpu = 0) { } static void vec_reset(VecEnv& ve) { - py::gil_scoped_release no_gil; + ve.require_open(); + // Keep the GIL: another Python thread must not close during this call. static_vec_reset(ve.vec); } static void cpu_vec_step_py(VecEnv& ve, long long actions_ptr) { + ve.require_open(); memcpy(ve.vec->actions, (void*)actions_ptr, (size_t)ve.total_agents * ve.num_atns * sizeof(float)); - { - py::gil_scoped_release no_gil; - cpu_vec_step(ve.vec); - } + // The vector can be closed by another Python thread only after this call. + cpu_vec_step(ve.vec); } static py::dict vec_log(VecEnv& ve) { - Dict* out = create_dict(32); - static_vec_log(ve.vec, out); + ve.require_open(); + auto out = owned_dict(32); + static_vec_log(ve.vec, out.get()); py::dict result; for (int i = 0; i < out->size; i++) result[out->items[i].key] = out->items[i].value; - free(out->items); - free(out); return result; } static void vec_close(VecEnv& ve) { - static_vec_close(ve.vec); - ve.vec = nullptr; + ve.close(); } // ============================================================================ @@ -161,6 +190,10 @@ static void vec_close(VecEnv& ve) { // ============================================================================ PYBIND11_MODULE(_C, m) { + // Private diagnostics count binding-owned resources, not total engine memory. + m.def("_resource_counts", []() { + return py::make_tuple(live_dicts.load(), live_vectors.load()); + }); m.attr("precision_bytes") = 4; m.attr("env_name") = PUFFER_STRINGIFY(ENV_NAME); m.attr("gpu") = 0; @@ -176,12 +209,12 @@ PYBIND11_MODULE(_C, m) { .def_readonly("obs_dtype", &VecEnv::obs_dtype) .def_readonly("obs_elem_size", &VecEnv::obs_elem_size) .def_property_readonly("gpu", [](VecEnv&) { return 0; }) - .def_property_readonly("obs_ptr", [](VecEnv& ve) { return (long long)ve.vec->observations; }) - .def_property_readonly("rewards_ptr", [](VecEnv& ve) { return (long long)ve.vec->rewards; }) - .def_property_readonly("terminals_ptr", [](VecEnv& ve) { return (long long)ve.vec->terminals; }) + .def_property_readonly("obs_ptr", [](VecEnv& ve) { ve.require_open(); return (long long)ve.vec->observations; }) + .def_property_readonly("rewards_ptr", [](VecEnv& ve) { ve.require_open(); return (long long)ve.vec->rewards; }) + .def_property_readonly("terminals_ptr", [](VecEnv& ve) { ve.require_open(); return (long long)ve.vec->terminals; }) .def("reset", &vec_reset) .def("cpu_step", &cpu_vec_step_py) - .def("render", [](VecEnv& ve, int env_id) { static_vec_render(ve.vec, env_id); }) + .def("render", [](VecEnv& ve, int env_id) { ve.require_open(); static_vec_render(ve.vec, env_id); }) .def("log", &vec_log) .def("close", &vec_close); } diff --git a/tests/test_pufferlib4_lifecycle.py b/tests/test_pufferlib4_lifecycle.py new file mode 100644 index 0000000000..15f96a862c --- /dev/null +++ b/tests/test_pufferlib4_lifecycle.py @@ -0,0 +1,83 @@ +"""Exercise native lifecycle failures in child processes, containing C crashes.""" +import subprocess +import sys +import textwrap +import unittest + + +class LifecycleTests(unittest.TestCase): + def run_child(self, body): + result = subprocess.run([sys.executable, '-c', textwrap.dedent(''' +from pufferlib import _C +args = {'vec': {'total_agents': 7, 'num_buffers': 1, 'num_threads': 1}, + 'env': {'max_years': 2, 'welfare_mode': 1}} +v = _C.create_vec(args, 0) +''') + textwrap.dedent(body)], capture_output=True, text=True, timeout=30) + self.assertEqual(result.returncode, 0, result.stdout + result.stderr) + + def test_close_twice(self): + self.run_child('v.close(); v.close()') + + def test_post_close_calls(self): + for operation in ['v.reset()', 'v.cpu_step(0)', 'v.log()', + 'v.render(0)', 'v.obs_ptr', 'v.rewards_ptr', 'v.terminals_ptr']: + with self.subTest(operation=operation): + self.run_child(f''' +v.close() +try: + {operation} +except RuntimeError as exc: + assert 'closed' in str(exc) +else: + raise AssertionError('post-close operation accepted') +''') + + def test_repeated_implicit_and_explicit_close(self): + self.run_child(''' +import gc +assert _C._resource_counts() == (0, 1) +for i in range(50): + v.reset() + if i % 2: + v.close() + del v + gc.collect() + assert _C._resource_counts() == (0, 0) + v = _C.create_vec(args, 0) +v.close() +assert _C._resource_counts() == (0, 0) +''') + + def test_constructor_exception_releases_dictionaries(self): + self.run_child(''' +v.close() +for _ in range(20): + bad = dict(args, env={1: 2}) + try: + _C.create_vec(bad, 0) + except TypeError: + pass + else: + raise AssertionError('invalid key accepted') + assert _C._resource_counts() == (0, 0) +''') + + def test_repeated_training_and_trainer_close_twice(self): + self.run_child(''' +v.close() +import contextlib +import gc +import io +import runpy +from pufferlib.torch_pufferl import PuffeRL +original_close = PuffeRL.close +def close_twice(self): + original_close(self) + original_close(self) +PuffeRL.close = close_twice +for _ in range(2): + with contextlib.redirect_stdout(io.StringIO()): + runpy.run_path('scripts/smoke_pufferlib4_training.py', run_name='__main__') + gc.collect() + assert _C._resource_counts() == (0, 0) +''') From fac46771f33243a4ae0f95e1fa7c06ba51b1a4b9 Mon Sep 17 00:00:00 2001 From: Mohit Kulkarni Date: Thu, 22 Oct 2026 15:08:40 -0400 Subject: [PATCH 2/2] Run CPU lifecycle regressions in training CI --- .github/workflows/install.yml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.github/workflows/install.yml b/.github/workflows/install.yml index 38c4a59b45..1d7eb74c69 100644 --- a/.github/workflows/install.yml +++ b/.github/workflows/install.yml @@ -98,5 +98,8 @@ jobs: - name: Check final-action rewards and next-state values run: python tests/test_pufferlib4_final_action_bootstrap.py + - name: Check CPU vector ownership and close behavior + run: python -m pytest tests/test_pufferlib4_lifecycle.py -q + - name: Check actual PufferLib 4 CPU training and weight reload run: python scripts/smoke_pufferlib4_training.py