andreatnvidia · Oct 27, 2026 · Oct 27, 2026 · Oct 27, 2026 · Oct 27, 2026 · Oct 27, 2026
@@ -9,7 +9,7 @@
 # Plugins
 /plugins/data-designer-docx/ @NVIDIA-NeMo/data_designer_reviewers @mvansegbroeck
 /plugins/data-designer-github/ @NVIDIA-NeMo/data_designer_reviewers @eric-tramel
-/plugins/data-designer-gym/ @andreatnvidia
+/plugins/data-designer-gym/ @NVIDIA-NeMo/data_designer_reviewers @andreatnvidia
 /plugins/data-designer-nemotron-indic-translation/ @NVIDIA-NeMo/data_designer_reviewers @kipraveen
 /plugins/data-designer-retrieval-sdg/ @NVIDIA-NeMo/data_designer_reviewers @shan-nvidia @oliverholworthy
 /plugins/data-designer-template/ @NVIDIA-NeMo/data_designer_reviewers
@@ -31,6 +31,22 @@ jobs:
       - run: make sync
       - run: make test
 
+  test-pi-adapter:
+    runs-on: ubuntu-latest
+    timeout-minutes: 10
+    defaults:
+      run:
+        working-directory: plugins/data-designer-gym/examples/pi_sessions
+    steps:
+      - uses: actions/checkout@v4
+      - uses: astral-sh/setup-uv@v4
+        with:
+          python-version: "3.14.4"
+          enable-cache: true
+          cache-dependency-glob: plugins/data-designer-gym/examples/pi_sessions/pyproject.toml
+      - run: make install
+      - run: make test
+
   validate:
     runs-on: ubuntu-latest
     steps:

@@ -33,6 +33,56 @@
         }
       ]
     },
+    {
+      "name": "data-designer-gym",
+      "description": "Execute Gym resources and harnesses, export tasks, and ingest rollouts with Data Designer",
+      "install": {
+        "requirement": "data-designer-gym",
+        "index_url": "https://nvidia-nemo.github.io/DataDesignerPlugins/simple/"
+      },
+      "compatibility": {
+        "python": {
+          "specifier": ">=3.10"
+        },
+        "data_designer": {
+          "requirement": "data-designer>=0.9.7",
+          "specifier": ">=0.9.7",
+          "marker": null
+        }
+      },
+      "docs": {
+        "url": "https://nvidia-nemo.github.io/DataDesignerPlugins/plugins/data-designer-gym/"
+      },
+      "plugins": [
+        {
+          "name": "gym",
+          "plugin_type": "processor",
+          "entry_point": {
+            "group": "data_designer.plugins",
+            "name": "gym",
+            "value": "data_designer_gym.plugin:plugin"
+          }
+        },
+        {
+          "name": "gym-harness",
+          "plugin_type": "column-generator",
+          "entry_point": {
+            "group": "data_designer.plugins",
+            "name": "gym-harness",
+            "value": "data_designer_gym.plugin:harness_plugin"
+          }
+        },
+        {
+          "name": "gym-resource",
+          "plugin_type": "column-generator",
+          "entry_point": {
+            "group": "data_designer.plugins",
+            "name": "gym-resource",
+            "value": "data_designer_gym.plugin:resource_plugin"
+          }
+        }
+      ]
+    },
     {
       "name": "data-designer-retrieval-sdg",
       "description": "Retriever SDG toolkit with Data Designer extensions, a multi-step QA generation pipeline, CLI, and Automodel-compatible data conversion",

@@ -2,6 +2,15 @@
 
 Experimental initial integration: `data-designer-gym` columns call operator-managed Gym services, preserve outputs and attempt evidence, and feed downstream processing. Task export and rollout ingestion remain available. The operator owns service/model installation, startup, readiness, placement, and shutdown.
 
+Once 0.1.0 is published, install into an active Python >=3.10 environment:
+
+```bash
+uv pip install \
+    --default-index https://pypi.org/simple/ \
+    --index https://nvidia-nemo.github.io/DataDesignerPlugins/simple/ \
+    data-designer-gym
+```
+
 For installation, tested runtime combinations, and a local resource-to-downstream-column run with cleanup and no model calls, follow the [package README](https://github.com/NVIDIA-NeMo/DataDesignerPlugins/blob/main/plugins/data-designer-gym/README.md#installation).
 
 ## Resource and harness columns

@@ -1,3 +1,3 @@
 # Owner(s) of this plugin - used to generate the root CODEOWNERS file.
 # GitHub accepts @username, @org/team, or email format.
-* @andreatnvidia
+* @NVIDIA-NeMo/data_designer_reviewers @andreatnvidia
@@ -6,20 +6,29 @@ The operator starts, checks, and stops Gym services and any model runtime. The p
 
 ## Installation
 
+Once 0.1.0 is published, install from the NVIDIA-NeMo plugin index into an active Python environment:
+
+```bash
+uv pip install \
+    --default-index https://pypi.org/simple/ \
+    --index https://nvidia-nemo.github.io/DataDesignerPlugins/simple/ \
+    data-designer-gym
+```
+
 From a source checkout containing this plugin, start in `plugins/data-designer-gym`:
 
 ```bash
 uv venv --python 3.12 .venv
-uv pip install --python .venv/bin/python 'data-designer==0.9.5' .
+uv pip install --python .venv/bin/python 'data-designer==0.9.7' .
 ```
 
-The package declares Python >=3.10 and DD >=0.9.1. The checked combinations below are narrower than those dependency ranges; other combinations require validation.
+The package declares Python >=3.10 and DD >=0.9.7. The checked combinations below are narrower than those dependency ranges; other combinations require validation.
 
 | Check | DD runtime | Gym runtime |
 | --- | --- | --- |
-| Plugin suite | DD 0.9.4 and 0.9.5, Python 3.12.12 | Simulated HTTP responses |
-| Installed wheel and local resource example | DD 0.9.5, Python 3.12.12 | Gym `58b3240e04f1fa2061328a5d9a8ec7581262f8eb`, Python 3.14.4 |
-| Optional Pi restart smoke | DD 0.9.5, Python 3.12.12 | Same Gym revision, Python 3.14.4, Pi 1.2.0, Node 25.9.0 |
+| Plugin suite | DD 0.9.7, Python 3.10.21 and 3.12.12 | Simulated HTTP responses |
+| Installed wheel and local resource example | DD 0.9.7, Python 3.10.21 | Gym `58b3240e04f1fa2061328a5d9a8ec7581262f8eb`, Python 3.14.4 |
+| Optional Pi restart smoke | DD 0.9.7, Python 3.10.21 | Same Gym revision, Python 3.14.4, Pi 1.2.0, Node 25.9.0 |
 
 ### First local run
 

@@ -2,6 +2,15 @@
 
 Experimental initial integration: `data-designer-gym` columns call operator-managed Gym services, preserve outputs and attempt evidence, and feed downstream processing. Task export and rollout ingestion remain available. The operator owns service/model installation, startup, readiness, placement, and shutdown.
 
+Once 0.1.0 is published, install into an active Python >=3.10 environment:
+
+```bash
+uv pip install \
+    --default-index https://pypi.org/simple/ \
+    --index https://nvidia-nemo.github.io/DataDesignerPlugins/simple/ \
+    data-designer-gym
+```
+
 For installation, tested runtime combinations, and a local resource-to-downstream-column run with cleanup and no model calls, follow the [package README](https://github.com/NVIDIA-NeMo/DataDesignerPlugins/blob/main/plugins/data-designer-gym/README.md#installation).
 
 ## Resource and harness columns

@@ -114,7 +114,7 @@ Omit the plugin's per-call `session` initialization/cleanup: this service owns t
 
 This extension currently uses local Pi with read/write/edit tools. Separate workspaces are not a sandbox isolation boundary. It does not enable upstream native sandbox sessions, `/run` verification, mid-tool checkpoint restoration, or distributed routing. Native agent-session initialization is rejected so a stale session cookie cannot silently select the upstream single-activation path. Sandbox adoption requires separate validation and is not implied by inheritance from `PiAgent`.
 
-After the workflow drains, verify `/stats` reports zero active Pi processes and stop the service with Ctrl+C. Keep the configured session root for retries or later turns; the operator removes it when the run's retention period ends. Shutdown cancels active requests and closes their Pi processes. Client timeouts alone do not guarantee that the remote turn has stopped.
+After the workflow drains, verify `/stats` reports zero active Pi processes and stop the service with Ctrl+C. Its `active_requests` count includes queued turns. Keep the configured session root for retries or later turns; the operator removes it when the run's retention period ends. Shutdown cancels active and queued requests and closes their Pi processes. Client timeouts alone do not guarantee that the remote turn has stopped.
 
 Slurm service placement, readiness, resource budgeting, and teardown require a separate runner integration. This example does not supply managed startup, sandbox provisioning, or Slurm acceptance evidence.
 

@@ -53,6 +53,7 @@ class PersistentPiAgent(PiAgent):
     _store: SessionStore = PrivateAttr()
     _requests: set[asyncio.Task] = PrivateAttr(default_factory=set)
     _processes: dict[str, int] = PrivateAttr(default_factory=dict)
+    _shutting_down: bool = PrivateAttr(default=False)
 
     def model_post_init(self, context) -> None:
         super().model_post_init(context)
@@ -78,14 +79,19 @@ async def lifespan(self, app):
             raise
         if version.returncode != 0 or stdout.decode().strip() != self.config.pi_version:
             raise ValueError(f"Install Pi {self.config.pi_version} before starting this adapter")
+        self._shutting_down = False
         try:
             yield
         finally:
             requests = list(self._requests)
-            for task in requests:
-                task.cancel()
+            self.cancel_requests()
             await asyncio.gather(*requests, return_exceptions=True)
 
+    def cancel_requests(self) -> None:
+        self._shutting_down = True
+        for task in list(self._requests):
+            task.cancel()
+
     def fingerprint(self) -> str:
         configuration = self.config.model_dump(
             mode="json",
@@ -106,6 +112,8 @@ async def responses(
         self, request: Request, body: NeMoGymResponseCreateParamsNonStreaming = Body()
     ) -> NeMoGymResponse:
         """Accept the existing DD user/metadata.turn convention, returning committed retries unchanged."""
+        if self._shutting_down:
+            raise HTTPException(503, "Pi service is shutting down")
         if self._agent_session_id_from_request(request) is not None:
             raise HTTPException(409, "Use conversation IDs, not native single-activation agent sessions")
         metadata = body.metadata or {}
@@ -120,6 +128,8 @@ async def responses(
         self._requests.add(task)
         try:
             async with self.sem:
+                if self._shutting_down:
+                    raise HTTPException(503, "Pi service is shutting down")
                 with self._store.turn(body.user, number, validated.model_dump(mode="json"), self.fingerprint()) as turn:
                     if turn.receipt is not None:
                         return NeMoGymResponse.model_validate(turn.receipt["response"])
@@ -230,7 +240,27 @@ async def run(self, request: Request, body: dict = Body()):
         raise HTTPException(422, "This answer-only adapter uses /v1/responses; verification belongs to the workflow")
 
     async def stats(self) -> dict:
-        return {"active_pi_processes": len(self._processes)}
+        return {"active_pi_processes": len(self._processes), "active_requests": len(self._requests)}
+
+
+class PersistentPiServer(uvicorn.Server):
+    """Cancel owned turns before Uvicorn begins draining requests."""
+
+    def __init__(self, agent: PersistentPiAgent):
+        self.agent = agent
+        super().__init__(
+            uvicorn.Config(
+                agent.setup_webserver(),
+                host=agent.config.host,
+                port=agent.config.port,
+                workers=1,
+                timeout_graceful_shutdown=0,
+            )
+        )
+
+    def handle_exit(self, sig, frame) -> None:
+        super().handle_exit(sig, frame)
+        self.agent.cancel_requests()
 
 
 def main() -> None:
@@ -246,7 +276,7 @@ def main() -> None:
             global_config_dict=OmegaConf.load(args.gym_config) if args.gym_config else OmegaConf.create({}),
         ),
     )
-    uvicorn.run(agent.setup_webserver(), host=config.host, port=config.port, workers=1)
+    PersistentPiServer(agent).run()
 
 
 if __name__ == "__main__":

@@ -20,6 +20,7 @@ def main():
     if "--version" in sys.argv:
         print("1.2.0")
         return
+    Path("pi.pid").write_text(str(os.getpid()))
     session = (
         Path(option("--session"))
         if "--session" in sys.argv
@@ -40,6 +41,9 @@ def main():
         text = command["message"]
         if text == "hang":
             time.sleep(30)
+        if text == "wait":
+            while not Path("release-turn").exists():
+                time.sleep(0.001)
         if text == "crash":
             os._exit(7)
         user = {"role": "user", "content": [{"type": "text", "text": text}]}

@@ -4,9 +4,15 @@
 import asyncio
 import hashlib
 import json
+import os
 import shlex
 import shutil
+import signal
+import socket
+import subprocess
 import sys
+import time
+from concurrent.futures import ThreadPoolExecutor
 from pathlib import Path
 
 import httpx
@@ -54,6 +60,98 @@ def request(turn=1, text="VALUE = 1", user="run-row-1"):
     return {"user": user, "metadata": {"turn": str(turn)}, "input": text}
 
 
+@pytest.fixture
+def server(config, tmp_path):
+    with socket.socket() as port:
+        port.bind((config.host, 0))
+        config = config.model_copy(update={"port": port.getsockname()[1], "timeout": 60, "concurrency": 1})
+    path = tmp_path / "server.json"
+    path.write_text(config.model_dump_json())
+    log = tmp_path / "server.log"
+    with log.open("w") as stream:
+        process = subprocess.Popen(
+            [sys.executable, "-m", "dd_gym_pi_sessions.app", "--config", str(path)],
+            stdout=stream,
+            stderr=subprocess.STDOUT,
+        )
+    try:
+        with httpx.Client(base_url=f"http://{config.host}:{config.port}", timeout=30, trust_env=False) as client:
+            for _ in range(300):
+                assert process.poll() is None, log.read_text()
+                try:
+                    if client.get("/stats").status_code == 200:
+                        break
+                except httpx.HTTPError:
+                    pass
+                time.sleep(0.1)
+            else:
+                pytest.fail("Adapter readiness timed out")
+            yield process, client, config
+    finally:
+        if process.poll() is None:
+            process.kill()
+            process.wait()
+        for path in Path(config.session_root).glob("*/workspace/pi.pid"):
+            try:
+                os.killpg(int(path.read_text()), signal.SIGKILL)
+            except ProcessLookupError:
+                pass
+
+
+@pytest.mark.parametrize("text", ["hang", "wait"])
+def test_sigterm_cancels_active_and_queued_turns(server, text):
+    process, client, config = server
+    root = Path(config.session_root)
+    completed_request = request(user="completed")
+    completed = client.post("/v1/responses", json=completed_request)
+    assert completed.status_code == 200
+    receipt = (root / "completed/turn-1.json").read_bytes()
+    with ThreadPoolExecutor(max_workers=2) as pool:
+        active = pool.submit(client.post, "/v1/responses", json=request(text=text, user="active"))
+        pid_file = root / "active/workspace/pi.pid"
+        for _ in range(300):
+            if pid_file.exists() and client.get("/stats").json()["active_pi_processes"] == 1:
+                break
+            time.sleep(0.01)
+        assert pid_file.exists()
+        pid = int(pid_file.read_text())
+        queued = pool.submit(client.post, "/v1/responses", json=request(user="queued"))
+        for _ in range(300):
+            if client.get("/stats").json()["active_requests"] == 2:
+                break
+            time.sleep(0.01)
+        assert client.get("/stats").json()["active_requests"] == 2
+        process.terminate()
+        if text == "wait":
+            log = root.parent / "server.log"
+            for _ in range(1000):
+                if "Shutting down" in log.read_text():
+                    break
+                time.sleep(0.001)
+            assert "Shutting down" in log.read_text()
+            (root / "active/workspace/release-turn").touch()
+        try:
+            process.wait(timeout=8)
+        except subprocess.TimeoutExpired:
+            process.kill()
+            process.wait()
+            pytest.fail("Adapter did not exit within the shutdown deadline")
+        for pending in (active, queued):
+            try:
+                assert pending.result(timeout=2).status_code >= 500
+            except httpx.HTTPError:
+                pass
+    with pytest.raises(ProcessLookupError):
+        os.kill(pid, 0)
+    assert json.loads((root / "active/manifest.json").read_text())["status"] == "interrupted"
+    assert not (root / "active/turn-1.json").exists()
+    assert not (root / "active/workspace/solution.py").exists()
+    assert not (root / "queued").exists()
+    assert (root / "completed/turn-1.json").read_bytes() == receipt
+    with TestClient(agent(config).setup_webserver()) as replay:
+        assert replay.post("/v1/responses", json=completed_request).json() == completed.json()
+
+
 def test_upstream_mapping_and_restart(config):
     assert issubclass(PersistentPiAgent, PiAgent)
     first_agent = agent(config)