|
| 1 | +# Data Designer Gym |
| 2 | + |
| 3 | +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. |
| 4 | + |
| 5 | +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). |
| 6 | + |
| 7 | +## Resource and harness columns |
| 8 | + |
| 9 | +For a two-turn Pi conversation with durable history and an agent-server restart, see the optional [persistent Pi service](https://github.com/NVIDIA-NeMo/DataDesignerPlugins/blob/main/plugins/data-designer-gym/examples/pi_sessions/README.md). It extends upstream Gym's Pi adapter in a separate runtime and works with the existing harness columns. |
| 10 | + |
| 11 | +The package also installs `gym-resource` and `gym-harness` column generators. These execute inside the DD DAG and declare dependencies on every request, session, identity, and provenance source path. |
| 12 | + |
| 13 | +```python |
| 14 | +from data_designer_gym.config import GymResourceColumnConfig, GymHarnessColumnConfig |
| 15 | + |
| 16 | +builder.add_column(GymResourceColumnConfig( |
| 17 | + name="execution", |
| 18 | + endpoint="http://gym-resource:8000", |
| 19 | + route="/execute", # An execution extension supplied by your service. |
| 20 | + row_id_column="row_id", |
| 21 | + request_columns={"database_path": "database_path", "sql": "candidate.sql"}, |
| 22 | + provenance_columns=["candidate.id", "database_variant"], |
| 23 | + max_parallel_requests=4, |
| 24 | + timeout=60, |
| 25 | +)) |
| 26 | + |
| 27 | +builder.add_column(GymHarnessColumnConfig( |
| 28 | + name="answer", |
| 29 | + evidence_column="harness_evidence", |
| 30 | + endpoint="http://gym-agent:8001", |
| 31 | + row_id_column="row_id", |
| 32 | + request_columns={"input": "question", "tools": "tools"}, |
| 33 | + provenance_columns=["candidate.id", "database_variant"], |
| 34 | + session={ |
| 35 | + "initialize": [{"endpoint": "http://gym-resource:8000", "route": "/seed_session", "request_column": "task"}], |
| 36 | + "cleanup": [{"endpoint": "http://gym-resource:8000", "route": "/cleanup"}], |
| 37 | + }, |
| 38 | +)) |
| 39 | +``` |
| 40 | + |
| 41 | +`request_columns` maps request field names to upstream paths. `request_values` adds literal JSON values. Use `request_column` instead for a preassembled JSON object, including a JSON string. Resource requests do not need `responses_create_params` or a reference answer. |
| 42 | + |
| 43 | +The resource cell is a structured envelope with `status`, `row_id`, `attempt_id`, `artifact`, failure fields, `provenance_json`, `cleanup_errors_json`, and `response_json`. Native request/response objects are stored under the DD dataset's `gym/` directory, including during `preview()`. Preview calls therefore write evidence to disk. JSON strings keep arbitrary payloads lossless and safe for Parquet, including heterogeneous SQL results. Downstream DD judges can reference `{{ execution.response_json }}`. A subsequent Gym request can resolve paths such as `execution.response_json.rows` using the normal nested source mapping. |
| 44 | + |
| 45 | +The harness cell contains only the final assistant text. Its declared evidence column links the full native response and available tool trace to the same row ID and provenance. The last output item must be a completed assistant message with text. An unset top-level status from Gym chat adapters is accepted only with that terminal message and no incomplete details or pending tools; pending tools, commentary, refusals without text, and incomplete responses produce a null answer with failure evidence. Adapters that require external tool continuation need a service-side completed-loop adapter. The plugin does not perform a caller-side tool loop. |
| 46 | + |
| 47 | +Supply a unique, stable `row_id_column`. Keep candidate ID, database variant, and source database provenance in workflow columns and list them in `provenance_columns`. Calls with repeated row IDs still get separate attempt artifacts. Prompts, model selection, SQL generation/reasoning, judges, and filtering belong in the workflow. |
| 48 | + |
| 49 | +### Lifecycle and failures |
| 50 | + |
| 51 | +Session routes differ between environments. Configure the initialization and cleanup routes supported by your deployment, in their required order. Initialization may include both resource seeding and agent setup. The `/cleanup` route above is an example, not a universal Gym API. Stateful initialization requires a cleanup configuration; stateless calls may omit `session`. |
| 52 | + |
| 53 | +Each row owns its HTTP client, cookies, and session headers. Connections are reused across that row's setup, execution, and cleanup. Cookies are explicitly forwarded between the trusted configured services. An initialization step can map a JSON token into a header with `response_headers={"X-Session-Token": "token"}`. Tokens and cookies are excluded from lifecycle artifacts. `bearer_token_env` reads a credential at generator startup and sends it only to that endpoint; another service may declare its own `bearer_token_env`. |
| 54 | + |
| 55 | +`max_parallel_requests` bounds active row lifecycles per generator, including cleanup. DD also receives service scheduling metadata. Limits are per worker/generator; aggregate concurrency across columns and workers belongs in the runner. Sharing a model endpoint does not share DD/Gym clients, model admission, retries, or usage accounting. `timeout` bounds setup plus execution, and `session.cleanup_timeout` bounds each cleanup step. Admission waits are cancellable without an execution deadline; set `queue_timeout` for a separate queue deadline. Python callers with a running event loop should use `await generator.agenerate(row)` for long calls. |
| 56 | + |
| 57 | +Failures are retained as data, with distinctions for request validation, admission timeout, execution timeout, HTTP rejection, server failure, transport failure, malformed response, unfinished harness output, cleanup, and artifact storage. HTTP success means the resource call completed. SQL errors, verification reward/mask/failure fields, and semantic rejection remain in the native response for downstream policies. The plugin does not retry state-changing requests. Cleanup runs after partial initialization, timeout, and cancellation, but a remote operation may continue after a client timeout; the service must enforce its own execution deadline. Configure a TTL or administrative cleanup for operations whose initialization fails before returning a session identifier. |
| 58 | + |
| 59 | +`failure_stage` identifies `request`, `admission`, `initialize`, `execute`, `cleanup`, or `storage`; `failure_route` identifies the configured HTTP route when applicable. Both are null on success. Initialization response bodies remain excluded because they may contain session credentials. Cleanup errors remain available separately when initialization or execution also fails. |
| 60 | + |
| 61 | +DD resume can re-execute a column whose dataset batch was not persisted, even if its remote operation already committed. Stable row IDs and attempt artifacts do not provide remote deduplication. Use the service's idempotency contract where available, and persist generated request inputs before a separate execution stage when re-generating them could change an already committed request. |
| 62 | + |
| 63 | +Cleanup failure marks the call failed and clears a harness answer; artifact-write failure also clears it. Cancellation writes attempt evidence and propagates to DD instead of returning an ordinary failed row. The workflow must decide which returned rows to keep or reject. |
| 64 | + |
| 65 | +Servers must implement their own execution restrictions, database allowlists, deadlines, and cleanup semantics. Run stateful Gym services with one worker unless the deployed revision explicitly supports more. The plugin uses direct HTTP connections (`trust_env=False`), so it does not inherit proxy environment variables. |
| 66 | + |
| 67 | +### Gym episodes and response projections |
| 68 | + |
| 69 | +Use `response_protocol="gym-episode"` on a resource column to interpret Gym's `result` / `failure` envelope. The existing `status` and `failure_kind` continue to describe the call and cleanup. The additional `episode_status`, `episode_failure_kind`, `episode_failure_reason`, `episode_failure_stage`, `episode_terminal`, `reward` and `mask_sample` describe the episode. A returned result is `completed` even when its reward is zero or its sample is masked; filtering and acceptance stay in the workflow. Malformed outcome fields produce a `response` failure, with the original response retained in the artifact. |
| 70 | + |
| 71 | +```python |
| 72 | +builder.add_column(GymResourceColumnConfig( |
| 73 | + name="episode", |
| 74 | + endpoint="http://gym-environment:8000", |
| 75 | + route="/run", |
| 76 | + row_id_column="row_id", |
| 77 | + request_column="episode_request", # A prepared Gym episode request, including identity. |
| 78 | + response_protocol="gym-episode", |
| 79 | + response_fields={"trajectory_id": "result.usersim_result.trajectory_id"}, |
| 80 | +)) |
| 81 | +``` |
| 82 | + |
| 83 | +`response_fields` maps output names to dot-separated JSON object paths. On a resource column it replaces the inline `response_json` with the projection. Missing optional paths become null; `{}` keeps an empty inline object. Omitting the option preserves the existing full resource response. The full response remains in the attempt artifact in every case. Harness columns can also set `response_fields` to add projected `response_json` to their existing evidence column without changing the answer text. |
| 84 | + |
| 85 | +This calls a whole Gym-hosted episode. It does not move UserSim's conversation loop into DD or add a sandbox client. Prepare the native episode request separately; the task exporter still uses its existing Responses-based task format. Align the column timeout with the deployed server's episode deadline and cleanup budget. |
| 86 | + |
| 87 | +### Live SQL example |
| 88 | + |
| 89 | +[execution_grounded_sql.py](https://github.com/NVIDIA-NeMo/DataDesignerPlugins/blob/main/plugins/data-designer-gym/examples/execution_grounded_sql.py) replays saved public SQL candidates against original/variant toy databases, runs a tool-assisted harness, applies ordinary DD judges, and exports accepted rows. This isolates integration parity from independently sampled generations. Replace the saved question/SQL/reasoning seed fields with normal DD generators for SDG. |
| 90 | + |
| 91 | +[toy_gym.py](https://github.com/NVIDIA-NeMo/DataDesignerPlugins/blob/main/plugins/data-designer-gym/examples/toy_gym.py) supplies the demo-only SQLite `/execute` extension and uses upstream Gym `SimpleAgent`, `SimpleModelServer`, and `ExampleMultiStepResourcesServer`. Standard BIRD/text-to-SQL verification requires reference SQL and does not replace execution-only SDG. |
| 92 | + |
| 93 | +Use the separate environments from the package README. Set `MODEL_ENDPOINT`, `MODEL_NAME`, and `MODEL_API_KEY` for an existing compatible model service in both terminals. Run from `plugins/data-designer-gym` in the source checkout: |
| 94 | + |
| 95 | +```bash |
| 96 | +# Terminal 1: |
| 97 | +.venv-gym/bin/python examples/toy_gym.py \ |
| 98 | + --database-dir /tmp/dd-gym-databases \ |
| 99 | + --model-endpoint "$MODEL_ENDPOINT" --model "$MODEL_NAME" |
| 100 | + |
| 101 | +# Terminal 2, after /stats responds on port 8765: |
| 102 | +.venv/bin/python examples/execution_grounded_sql.py \ |
| 103 | + --model-endpoint "$MODEL_ENDPOINT" --model "$MODEL_NAME" \ |
| 104 | + --output /tmp/dd-gym-smoke |
| 105 | +``` |
| 106 | + |
| 107 | +The model endpoint must support Responses (harness) and Chat Completions with structured output (DD judges). For a Chat Completions endpoint, pass `--model-adapter chat` to the demo service to use Gym's `VLLMModel` conversion adapter without deploying inference locally. Stop terminal 1 with Ctrl+C after inspecting the results and zero active sessions at `/stats`. On Slurm, service launch, placement, readiness, endpoint publication, and teardown stay in the runner; this example does not implement that orchestration. |
| 108 | + |
| 109 | +## Generic processor |
| 110 | + |
| 111 | +The processor can assemble Gym tasks directly from generated columns: |
| 112 | + |
| 113 | +```python |
| 114 | +from data_designer_gym.config import GymTaskProcessorConfig |
| 115 | + |
| 116 | +builder.add_processor( |
| 117 | + GymTaskProcessorConfig( |
| 118 | + name="gym_tasks", |
| 119 | + messages=[ |
| 120 | + {"role": "system", "content_column": "system_prompt"}, |
| 121 | + {"role": "user", "content_column": "user_query"}, |
| 122 | + ], |
| 123 | + tools_column="tools_json", |
| 124 | + task_columns={"ground_truth": "reference_plan.tool_calls"}, |
| 125 | + task_values={"environment_name": "my_environment"}, |
| 126 | + include="reference_plan.is_valid", |
| 127 | + provenance_columns=["seed_id"], |
| 128 | + ) |
| 129 | +) |
| 130 | +``` |
| 131 | + |
| 132 | +Nested paths map generated values into arbitrary task fields. A preassembled `task_column` remains available when an environment needs custom Python transformations. The processor owns only common-boundary validation, stable identity, and artifact serialization. |
| 133 | + |
| 134 | +Missing `include` values skip the row. A configured `scenario_id_column` must contain a non-missing ID. Task payloads and provenance reject non-finite or non-JSON values; use `None` for intentional JSON nulls. |
| 135 | + |
| 136 | +## From-scratch notebook |
| 137 | + |
| 138 | +The [Workplace Assistant notebook](https://github.com/NVIDIA-NeMo/DataDesignerPlugins/blob/main/plugins/data-designer-gym/examples/workplace-assistant.ipynb) rebuilds Gym's existing example from Data Designer primitives. Its seeds, prompts, schemas, and qualification stay in the notebook. `GymTaskProcessorConfig` declaratively assembles the task without a custom column. |
| 139 | + |
| 140 | +## Scenario bundle |
| 141 | + |
| 142 | +Each processor record contains a canonical scenario and its Gym-native task: |
| 143 | + |
| 144 | +```json |
| 145 | +{ |
| 146 | + "schema_version": "1", |
| 147 | + "scenario_id": "my-environment-1f4b...", |
| 148 | + "task": { |
| 149 | + "responses_create_params": {"input": []}, |
| 150 | + "ground_truth": {} |
| 151 | + }, |
| 152 | + "provenance": {"seed_id": 12} |
| 153 | +} |
| 154 | +``` |
| 155 | + |
| 156 | +Export and ingestion are available independently of the notebook: |
| 157 | + |
| 158 | +```bash |
| 159 | +data-designer-gym export scenarios.jsonl --output gym-tasks.jsonl |
| 160 | + |
| 161 | +data-designer-gym ingest \ |
| 162 | + --tasks gym-tasks.jsonl \ |
| 163 | + --rollouts rollouts.jsonl \ |
| 164 | + --failures rollouts_failures.jsonl \ |
| 165 | + --output normalized-rollouts.jsonl |
| 166 | +``` |
| 167 | + |
| 168 | +The normalized result retains the complete task and rollout and can seed Data Designer or another downstream workflow. |
| 169 | + |
| 170 | +Default scenario IDs hash the task content. Set `scenario_id_column` to a unique stable row ID when identical tasks have different provenance. Ingestion rejects conflicting tasks with the same index; identical duplicate task records are allowed. Omit `--failures` when no sidecar was produced; an explicitly supplied path must exist. |
| 171 | + |
| 172 | +Ingest also exposes `rollout_id`, `attempt`, `task_id` and `environment_server`, and maps native `_ng_failure_message` / `_ng_failure_stage` into `failure_reason` / `failure_stage`. Explicit legacy failure fields take precedence. Missing attempt IDs and failure kinds remain null rather than being inferred. Joining still requires the exported task index; native materialized-task export is separate work. |
| 173 | + |
| 174 | +## Limits and proposed follow-up |
| 175 | + |
| 176 | +The optional Pi service has its own runtime and lifecycle. A UserSim session client/probe, DD wrapper, recorder, and environment-specific cleanup shim are not shipped plugin APIs. Whole-episode HTTP calls do not establish support for a UserSim loop inside DD, sandbox sessions, or arbitrary Gym environments. Generated attempt IDs are retained in artifacts but are not automatically injected into requests. |
| 177 | + |
| 178 | +Proposed stages, with completion criteria, are documented in the [package roadmap](https://github.com/NVIDIA-NeMo/DataDesignerPlugins/blob/main/plugins/data-designer-gym/README.md#proposed-follow-up-stages): a supported UserSim session/probe and cleanup contract; shared model admission/accounting; managed worker startup/readiness/cleanup and Slurm resource planning; then one pinned sandbox/persistence workflow. These are discussion milestones, not approved commitments. |
0 commit comments