slurm-workflows 1.0.0__py3-none-any.whl

slurm_workflows/__init__.py

@@ -0,0 +1 @@
+from .slurm_pilot_executor import SlurmPilotExecutor, check_for_error

slurm_workflows/ds_service.py

@@ -0,0 +1,44 @@
+"""Run ds-service."""
+
+import shlex
+import subprocess
+
+from .utils import (
+    Closeable,
+    terminate_gracefully,
+    ignoring_sigint,
+)
+
+
+class DsService(Closeable):
+    """Run ds-service server locally"""
+
+    def __init__(
+        self,
+        host: str = "0.0.0.0",
+        port: int = 5051,
+        server_exe: str = "ds-service",
+    ):
+        self.address = f"{host}:{port}"
+        self.server_exe = server_exe
+        self._proc: subprocess.Popen | None = None
+
+    def start(self):
+        cmd = self.server_exe + f" --address {self.address}"
+        cmd = shlex.split(cmd)
+        print("Starting server ...")
+        print("executing: ", " ".join(cmd))
+
+        assert self._proc is None
+        with ignoring_sigint():
+            self._proc = subprocess.Popen(
+                cmd,
+                stdin=subprocess.DEVNULL,
+            )
+
+        print(f"server address: {self.address}")
+
+    def close(self):
+        if self._proc is not None:
+            terminate_gracefully(self._proc)
+            self._proc = None

slurm_workflows/run_jupyter.py

@@ -0,0 +1,57 @@
+"""Start a Jupyter Lab instance."""
+
+import sys
+import subprocess
+from datetime import datetime
+from pathlib import Path
+
+import click
+import platformdirs
+
+from .templates import render_template
+from .slurm_utils import submit_sbatch_job
+
+JUPYTER_EXE = "jupyter"
+
+
+@click.command()
+@click.option(
+    "--setup-script",
+    type=click.Path(exists=True, file_okay=True, dir_okay=False, path_type=Path),
+    required=True,
+    help="Path to setup script.",
+)
+@click.argument("sbatch-args", nargs=-1)
+def run_jupyter(
+    sbatch_args: list[str],
+    setup_script: Path,
+):
+    """Start a Jupyter Lab instance."""
+    print("Sbatch args: ", " ".join(sbatch_args))
+
+    name = "jupyter"
+
+    script = render_template(
+        "run_jupyter:script_template",
+        setup_script=setup_script,
+        jupyter_executable=JUPYTER_EXE,
+    )
+
+    now = datetime.now().isoformat()
+    work_dir = platformdirs.user_cache_path(appname=f"run-jupyter") / now
+    work_dir.mkdir(parents=True)
+
+    try:
+        job = submit_sbatch_job(
+            name=name, sbatch_args=sbatch_args, script=script, work_dir=work_dir
+        )
+
+        print(f"Job ID: {job.job_id}")
+        print(f"Output file: {job.output_file!s}")
+    except subprocess.CalledProcessError as cp:
+        print(f"Failed to submit job: {cp.returncode}")
+        if cp.stdout.strip():
+            print(cp.stdout)
+        if cp.stderr.strip():
+            print(cp.stderr)
+        sys.exit(1)

slurm_workflows/slurm_pilot_executor.py

@@ -0,0 +1,384 @@
+"""Pilot workers for slurm."""
+
+from __future__ import annotations
+
+import time
+import json
+import pickle
+import logging
+import subprocess
+from pathlib import Path
+from datetime import datetime
+from dataclasses import dataclass, field
+from typing import Callable, Iterable, Any, cast
+
+import platformdirs
+import cloudpickle
+from typeguard import typechecked
+from tqdm import tqdm
+from ds_service_client import Client, TaskState
+
+from .slurm_utils import (
+    get_running_jobids,
+    cancel_jobs,
+    submit_sbatch_job,
+    SlurmJob,
+)
+
+from .utils import (
+    RemoteExecutionError,
+    gen_random_string,
+    LOG_FORMAT,
+    LOG_LEVEL,
+)
+
+from .templates import render_template
+
+NoOutput = object()
+
+POLL_INTERVAL_S: float = 0.1
+
+
+@dataclass
+class Task:
+    task_id: str
+    queue: list[str]
+    priority: float
+    function: Callable | str
+    input: tuple
+    output: Any
+
+
+@dataclass
+class WorkerGroup:
+    name: str
+    sbatch_args: list[str]
+    is_batch_worker: bool
+    worker_exe: str
+    actor_class_name: str
+    setup_script: str
+    python_paths: list[str]
+    workers: dict[str, SlurmJob] = field(default_factory=dict, compare=False)
+    next_worker_index: int = field(default=0, compare=False)
+
+
+class SlurmPilotExecutor:
+    def __init__(
+        self,
+        server_address: str,
+        work_dir: Path | str | None = None,
+    ):
+        self.executor_id = gen_random_string()
+        self.next_task_index = 0
+
+        self.server_address = server_address
+        self.client = Client(server_address)
+
+        if work_dir is None:
+            now = datetime.now().isoformat()
+            work_dir = platformdirs.user_cache_path(appname=f"slurm-pilot") / now
+        self.work_dir = Path(work_dir)
+        self.work_dir.mkdir(parents=True, exist_ok=True)
+
+        self.logger = logging.getLogger("pilot_coordinator")
+        self.logger.setLevel(LOG_LEVEL)
+        handler = logging.FileHandler(self.work_dir / "coordinator.log", delay=True)
+        handler.setLevel(LOG_LEVEL)
+        formatter = logging.Formatter(LOG_FORMAT)
+        handler.setFormatter(formatter)
+        self.logger.addHandler(handler)
+
+        self.groups: dict[str, WorkerGroup] = {}
+
+    @typechecked
+    def define_worker(
+        self,
+        name: str,
+        sbatch_args: list[str],
+        setup_script: str,
+        worker_exe: str = "slurm-pilot-worker",
+        is_batch_worker: bool = False,
+        actor_class_name: str | None = None,
+        python_paths: list[str | Path] | None = None,
+        add_cwd_to_python_path: bool = True,
+    ) -> None:
+        python_str_paths: list[str] = []
+        if python_paths is not None:
+            for path in python_paths:
+                python_str_paths.append(str(path))
+        if add_cwd_to_python_path:
+            python_str_paths.append(str(Path.cwd()))
+
+        if actor_class_name is None:
+            actor_class_name = ""
+
+        group = WorkerGroup(
+            name=name,
+            sbatch_args=sbatch_args,
+            worker_exe=worker_exe,
+            is_batch_worker=is_batch_worker,
+            actor_class_name=actor_class_name,
+            setup_script=setup_script,
+            python_paths=python_str_paths,
+        )
+
+        if group.name in self.groups:
+            assert self.groups[group.name] == group
+        else:
+            self.groups[group.name] = group
+
+    def _add_worker(self, group: WorkerGroup) -> None:
+        worker_index = group.next_worker_index
+        group.next_worker_index += 1
+        name = f"slurm_pilot_worker.{group.name}.{worker_index}"
+
+        worker_script = render_template(
+            "slurm_pilot:worker_script",
+            group=group.name,
+            name=name,
+            server_address=self.server_address,
+            worker_exe=group.worker_exe,
+            work_dir=str(self.work_dir),
+            python_paths_json=json.dumps(group.python_paths),
+            setup_script=group.setup_script,
+            actor_class_name=group.actor_class_name,
+        )
+        worker_script_path = self.work_dir / f"{name}.sh"
+        worker_script_path.write_text(worker_script)
+        worker_script_path.chmod(0o755)
+
+        worker_sbatch_script = render_template(
+            "slurm_pilot:worker_sbatch_script",
+            is_batch_worker=group.is_batch_worker,
+            worker_script_path=worker_script_path,
+        )
+
+        self.logger.info("Starting worker %s", name)
+        try:
+            slurm_job = submit_sbatch_job(
+                name=name,
+                sbatch_args=group.sbatch_args,
+                script=worker_sbatch_script,
+                work_dir=self.work_dir,
+            )
+            group.workers[name] = slurm_job
+        except subprocess.CalledProcessError as cp:
+            print(f"Failed to submit slurm job: returncode={cp.returncode}")
+            if cp.stdout.strip():
+                print(cp.stdout)
+            if cp.stderr.strip():
+                print(cp.stderr)
+            raise cp
+
+    @typechecked
+    def scale_workers(self, name: str, count: int) -> None:
+        assert name in self.groups, "Unknown worker type"
+
+        group = self.groups[name]
+        if len(group.workers) < count:
+            to_hire = count - len(group.workers)
+            for _ in range(to_hire):
+                self._add_worker(group)
+
+        if len(group.workers) > count:
+            to_retire = len(group.workers) - count
+
+            try:
+                running_jobids = get_running_jobids()
+            except subprocess.CalledProcessError as cp:
+                print(
+                    f"Failed to get running slurm job ids: returncode={cp.returncode}"
+                )
+                if cp.stdout.strip():
+                    print(cp.stdout)
+                if cp.stderr.strip():
+                    print(cp.stderr)
+                raise RuntimeError("Failed to get running slurm job ids")
+            except Exception:
+                raise RuntimeError("Failed to get running slurm job ids")
+
+            to_cancel_jobids = []
+            for _ in range(to_retire):
+                _, worker = group.workers.popitem()
+                self.logger.info("Canceling worker: %s", worker.name)
+                if worker.job_id in running_jobids:
+                    to_cancel_jobids.append(worker.job_id)
+
+            if not to_cancel_jobids:
+                return
+
+            try:
+                cancel_jobs(to_cancel_jobids)
+            except subprocess.CalledProcessError as cp:
+                print(f"Failed to cancel slurm jobs: returncode={cp.returncode}")
+                if cp.stdout.strip():
+                    print(cp.stdout)
+                if cp.stderr.strip():
+                    print(cp.stderr)
+                raise RuntimeError("Failed to cancel slurm jobs")
+            except Exception:
+                raise RuntimeError("Failed to cancel slurm jobs")
+
+    def _submit(
+        self,
+        queue: list[str],
+        fn: Callable | str,
+        *args,
+        **kwargs,
+    ) -> Task:
+        priority = time.perf_counter()
+        function_bytes = cloudpickle.dumps(fn, protocol=pickle.HIGHEST_PROTOCOL)
+        input_bytes = cloudpickle.dumps(
+            (args, kwargs), protocol=pickle.HIGHEST_PROTOCOL
+        )
+
+        task_id = f"{self.executor_id}:{self.next_task_index}"
+        self.next_task_index += 1
+
+        task = Task(
+            task_id=task_id,
+            queue=queue,
+            priority=priority,
+            function=fn,
+            input=(args, kwargs),
+            output=NoOutput,
+        )
+
+        self.client.task_add(
+            task_id=task_id,
+            queue=queue,
+            priority=priority,
+            function=function_bytes,
+            input=input_bytes,
+        )
+        return task
+
+    @typechecked
+    def submit(
+        self, queue: str | list[str], fn: Callable | str, *args, **kwargs
+    ) -> Task:
+        if isinstance(queue, str):
+            queue = [queue]
+
+        return self._submit(queue, fn, *args, **kwargs)
+
+    def _as_completed(self, tasks: list[Task]) -> Iterable[Task]:
+        pending: list[Task] = []
+        for task in tasks:
+            if task.output is NoOutput:
+                pending.append(task)
+            else:
+                yield task
+
+        while pending:
+            # Status for every pending task comes back in a single request,
+            # in the same order as the ids we sent.
+            states = self.client.task_get_status([t.task_id for t in pending])
+            states = cast(list[Task], states)
+
+            next_pending: list[Task] = []
+            completed = 0
+            for task, state in zip(pending, states):
+                if state == TaskState.Complete:
+                    output = self.client.task_get_output(task.task_id)
+                    task.output = cloudpickle.loads(output)
+                    completed += 1
+                    yield task
+                elif state == TaskState.Undefined:
+                    raise RuntimeError(
+                        f"Task {task.task_id} is unknown to the task queue server"
+                    )
+                else:
+                    next_pending.append(task)
+
+            pending = next_pending
+            if pending and not completed:
+                time.sleep(POLL_INTERVAL_S)
+
+    @typechecked
+    def as_completed(
+        self, tasks: Iterable[Task], desc: str | None = None, unit: str = "task"
+    ) -> Iterable[Task]:
+        tasks = list(tasks)
+        iterable = self._as_completed(tasks)
+        iterable = tqdm(iterable, total=len(tasks), desc=desc, unit=unit)
+        return iterable
+
+    @typechecked
+    def wait(
+        self, tasks: Iterable[Task], desc: str | None = None, unit: str = "task"
+    ) -> None:
+        for _ in self.as_completed(tasks, desc, unit):
+            pass
+
+    def num_groups(self):
+        return len(self.groups)
+
+    def num_workers(self, detail: bool = False):
+        if detail:
+            return {g.name: len(g.workers) for g in self.groups.values()}
+        else:
+            return sum(len(g.workers) for g in self.groups.values())
+
+    def _cleanup_all_workers(self):
+        try:
+            job_ids = get_running_jobids()
+        except subprocess.CalledProcessError as cp:
+            print(f"Failed to get running slurm job ids: returncode={cp.returncode}")
+            if cp.stdout.strip():
+                print(cp.stdout)
+            if cp.stderr.strip():
+                print(cp.stderr)
+            job_ids = None
+        except Exception:
+            self.logger.exception("Failed to get running slurm job ids")
+            job_ids = None
+
+        if job_ids is None:
+            return
+
+        to_cancel_jobids = []
+        for group in self.groups.values():
+            for worker in group.workers.values():
+                if worker.job_id in job_ids:
+                    to_cancel_jobids.append(worker.job_id)
+
+        if not to_cancel_jobids:
+            return
+
+        try:
+            cancel_jobs(to_cancel_jobids)
+        except subprocess.CalledProcessError as cp:
+            print(f"Failed to cancel slurm jobs: returncode={cp.returncode}")
+            if cp.stdout.strip():
+                print(cp.stdout)
+            if cp.stderr.strip():
+                print(cp.stderr)
+        except Exception:
+            self.logger.exception("Failed to cancel slurm jobs")
+
+    def close(self):
+        self._cleanup_all_workers()
+        for group in self.groups.values():
+            group.workers.clear()
+
+        self.client.close()
+
+    def stop(self):
+        self._cleanup_all_workers()
+        for group in self.groups.values():
+            group.workers.clear()
+
+
+def check_for_error(tasks: list[Task], verbose: bool = True) -> list[Task]:
+    ret = []
+    for task in tasks:
+        if isinstance(task.output, RemoteExecutionError):
+            ret.append(task)
+
+            if verbose:
+                print(f"task_id={task.task_id}")
+                print(f"  error={task.output.error}")
+                print(f"  error_id={task.output.error_id}")
+
+    return ret

slurm_workflows/slurm_pilot_worker.py

@@ -0,0 +1,168 @@
+"""Pilot workers for Slurm pilot."""
+
+import os
+import sys
+import json
+import time
+import pickle
+import socket
+import logging
+import importlib
+from pathlib import Path
+from typing import Any
+
+import click
+import cloudpickle
+from ds_service_client import Client
+
+from .utils import gen_error_id, RemoteExecutionError, LOG_FORMAT, LOG_LEVEL
+
+NEXT_TASK_RETRY_TIME_S: float = 0.1
+
+
+class PilotWorkerProcess:
+    def __init__(
+        self,
+        group: str,
+        name: str,
+        actor_class_name: str,
+        server_address: str,
+        work_dir: Path,
+        slurm_job_id: int,
+        hostname: str,
+        pid: int,
+    ):
+        self.group = group
+        self.server_address = server_address
+        self.work_dir = work_dir
+
+        self.worker_id = "%s:%s:%s:%s:%s" % (group, name, slurm_job_id, hostname, pid)
+        self.logger = logging.getLogger("worker_process")
+        self.client = Client(self.server_address)
+
+        self.actor_instance: Any | None
+        if actor_class_name == "":
+            self.actor_instance = None
+        else:
+            class_name_parts = actor_class_name.split(".")
+            module_name = ".".join(class_name_parts[:-1])
+            class_name = class_name_parts[-1]
+
+            module = importlib.import_module(module_name)
+            klass = getattr(module, class_name)
+            self.actor_instance = klass()
+
+    def close(self):
+        self.client.close()
+        if self.actor_instance is not None:
+            if hasattr(self.actor_instance, "close"):
+                self.actor_instance.close()
+            self.actor_instance = None
+
+    def main(self):
+        self.logger.info("Starting worker: %s" % self.worker_id)
+
+        while True:
+            try:
+                task = self.client.task_get(self.worker_id, self.group)
+                try:
+                    self.logger.info(
+                        "task_id=%s: Deserializing function and inputs ...",
+                        task.task_id,
+                    )
+                    function = cloudpickle.loads(task.function)
+                    if self.actor_instance is not None:
+                        function = getattr(self.actor_instance, function)
+                    args, kwargs = cloudpickle.loads(task.input)
+
+                    self.logger.info("task_id=%s: Executing ...", task.task_id)
+                    retval = function(*args, **kwargs)
+
+                    self.logger.info("task_id=%s: Serializng output ...", task.task_id)
+                    output = cloudpickle.dumps(retval, protocol=pickle.HIGHEST_PROTOCOL)
+
+                    self.client.task_done(task.task_id, output)
+                except Exception as e:
+                    eid = gen_error_id()
+                    self.logger.exception(
+                        "Error executing %s: %s: %s", task.task_id, eid, e
+                    )
+
+                    retval = RemoteExecutionError(error=str(e), error_id=eid)
+                    output = cloudpickle.dumps(retval, protocol=pickle.HIGHEST_PROTOCOL)
+                    self.client.task_done(task.task_id, output)
+            except TimeoutError:
+                time.sleep(NEXT_TASK_RETRY_TIME_S)
+            except Exception:
+                self.logger.exception("Unexpected exception")
+
+
+@click.command()
+@click.option("--group", type=str, required=True, help="Worker group.")
+@click.option("--name", type=str, required=True, help="Worker job name.")
+@click.option(
+    "--actor-class-name",
+    type=str,
+    required=True,
+    help="Name for actor class in DS server store.",
+)
+@click.option("--server-address", type=str, required=True, help="Pilot server address.")
+@click.option(
+    "--work-dir",
+    type=click.Path(exists=True, file_okay=False, dir_okay=True, path_type=Path),
+    required=True,
+    help="Work directory.",
+)
+@click.option(
+    "--python-paths-json",
+    type=str,
+    required=True,
+    help="JSON encoded Python paths.",
+)
+def slurm_pilot_worker(
+    group: str,
+    name: str,
+    actor_class_name: str,
+    server_address: str,
+    work_dir: Path,
+    python_paths_json: str,
+):
+    """Start a slurm pilot worker."""
+    slurm_job_id = int(os.environ.get("SLURM_JOB_ID", -1))
+    hostname = socket.gethostname()
+    pid = os.getpid()
+    log_file = work_dir / f"{name}-{slurm_job_id}-{hostname}-{pid}.log"
+
+    print(f"Redirecting standard output and standard error to {log_file}")
+    sys.stdout.flush()
+    sys.stderr.flush()
+
+    with open(log_file, "wt") as fout:
+        sys.stdout = fout
+        sys.stderr = fout
+
+        logging.basicConfig(stream=fout, format=LOG_FORMAT, level=LOG_LEVEL)
+
+        os.environ["PILOT_WORKER_NAME"] = name
+        os.environ["PILOT_WORKER_GROUP"] = group
+        os.environ["DS_SERVER_ADDRESS"] = server_address
+
+        python_paths: list[str] = json.loads(python_paths_json)
+        for path in python_paths:
+            sys.path.insert(0, path)
+
+        worker = PilotWorkerProcess(
+            group=group,
+            name=name,
+            actor_class_name=actor_class_name,
+            server_address=server_address,
+            work_dir=work_dir,
+            slurm_job_id=slurm_job_id,
+            hostname=hostname,
+            pid=pid,
+        )
+
+        try:
+            worker.main()
+        finally:
+            worker.close()