statusing ¶

def clear_cache(self) -> None:
    """Wipes the internal trial cache to force a full disk re-scan."""
    self._cache.clear()

def reload_if_new_run(self) -> JobStatus | None:
    """
    Checks `workdir/JOB_STATUS_FNAME` for a `run_id` different from this instance's.

    If the user reran their analysis (e.g. after deleting the workdir), the file on disk now belongs to a different run and this instance's cached state (trial cache, trial_nums, generator, etc.) is stale. Returns a freshly-loaded `JobStatus` for the new run if so, or `None` if the file is missing, unreadable, or still belongs to this run.
    """
    path = self.workdir / JOB_STATUS_FNAME
    if not path.exists():
        return None

    try:
        on_disk = JobStatus.model_validate_json(path.read_text())
    except Exception:
        return None

    if on_disk.id == self.id:
        return None

    logger.info(f"Detected a new job run at {path} (run_id changed). Reloading.")
    return on_disk

def get_trial_status(self, trial_num: int) -> TrialStatus | None:
    """Returns the cached status for a trial, or None if not yet loaded."""
    return self._cache.get(trial_num)

def refresh_from_disk(
    self,
    n_proc: int = 1,
    progress_callback: Callable[[float], None] | None = None,
    force_refetch: bool = False,
    persist: bool = True,
) -> None:
    """
    Scans the workdir to identify which trials still need execution, but only for runs not already in the cache (i.e., completed).

    Args:
        n_proc: Number of worker threads used to scan trial status files.
        progress_callback: Optional callback invoked with scan progress (0-100).
        force_refetch: If `True`, clears the cache before scanning so every trial is re-read from disk.
        persist: If `True` (default), writes the refreshed status back to `workdir/JOB_STATUS_FNAME`. Read-only consumers (e.g. the Dojo monitor) should pass `False` so they don't overwrite the file the run owner manages.

    """
    if force_refetch:
        self.clear_cache()

    max_workers = max(1, n_proc)

    # identify work
    needed_tns = [
        tn
        for tn in self.trial_nums
        if tn not in self._cache
        or self._cache[tn].completion == Completion.INCOMPLETE
    ]

    if not needed_tns:
        if progress_callback:
            return progress_callback(100.0)
        return

    def _load_status(tn: int) -> tuple[int, TrialStatus | None]:
        path = self.trial_num_to_path(tn) / TRIAL_STATUS_FNAME
        if not path.exists():
            return tn, None
        try:
            return tn, TrialStatus.model_validate_json(path.read_text())
        except Exception:
            return tn, None

    # execute with real-time tracking
    total_tasks = len(needed_tns)
    completed_count = 0

    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        future_to_tn = {executor.submit(_load_status, tn): tn for tn in needed_tns}

        # as completed yields futures as soon as done
        for future in as_completed(future_to_tn):
            tn, status = future.result()

            # update the cache immediately
            if status is not None:
                with self._lock:
                    self._cache[tn] = status

            # broadcast status
            completed_count += 1
            if progress_callback:
                pct = (completed_count / total_tasks) * 100
                progress_callback(pct)

    # post calculations
    with self._lock:
        timed_durations = [
            s.td.total_seconds()
            for s in self._cache.values()
            if s.completion in _TIMED_COMPLETIONS
        ]
        if timed_durations:
            avg_sec = sum(timed_durations) / len(timed_durations)
            self.average_trial_duration = timedelta(seconds=avg_sec)

        if not self.is_done:
            self.elapsed = datetime.now(UTC) - self.start_time
    if persist:
        self.dump_to_path(self.workdir / JOB_STATUS_FNAME)

def get_trial_nums_with_status(
    self, *status: Completion, invert: bool = False
) -> list[int]:
    """Returns a list of trial numbers which have the provided status. If invert is True it will provide trial numbers which do not match."""
    to_check = [m for m in Completion if m not in status] if invert else status

    return sorted([tn for tn, comp in self._registry.items() if comp in to_check])

def update_trial(self, status: TrialStatus, save: bool = True) -> None:
    """Updates the internal registry and average trial duration and optionally persists the global status."""
    with self._lock:
        self._cache[status.trial_num] = status

        timed_durations = [
            s.td.total_seconds()
            for s in self._cache.values()
            if s.completion in _TIMED_COMPLETIONS and s.td is not None
        ]

        if timed_durations:
            self.average_trial_duration = timedelta(
                seconds=sum(timed_durations) / len(timed_durations)
            )

        self.elapsed = datetime.now(UTC) - self.start_time

        if save:
            self.dump_to_path(self.workdir / JOB_STATUS_FNAME)

def completion_timeline(self, n_bins: int = 40) -> list[TimelineBin]:
    """
    Reconstructs cumulative outcome counts over the job's wall-clock life, for the monitor's stacked progress histogram.

    Each trial's completion time is taken as the latest recorded step end (`started + elapsed`). A trial that is still `INCOMPLETE` but has begun generating or solving counts as "running" from that start time onward rather than "pending"; claimed-but-not-yet-started trials (only `pending.started` set) still count as pending. The axis runs from job start to now if any trial is currently running, or to the last completion event if the job is idle/abandoned (so a stalled job doesn't render a flat tail stretching to now). For each of `n_bins` evenly spaced timestamps in that range, the bin reports how many trials had succeeded, failed requirements, errored, or were running by that time; everything else is pending.
    """
    with self._lock:
        statuses = list(self._cache.values())

    events: list[tuple[datetime, Completion]] = []
    running_starts: list[datetime] = []
    for s in statuses:
        if s.completion == Completion.INCOMPLETE:
            run_start = s.generating.started or s.solving.started
            if run_start is not None:
                running_starts.append(run_start)
            continue
        step_ends = [
            step.started + timedelta(seconds=step.elapsed)
            for step in (s.pending, s.generating, s.solving)
            if step.started is not None and step.elapsed is not None
        ]
        if not step_ends:
            continue
        events.append((max(step_ends), s.completion))

    if not events and not running_starts:
        return []

    start = self.start_time
    if running_starts:
        # the job is actively progressing right now: let the axis track
        # "now" rather than freezing at the last completion event
        end = datetime.now(UTC)
    else:
        # nothing is running: cap the axis at the last completion event so
        # a job that stopped (or was abandoned) long ago doesn't render a
        # flat tail stretching to now; self.elapsed can't be used here
        # since refresh_from_disk keeps bumping it to now-start while the
        # job is not done
        end = min(datetime.now(UTC), max(ts for ts, _ in events))
    total_span = (end - start).total_seconds()
    if total_span <= 0:
        return []

    bins: list[TimelineBin] = []
    for i in range(1, n_bins + 1):
        t = start + timedelta(seconds=total_span * i / n_bins)
        n_s = sum(1 for ts, c in events if ts <= t and c == Completion.SUCCESS)
        n_f = sum(1 for ts, c in events if ts <= t and c in _FAILED_COMPLETIONS)
        n_e = sum(1 for ts, c in events if ts <= t and c == Completion.ERROR)
        n_r = sum(1 for rs in running_starts if rs <= t)
        bins.append(
            {
                "t": (t - start).total_seconds(),
                # microsecond precision: Plotly plots this as a categorical
                # x-axis, so two bins that round to the same whole second
                # (common with many bins over a short span) would collide
                # onto a single x position and appear summed/overlaid
                # instead of as separate bars
                "label": self._utc_to_local(t).strftime("%Y-%m-%d %H:%M:%S.%f"),
                "n_success": n_s,
                "n_failed": n_f,
                "n_error": n_e,
                "n_running": n_r,
                "n_pending": max(0, self.n_trial - n_s - n_f - n_e - n_r),
            }
        )
    return bins

def to_monitor_json(self, n_proc: int | None = None) -> dict:
    """Returns a lightweight summary optimized for the Alpine.js dashboard."""
    from mujoco_mojo.runtime.signal_manager import SignalManager

    # We trigger the disk refresh here so the data is fresh
    # obj = self.model_validate_json((self.workdir / JOB_STATUS_FNAME).read_text())
    # obj.refresh_from_disk(n_proc=obj.n_proc if n_proc is None else n_proc)

    avg_seconds = self.average_trial_duration.total_seconds()

    # Avoid divide by zero if no trials have finished yet
    throughput = 0
    if avg_seconds > 0:
        throughput = (60.0 / avg_seconds) * self.n_proc

    success_tns = self.success_trial_nums
    failed_tns = self.failed_trial_nums
    error_tns = self.error_trial_nums

    def _with_db(tns: list[int]) -> list[int]:
        return [
            tn
            for tn in tns
            if (
                self.trial_num_to_path(tn) / SignalManager.default_output_name()
            ).exists()
        ]

    return {
        "progress": self.progress * 100,
        "n_success": self.n_success,
        "n_failed": self.n_failed,
        "n_error": self.n_error,
        "n_trial": self.n_trial,
        "n_done": self.n_done,
        "n_remaining": self.n_remaining,
        "failure_rate": self.failure_rate,
        "error_rate": self.error_rate,
        "throughput": round(throughput, 1),
        "avg_duration": str(self.average_trial_duration).split(".")[0],
        "time_remaining": str(self.time_remaining_average_success).split(".")[0],
        "elapsed": str(self.elapsed).split(".")[0],
        "is_complete": self.progress >= 1.0,
        "success_tns": success_tns,
        "failure_tns": failed_tns,
        "error_tns": error_tns,
        "last_success_tn": max(success_tns) if success_tns else None,
        "last_failure_tn": max(failed_tns) if failed_tns else None,
        "last_error_tn": max(error_tns) if error_tns else None,
        "start_time": f"{self._utc_to_local(self.start_time).strftime('%Y-%m-%d %H:%M:%S')} {self.local_tzabbr}",
        "end_time": f"{self._utc_to_local(self.end_time).strftime('%Y-%m-%d %H:%M:%S')} {self.local_tzabbr}",
        "version": f"mujoco-mojo v{version('mujoco-mojo')}",
        "padding_style": self.padding_style,
        "failure_tns_with_db": _with_db(failed_tns),
        "error_tns_with_db": _with_db(error_tns),
        "timeline": self.completion_timeline(),
    }