runops 0.2.0__py3-none-any.whl

runops/slurm/query.py

@@ -0,0 +1,384 @@
+"""Slurm job state queries via squeue and sacct.
+
+Provides functions to query active and historical job states and map them to
+runops ``RunState`` values.  All subprocess calls go through an injectable
+``CommandRunner`` callable so that tests never invoke real Slurm commands.
+"""
+
+from __future__ import annotations
+
+import re
+from collections import OrderedDict
+from dataclasses import dataclass
+
+from runops.core.state import RunState
+from runops.slurm.submit import (
+    CommandResult,
+    CommandRunner,
+    SlurmNotFoundError,
+    _default_runner,
+)
+
+# ---------------------------------------------------------------------------
+# Slurm state -> runops RunState mapping
+# ---------------------------------------------------------------------------
+
+_SLURM_STATE_MAP: dict[str, RunState] = {
+    # Active / queued
+    "PENDING": RunState.SUBMITTED,
+    "CONFIGURING": RunState.RUNNING,
+    "RUNNING": RunState.RUNNING,
+    "COMPLETING": RunState.RUNNING,
+    "SUSPENDED": RunState.RUNNING,
+    "REQUEUED": RunState.SUBMITTED,
+    # Successful termination
+    "COMPLETED": RunState.COMPLETED,
+    # Failure modes
+    "FAILED": RunState.FAILED,
+    "NODE_FAIL": RunState.FAILED,
+    "OUT_OF_MEMORY": RunState.FAILED,
+    "TIMEOUT": RunState.FAILED,
+    "PREEMPTED": RunState.FAILED,
+    "BOOT_FAIL": RunState.FAILED,
+    "DEADLINE": RunState.FAILED,
+    # Cancellation
+    "CANCELLED": RunState.CANCELLED,
+}
+
+
+#: Maps Slurm failure states to human-readable failure reasons.
+_FAILURE_REASON_MAP: dict[str, str] = {
+    "TIMEOUT": "timeout",
+    "OUT_OF_MEMORY": "oom",
+    "NODE_FAIL": "node_fail",
+    "PREEMPTED": "preempted",
+    "BOOT_FAIL": "boot_fail",
+    "DEADLINE": "deadline",
+    "FAILED": "exit_error",
+}
+
+
+class SlurmQueryError(RuntimeError):
+    """Raised when a Slurm query command fails unexpectedly."""
+
+
+@dataclass(frozen=True)
+class PartitionInfo:
+    """Information about a Slurm partition from ``sinfo``.
+
+    Attributes:
+        name: Partition name.
+        avail: Availability status (e.g. ``"up"``).
+        timelimit: Raw time limit string from sinfo (e.g. ``"5-00:00:00"``).
+        timelimit_hours: Time limit in hours (for easy comparison).
+        nodes_total: Total node count in the partition.
+    """
+
+    name: str
+    avail: str
+    timelimit: str
+    timelimit_hours: float
+    nodes_total: int = 0
+
+
+@dataclass(frozen=True)
+class JobStatus:
+    """Result of a Slurm job status query.
+
+    Attributes:
+        run_state: Mapped runops RunState.
+        slurm_state: Raw Slurm state string.
+        failure_reason: Reason for failure (empty if not failed).
+        exit_code: Slurm exit code string (if available).
+    """
+
+    run_state: RunState
+    slurm_state: str
+    failure_reason: str = ""
+    exit_code: str = ""
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def map_slurm_state(slurm_state: str) -> RunState:
+    """Map a raw Slurm job state string to a runops ``RunState``.
+
+    Slurm may append qualifiers like ``CANCELLED by 1000`` -- only the first
+    word is used for the lookup.
+
+    Args:
+        slurm_state: Raw state string from squeue or sacct (e.g.
+            ``"RUNNING"``, ``"CANCELLED by 1000"``).
+
+    Returns:
+        The corresponding ``RunState``.
+
+    Raises:
+        SlurmQueryError: If the state is not recognised.
+    """
+    # Take first token to handle "CANCELLED by UID" variants
+    key = slurm_state.strip().split()[0].rstrip("+")
+    try:
+        return _SLURM_STATE_MAP[key]
+    except KeyError:
+        raise SlurmQueryError(f"Unknown Slurm job state: {slurm_state!r}") from None
+
+
+# ---------------------------------------------------------------------------
+# sinfo (partition queries)
+# ---------------------------------------------------------------------------
+
+_TIMELIMIT_RE = re.compile(r"(?:(\d+)-)?(\d+):(\d+):(\d+)")
+
+
+def _parse_timelimit(timelimit: str) -> float:
+    """Parse a Slurm time limit string to hours.
+
+    Supports formats like ``5-00:00:00``, ``120:00:00``, ``infinite``.
+
+    Args:
+        timelimit: Raw time limit string from sinfo.
+
+    Returns:
+        Time limit in hours, or ``float('inf')`` for unlimited.
+    """
+    if timelimit.lower() in ("infinite", "n/a"):
+        return float("inf")
+    m = _TIMELIMIT_RE.match(timelimit.strip())
+    if not m:
+        return float("inf")
+    day_str, hour_str, min_str, sec_str = m.groups()
+    day = int(day_str) if day_str else 0
+    hour = int(hour_str)
+    minutes = int(min_str)
+    sec = int(sec_str)
+    return 24.0 * day + hour + minutes / 60.0 + sec / 3600.0
+
+
+def sinfo_partitions(
+    *,
+    runner: CommandRunner | None = None,
+) -> OrderedDict[str, PartitionInfo]:
+    """Query ``sinfo`` for available partitions and their limits.
+
+    Args:
+        runner: Optional command runner for testing.
+
+    Returns:
+        Ordered dict mapping partition name to :class:`PartitionInfo`.
+
+    Raises:
+        SlurmNotFoundError: If ``sinfo`` is not on PATH.
+        SlurmQueryError: If sinfo returns an error.
+    """
+    run = runner or _default_runner
+    cmd = [
+        "sinfo",
+        "--noheader",
+        "--format=%P|%a|%l|%D",
+    ]
+
+    try:
+        result: CommandResult = run(cmd)
+    except SlurmNotFoundError:
+        raise
+
+    if result.returncode != 0:
+        raise SlurmQueryError(
+            f"sinfo failed (exit {result.returncode}):\n{result.stderr.strip()}"
+        )
+
+    partitions: OrderedDict[str, PartitionInfo] = OrderedDict()
+    for line in result.stdout.strip().splitlines():
+        parts = line.strip().split("|")
+        if len(parts) < 4:
+            continue
+        name = parts[0].rstrip("*")  # default partition has trailing '*'
+        avail = parts[1]
+        timelimit = parts[2]
+        try:
+            nodes_total = int(parts[3])
+        except ValueError:
+            nodes_total = 0
+
+        partitions[name] = PartitionInfo(
+            name=name,
+            avail=avail,
+            timelimit=timelimit,
+            timelimit_hours=_parse_timelimit(timelimit),
+            nodes_total=nodes_total,
+        )
+
+    return partitions
+
+
+# ---------------------------------------------------------------------------
+# squeue
+# ---------------------------------------------------------------------------
+
+
+def squeue_status(
+    job_id: str,
+    *,
+    runner: CommandRunner | None = None,
+) -> str | None:
+    """Query ``squeue`` for an active job's state.
+
+    Args:
+        job_id: Slurm job ID.
+        runner: Optional command runner for testing.
+
+    Returns:
+        The raw Slurm state string (e.g. ``"RUNNING"``) if the job is still
+        in the queue, or ``None`` if it has left the queue.
+
+    Raises:
+        SlurmNotFoundError: If ``squeue`` is not on PATH.
+        SlurmQueryError: If squeue returns a non-zero exit code.
+    """
+    run = runner or _default_runner
+    cmd = [
+        "squeue",
+        "--job",
+        job_id,
+        "--noheader",
+        "--format=%T",
+    ]
+
+    try:
+        result: CommandResult = run(cmd)
+    except SlurmNotFoundError:
+        raise
+
+    if result.returncode != 0:
+        # squeue returns non-zero when the job is not found on some clusters
+        if "Invalid job id" in result.stderr:
+            return None
+        raise SlurmQueryError(
+            f"squeue failed (exit {result.returncode}):\n{result.stderr.strip()}"
+        )
+
+    state = result.stdout.strip()
+    if not state:
+        return None
+    return state
+
+
+# ---------------------------------------------------------------------------
+# sacct
+# ---------------------------------------------------------------------------
+
+
+def sacct_status(
+    job_id: str,
+    *,
+    runner: CommandRunner | None = None,
+) -> dict[str, str] | None:
+    """Query ``sacct`` for a historical job's state and exit code.
+
+    Uses ``--parsable2 --noheader`` with explicit format fields for reliable
+    parsing.
+
+    Args:
+        job_id: Slurm job ID.
+        runner: Optional command runner for testing.
+
+    Returns:
+        A dictionary with keys ``"state"`` and ``"exit_code"`` if the job is
+        found, or ``None`` if sacct has no record.
+
+    Raises:
+        SlurmNotFoundError: If ``sacct`` is not on PATH.
+        SlurmQueryError: If sacct returns a non-zero exit code.
+    """
+    run = runner or _default_runner
+    cmd = [
+        "sacct",
+        "--jobs",
+        job_id,
+        "--noheader",
+        "--parsable2",
+        "--format=JobID,State,ExitCode",
+    ]
+
+    try:
+        result: CommandResult = run(cmd)
+    except SlurmNotFoundError:
+        raise
+
+    if result.returncode != 0:
+        raise SlurmQueryError(
+            f"sacct failed (exit {result.returncode}):\n{result.stderr.strip()}"
+        )
+
+    # sacct may return multiple lines (one per step).  We want the "batch"
+    # line or the main job line (the one whose JobID matches exactly).
+    for line in result.stdout.strip().splitlines():
+        parts = line.split("|")
+        if len(parts) < 3:
+            continue
+        sacct_job_id, state, exit_code = parts[0], parts[1], parts[2]
+        # Match the main job entry (not sub-steps like "12345.batch")
+        if sacct_job_id.strip() == job_id:
+            return {"state": state.strip(), "exit_code": exit_code.strip()}
+
+    return None
+
+
+# ---------------------------------------------------------------------------
+# Combined query
+# ---------------------------------------------------------------------------
+
+
+def query_job_status(
+    job_id: str,
+    *,
+    runner: CommandRunner | None = None,
+) -> JobStatus:
+    """Determine the current runops state of a Slurm job.
+
+    Strategy: try ``squeue`` first (cheap, covers active jobs).  If the job
+    is no longer in the queue, fall back to ``sacct`` (covers completed /
+    historical jobs).
+
+    Args:
+        job_id: Slurm job ID.
+        runner: Optional command runner for testing.
+
+    Returns:
+        A :class:`JobStatus` with the mapped state, raw Slurm state,
+        failure reason, and exit code.
+
+    Raises:
+        SlurmNotFoundError: If Slurm commands are not on PATH.
+        SlurmQueryError: If neither squeue nor sacct can find the job, or
+            if the returned state is unrecognised.
+    """
+    # 1. Try squeue (active jobs)
+    sq_state = squeue_status(job_id, runner=runner)
+    if sq_state is not None:
+        raw = sq_state.strip().split()[0].rstrip("+")
+        return JobStatus(
+            run_state=map_slurm_state(sq_state),
+            slurm_state=raw,
+        )
+
+    # 2. Fall back to sacct (historical jobs)
+    sa_info = sacct_status(job_id, runner=runner)
+    if sa_info is not None:
+        raw = sa_info["state"].strip().split()[0].rstrip("+")
+        run_state = map_slurm_state(sa_info["state"])
+        return JobStatus(
+            run_state=run_state,
+            slurm_state=raw,
+            failure_reason=_FAILURE_REASON_MAP.get(raw, ""),
+            exit_code=sa_info.get("exit_code", ""),
+        )
+
+    raise SlurmQueryError(
+        f"Job {job_id} not found in squeue or sacct. "
+        "It may have been purged from the Slurm database."
+    )

runops/slurm/submit.py

@@ -0,0 +1,203 @@
+"""Slurm sbatch submission.
+
+Provides functions to submit job scripts via sbatch and parse the resulting
+job ID.  All subprocess calls go through a single ``run_command`` callable
+so that tests can inject a mock without touching the real Slurm installation.
+"""
+
+from __future__ import annotations
+
+import re
+import subprocess
+from collections.abc import Callable
+from pathlib import Path
+from typing import NamedTuple
+
+# ---------------------------------------------------------------------------
+# Thin subprocess abstraction
+# ---------------------------------------------------------------------------
+
+_SBATCH_JOB_RE = re.compile(r"Submitted batch job (\d+)")
+
+
+class CommandResult(NamedTuple):
+    """Result of a shell command execution."""
+
+    returncode: int
+    stdout: str
+    stderr: str
+
+
+#: Type alias for the injectable command runner.
+CommandRunner = Callable[[list[str]], CommandResult]
+
+
+def _default_runner(cmd: list[str]) -> CommandResult:
+    """Run a command via ``subprocess.run``.
+
+    This is the production implementation used when no mock is injected.
+
+    Args:
+        cmd: Command and arguments to execute.
+
+    Returns:
+        A ``CommandResult`` with return code, stdout, and stderr.
+
+    Raises:
+        SlurmNotFoundError: If the command executable is not found on PATH.
+    """
+    try:
+        proc = subprocess.run(
+            cmd,
+            capture_output=True,
+            text=True,
+            check=False,
+            timeout=60,
+        )
+    except FileNotFoundError as exc:
+        raise SlurmNotFoundError(
+            f"Command not found: {cmd[0]!r}. Is Slurm installed and on PATH?"
+        ) from exc
+    except subprocess.TimeoutExpired as exc:
+        raise SlurmSubmitError(f"sbatch timed out after 60 seconds: {exc}") from exc
+    return CommandResult(
+        returncode=proc.returncode,
+        stdout=proc.stdout,
+        stderr=proc.stderr,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Exceptions
+# ---------------------------------------------------------------------------
+
+
+class SlurmNotFoundError(RuntimeError):
+    """Raised when the Slurm command is not found on PATH."""
+
+
+class SlurmSubmitError(RuntimeError):
+    """Raised when sbatch fails or returns unexpected output."""
+
+
+class SlurmCancelError(RuntimeError):
+    """Raised when scancel fails."""
+
+
+# ---------------------------------------------------------------------------
+# Public API
+# ---------------------------------------------------------------------------
+
+
+def parse_job_id(sbatch_stdout: str) -> str:
+    """Extract the job ID from sbatch standard output.
+
+    Expected format: ``Submitted batch job 12345``
+
+    Args:
+        sbatch_stdout: The captured stdout from an sbatch invocation.
+
+    Returns:
+        The numeric job ID as a string.
+
+    Raises:
+        SlurmSubmitError: If the output does not match the expected pattern.
+    """
+    match = _SBATCH_JOB_RE.search(sbatch_stdout)
+    if match is None:
+        raise SlurmSubmitError(
+            f"Could not parse job ID from sbatch output: {sbatch_stdout!r}"
+        )
+    return match.group(1)
+
+
+def sbatch_submit(
+    job_script: Path,
+    working_dir: Path,
+    *,
+    extra_args: list[str] | None = None,
+    afterok: str | None = None,
+    runner: CommandRunner | None = None,
+) -> str:
+    """Submit a job script via ``sbatch``.
+
+    Args:
+        job_script: Path to the job script file (e.g. ``submit/job.sh``).
+        working_dir: Working directory for the sbatch process (typically
+            the run directory's ``work/`` subdirectory).
+        extra_args: Additional sbatch arguments (e.g.
+            ``["--partition=gr10451a"]``).  These override script directives.
+        afterok: If set, add ``--dependency=afterok:<job_id>`` so this job
+            starts only after the specified job completes successfully.
+        runner: Optional callable that executes a command list and returns
+            a ``CommandResult``.  Defaults to the real subprocess runner.
+            Inject a mock here for testing.
+
+    Returns:
+        The Slurm job ID as a string.
+
+    Raises:
+        FileNotFoundError: If *job_script* does not exist.
+        SlurmNotFoundError: If ``sbatch`` is not on PATH.
+        SlurmSubmitError: If sbatch returns a non-zero exit code or its
+            output cannot be parsed.
+    """
+    if not job_script.exists():
+        raise FileNotFoundError(f"Job script not found: {job_script}")
+
+    run = runner or _default_runner
+    cmd = ["sbatch", f"--chdir={working_dir}"]
+    if afterok:
+        cmd.append(f"--dependency=afterok:{afterok}")
+    if extra_args:
+        cmd.extend(extra_args)
+    cmd.append(str(job_script))
+    result = run(cmd)
+
+    if result.returncode != 0:
+        raise SlurmSubmitError(
+            f"sbatch failed (exit {result.returncode}):\n{result.stderr.strip()}"
+        )
+
+    return parse_job_id(result.stdout)
+
+
+def scancel_job(
+    job_id: str,
+    *,
+    runner: CommandRunner | None = None,
+) -> None:
+    """Cancel a Slurm job via ``scancel``.
+
+    Args:
+        job_id: The Slurm job ID to cancel.
+        runner: Optional callable that executes a command list and returns
+            a ``CommandResult``.  Defaults to the real subprocess runner.
+
+    Raises:
+        SlurmNotFoundError: If ``scancel`` is not on PATH.
+        SlurmCancelError: If scancel returns a non-zero exit code.
+    """
+    run = runner or _default_runner
+    result = run(["scancel", job_id])
+    if result.returncode != 0:
+        raise SlurmCancelError(
+            f"scancel failed (exit {result.returncode}):\n{result.stderr.strip()}"
+        )
+
+
+# Keep the old name as an alias so the existing test import still works,
+# but point it at the new implementation.
+def sbatch(job_script: Path) -> str:
+    """Submit a job script via sbatch (legacy wrapper).
+
+    Deprecated: prefer ``sbatch_submit`` which accepts a working directory
+    and an injectable runner.
+
+    Args:
+        job_script: Path to the job.sh file.
+
+    Returns:
+        The Slurm job_id as a string.
+    """
+    return sbatch_submit(job_script, job_script.parent)

runops/templates/__init__.py

@@ -0,0 +1,29 @@
+"""Template loading utilities for runops."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import jinja2
+
+_TEMPLATES_DIR = Path(__file__).resolve().parent
+
+
+def get_jinja_env() -> jinja2.Environment:
+    """Return a Jinja2 environment that loads from runops/templates/."""
+    return jinja2.Environment(
+        loader=jinja2.FileSystemLoader(str(_TEMPLATES_DIR)),
+        keep_trailing_newline=True,
+        undefined=jinja2.StrictUndefined,
+    )
+
+
+def load_static(relative_path: str) -> str:
+    """Load a static (non-Jinja2) template file as text."""
+    return (_TEMPLATES_DIR / relative_path).read_text(encoding="utf-8")
+
+
+def render(template_path: str, **kwargs: object) -> str:
+    """Render a Jinja2 template with the given variables."""
+    env = get_jinja_env()
+    return env.get_template(template_path).render(**kwargs)

runops/templates/adapters/beach/agent_guide.md

@@ -0,0 +1,50 @@
+### BEACH (BEM + Accumulated CHarge Simulator)
+
+#### 概要
+BEACH は境界要素法 (BEM) ベースの表面帯電シミュレーション。
+MPI + OpenMP ハイブリッド並列で実行し、宇宙機表面の帯電現象を計算する。
+
+#### 入力ファイル
+- **`input/beach.toml`** — メイン設定ファイル (TOML 形式)
+  - `[sim]`: `dt`, `max_step`, `batch_count`, `field_solver`
+  - `[mesh]`: `obj_path` (OBJ メッシュファイルパス)
+  - `[environment]`: プラズマ環境パラメータ (密度, 温度, etc.)
+  - `[output]`: `dir` (出力ディレクトリ)
+
+#### 出力ファイル (`work/latest/` 以下)
+- `summary.txt` — 計算結果サマリー (完了時に生成)
+- `charges.csv` — 表面電荷分布
+- `charge_history.csv` — 電荷時間履歴
+- `potential_history.csv` — 電位時間履歴
+- `mesh_triangles.csv`, `mesh_sources.csv` — メッシュ情報
+- `performance_profile.csv` — 性能プロファイル
+
+#### 完了判定
+- `work/latest/summary.txt` が存在 → completed
+- stderr に "error", "fatal", "killed" → failed
+- `charges.csv` のみ存在 → running (途中)
+
+#### パラメータサーベイでよく変えるパラメータ
+- `sim.dt`, `sim.max_step`, `sim.batch_count`
+- `environment.electron_density`, `environment.electron_temperature`
+- `environment.ion_density`, `environment.ion_temperature`
+- `mesh.obj_path` (異なるジオメトリの比較)
+
+#### ドキュメント・参考
+- BEACH ソースリポジトリの README / docs/
+- OBJ メッシュファイルは Blender 等で作成
+- パラメータの dot 記法例: `sim.dt=1.0e-6`, `environment.electron_density=1.0e12`
+
+#### 実行コマンド
+```
+srun beach input/beach.toml
+```
+
+`beach.toml` の `output.dir` は `work/latest` に自動設定される。
+
+#### 環境変数 (OpenMP)
+```
+OMP_NUM_THREADS=${SLURM_CPUS_PER_TASK:-1}
+OMP_PROC_BIND=spread
+OMP_PLACES=cores
+```

runops/templates/adapters/beach/beach.toml

@@ -0,0 +1,19 @@
+# BEACH configuration
+# See BEACH documentation for full parameter reference
+
+[sim]
+dt = 1.0e-6
+max_step = 1000
+batch_count = 100
+
+[mesh]
+# obj_path = "mesh.obj"
+
+[environment]
+electron_density = 1.0e12
+electron_temperature = 1.0
+ion_density = 1.0e12
+ion_temperature = 1.0
+
+[output]
+# dir is set automatically by runops

runops/templates/adapters/beach/case.toml

@@ -0,0 +1,16 @@
+[case]
+name = ""
+simulator = "beach"
+launcher = "default"
+description = ""
+
+[params]
+# "sim.dt" = 1.0e-6
+# "sim.max_step" = 1000
+# "environment.electron_density" = 1.0e12
+
+[job]
+partition = ""
+nodes = 1
+ntasks = 1
+walltime = "01:00:00"