diff --git a/isaaclab_arena/tests/test_osmo_build_experiment_output.py b/isaaclab_arena/tests/test_osmo_build_experiment_output.py index eea2d47b15..48d421422c 100644 --- a/isaaclab_arena/tests/test_osmo_build_experiment_output.py +++ b/isaaclab_arena/tests/test_osmo_build_experiment_output.py @@ -46,6 +46,17 @@ def test_loads_experiment_runner_output_directories_as_paths(tmp_path): assert experiment_runner_output_directories_by_run_name == {"first": experiment_runner_output_directory} +def test_loads_empty_experiment_runner_output_directories(tmp_path): + experiment_runner_output_directories_file_path = tmp_path / "experiment-runner-output-directories.json" + experiment_runner_output_directories_file_path.write_text("{}", 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 == {} + + 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) @@ -95,3 +106,14 @@ def test_builds_experiment_output_from_separate_experiment_runner_outputs(tmp_pa assert "first" in report_contents assert "second" in report_contents assert "2 job(s)" in report_contents + + +def test_builds_empty_experiment_output_when_no_run_completed(tmp_path): + experiment_output_directory = tmp_path / "experiment-output" + + report_path = build_experiment_output({}, experiment_output_directory) + + assert report_path == experiment_output_directory / "index.html" + report_contents = report_path.read_text(encoding="utf-8") + assert "0 job(s)" in report_contents + assert "No results recorded yet." in report_contents diff --git a/isaaclab_arena/tests/test_osmo_coordinate_experiment_output.py b/isaaclab_arena/tests/test_osmo_coordinate_experiment_output.py new file mode 100644 index 0000000000..d4488f408a --- /dev/null +++ b/isaaclab_arena/tests/test_osmo_coordinate_experiment_output.py @@ -0,0 +1,332 @@ +# 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 coordinating one OSMO report from successful Experiment Runs.""" + +import json +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from osmo.scripts import coordinate_experiment_output as coordinator + +EXPERIMENT_RUNNER_TASK_NAMES_BY_RUN_NAME = { + "first": "experiment-runner-0", + "second": "experiment-runner-1", + "third": "experiment-runner-2", +} +SOURCE_WORKFLOW_OUTPUT_URL = "https://storage.example.com/v1/team/workflows/source-workflow-7" + + +def _workflow_with_runner_statuses( + workflow_status: str, + statuses_by_task_name: dict[str, str], + output_url: str = SOURCE_WORKFLOW_OUTPUT_URL, +) -> dict: + groups = [ + { + "name": f"arena-run-{run_index}", + "tasks": [{"name": task_name, "status": task_status}], + } + for run_index, (task_name, task_status) in enumerate(statuses_by_task_name.items()) + ] + return {"status": workflow_status, "outputs": output_url, "groups": groups} + + +@pytest.mark.parametrize( + ("statuses_by_task_name", "expected_completed_runs"), + [ + ( + { + "experiment-runner-0": "COMPLETED", + "experiment-runner-1": "COMPLETED", + "experiment-runner-2": "COMPLETED", + }, + EXPERIMENT_RUNNER_TASK_NAMES_BY_RUN_NAME, + ), + ( + { + "experiment-runner-0": "COMPLETED", + "experiment-runner-1": "FAILED", + "experiment-runner-2": "FAILED_IMAGE_PULL", + }, + {"first": "experiment-runner-0"}, + ), + ( + { + "experiment-runner-0": "FAILED_EVICTED", + "experiment-runner-2": "FAILED_CANCELED", + }, + {}, + ), + ], +) +def test_selects_only_completed_experiment_runners(statuses_by_task_name, expected_completed_runs): + workflow = _workflow_with_runner_statuses("FAILED", statuses_by_task_name) + + actual_statuses_by_task_name = coordinator.get_experiment_runner_statuses( + workflow, + EXPERIMENT_RUNNER_TASK_NAMES_BY_RUN_NAME, + ) + + assert actual_statuses_by_task_name == statuses_by_task_name + assert ( + coordinator.find_completed_experiment_runners( + EXPERIMENT_RUNNER_TASK_NAMES_BY_RUN_NAME, + actual_statuses_by_task_name, + ) + == expected_completed_runs + ) + + +def test_waits_for_source_workflow_terminal_status_instead_of_every_task(monkeypatch): + source_workflow_responses = iter([ + _workflow_with_runner_statuses( + "RUNNING", + {"experiment-runner-0": "COMPLETED", "experiment-runner-1": "RUNNING"}, + ), + _workflow_with_runner_statuses( + "FAILED_SUBMISSION", + {"experiment-runner-0": "COMPLETED"}, + ), + ]) + sleep_intervals = [] + monkeypatch.setattr(coordinator, "query_workflow", lambda _workflow_id: next(source_workflow_responses)) + monkeypatch.setattr(coordinator.time, "sleep", sleep_intervals.append) + + terminal_workflow = coordinator.wait_for_source_workflow_to_finish( + "source-workflow-7", + EXPERIMENT_RUNNER_TASK_NAMES_BY_RUN_NAME, + poll_interval_seconds=7, + ) + + assert terminal_workflow["status"] == "FAILED_SUBMISSION" + assert sleep_intervals == [7] + assert coordinator.find_completed_experiment_runners( + EXPERIMENT_RUNNER_TASK_NAMES_BY_RUN_NAME, + coordinator.get_experiment_runner_statuses( + terminal_workflow, + EXPERIMENT_RUNNER_TASK_NAMES_BY_RUN_NAME, + ), + ) == {"first": "experiment-runner-0"} + + +def test_retries_transient_osmo_query_failures(monkeypatch): + command_results = iter([ + SimpleNamespace(returncode=1, stdout="", stderr="temporary service error"), + SimpleNamespace(returncode=0, stdout='{"status": "RUNNING"}', stderr=""), + ]) + sleep_intervals = [] + monkeypatch.setattr(coordinator.subprocess, "run", lambda command, **kwargs: next(command_results)) + monkeypatch.setattr(coordinator.time, "sleep", sleep_intervals.append) + + response = coordinator.run_osmo_json_command( + ["osmo", "workflow", "query", "source-workflow-7"], + maximum_attempts=2, + retry_delay_seconds=3, + ) + + assert response == {"status": "RUNNING"} + assert sleep_intervals == [3] + + +def test_builds_completed_runner_output_url(): + assert ( + coordinator.experiment_runner_output_url( + f"{SOURCE_WORKFLOW_OUTPUT_URL}/", + "experiment-runner-2", + ) + == f"{SOURCE_WORKFLOW_OUTPUT_URL}/experiment-runner-2/" + ) + + with pytest.raises(AssertionError, match=r"HTTP\(S\)"): + coordinator.experiment_runner_output_url("swift://bucket/workflow", "experiment-runner-2") + + +def test_downloads_runner_output_without_preserving_remote_prefix(monkeypatch, tmp_path): + command_results = iter([ + SimpleNamespace(returncode=4, stdout="", stderr="temporary network error"), + SimpleNamespace(returncode=0, stdout="", stderr=""), + ]) + commands = [] + sleep_intervals = [] + + def capture_download(command, **kwargs): + commands.append(command) + return next(command_results) + + monkeypatch.setattr(coordinator.subprocess, "run", capture_download) + monkeypatch.setattr(coordinator.time, "sleep", sleep_intervals.append) + destination_directory = tmp_path / "runner-output" + + coordinator.download_experiment_runner_output( + f"{SOURCE_WORKFLOW_OUTPUT_URL}/experiment-runner-0/", + destination_directory, + maximum_attempts=2, + retry_delay_seconds=4, + ) + + assert destination_directory.is_dir() + assert commands[0] == commands[1] + assert commands[0][:4] == ["wget", "--recursive", "--no-parent", "--no-host-directories"] + assert "--cut-dirs=5" in commands[0] + assert f"--directory-prefix={destination_directory}" in commands[0] + assert commands[0][-1] == f"{SOURCE_WORKFLOW_OUTPUT_URL}/experiment-runner-0/" + assert sleep_intervals == [4] + + +def test_downloads_only_completed_runner_outputs(monkeypatch, tmp_path): + downloads = [] + + def capture_download(output_url, destination_directory): + downloads.append((output_url, destination_directory)) + run_name = "first" if output_url.endswith("experiment-runner-0/") else "third" + (destination_directory / run_name).mkdir(parents=True) + + monkeypatch.setattr(coordinator, "download_experiment_runner_output", capture_download) + + downloaded_directories = coordinator.download_completed_experiment_runner_outputs( + SOURCE_WORKFLOW_OUTPUT_URL, + {"first": "experiment-runner-0", "third": "experiment-runner-2"}, + tmp_path / "downloads", + ) + + assert downloads == [ + ( + f"{SOURCE_WORKFLOW_OUTPUT_URL}/experiment-runner-0/", + tmp_path / "downloads/experiment-runner-0", + ), + ( + f"{SOURCE_WORKFLOW_OUTPUT_URL}/experiment-runner-2/", + tmp_path / "downloads/experiment-runner-2", + ), + ] + assert downloaded_directories == { + "first": tmp_path / "downloads/experiment-runner-0", + "third": tmp_path / "downloads/experiment-runner-2", + } + + +def test_skips_unavailable_or_incomplete_completed_runner_outputs(monkeypatch, tmp_path): + def capture_download(output_url, destination_directory): + if output_url.endswith("experiment-runner-0/"): + raise RuntimeError("stored output is unavailable") + destination_directory.mkdir(parents=True) + if output_url.endswith("experiment-runner-2/"): + (destination_directory / "third").mkdir() + + monkeypatch.setattr(coordinator, "download_experiment_runner_output", capture_download) + + downloaded_directories = coordinator.download_completed_experiment_runner_outputs( + SOURCE_WORKFLOW_OUTPUT_URL, + EXPERIMENT_RUNNER_TASK_NAMES_BY_RUN_NAME, + tmp_path / "downloads", + ) + + assert downloaded_directories == {"third": tmp_path / "downloads/experiment-runner-2"} + + +def test_runs_builder_with_downloaded_output_mapping(monkeypatch, tmp_path): + captured_command = None + captured_output_directories = None + + def capture_builder(command, **kwargs): + nonlocal captured_command, captured_output_directories + captured_command = command + mapping_file_path = Path(command[command.index("--experiment-runner-output-directories-file") + 1]) + captured_output_directories = json.loads(mapping_file_path.read_text(encoding="utf-8")) + return SimpleNamespace(returncode=0, stdout="", stderr="") + + monkeypatch.setattr(coordinator.subprocess, "run", capture_builder) + experiment_output_directory = tmp_path / "combined-output" + builder_script_path = tmp_path / "build-output.py" + coordinator.build_experiment_output( + {"first": tmp_path / "runner-0", "third": tmp_path / "runner-2"}, + experiment_output_directory, + builder_script_path, + tmp_path, + ) + + assert captured_output_directories == { + "first": str(tmp_path / "runner-0"), + "third": str(tmp_path / "runner-2"), + } + assert captured_command[-2:] == ["--experiment-output-directory", str(experiment_output_directory)] + + +@pytest.mark.parametrize( + ("terminal_workflow", "expected_completed_runs"), + [ + ( + _workflow_with_runner_statuses( + "FAILED", + { + "experiment-runner-0": "COMPLETED", + "experiment-runner-1": "FAILED_IMAGE_PULL", + "experiment-runner-2": "COMPLETED", + }, + ), + {"first": "experiment-runner-0", "third": "experiment-runner-2"}, + ), + ( + _workflow_with_runner_statuses("FAILED_SUBMISSION", {}, output_url=""), + {}, + ), + ], +) +def test_coordinates_partial_or_empty_successes( + monkeypatch, + tmp_path, + terminal_workflow, + expected_completed_runs, +): + settings_file_path = tmp_path / "coordinator-settings.json" + settings_file_path.write_text( + json.dumps({ + "experiment_runner_task_names_by_run_name": EXPERIMENT_RUNNER_TASK_NAMES_BY_RUN_NAME, + "poll_interval_seconds": 5, + }), + encoding="utf-8", + ) + downloaded_runs = None + built_output_directories = None + + monkeypatch.setattr( + coordinator, + "wait_for_source_workflow_to_finish", + lambda *_args, **_kwargs: terminal_workflow, + ) + + def capture_downloads(source_output_url, completed_runs, download_root_directory): + nonlocal downloaded_runs + assert source_output_url == SOURCE_WORKFLOW_OUTPUT_URL + downloaded_runs = dict(completed_runs) + return {run_name: download_root_directory / run_name for run_name in completed_runs} + + def capture_build(output_directories, experiment_output_directory, build_script_path, temporary_directory): + nonlocal built_output_directories + built_output_directories = dict(output_directories) + assert experiment_output_directory == tmp_path / "combined-output" + assert build_script_path == tmp_path / "build-output.py" + assert temporary_directory.is_dir() + + monkeypatch.setattr(coordinator, "download_completed_experiment_runner_outputs", capture_downloads) + monkeypatch.setattr(coordinator, "build_experiment_output", capture_build) + + completed_runs = coordinator.coordinate_experiment_output( + "source-workflow-7", + settings_file_path, + tmp_path / "build-output.py", + tmp_path / "combined-output", + ) + + assert completed_runs == expected_completed_runs + if expected_completed_runs: + assert downloaded_runs == expected_completed_runs + assert set(built_output_directories) == set(expected_completed_runs) + else: + assert downloaded_runs is None + assert built_output_directories == {} diff --git a/isaaclab_arena/tests/test_osmo_experiment_workflow.py b/isaaclab_arena/tests/test_osmo_experiment_workflow.py index 487a3a9589..f0a031f2d1 100644 --- a/isaaclab_arena/tests/test_osmo_experiment_workflow.py +++ b/isaaclab_arena/tests/test_osmo_experiment_workflow.py @@ -29,16 +29,28 @@ submit_arena_experiment, ) from osmo.tasks.base_task import TaskCfg -from osmo.tasks.collect_experiment_outputs_task import ( +from osmo.tasks.experiment_output_coordinator_task import ( _REMOTE_BUILD_EXPERIMENT_OUTPUT_SCRIPT_PATH, - _REMOTE_EXPERIMENT_RUNNER_OUTPUT_DIRECTORIES_FILE_PATH, - experiment_runner_output_directory_input_token, + _REMOTE_COORDINATE_EXPERIMENT_OUTPUT_SCRIPT_PATH, + _REMOTE_COORDINATOR_SETTINGS_FILE_PATH, ) 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.arena_experiment_workflow import ( + _DRY_RUN_SOURCE_WORKFLOW_ID, + ExperimentOutputCoordinatorWorkflow, + Pi0ArenaExperimentRunsWorkflow, + Pi0ArenaExperimentWorkflow, + _experiment_output_coordinator_exec_timeout, + _osmo_duration_seconds, +) from osmo.workflows.workflow import WorkflowCfg -from osmo.workflows.workflow_constants import DATASET_SWIFT_URL, OSMO_TASK_OUTPUT_DIR, POLICY_SERVER_PORT +from osmo.workflows.workflow_constants import ( + DATASETS_HTTPS_URL, + DATASETS_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. @@ -100,8 +112,13 @@ def _embedded_experiment(task: dict) -> dict: return yaml.safe_load(experiment_file["contents"]) +def _rendered_workflows(output: str) -> list[dict]: + marker = "[dry-run] Rendered workflow YAML:\n\n" + return [yaml.safe_load(rendered_yaml) for rendered_yaml in output.split(marker)[1:]] + + def _rendered_workflow(output: str) -> dict: - return yaml.safe_load(output[output.index("version: 2\n") :]) + return _rendered_workflows(output)[0] def _workflow_groups(workflow: dict) -> list[dict]: @@ -190,10 +207,10 @@ def test_policy_server_rejects_workflow_fields(): ]) -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.""" +def test_fans_out_single_run_experiments_with_dedicated_pi0_servers(): + """Render one independent OSMO group for every Run.""" source_experiment_cfg = _pi0_experiment_cfg() - workflow = Pi0ArenaExperimentWorkflow( + workflow = Pi0ArenaExperimentRunsWorkflow( workflow_cfg=WorkflowCfg(workflow_name="pi0-experiment"), experiment_cfg=source_experiment_cfg, server_task_cfg=Pi0ServerTaskCfg(), @@ -206,7 +223,6 @@ def test_fans_out_single_run_experiments_with_dedicated_pi0_servers_and_one_expe "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)) @@ -224,7 +240,6 @@ def test_fans_out_single_run_experiments_with_dedicated_pi0_servers_and_one_expe [True, False], [True, False], [True], - [True], ] first_experiment = _embedded_experiment(first_tasks[0]) @@ -261,36 +276,88 @@ def test_fans_out_single_run_experiments_with_dedicated_pi0_servers_and_one_expe 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"] + +def test_coordinator_collects_stored_outputs_without_runner_dependencies(): + """Render an independent collector that reads OSMO outputs only after the source workflow finishes.""" + source_workflow_id = "pi0-experiment-7" + published_output_url = f"{DATASETS_SWIFT_URL}/{source_workflow_id}" + coordinator_workflow_cfg = WorkflowCfg( + workflow_name="arena-experiment-output-coordinator", + cpus=2, + gpus=0, + memory="8Gi", + exec_timeout="345900s", + ) + workflow = ExperimentOutputCoordinatorWorkflow( + workflow_cfg=coordinator_workflow_cfg, + task_cfg=ExperimentRunnerTaskCfg(), + source_workflow_id=source_workflow_id, + experiment_runner_task_names_by_run_name={ + "first": "experiment-runner-0", + "second": "experiment-runner-1", + "local": "experiment-runner-2", + }, + published_output_url=published_output_url, + ) + + rendered_workflow = workflow.generate_workflow() + assert [group["name"] for group in _workflow_groups(rendered_workflow)] == ["arena-experiment-output"] + experiment_output_coordinator_task = _workflow_tasks(rendered_workflow)[0] + assert experiment_output_coordinator_task["name"] == "collect-successful-experiment-outputs" + assert experiment_output_coordinator_task["inputs"] == [] + assert experiment_output_coordinator_task["outputs"] == [{"url": published_output_url}] + + coordinator_settings = json.loads( + _task_file(experiment_output_coordinator_task, _REMOTE_COORDINATOR_SETTINGS_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"), + assert coordinator_settings["experiment_runner_task_names_by_run_name"] == { + "first": "experiment-runner-0", + "second": "experiment-runner-1", + "local": "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 coordinator_settings["poll_interval_seconds"] == 30 + + coordinate_experiment_output_script_file = _task_file( + experiment_output_coordinator_task, + _REMOTE_COORDINATE_EXPERIMENT_OUTPUT_SCRIPT_PATH, + ) + assert "localpath" not in coordinate_experiment_output_script_file + assert "def find_completed_experiment_runners" in coordinate_experiment_output_script_file["contents"] + assert "osmo workflow submit" not in coordinate_experiment_output_script_file["contents"] + assert "wget" in coordinate_experiment_output_script_file["contents"] + build_experiment_output_script_file = _task_file( + experiment_output_coordinator_task, + _REMOTE_BUILD_EXPERIMENT_OUTPUT_SCRIPT_PATH, ) - 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 + assert "def build_experiment_output" in build_experiment_output_script_file["contents"] + coordinator_command = _task_file(experiment_output_coordinator_task, "/tmp/entry.sh")["contents"] + assert coordinator_command.startswith("set -euo pipefail") + assert _REMOTE_COORDINATE_EXPERIMENT_OUTPUT_SCRIPT_PATH in coordinator_command + assert f"--source-workflow-id {source_workflow_id}" in coordinator_command + assert _REMOTE_COORDINATOR_SETTINGS_FILE_PATH in coordinator_command + assert f"--experiment-output-directory '{OSMO_TASK_OUTPUT_DIR}'" in coordinator_command + coordinator_resource = rendered_workflow["workflow"]["resources"]["default"] + assert coordinator_resource == { + "cpu": 2, + "gpu": 0, + "memory": "8Gi", + "platform": "ovx-l40s", + "storage": "200Gi", + } + + +@pytest.mark.parametrize( + ("duration", "expected_seconds"), + [("1d", 86400), ("250ms", 0.25), ("PT1H30M", 5400), ("P2DT1.5S", 172801.5)], +) +def test_parses_osmo_timeout_formats(duration, expected_seconds): + """Keep derived collector timeouts compatible with OSMO's accepted duration forms.""" + assert _osmo_duration_seconds(duration) == expected_seconds + + +def test_coordinator_timeout_covers_source_queue_execution_and_collection(): + """Give the independent collector enough time to wait for the source and build its report.""" + assert _experiment_output_coordinator_exec_timeout(WorkflowCfg(queue_timeout="2d", exec_timeout="1d")) == "345900s" def test_embeds_effective_experiment_yaml(): @@ -374,7 +441,9 @@ def test_submission_composes_defaults_experiment_and_overrides(tmp_path, capsys) assert return_code == 0 rendered = capsys.readouterr().out assert "[dry-run] Rendered workflow YAML" in rendered - workflow = _rendered_workflow(rendered) + rendered_workflows = _rendered_workflows(rendered) + assert len(rendered_workflows) == 2 + workflow, coordinator_workflow = rendered_workflows assert workflow["workflow"]["name"] == "overridden-experiment" tasks = _workflow_tasks(workflow) assert [task["name"] for task in tasks] == ["experiment-runner-0", "policy-server-0"] @@ -395,12 +464,18 @@ def test_submission_composes_defaults_experiment_and_overrides(tmp_path, capsys) assert "--policy.config=overridden-pi0-config" in server_command assert "--policy.dir=gs://openpi-assets-simeval/pi05_droid_jointpos" in server_command + assert coordinator_workflow["workflow"]["name"] == "arena-experiment-output-coordinator" + coordinator_task = _workflow_tasks(coordinator_workflow)[0] + coordinator_command = _task_file(coordinator_task, "/tmp/entry.sh")["contents"] + assert f"--source-workflow-id {_DRY_RUN_SOURCE_WORKFLOW_ID}" in coordinator_command + assert coordinator_task["outputs"] == [{"url": f"{DATASETS_SWIFT_URL}/{_DRY_RUN_SOURCE_WORKFLOW_ID}"}] + def test_embedded_openpi_experiment_composes_through_experiment_runner_loader(tmp_path): """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( + workflow = Pi0ArenaExperimentRunsWorkflow( workflow_cfg=submission_cfg.osmo, experiment_cfg=submission_cfg.experiment_cfg, server_task_cfg=submission_cfg.policy_server, @@ -421,18 +496,16 @@ def test_embedded_openpi_experiment_composes_through_experiment_runner_loader(tm assert run_cfg.policy.ping_timeout == Pi0ServerTaskCfg.client_ping_timeout_s -def test_submission_overrides_osmo_resources(monkeypatch): +def test_submission_overrides_osmo_resources(monkeypatch, capsys): """Apply scheduler overrides after the typed workflow defaults.""" - submitted_command = None - submitted_resources = None + submitted_commands_and_workflows = [] def capture_submission(command, **kwargs): - nonlocal submitted_command, submitted_resources 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"] - return SimpleNamespace(returncode=0, stdout="") + submitted_commands_and_workflows.append((command, submitted_workflow)) + workflow_id = "source-workflow-7" if len(submitted_commands_and_workflows) == 1 else "coordinator-workflow-8" + return SimpleNamespace(returncode=0, stdout=f"Workflow ID - {workflow_id}\n") monkeypatch.setattr("osmo.workflows.workflow.subprocess.run", capture_submission) @@ -443,14 +516,28 @@ def capture_submission(command, **kwargs): ]) assert return_code == 0 - 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["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 + assert len(submitted_commands_and_workflows) == 2 + for submitted_command, _ in submitted_commands_and_workflows: + pool_flag_index = submitted_command.index("--pool") + assert submitted_command[pool_flag_index + 1] == "isaac-dev-l40-03" + + source_workflow = submitted_commands_and_workflows[0][1] + source_resource = source_workflow["workflow"]["resources"]["default"] + assert source_resource["platform"] == "ovx-l40" + assert source_resource["memory"] == "120Gi" + assert source_resource["gpu"] == 1 + + coordinator_workflow = submitted_commands_and_workflows[1][1] + coordinator_resource = coordinator_workflow["workflow"]["resources"]["default"] + assert coordinator_resource["platform"] == "ovx-l40" + assert coordinator_resource["memory"] == "8Gi" + assert coordinator_resource["gpu"] == 0 + coordinator_task = _workflow_tasks(coordinator_workflow)[0] + assert "--source-workflow-id source-workflow-7" in _task_file(coordinator_task, "/tmp/entry.sh")["contents"] + assert ( + f"Arena Experiment report (available after collection): {DATASETS_HTTPS_URL}/source-workflow-7/index.html" + in capsys.readouterr().out + ) def test_cli_requires_experiment_cfg_path_and_policy_server(capsys): diff --git a/osmo/scripts/build_experiment_output.py b/osmo/scripts/build_experiment_output.py index 47f30a1a2a..34e55aae28 100644 --- a/osmo/scripts/build_experiment_output.py +++ b/osmo/scripts/build_experiment_output.py @@ -38,9 +38,9 @@ def load_experiment_runner_output_directories_by_run_name( ) 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" + assert isinstance( + runner_output_directory_strings_by_run_name, dict + ), "Experiment Runner output directories must be a 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" @@ -62,7 +62,6 @@ def collect_run_outputs_into_experiment_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(), ( diff --git a/osmo/scripts/coordinate_experiment_output.py b/osmo/scripts/coordinate_experiment_output.py new file mode 100644 index 0000000000..d29e08ff73 --- /dev/null +++ b/osmo/scripts/coordinate_experiment_output.py @@ -0,0 +1,348 @@ +# 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 report from every successful Run in an OSMO Experiment workflow. + +The NVIDIA OSMO deployment exposes internally stored task outputs as HTTP directories under the workflow's +``outputs`` URL. The coordinator uses that directory tree to collect completed Runs without declaring them as task +dependencies. +""" + +from __future__ import annotations + +import argparse +import json +import shlex +import subprocess +import tempfile +import time +from collections.abc import Mapping +from pathlib import Path +from typing import Any +from urllib.parse import quote, urlsplit + +_ACTIVE_WORKFLOW_STATUSES = frozenset({"PENDING", "RUNNING", "WAITING"}) +_EXPERIMENT_RUNNER_OUTPUT_DIRECTORIES_FILE_NAME = "experiment_runner_output_directories.json" + + +def run_osmo_json_command( + command: list[str], + maximum_attempts: int = 5, + retry_delay_seconds: float = 5, +) -> dict[str, Any]: + """Run an OSMO command and retry transient failures before returning its JSON response.""" + assert maximum_attempts > 0, "OSMO command maximum attempts must be positive" + assert retry_delay_seconds >= 0, "OSMO command retry delay must not be negative" + + failure_message = "" + for attempt_number in range(1, maximum_attempts + 1): + result = subprocess.run(command, capture_output=True, text=True) + if result.returncode == 0: + try: + response = json.loads(result.stdout) + except json.JSONDecodeError: + failure_message = f"OSMO command returned invalid JSON: {shlex.join(command)}" + else: + assert isinstance(response, dict), f"OSMO command must return a JSON object: {shlex.join(command)}" + return response + else: + failure_message = ( + f"OSMO command failed with exit code {result.returncode}: {shlex.join(command)}\n" + f"{result.stderr.strip()}" + ) + + if attempt_number < maximum_attempts: + print( + f"OSMO query attempt {attempt_number} failed; retrying in {retry_delay_seconds:g} seconds.", + flush=True, + ) + time.sleep(retry_delay_seconds) + + raise RuntimeError(failure_message) + + +def query_workflow(workflow_id: str) -> dict[str, Any]: + """Return one workflow's current OSMO status response.""" + return run_osmo_json_command([ + "osmo", + "workflow", + "query", + workflow_id, + "--format-type", + "json", + ]) + + +def get_experiment_runner_statuses( + workflow: Mapping[str, Any], + experiment_runner_task_names_by_run_name: Mapping[str, str], +) -> dict[str, str]: + """Return current OSMO statuses for the expected Experiment Runner tasks.""" + expected_task_names = set(experiment_runner_task_names_by_run_name.values()) + statuses_by_task_name: dict[str, str] = {} + for group in workflow.get("groups", []): + for task in group.get("tasks", []): + task_name = task.get("name") + if task_name not in expected_task_names: + continue + assert task_name not in statuses_by_task_name, f"Workflow contains duplicate task name '{task_name}'" + task_status = task.get("status") + assert isinstance(task_status, str) and task_status, f"Task '{task_name}' has no OSMO status" + statuses_by_task_name[task_name] = task_status + return statuses_by_task_name + + +def find_completed_experiment_runners( + experiment_runner_task_names_by_run_name: Mapping[str, str], + statuses_by_task_name: Mapping[str, str], +) -> dict[str, str]: + """Return Run names mapped to Experiment Runner tasks that OSMO completed successfully.""" + return { + run_name: task_name + for run_name, task_name in experiment_runner_task_names_by_run_name.items() + if statuses_by_task_name.get(task_name) == "COMPLETED" + } + + +def wait_for_source_workflow_to_finish( + source_workflow_id: str, + experiment_runner_task_names_by_run_name: Mapping[str, str], + poll_interval_seconds: float, +) -> dict[str, Any]: + """Wait for the source workflow to finish and return its terminal OSMO response.""" + previous_workflow_status: str | None = None + previous_runner_statuses: dict[str, str] | None = None + while True: + workflow = query_workflow(source_workflow_id) + workflow_status = workflow.get("status") + assert isinstance(workflow_status, str) and workflow_status, "Source workflow has no OSMO status" + runner_statuses = get_experiment_runner_statuses(workflow, experiment_runner_task_names_by_run_name) + + if workflow_status != previous_workflow_status or runner_statuses != previous_runner_statuses: + runner_status_summary = ", ".join( + f"{run_name}={runner_statuses.get(task_name, 'NOT_FOUND')}" + for run_name, task_name in experiment_runner_task_names_by_run_name.items() + ) + print(f"Source workflow status: {workflow_status}; Runs: {runner_status_summary}", flush=True) + previous_workflow_status = workflow_status + previous_runner_statuses = runner_statuses + + if workflow_status not in _ACTIVE_WORKFLOW_STATUSES: + return workflow + time.sleep(poll_interval_seconds) + + +def experiment_runner_output_url(source_workflow_output_url: str, experiment_runner_task_name: str) -> str: + """Return the OSMO-hosted output URL for one completed Experiment Runner task.""" + parsed_source_url = urlsplit(source_workflow_output_url) + assert ( + parsed_source_url.scheme in {"http", "https"} and parsed_source_url.netloc + ), f"OSMO source workflow outputs must be exposed through an HTTP(S) URL; got '{source_workflow_output_url}'" + encoded_task_name = quote(experiment_runner_task_name, safe="-_.~") + return f"{source_workflow_output_url.rstrip('/')}/{encoded_task_name}/" + + +def download_experiment_runner_output( + experiment_runner_output_url: str, + destination_directory: Path, + maximum_attempts: int = 3, + retry_delay_seconds: float = 10, +) -> None: + """Recursively download one completed Experiment Runner output from OSMO storage.""" + assert maximum_attempts > 0, "Output download maximum attempts must be positive" + assert retry_delay_seconds >= 0, "Output download retry delay must not be negative" + parsed_output_url = urlsplit(experiment_runner_output_url) + path_segment_count = len([segment for segment in parsed_output_url.path.split("/") if segment]) + destination_directory.mkdir(parents=True, exist_ok=False) + command = [ + "wget", + "--recursive", + "--no-parent", + "--no-host-directories", + f"--cut-dirs={path_segment_count}", + f"--directory-prefix={destination_directory}", + "--continue", + "--quiet", + experiment_runner_output_url, + ] + + failure_message = "" + for attempt_number in range(1, maximum_attempts + 1): + result = subprocess.run(command, capture_output=True, text=True) + if result.returncode == 0: + return + failure_message = ( + f"Experiment Runner output download failed with exit code {result.returncode}: {shlex.join(command)}\n" + f"{result.stderr.strip()}" + ) + if attempt_number < maximum_attempts: + print( + f"Output download attempt {attempt_number} failed; retrying in {retry_delay_seconds:g} seconds.", + flush=True, + ) + time.sleep(retry_delay_seconds) + raise RuntimeError(failure_message) + + +def download_completed_experiment_runner_outputs( + source_workflow_output_url: str, + completed_experiment_runner_task_names_by_run_name: Mapping[str, str], + download_root_directory: Path, +) -> dict[str, Path]: + """Download each available completed Run output and skip outputs that cannot be collected.""" + downloaded_output_directories_by_run_name: dict[str, Path] = {} + for run_name, task_name in completed_experiment_runner_task_names_by_run_name.items(): + destination_directory = download_root_directory / task_name + output_url = experiment_runner_output_url(source_workflow_output_url, task_name) + print(f"Downloading completed Run '{run_name}' from {output_url}", flush=True) + try: + download_experiment_runner_output(output_url, destination_directory) + except RuntimeError as error: + print(f"Skipping Run '{run_name}' because its output could not be downloaded:\n{error}", flush=True) + continue + + expected_run_output_directory = destination_directory / run_name + if not expected_run_output_directory.is_dir(): + print( + f"Skipping Run '{run_name}' because its downloaded output does not contain " + f"'{expected_run_output_directory}'.", + flush=True, + ) + continue + downloaded_output_directories_by_run_name[run_name] = destination_directory + return downloaded_output_directories_by_run_name + + +def build_experiment_output( + experiment_runner_output_directories_by_run_name: Mapping[str, Path], + experiment_output_directory: Path, + build_experiment_output_script_path: Path, + temporary_directory: Path, +) -> None: + """Run the embedded builder over the downloaded successful Run outputs.""" + output_directories_file_path = temporary_directory / _EXPERIMENT_RUNNER_OUTPUT_DIRECTORIES_FILE_NAME + output_directories_file_path.write_text( + json.dumps( + {run_name: str(path) for run_name, path in experiment_runner_output_directories_by_run_name.items()} + ), + encoding="utf-8", + ) + command = [ + "/isaac-sim/python.sh", + str(build_experiment_output_script_path), + "--experiment-runner-output-directories-file", + str(output_directories_file_path), + "--experiment-output-directory", + str(experiment_output_directory), + ] + result = subprocess.run(command, capture_output=True, text=True) + if result.returncode != 0: + raise RuntimeError( + f"Experiment output builder failed with exit code {result.returncode}: {shlex.join(command)}\n" + f"{result.stderr.strip()}" + ) + + +def load_coordinator_settings(settings_file_path: Path) -> dict[str, Any]: + """Load the embedded Experiment output coordinator settings.""" + settings = json.loads(settings_file_path.read_text(encoding="utf-8")) + assert isinstance(settings, dict), "Experiment output coordinator settings must be a JSON object" + return settings + + +def coordinate_experiment_output( + source_workflow_id: str, + settings_file_path: Path, + build_experiment_output_script_path: Path, + experiment_output_directory: Path, +) -> dict[str, str]: + """Wait for all Runs, download successful outputs, and build their combined report.""" + settings = load_coordinator_settings(settings_file_path) + experiment_runner_task_names_by_run_name = settings["experiment_runner_task_names_by_run_name"] + poll_interval_seconds = settings["poll_interval_seconds"] + + terminal_source_workflow = wait_for_source_workflow_to_finish( + source_workflow_id, + experiment_runner_task_names_by_run_name, + poll_interval_seconds, + ) + terminal_runner_statuses = get_experiment_runner_statuses( + terminal_source_workflow, + experiment_runner_task_names_by_run_name, + ) + completed_experiment_runner_task_names_by_run_name = find_completed_experiment_runners( + experiment_runner_task_names_by_run_name, + terminal_runner_statuses, + ) + print( + f"Found {len(completed_experiment_runner_task_names_by_run_name)} of " + f"{len(experiment_runner_task_names_by_run_name)} completed Runs.", + flush=True, + ) + + with tempfile.TemporaryDirectory(prefix="arena_experiment_outputs_") as temporary_directory_string: + temporary_directory = Path(temporary_directory_string) + downloaded_output_directories_by_run_name: dict[str, Path] = {} + if completed_experiment_runner_task_names_by_run_name: + source_workflow_output_url = terminal_source_workflow.get("outputs") + assert ( + isinstance(source_workflow_output_url, str) and source_workflow_output_url + ), "OSMO source workflow has completed Run outputs but exposes no workflow output URL" + downloaded_output_directories_by_run_name = download_completed_experiment_runner_outputs( + source_workflow_output_url, + completed_experiment_runner_task_names_by_run_name, + temporary_directory / "downloads", + ) + collected_experiment_runner_task_names_by_run_name = { + run_name: completed_experiment_runner_task_names_by_run_name[run_name] + for run_name in downloaded_output_directories_by_run_name + } + print( + f"Building the report from {len(collected_experiment_runner_task_names_by_run_name)} available Run " + "outputs.", + flush=True, + ) + build_experiment_output( + downloaded_output_directories_by_run_name, + experiment_output_directory, + build_experiment_output_script_path, + temporary_directory, + ) + + return collected_experiment_runner_task_names_by_run_name + + +def _parse_arguments() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--source-workflow-id", required=True, help="OSMO workflow containing the Experiment Runs") + parser.add_argument("--settings-file", required=True, type=Path, help="embedded coordinator settings JSON") + parser.add_argument( + "--build-experiment-output-script", + required=True, + type=Path, + help="embedded script that builds the combined Experiment output", + ) + parser.add_argument( + "--experiment-output-directory", + required=True, + type=Path, + help="OSMO output directory for the combined Experiment report", + ) + return parser.parse_args() + + +def main() -> None: + """Coordinate one aggregated Experiment output from all successful Runs.""" + arguments = _parse_arguments() + coordinate_experiment_output( + arguments.source_workflow_id, + arguments.settings_file, + arguments.build_experiment_output_script, + arguments.experiment_output_directory, + ) + + +if __name__ == "__main__": + main() diff --git a/osmo/submit_arena_experiment.py b/osmo/submit_arena_experiment.py index aa55886b8c..747c93a8f2 100644 --- a/osmo/submit_arena_experiment.py +++ b/osmo/submit_arena_experiment.py @@ -52,7 +52,7 @@ class ArenaExperimentSubmissionCfg: def submit_arena_experiment(submission_cfg: ArenaExperimentSubmissionCfg) -> int: - """Build and submit the OSMO workflow described by ``submission_cfg``. + """Build and submit the OSMO workflows described by ``submission_cfg``. Args: submission_cfg: Composed Experiment, task, server, and OSMO configuration. @@ -122,7 +122,7 @@ def _create_argument_parser() -> argparse.ArgumentParser: policy_server_choices = ",".join(POLICY_SERVER_TASK_CFG_BY_NAME) parser = argparse.ArgumentParser( usage=f"%(prog)s [-h] --experiment_cfg PATH --policy_server {{{policy_server_choices}}} [OVERRIDE ...]", - description="Submit a typed Arena Experiment as an OSMO workflow.", + description="Submit a typed Arena Experiment to OSMO.", formatter_class=argparse.RawDescriptionHelpFormatter, epilog=r""" Example: @@ -157,7 +157,7 @@ def _create_argument_parser() -> argparse.ArgumentParser: def main(cli_args: list[str] | None = None) -> int: - """Load the Experiment, apply overrides, and submit its OSMO workflow.""" + """Load the Experiment, apply overrides, and submit its OSMO workflows.""" # Argparse resolves the Experiment path and server selector first; they determine # the concrete configs Hydra receives for its remaining overrides. parser = _create_argument_parser() diff --git a/osmo/tasks/experiment_output_coordinator_task.py b/osmo/tasks/experiment_output_coordinator_task.py new file mode 100644 index 0000000000..9d1d62b72f --- /dev/null +++ b/osmo/tasks/experiment_output_coordinator_task.py @@ -0,0 +1,103 @@ +# 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 coordinates one report from all successful Experiment Runs.""" + +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 OSMO_TASK_OUTPUT_DIR + +_LOCAL_COORDINATE_EXPERIMENT_OUTPUT_SCRIPT_PATH = ( + Path(__file__).parents[1] / "scripts" / "coordinate_experiment_output.py" +) +_LOCAL_BUILD_EXPERIMENT_OUTPUT_SCRIPT_PATH = Path(__file__).parents[1] / "scripts" / "build_experiment_output.py" +_REMOTE_COORDINATE_EXPERIMENT_OUTPUT_SCRIPT_PATH = "/tmp/arena_coordinate_experiment_output.py" +_REMOTE_BUILD_EXPERIMENT_OUTPUT_SCRIPT_PATH = "/tmp/arena_build_experiment_output.py" +_REMOTE_COORDINATOR_SETTINGS_FILE_PATH = "/tmp/arena_experiment_output_coordinator_settings.json" + + +class ExperimentOutputCoordinatorTask(BaseTask): + """Wait for every Run and build one report from the successful outputs.""" + + def __init__( + self, + image: str, + source_workflow_id: str, + experiment_runner_task_names_by_run_name: Mapping[str, str], + published_output_url: str, + poll_interval_seconds: int = 30, + lead: bool | None = None, + resource: str | None = None, + *, + task_name: str, + ) -> None: + assert source_workflow_id, "Experiment output coordinator requires a source workflow ID" + assert experiment_runner_task_names_by_run_name, "Experiment output coordinator requires at least one Run task" + assert published_output_url, "Experiment output coordinator requires a published output URL" + assert poll_interval_seconds > 0, "Experiment output coordinator poll interval must be positive" + super().__init__(task_name=task_name, lead=lead, resource=resource) + self.image = image + self.source_workflow_id = source_workflow_id + self.experiment_runner_task_names_by_run_name = dict(experiment_runner_task_names_by_run_name) + self.published_output_url = published_output_url + self.poll_interval_seconds = poll_interval_seconds + + def _get_image(self) -> str: + return self.image + + def _get_inputs(self) -> list[dict[str, Any]]: + """Start independently so a failed Run cannot cause ``FAILED_UPSTREAM``.""" + return [] + + def _get_outputs(self) -> list[dict[str, Any]]: + """Publish the report under the source Experiment workflow ID.""" + return [{"url": self.published_output_url}] + + def _get_files_to_create(self) -> list[dict[str, Any]]: + """Embed the self-contained coordinator, collector builder, and settings.""" + settings = { + "experiment_runner_task_names_by_run_name": self.experiment_runner_task_names_by_run_name, + "poll_interval_seconds": self.poll_interval_seconds, + } + return [ + *super()._get_files_to_create(), + { + "path": _REMOTE_COORDINATE_EXPERIMENT_OUTPUT_SCRIPT_PATH, + "contents": block_literal_str( + _LOCAL_COORDINATE_EXPERIMENT_OUTPUT_SCRIPT_PATH.read_text(encoding="utf-8") + ), + }, + { + "path": _REMOTE_BUILD_EXPERIMENT_OUTPUT_SCRIPT_PATH, + "contents": block_literal_str(_LOCAL_BUILD_EXPERIMENT_OUTPUT_SCRIPT_PATH.read_text(encoding="utf-8")), + }, + { + "path": _REMOTE_COORDINATOR_SETTINGS_FILE_PATH, + "contents": block_literal_str(json.dumps(settings, indent=2)), + }, + ] + + def _get_run_script(self) -> str: + coordinate_experiment_output_command = shlex.join([ + "/isaac-sim/python.sh", + _REMOTE_COORDINATE_EXPERIMENT_OUTPUT_SCRIPT_PATH, + "--source-workflow-id", + self.source_workflow_id, + "--settings-file", + _REMOTE_COORDINATOR_SETTINGS_FILE_PATH, + "--build-experiment-output-script", + _REMOTE_BUILD_EXPERIMENT_OUTPUT_SCRIPT_PATH, + "--experiment-output-directory", + OSMO_TASK_OUTPUT_DIR, + ]) + return f"set -euo pipefail\n{coordinate_experiment_output_command}\n" diff --git a/osmo/workflows/arena_experiment_workflow.py b/osmo/workflows/arena_experiment_workflow.py index b182131ca7..7b33a08f5b 100644 --- a/osmo/workflows/arena_experiment_workflow.py +++ b/osmo/workflows/arena_experiment_workflow.py @@ -7,29 +7,73 @@ from __future__ import annotations +import math +import re from copy import deepcopy +from dataclasses import replace 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_output_coordinator_task import ExperimentOutputCoordinatorTask 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 -from osmo.workflows.workflow_constants import POLICY_SERVER_PORT +from osmo.workflows.workflow import CompositeWorkflow, Workflow, WorkflowCfg, WorkflowSubmissionResult +from osmo.workflows.workflow_constants import DATASETS_HTTPS_URL, DATASETS_SWIFT_URL, POLICY_SERVER_PORT +_COORDINATOR_WORKFLOW_NAME = "arena-experiment-output-coordinator" +_DRY_RUN_SOURCE_WORKFLOW_ID = "dry-run-experiment-workflow-id" +_COORDINATOR_TIMEOUT_GRACE_SECONDS = 300 +_ISO_8601_DURATION_PATTERN = re.compile( + r"^P(?:(?P\d+)D)?(?:T(?:(?P\d+)H)?(?:(?P\d+)M)?" r"(?:(?P\d+(?:\.\d+)?)S)?)?$" +) +_SECONDS_BY_DURATION_UNIT = { + "d": 24 * 60 * 60, + "h": 60 * 60, + "m": 60, + "s": 1, + "ms": 1 / 1_000, + "us": 1 / 1_000_000, +} -class Pi0ArenaExperimentWorkflow(Workflow): + +def _osmo_duration_seconds(duration: str) -> float: + """Convert an OSMO duration string to seconds.""" + if duration.startswith("P"): + match = _ISO_8601_DURATION_PATTERN.fullmatch(duration) + assert match and any(match.groupdict().values()), f"Invalid OSMO duration '{duration}'" + return ( + int(match.group("days") or 0) * _SECONDS_BY_DURATION_UNIT["d"] + + int(match.group("hours") or 0) * _SECONDS_BY_DURATION_UNIT["h"] + + int(match.group("minutes") or 0) * _SECONDS_BY_DURATION_UNIT["m"] + + float(match.group("seconds") or 0) + ) + + duration_unit = next((unit for unit in ("ms", "us", "d", "h", "m", "s") if duration.endswith(unit)), None) + assert duration_unit is not None, f"Invalid OSMO duration '{duration}'" + duration_value = duration[: -len(duration_unit)] + assert duration_value.isdigit(), f"Invalid OSMO duration '{duration}'" + return int(duration_value) * _SECONDS_BY_DURATION_UNIT[duration_unit] + + +def _experiment_output_coordinator_exec_timeout(source_workflow_cfg: WorkflowCfg) -> str: + """Cover source queue/execution, report collection, and short orchestration delays.""" + source_queue_seconds = _osmo_duration_seconds(source_workflow_cfg.queue_timeout) + source_execution_seconds = _osmo_duration_seconds(source_workflow_cfg.exec_timeout) + coordinator_execution_seconds = ( + source_queue_seconds + 2 * source_execution_seconds + _COORDINATOR_TIMEOUT_GRACE_SECONDS + ) + return f"{math.ceil(coordinator_execution_seconds)}s" + + +class Pi0ArenaExperimentRunsWorkflow(Workflow): """Run every Arena Experiment Run in its own OSMO group.""" constructs_groups_directly = True task_cfg_type = ExperimentRunnerTaskCfg server_task_cfg_type = Pi0ServerTaskCfg - """Configuration type used by this policy-server workflow.""" - - experiment_output_resource_name = "experiment-output" def __init__( self, @@ -53,32 +97,28 @@ def __init__( 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 + @property + def experiment_runner_task_names_by_run_name(self) -> dict[str, str]: + """Map every Run name to its deterministic OSMO Experiment Runner task name.""" + return { + run_name: f"experiment-runner-{run_index}" for run_index, run_name in enumerate(self.experiment_cfg.runs) + } - 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 _get_group_dicts(self) -> list[dict[str, Any]]: + """Create one independently scheduled group per Run.""" + return [ + self._create_run_group_dict(run_index, run_name, run_config) + for run_index, (run_name, run_config) in enumerate(self.experiment_cfg.runs.items()) + ] def _create_run_group_dict( self, run_index: int, run_name: str, run_config: ArenaRunCfg, - ) -> tuple[dict[str, Any], str]: + ) -> dict[str, Any]: """Create one OSMO group that executes a single-Run Arena Experiment.""" - experiment_runner_task_name = f"experiment-runner-{run_index}" + experiment_runner_task_name = self.experiment_runner_task_names_by_run_name[run_name] single_run_experiment_config = ArenaExperimentCfg(runs={run_name: deepcopy(run_config)}) pi0_policy_server_tasks: list[BaseTask] = [] @@ -103,38 +143,10 @@ def _create_run_group_dict( published_output_url=None, ) run_group_tasks = [experiment_runner_task, *pi0_policy_server_tasks] - - run_group_dict = { + return { "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.""" @@ -169,3 +181,111 @@ def _configure_pi0_remote_policy_for_server( # 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 + + +class ExperimentOutputCoordinatorWorkflow(Workflow): + """Wait for one Runs workflow and collect only its successful task outputs.""" + + constructs_groups_directly = True + task_cfg_type = ExperimentRunnerTaskCfg + + def __init__( + self, + workflow_cfg: WorkflowCfg, + task_cfg: ExperimentRunnerTaskCfg, + source_workflow_id: str, + experiment_runner_task_names_by_run_name: dict[str, str], + published_output_url: str, + ) -> None: + self.source_workflow_id = source_workflow_id + self.experiment_runner_task_names_by_run_name = experiment_runner_task_names_by_run_name + self.published_output_url = published_output_url + super().__init__(workflow_cfg=workflow_cfg, task_cfg=task_cfg, group_name="arena-experiment-output") + + def _get_group_dicts(self) -> list[dict[str, Any]]: + """Create the independent collector group with no source task dependencies.""" + coordinator_task = ExperimentOutputCoordinatorTask( + task_name="collect-successful-experiment-outputs", + image=self.task_cfg.image, + source_workflow_id=self.source_workflow_id, + experiment_runner_task_names_by_run_name=self.experiment_runner_task_names_by_run_name, + published_output_url=self.published_output_url, + lead=True, + ) + return [{"name": self.group_name, "tasks": [coordinator_task.create_task_dict()]}] + + +class Pi0ArenaExperimentWorkflow(CompositeWorkflow): + """Submit Pi0 Arena Runs and their successful-output collector with one command.""" + + task_cfg_type = ExperimentRunnerTaskCfg + server_task_cfg_type = Pi0ServerTaskCfg + + def __init__( + self, + workflow_cfg: WorkflowCfg, + experiment_cfg: ArenaExperimentCfg, + server_task_cfg: Pi0ServerTaskCfg, + group_name: str = "arena", + task_cfg: ExperimentRunnerTaskCfg | None = None, + ) -> None: + experiment_runner_task_cfg = task_cfg or ExperimentRunnerTaskCfg() + super().__init__(workflow_cfg=workflow_cfg, task_cfg=experiment_runner_task_cfg) + self.experiment_runs_workflow = Pi0ArenaExperimentRunsWorkflow( + workflow_cfg=workflow_cfg, + experiment_cfg=experiment_cfg, + server_task_cfg=server_task_cfg, + group_name=group_name, + task_cfg=experiment_runner_task_cfg, + ) + + def _submit_steps(self) -> WorkflowSubmissionResult: + """Submit the Runs first, then an independent workflow that collects their successful outputs.""" + experiment_runs_result = self.experiment_runs_workflow.submit_workflow() + if experiment_runs_result.returncode != 0: + return experiment_runs_result + + if self.workflow_cfg.dry_run: + source_workflow_id = _DRY_RUN_SOURCE_WORKFLOW_ID + else: + assert experiment_runs_result.workflow_id, ( + "Could not parse the Arena Experiment workflow ID from the OSMO submission output. The Runs may have" + " been submitted anyway; check `osmo workflow list` before retrying." + ) + source_workflow_id = experiment_runs_result.workflow_id + print(f"Arena Experiment Runs workflow: {source_workflow_id}") + + published_output_url = f"{DATASETS_SWIFT_URL}/{source_workflow_id}" + experiment_report_url = f"{DATASETS_HTTPS_URL}/{source_workflow_id}/index.html" + coordinator_workflow = ExperimentOutputCoordinatorWorkflow( + workflow_cfg=self._create_coordinator_workflow_cfg(), + task_cfg=self.task_cfg, + source_workflow_id=source_workflow_id, + experiment_runner_task_names_by_run_name=( + self.experiment_runs_workflow.experiment_runner_task_names_by_run_name + ), + published_output_url=published_output_url, + ) + coordinator_result = coordinator_workflow.submit_workflow() + if coordinator_result.returncode != 0 and not self.workflow_cfg.dry_run: + print( + f"Output collector submission failed; the Arena Experiment Runs workflow {source_workflow_id}" + " is still running and its successful task outputs remain stored by OSMO." + ) + elif not self.workflow_cfg.dry_run: + print(f"Arena Experiment report (available after collection): {experiment_report_url}") + return WorkflowSubmissionResult( + returncode=coordinator_result.returncode, + workflow_id=experiment_runs_result.workflow_id, + ) + + def _create_coordinator_workflow_cfg(self) -> WorkflowCfg: + """Create a CPU-only collector config with enough time to wait for and collect the Runs.""" + return replace( + self.workflow_cfg, + workflow_name=_COORDINATOR_WORKFLOW_NAME, + cpus=2, + gpus=0, + memory="8Gi", + exec_timeout=_experiment_output_coordinator_exec_timeout(self.workflow_cfg), + )