From e95d33e5798c3f4230fa02a358b9c5a6bf128fc1 Mon Sep 17 00:00:00 2001 From: John Toman Date: Thu, 16 Jul 2026 22:53:33 -0700 Subject: [PATCH] Read job id from env Makes it significantly easier to learn the run id from the job list --- composer/diagnostics/timing.py | 7 ++++++- composer/pipeline/cli.py | 1 + 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/composer/diagnostics/timing.py b/composer/diagnostics/timing.py index 09a860e0..e99760ce 100644 --- a/composer/diagnostics/timing.py +++ b/composer/diagnostics/timing.py @@ -8,6 +8,7 @@ """ from contextlib import asynccontextmanager, contextmanager +import os from logging import Logger import time from contextvars import ContextVar @@ -75,6 +76,10 @@ class PhaseRecord: final_link: str | None = None # URL (cloud) or local results directory (local) of the last prover run. token_usage_by_model: dict[str, TokenTotals] = field(default_factory=dict) +def run_id_generator() -> str: + if (res := os.getenv("AUTOPROVER_JOB_ID")): + return res + return uuid.uuid4().hex @dataclass class RunSummary: @@ -85,7 +90,7 @@ class RunSummary: _active_prover_by_task: dict[str, tuple[float, int]] = field(default_factory=dict, repr=False) prover_reported_ms_total: int = 0 # Run-wide prover-REPORTED runtime, summed over every prover run. Distinct from prover_total_s, which is composer's client-side wall-clock (and so includes cloud queue / polling / result download). _active_prover_reported_by_task: dict[str, int] = field(default_factory=dict, repr=False) # Maps task_id -> prover-reported ms accumulated while the task is in flight. - run_id: str = field(default_factory=lambda: uuid.uuid4().hex) # Maps task_id -> (prover_s_accum, prover_calls) recorded while task is in flight. + run_id: str = field(default_factory=run_id_generator) # Maps task_id -> (prover_s_accum, prover_calls) recorded while task is in flight. _latest_link_by_task: dict[str, str] = field(default_factory=dict, repr=False) # Maps task_id -> link for the most recent prover run. token_usage_by_model: dict[str, TokenTotals] = field(default_factory=dict) # Maps model_name -> accumulated raw token counts across the whole run. _active_tokens_by_task: dict[str, dict[str, TokenTotals]] = field(default_factory=dict, repr=False) # Maps task_id -> {model_name -> token counts} accumulated while the task is in flight. diff --git a/composer/pipeline/cli.py b/composer/pipeline/cli.py index c398bf9f..ccdd16c5 100644 --- a/composer/pipeline/cli.py +++ b/composer/pipeline/cli.py @@ -155,6 +155,7 @@ async def cli_pipeline[P: enum.Enum, H]( model = get_model() text_log, events_log = setup_autoprove_logging(project_root, thread_id) print(f"autoprove logs: {text_log}\n events: {events_log}", file=sys.stderr) + print(f"Selected run id: {summary.run_id}") install_run_summary(summary) disc_cache_ns: tuple[str, ...] | None = (