Iris: compute job CPU wall-time two ways — (A) controller tasks table, (B) iris.task stats namespace. Reproduces PR marin-community/marin#5637 + adds true integrated cpu_seconds.

README.md

iris_job_cpu_time.py

Two ways to compute "CPU wall time" for an Iris job subtree, validated against the marin production cluster.

Background: marin-community/marin#5637 adds iris job cpu-time, which on the controller answers via:

SUM(t.finished_at_ms - t.started_at_ms)  over all leaf-job tasks

Variant A issues that SQL against the controller's tasks table. Variant B uses the controller only to enumerate leaves, then aggregates the iris.task finelog stats namespace (max(ts) − min(ts) per (task_id, attempt_id)).

Run

# Local tunnel from `gcloud compute ssh iris-controller-marin … -L 10001:localhost:10000`
python iris_job_cpu_time.py http://localhost:10001 /gonzalo/exp166-prod-v0.6 --include-failed

Example: /gonzalo/exp166-prod-v0.6 on marin (2026-05-11)

=== A) controller `tasks` table (PR 5637 logic) ===
  leaf tasks       : 5
  cpu_wall_ms      :     190,310,743  (52h51m50s)

=== B) iris.task stats (sample-interval approximation) ===
  leaf tasks       : 6
  approx wall_ms   :     211,671,289  (58h47m51s)
  integrated cpu_s : 237987.9  (true CPU consumed)

Why the variants disagree

metric	A (controller tasks)	B (iris.task stats)
leaf tasks counted	5	6
cpu_wall_ms	190,310,743	211,671,289 (≈ +11 %)
integrated cpu_seconds	—	237,987.9

The gap comes from one in-flight leaf, …/checkpoints-dna-bolinas-exp166-v0.1-p4b_de2b2599-6f71575c/train_lm, which is still RUNNING. Its tasks row has finished_at_ms IS NULL, so A drops it; B keeps streaming samples for it (~23,184 s). Subtracting that leaf gives 188,487 s ≈ 52.4 h — within 1 % of A's 52.9 h. The remaining slack is the iris.task sampling interval (~7 s per attempt × 21 attempts).

Per-leaf breakdown (variant B)

leaf job (…train_lm)	n_attempts	wall_s	cpu_s
…4cc0011d	7	84,021	58,586
…31d3c8a3	1	59,611	94,158
…851a0962	1	34,162	52,449
…6f71575c (running)	10	23,184	23,707
…f57f421f	1	10,680	9,048
…1416b708	1	13	40

The per-leaf cpu_s / wall_s ratio separates "16 h on a 4-core worker at 60 % avg utilization" from "16 h pegging one core" — information that wall-time-summed (A) collapses.

Caveats for variant B (iris.task only)

• No state column in iris.task (lib/iris/src/iris/cluster/worker/stats.py:81), so we can't filter to TASK_STATE_SUCCEEDED. Only the include-failed=True mode is meaningful.
• Running tasks included, unlike A which requires finished_at_ms IS NOT NULL.
• Inter-attempt gaps not captured: A uses one (started_at_ms, finished_at_ms) per task end-to-end (gaps from preempt + restart count); B sums per-attempt sample spans and ignores gaps.
• Sampling-interval resolution: ~7 s per attempt.

Channel used

Both variants speak to the controller URL directly. Variant B uses finelog.client.LogClient.connect(controller_url) — the controller mounts FinelogStatsServiceWSGIApplication at /finelog.stats.StatsService/* (controller/dashboard.py:549), forwarding to the in-process or external finelog server. The equivalent over the generic endpoint proxy is:

POST  <controller>/proxy/system.log-server/finelog.stats.StatsService/Query
Content-Type: application/json
Connect-Protocol-Version: 1
{"sql": "SELECT count(*) FROM \"iris.task\" WHERE ts >= now() - INTERVAL '5 minutes'"}

iris_job_cpu_time.py

#!/usr/bin/env python3
# Copyright The Marin Authors
# SPDX-License-Identifier: Apache-2.0
"""Two ways to compute "CPU wall time" for an Iris job subtree.

Background: PR https://github.com/marin-community/marin/pull/5637 adds
``iris job cpu-time`` which the controller answers via:

    SUM(t.finished_at_ms - t.started_at_ms)  over all leaf-job tasks

This script reproduces that against a running Iris controller two ways:

  (A) cpu_wall_ms_from_controller(...)
      Hits the controller's ``ExecuteRawQuery`` only. Walks the job tree
      client-side (one SELECT on ``jobs``), then issues one SELECT on
      ``tasks`` summing ``finished_at_ms - started_at_ms`` over the
      leaves. This matches PR 5637 exactly.

  (B) cpu_wall_ms_from_stats(...)
      Uses the controller only to enumerate leaves (one SELECT on
      ``jobs``), then aggregates from the finelog ``iris.task`` stats
      namespace via the StatsService that the controller exposes at
      ``/finelog.stats.StatsService/*``. Sums ``max(ts) - min(ts)`` per
      (task_id, attempt_id). Caveats vs (A):

      * iris.task carries no task ``state`` column, so we can't filter
        to TASK_STATE_SUCCEEDED — only the include_failed=True mode is
        meaningful.
      * Still-running tasks are included (PR (A) excludes them: their
        finished_at_ms is NULL).
      * Inter-attempt gaps (preempt + restart) are not captured.
      * Resolution is one sampling interval (~7 s) per attempt.

      In exchange you get true integrated CPU seconds via
      ``Σ avg(cpu_millicores) · Δt`` — a different metric, returned
      alongside.

Usage:
    # The controller URL is whatever ``iris cluster dashboard`` /
    # ``iris cluster dashboard-proxy`` print, or a local tunnel.
    python iris_job_cpu_time.py http://localhost:10001 /gonzalo/exp166-prod-v0.6
"""

from __future__ import annotations

import argparse
import contextlib
import json
import sys

# pip install marin-iris finelog pyarrow
from finelog.client import LogClient
from iris.cli.main import rpc_client
from iris.cluster.worker.stats import IrisTaskStat
from iris.rpc import query_pb2


# ---------------------------------------------------------------------------
# Shared: walk the job tree to find leaves (no WITH RECURSIVE — the
# controller's ExecuteRawQuery rejects non-SELECT prefixes).
# ---------------------------------------------------------------------------


def _exec_raw(client, sql: str) -> list[dict]:
    """Run a SELECT via the controller's ExecuteRawQuery RPC. Returns a
    list of row dicts."""
    resp = client.execute_raw_query(query_pb2.RawQueryRequest(sql=sql))
    headers = [c.name for c in resp.columns]
    return [dict(zip(headers, json.loads(r), strict=True)) for r in resp.rows]


def _leaf_job_ids(client, root_job_id: str) -> list[str]:
    """Return the subset of jobs under ``root_job_id`` that have no children."""
    # Prefix-match on job_id picks up the root and all descendants. job_id is
    # `/<user>/<name>[/<child>...]`, so the LIKE term must escape nothing.
    sql = (
        "SELECT job_id, parent_job_id FROM jobs "
        f"WHERE job_id = '{root_job_id}' OR job_id LIKE '{root_job_id}/%'"
    )
    rows = _exec_raw(client, sql)
    children: dict[str, list[str]] = {}
    ids: set[str] = set()
    for r in rows:
        ids.add(r["job_id"])
        children.setdefault(r["parent_job_id"], []).append(r["job_id"])
    return [j for j in ids if not children.get(j)]


def _sql_in_list(values: list[str]) -> str:
    return ", ".join("'" + v.replace("'", "''") + "'" for v in values)


# ---------------------------------------------------------------------------
# (A) Controller-only — matches PR 5637's logic.
# ---------------------------------------------------------------------------


def cpu_wall_ms_from_controller(
    controller_url: str,
    root_job_id: str,
    *,
    include_failed: bool = True,
) -> dict[str, int]:
    """Compute cpu_wall_ms by reading the controller's ``tasks`` table.

    Reproduces ``GetJobCpuTime`` from PR 5637 exactly (sans the server-side
    WITH RECURSIVE, which we do client-side).
    """
    with rpc_client(controller_url) as client:
        leaves = _leaf_job_ids(client, root_job_id)
        if not leaves:
            return {"leaf_tasks": 0, "cpu_wall_ms": 0}

        # state=4 is TASK_STATE_SUCCEEDED; see job.proto.
        state_filter = "" if include_failed else "AND t.state = 4"
        sql = f"""
            SELECT COUNT(DISTINCT t.task_id)                            AS leaf_tasks,
                   COALESCE(SUM(t.finished_at_ms - t.started_at_ms), 0) AS cpu_wall_ms
            FROM tasks t
            WHERE t.job_id IN ({_sql_in_list(leaves)})
              AND t.started_at_ms IS NOT NULL
              AND t.finished_at_ms IS NOT NULL
              AND t.finished_at_ms > t.started_at_ms
              {state_filter}
        """
        rows = _exec_raw(client, sql)
        return {
            "leaf_tasks": int(rows[0]["leaf_tasks"]),
            "cpu_wall_ms": int(rows[0]["cpu_wall_ms"]),
        }


# ---------------------------------------------------------------------------
# (B) Controller for leaves only — aggregate from iris.task stats.
# ---------------------------------------------------------------------------


def cpu_wall_ms_from_stats(
    controller_url: str,
    root_job_id: str,
) -> dict[str, int | float]:
    """Approximate cpu_wall_ms from the ``iris.task`` finelog namespace.

    Hits the controller for the leaf list (``jobs`` table), then issues
    one Stats Query to the finelog server (also reachable via the
    controller — it mounts ``/finelog.stats.StatsService/*``).
    """
    with rpc_client(controller_url) as client:
        leaves = _leaf_job_ids(client, root_job_id)
    if not leaves:
        return {"leaf_tasks": 0, "approx_cpu_wall_ms": 0, "integrated_cpu_seconds": 0.0}

    # Equivalent over the proxy channel:
    #   POST controller/proxy/system.log-server/finelog.stats.StatsService/Query
    # which forwards to /system/log-server. LogClient.connect against the
    # controller hits the mounted FinelogStatsServiceWSGIApplication directly.
    with contextlib.closing(LogClient.connect(controller_url, timeout_ms=60_000)) as c:
        tbl = c.get_table("iris.task", IrisTaskStat)
        sql = f"""
            WITH per_attempt AS (
              SELECT task_id,
                     attempt_id,
                     EXTRACT(EPOCH FROM (max(ts) - min(ts))) AS wall_s,
                     avg(cpu_millicores)                     AS avg_cpu_m
              FROM "iris.task"
              WHERE regexp_replace(task_id, '/[0-9]+$', '') IN ({_sql_in_list(leaves)})
              GROUP BY task_id, attempt_id
            )
            SELECT count(DISTINCT task_id)                       AS leaf_tasks,
                   round(sum(wall_s) * 1000)                     AS approx_cpu_wall_ms,
                   round(sum(wall_s * avg_cpu_m / 1000.0), 1)    AS integrated_cpu_seconds
            FROM per_attempt
        """
        t = tbl.query(sql, max_rows=10).to_pandas()
        row = t.iloc[0]
        return {
            "leaf_tasks": int(row["leaf_tasks"]),
            "approx_cpu_wall_ms": int(row["approx_cpu_wall_ms"]),
            "integrated_cpu_seconds": float(row["integrated_cpu_seconds"]),
        }


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------


def _fmt_ms(ms: int) -> str:
    h, rem = divmod(ms // 1000, 3600)
    m, s = divmod(rem, 60)
    return f"{h}h{m:02d}m{s:02d}s"


def main(argv: list[str] | None = None) -> int:
    p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    p.add_argument("controller_url", help="e.g. http://localhost:10001")
    p.add_argument("job_id", help="root job id, e.g. /alice/my-job")
    p.add_argument(
        "--include-failed",
        action="store_true",
        help="(controller variant) count all tasks with timestamps, not just SUCCEEDED",
    )
    args = p.parse_args(argv)

    a = cpu_wall_ms_from_controller(args.controller_url, args.job_id, include_failed=args.include_failed)
    print("=== A) controller `tasks` table (PR 5637 logic) ===")
    print(f"  leaf tasks       : {a['leaf_tasks']}")
    print(f"  cpu_wall_ms      : {a['cpu_wall_ms']:>15,d}  ({_fmt_ms(a['cpu_wall_ms'])})")

    b = cpu_wall_ms_from_stats(args.controller_url, args.job_id)
    print()
    print("=== B) iris.task stats (sample-interval approximation) ===")
    print(f"  leaf tasks       : {b['leaf_tasks']}")
    print(f"  approx wall_ms   : {b['approx_cpu_wall_ms']:>15,d}  ({_fmt_ms(b['approx_cpu_wall_ms'])})")
    print(f"  integrated cpu_s : {b['integrated_cpu_seconds']:.1f}  (true CPU consumed)")
    return 0


if __name__ == "__main__":
    sys.exit(main())