Release CPU vector resources and make close safe

.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

src/bindings_cpu.cpp

@@ -6,6 +6,8 @@
 #define _PUFFER_STRINGIFY(x) #x
 #define PUFFER_STRINGIFY(x) _PUFFER_STRINGIFY(x)
 #include <cstring>
+#include <memory>
+#include <atomic>
 
 // 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<double>();
 }
 
-static Dict* py_dict_to_c_dict(py::dict py_dict) {
-    Dict* c_dict = create_dict(py_dict.size());
+static std::atomic<int> live_dicts{0};
+static std::atomic<int> 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<Dict, DictDeleter>;
+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<double>()); }
+        if (!key) throw py::error_already_set();
+        try { dict_set(c_dict.get(), key, item.second.cast<double>()); }
         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<VecEnv> create_vec(py::dict args, int gpu = 0) {
     py::dict env_kwargs = args["env"].cast<py::dict>();
     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<VecEnv>();
     {
         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,40 +158,42 @@ static std::unique_ptr<VecEnv> 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();
 }
 
 // ============================================================================
 // Module
 // ============================================================================
 
 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);
 }

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)
+''')