code/scripts/launch_oven_fuller_logging_batch.py

from pathlib import Path
import argparse
import json
import os
import shutil
import subprocess
import sys
import time
from typing import Dict, List, Optional, Sequence


PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))


def _configure_thread_env() -> None:
    defaults = {
        "OMP_NUM_THREADS": "1",
        "OPENBLAS_NUM_THREADS": "1",
        "MKL_NUM_THREADS": "1",
        "NUMEXPR_NUM_THREADS": "1",
        "VECLIB_MAXIMUM_THREADS": "1",
        "BLIS_NUM_THREADS": "1",
        "MALLOC_ARENA_MAX": "2",
    }
    for key, value in defaults.items():
        os.environ.setdefault(key, value)


def _configure_coppeliasim_env() -> None:
    coppeliasim_root = os.environ.setdefault("COPPELIASIM_ROOT", "/workspace/coppelia_sim")
    ld_library_path_parts = [
        part for part in os.environ.get("LD_LIBRARY_PATH", "").split(":") if part
    ]
    if coppeliasim_root not in ld_library_path_parts:
        ld_library_path_parts.insert(0, coppeliasim_root)
    os.environ["LD_LIBRARY_PATH"] = ":".join(ld_library_path_parts)


_configure_thread_env()
_configure_coppeliasim_env()

from rr_label_study.oven_study import _aggregate_summary, _episode_dirs


def _select_episode_indices(
    total_episodes: int,
    episode_offset: int,
    max_episodes: Optional[int],
    episode_indices: Optional[Sequence[int]],
) -> List[int]:
    if episode_indices is not None:
        selected: List[int] = []
        seen = set()
        for raw_index in episode_indices:
            episode_index = int(raw_index)
            if not (0 <= episode_index < total_episodes):
                raise ValueError(
                    f"episode index {episode_index} outside available range 0..{total_episodes - 1}"
                )
            if episode_index in seen:
                continue
            selected.append(episode_index)
            seen.add(episode_index)
        return selected

    remaining = max(0, total_episodes - episode_offset)
    if max_episodes is not None:
        remaining = min(remaining, max_episodes)
    if remaining <= 0:
        return []
    return list(range(episode_offset, episode_offset + remaining))


def _is_complete_episode_dir(output_dir: Path, episode_name: str) -> bool:
    required = [
        output_dir.joinpath(f"{episode_name}.dense.csv"),
        output_dir.joinpath(f"{episode_name}.keyframes.csv"),
        output_dir.joinpath(f"{episode_name}.debug.jsonl"),
        output_dir.joinpath(f"{episode_name}.metrics.json"),
        output_dir.joinpath("summary.json"),
        output_dir.joinpath("templates.json"),
        output_dir.joinpath("templates.pkl"),
    ]
    return all(path.exists() for path in required)


def _load_metrics(result_dir: Path, episode_names: Sequence[str]) -> List[Dict[str, object]]:
    metrics: List[Dict[str, object]] = []
    for episode_name in episode_names:
        metrics_path = result_dir.joinpath(episode_name, f"{episode_name}.metrics.json")
        if metrics_path.exists():
            with metrics_path.open("r", encoding="utf-8") as handle:
                metrics.append(json.load(handle))
    return metrics


def _write_json(path: Path, payload: Dict[str, object]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as handle:
        json.dump(payload, handle, indent=2)


def _run_episode(
    dataset_root: Path,
    episode_dir: Path,
    output_dir: Path,
    checkpoint_stride: int,
    num_workers: int,
    base_display: int,
    templates_json: Path,
    stagger_seconds: float,
    thread_count: int,
    log_path: Path,
) -> int:
    env = os.environ.copy()
    thread_count_str = str(thread_count)
    env["OMP_NUM_THREADS"] = thread_count_str
    env["OPENBLAS_NUM_THREADS"] = thread_count_str
    env["MKL_NUM_THREADS"] = thread_count_str
    env["NUMEXPR_NUM_THREADS"] = thread_count_str
    env["VECLIB_MAXIMUM_THREADS"] = thread_count_str
    env["BLIS_NUM_THREADS"] = thread_count_str
    env["MALLOC_ARENA_MAX"] = "2"
    env["PYTHONUNBUFFERED"] = "1"
    with log_path.open("w", encoding="utf-8") as log_handle:
        process = subprocess.Popen(
            [
                sys.executable,
                str(PROJECT_ROOT.joinpath("scripts", "recompute_oven_episode_parallel.py")),
                "--dataset-root",
                str(dataset_root),
                "--episode-dir",
                str(episode_dir),
                "--output-dir",
                str(output_dir),
                "--checkpoint-stride",
                str(checkpoint_stride),
                "--num-workers",
                str(num_workers),
                "--base-display",
                str(base_display),
                "--templates-json",
                str(templates_json),
                "--stagger-seconds",
                str(stagger_seconds),
            ],
            stdout=log_handle,
            stderr=subprocess.STDOUT,
            cwd=str(PROJECT_ROOT),
            env=env,
        )
        return process.wait()


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--dataset-root",
        default="/workspace/data/bimanual_take_tray_out_of_oven_train_128",
    )
    parser.add_argument("--result-dir", required=True)
    parser.add_argument("--templates-json", required=True)
    parser.add_argument("--episode-offset", type=int, default=0)
    parser.add_argument("--max-episodes", type=int, default=100)
    parser.add_argument("--episode-indices")
    parser.add_argument("--checkpoint-stride", type=int, default=16)
    parser.add_argument("--num-workers", type=int, default=24)
    parser.add_argument("--base-display", type=int, default=900)
    parser.add_argument("--stagger-seconds", type=float, default=0.15)
    parser.add_argument("--thread-count", type=int, default=1)
    parser.add_argument("--max-retries", type=int, default=2)
    args = parser.parse_args()

    dataset_root = Path(args.dataset_root)
    result_dir = Path(args.result_dir)
    result_dir.mkdir(parents=True, exist_ok=True)
    templates_json = Path(args.templates_json)
    if not templates_json.exists():
        raise FileNotFoundError(f"missing templates json: {templates_json}")

    all_episode_dirs = _episode_dirs(dataset_root)
    explicit_episode_indices = None
    if args.episode_indices:
        explicit_episode_indices = [
            int(chunk.strip()) for chunk in args.episode_indices.split(",") if chunk.strip()
        ]
    selected_episode_indices = _select_episode_indices(
        total_episodes=len(all_episode_dirs),
        episode_offset=args.episode_offset,
        max_episodes=args.max_episodes,
        episode_indices=explicit_episode_indices,
    )
    selected_episode_names = [f"episode{index}" for index in selected_episode_indices]

    manifest = {
        "dataset_root": str(dataset_root.resolve()),
        "result_dir": str(result_dir.resolve()),
        "templates_json": str(templates_json.resolve()),
        "episode_indices": selected_episode_indices,
        "checkpoint_stride": args.checkpoint_stride,
        "num_workers": args.num_workers,
        "base_display": args.base_display,
        "stagger_seconds": args.stagger_seconds,
        "thread_count": args.thread_count,
        "max_retries": args.max_retries,
        "started_at_epoch": time.time(),
    }
    _write_json(result_dir.joinpath("run_manifest.json"), manifest)

    progress_path = result_dir.joinpath("progress.json")
    logs_dir = result_dir.joinpath("logs")
    logs_dir.mkdir(parents=True, exist_ok=True)

    completed: List[int] = []
    failed: List[Dict[str, object]] = []

    for episode_index in selected_episode_indices:
        episode_name = f"episode{episode_index}"
        episode_dir = all_episode_dirs[episode_index]
        final_output_dir = result_dir.joinpath(episode_name)
        if _is_complete_episode_dir(final_output_dir, episode_name):
            completed.append(episode_index)
            _write_json(
                progress_path,
                {
                    "current_episode": None,
                    "completed_episode_indices": completed,
                    "failed": failed,
                    "total_selected": len(selected_episode_indices),
                    "updated_at_epoch": time.time(),
                },
            )
            continue

        attempt_success = False
        current_failure: Optional[Dict[str, object]] = None
        for attempt_index in range(1, args.max_retries + 2):
            temp_output_dir = result_dir.joinpath(f".{episode_name}.tmp")
            if temp_output_dir.exists():
                shutil.rmtree(temp_output_dir)
            log_path = logs_dir.joinpath(f"{episode_name}.attempt{attempt_index:02d}.log")
            _write_json(
                progress_path,
                {
                    "current_episode": episode_name,
                    "current_attempt": attempt_index,
                    "completed_episode_indices": completed,
                    "failed": failed,
                    "total_selected": len(selected_episode_indices),
                    "updated_at_epoch": time.time(),
                },
            )
            return_code = _run_episode(
                dataset_root=dataset_root,
                episode_dir=episode_dir,
                output_dir=temp_output_dir,
                checkpoint_stride=args.checkpoint_stride,
                num_workers=args.num_workers,
                base_display=args.base_display,
                templates_json=templates_json,
                stagger_seconds=args.stagger_seconds,
                thread_count=args.thread_count,
                log_path=log_path,
            )
            if return_code == 0 and _is_complete_episode_dir(temp_output_dir, episode_name):
                if final_output_dir.exists():
                    shutil.rmtree(final_output_dir)
                temp_output_dir.rename(final_output_dir)
                completed.append(episode_index)
                attempt_success = True
                current_failure = None
                break

            current_failure = {
                "episode_index": episode_index,
                "episode_name": episode_name,
                "attempt": attempt_index,
                "return_code": return_code,
                "log_path": str(log_path),
                "updated_at_epoch": time.time(),
            }
            if temp_output_dir.exists():
                failed_dir = result_dir.joinpath("failed_attempts", f"{episode_name}.attempt{attempt_index:02d}")
                failed_dir.parent.mkdir(parents=True, exist_ok=True)
                if failed_dir.exists():
                    shutil.rmtree(failed_dir)
                temp_output_dir.rename(failed_dir)

        if not attempt_success:
            failed.append(current_failure or {"episode_index": episode_index, "episode_name": episode_name})
            _write_json(
                progress_path,
                {
                    "current_episode": None,
                    "completed_episode_indices": completed,
                    "failed": failed,
                    "total_selected": len(selected_episode_indices),
                    "updated_at_epoch": time.time(),
                },
            )
            raise RuntimeError(f"failed to produce complete result for {episode_name}")

        metrics = _load_metrics(result_dir, selected_episode_names)
        if metrics:
            _write_json(result_dir.joinpath("summary.json"), _aggregate_summary(metrics))
        _write_json(
            progress_path,
            {
                "current_episode": None,
                "completed_episode_indices": completed,
                "failed": failed,
                "total_selected": len(selected_episode_indices),
                "updated_at_epoch": time.time(),
            },
        )

    metrics = _load_metrics(result_dir, selected_episode_names)
    if metrics:
        _write_json(result_dir.joinpath("summary.json"), _aggregate_summary(metrics))
    _write_json(
        progress_path,
        {
            "current_episode": None,
            "completed_episode_indices": completed,
            "failed": failed,
            "total_selected": len(selected_episode_indices),
            "finished_at_epoch": time.time(),
        },
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())