kkrunner 0.1.0b1__py3-none-any.whl

kkrunner/__init__.py

@@ -0,0 +1,16 @@
+"""Backend-independent entry points for execution and durable run management."""
+
+from kkrunner._version import __version__
+from kkrunner.files import TaskFile
+from kkrunner.runs import ExecutionOutcome, ExecutionSpec, ProcessingPending, RunManager, RunState, RunStore
+
+__all__ = [
+    "ExecutionOutcome",
+    "ExecutionSpec",
+    "ProcessingPending",
+    "RunManager",
+    "RunState",
+    "RunStore",
+    "TaskFile",
+    "__version__",
+]

kkrunner/__main__.py

@@ -0,0 +1,3 @@
+from kkrunner.cli import main
+
+main()

kkrunner/_version.py

@@ -0,0 +1,5 @@
+"""Version of the installed KKRunner distribution."""
+
+from importlib.metadata import version
+
+__version__ = version("kkrunner")

kkrunner/backends/__init__.py

@@ -0,0 +1,35 @@
+"""Execution backends and their reusable endpoint identities."""
+
+from importlib import import_module
+from typing import TYPE_CHECKING
+
+from kkrunner.backends.endpoint import ExecutorEndpoint
+
+if TYPE_CHECKING:
+    from kkrunner.backends.local import LocalExecutor
+    from kkrunner.backends.pool import ExecutorPool, default_executor_pool
+    from kkrunner.backends.psij import PsijExecutor
+
+# Run contracts use endpoint metadata. Load implementations only on request so
+# importing those contracts does not load backends or create an import cycle.
+_EXPORTS = {
+    "LocalExecutor": "kkrunner.backends.local",
+    "PsijExecutor": "kkrunner.backends.psij",
+    "ExecutorPool": "kkrunner.backends.pool",
+    "default_executor_pool": "kkrunner.backends.pool",
+}
+
+__all__ = ["ExecutorEndpoint", "ExecutorPool", "LocalExecutor", "PsijExecutor", "default_executor_pool"]
+
+
+def __getattr__(name):
+    module = _EXPORTS.get(name)
+    if module is None:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+    value = getattr(import_module(module), name)
+    globals()[name] = value
+    return value
+
+
+def __dir__():
+    return sorted(set(globals()) | set(__all__))

kkrunner/backends/endpoint.py

@@ -0,0 +1,179 @@
+"""Serializable execution endpoint identity and reconstruction."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+
+@dataclass(frozen=True)
+class ExecutorEndpoint:
+    """Identity of a reusable execution backend.
+
+    The endpoint contains only connection/backend properties. Launch and
+    staging policy deliberately live elsewhere so one executor can be shared
+    by profiles that differ only in how individual calculations are launched.
+    The value is serializable and hashable, so it can also be persisted with a
+    run and used as an :class:`ExecutorPool` key after process restart.
+    """
+
+    executor: str = "local"
+    host: Any = None
+    ssh_options: Any = ()
+    remote_command: Any = ()
+    psij_backend: Any = None
+    psij_url: Any = None
+    ssh_zstd: bool = False
+    filesystem_id: str | None = None
+
+    def __post_init__(self):
+        object.__setattr__(self, "executor", str(self.executor or "local"))
+        if self.filesystem_id is not None and (not isinstance(self.filesystem_id, str) or not self.filesystem_id):
+            raise TypeError("filesystem_id must be a nonempty string or None")
+        for name in ("ssh_options", "remote_command"):
+            value = getattr(self, name)
+            if value is None:
+                value = ()
+            elif isinstance(value, str):
+                value = (value,)
+            else:
+                value = tuple(value)
+            object.__setattr__(self, name, value)
+
+        for name in ("host", "psij_backend", "psij_url"):
+            value = getattr(self, name)
+            if value is not None and not isinstance(value, str):
+                raise TypeError(f"ExecutorEndpoint.{name} must be a string or None")
+        for name in ("ssh_options", "remote_command"):
+            if any(not isinstance(item, str) for item in getattr(self, name)):
+                raise TypeError(f"ExecutorEndpoint.{name} must contain strings")
+        if type(self.ssh_zstd) is not bool:
+            raise TypeError("ExecutorEndpoint.ssh_zstd must be a bool")
+
+        # Identity describes the backend actually constructed below. Irrelevant
+        # profile fields must not split its pool key or collision namespace.
+        if self.executor in ("local", "psij"):
+            for name, value in (
+                ("host", None),
+                ("ssh_options", ()),
+                ("remote_command", ()),
+                ("ssh_zstd", False),
+            ):
+                object.__setattr__(self, name, value)
+        if self.executor == "local":
+            object.__setattr__(self, "psij_backend", None)
+            object.__setattr__(self, "psij_url", None)
+        elif self.executor == "psij":
+            object.__setattr__(self, "psij_backend", self.psij_backend or "slurm")
+        elif self.executor == "ssh":
+            if not self.remote_command:
+                object.__setattr__(self, "remote_command", ("kkrunner", "remote"))
+            if self.psij_backend is not None or self.psij_url is not None:
+                object.__setattr__(self, "psij_backend", self.psij_backend or "slurm")
+
+
+def endpoint_for_executor(executor) -> ExecutorEndpoint:
+    """Describe a live executor without persisting launch/staging policy."""
+    from kkrunner.backends.local import LocalExecutor
+    from kkrunner.backends.psij import PsijExecutor
+    from kkrunner.remote.executor import RemoteEndpointExecutor
+    from kkrunner.remote.transport import SshTransport
+
+    if isinstance(executor, LocalExecutor):
+        return ExecutorEndpoint()
+    if isinstance(executor, PsijExecutor):
+        return ExecutorEndpoint(
+            executor="psij",
+            psij_backend=executor.backend,
+            psij_url=executor.url,
+        )
+    if isinstance(executor, RemoteEndpointExecutor):
+        transport = executor.transport
+        psij_backend = None
+        psij_url = None
+        if executor.backend == "psij":
+            options = executor.backend_options or {}
+            # ``RemoteEndpointExecutor(..., backend="psij")`` is a complete
+            # public configuration even when no options are supplied: the
+            # target-side PSI/J backend then defaults to Slurm. Persist that
+            # default explicitly so restart reconstruction cannot silently
+            # turn the endpoint into the ordinary remote-local backend.
+            psij_backend = options.get("backend") or "slurm"
+            psij_url = options.get("url")
+
+        if not isinstance(transport, SshTransport):
+            return ExecutorEndpoint(
+                executor=(
+                    "custom:{}.{}:{}".format(
+                        type(executor).__module__,
+                        type(executor).__qualname__,
+                        executor.name or "remote",
+                    )
+                ),
+                host=getattr(transport, "host", None),
+                remote_command=executor.command,
+                filesystem_id=executor.filesystem_id,
+                psij_backend=psij_backend,
+                psij_url=psij_url,
+            )
+        return ExecutorEndpoint(
+            executor="ssh",
+            host=transport.host,
+            ssh_options=transport.options,
+            ssh_zstd=transport.zstd,
+            remote_command=executor.command,
+            filesystem_id=executor.filesystem_id,
+            psij_backend=psij_backend,
+            psij_url=psij_url,
+        )
+
+    # Custom executors can participate in duplicate-submit protection in this
+    # process. Reconstruction is intentionally limited to known endpoint types.
+    return ExecutorEndpoint(
+        executor="custom:{}.{}:{}".format(
+            type(executor).__module__,
+            type(executor).__qualname__,
+            executor.name or type(executor).__qualname__,
+        )
+    )
+
+
+def create_executor_from_endpoint(endpoint: ExecutorEndpoint):
+    """Reconstruct a standard executor from persisted endpoint identity."""
+    from kkrunner.backends.local import LocalExecutor
+    from kkrunner.backends.psij import PsijExecutor
+    from kkrunner.remote.executor import RemoteEndpointExecutor
+    from kkrunner.remote.transport import SshTransport
+
+    if isinstance(endpoint, dict):
+        endpoint = ExecutorEndpoint(**endpoint)
+
+    if endpoint.executor == "local":
+        return LocalExecutor()
+    if endpoint.executor == "psij":
+        return PsijExecutor(
+            endpoint.psij_backend or "slurm",
+            url=endpoint.psij_url,
+        )
+    if endpoint.executor == "ssh":
+        if not endpoint.host:
+            raise ValueError("SSH execution endpoint requires 'host'")
+        kwargs = {}
+        if endpoint.remote_command:
+            kwargs["command"] = endpoint.remote_command
+        if endpoint.psij_backend is not None or endpoint.psij_url is not None:
+            kwargs["backend"] = "psij"
+            kwargs["backend_options"] = {
+                "backend": endpoint.psij_backend or "slurm",
+                "url": endpoint.psij_url,
+            }
+        return RemoteEndpointExecutor(
+            SshTransport(
+                endpoint.host,
+                options=endpoint.ssh_options,
+                zstd=endpoint.ssh_zstd,
+            ),
+            filesystem_id=endpoint.filesystem_id,
+            **kwargs,
+        )
+    raise ValueError(f"Cannot reconstruct unsupported execution endpoint {endpoint.executor!r}")

kkrunner/backends/local.py

@@ -0,0 +1,328 @@
+"""Local subprocess execution backend."""
+
+from __future__ import annotations
+
+import asyncio
+from threading import Event, Lock, Thread
+
+from kkrunner.backends.local_process import materialize_local_spec, pump_stream, start_local_process, stop_local_process
+from kkrunner.files.availability import clear_declared_outputs, discover_local_files, validate_output_collisions
+from kkrunner.files.manifest import FileManifest
+from kkrunner.files.paths import file_path
+from kkrunner.runs.base import (
+    ExecutionOutcome,
+    ExecutionSpec,
+    Executor,
+    FailurePhase,
+    OutputObservers,
+    RunState,
+    SubmissionNotStarted,
+    close_output_files,
+    finish_monitor,
+    materialize_output_streams,
+    wait_for_monitor,
+)
+
+
+class LocalRunner:
+    """One local subprocess execution.
+
+    The runner owns only process lifecycle, output-file writing and raw output
+    observation. It deliberately performs no task-specific parsing.
+    """
+
+    def __init__(self, spec: ExecutionSpec, callback=None):
+        self.spec = spec
+        self.callback = callback
+        self.result = ExecutionOutcome(
+            output_file=file_path(spec.output_file, spec.directory),
+            input_file=file_path(spec.input_file, spec.directory),
+            error_file=file_path(spec.error_file, spec.directory),
+        )
+
+        self._state = RunState.NEW
+        self._output_observers = OutputObservers()
+        self._lock = Lock()
+        self._cancel_requested = Event()
+        self._started = Event()
+        self._done = Event()
+
+        self._thread = None
+        self._loop = None
+        self._process = None
+        self._pid = None
+        self._stop_task = None
+        self._exception = None
+        self._callback_exception = None
+        self._failure_phase = None
+        self._failure_message = None
+
+    @property
+    def state(self):
+        return self._state
+
+    @property
+    def exception(self):
+        return self._exception
+
+    @property
+    def backend_id(self):
+        return None if self._pid is None else str(self._pid)
+
+    @property
+    def failure_phase(self):
+        return self._failure_phase
+
+    @property
+    def failure_message(self):
+        return self._failure_message
+
+    @property
+    def return_code(self):
+        return self.result.return_code
+
+    def add_output_listener(self, callback):
+        self._output_observers.add(callback)
+
+    def add_output_error_listener(self, callback):
+        self._output_observers.add_error_listener(callback)
+
+    def _output_target(self, kind):
+        # A missing stderr target means live observation only, not an
+        # implicit merge into the saved stdout artifact.
+        return self.spec.error_file if kind == "err" else self.spec.output_file
+
+    def _forward_output(self, data, kind):
+        target = self._output_target(kind)
+        if target is not None:
+            target.write(data)
+            target.flush()
+        self._output_observers.emit_bytes(data, kind)
+
+    def _ensure_stop_task(self):
+        if self._process is None or self._process.returncode is not None:
+            return None
+        if self._stop_task is None or self._stop_task.done():
+            self._stop_task = asyncio.create_task(stop_local_process(self._process))
+        return self._stop_task
+
+    def cancel(self):
+        self._cancel_requested.set()
+        loop = self._loop
+        if loop is None:
+            if self._state == RunState.NEW:
+                self._state = RunState.CANCELLED
+            return
+        try:
+            loop.call_soon_threadsafe(self._ensure_stop_task)
+        except RuntimeError:
+            pass
+
+    async def _execute(self):
+        self._loop = asyncio.get_running_loop()
+        if self._cancel_requested.is_set():
+            self._state = RunState.CANCELLED
+            self._started.set()
+            return self.result
+
+        try:
+            protected = [
+                file_path(value, self.spec.directory)
+                for value in (self.spec.input_file, self.spec.stdin, *self.spec.input_files, *self.spec.cache_files)
+            ]
+            # Preserve already-open stdout/stderr handles: an unlink would
+            # disconnect their inode from the named destination.
+            outputs = [file_path(value, self.spec.directory) for value in self.spec.output_files]
+            validate_output_collisions(
+                (*outputs, self.result.output_file, self.result.error_file),
+                protected=protected,
+            )
+            self.spec = materialize_output_streams(self.spec)
+            clear_declared_outputs(
+                outputs,
+                protected=(*protected, self.result.output_file, self.result.error_file),
+            )
+            try:
+                self._process = await start_local_process(self.spec)
+            except BaseException as exc:
+                self._failure_phase = FailurePhase.LAUNCH
+                self._failure_message = str(exc)
+                self._state = RunState.FAILED
+                raise
+
+            self._pid = self._process.pid
+            self._state = RunState.RUNNING
+            self._started.set()
+
+            if self._cancel_requested.is_set():
+                stop_task = self._ensure_stop_task()
+                if stop_task is not None:
+                    await stop_task
+
+            wait_task = asyncio.create_task(self._process.wait())
+            stdout_task = asyncio.create_task(
+                pump_stream(
+                    self._process.stdout,
+                    lambda data: self._forward_output(data, "out"),
+                )
+            )
+            stderr_task = asyncio.create_task(
+                pump_stream(
+                    self._process.stderr,
+                    lambda data: self._forward_output(data, "err"),
+                )
+            )
+            tasks = (wait_task, stdout_task, stderr_task)
+            try:
+                return_code, _, _ = await asyncio.gather(*tasks)
+            finally:
+                for task in tasks:
+                    if not task.done():
+                        task.cancel()
+                await asyncio.gather(*tasks, return_exceptions=True)
+
+            self.result.return_code = return_code
+            if self._cancel_requested.is_set():
+                self._state = RunState.CANCELLED
+            elif return_code == 0:
+                self._state = RunState.FINISHED
+            else:
+                self._state = RunState.FAILED
+                self._failure_phase = FailurePhase.EXECUTION
+                self._failure_message = f"Process exited with status {return_code}"
+            if self._state == RunState.FINISHED:
+                files = discover_local_files(FileManifest.from_spec(self.spec))
+                if files.failed_artifacts:
+                    self._state = RunState.FAILED
+                    self._failure_phase = FailurePhase.ARTIFACTS
+                    self._failure_message = files.availability_message
+            return self.result
+        except BaseException:
+            self._started.set()
+            if self._state not in (RunState.CANCELLED, RunState.FAILED):
+                self._state = RunState.FAILED
+                self._failure_phase = FailurePhase.EXECUTION
+            if self._process is not None and self._process.returncode is None:
+                await stop_local_process(self._process)
+            raise
+        finally:
+            if self._stop_task is not None:
+                await asyncio.gather(self._stop_task, return_exceptions=True)
+            self._output_observers.finish_bytes()
+            close_output_files(self.spec.output_file, self.spec.error_file)
+            self._process = None
+            self._loop = None
+
+    def _thread_main(self):
+        try:
+            asyncio.run(self._execute())
+        except BaseException as exc:
+            self._exception = exc
+            if self._failure_message is None:
+                self._failure_message = str(exc)
+        finally:
+            finish_monitor(self)
+
+    def submit(self):
+        with self._lock:
+            if self._state == RunState.CANCELLED and self._thread is None:
+                raise RuntimeError("Run was cancelled before submission")
+            if self._state != RunState.NEW:
+                raise RuntimeError(f"Run is already {self._state.value}")
+            self._state = RunState.SUBMITTING
+            self._thread = Thread(target=self._thread_main, daemon=True)
+            self._thread.start()
+
+        self._started.wait()
+        if self._pid is None and self._state == RunState.FAILED:
+            self._done.wait()
+            error = self._exception or RuntimeError(self._failure_message or "Launch failed")
+            raise SubmissionNotStarted(str(error)) from self._exception
+
+    def wait(self):
+        if self._state == RunState.NEW:
+            raise RuntimeError("Run has not been submitted")
+        if self._state == RunState.CANCELLED and self._thread is None:
+            return self.result
+        return wait_for_monitor(self)
+
+    def run(self):
+        if self._state == RunState.NEW:
+            self.submit()
+        return self.wait()
+
+    async def run_async(self):
+        if self._state == RunState.NEW:
+            self.submit()
+        return await asyncio.to_thread(self.wait)
+
+
+class LocalExecutor(Executor):
+    """Create raw runners executed as local subprocesses."""
+
+    name = "local"
+
+    def create_run(self, spec: ExecutionSpec, callback=None):
+        return LocalRunner(materialize_local_spec(spec), callback=callback)
+
+    def create_managed_run(self, spec: ExecutionSpec, *, run_id: str, callback=None):
+        from kkrunner.backends.local_managed import ManagedLocalRunner
+
+        return ManagedLocalRunner(
+            self,
+            spec,
+            run_id=run_id,
+            callback=callback,
+        )
+
+    def recover_submission(self, run_id, submission, *, submitted_at=None):
+        from kkrunner.backends.local_managed import _local_status_from_facts
+        from kkrunner.remote.service import recover_remote
+
+        reply = recover_remote(
+            run_id,
+            {
+                "executor": "local",
+                "directory": submission.execution.directory.local,
+            },
+        )
+        if reply is None:
+            # The absence of a ticket is not proof that submission never
+            # happened. Never authorize an automatic duplicate launch.
+            raise NotImplementedError("Submission identity is not recoverable")
+        return _local_status_from_facts(
+            reply["remote_id"],
+            reply,
+            submission.files,
+        )
+
+    def attach(self, run, callback=None, output_listener=None, start=True):
+        from kkrunner.backends.local_managed import attach_local
+
+        return attach_local(
+            self,
+            run,
+            callback=callback,
+            output_listener=output_listener,
+            start=start,
+        )
+
+    def acknowledge_submission(self, run):
+        from kkrunner.remote.service import acknowledge_submission
+
+        acknowledge_submission(run.id)
+
+    def acknowledge(self, run):
+        from kkrunner.remote.service import acknowledge_remote
+
+        acknowledge_remote(run.execution.backend_id)
+
+    def status(self, *runs):
+        from kkrunner.backends.local_managed import local_status
+
+        statuses = []
+        for run in runs:
+            if run.execution.endpoint.executor != "local":
+                raise ValueError("Run does not belong to the local backend")
+            statuses.append(local_status(run.execution.backend_id, run.files))
+        return tuple(statuses)