feat(cli): config_composer for agent/dataset composition (RFC M3)

fern/versions/latest/pages/reference/cli-commands.mdx

@@ -36,6 +36,7 @@ gym dataset collate           # validate and collate a dataset
 # Environments
 gym env init                  # scaffold a new resources server
 gym env resolve               # resolve and print the final merged config
+gym env compose               # swap a benchmark's agent/dataset params and dump the merged config
 gym env validate              # validate a config (no Ray, no servers) — fast pre-flight check
 gym env packages              # list packages in a server's virtual environment
 gym env test                  # test resources server(s); all of them if none is given
@@ -361,6 +362,29 @@ gym env resolve \
     ++responses_create_params.temperature=0.6
 ```
 
+### `gym env compose`
+
+Rewrite a merged benchmark config along the **agent** and **dataset** axes, then dump the result as YAML. No servers are started and no secrets are read. The agent axis swaps the composable harness, carrying over the environment's `resources_server` / `model_server` / `datasets` wiring; the dataset axis edits the benchmark dataset's run params. Model selection stays with the `--model*` flags (applied before composition).
+
+| Option | Description |
+| --- | --- |
+| `--config PATH` | Config file to load. Repeatable. |
+| `--benchmark NAME` | Compose the named benchmark config. |
+| `--agent NAME` | Agent harness to compose in (`name` or `name/variant`). Must be composable (Pattern A); a self-contained Pattern B agent is rejected. |
+| `--num-repeats N` | Override `num_repeats` on the benchmark dataset. |
+| `--prompt-config PATH` | Prompt-config YAML to set on the benchmark dataset. |
+| `--search-dir DIR` | Extra root directory to search for named components. Repeatable. |
+
+```bash
+# Swap in a different agent harness and bump rollouts per task
+gym env compose --benchmark aime24 --agent simple_agent --num-repeats 4
+
+# Compose an explicit config with an agent variant
+gym env compose \
+    --config benchmarks/aime24/config.yaml \
+    --agent langgraph_agent/rewoo_agent
+```
+
 ### `gym env validate`
 
 Validate a config without starting Ray or any server subprocess — a fast pre-flight check that catches config mistakes (missing/malformed `config_paths`, unknown server cross-references, unset mandatory `???` values, schema errors) in well under a second instead of after a Ray bootstrap. Exits `0` when valid, or `1` with a clean message (no traceback) when not. A model config is **not** required — model interpolations resolve against a dummy model; pass one (or `--model-type`) if you want it validated too.

nemo_gym/cli/env.py

@@ -905,6 +905,66 @@ def dump_config():  # pragma: no cover
     print(OmegaConf.to_yaml(global_config_dict, resolve=True))
 
 
+def compose_config():  # pragma: no cover
+    """Compose a benchmark config: swap the agent harness and/or edit benchmark dataset params, then dump.
+
+    The agent axis swaps the composable harness (carrying the environment's resources_server /
+    model_server / datasets wiring); the dataset axis edits the benchmark dataset's run params. Model
+    selection stays with the `--model*` flags. The result is dumped as YAML (no servers are started).
+
+    Examples:
+
+    ```bash
+    gym env compose --benchmark aime24-x --agent simple_agent --num-repeats 4
+    gym env compose --config <benchmark.yaml> --agent langgraph_agent/rewoo_agent
+    ```
+    """
+    from nemo_gym.agent_registry import discover_agents
+    from nemo_gym.cli.main import _asset_config_path
+    from nemo_gym.config_composer import AgentNotComposableError, ComposeRequest, ConfigComposerError, compose
+
+    # NO_MODEL injects a dummy `policy_model` so a benchmark's model cross-refs resolve without real
+    # model creds — composition is about structure; model selection stays with the `--model*` flags.
+    global_config_dict = get_global_config_dict(
+        global_config_dict_parser_config=GlobalConfigDictParserConfig(
+            initial_global_config_dict=GlobalConfigDictParserConfig.NO_MODEL_GLOBAL_CONFIG_DICT,
+        ),
+    )
+    BaseNeMoGymCLIConfig.model_validate(global_config_dict)
+
+    # Pull the compose directives out so they don't leak into the dumped config (MISSING `???` preserved).
+    raw = OmegaConf.to_container(global_config_dict, resolve=False, throw_on_missing=False)
+    agent = raw.pop("compose_agent", None)
+    num_repeats = raw.pop("compose_num_repeats", None)
+    prompt_config = raw.pop("compose_prompt_config", None)
+    merged = OmegaConf.create(raw)
+
+    def _resolve_agent_config_path(name, require_composable=False):
+        # Reuse the unified CLI's asset resolver for name->path; add only the composability guard.
+        agent_dir = name.split("/", 1)[0]
+        if require_composable:
+            entry = discover_agents().get(agent_dir)
+            if entry is not None and entry.self_contained:
+                raise AgentNotComposableError(
+                    f"Agent '{agent_dir}' is self-contained (Pattern B) and cannot be composed into an "
+                    "environment; run it with its own config instead (see `gym list agents`)."
+                )
+        return _asset_config_path("agent", name)
+
+    request = ComposeRequest(
+        agent=agent,
+        num_repeats=int(num_repeats) if num_repeats is not None else None,
+        prompt_config=prompt_config,
+    )
+    try:
+        composed = compose(merged, request, resolve_agent_config_path=_resolve_agent_config_path)
+    except (ConfigComposerError, ValueError) as e:
+        rich.print(f"[red]Error:[/red] {escape(str(e))}")
+        raise SystemExit(1)
+
+    print(OmegaConf.to_yaml(composed, resolve=False))
+
+
 @exit_cleanly_on_config_error
 def validate():
     """Validate a config without starting Ray or any server subprocess.

nemo_gym/cli/main.py

@@ -130,6 +130,18 @@ def _bool_flag(name: str, hydra_key: str, flag_help: str) -> Flag:
 MODEL_URL = _value_flag("model-url", "policy_base_url", "Model server base URL.")
 MODEL_API_KEY = _value_flag("model-api-key", "policy_api_key", "Model server API key.")
 
+# Compose flags for `gym env compose`: rewrite a merged benchmark config along the agent/dataset axes.
+# Namespaced (`compose_*`) so they don't collide with server-config keys or `eval run`'s --agent/agent_name.
+COMPOSE_AGENT = _value_flag(
+    "agent", "compose_agent", "Agent harness to compose into the environment (name or name/variant)."
+)
+COMPOSE_NUM_REPEATS = _value_flag(
+    "num-repeats", "compose_num_repeats", "Override num_repeats on the benchmark dataset."
+)
+COMPOSE_PROMPT_CONFIG = _value_flag(
+    "prompt-config", "compose_prompt_config", "Prompt-config path to set on the benchmark dataset."
+)
+
 # Shared flag: select a single resources server by name. Reused by `env test`, `env init`, and `env packages`.
 RESOURCES_SERVER = Flag(
     register=lambda p: p.add_argument("--resources-server", metavar="NAME", help="Name of the resources server."),
@@ -157,6 +169,7 @@ def _bool_flag(name: str, hydra_key: str, flag_help: str) -> Flag:
     "environment": ("environments", "", "config"),
     "resources-server": ("resources_servers", "configs", None),
     "model-type": ("responses_api_models", "configs", None),
+    "agent": ("responses_api_agents", "configs", None),
 }
 
 
@@ -412,6 +425,11 @@ def _dataset_download(args: argparse.Namespace, overrides: list[str]) -> None:
         summary="Resolve the final config from configs, flags, and overrides.",
         flags=(CONFIG,),
     ),
+    "env compose": Command(
+        target="nemo_gym.cli.env:compose_config",
+        summary="Compose a benchmark config with a swapped agent and/or dataset params; dump the result.",
+        flags=(CONFIG, BENCHMARK, SEARCH_DIR, COMPOSE_AGENT, COMPOSE_NUM_REPEATS, COMPOSE_PROMPT_CONFIG),
+    ),
     "env validate": Command(
         target="nemo_gym.cli.env:validate",
         summary="Validate a config (paths, cross-refs, ??? values, servers) fast — no Ray, no servers.",