diff --git a/README.md b/README.md index a03b3eb..91e8375 100644 --- a/README.md +++ b/README.md @@ -1,23 +1,25 @@ +

+ Slurm Manager +

+ # slurm_manager -Submit Slurm jobs from Python and receive live status, logs, and heartbeats through Redis. +Submit Slurm jobs through slurmrestd and receive live logs, status, and heartbeats through Redis. ## Setup ```bash pip install -e . +cp .env.example .env ``` -The submitter needs Python and the Slurm client commands `sbatch`, `scontrol`, -`scancel`, `squeue`, and `sacct`. Compute nodes use the Bash wrapper and Redis worker; -they do not need this Python package or the submitter's Python environment. +Set `SLURM_JWT` in your local `.env` file. The example contains the RA REST URL +and API version. This credential file is ignored by Git and never uploaded. +The submitter needs access to slurmrestd and Redis. -Compute nodes need `redis-cli`, `jq`, and `stdbuf`, with the site's Redis module -available. The wrapper, Redis worker, workload, and activation script must be -accessible at the paths passed to Slurm. The Redis worker currently publishes to -`bec-slurm.psi.ch`. +## Simple job -## Submit a job +Run the included ten-second Bash script using the cluster's default resources: ```python from slurm_manager.slurm_manager import SlurmManager @@ -25,30 +27,64 @@ from slurm_manager.slurm_manager import SlurmManager manager = SlurmManager(redis_host="bec-slurm.psi.ch") try: job = manager.submit_job( - script_path="/shared/run_task.sh", env_path="/shared/activate_workload.sh" + script_path="examples/print_for_10s.sh", working_directory="/das/home/your_username" ) - job.listen("log", lambda message: print(message.log)) - job.wait(timeout=60) + job.wait() finally: manager.shutdown() ``` -The activation script selects the workload's environment on the compute node. -Jobs are submitted on hold and released after Redis acknowledges the monitoring -subscriptions. Missing heartbeats trigger Slurm queries; queued jobs remain pending. +`script_path` refers to a local Bash file whose contents are uploaded with the job. +For optional module loads or environment activation, pass `env_path="activate.sh"` +with a local Bash activation script. Additional files used by the workload are not uploaded. +Replace `working_directory` with an existing directory accessible on the compute node. -`job.wait()` raises on failure or timeout. A timeout does not cancel the job; -use `job.cancel()` for that. `job.listen()` accepts `status`, `log`, and `heartbeat`; -remove a callback with `job.unlisten(callback_id)`. Manager shutdown disconnects the -client without cancelling remote jobs. +## Customized job -## Development +Choose resources and log files, and receive live output and status updates: -```bash -pip install -e '.[dev]' -ruff check . -ruff format --check . -python -m pytest --random-order -q test +```python +from slurm_manager.job_options import JobOptions +from slurm_manager.slurm_manager import SlurmManager + +manager = SlurmManager(redis_host="bec-slurm.psi.ch") +try: + job = manager.submit_job( + script_path="examples/print_for_10s.sh", + working_directory="/das/home/your_username", + job_options=JobOptions( + name="my-analysis", + partition="day", + tasks=1, + time_limit_minutes=120, + standard_output="/das/home/your_username/slurm_logs/%j.out", + standard_error="/das/home/your_username/slurm_logs/%j.err", + ), + callbacks={ + "log": lambda message: print(message.log, end="", flush=True), + "status": lambda message: print(f"Job status: {message.status}"), + }, + ) + try: + job.wait(timeout=180) + except TimeoutError: + job.cancel() + raise +finally: + manager.shutdown() ``` -Unit tests run on Python 3.12, 3.13, and 3.14. Other CI checks use Python 3.13. +The log directory must already exist on the compute node; `%j` expands to the Slurm +job ID. Output is also forwarded through Redis. Workloads must flush their output +buffers for live feedback. + +Supported partitions are `hour`, `day`, `week`, and `gpu`. Omit `job_options` to use +one task and the cluster's default partition and time limit. `JobOptions` also +accepts `account` and `qos`. Requesting tasks reserves a Slurm allocation; the +workload script runs once and must launch any parallel processes itself. + +`time_limit_minutes` limits job runtime in minutes. `job.wait(timeout=...)` limits +how long the client waits, in seconds; it raises on failure or timeout. A timeout +does not cancel the job automatically, so the example explicitly calls +`job.cancel()`. `manager.shutdown()` disconnects without cancelling jobs. +Callbacks support `log`, `status`, and `heartbeat` and are attached before execution. diff --git a/assets/slurm-manager.png b/assets/slurm-manager.png new file mode 100644 index 0000000..be58ad4 Binary files /dev/null and b/assets/slurm-manager.png differ diff --git a/examples/submit_job.py b/examples/submit_job.py new file mode 100644 index 0000000..58919ac --- /dev/null +++ b/examples/submit_job.py @@ -0,0 +1,76 @@ +"""Submit local Bash files and verify ten stdout lines through Redis. + +Run from the repository root with credentials in .env: + python examples/submit_job.py + +Replace your_username in the paths below before running. +The remote working and log directories must already exist. One job is submitted; +there are no automatic submission retries. +""" + +from pathlib import Path + +from slurm_manager.event_messages import LogMessage +from slurm_manager.job_options import JobOptions +from slurm_manager.slurm_manager import SlurmManager + + +def main( + *, script_path: Path, working_directory: str, log_directory: str, worker_path: str +) -> None: + """Upload local Bash files and verify ten stdout lines through Redis. + + Args: + script_path (Path): Local Bash script that prints once per second. + working_directory (str): Existing absolute compute-node working directory. + log_directory (str): Existing compute-node directory for Slurm output files. + worker_path (str): Absolute compute-node path to the Go worker executable. + + Raises: + TimeoutError: The job exceeds the 90-second caller wait and is cancelled. + Exception: Submission, execution, or callback verification fails. + """ + lines = [] + + def log(message: LogMessage) -> None: + """Record and print a message delivered through the job's Redis log channel. + + Args: + message (LogMessage): Parsed worker stdout or stderr message. + """ + lines.append(message.log) + print(message.log, flush=True) + + manager = SlurmManager(redis_host="bec-slurm.psi.ch", worker_path=worker_path) + try: + job = manager.submit_job( + script_path=script_path, + working_directory=working_directory, + callbacks={"log": log}, + job_options=JobOptions( + name="slurm-manager-redis-stdout-smoke", + tasks=1, + time_limit_minutes=1, + standard_output=f"{log_directory}/redis-stdout-%j.out", + standard_error=f"{log_directory}/redis-stdout-%j.err", + ), + ) + print(f"Submitted job {job.slurm_job_id}", flush=True) + try: + job.wait(timeout=90) + except TimeoutError: + job.cancel() + raise + finally: + manager.shutdown(timeout=65) + assert sorted(lines) == [f"stdout-test {second:02d}/10" for second in range(1, 11)] + print("All ten stdout lines received through Redis.") + + +if __name__ == "__main__": + main( + script_path=Path(__file__).with_name("print_for_10s.sh"), + working_directory="/das/home/your_username", + log_directory="/das/home/your_username/slurm_logs", + worker_path="/das/home/your_username/slurm-worker", + )