Runner

def run(
    self,
    generator: MojoGenerator,
    runtime: MojoRuntime | None,
    seed: int | None,
    overrides: NamedValueDict[NDArray],
    gen_args: list[Any],
    gen_kwargs: dict[str, Any],
    run_args: list[Any],
    run_kwargs: dict[str, Any],
) -> tuple[MojoModel | None, TrialStatus, MjState | None]:
    """
    Executes the complete simulation pipeline for this trial.

    This method coordinates three main phases:
    1.  **Generation**: Calls the user-provided generator to build a `MojoModel` model.
    2.  **Persistence**: Creates the workspace and writes the model/config to disk.
    3.  **Execution**: Triggers the physics runtime if one is provided.

    Args:
        generator: Function that returns a `MojoModel` instance.
        runtime: Optional function to run the simulation (MuJoCo).
        seed: Seed to use to define the trial.
        overrides: Key-value pairs that override random distributions.
        gen_args: Positional arguments for the generator.
        gen_kwargs: Keyword arguments for the generator.
        run_args: Positional arguments for the runtime.
        run_kwargs: Keyword arguments for the runtime.

    Returns:
        The `MojoModel` object for the trial or None if there was a failure prior to generating the MojoModel and the status of the trial.

    """
    status = TrialStatus(trial_num=self.trial_num)
    status._path = self.trial_dir / TRIAL_STATUS_FNAME

    # set up a per-trial log file capturing everything logged during this trial
    self.trial_dir.mkdir(parents=True, exist_ok=True)
    trial_log_handler = get_trial_log_handler(self.trial_dir / "mojo.log")
    root_logger = logging.getLogger()
    root_logger.addHandler(trial_log_handler)

    with status.record_step(step_name="pending"):
        pass

    from mujoco_mojo.runtime.requirements_manager import RequirementSatisfied
    from mujoco_mojo.runtime.runtime_manager import SimulationStopped

    result = None
    state = None
    runtime_manager: RuntimeManager | None = None
    try:
        # 1. Generate
        with status.record_step(step_name="generating"):
            logger.info(f"Generating trial_num={self.trial_num}")
            mojo_model = (
                MojoModel()
                .with_overrides(overrides=overrides)
                .with_seed(seed=seed)
                .with_trial_num(self.trial_num)
            )
            mojo_model = generator(mojo_model, overrides, *gen_args, **gen_kwargs)
            mojo_model._trial_dir = self.trial_dir

            # 2. Setup Workspace & Save Metadata
            logger.info(f"Saving trial_num={self.trial_num} to {self.trial_dir}")

            # bundle assets, this remaps DepPath attributes to point to the shared asset dir
            rel_to_xml = Path(
                os.path.relpath(self.shared_asset_dir, self.trial_dir)
            )
            mojo_model.mjcf.bundle_assets(
                target_dir=self.shared_asset_dir, rel_to_xml=rel_to_xml
            )

            # save XML (with modified DepPath)
            if runtime is None:
                mojo_model.mjcf.write_xml(self.xml_path)
            if self.model_config_name:
                mojo_model.dump_to_path(self.model_config_path)
            self.named_value_path.write_text(mojo_model.named.model_dump_json())

        with status.record_step(step_name="solving"):
            # 3. Execute (if runtime provided)
            if runtime is not None:
                logger.info(f"Executing trial_num={self.trial_num} runtime")
                import mujoco_mojo.runtime as rt

                runtime_manager = rt.RuntimeManager(
                    signal_manager=rt.SignalManager(
                        export_path=self.trial_dir
                        / rt.SignalManager.default_output_name(),
                        unit_system=mojo_model.us,
                    )
                )
                runtime_manager._mojo_model = mojo_model
                state = mojo_model.mjcf.prep_for_sim(
                    self.xml_path, unit_system=mojo_model.us
                )
                result = runtime(
                    mojo_model,
                    runtime_manager,
                    state,
                    *run_args,
                    **run_kwargs,
                )
                status.requirements = runtime_manager.requirement_results
            else:
                logger.info(
                    f"No runtime definition was provided for trial_num={self.trial_num} so MuJoCo will not be run."
                )
                result = mojo_model
                state = None

        # serialize again in case new named values were added during the run
        if self.model_config_name:
            mojo_model.dump_to_path(self.model_config_path)
        self.named_value_path.write_text(mojo_model.named.model_dump_json())

        # clear unpicklable collision managers before returing to avoid multiprocessing serialization erros
        mojo_model.clear_unpickleable_data()

        status.step = "done"
        if status.requirements and not all(r.passed for r in status.requirements):
            status.completion = Completion.FAILURE
        else:
            status.completion = Completion.SUCCESS

    except (BdbQuit, KeyboardInterrupt):
        logger.warning("Quit command detected. Exiting execution...")
        raise
    except RequirementSatisfied:
        # a live requirement ended the trial early as a success: complete
        # normally, with the outcome decided by the requirement results
        status.step = "done"
        if runtime_manager is not None:
            status.requirements = runtime_manager.requirement_results
        if status.requirements and not all(r.passed for r in status.requirements):
            status.completion = Completion.FAILURE
        else:
            status.completion = Completion.SUCCESS
    except SimulationStopped:
        status.step = "done"
        status.completion = Completion.TERMINATED
        if runtime_manager is not None:
            status.requirements = runtime_manager.requirement_results
    except Exception as e:
        status.step = "done"
        status.completion = Completion.ERROR
        logger.exception(
            f"Trial {self.trial_num} failed with the following error: {e}"
        )
    finally:
        status.dump_to_path(status._path)
        root_logger.removeHandler(trial_log_handler)
        trial_log_handler.close()

    return result, status, state

def run(
    self,
    global_overrides: NamedValueDict[NDArray] = NamedValueDict[NDArray](),
    clean_workdir: bool = False,
    cleanup_delay: int = 10,
    execution_mode: ExecutionMode = ExecutionMode.LOCAL,
    trial_nums: list[int] | None = None,
) -> bool:
    """Vectors a job to be either computed locally or to be orchestrated by SLURM."""
    funcs = {
        "generator": self.generator,
        "runtime": self.runtime,
        "objective": self.objective,
    }
    non_none = {k: v for k, v in funcs.items() if v is not None}
    seen: dict[int, str] = {}
    for role, func in non_none.items():
        fid = id(func)
        if fid in seen:
            msg = f"'{getattr(func, '__qualname__', func)}' is assigned to both '{seen[fid]}' and '{role}'. Each function must be unique."
            logger.error(msg)
            raise ValueError(msg)
        seen[fid] = role

    if clean_workdir:
        if self.config.resume:
            msg = "clean_workdir and resume are mutually exclusive with one another. Use one or the other."
            logger.error(msg)
            raise ValueError(msg)

        self.force_remove_dir(countdown_from=cleanup_delay, path=self.workdir)

    self.workdir.mkdir(parents=True, exist_ok=True)
    if not (self.workdir / ".gitignore").exists():
        (self.workdir / ".gitignore").write_text("*", encoding="utf-8")

    # SLURM workers re-enter run() with execution_mode=LOCAL; these only need
    # to happen once, during orchestration, to avoid every worker racing to
    # write the same shared files
    if self.slurm_trial_id is None:
        write_dojo_script(self.workdir)
        self.capture_environment()

    match execution_mode:
        case ExecutionMode.LOCAL:
            return self.run_local(
                global_overrides=global_overrides,
                trial_nums=trial_nums,
            )
        case ExecutionMode.SLURM:
            return self.orchestrate_slurm(
                global_overrides=global_overrides,
                trial_nums=trial_nums,
            )
        case _:
            msg = f"No run command has been configured for execution mode {execution_mode}"
            logger.error(msg)
            raise NotImplementedError(msg)

def orchestrate_slurm(
    self,
    global_overrides: NamedValueDict[NDArray],
    trial_nums: list[int] | None = None,
) -> bool:
    """Generates an sbatch script and submits the job array to SLURM for a given config."""
    (self.workdir / "logs").mkdir(exist_ok=True)

    from rich.console import Console

    Console().print(
        "[dim]Tip: pass --slurm-config/-sc with a flat JSON file to auto-inject extra "
        "#SBATCH lines (prefix a key with 'sbatch.', e.g. 'sbatch.account') and/or "
        "environment variables (any other key) into the generated submission script.[/dim]"
    )

    if isinstance(self.config, MonteCarloConfig):
        try:
            had_fails = self.orchestrate_slurm_monte_carlo(
                global_overrides=global_overrides,
                trial_nums=trial_nums,
            )
        except (BdbQuit, KeyboardInterrupt):
            print("\n")
            logger.error("Aborted SLURM orchestration!")
            had_fails = True
    else:
        msg = f"A SLURM configuration for {self.config.__class__.__name__} has not been implemented."
        logger.error(msg)
        raise NotImplementedError(msg)
    return had_fails

def execute_single_trial(
    self, trial_num: int, seed: int | None, overrides_payload: dict
) -> tuple[MojoModel | None, TrialStatus, MjState | None]:
    """
    Helper to package a Trial and run it.

    The model and state are live in-process objects; callers that need to cross a ProcessPoolExecutor boundary (i.e. parallel MC) must use `_execute_trial_subprocess` instead, which drops the non-picklable objects before the result is serialized back.

    Returns `(mojo_model, trial_status, state)`.
    """
    overrides = NamedValueDict[NDArray].model_validate(overrides_payload)

    trial = Trial(
        trial_num=trial_num,
        base_dir=self.workdir,
        xml_name=self.xml_name,
        model_config_name=self.model_config_name,
        padding_style=self.config.padding_style,
    )

    mojo_model, trial_status, state = trial.run(
        generator=self.generator,
        runtime=self.runtime,
        seed=seed,
        overrides=overrides,
        gen_args=self.gen_args,
        gen_kwargs=self.gen_kwargs,
        run_args=self.run_args,
        run_kwargs=self.run_kwargs,
    )

    # write distribution tables from this process so they never need to
    # cross the pickle boundary; FileLock serializes concurrent workers,
    # and the inner exists() check ensures only the first one writes
    if mojo_model is not None and mojo_model.dists:
        stochas_dir = self.workdir / STOCHAS_DIR_NAME
        stochas_dir.mkdir(exist_ok=True)
        dists_json_path = stochas_dir / STOCHAS_DISTS_FNAME
        with FileLock(stochas_dir / ".dists.lock"):
            if not dists_json_path.exists():
                mojo_model.dists.to_tables(stochas_dir)
                tmp = stochas_dir / "dists.tmp.json"
                tmp.write_text(mojo_model.dists.model_dump_json(), encoding="utf-8")
                tmp.replace(dists_json_path)

    return mojo_model, trial_status, state

def run_monte_carlo(
    self,
    global_overrides: NamedValueDict[NDArray] = NamedValueDict[NDArray](),
    trial_nums: list[int] | None = None,
) -> bool:
    """Orchestrates a Monte Carlo job."""
    if self.slurm_trial_id is not None:
        tn = self.slurm_trial_id
        logger.info(f"SLURM Worker detected. Executing Trial {tn}")
        _mojo_model, trial_status, _ = self.execute_single_trial(
            trial_num=tn,
            seed=self.seed,
            overrides_payload=global_overrides.model_dump(),
        )

        return trial_status.completion == Completion.ERROR

    job_trial_nums = trial_nums if trial_nums else self.config.trial_nums

    # initialize the status tracker
    status_tracker = JobStatus(
        workdir=self.workdir.resolve(),
        job_type=JobType.MONTE_CARLO,
        execution_mode=ExecutionMode.LOCAL,
        n_proc=self.config.n_proc,
        seed=self.seed,
        padding_style=self.config.padding_style,
        generator=MojoRunner.inspect_protocol(self.generator),
        runtime=MojoRunner.inspect_protocol(self.runtime),
        objective=MojoRunner.inspect_protocol(self.objective),
        gen_args_used=bool(self.gen_args),
        gen_kwargs_used=bool(self.gen_kwargs),
        run_args_used=bool(self.run_args),
        run_kwargs_used=bool(self.run_kwargs),
        trial_nums=job_trial_nums,
    )

    # decide which trials to execute
    if self.config.resume:
        self._renumber_trial_folders(
            self.workdir.resolve(), self.config.padding_style
        )
        status_tracker.refresh_from_disk(n_proc=self.config.n_proc)
    status_tracker.dump_to_path(self.workdir / JOB_STATUS_FNAME)

    to_run = status_tracker.pending_trial_nums

    if not to_run:
        logger.info("All trials were already completed. Nothing to do.")
        return bool(status_tracker.unsuccessful_trial_nums)

    if self.config.is_parallel:
        logger.info(
            f"Running {len(to_run)} trials with {self.config.n_proc} processors. {status_tracker.n_done}/{self.config.n_trial} ({status_tracker.progress:.2%}) trials completed."
        )
        # needed for logging on Windows
        with multiprocessing.Manager() as m:
            log_queue = m.Queue()

            parent_log_level = logging.getLogger().getEffectiveLevel()
            listener = QueueListener(log_queue, *logging.getLogger().handlers)
            listener.start()

            executor = ProcessPoolExecutor(
                max_workers=self.config.n_proc,
                initializer=worker_init,
                initargs=(log_queue, parent_log_level),
            )
            try:
                future_to_tn = {
                    executor.submit(
                        self._execute_trial_subprocess,
                        tn,
                        self.seed,
                        global_overrides.model_dump(),
                    ): tn
                    for tn in to_run
                }
                for f in as_completed(future_to_tn):
                    tn = future_to_tn[f]
                    try:
                        trial_status = f.result()
                        status_tracker.update_trial(status=trial_status)
                    except (BdbQuit, KeyboardInterrupt):
                        # user is quitting from breakpoint() or CTRL+C
                        raise
                    except Exception as e:
                        logger.exception(f"Trial {tn} failed: {e}")
                        status_tracker.update_trial(
                            status=TrialStatus(
                                trial_num=tn, completion=Completion.ERROR
                            )
                        )
                    status_tracker.generate_report()
            except (BdbQuit, KeyboardInterrupt):
                # allows killing the job with one CTRL+C
                logger.warning("Interrupt recieved. Stopping all trials.")
                executor.shutdown(wait=False, cancel_futures=True)
                raise
            finally:
                listener.stop()
                executor.shutdown(wait=True)
    else:
        for tn in to_run:
            try:
                _, trial_status, _ = self.execute_single_trial(
                    trial_num=tn,
                    seed=self.seed,
                    overrides_payload=global_overrides.model_dump(),
                )
                status_tracker.update_trial(
                    status=trial_status,
                )
            except (BdbQuit, KeyboardInterrupt):
                # user is quitting from breakpoint() or CTRL+C
                raise
            except Exception as e:
                logger.exception(f"A trial failed with error: {e}")
                status_tracker.update_trial(
                    status=TrialStatus(trial_num=tn, completion=Completion.ERROR)
                )
            status_tracker.generate_report()

    status_tracker.generate_report(alert_generation=True)
    return bool(status_tracker.unsuccessful_trial_nums)

@staticmethod
def get_slurm_array_string(ids: list[int]) -> str:
    """Collapses [0, 1, 2, 5, 6] into '0-2,5-6' for SLURM."""
    if not ids:
        return ""

    ranges = []
    # Identify groups of consecutive integers
    for _, group in itertools.groupby(
        enumerate(sorted(ids)), lambda x: x[1] - x[0]
    ):
        group = list(group)
        start = group[0][1]
        end = group[-1][1]
        if start == end:
            ranges.append(str(start))
        else:
            ranges.append(f"{start}-{end}")

    return ",".join(ranges)

@staticmethod
def get_slurm_partitions() -> tuple[list[str], str | None]:
    """Queries sinfo for available partitions and identifies the default."""
    try:
        result = subprocess.run(
            ["sinfo", "-h", "--format=%P"],
            capture_output=True,
            text=True,
            timeout=5,
        )
        if result.returncode == 0:
            raw_partitions = [
                p.strip() for p in result.stdout.splitlines() if p.strip()
            ]

            default_partition = None
            clean_partitions = []

            for p in raw_partitions:
                if p.endswith("*"):
                    name = p.replace("*", "")
                    default_partition = name
                    clean_partitions.append(name)
                else:
                    clean_partitions.append(p)

            return sorted(list(set(clean_partitions))), default_partition
    except Exception:
        pass
    return [], None

@staticmethod
def normalize_to_mb(mem_str: str) -> int:
    """Converts SLURM memory strings (e.g., '1000', '1G', '1024M') to integer MB."""
    # Split the number from the unit (e.g., '1024M' -> '1024', 'M')
    match = re.match(r"(\d+)([KMGTP]?)", mem_str.upper())
    if not match:
        return 0

    value, unit = match.groups()
    value = int(value)

    # SLURM units are powers of 1024
    multipliers = {
        "K": 1 / 1024,  # Kilobytes to MB
        "M": 1,  # Megabytes
        "G": 1024,  # Gigabytes to MB
        "T": 1024**2,  # Terabytes to MB
        "P": 1024**3,  # Petabytes to MB
    }

    return int(value * multipliers.get(unit or "M", 1))

@staticmethod
def format_bytes(mb_value: int) -> str:
    """Scales MB back up to the most readable unit (G, T, etc.)."""
    units = ["M", "G", "T", "P"]
    value = float(mb_value)
    unit_index = 0

    # Keep dividing by 1024 as long as it's a clean multiple
    while value >= 1024 and unit_index < len(units) - 1:
        value /= 1024
        unit_index += 1

    # If it's a whole number (like 1.0G), show it as 1G.
    # Otherwise, show one decimal place (like 1.5G).
    if value.is_integer():
        return f"{int(value)}{units[unit_index]}"
    return f"{value:.1f}{units[unit_index]}"

@staticmethod
def get_slurm_node_mem_limit(partition_name: str) -> str:
    """Finds the MINIMUM RealMemory limit, normalized to MB."""
    try:
        # 1. Get nodes from partition
        part_info = subprocess.run(
            ["scontrol", "show", "partition", partition_name, "-o"],
            capture_output=True,
            text=True,
        ).stdout
        node_match = re.search(r"\bNodes=(\S+)", part_info)
        if not node_match:
            return "<UNKNOWN>"

        # 2. Get node info (can handle ranges like c[1-2])
        node_info = subprocess.run(
            ["scontrol", "show", "node", node_match.group(1), "-o"],
            capture_output=True,
            text=True,
        ).stdout

        # 3. Find all RealMemory values (capturing optional suffixes)
        # Regex captures digits and any trailing letters (like G or M)
        mem_matches = re.findall(r"\bRealMemory=(\d+[KMGTP]?)", node_info)
        if mem_matches:
            # Normalize every match to MB and find the lowest one
            min_mb = min(MojoRunner.normalize_to_mb(m) for m in mem_matches)
            return MojoRunner.format_bytes(min_mb)

    except Exception:
        pass

    return "<UNKNOWN>"

@staticmethod
def get_slurm_cpu_limit(partition_name: str) -> str:
    """Finds the minimum of the max allowed CPUs (CPUTot) among nodes in a partition."""
    try:
        # 1. Get the nodes in the partition
        part_cmd = ["scontrol", "show", "partition", partition_name, "-o"]
        part_info = subprocess.run(part_cmd, capture_output=True, text=True).stdout

        node_match = re.search(r"\bNodes=(\S+)", part_info)
        if not node_match:
            return "<UNKNOWN>"

        # 2. Get the node info
        node_cmd = ["scontrol", "show", "node", node_match.group(1), "-o"]
        nodes_info = subprocess.run(node_cmd, capture_output=True, text=True).stdout

        # 3. Find all CPUTot values (The total physical/logical CPUs on the node)
        cpu_values = re.findall(r"\bCPUTot=(\d+)", nodes_info)

        if cpu_values:
            # Find the minimum to ensure any node can handle the task
            return str(min(int(v) for v in cpu_values))

    except Exception:
        pass

    return "<UNKNOWN>"

@staticmethod
def slurm_time_to_seconds(time_str: str) -> int:
    """Converts SLURM time (D-HH:MM:SS or HH:MM:SS) to total seconds."""
    if time_str.upper() == "UNLIMITED":
        return -1

    # Format: Days-Hours:Minutes:Seconds
    days = 0
    if "-" in time_str:
        days_part, time_str = time_str.split("-")
        days = int(days_part)

    parts = list(map(int, time_str.split(":")))
    if len(parts) == 3:  # HH:MM:SS
        return days * 86400 + parts[0] * 3600 + parts[1] * 60 + parts[2]
    if len(parts) == 2:  # MM:SS
        return days * 86400 + parts[0] * 60 + parts[1]
    if len(parts) == 1:  # MM
        return days * 86400 + parts[0] * 60
    return 0

@staticmethod
def get_slurm_time_limit(partition_name: str) -> str:
    """Finds the MaxTime limit for a specific partition."""
    try:
        # Get partition info
        part_cmd = ["scontrol", "show", "partition", partition_name, "-o"]
        part_info = subprocess.run(part_cmd, capture_output=True, text=True).stdout

        # Look for MaxTime= followed by the time string (e.g., 1-00:00:00 or UNLIMITED)
        time_match = re.search(r"\bMaxTime=(\S+)", part_info)
        if time_match:
            return time_match.group(1)

    except Exception:
        pass

    return "<UNKNOWN>"

@staticmethod
def get_max_array_size() -> int:
    """Queries the global SLURM configuration for MaxArraySize."""
    try:
        result = subprocess.run(
            ["scontrol", "show", "config"], capture_output=True, text=True
        ).stdout
        match = re.search(r"MaxArraySize=(\d+)", result)
        return int(match.group(1)) if match else 1001
    except Exception:
        return 1001