Add Osmo Arena Experiment output downloader

isaaclab_arena/tests/test_osmo_download_experiment_output.py

@@ -0,0 +1,114 @@
+# 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 downloading complete Arena Experiment outputs from OSMO object storage."""
+
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+
+from osmo.scripts.download_experiment_output import download_experiment_output, main
+from osmo.workflows.workflow_constants import DATASETS_SWIFT_URL
+
+
+def test_downloads_experiment_output_to_default_directory(monkeypatch):
+    captured_download = None
+
+    def capture_download(workflow_id, output_directory, remote_base_uri):
+        nonlocal captured_download
+        captured_download = (workflow_id, output_directory, remote_base_uri)
+        return 0
+
+    monkeypatch.setattr("osmo.scripts.download_experiment_output.download_experiment_output", capture_download)
+
+    return_code = main(["arena-experiment-123"])
+
+    assert return_code == 0
+    assert captured_download == (
+        "arena-experiment-123",
+        Path("/eval/arena-experiment-123"),
+        DATASETS_SWIFT_URL,
+    )
+
+
+def test_help_states_default_output_directory(capsys):
+    with pytest.raises(SystemExit) as help_exit:
+        main(["--help"])
+
+    assert help_exit.value.code == 0
+    assert "/eval/<workflow-id>" in capsys.readouterr().out
+
+
+def test_downloads_from_explicit_remote_and_output_bases_without_shell_splitting(monkeypatch, tmp_path):
+    output_base_directory = tmp_path / "experiment output"
+    expected_output_directory = output_base_directory / "arena-experiment-123"
+    remote_base_uri = "s3://my-bucket/experiment-outputs/"
+    captured_command = None
+
+    def capture_download(command):
+        nonlocal captured_command
+        captured_command = command
+        return SimpleNamespace(returncode=0)
+
+    monkeypatch.setattr("osmo.scripts.download_experiment_output.subprocess.run", capture_download)
+
+    return_code = main([
+        "arena-experiment-123",
+        "--output-base-directory",
+        str(output_base_directory),
+        "--remote-base-uri",
+        remote_base_uri,
+    ])
+
+    assert return_code == 0
+    assert captured_command == [
+        "osmo",
+        "data",
+        "download",
+        "s3://my-bucket/experiment-outputs/arena-experiment-123/",
+        str(expected_output_directory),
+    ]
+    assert expected_output_directory.is_dir()
+
+
+@pytest.mark.parametrize("workflow_id", ["", ".", "..", "workflow/name", "workflow name"])
+def test_rejects_invalid_workflow_id_at_both_entry_points(monkeypatch, tmp_path, workflow_id):
+    def fail_if_called(command):
+        pytest.fail(f"Unexpected download command: {command}")
+
+    monkeypatch.setattr("osmo.scripts.download_experiment_output.subprocess.run", fail_if_called)
+
+    with pytest.raises(SystemExit) as invalid_argument_exit:
+        main([workflow_id])
+    assert invalid_argument_exit.value.code == 2
+
+    with pytest.raises(AssertionError, match="Invalid OSMO workflow ID"):
+        download_experiment_output(workflow_id, tmp_path / "output", DATASETS_SWIFT_URL)
+
+
+def test_rejects_nonempty_output_directory_before_download(monkeypatch, tmp_path):
+    output_directory = tmp_path / "existing-output"
+    output_directory.mkdir()
+    (output_directory / "stale-results.jsonl").write_text("stale", encoding="utf-8")
+
+    def fail_if_called(command):
+        pytest.fail(f"Unexpected download command: {command}")
+
+    monkeypatch.setattr("osmo.scripts.download_experiment_output.subprocess.run", fail_if_called)
+
+    with pytest.raises(AssertionError, match="Experiment output directory must be empty"):
+        download_experiment_output("arena-experiment-123", output_directory, DATASETS_SWIFT_URL)
+
+
+def test_propagates_osmo_download_failure(monkeypatch, tmp_path):
+    monkeypatch.setattr(
+        "osmo.scripts.download_experiment_output.subprocess.run",
+        lambda command: SimpleNamespace(returncode=23),
+    )
+
+    return_code = download_experiment_output("arena-experiment-123", tmp_path / "output", DATASETS_SWIFT_URL)
+
+    assert return_code == 23

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

osmo/scripts/download_experiment_output.py

@@ -0,0 +1,100 @@
+# 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
+
+"""Download one Arena Experiment output from OSMO object storage."""
+
+from __future__ import annotations
+
+import argparse
+import shlex
+import subprocess
+from pathlib import Path
+
+from osmo.workflows.utils.workflow_id import is_valid_workflow_id
+from osmo.workflows.workflow_constants import DATASETS_SWIFT_URL
+
+
+def _is_safe_workflow_id(value: str) -> bool:
+    """Return whether a workflow ID is safe to use as a remote and local path component."""
+    return is_valid_workflow_id(value) and value not in {".", ".."}
+
+
+def _workflow_id_argument(value: str) -> str:
+    """Parse a workflow ID that is safe to use as a remote and local path component."""
+    if not _is_safe_workflow_id(value):
+        raise argparse.ArgumentTypeError(f"invalid OSMO workflow ID: {value!r}")
+    return value
+
+
+def download_experiment_output(workflow_id: str, output_directory: Path, remote_base_uri: str) -> int:
+    """Download one complete Experiment output and return the OSMO process status.
+
+    Args:
+        workflow_id: OSMO workflow ID naming the published Experiment output.
+        output_directory: Exact local destination for the Experiment output.
+        remote_base_uri: Object-storage base URI containing workflow outputs.
+
+    Returns:
+        The ``osmo data download`` process status.
+    """
+    assert _is_safe_workflow_id(workflow_id), f"Invalid OSMO workflow ID: {workflow_id!r}"
+    output_directory = output_directory.expanduser()
+    output_directory.mkdir(parents=True, exist_ok=True)
+    assert not any(output_directory.iterdir()), f"Experiment output directory must be empty: '{output_directory}'"
+    remote_uri = f"{remote_base_uri.rstrip('/')}/{workflow_id}/"
+    command = ["osmo", "data", "download", remote_uri, output_directory.as_posix()]
+    print(f"$ {shlex.join(command)}", flush=True)
+    result = subprocess.run(command)
+    if result.returncode == 0:
+        print(f"Experiment output downloaded to '{output_directory}'.")
+        print(f"Open '{output_directory / 'index.html'}' to view the report.")
+    return result.returncode
+
+
+def _create_argument_parser() -> argparse.ArgumentParser:
+    """Create the Experiment-output download command-line parser."""
+    parser = argparse.ArgumentParser(
+        description=(
+            "Download one complete Arena Experiment output, including its report, per-Run results, JSONL outcomes, "
+            "and videos."
+        ),
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+
+  python3 -m osmo.scripts.download_experiment_output arena-experiment-123
+  python3 -m osmo.scripts.download_experiment_output arena-experiment-123 --output-base-directory ./my-output
+  python3 -m osmo.scripts.download_experiment_output arena-experiment-123 --remote-base-uri s3://my-bucket/outputs
+""",
+    )
+    parser.add_argument(
+        "workflow_id",
+        type=_workflow_id_argument,
+        help="OSMO workflow ID printed by the Arena Experiment submission command",
+    )
+    parser.add_argument(
+        "--remote-base-uri",
+        default=DATASETS_SWIFT_URL,
+        help="object-storage base URI containing workflow outputs (default: %(default)s)",
+    )
+    parser.add_argument(
+        "--output-base-directory",
+        type=Path,
+        default=Path("/eval"),
+        help="local base directory (default destination: %(default)s/<workflow-id>)",
+    )
+    parser.allow_abbrev = False
+    return parser
+
+
+def main(cli_args: list[str] | None = None) -> int:
+    """Download the Experiment output described on the command line."""
+    parsed_arguments = _create_argument_parser().parse_args(cli_args)
+    output_directory = parsed_arguments.output_base_directory / parsed_arguments.workflow_id
+    return download_experiment_output(parsed_arguments.workflow_id, output_directory, parsed_arguments.remote_base_uri)
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())