Support iris runners other than ryan

marin/train_tomat_tpu.py

@@ -44,17 +44,16 @@
 
 Prereqs:
 - `gs://marin-eu-west4/tomat/tokenized/<label>/worker-*/*.parquet` populated
-- ADC refreshed for `ryan.williams@openathena.ai` on hai-gcp-models.
+- ADC refreshed for your `@openathena.ai` account on hai-gcp-models
+  (`gcloud auth application-default login`).
 """
 
 from __future__ import annotations
 
-import dataclasses
 import json
 import os
 import sys
 from datetime import timedelta
-from pathlib import Path
 
 # Multihost-capable JAX init. Historically we called `jax.distributed.initialize()`
 # up-front here because Levanter's `WandbConfig.init` would call
@@ -546,7 +545,7 @@ def _assert_cache_local(cache_dir: str) -> None:
     or `TOMAT_ALLOW_XREG_CACHE=1` is set.
     """
     if os.environ.get("TOMAT_ALLOW_XREG_CACHE") == "1":
-        print(f"[tomat-tpu] TOMAT_ALLOW_XREG_CACHE=1 → skipping x-reg check", flush=True)
+        print("[tomat-tpu] TOMAT_ALLOW_XREG_CACHE=1 → skipping x-reg check", flush=True)
         return
     cache_region = _bucket_region(cache_dir)
     worker_region = _detect_gce_region()
@@ -1046,7 +1045,7 @@ def _has_ckpts(_bucket: str) -> bool:
               f"window_blocks={shuffle_window_blocks})")
     else:
         shuffle_cfg = False
-        print(f"[tomat-tpu] shuffle: OFF (TOMAT_SHUFFLE_WINDOW_BLOCKS=0)")
+        print("[tomat-tpu] shuffle: OFF (TOMAT_SHUFFLE_WINDOW_BLOCKS=0)")
 
     # `data_cfg_cls` selects between standard `LmDataConfig` and the F1
     # subclass that routes `F1PrebuiltLmDatasetFormat` components through

scripts/cron_iris_sync_modal.py

@@ -82,18 +82,32 @@
     # complete here.
     .pip_install(
         "boto3",
-        "click>=8.3.1", "cloudpickle>=3.1.2", "connect-python>=0.9.0",
+        "click>=8.3.1",
+        "cloudpickle>=3.1.2",
+        "connect-python>=0.9.0",
         # fsspec/gcsfs/s3fs pin to matching minor versions — keep all three
         # in sync (mirror what the laptop venv has).
-        "fsspec==2025.3.0", "gcsfs==2025.3.0", "google-auth>=2.0",
-        "google-cloud-tpu>=1.18.0", "grpcio>=1.76.0", "httpx>=0.28.1",
-        "humanfriendly>=10.0", "pydantic>=2.0", "pyjwt>=2.12.0",
-        "pyyaml>=6.0", "s3fs>=2024.0.0", "starlette>=0.50.0",
-        "tabulate>=0.9.0", "typing-extensions>=4.0",
-        "uvicorn[standard]>=0.23.0", "zstandard>=0.22.0",
+        "fsspec==2025.3.0",
+        "gcsfs==2025.3.0",
+        "google-auth>=2.0",
+        "google-cloud-tpu>=1.18.0",
+        "grpcio>=1.76.0",
+        "httpx>=0.28.1",
+        "humanfriendly>=10.0",
+        "pydantic>=2.0",
+        "pyjwt>=2.12.0",
+        "pyyaml>=6.0",
+        "s3fs>=2024.0.0",
+        "starlette>=0.50.0",
+        "tabulate>=0.9.0",
+        "typing-extensions>=4.0",
+        "uvicorn[standard]>=0.23.0",
+        "zstandard>=0.22.0",
         "s3fs==2025.3.0",
         # marin-finelog's deps not yet covered above:
-        "duckdb>=1.0.0", "protobuf>=5.0", "pyarrow>=19.0.0",
+        "duckdb>=1.0.0",
+        "protobuf>=5.0",
+        "pyarrow>=19.0.0",
         # finelog/store/log_namespace.py imports numpy at module level:
         "numpy>=1.26",
     )
@@ -103,15 +117,23 @@
         f"marin-rigging @ git+{MARIN_REPO}@{MARIN_SHA}#subdirectory=lib/rigging",
         extra_options="--no-deps",
     )
+    # Ship the canonical runner roster (scripts/iris_runners.py) into the
+    # image so `_sync` can import it at runtime (it isn't a pip dep).
+    .add_local_python_source("iris_runners")
 )
 
 adc_secret = modal.Secret.from_name("marin-iris-adc")  # GCP ADC for iris RPC
 r2_secret = modal.Secret.from_name("oa-r2-write")  # AWS_ACCESS_KEY_ID/SECRET
 
 app = modal.App("tomat-iris-sync-cron", image=image)
 
-# Same prefix set as `tomat iris sync` uses on the laptop.
-PREFIXES = ["/ryan/tomat", "/ryan/train", "/ryan/eval"]
+# The runner roster / prefix set comes from scripts/iris_runners.py (shipped
+# into the image above); `_sync` builds the prefixes from it at runtime.
+
+# Bound on simultaneous iris RPCs (each opens its own tunnel to the
+# controller). The cron fires every minute, so fan the per-prefix RPCs out
+# concurrently rather than sequentially (8×~60s would overrun the interval).
+MAX_SYNC_WORKERS = 4
 
 R2_ACCOUNT_ID = "43a6f2d588b1483733189d39418ec5be"
 R2_ENDPOINT = f"https://{R2_ACCOUNT_ID}.r2.cloudflarestorage.com"
@@ -126,6 +148,7 @@ def _materialize_adc() -> str:
     created the secret (`GOOGLE_APPLICATION_CREDENTIALS_JSON`).
     """
     import os
+
     raw = os.environ["GOOGLE_APPLICATION_CREDENTIALS_JSON"]
     path = "/tmp/adc.json"
     with open(path, "w") as f:
@@ -144,16 +167,27 @@ def _iris_job_list_json(prefix: str, timeout: int = 60) -> list[dict]:
     import sys
     import time
 
-    cmd = ["iris", "--cluster=marin", "job", "list",
-           "--prefix", prefix, "--json"]
+    cmd = ["iris", "--cluster=marin", "job", "list", "--prefix", prefix, "--json"]
     t0 = time.monotonic()
-    r = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)
+    try:
+        r = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)
+    except subprocess.TimeoutExpired:
+        # The function docstring promises a slow call just skips this prefix's
+        # upload — honor that instead of crashing the whole Modal cron run.
+        elapsed_ms = int((time.monotonic() - t0) * 1000)
+        print(
+            f"[iris-rpc {prefix} {elapsed_ms}ms] TIMEOUT (>{timeout}s, skip)",
+            file=sys.stderr,
+        )
+        return []
     elapsed_ms = int((time.monotonic() - t0) * 1000)
     tag = "SLOW" if elapsed_ms > 10_000 else ""
     print(f"[iris-rpc {prefix} {elapsed_ms}ms] {tag}", file=sys.stderr)
     if r.returncode != 0:
-        print(f"iris job list failed: rc={r.returncode}\n{r.stderr[-500:]}",
-              file=sys.stderr)
+        print(
+            f"iris job list failed: rc={r.returncode}\n{r.stderr[-500:]}",
+            file=sys.stderr,
+        )
         return []
     # Iris CLI emits log lines before the JSON; find the JSON array start.
     out = r.stdout
@@ -170,6 +204,7 @@ def _build_payload(rows_by_prefix: dict[str, list[dict]]) -> dict:
     crash-loops as healthy RUNNING jobs.
     """
     import datetime
+
     jobs: dict[str, dict] = {}
     for rows in rows_by_prefix.values():
         for row in rows:
@@ -192,15 +227,24 @@ def _build_payload(rows_by_prefix: dict[str, list[dict]]) -> dict:
                 "failures": int(row.get("failure_count") or 0),
                 "error": row.get("error") or None,
                 "exit_code": row.get("exit_code"),
-                "submitted_at_ms": int((row.get("submitted_at") or {}).get("epoch_ms") or 0) or None,
-                "started_at_ms": int((row.get("started_at") or {}).get("epoch_ms") or 0) or None,
-                "finished_at_ms": int((row.get("finished_at") or {}).get("epoch_ms") or 0) or None,
+                "submitted_at_ms": int(
+                    (row.get("submitted_at") or {}).get("epoch_ms") or 0
+                )
+                or None,
+                "started_at_ms": int((row.get("started_at") or {}).get("epoch_ms") or 0)
+                or None,
+                "finished_at_ms": int(
+                    (row.get("finished_at") or {}).get("epoch_ms") or 0
+                )
+                or None,
                 "num_tasks": int(row.get("task_count") or 0),
                 "task_state_counts": tsc,
             }
     return {
         "schema_version": 1,
-        "synced_at": datetime.datetime.now(datetime.timezone.utc).isoformat(timespec="seconds"),
+        "synced_at": datetime.datetime.now(datetime.timezone.utc).isoformat(
+            timespec="seconds"
+        ),
         "count": len(jobs),
         "jobs": jobs,
     }
@@ -210,6 +254,7 @@ def _r2_put_json(payload: dict) -> int:
     import json
     import os
     import boto3
+
     data = json.dumps(payload, indent=1).encode("utf-8")
     s3 = boto3.client(
         "s3",
@@ -218,34 +263,56 @@ def _r2_put_json(payload: dict) -> int:
         aws_secret_access_key=os.environ["AWS_SECRET_ACCESS_KEY"],
         region_name="auto",
     )
-    s3.put_object(Bucket=R2_BUCKET, Key=R2_KEY, Body=data,
-                  ContentType="application/json")
+    s3.put_object(
+        Bucket=R2_BUCKET, Key=R2_KEY, Body=data, ContentType="application/json"
+    )
     return len(data)
 
 
 def _sync() -> dict:
     """The actual work — runs in both the cron and one-shot entrypoints."""
     import sys
+    from concurrent.futures import ThreadPoolExecutor
+
+    from iris_runners import iris_prefixes  # shipped via add_local_python_source
+
     _materialize_adc()
-    rows_by_prefix = {p: _iris_job_list_json(p) for p in PREFIXES}
+    prefixes = iris_prefixes()
+    # Concurrent fan-out (bounded) — sequential per-prefix RPCs would overrun
+    # the 1-min cron. ex.map preserves input order, so zip pairs cleanly.
+    with ThreadPoolExecutor(max_workers=MAX_SYNC_WORKERS) as ex:
+        rows_by_prefix = dict(zip(prefixes, ex.map(_iris_job_list_json, prefixes)))
     payload = _build_payload(rows_by_prefix)
     # Safety: never upload an empty state. A broken iris CLI / missing dep /
     # auth failure should NOT clobber the existing R2 snapshot — let the
     # dashboard keep showing the last-good state until we fix the cron.
     if payload["count"] == 0:
-        print("[iris sync] 0 jobs across all prefixes — refusing to "
-              "overwrite R2 (likely iris/auth failure)", file=sys.stderr)
+        print(
+            "[iris sync] 0 jobs across all prefixes — refusing to "
+            "overwrite R2 (likely iris/auth failure)",
+            file=sys.stderr,
+        )
         return {"jobs": 0, "bytes": 0, "skipped": True}
     n_bytes = _r2_put_json(payload)
-    print(f"[iris sync] {payload['count']} jobs → "
-          f"s3://{R2_BUCKET}/{R2_KEY} ({n_bytes} bytes)", file=sys.stderr)
+    print(
+        f"[iris sync] {payload['count']} jobs → "
+        f"s3://{R2_BUCKET}/{R2_KEY} ({n_bytes} bytes)",
+        file=sys.stderr,
+    )
     return {"jobs": payload["count"], "bytes": n_bytes}
 
 
 @app.function(
     cpu=1,
     memory=1024,
-    timeout=120,
+    # 240s ceiling: at MAX_SYNC_WORKERS=4 and len(PREFIXES)=8 the worst-case
+    # wall-clock is ceil(8/4) × per-prefix_timeout = 2 × 60s = 120s with zero
+    # overhead headroom — equal to the function's own timeout, so a sibling
+    # prefix would always race the cron's own cutoff. 240s gives us a clean
+    # 2× margin and still leaves the safety check (`count == 0` → skip
+    # upload) free to bail if iris is wedged. Bump again if we add more
+    # runners (the worst-case scales with ceil(len(PREFIXES) / 4)).
+    timeout=240,
     secrets=[adc_secret, r2_secret],
     schedule=modal.Cron("* * * * *"),  # every minute UTC
 )
@@ -256,7 +323,14 @@ def cron_sync() -> dict:
 @app.function(
     cpu=1,
     memory=1024,
-    timeout=120,
+    # 240s ceiling: at MAX_SYNC_WORKERS=4 and len(PREFIXES)=8 the worst-case
+    # wall-clock is ceil(8/4) × per-prefix_timeout = 2 × 60s = 120s with zero
+    # overhead headroom — equal to the function's own timeout, so a sibling
+    # prefix would always race the cron's own cutoff. 240s gives us a clean
+    # 2× margin and still leaves the safety check (`count == 0` → skip
+    # upload) free to bail if iris is wedged. Bump again if we add more
+    # runners (the worst-case scales with ceil(len(PREFIXES) / 4)).
+    timeout=240,
     secrets=[adc_secret, r2_secret],
 )
 def sync_once() -> dict: