diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 40b6973..54fcc06 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -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 diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index b6febe6..fcf8c52 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -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: diff --git a/catalog/plugins.json b/catalog/plugins.json index 90b9940..4d403d5 100644 --- a/catalog/plugins.json +++ b/catalog/plugins.json @@ -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", diff --git a/docs/plugins/data-designer-gym/index.md b/docs/plugins/data-designer-gym/index.md index 7bca40d..6226e7e 100644 --- a/docs/plugins/data-designer-gym/index.md +++ b/docs/plugins/data-designer-gym/index.md @@ -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 diff --git a/plugins/data-designer-gym/CODEOWNERS b/plugins/data-designer-gym/CODEOWNERS index 24b97fc..61269b5 100644 --- a/plugins/data-designer-gym/CODEOWNERS +++ b/plugins/data-designer-gym/CODEOWNERS @@ -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 diff --git a/plugins/data-designer-gym/README.md b/plugins/data-designer-gym/README.md index da91d65..bbe2b1f 100644 --- a/plugins/data-designer-gym/README.md +++ b/plugins/data-designer-gym/README.md @@ -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 diff --git a/plugins/data-designer-gym/docs/index.md b/plugins/data-designer-gym/docs/index.md index 7bca40d..6226e7e 100644 --- a/plugins/data-designer-gym/docs/index.md +++ b/plugins/data-designer-gym/docs/index.md @@ -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 diff --git a/plugins/data-designer-gym/examples/pi_sessions/README.md b/plugins/data-designer-gym/examples/pi_sessions/README.md index e832923..44b15b2 100644 --- a/plugins/data-designer-gym/examples/pi_sessions/README.md +++ b/plugins/data-designer-gym/examples/pi_sessions/README.md @@ -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. diff --git a/plugins/data-designer-gym/examples/pi_sessions/src/dd_gym_pi_sessions/app.py b/plugins/data-designer-gym/examples/pi_sessions/src/dd_gym_pi_sessions/app.py index a1ca07c..339e3bc 100644 --- a/plugins/data-designer-gym/examples/pi_sessions/src/dd_gym_pi_sessions/app.py +++ b/plugins/data-designer-gym/examples/pi_sessions/src/dd_gym_pi_sessions/app.py @@ -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__": diff --git a/plugins/data-designer-gym/examples/pi_sessions/tests/fake_pi.py b/plugins/data-designer-gym/examples/pi_sessions/tests/fake_pi.py index 860025a..56c2b57 100644 --- a/plugins/data-designer-gym/examples/pi_sessions/tests/fake_pi.py +++ b/plugins/data-designer-gym/examples/pi_sessions/tests/fake_pi.py @@ -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}]} diff --git a/plugins/data-designer-gym/examples/pi_sessions/tests/test_adapter.py b/plugins/data-designer-gym/examples/pi_sessions/tests/test_adapter.py index 1f6f203..145c630 100644 --- a/plugins/data-designer-gym/examples/pi_sessions/tests/test_adapter.py +++ b/plugins/data-designer-gym/examples/pi_sessions/tests/test_adapter.py @@ -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) diff --git a/plugins/data-designer-gym/pyproject.toml b/plugins/data-designer-gym/pyproject.toml index e53b6cf..c1fdfde 100644 --- a/plugins/data-designer-gym/pyproject.toml +++ b/plugins/data-designer-gym/pyproject.toml @@ -7,7 +7,7 @@ version = "0.1.0" description = "Execute Gym resources and harnesses, export tasks, and ingest rollouts with Data Designer" requires-python = ">=3.10" dependencies = [ - "data-designer>=0.9.1", + "data-designer>=0.9.7", "httpx>=0.27,<1", ] license = "Apache-2.0" @@ -38,5 +38,12 @@ build-backend = "hatchling.build" [tool.hatch.build.targets.wheel] packages = ["src/data_designer_gym"] +[tool.hatch.build.targets.sdist] +exclude = [ + "/.gym-demo", + "/examples/pi_sessions/.pytest-tmp", + "/examples/pi_sessions/.runtime", +] + [tool.ruff] extend = "../../pyproject.toml" diff --git a/plugins/data-designer-gym/src/data_designer_gym/columns.py b/plugins/data-designer-gym/src/data_designer_gym/columns.py index 3f010a1..93888e3 100644 --- a/plugins/data-designer-gym/src/data_designer_gym/columns.py +++ b/plugins/data-designer-gym/src/data_designer_gym/columns.py @@ -8,7 +8,9 @@ import json import logging import os +from http.cookiejar import parse_ns_headers from threading import BoundedSemaphore +from time import time from typing import Any from uuid import uuid4 @@ -358,6 +360,30 @@ async def _post(self, client, step, body, cookies, headers, *, allow_empty=False token_env = self.config.bearer_token_env auth = {"Authorization": f"Bearer {self._tokens[token_env]}"} if token_env else {} response = await client.request(step.method, endpoint + step.route, json=body, headers=headers | auth) + # Deleted cookies are absent from HTTPX's response cookie jar. + for parts in parse_ns_headers(response.headers.get_list("set-cookie")): + name, _ = parts[0] + attributes = dict(parts[1:]) + if "max-age" in attributes: + try: + expired = int(attributes["max-age"]) <= 0 + except (TypeError, ValueError): + continue + else: + expires = attributes.get("expires") + expired = expires is not None and expires <= time() + if expired: + # Let HTTPX validate the deletion's origin as a live cookie. + live = [parts[0], ("max-age", "60")] + live.extend((key, value) for key, value in parts[1:] if key not in {"max-age", "expires"}) + header = "; ".join(key if value is None else f"{key}={value}" for key, value in live) + validation = httpx.Response( + response.status_code, + headers=[(b"set-cookie", header.encode(response.headers.encoding))], + request=response.request, + ) + if name in validation.cookies: + cookies.pop(name, None) cookies.update({cookie.name: cookie.value for cookie in response.cookies.jar}) try: value = response.json() diff --git a/plugins/data-designer-gym/tests/test_columns.py b/plugins/data-designer-gym/tests/test_columns.py index 4e1bf15..f0b3c06 100644 --- a/plugins/data-designer-gym/tests/test_columns.py +++ b/plugins/data-designer-gym/tests/test_columns.py @@ -7,6 +7,7 @@ import json from dataclasses import dataclass, field from functools import partial +from http.cookies import SimpleCookie from pathlib import Path from types import SimpleNamespace @@ -53,6 +54,7 @@ class GymHTTP: seed_status: int = 200 seed_delay: float = 0 seed_response: dict = field(default_factory=lambda: {"token": "row-token"}) + response_headers: dict = field(default_factory=dict) responses_by_sql: dict = field(default_factory=dict) calls: list = field(default_factory=list) active: int = 0 @@ -61,10 +63,13 @@ class GymHTTP: async def __call__(self, request: httpx.Request) -> httpx.Response: body = json.loads(request.content) self.calls.append((request.url.path, body, dict(request.headers))) + headers = self.response_headers.get(request.url.path, {}) if request.url.path == "/seed_session": await asyncio.sleep(self.seed_delay) return httpx.Response( - self.seed_status, json=self.seed_response, headers={"set-cookie": f"episode={body['id']}; Path=/"} + self.seed_status, + json=self.seed_response, + headers={"set-cookie": f"episode={body['id']}; Path=/"} | headers, ) if request.url.path == "/cleanup": await asyncio.sleep(self.cleanup_delay) @@ -82,7 +87,7 @@ async def __call__(self, request: httpx.Request) -> httpx.Response: if body.get("sql") in self.responses_by_sql: status, response = self.responses_by_sql[body["sql"]] return httpx.Response(status, json=response) - return httpx.Response(self.status, json=self.response) + return httpx.Response(self.status, json=self.response, headers=headers) finally: self.active -= 1 @@ -413,6 +418,55 @@ def test_bearer_tokens_are_scoped_to_step_endpoint(service, resources, config, r assert "authorization" not in service.calls[-1][2] +@pytest.mark.parametrize("route", ["/execute", "/session_transition"]) +@pytest.mark.parametrize( + ("cookie", "expected"), + [ + ("episode=; Max-Age=0; Path=/", None), + ("episode=; Expires=Thu, 01 Jan 1970 00:00:00 GMT; Path=/", None), + ("episode=; Max-Age=0; Domain=gym.example; Path=/", None), + ("episode=; Expires=Thu, 01 Jan 1970 00:00:00 GMT; Domain=.GYM.EXAMPLE; Path=/", None), + ("episode=; Max-Age=0; Priority=High; Domain=gym.example; Path=/", None), + ("episode=; Expires=Thu, 01 Jan 1970 00:00:00 GMT; Priority=High; Domain=gym.example; Path=/", None), + ('episode="é"; Max-Age=0; Domain=gym.example; Path=/', None), + ("episode=replaced; Path=/", "replaced"), + ("episode=replaced; Domain=.GYM.EXAMPLE; Path=/", "replaced"), + ("other=extra; Path=/", "unchanged"), + ("episode=replaced; Domain=unrelated.example; Path=/", "unchanged"), + ("episode=; Max-Age=0; Domain=unrelated.example; Path=/", "unchanged"), + ("episode=; Expires=Thu, 01 Jan 1970 00:00:00 GMT; Domain=unrelated.example; Path=/", "unchanged"), + ("episode=; Max-Age=0; Priority=High; Domain=unrelated.example; Path=/", "unchanged"), + ( + "episode=; Expires=Thu, 01 Jan 1970 00:00:00 GMT; Priority=High; Domain=unrelated.example; Path=/", + "unchanged", + ), + ('episode="é"; Max-Age=0; Domain=unrelated.example; Path=/', "unchanged"), + ("episode=replaced; Max-Age=invalid; Path=/", "unchanged"), + ("episode=replaced; Max-Age=60; Expires=Thu, 01 Jan 1970 00:00:00 GMT; Path=/", "replaced"), + ], +) +def test_lifecycle_cookie_changes_and_row_isolation(service, resources, config, row, session, route, cookie, expected): + service.response_headers[route] = httpx.Headers({"set-cookie": cookie}, encoding="utf-8") + if route == "/session_transition": + session["initialize"].append({"endpoint": "http://transition.gym.example:9001", "route": route}) + config = GymResourceColumnConfig.model_validate( + config.model_dump() | {"endpoint": "http://resource.gym.example:8000", "session": session} + ) + generator = GymResourceColumnGenerator(config, resources) + for row_id in ("first", "second"): + service.calls.clear() + assert generator.generate({**row, "id": row_id})["execution"]["status"] == "completed" + assert "cookie" not in service.calls[0][2] + transition = next(index for index, call in enumerate(service.calls) if call[0] == route) + assert SimpleCookie(service.calls[transition][2]["cookie"])["episode"].value == row_id + for _, _, headers in service.calls[transition + 1 :]: + sent = SimpleCookie(headers.get("cookie", "")) + if expected is None: + assert "episode" not in sent + else: + assert sent["episode"].value == (row_id if expected == "unchanged" else expected) + + def test_queue_wait_does_not_use_execution_timeout(service, resources, config, row) -> None: service.delay = 0.03 config = GymResourceColumnConfig.model_validate(config.model_dump() | {"max_parallel_requests": 1, "timeout": 0.06}) diff --git a/uv.lock b/uv.lock index 86f677e..476498b 100644 --- a/uv.lock +++ b/uv.lock @@ -633,7 +633,7 @@ dependencies = [ [package.metadata] requires-dist = [ - { name = "data-designer", specifier = ">=0.9.1" }, + { name = "data-designer", specifier = ">=0.9.7" }, { name = "httpx", specifier = ">=0.27,<1" }, ]