diff --git a/isaaclab_arena/tests/test_osmo_download_experiment_output.py b/isaaclab_arena/tests/test_osmo_download_experiment_output.py new file mode 100644 index 0000000000..8a28a76166 --- /dev/null +++ b/isaaclab_arena/tests/test_osmo_download_experiment_output.py @@ -0,0 +1,309 @@ +# 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.""" + +import subprocess +import sys +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from osmo.scripts.download_experiment_output import _parse_listed_output_paths, 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_cli_help_states_default_output_directory(): + result = subprocess.run( + [sys.executable, "-m", "osmo.scripts.download_experiment_output", "--help"], + capture_output=True, + text=True, + check=False, + timeout=60, + ) + + assert result.returncode == 0 + assert "/eval/" in result.stdout + + +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/" + remote_uri = "s3://my-bucket/experiment-outputs/arena-experiment-123" + list_outputs = iter([ + "experiment-outputs/arena-experiment-123/run one/episode results.jsonl\n\nTotal 1 objects found", + "experiment-outputs/arena-experiment-123/index.html\n\nTotal 1 objects found", + "experiment-outputs/arena-experiment-123/index.html\n\nTotal 1 objects found", + ]) + captured_commands = [] + + def capture_command(command, **kwargs): + captured_commands.append(command) + if command[2] == "list": + assert kwargs == {"capture_output": True, "text": True} + listed_objects, reported_total = next(list_outputs).rsplit("\n\n", maxsplit=1) + Path(command[5]).write_text(listed_objects, encoding="utf-8") + return SimpleNamespace(returncode=0, stdout=reported_total, stderr="") + assert kwargs == {} + if command[3] == remote_uri: + (expected_output_directory / "index.html").write_text("report", encoding="utf-8") + else: + local_directory = Path(command[4]) + (local_directory / "episode results.jsonl").write_text("{}\n", encoding="utf-8") + return SimpleNamespace(returncode=0) + + monkeypatch.setattr("osmo.scripts.download_experiment_output.subprocess.run", capture_command) + + return_code = main([ + "arena-experiment-123", + "--output-base-directory", + str(output_base_directory), + "--remote-base-uri", + remote_base_uri, + ]) + + assert return_code == 0 + assert len(captured_commands) == 5 + for list_command in captured_commands[:3]: + assert list_command[:5] == ["osmo", "data", "list", "--recursive", remote_uri] + assert Path(list_command[5]).name.startswith("objects-") + assert captured_commands[3:] == [ + ["osmo", "data", "download", remote_uri, str(expected_output_directory)], + [ + "osmo", + "data", + "download", + f"{remote_uri}/run one/episode results.jsonl", + str(expected_output_directory / "run one"), + ], + ] + assert sorted( + path.relative_to(expected_output_directory).as_posix() for path in expected_output_directory.rglob("*") + ) == [ + "index.html", + "run one", + "run one/episode results.jsonl", + ] + + +@pytest.mark.parametrize("workflow_id", ["", ".", "..", "workflow/name", "workflow name"]) +def test_cli_rejects_invalid_workflow_id(workflow_id): + result = subprocess.run( + [sys.executable, "-m", "osmo.scripts.download_experiment_output", workflow_id], + capture_output=True, + text=True, + check=False, + timeout=60, + ) + + assert result.returncode == 2 + assert "invalid OSMO workflow ID" in result.stderr + + +@pytest.mark.parametrize("workflow_id", ["", ".", "..", "workflow/name", "workflow name"]) +def test_download_rejects_invalid_workflow_id(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(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_list_failure(monkeypatch, tmp_path): + monkeypatch.setattr( + "osmo.scripts.download_experiment_output.subprocess.run", + lambda command, **kwargs: SimpleNamespace(returncode=23, stdout="", stderr="listing failed\n"), + ) + + return_code = download_experiment_output("arena-experiment-123", tmp_path / "output", DATASETS_SWIFT_URL) + + assert return_code == 23 + + +def test_propagates_osmo_download_failure(monkeypatch, tmp_path): + def fail_bulk_download(command, **kwargs): + if command[2] == "list": + Path(command[5]).write_text("workflows/arena-experiment-123/index.html\n", encoding="utf-8") + return SimpleNamespace( + returncode=0, + stdout="Total 1 object found\n", + stderr="", + ) + return SimpleNamespace(returncode=23) + + monkeypatch.setattr("osmo.scripts.download_experiment_output.subprocess.run", fail_bulk_download) + + return_code = download_experiment_output("arena-experiment-123", tmp_path / "output", DATASETS_SWIFT_URL) + + assert return_code == 23 + + +def test_retries_and_rejects_listing_without_report(monkeypatch, tmp_path, capsys): + captured_commands = [] + + def list_without_report(command, **kwargs): + captured_commands.append(command) + Path(command[5]).write_text("workflows/arena-experiment-123/run/episode.jsonl\n", encoding="utf-8") + return SimpleNamespace( + returncode=0, + stdout="Total 1 object found\n", + stderr="", + ) + + monkeypatch.setattr("osmo.scripts.download_experiment_output.subprocess.run", list_without_report) + + return_code = download_experiment_output("arena-experiment-123", tmp_path / "output", DATASETS_SWIFT_URL) + + assert return_code == 1 + assert len(captured_commands) == 3 + assert "index.html is missing" in capsys.readouterr().err + + +def test_rejects_successful_exact_download_that_does_not_create_file(monkeypatch, tmp_path, capsys): + output_directory = tmp_path / "output" + + def omit_listed_object(command, **kwargs): + if command[2] == "list": + Path(command[5]).write_text( + "\n".join([ + "workflows/arena-experiment-123/index.html", + "workflows/arena-experiment-123/run/episode.jsonl", + ]), + encoding="utf-8", + ) + return SimpleNamespace( + returncode=0, + stdout="Total 2 objects found\n", + stderr="", + ) + if command[3].endswith("arena-experiment-123"): + (output_directory / "index.html").write_text("report", encoding="utf-8") + return SimpleNamespace(returncode=0) + + monkeypatch.setattr("osmo.scripts.download_experiment_output.subprocess.run", omit_listed_object) + + return_code = download_experiment_output("arena-experiment-123", output_directory, DATASETS_SWIFT_URL) + + assert return_code == 1 + assert "reported success without downloading 'run/episode.jsonl'" in capsys.readouterr().err + + +def test_propagates_exact_object_download_failure(monkeypatch, tmp_path): + output_directory = tmp_path / "output" + + def fail_exact_download(command, **kwargs): + if command[2] == "list": + Path(command[5]).write_text( + "\n".join([ + "workflows/arena-experiment-123/index.html", + "workflows/arena-experiment-123/run/episode.jsonl", + ]), + encoding="utf-8", + ) + return SimpleNamespace( + returncode=0, + stdout="Total 2 objects found\n", + stderr="", + ) + if command[3].endswith("arena-experiment-123"): + (output_directory / "index.html").write_text("report", encoding="utf-8") + return SimpleNamespace(returncode=0) + return SimpleNamespace(returncode=24) + + monkeypatch.setattr("osmo.scripts.download_experiment_output.subprocess.run", fail_exact_download) + + return_code = download_experiment_output("arena-experiment-123", output_directory, DATASETS_SWIFT_URL) + + assert return_code == 24 + + +def test_rejects_unexpected_downloaded_file(monkeypatch, tmp_path, capsys): + output_directory = tmp_path / "output" + + def download_with_unexpected_file(command, **kwargs): + if command[2] == "list": + Path(command[5]).write_text("workflows/arena-experiment-123/index.html\n", encoding="utf-8") + return SimpleNamespace( + returncode=0, + stdout="Total 1 object found\n", + stderr="", + ) + (output_directory / "index.html").write_text("report", encoding="utf-8") + (output_directory / "unexpected.jsonl").write_text("{}\n", encoding="utf-8") + return SimpleNamespace(returncode=0) + + monkeypatch.setattr("osmo.scripts.download_experiment_output.subprocess.run", download_with_unexpected_file) + + return_code = download_experiment_output("arena-experiment-123", output_directory, DATASETS_SWIFT_URL) + + assert return_code == 1 + assert "contains unexpected files: unexpected.jsonl" in capsys.readouterr().err + + +@pytest.mark.parametrize( + "listed_path", + [ + "/workflows/arena-experiment-123/index.html", + "workflows/arena-experiment-123/../index.html", + "workflows/arena-experiment-123//index.html", + "workflows/./arena-experiment-123/index.html", + ], +) +def test_rejects_unsafe_listed_object_path(listed_path): + list_output = f"{listed_path}\n\nTotal 1 object found\n" + + with pytest.raises(ValueError, match="Unsafe remote object path"): + _parse_listed_output_paths( + list_output, + "swift://pdx.s8k.io/AUTH_team-isaac/isaaclab_arena/workflows/arena-experiment-123", + ) + + +def test_rejects_list_output_whose_total_does_not_match_object_keys(): + list_output = "workflows/arena-experiment-123/index.html\n\nTotal 2 objects found\n" + + with pytest.raises(ValueError, match="listed 1 object keys but reported 2"): + _parse_listed_output_paths( + list_output, + "swift://pdx.s8k.io/AUTH_team-isaac/isaaclab_arena/workflows/arena-experiment-123", + ) 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/download_experiment_output.py b/osmo/scripts/download_experiment_output.py new file mode 100644 index 0000000000..d85828b660 --- /dev/null +++ b/osmo/scripts/download_experiment_output.py @@ -0,0 +1,231 @@ +# 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 re +import shlex +import subprocess +import sys +import tempfile +from pathlib import Path, PurePosixPath +from urllib.parse import urlsplit + +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 _parse_listed_output_paths(list_output: str, remote_uri: str) -> list[PurePosixPath]: + """Return safe workflow-relative object paths from ``osmo data list`` output.""" + remote_path_parts = tuple(part for part in urlsplit(remote_uri).path.split("/") if part) + if not remote_path_parts: + raise ValueError(f"Remote URI has no object path: {remote_uri!r}") + + object_keys = [] + listed_object_count = None + for line in list_output.splitlines(): + object_key = line.strip() + if not object_key: + continue + count_match = re.fullmatch(r"Total (\d+) objects? found", object_key) + if count_match: + if listed_object_count is not None: + raise ValueError("OSMO object listing contains multiple totals") + listed_object_count = int(count_match.group(1)) + else: + object_keys.append(object_key) + if listed_object_count is None: + raise ValueError("OSMO object listing has no total") + if listed_object_count != len(object_keys): + raise ValueError(f"OSMO listed {len(object_keys)} object keys but reported {listed_object_count}") + + relative_paths: set[PurePosixPath] = set() + workflow_prefix: tuple[str, ...] | None = None + for object_key in object_keys: + object_parts = object_key.split("/") + if any(part in {"", ".", ".."} for part in object_parts): + raise ValueError(f"Unsafe remote object path: {object_key!r}") + matching_prefix_lengths = [ + prefix_length + for prefix_length in range(1, min(len(remote_path_parts), len(object_parts)) + 1) + if tuple(object_parts[:prefix_length]) == remote_path_parts[-prefix_length:] + ] + if not matching_prefix_lengths: + raise ValueError(f"Remote object is outside {remote_uri!r}: {object_key!r}") + prefix_length = max(matching_prefix_lengths) + listed_workflow_prefix = tuple(object_parts[:prefix_length]) + if workflow_prefix is None: + workflow_prefix = listed_workflow_prefix + if listed_workflow_prefix != workflow_prefix: + raise ValueError(f"Inconsistent remote object prefix: {object_key!r}") + relative_parts = object_parts[prefix_length:] + if not relative_parts: + raise ValueError(f"Remote listing contains the workflow prefix as an object: {object_key!r}") + relative_path = PurePosixPath(*relative_parts) + if relative_path in relative_paths: + raise ValueError(f"Duplicate remote object path: {object_key!r}") + relative_paths.add(relative_path) + return sorted(relative_paths, key=lambda path: path.as_posix()) + + +def _local_output_paths(output_directory: Path) -> set[PurePosixPath]: + """Return file paths below an Experiment output directory.""" + return { + PurePosixPath(path.relative_to(output_directory).as_posix()) + for path in output_directory.rglob("*") + if path.is_file() + } + + +def download_experiment_output(workflow_id: str, output_directory: Path, remote_base_uri: str) -> int: + """Download one complete Experiment output and verify every listed object. + + 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: + Zero for a complete output, otherwise the failing OSMO or validation 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}" + expected_path_set: set[PurePosixPath] = set() + # OSMO 6.3.1 can truncate terminal listings, while its file-output mode materializes the complete inventory. + with tempfile.TemporaryDirectory(prefix="arena-osmo-list-") as inventory_directory: + for inventory_attempt in range(3): + inventory_path = Path(inventory_directory) / f"objects-{inventory_attempt}.txt" + list_command = ["osmo", "data", "list", "--recursive", remote_uri, inventory_path.as_posix()] + print(f"$ {shlex.join(list_command)}", flush=True) + list_result = subprocess.run(list_command, capture_output=True, text=True) + if list_result.returncode != 0: + if list_result.stdout: + print(list_result.stdout, end="") + if list_result.stderr: + print(list_result.stderr, end="", file=sys.stderr) + return list_result.returncode + if not inventory_path.is_file(): + print("OSMO reported success without creating the object inventory.", file=sys.stderr) + return 1 + list_output = f"{inventory_path.read_text(encoding='utf-8')}\n{list_result.stdout}" + try: + expected_path_set.update(_parse_listed_output_paths(list_output, remote_uri)) + except ValueError as error: + print(f"Invalid OSMO object listing: {error}", file=sys.stderr) + return 1 + expected_paths = sorted(expected_path_set, key=lambda path: path.as_posix()) + if PurePosixPath("index.html") not in expected_paths: + print("OSMO object listing is incomplete: index.html is missing.", file=sys.stderr) + return 1 + + # OSMO 6.3.1 can omit objects from a directory download while returning success. Preserve the efficient bulk + # operation, repair omitted objects through exact-object downloads, and verify against the materialized listing. + download_command = ["osmo", "data", "download", remote_uri, output_directory.as_posix()] + print(f"$ {shlex.join(download_command)}", flush=True) + download_result = subprocess.run(download_command) + if download_result.returncode != 0: + return download_result.returncode + + downloaded_paths = _local_output_paths(output_directory) + missing_paths = sorted(set(expected_paths) - downloaded_paths, key=lambda path: path.as_posix()) + if missing_paths: + print(f"OSMO directory download omitted {len(missing_paths)} object(s); downloading them individually.") + for relative_path in missing_paths: + local_path = output_directory.joinpath(*relative_path.parts) + local_path.parent.mkdir(parents=True, exist_ok=True) + remote_object_uri = f"{remote_uri}/{relative_path.as_posix()}" + object_download_command = ["osmo", "data", "download", remote_object_uri, local_path.parent.as_posix()] + print(f"$ {shlex.join(object_download_command)}", flush=True) + object_download_result = subprocess.run(object_download_command) + if object_download_result.returncode != 0: + return object_download_result.returncode + if not local_path.is_file(): + print(f"OSMO reported success without downloading '{relative_path.as_posix()}'.", file=sys.stderr) + return 1 + + downloaded_paths = _local_output_paths(output_directory) + if downloaded_paths != set(expected_paths): + missing_paths = sorted(set(expected_paths) - downloaded_paths, key=lambda path: path.as_posix()) + unexpected_paths = sorted(downloaded_paths - set(expected_paths), key=lambda path: path.as_posix()) + if missing_paths: + print( + f"Experiment output is missing: {', '.join(path.as_posix() for path in missing_paths)}", file=sys.stderr + ) + if unexpected_paths: + print( + "Experiment output contains unexpected files:" + f" {', '.join(path.as_posix() for path in unexpected_paths)}", + file=sys.stderr, + ) + return 1 + + print(f"Experiment output downloaded to '{output_directory}'.") + print(f"Open '{output_directory / 'index.html'}' to view the report.") + return 0 + + +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/)", + ) + 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())