clusterpilot 0.1.0__py3-none-any.whl

clusterpilot/__init__.py

@@ -0,0 +1,3 @@
+"""ClusterPilot — AI-assisted HPC workflow manager."""
+
+__version__ = "0.1.0"

clusterpilot/__main__.py

@@ -0,0 +1,139 @@
+"""ClusterPilot entry point.
+
+Usage
+-----
+    clusterpilot                    # launch TUI
+    clusterpilot init               # create starter ~/.config/clusterpilot/config.toml
+    clusterpilot daemon run         # run poll daemon in foreground (no TUI)
+    clusterpilot daemon install     # install systemd user service
+"""
+from __future__ import annotations
+
+import argparse
+import sys
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        prog="clusterpilot",
+        description="AI-assisted HPC workflow manager",
+    )
+    sub = parser.add_subparsers(dest="cmd")
+
+    sub.add_parser("init", help="Create starter config at ~/.config/clusterpilot/config.toml")
+
+    daemon_p = sub.add_parser("daemon", help="Daemon management")
+    daemon_sub = daemon_p.add_subparsers(dest="daemon_cmd")
+    daemon_sub.add_parser("run", help="Run poll daemon in foreground")
+    daemon_sub.add_parser("install", help="Install systemd user service unit")
+
+    args = parser.parse_args()
+
+    if args.cmd == "init":
+        _cmd_init()
+    elif args.cmd == "daemon":
+        if args.daemon_cmd == "run":
+            _cmd_daemon_run()
+        elif args.daemon_cmd == "install":
+            _cmd_daemon_install()
+        else:
+            daemon_p.print_help()
+    else:
+        _cmd_tui()
+
+
+# ── Subcommands ───────────────────────────────────────────────────────────────
+
+def _cmd_init() -> None:
+    from clusterpilot.config import CONFIG_PATH, write_default_config
+    if CONFIG_PATH.exists():
+        print(f"Config already exists: {CONFIG_PATH}")
+        return
+    write_default_config()
+    print(f"Config written to {CONFIG_PATH}")
+    print("Edit it to set your cluster username and account, then run: clusterpilot")
+
+
+def _cmd_daemon_run() -> None:
+    import asyncio
+    import aiosqlite
+    from clusterpilot.config import ConfigError, load_config
+    from clusterpilot.db import DB_PATH, init_db
+    from clusterpilot.jobs.daemon import PollDaemon
+
+    try:
+        config = load_config()
+    except ConfigError as exc:
+        print(f"Error: {exc}", file=sys.stderr)
+        sys.exit(1)
+
+    async def _run() -> None:
+        DB_PATH.parent.mkdir(parents=True, exist_ok=True)
+        async with aiosqlite.connect(DB_PATH) as db:
+            await init_db(db)
+        daemon = PollDaemon(config, DB_PATH)
+        await daemon.run_forever()
+
+    print("ClusterPilot daemon running. Press Ctrl-C to stop.")
+    try:
+        asyncio.run(_run())
+    except KeyboardInterrupt:
+        print("\nDaemon stopped.")
+
+
+def _cmd_daemon_install() -> None:
+    from clusterpilot.jobs.daemon import write_service_file
+    path = write_service_file()
+    print(f"Service file written to: {path}")
+    print()
+    print("Enable and start with:")
+    print("  systemctl --user daemon-reload")
+    print("  systemctl --user enable --now clusterpilot-poll.service")
+
+
+def _cmd_tui() -> None:
+    import logging
+    from clusterpilot.config import ConfigError, load_config, write_default_config
+    from clusterpilot.tui.app import ClusterPilotApp
+
+    logging.basicConfig(
+        level=logging.WARNING,
+        format="%(levelname)s %(name)s: %(message)s",
+    )
+
+    try:
+        config = load_config()
+    except ConfigError:
+        # First run — write default config and guide the user.
+        write_default_config()
+        from clusterpilot.config import CONFIG_PATH
+        print("Welcome to ClusterPilot!")
+        print()
+        print(f"A starter config has been written to:\n  {CONFIG_PATH}")
+        print()
+        print("Edit it to add your cluster username and account,")
+        print("then run 'clusterpilot' again.")
+        sys.exit(0)
+
+    if not config.clusters:
+        print("No clusters defined in config. Edit ~/.config/clusterpilot/config.toml.")
+        sys.exit(1)
+
+    if not config.api_key:
+        import os
+        if not os.environ.get("ANTHROPIC_API_KEY"):
+            print(
+                "Warning: no API key configured. "
+                "Script generation will fail.\n"
+                "Set api_key in config.toml or export ANTHROPIC_API_KEY."
+            )
+
+    from clusterpilot.db import DB_PATH
+    DB_PATH.parent.mkdir(parents=True, exist_ok=True)
+
+    app = ClusterPilotApp(config)
+    app.run()
+
+
+if __name__ == "__main__":
+    main()

clusterpilot/cluster/__init__.py

@@ -0,0 +1,33 @@
+from clusterpilot.cluster.probe import (
+    ClusterProbe,
+    PartitionInfo,
+    load_cache,
+    probe_cluster,
+    save_cache,
+)
+from clusterpilot.cluster.slurm import (
+    TERMINAL_STATES,
+    SlurmError,
+    cancel,
+    find_log,
+    job_status,
+    submit,
+    tail_log,
+)
+
+__all__ = [
+    # probe
+    "ClusterProbe",
+    "PartitionInfo",
+    "load_cache",
+    "probe_cluster",
+    "save_cache",
+    # slurm
+    "TERMINAL_STATES",
+    "SlurmError",
+    "cancel",
+    "find_log",
+    "job_status",
+    "submit",
+    "tail_log",
+]

clusterpilot/cluster/probe.py

@@ -0,0 +1,230 @@
+"""Cluster probe: query sinfo, module avail, and sacctmgr; cache 24h.
+
+Results are stored in ~/.cache/clusterpilot/<cluster_name>/probe.json and
+returned from cache on subsequent calls until the TTL expires or force=True.
+
+Parsed output is based on confirmed Grex (yak.hpc.umanitoba.ca) format.
+"""
+from __future__ import annotations
+
+import asyncio
+import json
+import time
+from dataclasses import asdict, dataclass, field
+from pathlib import Path
+
+from clusterpilot.ssh.connection import run_remote
+
+_CACHE_ROOT = Path.home() / ".cache" / "clusterpilot"
+_CACHE_TTL = 24 * 3600   # seconds
+
+
+# ── Data classes ──────────────────────────────────────────────────────────────
+
+@dataclass
+class PartitionInfo:
+    name: str
+    max_time: str    # e.g. "7-00:00:00" or "21-00:00:00"
+    gres: str        # e.g. "gpu:v100:4" or "" for CPU-only
+    nodes: int
+    is_default: bool
+
+
+@dataclass
+class ClusterProbe:
+    cluster_name: str
+    probed_at: float           # Unix timestamp
+    partitions: list[PartitionInfo]
+    julia_versions: list[str]  # e.g. ["julia/1.10.3", "julia/1.11.3"]
+    accounts: list[str]        # e.g. ["def-stamps"]
+    account_max_wall: dict[str, str]   # account → max walltime, "" = no limit
+    python_versions: list[str] = field(default_factory=list)  # e.g. ["python/3.11.5"]
+
+    def gpu_partitions(self) -> list[PartitionInfo]:
+        """Return partitions that have GPU GRES."""
+        return [p for p in self.partitions if p.gres.startswith("gpu:")]
+
+    def cpu_partitions(self) -> list[PartitionInfo]:
+        """Return CPU-only partitions."""
+        return [p for p in self.partitions if not p.gres]
+
+    def default_partition(self) -> PartitionInfo | None:
+        for p in self.partitions:
+            if p.is_default:
+                return p
+        return None
+
+
+# ── Public API ────────────────────────────────────────────────────────────────
+
+def load_cache(cluster_name: str) -> ClusterProbe | None:
+    """Return cached probe if it exists and is younger than 24h, else None."""
+    path = _cache_path(cluster_name)
+    if not path.exists():
+        return None
+    try:
+        data = json.loads(path.read_text())
+        if time.time() - data["probed_at"] > _CACHE_TTL:
+            return None
+        return _from_dict(data)
+    except (KeyError, ValueError):
+        return None
+
+
+def save_cache(probe: ClusterProbe) -> None:
+    """Write probe data to ~/.cache/clusterpilot/<cluster>/probe.json."""
+    path = _cache_path(probe.cluster_name)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(json.dumps(asdict(probe), indent=2))
+
+
+async def probe_cluster(
+    cluster_name: str,
+    host: str,
+    user: str,
+    *,
+    force: bool = False,
+) -> ClusterProbe:
+    """Query sinfo, module avail, and sacctmgr on host.
+
+    Returns cached data if < 24h old (unless force=True).
+    Saves fresh results to cache before returning.
+
+    Requires an active SSH ControlMaster socket (call open_connection first).
+    """
+    if not force:
+        cached = load_cache(cluster_name)
+        if cached is not None:
+            return cached
+
+    sinfo_out, julia_out, python_out, sacctmgr_out = await _fetch_all(host, user)
+
+    result = ClusterProbe(
+        cluster_name=cluster_name,
+        probed_at=time.time(),
+        partitions=_parse_sinfo(sinfo_out),
+        julia_versions=_parse_julia_modules(julia_out),
+        python_versions=_parse_python_modules(python_out),
+        accounts=_parse_accounts(sacctmgr_out),
+        account_max_wall=_parse_max_wall(sacctmgr_out),
+    )
+    save_cache(result)
+    return result
+
+
+# ── Remote fetching ───────────────────────────────────────────────────────────
+
+async def _fetch_all(host: str, user: str) -> tuple[str, str, str, str]:
+    """Run all four probe commands concurrently."""
+    return await asyncio.gather(
+        run_remote(host, user, "sinfo -o '%P %l %G %D' --noheader"),
+        run_remote(host, user, "module avail julia 2>&1"),
+        run_remote(host, user, "module avail python 2>&1"),
+        run_remote(
+            host, user,
+            f"sacctmgr show user {user} withassoc "
+            f"format=account,maxjobs,maxwall -p --noheader",
+        ),
+    )
+
+
+# ── Parsers ───────────────────────────────────────────────────────────────────
+
+def _parse_sinfo(output: str) -> list[PartitionInfo]:
+    """Parse `sinfo -o '%P %l %G %D' --noheader` output.
+
+    Example line: "stamps 21-00:00:00 gpu:v100:4(S:0-1) 3"
+    """
+    partitions = []
+    for line in output.splitlines():
+        parts = line.split()
+        if len(parts) < 4:
+            continue
+        name_raw, max_time, gres_raw, nodes_str = parts[0], parts[1], parts[2], parts[3]
+        is_default = name_raw.endswith("*")
+        name = name_raw.rstrip("*")
+        # Strip socket-affinity suffix: "gpu:v100:4(S:0-1)" → "gpu:v100:4"
+        gres = gres_raw.split("(")[0] if gres_raw != "(null)" else ""
+        try:
+            nodes = int(nodes_str)
+        except ValueError:
+            nodes = 0
+        partitions.append(PartitionInfo(
+            name=name,
+            max_time=max_time,
+            gres=gres,
+            nodes=nodes,
+            is_default=is_default,
+        ))
+    return partitions
+
+
+def _parse_julia_modules(output: str) -> list[str]:
+    """Extract julia/X.Y.Z tokens from `module avail julia 2>&1` output.
+
+    Example: "   julia/1.10.3    julia/1.11.3 (D)"
+    Tokens "(D)" are separate and filtered naturally by the startswith check.
+    """
+    versions: set[str] = set()
+    for line in output.splitlines():
+        for token in line.split():
+            if token.startswith("julia/"):
+                versions.add(token)
+    return sorted(versions)
+
+
+def _parse_python_modules(output: str) -> list[str]:
+    """Extract python/X.Y.Z tokens from `module avail python 2>&1` output."""
+    versions: set[str] = set()
+    for line in output.splitlines():
+        for token in line.split():
+            if token.startswith("python/"):
+                versions.add(token)
+    return sorted(versions)
+
+
+def _parse_accounts(output: str) -> list[str]:
+    """Extract account names from pipe-delimited sacctmgr output."""
+    accounts = []
+    for line in output.splitlines():
+        if "|" not in line:
+            continue
+        account = line.split("|")[0].strip()
+        if account and account.lower() != "account":
+            accounts.append(account)
+    return accounts
+
+
+def _parse_max_wall(output: str) -> dict[str, str]:
+    """Extract account → max_walltime mapping from sacctmgr output.
+
+    Empty string means no limit set at the account level.
+    """
+    result: dict[str, str] = {}
+    for line in output.splitlines():
+        parts = line.split("|")
+        if len(parts) < 3:
+            continue
+        account = parts[0].strip()
+        max_wall = parts[2].strip()
+        if account and account.lower() != "account":
+            result[account] = max_wall
+    return result
+
+
+# ── Cache helpers ─────────────────────────────────────────────────────────────
+
+def _cache_path(cluster_name: str) -> Path:
+    return _CACHE_ROOT / cluster_name / "probe.json"
+
+
+def _from_dict(data: dict) -> ClusterProbe:
+    return ClusterProbe(
+        cluster_name=data["cluster_name"],
+        probed_at=data["probed_at"],
+        partitions=[PartitionInfo(**p) for p in data["partitions"]],
+        julia_versions=data["julia_versions"],
+        python_versions=data.get("python_versions", []),   # backwards-compat
+        accounts=data["accounts"],
+        account_max_wall=data["account_max_wall"],
+    )

clusterpilot/cluster/slurm.py

@@ -0,0 +1,174 @@
+"""SLURM commands: submit, poll status, cancel, fetch log output.
+
+All functions require an active SSH ControlMaster socket.
+"""
+from __future__ import annotations
+
+import re
+
+from clusterpilot.ssh.connection import SSHError, run_remote
+
+_SUBMITTED_RE = re.compile(r"Submitted batch job (\d+)")
+
+# States that mean the job will never run again.
+TERMINAL_STATES = frozenset({
+    "COMPLETED", "FAILED", "CANCELLED", "TIMEOUT",
+    "OUT_OF_MEMORY", "NODE_FAIL",
+})
+
+
+class SlurmError(SSHError):
+    """Raised when a SLURM command fails unexpectedly."""
+
+
+# ── Job submission ────────────────────────────────────────────────────────────
+
+async def submit(
+    host: str,
+    user: str,
+    remote_script_path: str,
+    *,
+    working_dir: str | None = None,
+) -> str:
+    """Run sbatch on remote_script_path. Returns the numeric job ID string.
+
+    Args:
+        host:               SSH hostname.
+        user:               Remote username.
+        remote_script_path: Absolute path to the .sh script on the cluster.
+        working_dir:        If given, cd here before running sbatch.
+
+    Raises:
+        SlurmError: if sbatch output doesn't contain "Submitted batch job NNN".
+    """
+    cmd = f"sbatch {remote_script_path}"
+    if working_dir:
+        cmd = f"cd {working_dir} && {cmd}"
+    try:
+        output = await run_remote(host, user, cmd)
+    except SSHError as exc:
+        raise SlurmError(f"sbatch failed: {exc}") from exc
+
+    match = _SUBMITTED_RE.search(output)
+    if not match:
+        raise SlurmError(f"Unexpected sbatch output: {output!r}")
+    return match.group(1)
+
+
+# ── Status polling ────────────────────────────────────────────────────────────
+
+async def job_status(host: str, user: str, job_id: str) -> str | None:
+    """Return the SLURM state for job_id, or None if the job cannot be found.
+
+    Strategy:
+    1. squeue (fast, in-memory) — works while the job is queued or running.
+    2. sacct (historical records) — works after the job has left the queue.
+
+    Common return values: PENDING, RUNNING, COMPLETED, FAILED,
+    CANCELLED, TIMEOUT, OUT_OF_MEMORY.
+    """
+    # 1. squeue — job still in queue
+    try:
+        out = await run_remote(
+            host, user,
+            f"squeue -j {job_id} -h -o '%T' 2>/dev/null",
+        )
+        state = out.strip()
+        if state:
+            return state
+    except SSHError:
+        pass
+
+    # 2. sacct — job already finished; -X = summary record only (no steps)
+    try:
+        out = await run_remote(
+            host, user,
+            f"sacct -j {job_id} -n -X -o State --parsable2 2>/dev/null",
+        )
+        for line in out.strip().splitlines():
+            # sacct can append "+" for job-step aggregates; strip it.
+            # "CANCELLED by 12345" → "CANCELLED"
+            state = line.strip().split("+")[0].split()[0]
+            if state:
+                return state
+    except SSHError:
+        pass
+
+    return None
+
+
+# ── Job control ───────────────────────────────────────────────────────────────
+
+async def cancel(host: str, user: str, job_id: str) -> None:
+    """Cancel a queued or running SLURM job via scancel."""
+    try:
+        await run_remote(host, user, f"scancel {job_id}")
+    except SSHError as exc:
+        raise SlurmError(f"scancel failed for job {job_id}: {exc}") from exc
+
+
+# ── Log access ────────────────────────────────────────────────────────────────
+
+async def tail_log(
+    host: str,
+    user: str,
+    remote_log_path: str,
+    n_lines: int = 50,
+) -> str:
+    """Return the last n_lines of a remote file. Empty string if not found."""
+    try:
+        return await run_remote(
+            host, user,
+            f"tail -n {n_lines} {remote_log_path} 2>/dev/null",
+        )
+    except SSHError:
+        return ""
+
+
+async def cat_log(
+    host: str,
+    user: str,
+    remote_log_path: str,
+) -> str:
+    """Return the full contents of a remote log file. Empty string if not found."""
+    try:
+        return await run_remote(
+            host, user,
+            f"cat {remote_log_path} 2>/dev/null",
+        )
+    except SSHError:
+        return ""
+
+
+async def find_log(
+    host: str,
+    user: str,
+    job_name: str,
+    job_id: str,
+    working_dir: str,
+) -> str | None:
+    """Locate the SLURM stdout log for this job on the remote host.
+
+    Tries common naming patterns in order:
+      <working_dir>/<job_name>-<job_id>.out   (ClusterPilot default: %x-%j.out)
+      <working_dir>/slurm-<job_id>.out        (SLURM default)
+      <working_dir>/<job_id>.out
+
+    Returns the first path that exists, or None.
+    """
+    candidates = [
+        f"{working_dir}/{job_name}-{job_id}.out",
+        f"{working_dir}/slurm-{job_id}.out",
+        f"{working_dir}/{job_id}.out",
+    ]
+    for path in candidates:
+        try:
+            out = await run_remote(
+                host, user,
+                f"test -f {path} && echo exists",
+            )
+            if out.strip() == "exists":
+                return path
+        except SSHError:
+            continue
+    return None