diff --git a/isaaclab_arena/evaluation/experiment_runner.py b/isaaclab_arena/evaluation/experiment_runner.py
index 71cd331fbc..965e19aab2 100644
--- a/isaaclab_arena/evaluation/experiment_runner.py
+++ b/isaaclab_arena/evaluation/experiment_runner.py
@@ -3,7 +3,6 @@
#
# SPDX-License-Identifier: Apache-2.0
-import os
from pathlib import Path
from isaaclab_arena.evaluation.arena_experiment import ArenaExperimentCfg
@@ -54,6 +53,19 @@ def _assert_camera_support_enabled(experiment_cfg: ArenaExperimentCfg, enable_ca
)
+def _assert_exact_experiment_output_directory_is_available(experiment_output_directory: Path) -> None:
+ """Check that an exact Experiment output path is missing or empty."""
+ if experiment_output_directory.exists():
+ assert (
+ experiment_output_directory.is_dir()
+ ), f"Experiment output path exists but is not a directory: '{experiment_output_directory}'"
+ existing_experiment_output_paths = list(experiment_output_directory.iterdir())
+ assert len(existing_experiment_output_paths) == 0, (
+ f"Experiment output directory '{experiment_output_directory}' is not empty. Choose another directory,"
+ " clear it, or use --output_base_dir to create a timestamped Experiment directory."
+ )
+
+
def main():
args_cli, experiment_overrides = parse_experiment_runner_args()
experiment_config_path = validate_experiment_config_path(args_cli.experiment_config)
@@ -87,9 +99,22 @@ def main():
assert legacy_experiment_config is not None, "--chunk_size currently supports only legacy JSON Experiments"
if len(legacy_experiment_config["jobs"]) > args_cli.chunk_size:
+ # TODO(alexmillane): Choose one timestamped Experiment output directory in the parent and pass that
+ # exact path to every legacy chunk worker. Each worker currently creates its own timestamped directory
+ # from --output_base_dir.
+ assert (
+ args_cli.experiment_output_directory is None
+ ), "--experiment_output_directory is not supported when --chunk_size dispatches multiple chunks"
run_legacy_json_in_chunks(args_cli, legacy_experiment_config)
return
+ if args_cli.experiment_output_directory is not None:
+ experiment_output_directory = args_cli.experiment_output_directory
+ _assert_exact_experiment_output_directory_is_available(experiment_output_directory)
+ else:
+ experiment_output_directory = Path(timestamped_run_dir(args_cli.output_base_dir))
+ experiment_output_directory.mkdir(parents=True, exist_ok=True)
+
with SimulationAppContext(args_cli):
experiment_cfg = load_arena_experiment_from_config_file(
experiment_config_path,
@@ -101,19 +126,12 @@ def main():
print(build_runs_info_table(experiment_cfg.runs.values(), []))
- # One reverse-dated output directory for the Experiment, with one subdirectory
- # per Run. Always date it so each invocation produces its own report directory.
- # TODO(alexmillane): Currently each chunk produces its own output directory.
- # We should use the same output directory for all chunks in the future.
- experiment_output_dir = Path(timestamped_run_dir(args_cli.output_base_dir))
-
if args_cli.record_viewport_video:
- os.makedirs(experiment_output_dir, exist_ok=True)
- print(f"[INFO] Video recording enabled. Videos will be saved to: {experiment_output_dir}")
+ print(f"[INFO] Video recording enabled. Videos will be saved to: {experiment_output_directory}")
results = execute_experiment(
experiment_cfg,
- output_dir=experiment_output_dir,
+ output_dir=experiment_output_directory,
record_viewport_video=args_cli.record_viewport_video,
record_camera_video=args_cli.record_camera_video,
continue_on_error=args_cli.continue_on_error,
@@ -126,7 +144,7 @@ def main():
metrics_logger.print_metrics()
# Write HTML report.
- report_path = build_report(experiment_output_dir)
+ report_path = build_report(experiment_output_directory)
if args_cli.serve_evaluation_report:
serve_until_ctrl_c(report_path.parent, args_cli.evaluation_report_port, report_path.name)
diff --git a/isaaclab_arena/evaluation/experiment_runner_cli.py b/isaaclab_arena/evaluation/experiment_runner_cli.py
index 5d2128ad2f..7718ef75cf 100644
--- a/isaaclab_arena/evaluation/experiment_runner_cli.py
+++ b/isaaclab_arena/evaluation/experiment_runner_cli.py
@@ -4,11 +4,13 @@
# SPDX-License-Identifier: Apache-2.0
import argparse
+from pathlib import Path
from isaaclab_arena.cli.isaaclab_arena_cli import get_isaaclab_arena_cli_parser
from isaaclab_arena.utils.hydra_overrides import assert_hydra_overrides
_DEFAULT_EXPERIMENT_CONFIG_PATH = "isaaclab_arena_environments/eval_jobs_configs/zero_action_jobs_config.json"
+_DEFAULT_EXPERIMENT_OUTPUT_BASE_DIRECTORY = "outputs"
def add_experiment_runner_arguments(parser: argparse.ArgumentParser) -> None:
@@ -38,15 +40,32 @@ def add_experiment_runner_arguments(parser: argparse.ArgumentParser) -> None:
default=False,
help="Record one mp4 per (env, camera, episode) from obs['camera_obs'] for each Run.",
)
- parser.add_argument(
+ # Keep existing Experiment Runner commands backward compatible:
+ # --output_base_dir writes to /.
+ # OSMO workflow tasks use --experiment_output_directory because each task
+ # must write directly to the exact {{output}} directory allocated by OSMO.
+ # TODO(cvolk): Replace these two path options with one path and an explicit
+ # timestamped-or-exact mode after existing --output_base_dir callers migrate.
+ output_directory_group = parser.add_mutually_exclusive_group()
+ output_directory_group.add_argument(
"--output_base_dir",
type=str,
- default="outputs",
+ default=_DEFAULT_EXPERIMENT_OUTPUT_BASE_DIRECTORY,
help=(
"Base directory for evaluation outputs (videos, per-episode results, report); a"
" reverse-dated Experiment subdirectory and per-Run subdirectory are added."
),
)
+ output_directory_group.add_argument(
+ "--experiment_output_directory",
+ type=Path,
+ default=None,
+ help=(
+ "Exact directory that will contain this Experiment's report and one subdirectory per Run."
+ " The directory must be missing or empty. Managed execution can use this instead of a timestamped"
+ " directory."
+ ),
+ )
parser.add_argument(
"--serve_evaluation_report",
action="store_true",
diff --git a/isaaclab_arena/tests/test_arena_experiment_config_loader.py b/isaaclab_arena/tests/test_arena_experiment_config_loader.py
index b3e04a58c5..e1d07d9763 100644
--- a/isaaclab_arena/tests/test_arena_experiment_config_loader.py
+++ b/isaaclab_arena/tests/test_arena_experiment_config_loader.py
@@ -169,6 +169,43 @@ def test_experiment_runner_rejects_yaml_chunking_before_starting_simulation(monk
experiment_runner.main()
+def test_experiment_runner_rejects_exact_output_for_multiple_legacy_chunks(monkeypatch, tmp_path):
+ monkeypatch.setattr(
+ "sys.argv",
+ [
+ "experiment_runner.py",
+ "--experiment_config",
+ str(GETTING_STARTED_JSON_PATH),
+ "--chunk_size",
+ "1",
+ "--experiment_output_directory",
+ str(tmp_path / "exact-experiment-output"),
+ ],
+ )
+
+ with pytest.raises(AssertionError, match="not supported when --chunk_size dispatches multiple chunks"):
+ experiment_runner.main()
+
+
+def test_experiment_runner_rejects_nonempty_exact_output_before_starting_simulation(monkeypatch, tmp_path):
+ exact_experiment_output_directory = tmp_path / "existing-experiment-output"
+ exact_experiment_output_directory.mkdir()
+ (exact_experiment_output_directory / "existing-result.jsonl").write_text("{}\n", encoding="utf-8")
+ monkeypatch.setattr(
+ "sys.argv",
+ [
+ "experiment_runner.py",
+ "--experiment_config",
+ str(GETTING_STARTED_YAML_PATH),
+ "--experiment_output_directory",
+ str(exact_experiment_output_directory),
+ ],
+ )
+
+ with pytest.raises(AssertionError, match="is not empty.*--output_base_dir"):
+ experiment_runner.main()
+
+
def test_legacy_json_experiment_rejects_hydra_overrides():
with pytest.raises(AssertionError, match="only for typed YAML"):
load_arena_experiment_from_config_file(
diff --git a/isaaclab_arena/tests/test_experiment_runner.py b/isaaclab_arena/tests/test_experiment_runner.py
index 7bc1ea7beb..3cda071cbc 100644
--- a/isaaclab_arena/tests/test_experiment_runner.py
+++ b/isaaclab_arena/tests/test_experiment_runner.py
@@ -33,6 +33,57 @@ def test_experiment_runner_parses_native_hydra_overrides():
]
+def test_experiment_runner_parses_timestamped_base_or_exact_output_directory(tmp_path):
+ exact_experiment_output_directory = tmp_path / "exact-experiment-output"
+ timestamped_experiment_output_base_directory = tmp_path / "timestamped-experiment-outputs"
+
+ default_arguments, default_experiment_overrides = parse_experiment_runner_args([
+ "--experiment_config",
+ "experiment.yaml",
+ ])
+ assert default_arguments.output_base_dir == "outputs"
+ assert default_arguments.experiment_output_directory is None
+ assert default_experiment_overrides == []
+
+ timestamped_output_arguments, timestamped_output_experiment_overrides = parse_experiment_runner_args([
+ "--output_base_dir",
+ str(timestamped_experiment_output_base_directory),
+ ])
+ assert timestamped_output_arguments.output_base_dir == str(timestamped_experiment_output_base_directory)
+ assert timestamped_output_arguments.experiment_output_directory is None
+ assert timestamped_output_experiment_overrides == []
+
+ exact_output_arguments, exact_output_experiment_overrides = parse_experiment_runner_args([
+ "--experiment_config",
+ "experiment.yaml",
+ "--experiment_output_directory",
+ str(exact_experiment_output_directory),
+ ])
+ assert exact_output_arguments.experiment_output_directory == exact_experiment_output_directory
+ assert exact_output_experiment_overrides == []
+
+
+def test_experiment_runner_rejects_timestamped_base_with_exact_output_directory(tmp_path):
+ """Reject mutually exclusive output directory flags in a fresh process."""
+ result = subprocess.run(
+ [
+ TestConstants.python_path,
+ f"{TestConstants.evaluation_dir}/experiment_runner.py",
+ "--output_base_dir",
+ str(tmp_path / "timestamped-experiment-outputs"),
+ "--experiment_output_directory",
+ str(tmp_path / "exact-experiment-output"),
+ ],
+ capture_output=True,
+ text=True,
+ check=False,
+ timeout=60,
+ )
+
+ assert result.returncode != 0
+ assert "argument --experiment_output_directory: not allowed with argument --output_base_dir" in result.stderr
+
+
@pytest.mark.with_subprocess
def test_experiment_runner_rejects_unknown_non_hydra_arguments():
"""Reject misspelled CLI flags in a fresh process."""
@@ -113,7 +164,7 @@ def test_experiment_runner_from_typed_yaml(tmp_path):
str(experiment_config_path),
config_option="--experiment_config",
extra_args=[
- "--output_base_dir",
+ "--experiment_output_directory",
str(tmp_path / "output"),
"runs.yaml_baseline.rollout_limit.num_steps=2",
],
@@ -123,6 +174,8 @@ def test_experiment_runner_from_typed_yaml(tmp_path):
run_row = next(line for line in result.stdout.splitlines() if "yaml_baseline" in line and "pending" in line)
run_cells = [cell.strip() for cell in run_row.split("|")[1:-1]]
assert run_cells[4] == "2"
+ assert (tmp_path / "output/index.html").is_file()
+ assert (tmp_path / "output/yaml_baseline/episode_results_rebuild0.jsonl").is_file()
@pytest.mark.with_subprocess
diff --git a/isaaclab_arena/tests/test_osmo_build_experiment_output.py b/isaaclab_arena/tests/test_osmo_build_experiment_output.py
new file mode 100644
index 0000000000..eea2d47b15
--- /dev/null
+++ b/isaaclab_arena/tests/test_osmo_build_experiment_output.py
@@ -0,0 +1,97 @@
+# Copyright (c) 2026, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md).
+# All rights reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Verify building one Experiment output from exact Experiment Runner task outputs."""
+
+import json
+from pathlib import Path
+
+import pytest
+
+from osmo.scripts.build_experiment_output import (
+ build_experiment_output,
+ collect_run_outputs_into_experiment_output,
+ load_experiment_runner_output_directories_by_run_name,
+)
+
+
+def _write_run_output(run_output_directory: Path, run_name: str, success: bool) -> None:
+ run_output_directory.mkdir(parents=True)
+ episode_result = {
+ "job_name": run_name,
+ "env_id": 0,
+ "episode_in_env": 0,
+ "success": success,
+ }
+ (run_output_directory / "episode_results_rebuild0.jsonl").write_text(
+ json.dumps(episode_result) + "\n",
+ encoding="utf-8",
+ )
+
+
+def test_loads_experiment_runner_output_directories_as_paths(tmp_path):
+ experiment_runner_output_directories_file_path = tmp_path / "experiment-runner-output-directories.json"
+ experiment_runner_output_directory = tmp_path / "experiment-runner-0-output"
+ experiment_runner_output_directories_file_path.write_text(
+ json.dumps({"first": str(experiment_runner_output_directory)}),
+ encoding="utf-8",
+ )
+
+ experiment_runner_output_directories_by_run_name = load_experiment_runner_output_directories_by_run_name(
+ experiment_runner_output_directories_file_path
+ )
+
+ assert experiment_runner_output_directories_by_run_name == {"first": experiment_runner_output_directory}
+
+
+def test_rejects_experiment_runner_output_without_the_requested_run(tmp_path):
+ experiment_runner_output_directory = tmp_path / "experiment-runner-0-output"
+ (experiment_runner_output_directory / "another-run").mkdir(parents=True)
+
+ with pytest.raises(AssertionError, match="Expected Run output directory for Run 'first'.*first"):
+ collect_run_outputs_into_experiment_output(
+ {"first": experiment_runner_output_directory},
+ tmp_path / "experiment-output",
+ )
+
+
+def test_collects_run_outputs_without_building_report(tmp_path):
+ experiment_runner_output_directory = tmp_path / "experiment-runner-0-output"
+ _write_run_output(experiment_runner_output_directory / "first", "first", True)
+ experiment_output_directory = tmp_path / "experiment-output"
+
+ collect_run_outputs_into_experiment_output(
+ {"first": experiment_runner_output_directory},
+ experiment_output_directory,
+ )
+
+ assert (experiment_output_directory / "first/episode_results_rebuild0.jsonl").is_file()
+ assert not (experiment_output_directory / "index.html").exists()
+
+
+def test_builds_experiment_output_from_separate_experiment_runner_outputs(tmp_path):
+ first_experiment_runner_output_directory = tmp_path / "experiment-runner-0-output"
+ second_experiment_runner_output_directory = tmp_path / "experiment-runner-1-output"
+ first_run_output_directory = first_experiment_runner_output_directory / "first"
+ second_run_output_directory = second_experiment_runner_output_directory / "second"
+ _write_run_output(first_run_output_directory, "first", True)
+ _write_run_output(second_run_output_directory, "second", False)
+ experiment_output_directory = tmp_path / "experiment-output"
+
+ report_path = build_experiment_output(
+ {
+ "first": first_experiment_runner_output_directory,
+ "second": second_experiment_runner_output_directory,
+ },
+ experiment_output_directory,
+ )
+
+ assert report_path == experiment_output_directory / "index.html"
+ assert (experiment_output_directory / "first/episode_results_rebuild0.jsonl").is_file()
+ assert (experiment_output_directory / "second/episode_results_rebuild0.jsonl").is_file()
+ report_contents = report_path.read_text(encoding="utf-8")
+ assert "first" in report_contents
+ assert "second" in report_contents
+ assert "2 job(s)" in report_contents
diff --git a/isaaclab_arena/tests/test_osmo_experiment_workflow.py b/isaaclab_arena/tests/test_osmo_experiment_workflow.py
index 2df838e52e..487a3a9589 100644
--- a/isaaclab_arena/tests/test_osmo_experiment_workflow.py
+++ b/isaaclab_arena/tests/test_osmo_experiment_workflow.py
@@ -3,8 +3,9 @@
#
# SPDX-License-Identifier: Apache-2.0
-"""Verify OSMO workflows for complete Arena Experiments."""
+"""Verify distributed OSMO workflows for Arena Experiments."""
+import json
import yaml
from pathlib import Path
from types import SimpleNamespace
@@ -28,11 +29,16 @@
submit_arena_experiment,
)
from osmo.tasks.base_task import TaskCfg
+from osmo.tasks.collect_experiment_outputs_task import (
+ _REMOTE_BUILD_EXPERIMENT_OUTPUT_SCRIPT_PATH,
+ _REMOTE_EXPERIMENT_RUNNER_OUTPUT_DIRECTORIES_FILE_PATH,
+ experiment_runner_output_directory_input_token,
+)
from osmo.tasks.experiment_runner_task import REMOTE_EXPERIMENT_PATH, ExperimentRunnerTask, ExperimentRunnerTaskCfg
from osmo.tasks.pi0_server_task import Pi0ServerTask, Pi0ServerTaskCfg
from osmo.workflows.arena_experiment_workflow import Pi0ArenaExperimentWorkflow
from osmo.workflows.workflow import WorkflowCfg
-from osmo.workflows.workflow_constants import POLICY_SERVER_PORT
+from osmo.workflows.workflow_constants import DATASET_SWIFT_URL, OSMO_TASK_OUTPUT_DIR, POLICY_SERVER_PORT
# Composing complete Arena Experiments loads Isaac runtime modules, so these tests
# must not share a pytest process with the persistent SimulationApp tests.
@@ -98,8 +104,12 @@ def _rendered_workflow(output: str) -> dict:
return yaml.safe_load(output[output.index("version: 2\n") :])
-def _workflow_tasks(workflow: dict) -> list[dict]:
- return workflow["workflow"]["groups"][0]["tasks"]
+def _workflow_groups(workflow: dict) -> list[dict]:
+ return workflow["workflow"]["groups"]
+
+
+def _workflow_tasks(workflow: dict, group_index: int = 0) -> list[dict]:
+ return _workflow_groups(workflow)[group_index]["tasks"]
def _compose_submission(
@@ -180,8 +190,8 @@ def test_policy_server_rejects_workflow_fields():
])
-def test_renders_experiment_runner_and_shared_pi0_server_with_effective_endpoints():
- """Wire every matching pi0 Run in the effective Experiment to one server."""
+def test_fans_out_single_run_experiments_with_dedicated_pi0_servers_and_one_experiment_output():
+ """Render one independent Run group per Run and collect their outputs into one Experiment output."""
source_experiment_cfg = _pi0_experiment_cfg()
workflow = Pi0ArenaExperimentWorkflow(
workflow_cfg=WorkflowCfg(workflow_name="pi0-experiment"),
@@ -189,35 +199,99 @@ def test_renders_experiment_runner_and_shared_pi0_server_with_effective_endpoint
server_task_cfg=Pi0ServerTaskCfg(),
)
- tasks = _workflow_tasks(workflow.generate_workflow())
- assert [task["name"] for task in tasks] == ["experiment_runner", "policy_server"]
- assert [task["lead"] for task in tasks] == [True, False]
-
- eval_task = tasks[0]
- experiment = _embedded_experiment(eval_task)
- server_host = Pi0ServerTask.host_token()
- assert experiment["runs"]["first"]["policy"]["remote_host"] == server_host
- assert experiment["runs"]["first"]["policy"]["remote_port"] == POLICY_SERVER_PORT
- assert experiment["runs"]["first"]["policy"]["ping_timeout"] == Pi0ServerTaskCfg.client_ping_timeout_s
- assert experiment["runs"]["second"]["policy"]["remote_host"] == server_host
- assert experiment["runs"]["second"]["policy"]["remote_port"] == POLICY_SERVER_PORT
- assert experiment["runs"]["second"]["policy"]["ping_timeout"] == Pi0ServerTaskCfg.client_ping_timeout_s
- assert "remote_host" not in experiment["runs"]["local"]["policy"]
- assert "remote_port" not in experiment["runs"]["local"]["policy"]
+ rendered_workflow = workflow.generate_workflow()
+ groups = _workflow_groups(rendered_workflow)
+ assert rendered_workflow == workflow.generate_workflow()
+ assert [group["name"] for group in groups] == [
+ "arena-run-0",
+ "arena-run-1",
+ "arena-run-2",
+ "arena-experiment-output",
+ ]
+ task_names = [task["name"] for group in groups for task in group["tasks"]]
+ assert len(task_names) == len(set(task_names))
+ workflow_names = [group["name"] for group in groups] + task_names
+ normalized_workflow_names = [name.lower().replace("_", "-") for name in workflow_names]
+ assert len(normalized_workflow_names) == len(set(normalized_workflow_names))
+
+ first_tasks = groups[0]["tasks"]
+ second_tasks = groups[1]["tasks"]
+ local_tasks = groups[2]["tasks"]
+ assert [task["name"] for task in first_tasks] == ["experiment-runner-0", "policy-server-0"]
+ assert [task["name"] for task in second_tasks] == ["experiment-runner-1", "policy-server-1"]
+ assert [task["name"] for task in local_tasks] == ["experiment-runner-2"]
+ assert [[task["lead"] for task in group["tasks"]] for group in groups] == [
+ [True, False],
+ [True, False],
+ [True],
+ [True],
+ ]
+
+ first_experiment = _embedded_experiment(first_tasks[0])
+ second_experiment = _embedded_experiment(second_tasks[0])
+ local_experiment = _embedded_experiment(local_tasks[0])
+ assert list(first_experiment["runs"]) == ["first"]
+ assert list(second_experiment["runs"]) == ["second"]
+ assert list(local_experiment["runs"]) == ["local"]
+ assert first_experiment["runs"]["first"]["policy"]["remote_host"] == Pi0ServerTask.host_token("policy-server-0")
+ assert second_experiment["runs"]["second"]["policy"]["remote_host"] == Pi0ServerTask.host_token("policy-server-1")
+ for run_name, experiment in (("first", first_experiment), ("second", second_experiment)):
+ policy = experiment["runs"][run_name]["policy"]
+ assert policy["remote_port"] == POLICY_SERVER_PORT
+ assert policy["ping_timeout"] == Pi0ServerTaskCfg.client_ping_timeout_s
+ assert "remote_host" not in local_experiment["runs"]["local"]["policy"]
+ assert "remote_port" not in local_experiment["runs"]["local"]["policy"]
assert source_experiment_cfg.runs["first"].policy.remote_host == "user-host"
assert source_experiment_cfg.runs["first"].policy.ping_timeout == 10
- command = _task_file(eval_task, "/tmp/entry.sh")["contents"]
- assert "experiment_runner.py" in command
- assert f"--experiment_config {REMOTE_EXPERIMENT_PATH}" in command
- assert "--enable_cameras" in command
- assert "policy_runner.py" not in command
- assert "runs." not in command
+ experiment_runner_command = _task_file(first_tasks[0], "/tmp/entry.sh")["contents"]
+ assert "experiment_runner.py" in experiment_runner_command
+ assert f"--experiment_config {REMOTE_EXPERIMENT_PATH}" in experiment_runner_command
+ assert f"--experiment_output_directory '{OSMO_TASK_OUTPUT_DIR}'" in experiment_runner_command
+ assert "--output_base_dir" not in experiment_runner_command
+ assert "--enable_cameras" in experiment_runner_command
+ assert "policy_runner.py" not in experiment_runner_command
+ assert "runs." not in experiment_runner_command
- server_command = _task_file(tasks[1], "/tmp/entry.sh")["contents"]
+ assert first_tasks[0]["outputs"] == []
+ assert second_tasks[0]["outputs"] == []
+ assert local_tasks[0]["outputs"] == []
+
+ server_command = _task_file(first_tasks[1], "/tmp/entry.sh")["contents"]
assert f"scripts/serve_policy.py --port={POLICY_SERVER_PORT} policy:checkpoint" in server_command
assert "--policy.config=pi05_droid_jointpos_polaris" in server_command
+ experiment_output_task = groups[3]["tasks"][0]
+ assert experiment_output_task["name"] == "collect-experiment-outputs"
+ assert experiment_output_task["resource"] == "experiment-output"
+ assert experiment_output_task["inputs"] == [
+ {"task": "experiment-runner-0"},
+ {"task": "experiment-runner-1"},
+ {"task": "experiment-runner-2"},
+ ]
+ assert experiment_output_task["outputs"] == [{"url": DATASET_SWIFT_URL}]
+ experiment_runner_output_directories_by_run_name = json.loads(
+ _task_file(experiment_output_task, _REMOTE_EXPERIMENT_RUNNER_OUTPUT_DIRECTORIES_FILE_PATH)["contents"]
+ )
+ assert experiment_runner_output_directories_by_run_name == {
+ "first": experiment_runner_output_directory_input_token("experiment-runner-0"),
+ "second": experiment_runner_output_directory_input_token("experiment-runner-1"),
+ "local": experiment_runner_output_directory_input_token("experiment-runner-2"),
+ }
+ experiment_output_script_file = _task_file(experiment_output_task, _REMOTE_BUILD_EXPERIMENT_OUTPUT_SCRIPT_PATH)
+ assert "localpath" not in experiment_output_script_file
+ assert "def build_experiment_output" in experiment_output_script_file["contents"]
+ experiment_output_command = _task_file(experiment_output_task, "/tmp/entry.sh")["contents"]
+ assert experiment_output_command.startswith("set -euo pipefail")
+ assert _REMOTE_BUILD_EXPERIMENT_OUTPUT_SCRIPT_PATH in experiment_output_command
+ assert (
+ f"--experiment-runner-output-directories-file {_REMOTE_EXPERIMENT_RUNNER_OUTPUT_DIRECTORIES_FILE_PATH}"
+ in experiment_output_command
+ )
+ assert "--experiment-output-directory" in experiment_output_command
+ assert OSMO_TASK_OUTPUT_DIR in experiment_output_command
+ assert rendered_workflow["workflow"]["resources"]["experiment-output"]["gpu"] == 0
+
def test_embeds_effective_experiment_yaml():
"""Embed the composed Experiment instead of staging its source file."""
@@ -225,6 +299,7 @@ def test_embeds_effective_experiment_yaml():
task_cfg=ExperimentRunnerTaskCfg(image="registry.example.com/evaluator:typed-api"),
experiment_cfg=_zero_action_experiment_cfg(),
lead=True,
+ task_name="experiment-runner",
)
eval_task = experiment_runner_task.create_task_dict()
@@ -302,7 +377,7 @@ def test_submission_composes_defaults_experiment_and_overrides(tmp_path, capsys)
workflow = _rendered_workflow(rendered)
assert workflow["workflow"]["name"] == "overridden-experiment"
tasks = _workflow_tasks(workflow)
- assert [task["name"] for task in tasks] == ["experiment_runner", "policy_server"]
+ assert [task["name"] for task in tasks] == ["experiment-runner-0", "policy-server-0"]
assert tasks[0]["image"] == "registry.example.com/evaluator:branch"
assert tasks[1]["image"] == "registry.example.com/openpi:overridden"
@@ -311,7 +386,7 @@ def test_submission_composes_defaults_experiment_and_overrides(tmp_path, capsys)
assert experiment["runs"]["openpi_maple_table"]["rollout_limit"]["num_episodes"] == 4
assert experiment["runs"]["openpi_maple_table"]["environment_builder"]["num_envs"] == 2
assert policy["ping_interval"] == 33.0
- assert policy["remote_host"] == Pi0ServerTask.host_token()
+ assert policy["remote_host"] == Pi0ServerTask.host_token("policy-server-0")
assert policy["remote_port"] == POLICY_SERVER_PORT
assert policy["ping_timeout"] == 600.0
assert "experiment_cfg.runs" not in _task_file(tasks[0], "/tmp/entry.sh")["contents"]
@@ -322,7 +397,7 @@ def test_submission_composes_defaults_experiment_and_overrides(tmp_path, capsys)
def test_embedded_openpi_experiment_composes_through_experiment_runner_loader(tmp_path):
- """Keep the rendered OSMO handoff compatible with the Experiment Runner loader."""
+ """Keep every single-Run OSMO handoff compatible with the Experiment Runner loader."""
submission_cfg = _compose_submission()
assert isinstance(submission_cfg.policy_server, Pi0ServerTaskCfg)
workflow = Pi0ArenaExperimentWorkflow(
@@ -331,18 +406,19 @@ def test_embedded_openpi_experiment_composes_through_experiment_runner_loader(tm
server_task_cfg=submission_cfg.policy_server,
task_cfg=submission_cfg.experiment_runner,
)
- experiment_path = tmp_path / "effective_experiment.yaml"
- experiment_file = _task_file(_workflow_tasks(workflow.generate_workflow())[0], REMOTE_EXPERIMENT_PATH)
- experiment_path.write_text(experiment_file["contents"], encoding="utf-8")
-
- experiment_cfg = load_arena_experiment_from_config_file(experiment_path, device="cuda:0")
- run_cfg = experiment_cfg.runs[OPENPI_RUN_NAME]
+ rendered_workflow = workflow.generate_workflow()
+ for index, run_name in enumerate(submission_cfg.experiment_cfg.runs):
+ experiment_path = tmp_path / f"effective_experiment_{index}.yaml"
+ experiment_file = _task_file(_workflow_tasks(rendered_workflow, index)[0], REMOTE_EXPERIMENT_PATH)
+ experiment_path.write_text(experiment_file["contents"], encoding="utf-8")
- assert list(experiment_cfg.runs) == list(submission_cfg.experiment_cfg.runs)
- assert isinstance(run_cfg.policy, Pi0RemotePolicyCfg)
- assert run_cfg.policy.remote_host == Pi0ServerTask.host_token()
- assert run_cfg.policy.remote_port == POLICY_SERVER_PORT
- assert run_cfg.policy.ping_timeout == Pi0ServerTaskCfg.client_ping_timeout_s
+ experiment_cfg = load_arena_experiment_from_config_file(experiment_path, device="cuda:0")
+ assert list(experiment_cfg.runs) == [run_name]
+ run_cfg = experiment_cfg.runs[run_name]
+ assert isinstance(run_cfg.policy, Pi0RemotePolicyCfg)
+ assert run_cfg.policy.remote_host == Pi0ServerTask.host_token(f"policy-server-{index}")
+ assert run_cfg.policy.remote_port == POLICY_SERVER_PORT
+ assert run_cfg.policy.ping_timeout == Pi0ServerTaskCfg.client_ping_timeout_s
def test_submission_overrides_osmo_resources(monkeypatch):
@@ -355,7 +431,7 @@ def capture_submission(command, **kwargs):
assert kwargs["text"] is True
submitted_command = command
submitted_workflow = yaml.safe_load(Path(command[3]).read_text(encoding="utf-8"))
- submitted_resources = submitted_workflow["workflow"]["resources"]["default"]
+ submitted_resources = submitted_workflow["workflow"]["resources"]
return SimpleNamespace(returncode=0, stdout="")
monkeypatch.setattr("osmo.workflows.workflow.subprocess.run", capture_submission)
@@ -370,8 +446,11 @@ def capture_submission(command, **kwargs):
assert submitted_command is not None
pool_flag_index = submitted_command.index("--pool")
assert submitted_command[pool_flag_index + 1] == "isaac-dev-l40-03"
- assert submitted_resources["platform"] == "ovx-l40"
- assert submitted_resources["memory"] == "120Gi"
+ assert submitted_resources["default"]["platform"] == "ovx-l40"
+ assert submitted_resources["default"]["memory"] == "120Gi"
+ assert submitted_resources["experiment-output"]["platform"] == "ovx-l40"
+ assert submitted_resources["experiment-output"]["memory"] == "120Gi"
+ assert submitted_resources["experiment-output"]["gpu"] == 0
def test_cli_requires_experiment_cfg_path_and_policy_server(capsys):
@@ -441,7 +520,7 @@ def test_cli_accepts_arbitrary_paths_and_trailing_overrides(tmp_path, capsys):
workflow = _rendered_workflow(capsys.readouterr().out)
assert workflow["workflow"]["name"] == "path-based-submission"
tasks = _workflow_tasks(workflow)
- assert [task["name"] for task in tasks] == ["experiment_runner", "policy_server"]
+ assert [task["name"] for task in tasks] == ["experiment-runner-0", "policy-server-0"]
assert tasks[0]["image"] == "registry.example.com/evaluator:cli"
@@ -494,7 +573,8 @@ def test_pi0_server_quotes_configurable_shell_values():
Pi0ServerTaskCfg(
policy_config="config with spaces",
policy_dir="gs://bucket/checkpoint; false",
- )
+ ),
+ task_name="policy-server",
)
command = task._get_run_script()
diff --git a/isaaclab_arena/tests/test_osmo_workflow.py b/isaaclab_arena/tests/test_osmo_workflow.py
index fc6187b535..4c66a8ce05 100644
--- a/isaaclab_arena/tests/test_osmo_workflow.py
+++ b/isaaclab_arena/tests/test_osmo_workflow.py
@@ -5,12 +5,23 @@
"""Verify typed OSMO workflow construction and its compatibility CLI."""
+import pytest
+
from osmo.submit_evaluation_workflow import main
+from osmo.tasks.dreamzero_policy_runner_task import DreamZeroPolicyRunnerTaskCfg
+from osmo.tasks.pi0_server_task import Pi0ServerTask, Pi0ServerTaskCfg
from osmo.tasks.policy_runner_task import PolicyRunnerTaskCfg
+from osmo.workflows.dreamzero_split_workflows import DreamZeroPolicyRunnerWorkflow
from osmo.workflows.server_plus_policy_runner_workflow import Pi0PlusPolicyRunnerWorkflow
from osmo.workflows.workflow import WorkflowCfg
+def test_task_name_is_a_required_keyword_argument():
+ """Reject construction when a workflow does not name its task instance."""
+ with pytest.raises(TypeError, match="task_name"):
+ Pi0ServerTask(Pi0ServerTaskCfg())
+
+
def test_typed_workflow_config_renders_policy_runner_and_server():
"""Construct a multi-task workflow without passing through argparse."""
workflow = Pi0PlusPolicyRunnerWorkflow(
@@ -34,6 +45,36 @@ def test_typed_workflow_config_renders_policy_runner_and_server():
assert "example_environment light.hdr_image.enabled=true" in policy_runner_command
+def test_static_workflow_threads_declared_task_names_into_host_token():
+ """Use the same explicit server name for its task and the runner's host token."""
+
+ class CustomNamedPi0Workflow(Pi0PlusPolicyRunnerWorkflow):
+ task_names = ["custom-runner", "custom-server"]
+
+ workflow = CustomNamedPi0Workflow(
+ workflow_cfg=WorkflowCfg(),
+ task_cfg=PolicyRunnerTaskCfg(arena_env="example_environment"),
+ )
+
+ tasks = workflow.generate_workflow()["workflow"]["groups"][0]["tasks"]
+ assert [task["name"] for task in tasks] == ["custom-runner", "custom-server"]
+ assert "--remote_host {{host:custom-server}}" in tasks[0]["files"][0]["contents"]
+
+
+def test_dreamzero_runner_quotes_explicit_server_task_name():
+ """Quote the submitted server task name in DreamZero's port-forward command."""
+ workflow = DreamZeroPolicyRunnerWorkflow(
+ workflow_cfg=WorkflowCfg(),
+ task_cfg=DreamZeroPolicyRunnerTaskCfg(arena_env="example_environment"),
+ server_workflow_id="server-workflow-id",
+ server_task_name="custom-server; false",
+ )
+
+ task = workflow.generate_workflow()["workflow"]["groups"][0]["tasks"][0]
+ assert task["name"] == "policy_runner"
+ assert "port-forward server-workflow-id 'custom-server; false'" in task["files"][0]["contents"]
+
+
def test_compatibility_cli_builds_typed_config(capsys):
"""Keep the submission CLI as a thin adapter around typed workflow configs."""
return_code = main([
diff --git a/osmo/scripts/__init__.py b/osmo/scripts/__init__.py
new file mode 100644
index 0000000000..16ea4c2183
--- /dev/null
+++ b/osmo/scripts/__init__.py
@@ -0,0 +1,4 @@
+# Copyright (c) 2026, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md).
+# All rights reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
diff --git a/osmo/scripts/build_experiment_output.py b/osmo/scripts/build_experiment_output.py
new file mode 100644
index 0000000000..47f30a1a2a
--- /dev/null
+++ b/osmo/scripts/build_experiment_output.py
@@ -0,0 +1,130 @@
+# Copyright (c) 2026, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md).
+# All rights reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Build one Arena Experiment output from independently executed Experiment Runner tasks.
+
+The input JSON maps each Run name to its Experiment Runner task's output directory. Each task output must contain
+``/...``. Those Run directories are copied into the ``/`` layout, where one
+``index.html`` report is generated.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import shutil
+from collections.abc import Mapping
+from pathlib import Path
+
+from isaaclab_arena.visualization.report import build_report
+
+
+def load_experiment_runner_output_directories_by_run_name(
+ experiment_runner_output_directories_file_path: Path,
+) -> dict[str, Path]:
+ """Load ``run-name -> Experiment Runner task output directory`` from JSON.
+
+ Args:
+ experiment_runner_output_directories_file_path: JSON file containing one Experiment Runner output directory
+ per Run.
+
+ Returns:
+ Run names mapped to Experiment Runner output directories.
+ """
+ with experiment_runner_output_directories_file_path.open(
+ encoding="utf-8"
+ ) as experiment_runner_output_directories_file:
+ runner_output_directory_strings_by_run_name = json.load(experiment_runner_output_directories_file)
+
+ assert (
+ isinstance(runner_output_directory_strings_by_run_name, dict) and runner_output_directory_strings_by_run_name
+ ), "Experiment Runner output directories must be a non-empty JSON mapping"
+ experiment_runner_output_directories_by_run_name: dict[str, Path] = {}
+ for run_name, runner_output_directory_string in runner_output_directory_strings_by_run_name.items():
+ assert isinstance(run_name, str) and run_name, "Run names must be non-empty strings"
+ assert (
+ isinstance(runner_output_directory_string, str) and runner_output_directory_string
+ ), f"Experiment Runner output directory for Run '{run_name}' must be a non-empty string"
+ experiment_runner_output_directories_by_run_name[run_name] = Path(runner_output_directory_string)
+ return experiment_runner_output_directories_by_run_name
+
+
+def collect_run_outputs_into_experiment_output(
+ experiment_runner_output_directories_by_run_name: Mapping[str, Path],
+ experiment_output_directory: Path,
+) -> None:
+ """Collect each Experiment Runner task's Run directory into one Experiment output directory.
+
+ Args:
+ experiment_runner_output_directories_by_run_name: Run names mapped to Experiment Runner task output
+ directories. Each task output directory must contain a child directory with the corresponding Run name.
+ experiment_output_directory: Destination Experiment directory containing one subdirectory per Run.
+ """
+ assert experiment_runner_output_directories_by_run_name, "At least one Experiment Runner output is required"
+ for run_name, experiment_runner_output_directory in experiment_runner_output_directories_by_run_name.items():
+ source_run_output_directory = experiment_runner_output_directory / run_name
+ assert source_run_output_directory.is_dir(), (
+ f"Expected Run output directory for Run '{run_name}' does not exist or is not a directory: "
+ f"'{source_run_output_directory}'"
+ )
+ shutil.copytree(
+ source_run_output_directory,
+ experiment_output_directory / run_name,
+ )
+
+
+def build_experiment_output(
+ experiment_runner_output_directories_by_run_name: Mapping[str, Path],
+ experiment_output_directory: Path,
+) -> Path:
+ """Build one complete Experiment output from Experiment Runner task outputs.
+
+ Args:
+ experiment_runner_output_directories_by_run_name: Run names mapped to Experiment Runner task output
+ directories.
+ experiment_output_directory: Experiment output directory containing one subdirectory per Run and
+ ``index.html``.
+
+ Returns:
+ Path to the generated Experiment report.
+ """
+ collect_run_outputs_into_experiment_output(
+ experiment_runner_output_directories_by_run_name,
+ experiment_output_directory,
+ )
+ return build_report(experiment_output_directory)
+
+
+def _parse_arguments() -> argparse.Namespace:
+ parser = argparse.ArgumentParser(description=__doc__)
+ parser.add_argument(
+ "--experiment-runner-output-directories-file",
+ required=True,
+ type=Path,
+ help="JSON mapping of each Run name to its Experiment Runner task output directory",
+ )
+ parser.add_argument(
+ "--experiment-output-directory",
+ required=True,
+ type=Path,
+ help="Arena Experiment output containing one directory per Run and index.html",
+ )
+ return parser.parse_args()
+
+
+def main() -> None:
+ """Build one Experiment output from the Experiment Runner outputs described on the command line."""
+ parsed_arguments = _parse_arguments()
+ experiment_runner_output_directories_by_run_name = load_experiment_runner_output_directories_by_run_name(
+ parsed_arguments.experiment_runner_output_directories_file
+ )
+ build_experiment_output(
+ experiment_runner_output_directories_by_run_name,
+ parsed_arguments.experiment_output_directory,
+ )
+
+
+if __name__ == "__main__":
+ main()
diff --git a/osmo/tasks/base_task.py b/osmo/tasks/base_task.py
index e5a21656e3..0f479d2995 100644
--- a/osmo/tasks/base_task.py
+++ b/osmo/tasks/base_task.py
@@ -28,14 +28,20 @@ def __init__(
self,
task_cfg: TaskCfg | None = None,
lead: bool | None = None,
+ *,
+ task_name: str,
+ resource: str | None = None,
) -> None:
+ assert isinstance(task_name, str) and task_name, "Task name must be a non-empty string"
+ self.task_name = task_name
self.task_cfg = task_cfg
self.lead = lead
+ self.resource = resource
def create_task_dict(self) -> dict[str, Any]:
"""Assemble the task dict consumed by OSMO."""
- return {
- "name": self.get_task_name(),
+ task = {
+ "name": self.task_name,
"args": ["/tmp/entry.sh"],
"command": ["bash"],
"credentials": self._get_credentials(),
@@ -47,6 +53,9 @@ def create_task_dict(self) -> dict[str, Any]:
"outputs": self._get_outputs(),
"lead": self.lead,
}
+ if self.resource is not None:
+ task["resource"] = self.resource
+ return task
def _get_files_to_create(self) -> list[dict[str, Any]]:
"""Return files OSMO creates in the task container before starting it."""
@@ -72,14 +81,10 @@ def _get_credentials(self) -> dict[str, dict[str, str]]:
}
@staticmethod
- @abstractmethod
- def get_task_name() -> str:
- """Return the task name."""
-
- @classmethod
- def host_token(cls) -> str:
+ def host_token(task_name: str) -> str:
"""Return the OSMO ``{{host:}}`` token that resolves to this task's runtime host/IP."""
- return "{{host:" + cls.get_task_name() + "}}"
+ assert isinstance(task_name, str) and task_name, "Host token task name must be a non-empty string"
+ return "{{host:" + task_name + "}}"
@abstractmethod
def _get_image(self) -> str:
diff --git a/osmo/tasks/collect_experiment_outputs_task.py b/osmo/tasks/collect_experiment_outputs_task.py
new file mode 100644
index 0000000000..a23e8fcfb7
--- /dev/null
+++ b/osmo/tasks/collect_experiment_outputs_task.py
@@ -0,0 +1,96 @@
+# Copyright (c) 2026, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md).
+# All rights reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""OSMO task that collects Experiment Runner outputs into one published Arena Experiment output."""
+
+from __future__ import annotations
+
+import json
+import shlex
+from collections.abc import Mapping
+from pathlib import Path
+from typing import Any
+
+from osmo.tasks.base_task import BaseTask
+from osmo.workflows.utils.yaml_utils import block_literal_str
+from osmo.workflows.workflow_constants import DATASET_SWIFT_URL, OSMO_TASK_OUTPUT_DIR
+
+_LOCAL_BUILD_EXPERIMENT_OUTPUT_SCRIPT_PATH = Path(__file__).parents[1] / "scripts" / "build_experiment_output.py"
+_REMOTE_BUILD_EXPERIMENT_OUTPUT_SCRIPT_PATH = "/tmp/arena_build_experiment_output.py"
+_REMOTE_EXPERIMENT_RUNNER_OUTPUT_DIRECTORIES_FILE_PATH = "/tmp/arena_experiment_runner_output_directories.json"
+
+
+def experiment_runner_output_directory_input_token(experiment_runner_task_name: str) -> str:
+ """Return the OSMO input token that resolves to an Experiment Runner task's output directory."""
+ return "{{input:" + experiment_runner_task_name + "}}"
+
+
+class CollectExperimentOutputsTask(BaseTask):
+ """Collect and publish one Experiment output from the Experiment Runner task outputs.
+
+ For each Run, OSMO exposes its Experiment Runner task output at ``{{input:}}``. The embedded script
+ copies ``{{input:}}//...`` to ``{{output}}//...`` and writes
+ ``{{output}}/index.html``. Only this final task output is published to Swift; the Experiment Runner task outputs
+ remain workflow-local.
+ """
+
+ def __init__(
+ self,
+ image: str,
+ experiment_runner_task_names_by_run_name: Mapping[str, str],
+ lead: bool | None = None,
+ resource: str | None = None,
+ *,
+ task_name: str,
+ ) -> None:
+ assert experiment_runner_task_names_by_run_name, "Experiment output requires at least one Run task"
+ super().__init__(task_name=task_name, lead=lead, resource=resource)
+ self.image = image
+ self.experiment_runner_task_names_by_run_name = dict(experiment_runner_task_names_by_run_name)
+
+ def _get_image(self) -> str:
+ return self.image
+
+ def _get_inputs(self) -> list[dict[str, Any]]:
+ """Make every Experiment Runner task's workflow-local output available to this task."""
+ return [
+ {"task": experiment_runner_task_name}
+ for experiment_runner_task_name in self.experiment_runner_task_names_by_run_name.values()
+ ]
+
+ def _get_outputs(self) -> list[dict[str, Any]]:
+ """Publish the final Experiment directory, including all Runs and ``index.html``."""
+ return [{"url": DATASET_SWIFT_URL}]
+
+ def _get_files_to_create(self) -> list[dict[str, Any]]:
+ """Embed the output-building script and its ``run-name -> Experiment Runner output`` JSON input."""
+ experiment_runner_output_directory_tokens_by_run_name = {
+ run_name: experiment_runner_output_directory_input_token(experiment_runner_task_name)
+ for run_name, experiment_runner_task_name in self.experiment_runner_task_names_by_run_name.items()
+ }
+ return [
+ *super()._get_files_to_create(),
+ {
+ "path": _REMOTE_BUILD_EXPERIMENT_OUTPUT_SCRIPT_PATH,
+ "contents": block_literal_str(_LOCAL_BUILD_EXPERIMENT_OUTPUT_SCRIPT_PATH.read_text(encoding="utf-8")),
+ },
+ {
+ "path": _REMOTE_EXPERIMENT_RUNNER_OUTPUT_DIRECTORIES_FILE_PATH,
+ "contents": block_literal_str(
+ json.dumps(experiment_runner_output_directory_tokens_by_run_name, indent=2)
+ ),
+ },
+ ]
+
+ def _get_run_script(self) -> str:
+ build_experiment_output_command = shlex.join([
+ "/isaac-sim/python.sh",
+ _REMOTE_BUILD_EXPERIMENT_OUTPUT_SCRIPT_PATH,
+ "--experiment-runner-output-directories-file",
+ _REMOTE_EXPERIMENT_RUNNER_OUTPUT_DIRECTORIES_FILE_PATH,
+ "--experiment-output-directory",
+ OSMO_TASK_OUTPUT_DIR,
+ ])
+ return f"set -euo pipefail\n{build_experiment_output_command}\n"
diff --git a/osmo/tasks/dreamzero_policy_runner_task.py b/osmo/tasks/dreamzero_policy_runner_task.py
index 5299d90e16..ca6049ed15 100644
--- a/osmo/tasks/dreamzero_policy_runner_task.py
+++ b/osmo/tasks/dreamzero_policy_runner_task.py
@@ -5,6 +5,7 @@
"""DreamZero policy-runner task for the Isaac Lab Arena OSMO workflow."""
+import shlex
from dataclasses import dataclass
from osmo.tasks.policy_runner_task import PolicyRunnerTask, PolicyRunnerTaskCfg
@@ -52,12 +53,17 @@ def __init__(
task_cfg: DreamZeroPolicyRunnerTaskCfg,
server_workflow_id: str,
lead: bool | None = None,
+ *,
+ task_name: str,
+ server_task_name: str,
) -> None:
- super().__init__(task_cfg=task_cfg, lead=lead)
+ super().__init__(task_name=task_name, task_cfg=task_cfg, lead=lead)
# This ID is spliced unquoted into the generated bash entry script, so reject any
# value that could word-split or inject there (see is_valid_workflow_id).
assert is_valid_workflow_id(server_workflow_id), f"Invalid OSMO workflow ID: {server_workflow_id!r}"
+ assert isinstance(server_task_name, str) and server_task_name, "Server task name must be a non-empty string"
self.server_workflow_id = server_workflow_id
+ self.server_task_name = server_task_name
def _get_credentials(self) -> dict[str, dict[str, str]]:
return {
@@ -92,7 +98,7 @@ def _get_run_script(self) -> str:
"set -x\n"
"(\n"
" while true; do\n"
- f" osmo workflow port-forward {self.server_workflow_id} {self.get_server_task_name()}"
+ f" osmo workflow port-forward {self.server_workflow_id} {shlex.quote(self.server_task_name)}"
f" --port {POLICY_SERVER_PORT} || true\n"
' echo "Tunnel to DreamZero server exited; restarting in'
f' {SERVER_WAIT_INTERVAL_SECONDS}s ..."\n'
@@ -121,10 +127,3 @@ def _get_policy_args(self) -> list[str]:
"--initial_connect_wait_s",
str(INITIAL_CONNECT_WAIT_SECONDS),
]
-
- @staticmethod
- def get_server_task_name() -> str:
- """Name of the server task inside the server workflow this runner tunnels to."""
- from osmo.tasks.dreamzero_server_task import DreamZeroServerTask
-
- return DreamZeroServerTask.get_task_name()
diff --git a/osmo/tasks/dreamzero_server_task.py b/osmo/tasks/dreamzero_server_task.py
index d610b64e4b..2e9f1d8683 100644
--- a/osmo/tasks/dreamzero_server_task.py
+++ b/osmo/tasks/dreamzero_server_task.py
@@ -32,12 +32,10 @@ def __init__(
self,
task_cfg: DreamZeroServerTaskCfg | None = None,
lead: bool | None = None,
+ *,
+ task_name: str,
) -> None:
- super().__init__(task_cfg=task_cfg or DreamZeroServerTaskCfg(), lead=lead)
-
- @staticmethod
- def get_task_name() -> str:
- return "dreamzero_server"
+ super().__init__(task_name=task_name, task_cfg=task_cfg or DreamZeroServerTaskCfg(), lead=lead)
def _get_image(self) -> str:
return self.task_cfg.image
diff --git a/osmo/tasks/experiment_runner_task.py b/osmo/tasks/experiment_runner_task.py
index f7ea5f31d1..1d9c34297f 100644
--- a/osmo/tasks/experiment_runner_task.py
+++ b/osmo/tasks/experiment_runner_task.py
@@ -42,14 +42,14 @@ def __init__(
task_cfg: ExperimentRunnerTaskCfg,
experiment_cfg: ArenaExperimentCfg,
lead: bool | None = None,
+ *,
+ task_name: str,
+ published_output_url: str | None = DATASET_SWIFT_URL,
) -> None:
- super().__init__(task_cfg=task_cfg, lead=lead)
+ super().__init__(task_name=task_name, task_cfg=task_cfg, lead=lead)
assert isinstance(experiment_cfg, ArenaExperimentCfg)
self.experiment_cfg = deepcopy(experiment_cfg)
-
- @staticmethod
- def get_task_name() -> str:
- return "experiment_runner"
+ self.published_output_url = published_output_url
def _get_image(self) -> str:
return self.task_cfg.image
@@ -58,7 +58,8 @@ def _get_inputs(self) -> list[dict[str, Any]]:
return []
def _get_outputs(self) -> list[dict[str, Any]]:
- return [{"url": DATASET_SWIFT_URL}]
+ """Publish this output externally, or leave it workflow-local for a downstream task."""
+ return [] if self.published_output_url is None else [{"url": self.published_output_url}]
def _get_files_to_create(self) -> list[dict[str, Any]]:
"""Embed the effective Experiment at the path consumed by ``experiment_runner.py``."""
@@ -74,7 +75,7 @@ def _get_run_script(self) -> str:
EXPERIMENT_RUNNER_SCRIPT,
"--experiment_config",
REMOTE_EXPERIMENT_PATH,
- "--output_base_dir",
+ "--experiment_output_directory",
OSMO_TASK_OUTPUT_DIR,
"--viz",
"none",
diff --git a/osmo/tasks/gr00t_policy_runner_task.py b/osmo/tasks/gr00t_policy_runner_task.py
index c173efbf2c..24397f194c 100644
--- a/osmo/tasks/gr00t_policy_runner_task.py
+++ b/osmo/tasks/gr00t_policy_runner_task.py
@@ -29,8 +29,10 @@ def __init__(
task_cfg: Gr00tPolicyRunnerTaskCfg,
remote_host: str,
lead: bool | None = None,
+ *,
+ task_name: str,
) -> None:
- super().__init__(task_cfg=task_cfg, lead=lead)
+ super().__init__(task_name=task_name, task_cfg=task_cfg, lead=lead)
# Host of the GR00T server this runner connects to; the workflow resolves it from the server task.
self.remote_host = remote_host
diff --git a/osmo/tasks/gr00t_server_task.py b/osmo/tasks/gr00t_server_task.py
index 50801ec14e..6b8f2197fe 100644
--- a/osmo/tasks/gr00t_server_task.py
+++ b/osmo/tasks/gr00t_server_task.py
@@ -33,12 +33,10 @@ def __init__(
self,
task_cfg: Gr00tServerTaskCfg | None = None,
lead: bool | None = None,
+ *,
+ task_name: str,
) -> None:
- super().__init__(task_cfg=task_cfg or Gr00tServerTaskCfg(), lead=lead)
-
- @staticmethod
- def get_task_name() -> str:
- return "gr00t_server"
+ super().__init__(task_name=task_name, task_cfg=task_cfg or Gr00tServerTaskCfg(), lead=lead)
def _get_image(self) -> str:
return self.task_cfg.image
diff --git a/osmo/tasks/pi0_remote_policy_runner_task.py b/osmo/tasks/pi0_remote_policy_runner_task.py
index 87fadfa181..29fe9fc236 100644
--- a/osmo/tasks/pi0_remote_policy_runner_task.py
+++ b/osmo/tasks/pi0_remote_policy_runner_task.py
@@ -17,8 +17,10 @@ def __init__(
task_cfg: PolicyRunnerTaskCfg,
remote_host: str,
lead: bool | None = None,
+ *,
+ task_name: str,
) -> None:
- super().__init__(task_cfg=task_cfg, lead=lead)
+ super().__init__(task_name=task_name, task_cfg=task_cfg, lead=lead)
# Host of the pi0 server this runner connects to; the workflow resolves it from the server task.
self.remote_host = remote_host
diff --git a/osmo/tasks/pi0_server_task.py b/osmo/tasks/pi0_server_task.py
index a7201b828d..effa96df50 100644
--- a/osmo/tasks/pi0_server_task.py
+++ b/osmo/tasks/pi0_server_task.py
@@ -67,12 +67,10 @@ def __init__(
self,
task_cfg: Pi0ServerTaskCfg | None = None,
lead: bool | None = None,
+ *,
+ task_name: str,
) -> None:
- super().__init__(task_cfg=task_cfg or Pi0ServerTaskCfg(), lead=lead)
-
- @staticmethod
- def get_task_name() -> str:
- return "policy_server"
+ super().__init__(task_name=task_name, task_cfg=task_cfg or Pi0ServerTaskCfg(), lead=lead)
def _get_image(self) -> str:
return self.task_cfg.image
diff --git a/osmo/tasks/policy_runner_task.py b/osmo/tasks/policy_runner_task.py
index 3150ed8d5d..ecdf97a589 100644
--- a/osmo/tasks/policy_runner_task.py
+++ b/osmo/tasks/policy_runner_task.py
@@ -53,12 +53,10 @@ def __init__(
self,
task_cfg: PolicyRunnerTaskCfg,
lead: bool | None = None,
+ *,
+ task_name: str,
) -> None:
- super().__init__(task_cfg=task_cfg, lead=lead)
-
- @staticmethod
- def get_task_name() -> str:
- return "policy_runner"
+ super().__init__(task_name=task_name, task_cfg=task_cfg, lead=lead)
def _get_image(self) -> str:
return self.task_cfg.image
diff --git a/osmo/workflows/arena_experiment_workflow.py b/osmo/workflows/arena_experiment_workflow.py
index 7b54cc566e..b182131ca7 100644
--- a/osmo/workflows/arena_experiment_workflow.py
+++ b/osmo/workflows/arena_experiment_workflow.py
@@ -7,12 +7,14 @@
from __future__ import annotations
-from collections.abc import Sequence
from copy import deepcopy
+from typing import Any
from isaaclab_arena.evaluation.arena_experiment import ArenaExperimentCfg
+from isaaclab_arena.evaluation.arena_run import ArenaRunCfg
from isaaclab_arena_openpi.policy.pi0_remote_config import Pi0RemotePolicyCfg
from osmo.tasks.base_task import BaseTask
+from osmo.tasks.collect_experiment_outputs_task import CollectExperimentOutputsTask
from osmo.tasks.experiment_runner_task import ExperimentRunnerTask, ExperimentRunnerTaskCfg
from osmo.tasks.pi0_server_task import Pi0ServerTask, Pi0ServerTaskCfg
from osmo.workflows.workflow import Workflow, WorkflowCfg
@@ -20,14 +22,14 @@
class Pi0ArenaExperimentWorkflow(Workflow):
- """Run an Arena Experiment with one shared pi0 policy server."""
+ """Run every Arena Experiment Run in its own OSMO group."""
- task_cls_list = [ExperimentRunnerTask, Pi0ServerTask]
+ constructs_groups_directly = True
task_cfg_type = ExperimentRunnerTaskCfg
server_task_cfg_type = Pi0ServerTaskCfg
"""Configuration type used by this policy-server workflow."""
- lead_list = [True, False]
+ experiment_output_resource_name = "experiment-output"
def __init__(
self,
@@ -46,54 +48,124 @@ def __init__(
group_name=group_name,
)
- # Each Experiment Run selects the policy client that Experiment Runner instantiates.
- # This workflow separately launches one Pi0 inference-server task. Verify that every
- # Pi0RemotePolicy Run requests the variant served by that task before connecting it.
- pi0_run_variants = self._get_pi0_run_variants()
- self._assert_pi0_server_compatible(pi0_run_variants)
- self._connect_pi0_runs(list(pi0_run_variants))
-
- def _get_tasks(self) -> list[BaseTask]:
- """Create the lead evaluator and non-lead pi0 server."""
- return [
- ExperimentRunnerTask(
- task_cfg=self.task_cfg,
- experiment_cfg=self.experiment_cfg,
- lead=self.lead_flags[0],
- ),
- Pi0ServerTask(self.pi0_server_task_cfg, lead=self.lead_flags[1]),
- ]
-
- def _get_pi0_run_variants(self) -> dict[str, str]:
+ # Every Pi0 Run gets a dedicated server task. Verify that all of those
+ # clients request the variant configured for the server deployment.
+ pi0_policy_variants_by_run = self._get_pi0_policy_variants_by_run()
+ self._assert_pi0_server_compatible(pi0_policy_variants_by_run)
+
+ def _get_group_dicts(self) -> list[dict[str, Any]]:
+ """Create one independently scheduled group per Run, then collect their outputs into one Experiment output."""
+ run_group_dicts: list[dict[str, Any]] = []
+ experiment_runner_task_names_by_run_name: dict[str, str] = {}
+ for run_index, (run_name, run_config) in enumerate(self.experiment_cfg.runs.items()):
+ run_group_dict, experiment_runner_task_name = self._create_run_group_dict(
+ run_index,
+ run_name,
+ run_config,
+ )
+ run_group_dicts.append(run_group_dict)
+ experiment_runner_task_names_by_run_name[run_name] = experiment_runner_task_name
+
+ experiment_output_group_dict = self._create_experiment_output_group_dict(
+ experiment_runner_task_names_by_run_name
+ )
+ return [*run_group_dicts, experiment_output_group_dict]
+
+ def _create_run_group_dict(
+ self,
+ run_index: int,
+ run_name: str,
+ run_config: ArenaRunCfg,
+ ) -> tuple[dict[str, Any], str]:
+ """Create one OSMO group that executes a single-Run Arena Experiment."""
+ experiment_runner_task_name = f"experiment-runner-{run_index}"
+ single_run_experiment_config = ArenaExperimentCfg(runs={run_name: deepcopy(run_config)})
+
+ pi0_policy_server_tasks: list[BaseTask] = []
+ run_policy_config = single_run_experiment_config.runs[run_name].policy
+ if isinstance(run_policy_config, Pi0RemotePolicyCfg):
+ pi0_server_task_name = f"policy-server-{run_index}"
+ self._configure_pi0_remote_policy_for_server(run_policy_config, pi0_server_task_name)
+ pi0_policy_server_tasks.append(
+ Pi0ServerTask(
+ self.pi0_server_task_cfg,
+ lead=False,
+ task_name=pi0_server_task_name,
+ )
+ )
+
+ # Construct this after connecting the policy because the task snapshots the Experiment.
+ experiment_runner_task = ExperimentRunnerTask(
+ task_cfg=self.task_cfg,
+ experiment_cfg=single_run_experiment_config,
+ lead=True,
+ task_name=experiment_runner_task_name,
+ published_output_url=None,
+ )
+ run_group_tasks = [experiment_runner_task, *pi0_policy_server_tasks]
+
+ run_group_dict = {
+ "name": f"arena-run-{run_index}",
+ "tasks": [run_group_task.create_task_dict() for run_group_task in run_group_tasks],
+ }
+ return run_group_dict, experiment_runner_task_name
+
+ def _create_experiment_output_group_dict(
+ self,
+ experiment_runner_task_names_by_run_name: dict[str, str],
+ ) -> dict[str, Any]:
+ """Collect every Experiment Runner task output into one published Experiment output."""
+ collect_experiment_outputs_task = CollectExperimentOutputsTask(
+ task_name="collect-experiment-outputs",
+ image=self.task_cfg.image,
+ experiment_runner_task_names_by_run_name=experiment_runner_task_names_by_run_name,
+ lead=True,
+ resource=self.experiment_output_resource_name,
+ )
+ return {
+ "name": "arena-experiment-output",
+ "tasks": [collect_experiment_outputs_task.create_task_dict()],
+ }
+
+ def _create_resources_dict(self) -> dict[str, dict[str, Any]]:
+ """Use configured resources for Runs and a CPU-only resource for collecting the Experiment output."""
+ run_task_resource = self._create_resource_dict()
+ experiment_output_task_resource = {**run_task_resource, "gpu": 0}
+ return {
+ "default": run_task_resource,
+ self.experiment_output_resource_name: experiment_output_task_resource,
+ }
+
+ def _get_pi0_policy_variants_by_run(self) -> dict[str, str]:
"""Return effective pi0-remote Run variants needed for compatibility checks."""
- pi0_run_variants = {}
- for run_name, run_cfg in self.experiment_cfg.runs.items():
- if not isinstance(run_cfg.policy, Pi0RemotePolicyCfg):
+ pi0_policy_variants_by_run = {}
+ for run_name, run_config in self.experiment_cfg.runs.items():
+ if not isinstance(run_config.policy, Pi0RemotePolicyCfg):
continue
- pi0_run_variants[run_name] = run_cfg.policy.policy_variant
- return pi0_run_variants
+ pi0_policy_variants_by_run[run_name] = run_config.policy.policy_variant
+ return pi0_policy_variants_by_run
- def _assert_pi0_server_compatible(self, pi0_run_variants: dict[str, str]) -> None:
+ def _assert_pi0_server_compatible(self, pi0_policy_variants_by_run: dict[str, str]) -> None:
"""Require Pi0RemotePolicy Runs whose variants match the deployed server."""
- assert pi0_run_variants, "pi0 server requires at least one Run using Pi0RemotePolicy"
- incompatible_runs = {
- run_name: variant
- for run_name, variant in pi0_run_variants.items()
- if variant != self.pi0_server_task_cfg.policy_variant
+ assert pi0_policy_variants_by_run, "pi0 server requires at least one Run using Pi0RemotePolicy"
+ incompatible_policy_variants_by_run = {
+ run_name: policy_variant
+ for run_name, policy_variant in pi0_policy_variants_by_run.items()
+ if policy_variant != self.pi0_server_task_cfg.policy_variant
}
- assert not incompatible_runs, (
- f"pi0_remote Runs require variants {incompatible_runs}, but the pi0 server is configured for "
- f"'{self.pi0_server_task_cfg.policy_variant}'"
+ assert not incompatible_policy_variants_by_run, (
+ f"pi0_remote Runs require variants {incompatible_policy_variants_by_run}, but the pi0 server is configured"
+ f" for '{self.pi0_server_task_cfg.policy_variant}'"
)
- def _connect_pi0_runs(self, run_names: Sequence[str]) -> None:
- """Connect matching pi0 Runs to the shared server task."""
- host_token = Pi0ServerTask.host_token()
- for run_name in run_names:
- policy_cfg = self.experiment_cfg.runs[run_name].policy
- assert isinstance(policy_cfg, Pi0RemotePolicyCfg)
- policy_cfg.remote_host = host_token
- policy_cfg.remote_port = POLICY_SERVER_PORT
- # The first OSMO inference may compile longer than the policy's normal
- # keepalive timeout. Use the timeout owned by this server deployment.
- policy_cfg.ping_timeout = self.pi0_server_task_cfg.client_ping_timeout_s
+ def _configure_pi0_remote_policy_for_server(
+ self,
+ pi0_remote_policy_config: Pi0RemotePolicyCfg,
+ pi0_server_task_name: str,
+ ) -> None:
+ """Configure a Pi0 remote policy to use its dedicated OSMO server task."""
+ pi0_remote_policy_config.remote_host = Pi0ServerTask.host_token(pi0_server_task_name)
+ pi0_remote_policy_config.remote_port = POLICY_SERVER_PORT
+ # The first OSMO inference may compile longer than the policy's normal
+ # keepalive timeout. Use the timeout owned by this server deployment.
+ pi0_remote_policy_config.ping_timeout = self.pi0_server_task_cfg.client_ping_timeout_s
diff --git a/osmo/workflows/dreamzero_split_workflows.py b/osmo/workflows/dreamzero_split_workflows.py
index d51b1347a4..859997c1cd 100644
--- a/osmo/workflows/dreamzero_split_workflows.py
+++ b/osmo/workflows/dreamzero_split_workflows.py
@@ -25,6 +25,9 @@
from osmo.tasks.dreamzero_server_task import DreamZeroServerTask, DreamZeroServerTaskCfg
from osmo.workflows.workflow import CompositeWorkflow, Workflow, WorkflowCfg, WorkflowSubmissionResult
+DREAMZERO_SERVER_TASK_NAME = "dreamzero_server"
+DREAMZERO_POLICY_RUNNER_TASK_NAME = "policy_runner"
+
@dataclass
class DreamZeroWorkflowCfg(WorkflowCfg):
@@ -85,6 +88,7 @@ class DreamZeroServerWorkflow(Workflow):
"""Workflow containing only the DreamZero inference server, for cross-pool evaluation."""
task_cls_list = [DreamZeroServerTask]
+ task_names = [DREAMZERO_SERVER_TASK_NAME]
task_cfg_type = DreamZeroServerTaskCfg
workflow_cfg_type = DreamZeroServerWorkflowCfg
@@ -97,6 +101,7 @@ class DreamZeroPolicyRunnerWorkflow(Workflow):
"""
task_cls_list = [DreamZeroPolicyRunnerTask]
+ task_names = [DREAMZERO_POLICY_RUNNER_TASK_NAME]
task_cfg_type = DreamZeroPolicyRunnerTaskCfg
def __init__(
@@ -105,12 +110,23 @@ def __init__(
task_cfg: DreamZeroPolicyRunnerTaskCfg,
server_workflow_id: str,
group_name: str = "arena",
+ *,
+ server_task_name: str,
) -> None:
super().__init__(workflow_cfg=workflow_cfg, task_cfg=task_cfg, group_name=group_name)
self.server_workflow_id = server_workflow_id
+ self.server_task_name = server_task_name
def _get_tasks(self) -> list[BaseTask]:
- return [DreamZeroPolicyRunnerTask(self.task_cfg, server_workflow_id=self.server_workflow_id, lead=True)]
+ return [
+ DreamZeroPolicyRunnerTask(
+ task_name=self.task_names[0],
+ task_cfg=self.task_cfg,
+ server_workflow_id=self.server_workflow_id,
+ server_task_name=self.server_task_name,
+ lead=True,
+ )
+ ]
class DreamZeroEvaluationWorkflow(CompositeWorkflow):
@@ -152,7 +168,10 @@ def _submit_steps(self) -> WorkflowSubmissionResult:
print(f"DreamZero server workflow: {server_workflow_id}")
runner_workflow = DreamZeroPolicyRunnerWorkflow(
- self.workflow_cfg, self.task_cfg, server_workflow_id=server_workflow_id
+ self.workflow_cfg,
+ self.task_cfg,
+ server_workflow_id=server_workflow_id,
+ server_task_name=server_workflow.task_names[0],
)
runner_result = runner_workflow.submit_workflow()
if runner_result.returncode != 0 and not self.workflow_cfg.dry_run:
diff --git a/osmo/workflows/server_plus_policy_runner_workflow.py b/osmo/workflows/server_plus_policy_runner_workflow.py
index 2da34677d2..a0865de91a 100644
--- a/osmo/workflows/server_plus_policy_runner_workflow.py
+++ b/osmo/workflows/server_plus_policy_runner_workflow.py
@@ -19,18 +19,24 @@
class ServerPlusPolicyRunnerWorkflow(Workflow):
"""Two-task workflow: a policy-runner (lead) plus the inference server it connects to.
- Subclasses declare ``task_cls_list = [runner_cls, server_cls]``; the runner is wired to the
- server via the server task's OSMO host token so the two stay in sync with the task name.
+ Subclasses declare their runner and server classes and explicit OSMO task names. The runner
+ is wired to the server with that same server task name.
"""
lead_list = [True, False]
def _get_tasks(self) -> list[BaseTask]:
runner_cls, server_cls = self.task_cls_list
+ runner_task_name, server_task_name = self.task_names
runner_lead, server_lead = self.lead_flags
return [
- runner_cls(self.task_cfg, remote_host=server_cls.host_token(), lead=runner_lead),
- server_cls(lead=server_lead),
+ runner_cls(
+ task_name=runner_task_name,
+ task_cfg=self.task_cfg,
+ remote_host=server_cls.host_token(server_task_name),
+ lead=runner_lead,
+ ),
+ server_cls(task_name=server_task_name, lead=server_lead),
]
@@ -38,6 +44,7 @@ class Gr00tPolicyRunnerWorkflow(ServerPlusPolicyRunnerWorkflow):
"""Two-task workflow: a GR00T server plus the lead policy-runner eval task."""
task_cls_list = [Gr00tPolicyRunnerTask, Gr00tServerTask]
+ task_names = ["policy_runner", "gr00t_server"]
task_cfg_type = Gr00tPolicyRunnerTaskCfg
@@ -45,4 +52,5 @@ class Pi0PlusPolicyRunnerWorkflow(ServerPlusPolicyRunnerWorkflow):
"""Workflow containing a policy-runner task and its pi0 policy server."""
task_cls_list = [Pi0RemotePolicyRunnerTask, Pi0ServerTask]
+ task_names = ["policy_runner", "policy_server"]
task_cfg_type = PolicyRunnerTaskCfg
diff --git a/osmo/workflows/workflow.py b/osmo/workflows/workflow.py
index d87c91685f..06bc8db048 100644
--- a/osmo/workflows/workflow.py
+++ b/osmo/workflows/workflow.py
@@ -113,7 +113,13 @@ class Workflow(SubmittableWorkflow):
"""Builds, renders, and submits a single Arena OSMO workflow."""
task_cls_list: list[type[BaseTask]] = []
- """Task classes that make up this workflow, in group order. Subclasses must set this."""
+ """Task classes that make up a static workflow, in group order."""
+
+ task_names: list[str] = []
+ """OSMO names for the task instances, in the same order as ``task_cls_list``."""
+
+ constructs_groups_directly: bool = False
+ """Whether the workflow creates groups without the static task declarations above."""
lead_list: list[bool] | None = None
"""Per-task lead flags; ``None`` lets a single-task workflow default its task to lead."""
@@ -124,11 +130,22 @@ def __init__(
task_cfg: TaskCfg,
group_name: str = "arena",
) -> None:
- assert len(self.task_cls_list) > 0, "Workflow subclasses must set task_cls_list"
super().__init__(workflow_cfg=workflow_cfg, task_cfg=task_cfg)
- # Single-task workflows may leave lead_list unset; that task defaults to lead.
- self.lead_flags = self.lead_list if self.lead_list is not None else [True]
- self._assert_single_lead_task(self.lead_flags)
+ if self.task_cls_list:
+ assert len(self.task_names) == len(self.task_cls_list), "Each task class requires one explicit task name"
+ assert all(
+ isinstance(task_name, str) and task_name for task_name in self.task_names
+ ), "Task names must be non-empty strings"
+ assert len(set(self.task_names)) == len(self.task_names), "Task names must be unique within a group"
+ # Single-task workflows may leave lead_list unset; that task defaults to lead.
+ self.lead_flags = self.lead_list if self.lead_list is not None else [True]
+ self._assert_single_lead_task(self.lead_flags)
+ else:
+ # Workflows with custom group construction create and name their tasks directly.
+ assert (
+ self.constructs_groups_directly
+ ), "Workflow subclasses must declare tasks or set constructs_groups_directly=True"
+ self.lead_flags = []
self.group_name = group_name
def _assert_single_lead_task(self, lead_flags: list[bool]) -> None:
@@ -146,11 +163,8 @@ def create_workflow_dict(self) -> dict[str, Any]:
"version": 2,
"workflow": {
"name": self.workflow_cfg.workflow_name,
- "groups": [{
- "name": self.group_name,
- "tasks": [task.create_task_dict() for task in self._get_tasks()],
- }],
- "resources": {"default": self._create_resource_dict()},
+ "groups": self._get_group_dicts(),
+ "resources": self._create_resources_dict(),
"timeout": {
"exec_timeout": self.workflow_cfg.exec_timeout,
"queue_timeout": self.workflow_cfg.queue_timeout,
@@ -158,6 +172,17 @@ def create_workflow_dict(self) -> dict[str, Any]:
},
}
+ def _get_group_dicts(self) -> list[dict[str, Any]]:
+ """Create the OSMO groups in this workflow."""
+ return [{
+ "name": self.group_name,
+ "tasks": [task.create_task_dict() for task in self._get_tasks()],
+ }]
+
+ def _create_resources_dict(self) -> dict[str, dict[str, Any]]:
+ """Create the named OSMO task resources used by this workflow."""
+ return {"default": self._create_resource_dict()}
+
def render_yaml(self) -> str:
"""Render the workflow dict to YAML text."""
return yaml.dump(
@@ -179,10 +204,11 @@ def submit_workflow(self) -> WorkflowSubmissionResult:
def _get_tasks(self) -> list[BaseTask]:
"""Instantiate task objects for this workflow."""
+ assert self.task_cls_list, "Static workflows must declare task_cls_list"
tasks = []
- for task_cls, lead in zip(self.task_cls_list, self.lead_flags):
+ for task_cls, task_name, lead in zip(self.task_cls_list, self.task_names, self.lead_flags):
assert issubclass(task_cls, BaseTask)
- tasks.append(task_cls(self.task_cfg, lead=lead))
+ tasks.append(task_cls(task_name=task_name, task_cfg=self.task_cfg, lead=lead))
return tasks
def _submit_rendered_workflow(self, rendered: str) -> WorkflowSubmissionResult:
diff --git a/osmo/workflows/workflow_constants.py b/osmo/workflows/workflow_constants.py
index 05651ae20e..02f68c78af 100644
--- a/osmo/workflows/workflow_constants.py
+++ b/osmo/workflows/workflow_constants.py
@@ -3,7 +3,7 @@
#
# SPDX-License-Identifier: Apache-2.0
-# tag denoting the output folder on OSMO
+# OSMO template token resolved at runtime to the current task's writable output directory.
OSMO_TASK_OUTPUT_DIR = "{{output}}"
diff --git a/osmo/workflows/zero_action_policy_runner_workflow.py b/osmo/workflows/zero_action_policy_runner_workflow.py
index 18d6ff02fc..e85e3ee356 100644
--- a/osmo/workflows/zero_action_policy_runner_workflow.py
+++ b/osmo/workflows/zero_action_policy_runner_workflow.py
@@ -16,4 +16,5 @@ class ZeroActionPolicyRunnerWorkflow(Workflow):
"""Workflow containing one zero-action policy-runner task."""
task_cls_list = [ZeroActionPolicyRunnerTask]
+ task_names = ["policy_runner"]
task_cfg_type = PolicyRunnerTaskCfg