simrun 0.1.0__py3-none-any.whl

simrun/mcp.py

@@ -0,0 +1,161 @@
+"""The MCP server of simrun, a local stdio process of an agent.
+
+The tools go through the `Client`; they return the JSON of the API and `store`, the
+absolute local prefix of the paths in it, which are relative to the store, so an
+agent opens a result with `store / path`. A refusal is returned as
+`{"error": detail, "status": code}` instead of raised, so an agent can read it.
+"""
+
+import logging
+import subprocess
+from collections.abc import Callable
+from pathlib import Path
+from typing import Any
+
+import httpx
+from fastmcp import FastMCP
+from pydantic import BaseModel
+
+from simrun.client import Client, ClientError
+from simrun.service.logs import DEFAULT_TAIL_LINES
+from simrun.store import RunState
+from simrun.workflow import SplitMode
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_WAIT_S = 600.0
+# status of an error which did not come from the service
+CLIENT_SIDE = 400
+UNREACHABLE = 503
+
+
+def _plain(value: Any) -> Any:
+    if isinstance(value, BaseModel):
+        return value.model_dump(mode="json")
+    if isinstance(value, dict):
+        return {k: _plain(v) for k, v in value.items()}
+    if isinstance(value, list):
+        return [_plain(v) for v in value]
+    return value
+
+
+def build_server(client: Client) -> FastMCP:
+    """The MCP server whose tools call the service through `client`."""
+    server = FastMCP("simrun")
+    store = str(client.store.resolve())
+
+    def answer(call: Callable[[], dict[str, Any]]) -> dict[str, Any]:
+        try:
+            return {**_plain(call()), "store": store}
+        except ClientError as error:
+            return {"error": error.detail, "status": error.status}
+        except httpx.TransportError as error:
+            return {
+                "error": f"the service is not reachable: {error}",
+                "status": UNREACHABLE,
+            }
+        except (ValueError, RuntimeError, OSError, subprocess.SubprocessError) as error:
+            return {"error": str(error), "status": CLIENT_SIDE}
+
+    @server.tool
+    def submit_workflow(
+        workflow_path: str,
+        sources: dict[str, str] | None = None,
+        group: str | None = None,
+        single_server: bool = False,
+    ) -> dict[str, Any]:
+        """Snapshot the working trees and submit a workflow.
+
+        `sources` maps a source name to the path of its git working tree, the model
+        first. Returns the summary of the run at once, in the state `submitting`
+        while its environment is resolved and built; follow it with `wait_for_run`.
+        """
+        paths = {name: Path(p).expanduser() for name, p in (sources or {}).items()}
+        split = SplitMode.SINGLE_SERVER if single_server else None
+        return answer(
+            lambda: {
+                "run": client.submit(
+                    Path(workflow_path).expanduser(), paths, group, split
+                )
+            }
+        )
+
+    @server.tool
+    def wait_for_run(run_id: str, timeout_s: float = DEFAULT_WAIT_S) -> dict[str, Any]:
+        """Wait for a run to end or for `timeout_s` seconds; returns its summary."""
+        return answer(lambda: {"run": client.wait(run_id, timeout_s)})
+
+    @server.tool
+    def run_status(run_id: str) -> dict[str, Any]:
+        """The record of a run: state, tasks, plans and the provenance of its sources."""
+        return answer(lambda: {"run": client.status(run_id)})
+
+    @server.tool
+    def run_progress(run_id: str) -> dict[str, Any]:
+        """The progress of every task of a run: ledger counts, first and last progress lines."""
+        return answer(lambda: {"progress": client.progress(run_id)})
+
+    @server.tool
+    def list_runs(
+        group: str | None = None,
+        agent: str | None = None,
+        state: str | None = None,
+        owner: str | None = None,
+    ) -> dict[str, Any]:
+        """The runs, filtered by group, agent, state and owner.
+
+        The runs of an owner include the ones of their agents.
+        """
+        return answer(
+            lambda: {
+                "runs": client.list_runs(
+                    group, agent, RunState(state) if state else None, owner
+                )
+            }
+        )
+
+    @server.tool
+    def compare_runs(run_ids: list[str]) -> dict[str, Any]:
+        """Provenance and task summaries of runs side by side."""
+        return answer(lambda: {"table": client.compare(run_ids)})
+
+    @server.tool
+    def cancel_run(run_id: str) -> dict[str, Any]:
+        """Cancel a run."""
+        return answer(lambda: {"run": client.cancel(run_id)})
+
+    @server.tool
+    def run_logs(
+        run_id: str, task_id: str, tail_lines: int = DEFAULT_TAIL_LINES
+    ) -> dict[str, Any]:
+        """The logs of the pods of a task, the last `tail_lines` of each (0: all)."""
+        return answer(lambda: {"logs": client.logs(run_id, task_id, tail_lines)})
+
+    @server.tool
+    def run_outputs(run_id: str, task_id: str) -> dict[str, Any]:
+        """The output files of a task, paths relative to `store`."""
+        return answer(lambda: {"files": client.outputs(run_id, task_id)})
+
+    @server.tool
+    def estimate_workflow(
+        workflow_path: str,
+        sources: dict[str, str] | None = None,
+        single_server: bool = False,
+    ) -> dict[str, Any]:
+        """The plan and the expected duration of every task of a submission."""
+        paths = {name: Path(p).expanduser() for name, p in (sources or {}).items()}
+        split = SplitMode.SINGLE_SERVER if single_server else None
+        return answer(
+            lambda: {
+                "tasks": client.estimate(
+                    Path(workflow_path).expanduser(), paths, None, split
+                )
+            }
+        )
+
+    return server
+
+
+def main() -> None:
+    """Run the server on stdio with the client of the environment."""
+    build_server(Client.from_env()).run()

simrun/planner.py

@@ -0,0 +1,303 @@
+"""The planner: the shape of a worker pod and the number of pieces of a task."""
+
+import logging
+import math
+from collections import Counter
+from dataclasses import dataclass
+
+from simrun.plugins import TaskType
+from simrun.plugins_builtin import COMMAND, CommandSettings
+from simrun.resources import GIB, RETRY_FACTOR, ResolvedResources
+from simrun.worker import processes
+from simrun.workflow import Priority, SplitMode, TaskSpec
+
+logger = logging.getLogger(__name__)
+
+MAX_PIECES = 8
+PRIORITY_VALUES = {
+    "simrun-owner-base": 1000,
+    "simrun-owner-extra": 900,
+    "simrun-standard-base": 500,
+    "simrun-standard-extra": 400,
+    "simrun-agent-base": 200,
+    "simrun-agent-extra": 100,
+}
+
+
+class PlanError(ValueError):
+    """A task which cannot be planned on the nodes."""
+
+
+@dataclass(frozen=True, init=False)
+class Node:
+    """A node of the cluster.
+
+    Attributes:
+        name: Name of the node.
+        cores: The whole cores a worker pod of a run can have on the node, what the
+            planner plans with.
+        memory_gb: The GiB of memory a worker pod of a run can have on the node,
+            what the planner plans with.
+        gpus: Number of GPUs.
+        reserved_by: Name of the reservation holding the node, None if free.
+        allocatable_cores: The cores the scheduler gives the pods of the runs
+            (sidecars included), as it counts their requests.
+        allocatable_memory_gb: The GiB of memory the scheduler gives the pods of
+            the runs, as it counts their requests.
+    """
+
+    name: str
+    cores: int
+    memory_gb: float
+    gpus: int
+    reserved_by: str | None
+    allocatable_cores: float
+    allocatable_memory_gb: float
+
+    def __init__(
+        self,
+        name: str,
+        cores: int,
+        memory_gb: float,
+        gpus: int = 0,
+        reserved_by: str | None = None,
+        allocatable_cores: float | None = None,
+        allocatable_memory_gb: float | None = None,
+    ) -> None:
+        """Create the node; the allocatable cores and memory default to the others."""
+        values = {
+            "name": name,
+            "cores": cores,
+            "memory_gb": memory_gb,
+            "gpus": gpus,
+            "reserved_by": reserved_by,
+            "allocatable_cores": float(
+                cores if allocatable_cores is None else allocatable_cores
+            ),
+            "allocatable_memory_gb": float(
+                memory_gb if allocatable_memory_gb is None else allocatable_memory_gb
+            ),
+        }
+        for key, value in values.items():
+            # a frozen dataclass sets its fields past its own __setattr__
+            object.__setattr__(self, key, value)
+
+
+@dataclass(frozen=True)
+class PodShape:
+    """The resources of a worker pod.
+
+    Attributes:
+        cores: Number of cores.
+        memory_gb: Memory in GiB, rounded up.
+        processes: Number of worker processes.
+        per_worker_gb: Memory per worker process in GiB.
+        retry_memory_gb: Memory of a retry in GiB.
+        retry_per_worker_gb: Memory per worker process of a retry in GiB.
+        scratch_gb: Local scratch in GiB.
+        gpus: Number of GPUs.
+    """
+
+    cores: int
+    memory_gb: int
+    processes: int
+    per_worker_gb: float
+    retry_memory_gb: int
+    retry_per_worker_gb: float
+    scratch_gb: float
+    gpus: int
+
+
+@dataclass(frozen=True)
+class TaskPlan:
+    """The plan of a task.
+
+    Attributes:
+        task_id: Id of the task.
+        pods: Number of worker pods.
+        shape: Shape of a pod.
+        base_priority: Priority class of the first pods.
+        extra_priority: Priority class of the further pods.
+        units: Number of units, None for an unsplit task.
+    """
+
+    task_id: str
+    pods: int
+    shape: PodShape
+    base_priority: str
+    extra_priority: str
+    units: int | None
+
+
+def _usable(nodes: list[Node], gpus: int) -> list[Node]:
+    return [n for n in nodes if n.reserved_by is None and n.gpus >= gpus]
+
+
+def _common_node(nodes: list[Node], gpus: int) -> Node:
+    """Returns the most common node by cores and memory, the larger on ties."""
+    usable = _usable(nodes, gpus)
+    if not usable:
+        raise PlanError(f"no node which is not reserved has {gpus} GPUs")
+    counts = Counter((n.cores, n.memory_gb) for n in usable)
+    cores, memory = max(counts, key=lambda k: (counts[k], k[1], k[0]))
+    return next(n for n in usable if (n.cores, n.memory_gb) == (cores, memory))
+
+
+def _largest(nodes: list[Node], gpus: int) -> Node:
+    usable = _usable(nodes, gpus)
+    if not usable:
+        raise PlanError(f"no node which is not reserved has {gpus} GPUs")
+    return max(usable, key=lambda n: (n.memory_gb, n.cores))
+
+
+def _fit(node: Node, base: float, per_worker: float) -> int | None:
+    """Returns the number of processes on the node, None if one does not fit."""
+    try:
+        return processes(node.cores, int(node.memory_gb * GIB), base, per_worker)
+    except ValueError:
+        return None
+
+
+def shape(
+    resources: ResolvedResources, nodes: list[Node], max_processes: int | None = None
+) -> PodShape:
+    """Sizes a pod for the most common node which is not reserved.
+
+    When one worker does not fit that node, the largest usable node is taken.
+
+    Args:
+        resources: The resources of the task.
+        nodes: The nodes of the cluster.
+        max_processes: At most this many processes, as many as fit by default.
+
+    Returns:
+        The shape of a pod.
+
+    Raises:
+        PlanError: One worker fits on no node.
+    """
+    base = resources.memory_base_gb
+    per_worker = resources.memory_per_worker_gb
+    node = _common_node(nodes, resources.gpus)
+    n = _fit(node, base, per_worker)
+    if n is None:
+        node = _largest(nodes, resources.gpus)
+        n = _fit(node, base, per_worker)
+    if n is None:
+        raise PlanError(
+            f"one worker needs {per_worker} GiB besides {base} GiB base, more than "
+            f"the largest node {node.name} with {node.memory_gb} GiB"
+        )
+    if max_processes is not None:
+        n = min(n, max_processes)
+    retry_per_worker = per_worker * RETRY_FACTOR
+    retry_node = node
+    retry_n = _fit(retry_node, base, retry_per_worker)
+    if retry_n is None:
+        retry_node = _largest(nodes, resources.gpus)
+        retry_n = _fit(retry_node, base, retry_per_worker)
+    if retry_n is not None:
+        retry_n = min(retry_n, n)
+        retry_memory = math.ceil(base + retry_per_worker * retry_n)
+    else:
+        logger.warning(
+            "a retry with %s GiB per worker fits on no node, "
+            "it takes the whole memory of %s with one process",
+            retry_per_worker,
+            retry_node.name,
+        )
+        retry_memory = int(retry_node.memory_gb)
+        retry_per_worker = max(0.0, retry_memory - base)
+    return PodShape(
+        cores=n,
+        memory_gb=math.ceil(base + per_worker * n),
+        processes=n,
+        per_worker_gb=per_worker,
+        retry_memory_gb=retry_memory,
+        retry_per_worker_gb=retry_per_worker,
+        scratch_gb=resources.scratch_gb,
+        gpus=resources.gpus,
+    )
+
+
+def _command_shape(settings: CommandSettings, nodes: list[Node]) -> PodShape:
+    memory = math.ceil(settings.memory_gb)
+    fitting = [
+        n
+        for n in _usable(nodes, settings.gpus)
+        if n.cores >= settings.cores and n.memory_gb >= memory
+    ]
+    if not fitting:
+        largest = _largest(nodes, settings.gpus)
+        raise PlanError(
+            f"the command needs {settings.cores} cores, {memory} GiB and "
+            f"{settings.gpus} GPUs, more than the largest node {largest.name} "
+            f"with {largest.cores} cores and {largest.memory_gb} GiB"
+        )
+    node_memory = max(n.memory_gb for n in fitting)
+    retry = max(
+        memory, min(math.ceil(settings.memory_gb * RETRY_FACTOR), int(node_memory))
+    )
+    return PodShape(
+        cores=settings.cores,
+        memory_gb=memory,
+        processes=1,
+        per_worker_gb=settings.memory_gb,
+        retry_memory_gb=retry,
+        # processes is 1, so the memory per worker of a retry is its memory
+        retry_per_worker_gb=float(retry),
+        scratch_gb=settings.scratch_gb,
+        gpus=settings.gpus,
+    )
+
+
+def _slots(pod: PodShape, nodes: list[Node]) -> int:
+    total = 0
+    for n in _usable(nodes, pod.gpus):
+        fit = min(int(n.cores // pod.cores), int(n.memory_gb // pod.memory_gb))
+        if pod.gpus:
+            fit = min(fit, n.gpus // pod.gpus)
+        total += max(fit, 0)
+    return total
+
+
+def plan_task(
+    task: TaskSpec,
+    task_type: TaskType,
+    resources: ResolvedResources,
+    nodes: list[Node],
+    mode: SplitMode,
+    priority: Priority,
+    units: int | None,
+) -> TaskPlan:
+    """Plans the pods of a task.
+
+    Args:
+        task: The task.
+        task_type: Its type.
+        resources: The resolved resources.
+        nodes: The nodes of the cluster.
+        mode: The split mode.
+        priority: The priority of the run.
+        units: The number of units, None for an unsplit task.
+
+    Returns:
+        The plan; extra pods are not limited by idle capacity, Kueue queues them.
+
+    Raises:
+        PlanError: The task fits on no node.
+    """
+    base = f"simrun-{priority.value}-base"
+    extra = f"simrun-{priority.value}-extra"
+    if not task_type.splittable:
+        if task_type.name == COMMAND.name:
+            pod = _command_shape(CommandSettings.model_validate(task.settings), nodes)
+        else:
+            # a single command is one process, unless it uses the cores of its pod
+            pod = shape(resources, nodes, None if task_type.uses_cores else 1)
+        return TaskPlan(task.id, 1, pod, base, extra, None)
+    pod = shape(resources, nodes)
+    pods = 1
+    if mode == SplitMode.FASTEST and units is not None:
+        pods = min(MAX_PIECES, units, _slots(pod, nodes))
+    return TaskPlan(task.id, max(1, pods), pod, base, extra, units)

simrun/plugins.py

@@ -0,0 +1,149 @@
+"""The plugin interface: task types, their resource hints and their discovery."""
+
+from collections.abc import Callable
+from dataclasses import dataclass
+from importlib import import_module
+from importlib.metadata import entry_points
+from pathlib import Path
+from typing import Any
+
+from pydantic import BaseModel
+
+ENTRY_POINT_GROUP = "simrun.task_types"
+
+
+@dataclass(frozen=True)
+class ResourceHint:
+    """Default resources of a task type.
+
+    Attributes:
+        memory_base_gb: Base memory of a task in GiB.
+        memory_per_worker_gb: Memory per worker in GiB.
+        scratch_gb: Local scratch in GiB.
+        gpus: Number of GPUs.
+    """
+
+    memory_base_gb: float = 2.0
+    memory_per_worker_gb: float = 6.0
+    scratch_gb: float = 0.0
+    gpus: int = 0
+
+
+@dataclass(frozen=True)
+class TaskPaths:
+    """The locations a task works in.
+
+    Attributes:
+        store: Root of the run store.
+        run_dir: Directory of the run.
+        task_dir: Directory of the task.
+        out: Output directory of the task, `task_dir / "out"`.
+        inputs: Input port to the directory of the output it reads.
+        scratch: Local scratch directory.
+    """
+
+    store: Path
+    run_dir: Path
+    task_dir: Path
+    out: Path
+    inputs: dict[str, Path]
+    scratch: Path
+
+
+@dataclass(frozen=True)
+class TaskType:
+    """A type of task, registered as an entry point of `ENTRY_POINT_GROUP`.
+
+    The callables are `"module:function"` strings, so a worker process resolves
+    them by import, see `resolve`. Signatures: `count_units(settings) -> int`,
+    `setup(settings, paths) -> None` (once per process),
+    `compute_unit(settings, index, paths) -> None`, `merge(settings, paths) -> None`
+    and `plan(settings) -> settings`, which the client calls before a submission
+    when the type is installed locally (`simrun.client.plan_tasks`).
+
+    A unit may be computed more than once: by a worker which is evicted before it
+    completes the unit, or by a restarted incarnation of a worker whose earlier one
+    still runs. So `compute_unit` writes outputs which depend on the settings and the
+    index only (deterministic), each written atomically (a temporary file in the same
+    directory, then renamed), so a reader never sees a partial unit.
+
+    Attributes:
+        name: Name of the type, as used in a workflow.
+        settings: Pydantic model of the settings.
+        inputs: Names of the input ports.
+        outputs: Names of the output ports.
+        hint: Default resources.
+        checkpoints: Whether the task resumes from checkpoints.
+        command: Argv template of unsplit tasks with `{settings}`, `{out}` and
+            `{input}`: a part with `{input}` is repeated once per input port of the
+            task, sorted by port, with `<port>=<path>` (`<scratch>/<port>` for a
+            staged port), and dropped without inputs.
+            It runs through `simrun-exec` (see `simrun.worker.exec_main`).
+        uses_cores: Whether the command of an unsplit task uses the cores of its pod
+            (a pool of processes); else the pod is sized for one process.
+        count_units: Reference to the function counting the units of work.
+        setup: Reference to the function preparing a process.
+        compute_unit: Reference to the function computing one unit.
+        merge: Reference to the function merging the units.
+        plan: Reference to the function which completes the settings of a task
+            before it is submitted, e.g. with values derived from its model.
+        image: Image of the task, None for the environment image of the run.
+        cost_setting: Name of the setting which changes the cost of a task, e.g. the
+            subset of the data of a fit; the resources are learned per value of it
+            (`simrun.resources.resource_key`).
+        description: What a task of the type does in a few words, the subtitle of
+            the type in the palette of the Editor.
+    """
+
+    name: str
+    settings: type[BaseModel]
+    inputs: tuple[str, ...] = ()
+    outputs: tuple[str, ...] = ("out",)
+    hint: ResourceHint = ResourceHint()
+    checkpoints: bool = False
+    command: str | None = None
+    uses_cores: bool = False
+    count_units: str | None = None
+    setup: str | None = None
+    compute_unit: str | None = None
+    merge: str | None = None
+    plan: str | None = None
+    image: str | None = None
+    cost_setting: str | None = None
+    description: str = ""
+
+    @property
+    def splittable(self) -> bool:
+        """Whether the task can be split into pieces."""
+        return None not in (self.count_units, self.compute_unit, self.merge)
+
+
+def resolve(ref: str) -> Callable[..., Any]:
+    """Import the function a `"module:function"` string names.
+
+    Args:
+        ref: Reference of the form `module:function`.
+
+    Returns:
+        The function.
+    """
+    module, _, name = ref.partition(":")
+    return getattr(import_module(module), name)
+
+
+def load_task_types() -> dict[str, TaskType]:
+    """Load the task types of all entry points.
+
+    Returns:
+        The task types by name.
+
+    Raises:
+        ValueError: A name is registered twice.
+    """
+    types: dict[str, TaskType] = {}
+    for entry in entry_points(group=ENTRY_POINT_GROUP):
+        task_type = entry.load()
+        if task_type.name in types:
+            raise ValueError(f"duplicate task type {task_type.name!r}")
+        types[task_type.name] = task_type
+    return types

simrun/plugins_builtin.py

@@ -0,0 +1,33 @@
+"""The built-in task types."""
+
+from pydantic import BaseModel
+
+from simrun.plugins import TaskType
+
+
+class CommandSettings(BaseModel):
+    """Settings of a plain command task.
+
+    Attributes:
+        image: Image the command runs in.
+        args: Argument vector of the command.
+        cores: Number of cores.
+        memory_gb: Memory in GiB.
+        gpus: Number of GPUs.
+        scratch_gb: Local scratch in GiB.
+    """
+
+    image: str
+    args: list[str]
+    cores: int
+    memory_gb: float
+    gpus: int = 0
+    scratch_gb: float = 0
+
+
+# `command` is None: the settings carry the image and the args.
+COMMAND = TaskType(
+    name="command",
+    settings=CommandSettings,
+    description="A command in the image of the run",
+)