utrain 0.0.4__py3-none-any.whl

utrain/cli/attempts.py

@@ -0,0 +1,116 @@
+import pathlib
+
+import sqlalchemy
+import sqlalchemy.orm
+
+from . import db as dbmod
+from . import exceptions, output, reconcile
+
+
+def list_attempts(run_id_prefix: str, session: sqlalchemy.orm.Session) -> None:
+    run_id = dbmod.resolve_run_id(run_id_prefix, session)
+
+    attempt_n = dbmod.latest_attempt(run_id, session)
+    if attempt_n is not None:
+        reconcile.reconcile_attempt(run_id, attempt_n, session)
+
+    rows = (
+        session.execute(
+            sqlalchemy.select(dbmod.run_attempts)
+            .where(dbmod.run_attempts.c.run_id == run_id)
+            .order_by(dbmod.run_attempts.c.attempt.desc())
+        )
+        .mappings()
+        .fetchall()
+    )
+
+    rid = dbmod.short_run_id(run_id, session)
+
+    headers = ["ATTEMPT", "FROM_PHASE", "STATUS", "STARTED", "ENDED"]
+    table_rows = [
+        [
+            f"{rid}/{r['attempt']}",
+            str(r["from_phase"]) if r["from_phase"] else "--",
+            str(r["status"]),
+            output.format_time(r["started_at"]),
+            output.format_time(r["ended_at"]),
+        ]
+        for r in rows
+    ]
+    print(output.format_table(headers, table_rows))
+
+
+def list_attempt_ids(run_id_prefix: str, session: sqlalchemy.orm.Session) -> list[str]:
+    run_id = dbmod.resolve_run_id(run_id_prefix, session)
+
+    rows = (
+        session.execute(
+            sqlalchemy.select(dbmod.run_attempts.c.attempt)
+            .where(dbmod.run_attempts.c.run_id == run_id)
+            .order_by(dbmod.run_attempts.c.attempt.desc())
+        )
+        .scalars()
+        .fetchall()
+    )
+    return [f"{run_id}/{attempt}" for attempt in rows]
+
+
+def show_attempt(run_id_prefix: str, attempt_n: int, session: sqlalchemy.orm.Session) -> None:
+    run_id = dbmod.resolve_run_id(run_id_prefix, session)
+    reconcile.reconcile_attempt(run_id, attempt_n, session)
+
+    run_row = dbmod.get_run(run_id, session)
+    run_dir = pathlib.Path(str(run_row["run_dir"]))
+
+    attempt_row = (
+        session.execute(
+            sqlalchemy.select(dbmod.run_attempts).where(
+                (dbmod.run_attempts.c.run_id == run_id)
+                & (dbmod.run_attempts.c.attempt == attempt_n)
+            )
+        )
+        .mappings()
+        .fetchone()
+    )
+    if attempt_row is None:
+        raise exceptions.UI(f"abort: attempt {attempt_n} not found for run '{run_id}'")
+
+    phase_rows = (
+        session.execute(
+            sqlalchemy.select(dbmod.run_phases)
+            .where(
+                (dbmod.run_phases.c.run_id == run_id) & (dbmod.run_phases.c.attempt == attempt_n)
+            )
+            .order_by(dbmod.run_phases.c.phase_order)
+        )
+        .mappings()
+        .fetchall()
+    )
+
+    print(f"run:        {run_id} ({run_row['name']})")
+    print(f"attempt:    {attempt_n}")
+    print(f"from_phase: {attempt_row['from_phase'] or '--'}")
+    print(f"status:     {attempt_row['status']}")
+    print(f"started:    {output.format_time(attempt_row['started_at'])}")
+    print(f"ended:      {output.format_time(attempt_row['ended_at'])}")
+
+    if phase_rows:
+        print()
+        headers = ["PHASE", "ORDER", "STATUS", "STARTED", "ENDED"]
+        table_rows = [
+            [
+                str(p["phase"]),
+                str(p["phase_order"]),
+                str(p["status"]),
+                output.format_time(p["started_at"]),
+                output.format_time(p["ended_at"]),
+            ]
+            for p in phase_rows
+        ]
+        print(output.format_table(headers, table_rows))
+
+    logs_dir = run_dir / "attempt" / str(attempt_n) / "logs"
+    wandb_dir = run_dir / "attempt" / str(attempt_n) / "wandb"
+    print()
+    print(f"logs: {logs_dir}")
+    print(f"data: {wandb_dir}")

utrain/cli/compute.py

@@ -0,0 +1,172 @@
+import json
+import os
+import pathlib
+import subprocess
+
+
+class CpuInfo:
+    name: str
+    cores: int
+    mem_total_gb: float
+    mem_available_gb: float
+
+    def __init__(
+        self,
+        name: str,
+        cores: int,
+        mem_total_gb: float,
+        mem_available_gb: float,
+    ) -> None:
+        self.name = name
+        self.cores = cores
+        self.mem_total_gb = mem_total_gb
+        self.mem_available_gb = mem_available_gb
+
+
+class GpuInfo:
+    index: int
+    name: str
+    power_draw: float
+    power_limit: float
+    util: int
+    mem_used_mb: float
+    mem_total_mb: float
+
+    def __init__(
+        self,
+        index: int,
+        name: str,
+        power_draw: float,
+        power_limit: float,
+        util: int,
+        mem_used_mb: float,
+        mem_total_mb: float,
+    ) -> None:
+        self.index = index
+        self.name = name
+        self.power_draw = power_draw
+        self.power_limit = power_limit
+        self.util = util
+        self.mem_used_mb = mem_used_mb
+        self.mem_total_mb = mem_total_mb
+
+
+class ComputeInfo:
+    cpu: CpuInfo
+    gpus: list[GpuInfo]
+
+    def __init__(self, cpu: CpuInfo, gpus: list[GpuInfo]) -> None:
+        self.cpu = cpu
+        self.gpus = gpus
+
+
+def _read_cpu_name() -> str:
+    try:
+        text = pathlib.Path("/proc/cpuinfo").read_text()
+        for line in text.splitlines():
+            if line.startswith("model name"):
+                return line.split(":", 1)[1].strip()
+    except OSError:
+        pass
+    return "Unknown CPU"
+
+
+def _read_meminfo() -> tuple[float, float]:
+    try:
+        text = pathlib.Path("/proc/meminfo").read_text()
+        total_kb = 0
+        available_kb = 0
+        for line in text.splitlines():
+            if line.startswith("MemTotal:"):
+                total_kb = int(line.split()[1])
+            elif line.startswith("MemAvailable:"):
+                available_kb = int(line.split()[1])
+        return total_kb / (1024 * 1024), available_kb / (1024 * 1024)
+    except OSError:
+        return 0.0, 0.0
+
+
+def _read_gpus() -> list[GpuInfo]:
+    try:
+        result = subprocess.run(
+            [
+                "nvidia-smi",
+                "--query-gpu=index,name,power.draw,enforced.power.limit,memory.used,memory.total,utilization.gpu",
+                "--format=csv,noheader,nounits",
+            ],
+            capture_output=True,
+            text=True,
+            timeout=10,
+        )
+        if result.returncode != 0:
+            return []
+        gpus: list[GpuInfo] = []
+        for line in result.stdout.strip().splitlines():
+            parts = [p.strip() for p in line.split(",")]
+            if len(parts) < 7:
+                continue
+
+            def _float(s: str) -> float:
+                try:
+                    return float(s)
+                except ValueError:
+                    return 0.0
+
+            def _int(s: str) -> int:
+                try:
+                    return int(s)
+                except ValueError:
+                    return 0
+
+            gpus.append(
+                GpuInfo(
+                    index=_int(parts[0]),
+                    name=parts[1],
+                    power_draw=_float(parts[2]),
+                    power_limit=_float(parts[3]),
+                    mem_used_mb=_float(parts[4]),
+                    mem_total_mb=_float(parts[5]),
+                    util=_int(parts[6]),
+                )
+            )
+        return gpus
+    except (FileNotFoundError, subprocess.TimeoutExpired):
+        return []
+
+
+def collect_compute() -> ComputeInfo:
+    fixture = os.environ.get("UTRAIN_COMPUTE_FIXTURE")
+    if fixture:
+        data: dict[str, object] = json.loads(pathlib.Path(fixture).read_text())
+        cpu_data = data["cpu"]
+        assert isinstance(cpu_data, dict)
+        cpu = CpuInfo(
+            name=str(cpu_data["name"]),
+            cores=int(cpu_data["cores"]),  # type: ignore[arg-type]
+            mem_total_gb=float(cpu_data["mem_total_gb"]),  # type: ignore[arg-type]
+            mem_available_gb=float(cpu_data["mem_available_gb"]),  # type: ignore[arg-type]
+        )
+        gpus: list[GpuInfo] = []
+        for i, g in enumerate(data.get("gpus", [])):  # type: ignore[union-attr]
+            assert isinstance(g, dict)
+            gpus.append(
+                GpuInfo(
+                    index=int(g.get("index", i)),  # type: ignore[arg-type]
+                    name=str(g["name"]),
+                    power_draw=float(g["power_draw"]),  # type: ignore[arg-type]
+                    power_limit=float(g["power_limit"]),  # type: ignore[arg-type]
+                    util=int(g["util"]),  # type: ignore[arg-type]
+                    mem_used_mb=float(g["mem_used_mb"]),  # type: ignore[arg-type]
+                    mem_total_mb=float(g["mem_total_mb"]),  # type: ignore[arg-type]
+                )
+            )
+        return ComputeInfo(cpu=cpu, gpus=gpus)
+
+    mem_total, mem_available = _read_meminfo()
+    cpu = CpuInfo(
+        name=_read_cpu_name(),
+        cores=os.cpu_count() or 1,
+        mem_total_gb=mem_total,
+        mem_available_gb=mem_available,
+    )
+    return ComputeInfo(cpu=cpu, gpus=_read_gpus())

utrain/cli/db.py

@@ -0,0 +1,178 @@
+import collections.abc
+import contextlib
+import sqlite3
+
+import sqlalchemy
+import sqlalchemy.event
+import sqlalchemy.orm
+
+from .. import config
+from . import exceptions
+
+metadata = sqlalchemy.MetaData()
+
+runs = sqlalchemy.Table(
+    "runs",
+    metadata,
+    sqlalchemy.Column("id", sqlalchemy.Text, primary_key=True),
+    sqlalchemy.Column("name", sqlalchemy.Text, nullable=False),
+    sqlalchemy.Column("image", sqlalchemy.Text, nullable=False),
+    sqlalchemy.Column("compute", sqlalchemy.Text, nullable=False),
+    sqlalchemy.Column("run_dir", sqlalchemy.Text, nullable=False),
+    sqlalchemy.Column("status", sqlalchemy.Text, nullable=False, default="configuring"),
+    sqlalchemy.Column("config_hash", sqlalchemy.Text, nullable=True),
+    sqlalchemy.Column("created_at", sqlalchemy.Float, nullable=False),
+)
+
+run_attempts = sqlalchemy.Table(
+    "run_attempts",
+    metadata,
+    sqlalchemy.Column("run_id", sqlalchemy.Text, nullable=False),
+    sqlalchemy.Column("attempt", sqlalchemy.Integer, nullable=False),
+    sqlalchemy.Column("from_phase", sqlalchemy.Text, nullable=True),
+    sqlalchemy.Column("status", sqlalchemy.Text, nullable=False),
+    sqlalchemy.Column("pid", sqlalchemy.Integer, nullable=True),
+    sqlalchemy.Column("started_at", sqlalchemy.Float, nullable=False),
+    sqlalchemy.Column("ended_at", sqlalchemy.Float, nullable=True),
+    sqlalchemy.PrimaryKeyConstraint("run_id", "attempt"),
+)
+
+run_phases = sqlalchemy.Table(
+    "run_phases",
+    metadata,
+    sqlalchemy.Column("run_id", sqlalchemy.Text, nullable=False),
+    sqlalchemy.Column("attempt", sqlalchemy.Integer, nullable=False),
+    sqlalchemy.Column("phase", sqlalchemy.Text, nullable=False),
+    sqlalchemy.Column("phase_order", sqlalchemy.Integer, nullable=False),
+    sqlalchemy.Column("status", sqlalchemy.Text, nullable=False),
+    sqlalchemy.Column("started_at", sqlalchemy.Float, nullable=True),
+    sqlalchemy.Column("ended_at", sqlalchemy.Float, nullable=True),
+    sqlalchemy.PrimaryKeyConstraint("run_id", "attempt", "phase"),
+)
+
+
+def _set_sqlite_pragmas(dbapi_conn: sqlite3.Connection, _record: object) -> None:
+    # A detached orchestrator writes the DB concurrently with foreground read
+    # commands (notably `run show --wait`). Wait up to 5s for a held lock instead
+    # of failing immediately with "database is locked".
+    cursor = dbapi_conn.cursor()
+    cursor.execute("PRAGMA busy_timeout=5000")
+    cursor.close()
+
+
+def create_engine(settings: config.Settings) -> sqlalchemy.Engine:
+    url = f"sqlite:///{settings.db_path}"
+    engine = sqlalchemy.create_engine(url, connect_args={"check_same_thread": False})
+    sqlalchemy.event.listen(engine, "connect", _set_sqlite_pragmas)
+    return engine
+
+
+def init_db(engine: sqlalchemy.Engine) -> None:
+    metadata.create_all(engine)
+    _migrate(engine)
+
+
+def _migrate(engine: sqlalchemy.Engine) -> None:
+    with engine.begin() as conn:
+        inspector = sqlalchemy.inspect(engine)
+        existing = {
+            t: {c["name"] for c in inspector.get_columns(t)} for t in inspector.get_table_names()
+        }
+
+        # Rename gpu column to compute if it exists and compute doesn't
+        if "runs" in existing and "gpu" in existing["runs"] and "compute" not in existing["runs"]:
+            conn.execute(sqlalchemy.text("ALTER TABLE runs RENAME COLUMN gpu TO compute"))
+            # Refresh existing set after the rename
+            existing["runs"].discard("gpu")
+            existing["runs"].add("compute")
+
+        # Add new columns to legacy 'runs' table if it came from the old schema
+        if "runs" in existing:
+            for col, ddl in [
+                ("name", "TEXT NOT NULL DEFAULT ''"),
+                ("image", "TEXT NOT NULL DEFAULT ''"),
+                ("compute", "TEXT NOT NULL DEFAULT 'cpu'"),
+                ("config_hash", "TEXT"),
+                ("created_at", "REAL NOT NULL DEFAULT 0"),
+            ]:
+                if col not in existing["runs"]:
+                    conn.execute(sqlalchemy.text(f"ALTER TABLE runs ADD COLUMN {col} {ddl}"))
+
+        # Drop columns that don't belong in the new schema (SQLite can't DROP columns
+        # before 3.35; skip silently — the extra columns are harmless).
+
+
+@contextlib.contextmanager
+def with_db(
+    settings: config.Settings,
+) -> collections.abc.Generator[sqlalchemy.orm.Session, None, None]:
+    engine = create_engine(settings)
+    init_db(engine)
+    with sqlalchemy.orm.Session(engine) as session:
+        try:
+            yield session
+            session.commit()
+        except Exception:
+            session.rollback()
+            raise
+
+
+def resolve_run_id(prefix: str, session: sqlalchemy.orm.Session) -> str:
+    rows = (
+        session.execute(sqlalchemy.select(runs.c.id).where(runs.c.id.like(f"{prefix}%")))
+        .scalars()
+        .fetchall()
+    )
+    if not rows:
+        raise exceptions.UI(f"abort: run '{prefix}' not found")
+    if len(rows) > 1:
+        matches = ", ".join(str(r) for r in rows[:4])
+        raise exceptions.UI(f"abort: id prefix '{prefix}' is ambiguous (matches: {matches})")
+    return str(rows[0])
+
+
+def short_run_id(run_id: str, session: sqlalchemy.orm.Session) -> str:
+    """Shortest prefix of run_id that is unique across the runs table.
+
+    Uses run_id's lexicographic neighbors (PK-indexed), so it avoids scanning the
+    whole table. The result may be one char shorter than the uniform width used by
+    'run list', but still resolves unambiguously through resolve_run_id.
+    """
+    pred = session.execute(
+        sqlalchemy.select(runs.c.id).where(runs.c.id < run_id).order_by(runs.c.id.desc()).limit(1)
+    ).scalar_one_or_none()
+    succ = session.execute(
+        sqlalchemy.select(runs.c.id).where(runs.c.id > run_id).order_by(runs.c.id.asc()).limit(1)
+    ).scalar_one_or_none()
+
+    def _lcp(a: str, b: str | None) -> int:
+        if b is None:
+            return 0
+        n = 0
+        for ca, cb in zip(a, b):
+            if ca != cb:
+                break
+            n += 1
+        return n
+
+    plen = 1 + max(
+        _lcp(run_id, str(pred) if pred is not None else None),
+        _lcp(run_id, str(succ) if succ is not None else None),
+    )
+    return run_id[:plen]
+
+
+def latest_attempt(run_id: str, session: sqlalchemy.orm.Session) -> int | None:
+    result = session.execute(
+        sqlalchemy.select(sqlalchemy.func.max(run_attempts.c.attempt)).where(
+            run_attempts.c.run_id == run_id
+        )
+    ).scalar_one_or_none()
+    return int(result) if result is not None else None
+
+
+def get_run(run_id: str, session: sqlalchemy.orm.Session) -> sqlalchemy.engine.RowMapping:
+    row = session.execute(sqlalchemy.select(runs).where(runs.c.id == run_id)).mappings().fetchone()
+    if row is None:
+        raise exceptions.UI(f"abort: run '{run_id}' not found")
+    return row

utrain/cli/debug.py

@@ -0,0 +1,38 @@
+import logging
+import os
+import sys
+
+FORMAT = "%(asctime)s:%(levelname)s:%(module)s.%(funcName)s:%(message)s"
+DATEFMT = "%H:%M:%S"
+
+
+def setup(debug: int, log_filename: str) -> None:
+    log_level = os.getenv("UTRAIN_LOG_LEVEL")
+    match log_level:
+        case "DEBUG":
+            debug = 3
+        case "INFO":
+            debug = 2
+        case "WARNING":
+            debug = 1
+        case None:
+            pass
+        case _:
+            try:
+                debug = int(log_level)
+            except ValueError:
+                pass
+    if debug == 0:
+        return
+    match debug:
+        case 1:
+            level = logging.WARNING
+        case 2:
+            level = logging.INFO
+        case _:
+            level = logging.DEBUG
+    try:
+        f = open(log_filename, "a", buffering=1)
+    except Exception:
+        f = sys.stdout
+    logging.basicConfig(stream=f, level=level, format=FORMAT, datefmt=DATEFMT, force=True)

utrain/cli/exceptions.py

@@ -0,0 +1,2 @@
+class UI(Exception):
+    pass

utrain/cli/images.py

@@ -0,0 +1,121 @@
+import subprocess
+
+import sqlalchemy
+import sqlalchemy.orm
+
+from .. import container
+from . import db as dbmod
+from . import exceptions
+
+
+class ImageInfo:
+    name: str
+    size_str: str
+    run_count: int
+
+    def __init__(self, name: str, size_str: str, run_count: int) -> None:
+        self.name = name
+        self.size_str = size_str
+        self.run_count = run_count
+
+
+def _image_exists(ref: str) -> bool:
+    result = subprocess.run(["podman", "image", "exists", ref])
+    return result.returncode == 0
+
+
+def _format_size(num_bytes: int) -> str:
+    size = float(num_bytes)
+    for unit in ("B", "K", "M", "G"):
+        if size < 1024 or unit == "G":
+            return f"{size:.0f}{unit}"
+        size /= 1024
+    return "?"
+
+
+def _image_size(name: str) -> str:
+    try:
+        result = subprocess.run(
+            [
+                "podman",
+                "image",
+                "inspect",
+                container.podman.image_ref(name),
+                "--format",
+                "{{.Size}}",
+            ],
+            capture_output=True,
+            text=True,
+            timeout=30,
+        )
+        if result.returncode == 0:
+            return _format_size(int(result.stdout.strip()))
+    except (FileNotFoundError, ValueError, subprocess.TimeoutExpired):
+        pass
+    return "?"
+
+
+def list_images(session: sqlalchemy.orm.Session) -> list[ImageInfo]:
+    presets = container.podman.list_presets()
+    images: list[ImageInfo] = []
+    for name in sorted(presets):
+        run_count = session.execute(
+            sqlalchemy.select(sqlalchemy.func.count()).where(
+                (dbmod.runs.c.image == name) & (dbmod.runs.c.status != "deleted")
+            )
+        ).scalar_one()
+        images.append(ImageInfo(name=name, size_str=_image_size(name), run_count=int(run_count)))
+    return images
+
+
+def add_image(url: str) -> str:
+    # Derive <name> from the last path component of the URL (strip scheme and tag).
+    path_part = url.split("://", 1)[-1]
+    base = path_part.split("/")[-1].split(":")[0]
+    name = container.podman.preset_key(base)
+
+    # `podman://` is an enroot transport, not a podman one; it means "already in
+    # the local store", so there is nothing to pull. Everything else goes to
+    # `podman pull` with its scheme intact (docker://, or a bare registry ref).
+    if not url.startswith("podman://") and not _image_exists(path_part):
+        result = subprocess.run(["podman", "pull", url])
+        if result.returncode != 0:
+            raise exceptions.UI(f"abort: podman pull failed (exit {result.returncode})")
+
+    # Tag from the scheme-stripped ref: that is the name a pull stores locally,
+    # and `podman tag` rejects a transport prefix.
+    result = subprocess.run(
+        ["podman", "tag", path_part, container.podman.image_ref(name)],
+        capture_output=True,
+        text=True,
+    )
+    if result.returncode != 0:
+        raise exceptions.UI(f"abort: podman tag failed: {result.stderr.strip()}")
+
+    return name
+
+
+def remove_image(name: str, session: sqlalchemy.orm.Session, force: bool = False) -> None:
+    if name not in container.podman.list_presets():
+        raise exceptions.UI(f"abort: image '{name}' not found")
+
+    run_count = session.execute(
+        sqlalchemy.select(sqlalchemy.func.count()).where(
+            (dbmod.runs.c.image == name) & (dbmod.runs.c.status != "deleted")
+        )
+    ).scalar_one()
+
+    if int(run_count) > 0 and not force:
+        raise exceptions.UI(
+            f"abort: image '{name}' is used by {run_count} run(s); use --force to remove anyway"
+        )
+
+    # Quiet: `podman rmi` reports every tag it drops ("Untagged: ..."), which is
+    # podman's bookkeeping, not utrain's output. Removal is silent on success.
+    result = subprocess.run(
+        ["podman", "rmi", container.podman.image_ref(name)],
+        capture_output=True,
+        text=True,
+    )
+    if result.returncode != 0:
+        raise exceptions.UI(f"abort: podman rmi failed: {result.stderr.strip()}")