trainmeter 0.0.2__py3-none-any.whl

trainmeter/cli.py

@@ -0,0 +1,382 @@
+"""The `tm` command: `tm [options] -- <command>` or `tm script.py [args]`."""
+
+from __future__ import annotations
+
+import os
+import re
+import socket
+import sys
+import time
+import warnings
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+from . import __version__
+from .config import VALUE_FLAGS, declared_facts
+from .engine import Engine
+from .records import make_record, redact_argv
+from .report import format_elapsed, format_summary
+from .sources.gpu import GpuSampler, valid_gpus_spec
+from .sources.host import HostSampler, read_uptime
+from .supervisor.ingest import Pipeline
+from .supervisor.launcher import launch
+from .supervisor.live import LiveState
+from .supervisor.store import RunStore
+from .web.server import DEFAULT_PORT, WebServer
+
+USAGE = """\
+usage: tm [options] -- <command> [args...]
+       tm [options] script.py [args...]
+       tm doctor [--gpus LIST]      what tm can see on this node
+       tm passport [RUN]            the run passport of a finished run, for a model card
+       tm export wandb [RUN] --project P   upload a finished run to Weights & Biases
+       tm view [RUN]                a finished run in the dashboard
+       tm attach PID                watch a program that is already running, from outside
+
+Runs the command untouched and records the run in .trainmeter/runs/.
+
+options:
+  --name NAME          label added to the run directory name
+  --gpus LIST          GPUs to sample, e.g. 0,1 (default: CUDA_VISIBLE_DEVICES, else all)
+  --port N             first port to try for the dashboard (default 8765)
+  --no-web             do not serve the dashboard
+  -q, --quiet          print no banner and no summary line
+  --no-agent           do not inject the agent into the training process
+  --no-tap             do not pick up series from wandb, TensorBoard and MLflow
+  --version            print the version and exit
+  -h, --help           print this help and exit
+
+declare what only you know (20B, 1.5e9 and 300M are accepted):
+  --budget-tokens N    --budget-steps N     --budget-flops N
+  --tokens-per-step N  --flops-per-token N  --convention NAME
+  --params N           --embedding-params N --layers N
+  --d-model N          --seq-len N          --dp N
+  --plan-mfu X         --plan-hours H       what you planned, for the Plan vs actual view
+"""
+
+RUN_ID_ENV = "TRAINMETER_RUN_ID"
+SOCK_ENV = "TRAINMETER_SOCK"
+AGENT_ENV = "TRAINMETER_AGENT"
+
+
+@dataclass(frozen=True)
+class Parsed:
+    command: list[str]
+    name: str | None = None
+    quiet: bool = False
+    agent: bool = True
+    declared: dict[str, Any] = field(default_factory=dict)
+    gpus: str | None = None
+    port: int = DEFAULT_PORT
+    web: bool = True
+    tap: bool = True
+
+
+def parse_args(argv: Sequence[str]) -> Parsed | int:
+    """Parse the options before the command. Returns an exit status for help, version, errors."""
+    opts: dict[str, Any] = {
+        "name": None,
+        "gpus": None,
+        "port": DEFAULT_PORT,
+        "web": True,
+        "tap": True,
+        "quiet": False,
+        "agent": True,
+        "values": {},
+    }
+    i = 0
+    while i < len(argv):
+        token = argv[i]
+        flag, eq, inline = token.partition("=")
+        if token == "--":
+            return _finish(list(argv[i + 1 :]), opts)
+        if token in ("-q", "--quiet"):
+            opts["quiet"] = True
+        elif token == "--no-agent":
+            opts["agent"] = False
+        elif token == "--no-tap":
+            opts["tap"] = False
+        elif token == "--no-web":
+            opts["web"] = False
+        elif flag in ("--name", "--gpus", "--port", *VALUE_FLAGS):
+            if eq:
+                value = inline
+            else:
+                i += 1
+                if i >= len(argv):
+                    return _fail(f"{flag} needs a value")
+                value = argv[i]
+            if flag == "--name":
+                opts["name"] = value
+            elif flag == "--port":
+                if not value.isdigit() or not 0 <= int(value) <= 65535:
+                    return _fail("--port: expected a port number, 0 picks any free port")
+                opts["port"] = int(value)
+            elif flag == "--gpus":
+                if not valid_gpus_spec(value):
+                    return _fail("--gpus: expected indices such as 0,1 or GPU UUIDs")
+                opts["gpus"] = value
+            else:
+                opts["values"][flag] = value
+        elif token in ("-h", "--help"):
+            print(USAGE, end="")
+            return 0
+        elif token == "--version":
+            print(f"trainmeter {__version__}")
+            return 0
+        elif token.startswith("-"):
+            return _fail(f"unknown option {token}")
+        else:
+            return _finish(list(argv[i:]), opts)
+        i += 1
+    return _finish([], opts)
+
+
+def _fail(message: str) -> int:
+    print(f"tm: {message}\n\n{USAGE}", end="", file=sys.stderr)
+    return 2
+
+
+def _finish(command: list[str], opts: dict[str, Any]) -> Parsed | int:
+    if not command:
+        return _fail("no command given")
+    try:
+        declared = declared_facts(opts["values"])
+    except ValueError as exc:
+        return _fail(str(exc))
+    if command[0].endswith(".py") and Path(command[0]).is_file():
+        command = [sys.executable, *command]  # shorthand: run the script with this interpreter
+    return Parsed(
+        command,
+        opts["name"],
+        opts["quiet"],
+        opts["agent"],
+        declared,
+        opts["gpus"],
+        opts["port"],
+        opts["web"],
+        opts["tap"],
+    )
+
+
+def default_name(command: list[str]) -> str:
+    """Label for the run directory: the script for `python train.py`, else the program."""
+    first = Path(command[0]).name
+    if (
+        re.fullmatch(r"python[0-9.]*", first)
+        and len(command) > 1
+        and not command[1].startswith("-")
+    ):
+        return Path(command[1]).stem
+    return Path(first).stem
+
+
+def _display(path: Path) -> str:
+    try:
+        return str(path.relative_to(Path.cwd()))
+    except ValueError:
+        return str(path)
+
+
+def _bootstrap_dir() -> str:
+    return str(Path(__file__).parent / "agent" / "bootstrap")
+
+
+def main(argv: Sequence[str] | None = None, env: Mapping[str, str] | None = None) -> int:
+    args = list(sys.argv[1:] if argv is None else argv)
+    if args[:1] == ["passport"]:
+        from .commands import passport_main
+
+        return passport_main(args[1:])
+    if args[:1] == ["export"]:
+        from .commands import export_main
+
+        return export_main(args[1:])
+    if args[:1] == ["view"]:
+        from .viewer import view_main
+
+        return view_main(args[1:])
+    if args[:1] == ["attach"]:
+        from .viewer import attach_main
+
+        return attach_main(args[1:])
+    if args[:1] == ["doctor"]:
+        from .doctor import main as doctor_main
+
+        return doctor_main(args[1:], env)
+    parsed = parse_args(args)
+    if isinstance(parsed, int):
+        return parsed
+    environ = dict(os.environ if env is None else env)
+
+    if RUN_ID_ENV in environ:
+        print("tm: already running under tm, running the command untouched", file=sys.stderr)
+        try:
+            os.execvpe(parsed.command[0], parsed.command, environ)
+        except OSError as exc:
+            print(f"tm: cannot run {parsed.command[0]!r}: {exc}", file=sys.stderr)
+            return 127 if isinstance(exc, FileNotFoundError) else 126
+
+    store = _open_store(parsed.name or default_name(parsed.command))
+    engine = Engine()
+    live = LiveState(engine, _header(parsed))  # feeds the engine under a lock, for every thread
+
+    def handle(record: dict[str, Any]) -> None:
+        if store is not None:
+            store.append_record(record)
+        live.feed(record)
+
+    pipeline = _open_pipeline(handle)
+    if store is not None:
+        environ[RUN_ID_ENV] = store.run_id
+        if not parsed.quiet:
+            print(f"trainmeter: run {_display(store.path)}", file=sys.stderr)
+        _write_meta(store, live.header)
+    handle(make_record("run.start", "supervisor", time.monotonic(), live.header, pid=os.getpid()))
+    _declare(handle, parsed.declared)
+    if pipeline is not None:
+        environ[SOCK_ENV] = pipeline.path
+        if parsed.agent:
+            environ["PYTHONPATH"] = os.pathsep.join(
+                p for p in (_bootstrap_dir(), environ.get("PYTHONPATH", "")) if p
+            )
+    if not parsed.agent:
+        environ[AGENT_ENV] = "0"
+    if not parsed.tap:
+        environ["TRAINMETER_TAP"] = "0"
+
+    web = WebServer(live, parsed.port) if parsed.web else None
+    url = web.start() if web is not None else None
+    if url is not None:
+        live.start()
+        if not parsed.quiet:
+            print(f"trainmeter: dashboard {url}", file=sys.stderr)
+    gpu_sampler = GpuSampler(
+        lambda d: handle(make_record("sample.gpu", "gpu", time.monotonic(), d)),
+        gpus=parsed.gpus,
+        cuda_visible=environ.get("CUDA_VISIBLE_DEVICES"),
+    )
+    gpu_sampler.start()
+    sampler: HostSampler | None = None
+
+    def on_spawn(pid: int) -> None:
+        nonlocal sampler
+        sampler = HostSampler(
+            pid, lambda d: handle(make_record("sample.host", "host", time.monotonic(), d))
+        )
+        sampler.start()
+
+    try:
+        outcome = launch(parsed.command, env=environ, on_spawn=on_spawn)
+    finally:
+        gpu_sampler.stop()
+        if sampler is not None:
+            sampler.stop()
+        if pipeline is not None:
+            pipeline.close()
+
+    end = {"elapsed_s": outcome.elapsed_s, "exit_status": outcome.status}
+    if outcome.signal is not None:
+        end["signal"] = outcome.signal
+    handle(make_record("run.end", "supervisor", time.monotonic(), end, pid=os.getpid()))
+    live.finish({"exit_status": outcome.status, "signal": outcome.signal})
+    live.stop()
+    if web is not None:
+        web.stop()
+    if outcome.error is not None:
+        print(f"tm: cannot run {parsed.command[0]!r}: {outcome.error}", file=sys.stderr)
+    snap = _snapshot(engine)
+    _finish_run(store, end, snap, pipeline.dropped if pipeline is not None else 0)
+    _write_passport(store, snap, live, engine, end)
+    if not parsed.quiet:
+        print(format_summary(outcome.elapsed_s, snap), file=sys.stderr)
+    return outcome.status
+
+
+def _snapshot(engine: Engine) -> dict[str, Any]:
+    try:
+        return engine.snapshot()
+    except Exception as exc:  # noqa: BLE001 - the exit status is the child's, whatever happens here
+        warnings.warn(f"trainmeter: could not compute the summary: {exc!r}", stacklevel=2)
+        return {}
+
+
+def _open_store(name: str) -> RunStore | None:
+    try:
+        return RunStore.create(Path.cwd(), name)
+    except OSError as exc:
+        warnings.warn(f"trainmeter: no run log, the command runs anyway: {exc}", stacklevel=2)
+        return None
+
+
+def _open_pipeline(handle) -> Pipeline | None:
+    try:
+        pipeline = Pipeline(handle)
+        pipeline.start()
+        return pipeline
+    except OSError as exc:
+        warnings.warn(f"trainmeter: no event socket, the command runs anyway: {exc}", stacklevel=2)
+        return None
+
+
+def _declare(handle, declared: dict[str, Any]) -> None:
+    for key, value in declared.items():
+        handle(
+            make_record(
+                "fact", "supervisor", 0.0, {"key": key, "value": value, "source": "declared"}
+            )
+        )
+
+
+def _header(parsed: Parsed) -> dict[str, Any]:
+    return {
+        "argv": redact_argv(parsed.command),
+        "cwd": os.getcwd(),
+        "host": socket.gethostname(),
+        "tm_version": __version__,
+        "declared": parsed.declared,
+        "node_uptime_s": read_uptime(),
+    }
+
+
+def _write_meta(store: RunStore, header: dict[str, Any]) -> None:
+    try:
+        store.write_json("meta.json", {**header, "mono0": store.mono0, "wall0": store.wall0})
+    except Exception as exc:  # noqa: BLE001 - the run must go on
+        warnings.warn(f"trainmeter: could not record the start: {exc!r}", stacklevel=2)
+
+
+def _finish_run(
+    store: RunStore | None, end: dict[str, Any], snap: dict[str, Any], dropped: int
+) -> None:
+    if store is None:
+        return
+    try:
+        store.write_json("summary.json", {**end, "dropped_datagrams": dropped, **snap})
+    except Exception as exc:  # noqa: BLE001
+        warnings.warn(f"trainmeter: could not record the end: {exc!r}", stacklevel=2)
+    finally:
+        store.close()
+
+
+def _write_passport(
+    store: RunStore | None,
+    snap: dict[str, Any],
+    live: LiveState,
+    engine: Engine,
+    end: dict[str, Any],
+) -> None:
+    if store is None or not snap:
+        return
+    try:
+        from .export.files import write_passport
+        from .passport import build
+
+        write_passport(store.path, build(snap, live.header, live.facts(), engine.hello, end))
+    except Exception as exc:  # noqa: BLE001 - a passport must never cost the exit status
+        warnings.warn(f"trainmeter: could not write the passport: {exc!r}", stacklevel=2)
+
+
+__all__ = ["main", "parse_args", "default_name", "format_elapsed"]

trainmeter/commands.py

@@ -0,0 +1,99 @@
+"""Subcommands that read a finished run: `tm passport` and `tm export wandb`."""
+
+from __future__ import annotations
+
+import json
+import sys
+from pathlib import Path
+
+from .export.files import write_passport
+from .passport import to_markdown
+from .replay import find_run, replay
+
+PASSPORT_USAGE = """\
+usage: tm passport [RUN] [--json] [--write]
+
+Prints the run passport of RUN (a run directory or its name; default: the latest run) as Markdown.
+  --json    print JSON instead
+  --write   write passport.json and passport.md into the run directory
+"""
+
+EXPORT_USAGE = """\
+usage: tm export wandb [RUN] --project PROJECT [--entity ENTITY] [--name NAME]
+
+Uploads a finished run to Weights & Biases. Nothing is sent unless you run this command.
+Needs `pip install wandb` and a logged-in wandb, or WANDB_MODE=offline.
+"""
+
+
+def _fail(message: str, usage: str) -> int:
+    print(f"tm: {message}\n\n{usage}", end="", file=sys.stderr)
+    return 2
+
+
+def passport_main(args: list[str]) -> int:
+    as_json = write = False
+    name = None
+    for arg in args:
+        if arg == "--json":
+            as_json = True
+        elif arg == "--write":
+            write = True
+        elif arg in ("-h", "--help"):
+            print(PASSPORT_USAGE, end="")
+            return 0
+        elif arg.startswith("-") or name is not None:
+            return _fail(f"unexpected argument {arg!r}", PASSPORT_USAGE)
+        else:
+            name = arg
+    try:
+        replayed = replay(find_run(Path.cwd(), name))
+    except FileNotFoundError as exc:
+        print(f"tm: {exc}", file=sys.stderr)
+        return 1
+    data = replayed.passport()
+    if write:
+        for path in write_passport(replayed.path, data):
+            print(f"wrote {path}", file=sys.stderr)
+    sys.stdout.write(json.dumps(data, indent=2) + "\n" if as_json else to_markdown(data))
+    return 0
+
+
+def export_main(args: list[str]) -> int:
+    if args[:1] != ["wandb"]:
+        return _fail("expected `tm export wandb`", EXPORT_USAGE)
+    opts: dict[str, str] = {}
+    name = None
+    rest = args[1:]
+    i = 0
+    while i < len(rest):
+        arg = rest[i]
+        if arg in ("--project", "--entity", "--name"):
+            i += 1
+            if i >= len(rest):
+                return _fail(f"{arg} needs a value", EXPORT_USAGE)
+            opts[arg] = rest[i]
+        elif arg in ("-h", "--help"):
+            print(EXPORT_USAGE, end="")
+            return 0
+        elif arg.startswith("-") or name is not None:
+            return _fail(f"unexpected argument {arg!r}", EXPORT_USAGE)
+        else:
+            name = arg
+        i += 1
+    if "--project" not in opts:
+        return _fail("--project is required", EXPORT_USAGE)
+    try:
+        replayed = replay(find_run(Path.cwd(), name))
+        write_passport(replayed.path, replayed.passport())
+        from .export.wandb import export
+
+        where = export(replayed, opts["--project"], opts.get("--entity"), opts.get("--name"))
+    except FileNotFoundError as exc:
+        print(f"tm: {exc}", file=sys.stderr)
+        return 1
+    except ImportError:
+        print("tm: wandb is not installed here (pip install wandb)", file=sys.stderr)
+        return 1
+    print(f"exported {replayed.path.name} to {where}", file=sys.stderr)
+    return 0

trainmeter/config.py

@@ -0,0 +1,65 @@
+"""Declared facts: the command-line flags that describe the run, turned into facts.
+
+Pure. `tm.toml` is not read yet: it needs `tomllib`, which Python 3.10 lacks.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from .flops import CONVENTIONS
+
+_SUFFIX = {"k": 1e3, "m": 1e6, "b": 1e9, "t": 1e12}
+
+# flag -> (fact key, kind). "count" accepts 20B, 1.5e9, 300M; "int" a whole number.
+VALUE_FLAGS: dict[str, tuple[str, str]] = {
+    "--budget-tokens": ("budget_tokens", "count"),
+    "--budget-steps": ("budget_steps", "count"),
+    "--budget-flops": ("budget_flops", "count"),
+    "--tokens-per-step": ("tokens_per_step", "count"),
+    "--flops-per-token": ("flops_per_token", "count"),
+    "--convention": ("convention", "convention"),
+    "--params": ("n_params", "count"),
+    "--embedding-params": ("n_embedding_params", "count"),
+    "--layers": ("n_layers", "int"),
+    "--d-model": ("d_model", "int"),
+    "--seq-len": ("seq_len", "int"),
+    "--dp": ("dp", "int"),
+    "--plan-mfu": ("plan_mfu", "fraction"),
+    "--plan-hours": ("plan_hours", "positive"),
+}
+
+
+def parse_count(text: str) -> float:
+    """`20B` -> 2e10, `1.5e9` -> 1.5e9, `300M` -> 3e8. Raises ValueError otherwise."""
+    raw = text.strip().lower().replace("_", "")
+    if raw and raw[-1] in _SUFFIX:
+        return float(raw[:-1]) * _SUFFIX[raw[-1]]
+    return float(raw)
+
+
+def declared_facts(options: dict[str, str]) -> dict[str, Any]:
+    """Fact key -> value for the value flags in `options`. Raises ValueError with a message."""
+    facts: dict[str, Any] = {}
+    for flag, raw in options.items():
+        key, kind = VALUE_FLAGS[flag]
+        try:
+            if kind == "convention":
+                if raw not in CONVENTIONS:
+                    raise ValueError(f"expected one of {', '.join(CONVENTIONS)}")
+                facts[key] = raw
+                continue
+            if kind == "fraction":  # 0.4 or 40%
+                text = raw.strip()
+                value = float(text[:-1]) / 100 if text.endswith("%") else float(text)
+                if not 0 < value <= 1:
+                    raise ValueError("expected a fraction in (0, 1] such as 0.4 or 40%")
+                facts[key] = value
+                continue
+            number = parse_count(raw)
+        except ValueError as exc:
+            raise ValueError(f"{flag}: {exc}") from exc
+        if number <= 0:
+            raise ValueError(f"{flag}: must be positive")
+        facts[key] = int(number) if kind == "int" or number == int(number) else number
+    return facts

trainmeter/doctor.py

@@ -0,0 +1,118 @@
+"""`tm doctor`: what trainmeter can see on this node, and why anything is missing.
+
+`run_checks` is pure over a backend and an environment, so it is tested with fakes. The live GPM
+read at the end is the one line that proves the counters work; on real hardware only.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import os
+import sys
+import time
+from collections.abc import Callable, Mapping
+from dataclasses import dataclass
+from pathlib import Path
+
+from . import __version__
+from .peaks import lookup
+from .sources.gpu import Backend, GpuUnavailable, select
+
+OK, WARN, FAIL = "ok", "warn", "FAIL"
+
+
+@dataclass(frozen=True)
+class Check:
+    status: str
+    name: str
+    detail: str
+
+
+def run_checks(
+    backend: Backend,
+    env: Mapping[str, str],
+    *,
+    gpus: str | None = None,
+    proc_root: Path = Path("/proc"),
+    sleep: Callable[[float], None] = time.sleep,
+) -> list[Check]:
+    checks = [Check(OK, "trainmeter", f"{__version__}, python {sys.version.split()[0]}")]
+    checks.append(
+        Check(OK if (proc_root / "uptime").exists() else WARN, "host stats", "read from /proc")
+    )
+    torch_found = importlib.util.find_spec("torch") is not None
+    checks.append(
+        Check(
+            OK if torch_found else WARN,
+            "torch",
+            "importable, the agent can attach" if torch_found else "not importable here",
+        )
+    )
+    try:
+        backend.open()
+    except GpuUnavailable as exc:
+        return [*checks, Check(FAIL, "NVML", f"{exc}; OFU and observed FLOPs will be absent")]
+    try:
+        versions = backend.versions()
+        detail = ", ".join(f"{k} {v}" for k, v in versions.items()) or "initialized"
+        checks.append(Check(OK, "NVML", detail))
+        idents = backend.identities()
+        indices, how, warning = select(idents, gpus, env.get("CUDA_VISIBLE_DEVICES"))
+        label = {"flag": "--gpus", "env": "CUDA_VISIBLE_DEVICES", "all": "all GPUs"}[how]
+        picked = ", ".join(map(str, indices)) or "none"
+        checks.append(
+            Check(WARN if warning or not indices else OK, "selection", f"{label}: {picked}")
+        )
+        if warning:
+            checks.append(Check(WARN, "selection", warning))
+        for i in indices:
+            checks.extend(_device(backend, i, idents[i], sleep))
+    finally:
+        backend.close()
+    return checks
+
+
+def _device(backend: Backend, i: int, ident: dict, sleep) -> list[Check]:
+    name = ident.get("name", "?")
+    peak = lookup(name)
+    out = [
+        Check(
+            OK if peak else WARN,
+            f"GPU {i}",
+            f"{name}, dense BF16 peak {peak.bf16_dense_flops / 1e12:.1f} TFLOP/s ({peak.name})"
+            if peak
+            else f"{name}: not in the peak table, MFU and observed FLOPs will be absent",
+        )
+    ]
+    reason = backend.gpm_support(i)
+    if reason is not None:
+        return [*out, Check(WARN, f"GPU {i} GPM", f"{reason}; OFU and arithmetic intensity absent")]
+    backend.gpm(i)  # primes the first sample
+    sleep(1.0)
+    counters = backend.gpm(i)
+    if not counters or "tensor_active" not in counters:
+        return [
+            *out,
+            Check(WARN, f"GPU {i} GPM", "supported, but a live read returned no counters"),
+        ]
+    shown = ", ".join(f"{k} {v:.3g}" for k, v in sorted(counters.items()))
+    return [*out, Check(OK, f"GPU {i} GPM", f"live read: {shown} (idle GPU: zeros are normal)")]
+
+
+def format_checks(checks: list[Check]) -> str:
+    width = max(len(c.name) for c in checks)
+    return "\n".join(f"[{c.status:>4}] {c.name:<{width}}  {c.detail}" for c in checks) + "\n"
+
+
+def main(argv: list[str], env: Mapping[str, str] | None = None) -> int:
+    from .sources.gpu import default_backend
+
+    gpus = None
+    if argv[:1] == ["--gpus"] and len(argv) == 2:
+        gpus = argv[1]
+    elif argv:
+        print("usage: tm doctor [--gpus LIST]", file=sys.stderr)
+        return 2
+    checks = run_checks(default_backend(), os.environ if env is None else env, gpus=gpus)
+    sys.stdout.write(format_checks(checks))
+    return 1 if any(c.status == FAIL for c in checks) else 0