slurm-workflows 1.0.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- slurm_workflows-1.0.0/LICENSE +19 -0
- slurm_workflows-1.0.0/PKG-INFO +315 -0
- slurm_workflows-1.0.0/README.md +291 -0
- slurm_workflows-1.0.0/pyproject.toml +49 -0
- slurm_workflows-1.0.0/setup.cfg +4 -0
- slurm_workflows-1.0.0/src/slurm_workflows/__init__.py +1 -0
- slurm_workflows-1.0.0/src/slurm_workflows/ds_service.py +44 -0
- slurm_workflows-1.0.0/src/slurm_workflows/run_jupyter.py +57 -0
- slurm_workflows-1.0.0/src/slurm_workflows/slurm_pilot_executor.py +384 -0
- slurm_workflows-1.0.0/src/slurm_workflows/slurm_pilot_worker.py +168 -0
- slurm_workflows-1.0.0/src/slurm_workflows/slurm_utils.py +141 -0
- slurm_workflows-1.0.0/src/slurm_workflows/templates/__init__.py +143 -0
- slurm_workflows-1.0.0/src/slurm_workflows/templates/run_jupyter.jinja +18 -0
- slurm_workflows-1.0.0/src/slurm_workflows/templates/slurm_pilot.jinja +25 -0
- slurm_workflows-1.0.0/src/slurm_workflows/templates/slurm_utils.jinja +10 -0
- slurm_workflows-1.0.0/src/slurm_workflows/utils.py +72 -0
- slurm_workflows-1.0.0/src/slurm_workflows.egg-info/SOURCES.txt +14 -0
|
@@ -0,0 +1,19 @@
|
|
|
1
|
+
Copyright (C) 2026 Rector and Visitors of the University of Virginia
|
|
2
|
+
|
|
3
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
4
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
5
|
+
in the Software without restriction, including without limitation the rights
|
|
6
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
7
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
8
|
+
furnished to do so, subject to the following conditions:
|
|
9
|
+
|
|
10
|
+
The above copyright notice and this permission notice shall be included in all
|
|
11
|
+
copies or substantial portions of the Software.
|
|
12
|
+
|
|
13
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
14
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
15
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
16
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
17
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
18
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
19
|
+
SOFTWARE.
|
|
@@ -0,0 +1,315 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: slurm-workflows
|
|
3
|
+
Version: 1.0.0
|
|
4
|
+
Summary: HPC workflow helpers for Slurm clusters.
|
|
5
|
+
Author-email: Parantapa Bhattacharya <parantapa@virginia.edu>
|
|
6
|
+
Project-URL: Homepage, https://github.com/parantapa/slurm-hpc-workflows
|
|
7
|
+
Classifier: Programming Language :: Python :: 3
|
|
8
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
9
|
+
Classifier: Operating System :: OS Independent
|
|
10
|
+
Requires-Python: >=3.12
|
|
11
|
+
Description-Content-Type: text/markdown
|
|
12
|
+
License-File: LICENSE
|
|
13
|
+
Requires-Dist: click
|
|
14
|
+
Requires-Dist: tqdm
|
|
15
|
+
Requires-Dist: jinja2
|
|
16
|
+
Requires-Dist: json5
|
|
17
|
+
Requires-Dist: platformdirs
|
|
18
|
+
Requires-Dist: cloudpickle
|
|
19
|
+
Requires-Dist: typeguard
|
|
20
|
+
Requires-Dist: ds-service-client
|
|
21
|
+
Provides-Extra: test
|
|
22
|
+
Requires-Dist: pytest; extra == "test"
|
|
23
|
+
Dynamic: license-file
|
|
24
|
+
|
|
25
|
+
# slurm-workflows: HPC workflow helpers for Slurm clusters.
|
|
26
|
+
|
|
27
|
+

|
|
28
|
+
|
|
29
|
+
`slurm-workflows` lets you run Python functions on a Slurm cluster
|
|
30
|
+
without writing sbatch scripts by hand.
|
|
31
|
+
It provides a [`concurrent.futures`](https://docs.python.org/3/library/concurrent.futures.html)-inspired
|
|
32
|
+
interface that launches long-lived **pilot jobs** and dispatches tasks to them,
|
|
33
|
+
so Slurm's queueing latency is paid once per worker instead of once per task.
|
|
34
|
+
|
|
35
|
+
### Features
|
|
36
|
+
|
|
37
|
+
- **Pilot-job task execution** — pay the queue wait once,
|
|
38
|
+
then dispatch tasks at queue-latency speed.
|
|
39
|
+
- **Dynamic scaling** — grow or shrink a pool of workers at runtime
|
|
40
|
+
with `scale_workers`.
|
|
41
|
+
- **Stateful actors** — keep expensive per-worker state
|
|
42
|
+
(loaded models, DB connections) warm across many tasks.
|
|
43
|
+
- **Transparent serialization** — functions, arguments, and return values
|
|
44
|
+
are [cloudpickled](https://github.com/cloudpipe/cloudpickle),
|
|
45
|
+
so closures and lambdas work.
|
|
46
|
+
- **Non-fatal remote errors** — an exception on a worker
|
|
47
|
+
doesn't kill the driver script; it comes back as the task's result.
|
|
48
|
+
|
|
49
|
+
## Requirements
|
|
50
|
+
|
|
51
|
+
- Python >= 3.12
|
|
52
|
+
- Access to a Slurm cluster (`sbatch`, `squeue`, `scancel` on `PATH`)
|
|
53
|
+
- A running [`ds-service`](https://github.com/parantapa/ds-service) server,
|
|
54
|
+
reachable from the login node *and* the compute nodes.
|
|
55
|
+
The client library is installed as a dependency;
|
|
56
|
+
the server is a separate install.
|
|
57
|
+
|
|
58
|
+
## Installation
|
|
59
|
+
|
|
60
|
+
```sh
|
|
61
|
+
git clone https://github.com/parantapa/slurm-hpc-workflows.git
|
|
62
|
+
cd slurm-hpc-workflows
|
|
63
|
+
pip install -ve .
|
|
64
|
+
```
|
|
65
|
+
|
|
66
|
+
## Concepts
|
|
67
|
+
|
|
68
|
+
**Setup script.** Every worker sources a shell script
|
|
69
|
+
on its compute node before starting.
|
|
70
|
+
This is how your environment (`module load`, `conda activate`)
|
|
71
|
+
reaches the compute node — nothing is inherited from the login node.
|
|
72
|
+
|
|
73
|
+
**Worker group.** A named recipe for a worker:
|
|
74
|
+
sbatch arguments, setup script, optional actor class.
|
|
75
|
+
Defining a group does not launch workers.
|
|
76
|
+
`scale_workers` method is used to start/stop workers.
|
|
77
|
+
|
|
78
|
+
**Queue.** Tasks are submitted to a named queue,
|
|
79
|
+
and **a worker group pulls from the queue matching its own name**.
|
|
80
|
+
So `submit("gpu", ...)` is served by workers from the group named `gpu`.
|
|
81
|
+
|
|
82
|
+
## Quick start
|
|
83
|
+
|
|
84
|
+
Create a setup script, for example `setup.sh`:
|
|
85
|
+
|
|
86
|
+
```sh
|
|
87
|
+
module load gcc/14.2.0
|
|
88
|
+
conda activate my-env
|
|
89
|
+
```
|
|
90
|
+
|
|
91
|
+
Then run tasks against a pilot pool:
|
|
92
|
+
|
|
93
|
+
```python
|
|
94
|
+
from slurm_workflows import SlurmPilotExecutor, check_for_error
|
|
95
|
+
|
|
96
|
+
def square(x):
|
|
97
|
+
return x * x
|
|
98
|
+
|
|
99
|
+
DS_SERVICE_ADDRESS = "HOST-IP:5051"
|
|
100
|
+
|
|
101
|
+
# server_address points at your running ds-service instance.
|
|
102
|
+
executor = SlurmPilotExecutor(server_address=DS_SERVICE_ADDRESS)
|
|
103
|
+
|
|
104
|
+
# 1. Describe a kind of worker (nothing is launched yet).
|
|
105
|
+
executor.define_worker(
|
|
106
|
+
name="cpu",
|
|
107
|
+
sbatch_args=["-A my_alloc", "-p standard", "--cpus-per-task=4", "-t 01:00:00"],
|
|
108
|
+
setup_script="setup.sh",
|
|
109
|
+
)
|
|
110
|
+
|
|
111
|
+
# 2. Launch 4 pilot jobs of that kind.
|
|
112
|
+
executor.scale_workers("cpu", 4)
|
|
113
|
+
|
|
114
|
+
# 3. Submit tasks to a named queue; workers of that group pull from it.
|
|
115
|
+
tasks = [executor.submit("cpu", square, i) for i in range(100)]
|
|
116
|
+
|
|
117
|
+
# 4. Collect results as they complete (tqdm progress bar included).
|
|
118
|
+
for task in executor.as_completed(tasks, desc="squaring"):
|
|
119
|
+
... # task.output holds the return value
|
|
120
|
+
|
|
121
|
+
# Surface any tasks that raised on the worker.
|
|
122
|
+
for task in check_for_error(tasks):
|
|
123
|
+
print(task.task_id, task.output.error_id)
|
|
124
|
+
|
|
125
|
+
# 5. Cancel all pilot jobs when done.
|
|
126
|
+
executor.close()
|
|
127
|
+
```
|
|
128
|
+
|
|
129
|
+
`sbatch_args` are passed straight through to `sbatch`,
|
|
130
|
+
so any Slurm option works.
|
|
131
|
+
`submit` returns immediately with a `Task` handle;
|
|
132
|
+
`as_completed(tasks)` (or `wait(tasks)`) blocks until results are ready.
|
|
133
|
+
|
|
134
|
+
You don't have to wait for workers before submitting ---
|
|
135
|
+
tasks queue up and are picked up as pilot jobs start running.
|
|
136
|
+
|
|
137
|
+
### Stateful actors
|
|
138
|
+
|
|
139
|
+
To keep per-worker state warm across tasks,
|
|
140
|
+
register an actor class by its importable name.
|
|
141
|
+
Each worker instantiates it once at startup,
|
|
142
|
+
and you dispatch **method names** (as strings) instead of functions:
|
|
143
|
+
|
|
144
|
+
```python
|
|
145
|
+
# my_pkg/model.py
|
|
146
|
+
class Model:
|
|
147
|
+
def __init__(self):
|
|
148
|
+
self.model = load_expensive_model() # runs once per worker
|
|
149
|
+
|
|
150
|
+
def predict(self, x):
|
|
151
|
+
return self.model(x)
|
|
152
|
+
|
|
153
|
+
def close(self): # optional cleanup hook
|
|
154
|
+
self.model.release()
|
|
155
|
+
```
|
|
156
|
+
|
|
157
|
+
```python
|
|
158
|
+
executor.define_worker(
|
|
159
|
+
name="gpu",
|
|
160
|
+
sbatch_args=["-A my_alloc", "-p gpu", "--gres=gpu:1", "-t 02:00:00"],
|
|
161
|
+
setup_script="setup.sh",
|
|
162
|
+
actor_class_name="my_pkg.model.Model",
|
|
163
|
+
)
|
|
164
|
+
executor.scale_workers("gpu", 2)
|
|
165
|
+
|
|
166
|
+
tasks = [executor.submit("gpu", "predict", item) for item in dataset]
|
|
167
|
+
executor.wait(tasks)
|
|
168
|
+
```
|
|
169
|
+
|
|
170
|
+
The class must be importable on the compute node.
|
|
171
|
+
By default the executors's current working directory
|
|
172
|
+
is added to the workers' `sys.path`; add more with `python_paths=[...]`.
|
|
173
|
+
|
|
174
|
+
### One worker per job, or one per task
|
|
175
|
+
|
|
176
|
+
`is_batch_worker` controls how many worker processes each Slurm job starts:
|
|
177
|
+
|
|
178
|
+
| Setting | Script is run with | Workers per job |
|
|
179
|
+
| --- | --- | --- |
|
|
180
|
+
| `is_batch_worker=False` (default) | `srun` | one per Slurm task in the allocation |
|
|
181
|
+
| `is_batch_worker=True` | sourced directly | one, on the batch node |
|
|
182
|
+
|
|
183
|
+
So with the default, `--nodes=4 --ntasks-per-node=2`
|
|
184
|
+
gives you 8 worker processes from a single `scale_workers(..., 1)` call.
|
|
185
|
+
Use `is_batch_worker=True` when you want a single process
|
|
186
|
+
that owns the whole allocation (e.g. an MPI-style or whole-node job).
|
|
187
|
+
|
|
188
|
+
### Running the task-queue server
|
|
189
|
+
|
|
190
|
+
The executor and workers communicate only through a `ds-service` server
|
|
191
|
+
— they never talk to each other directly.
|
|
192
|
+
You can start one on the login node:
|
|
193
|
+
|
|
194
|
+
```python
|
|
195
|
+
from slurm_workflows.ds_service import DsService
|
|
196
|
+
|
|
197
|
+
with DsService(host="0.0.0.0", port=5051) as ds:
|
|
198
|
+
executor = SlurmPilotExecutor(server_address=ds.address)
|
|
199
|
+
...
|
|
200
|
+
```
|
|
201
|
+
|
|
202
|
+
The server must be reachable from the compute nodes,
|
|
203
|
+
so bind it to an address the workers can route to (`0.0.0.0` above),
|
|
204
|
+
and pass workers a routable host
|
|
205
|
+
— a login node's cluster-internal IP, not `localhost`.
|
|
206
|
+
|
|
207
|
+
## API reference
|
|
208
|
+
|
|
209
|
+
Import from the package root:
|
|
210
|
+
`from slurm_workflows import SlurmPilotExecutor, check_for_error`.
|
|
211
|
+
|
|
212
|
+
### `SlurmPilotExecutor(server_address, work_dir=None)`
|
|
213
|
+
|
|
214
|
+
`server_address` is the `host:port` of the `ds-service` server.
|
|
215
|
+
`work_dir` defaults to a timestamped directory under the platform cache dir
|
|
216
|
+
(`XDG_CACHE_HOME`-driven on Linux); generated scripts and all logs land there.
|
|
217
|
+
|
|
218
|
+
| Method | Purpose |
|
|
219
|
+
| --- | --- |
|
|
220
|
+
| `define_worker(name, sbatch_args, setup_script, ...)` | Register a worker group. Idempotent — redefining a group identically is a no-op, redefining it differently asserts. |
|
|
221
|
+
| `scale_workers(name, count)` | Submit or cancel pilot jobs so the group has `count` jobs. |
|
|
222
|
+
| `submit(queue, fn, *args, **kwargs) -> Task` | Enqueue a task. `queue` is a group name or a list of them; `fn` is a callable, or a method name (`str`) for actor workers. |
|
|
223
|
+
| `as_completed(tasks, desc=None, unit="task")` | Yield tasks as their results arrive, wrapped in a tqdm bar. |
|
|
224
|
+
| `wait(tasks, desc=None, unit="task")` | Same, but discards the iterator — just block until all are done. |
|
|
225
|
+
| `num_groups()` / `num_workers(detail=False)` | Counts of defined groups and submitted workers; `detail=True` returns a per-group dict. |
|
|
226
|
+
| `stop()` | Cancel all pilot jobs, keep the executor usable. |
|
|
227
|
+
| `close()` | Cancel all pilot jobs and close the queue-server connection. |
|
|
228
|
+
|
|
229
|
+
Remaining `define_worker` options:
|
|
230
|
+
|
|
231
|
+
| Argument | Default | Meaning |
|
|
232
|
+
| --- | --- | --- |
|
|
233
|
+
| `is_batch_worker` | `False` | See [above](#one-worker-per-job-or-one-per-task). |
|
|
234
|
+
| `actor_class_name` | `None` | Fully qualified class name to instantiate once per worker. |
|
|
235
|
+
| `python_paths` | `None` | Extra paths prepended to the workers' `sys.path`. |
|
|
236
|
+
| `add_cwd_to_python_path` | `True` | Also add the coordinator's cwd. |
|
|
237
|
+
| `worker_exe` | `"slurm-pilot-worker"` | Worker entry point, if you've wrapped or renamed it. |
|
|
238
|
+
|
|
239
|
+
### `Task`
|
|
240
|
+
|
|
241
|
+
`submit` returns a `Task` with `task_id`, `queue`, `priority`, `function`,
|
|
242
|
+
`input`, and `output`. `output` is a sentinel until the task completes; after
|
|
243
|
+
that it holds the return value — or a `RemoteExecutionError(error, error_id)`
|
|
244
|
+
if the worker raised.
|
|
245
|
+
|
|
246
|
+
### `check_for_error(tasks, verbose=True)`
|
|
247
|
+
|
|
248
|
+
Returns the subset of `tasks` whose `output` is a `RemoteExecutionError`,
|
|
249
|
+
printing each one's `error` and `error_id` unless `verbose=False`.
|
|
250
|
+
|
|
251
|
+
### Worker environment
|
|
252
|
+
|
|
253
|
+
Inside a task, these environment variables are set:
|
|
254
|
+
|
|
255
|
+
- `PILOT_WORKER_NAME` — e.g. `slurm_pilot_worker.cpu.0`
|
|
256
|
+
- `PILOT_WORKER_GROUP` — the group name
|
|
257
|
+
- `DS_SERVER_ADDRESS` — the queue server address
|
|
258
|
+
- plus the usual Slurm variables (`SLURM_JOB_ID`, …)
|
|
259
|
+
|
|
260
|
+
| Process | Runs on | Role |
|
|
261
|
+
| --- | --- | --- |
|
|
262
|
+
| Coordinator (`SlurmPilotExecutor`) | login node | defines worker groups, scales pilot jobs, submits tasks |
|
|
263
|
+
| `ds-service` | login node (or elsewhere) | holds tasks on named queues |
|
|
264
|
+
| Pilot workers | compute nodes | pull tasks, execute them, return results |
|
|
265
|
+
|
|
266
|
+
`scale_workers` renders a shell script and an sbatch wrapper
|
|
267
|
+
from Jinja templates and submits them.
|
|
268
|
+
Each job sources your setup script and launches `slurm-pilot-worker`,
|
|
269
|
+
which loops forever: fetch a task from its group's queue,
|
|
270
|
+
cloudpickle-load the function, run it, post the cloudpickled result back.
|
|
271
|
+
|
|
272
|
+
Two details worth knowing:
|
|
273
|
+
|
|
274
|
+
- **Exceptions are values.** A task that raises on a worker
|
|
275
|
+
does not propagate to the coordinator.
|
|
276
|
+
The worker catches it, logs the traceback under a generated `error_id`,
|
|
277
|
+
and returns a `RemoteExecutionError` as the task's `output`.
|
|
278
|
+
Always run `check_for_error` over a completed batch.
|
|
279
|
+
- **Submitting from inside a job works.** `sbatch` is invoked
|
|
280
|
+
with all `SLURM_*` / `PMI_*` / `SRUN_*` variables
|
|
281
|
+
stripped from the environment,
|
|
282
|
+
so a coordinator running inside a Slurm allocation
|
|
283
|
+
can still submit pilot jobs.
|
|
284
|
+
|
|
285
|
+
## Logs and troubleshooting
|
|
286
|
+
|
|
287
|
+
Everything for a run lives under the executor's `work_dir`
|
|
288
|
+
(printed as `executor.work_dir`):
|
|
289
|
+
|
|
290
|
+
| File | Contents |
|
|
291
|
+
| --- | --- |
|
|
292
|
+
| `coordinator.log` | Worker submission and cancellation from the executor's side |
|
|
293
|
+
| `<worker-name>.sh`, `<worker-name>.sbatch` | The generated scripts — read these first when a job dies immediately |
|
|
294
|
+
| `<worker-name>-<jobid>.out` | Slurm's stdout/stderr for the job, including setup-script failures |
|
|
295
|
+
| `<worker-name>-<jobid>-<host>-<pid>.log` | The worker process's own log: task-by-task progress and full tracebacks |
|
|
296
|
+
|
|
297
|
+
The `error_id` inside a `RemoteExecutionError` appears verbatim in the worker
|
|
298
|
+
log next to the traceback — grep for it across the work dir to find the failing
|
|
299
|
+
task's stack.
|
|
300
|
+
|
|
301
|
+
Common failure modes:
|
|
302
|
+
|
|
303
|
+
- **Tasks never complete, jobs are running.**
|
|
304
|
+
The queue name doesn't match a worker group name,
|
|
305
|
+
or the workers can't reach `ds-service` from the compute nodes.
|
|
306
|
+
Check the worker's `.log` file.
|
|
307
|
+
- **Jobs start and exit within seconds.**
|
|
308
|
+
The setup script failed. Check the `.out` file.
|
|
309
|
+
- **`ModuleNotFoundError` on a worker.**
|
|
310
|
+
The module isn't importable on the compute node
|
|
311
|
+
— add `python_paths=[...]` or install it into the environment the setup script activates.
|
|
312
|
+
|
|
313
|
+
## License
|
|
314
|
+
|
|
315
|
+
MIT — see [LICENSE](LICENSE).
|
|
@@ -0,0 +1,291 @@
|
|
|
1
|
+
# slurm-workflows: HPC workflow helpers for Slurm clusters.
|
|
2
|
+
|
|
3
|
+

|
|
4
|
+
|
|
5
|
+
`slurm-workflows` lets you run Python functions on a Slurm cluster
|
|
6
|
+
without writing sbatch scripts by hand.
|
|
7
|
+
It provides a [`concurrent.futures`](https://docs.python.org/3/library/concurrent.futures.html)-inspired
|
|
8
|
+
interface that launches long-lived **pilot jobs** and dispatches tasks to them,
|
|
9
|
+
so Slurm's queueing latency is paid once per worker instead of once per task.
|
|
10
|
+
|
|
11
|
+
### Features
|
|
12
|
+
|
|
13
|
+
- **Pilot-job task execution** — pay the queue wait once,
|
|
14
|
+
then dispatch tasks at queue-latency speed.
|
|
15
|
+
- **Dynamic scaling** — grow or shrink a pool of workers at runtime
|
|
16
|
+
with `scale_workers`.
|
|
17
|
+
- **Stateful actors** — keep expensive per-worker state
|
|
18
|
+
(loaded models, DB connections) warm across many tasks.
|
|
19
|
+
- **Transparent serialization** — functions, arguments, and return values
|
|
20
|
+
are [cloudpickled](https://github.com/cloudpipe/cloudpickle),
|
|
21
|
+
so closures and lambdas work.
|
|
22
|
+
- **Non-fatal remote errors** — an exception on a worker
|
|
23
|
+
doesn't kill the driver script; it comes back as the task's result.
|
|
24
|
+
|
|
25
|
+
## Requirements
|
|
26
|
+
|
|
27
|
+
- Python >= 3.12
|
|
28
|
+
- Access to a Slurm cluster (`sbatch`, `squeue`, `scancel` on `PATH`)
|
|
29
|
+
- A running [`ds-service`](https://github.com/parantapa/ds-service) server,
|
|
30
|
+
reachable from the login node *and* the compute nodes.
|
|
31
|
+
The client library is installed as a dependency;
|
|
32
|
+
the server is a separate install.
|
|
33
|
+
|
|
34
|
+
## Installation
|
|
35
|
+
|
|
36
|
+
```sh
|
|
37
|
+
git clone https://github.com/parantapa/slurm-hpc-workflows.git
|
|
38
|
+
cd slurm-hpc-workflows
|
|
39
|
+
pip install -ve .
|
|
40
|
+
```
|
|
41
|
+
|
|
42
|
+
## Concepts
|
|
43
|
+
|
|
44
|
+
**Setup script.** Every worker sources a shell script
|
|
45
|
+
on its compute node before starting.
|
|
46
|
+
This is how your environment (`module load`, `conda activate`)
|
|
47
|
+
reaches the compute node — nothing is inherited from the login node.
|
|
48
|
+
|
|
49
|
+
**Worker group.** A named recipe for a worker:
|
|
50
|
+
sbatch arguments, setup script, optional actor class.
|
|
51
|
+
Defining a group does not launch workers.
|
|
52
|
+
`scale_workers` method is used to start/stop workers.
|
|
53
|
+
|
|
54
|
+
**Queue.** Tasks are submitted to a named queue,
|
|
55
|
+
and **a worker group pulls from the queue matching its own name**.
|
|
56
|
+
So `submit("gpu", ...)` is served by workers from the group named `gpu`.
|
|
57
|
+
|
|
58
|
+
## Quick start
|
|
59
|
+
|
|
60
|
+
Create a setup script, for example `setup.sh`:
|
|
61
|
+
|
|
62
|
+
```sh
|
|
63
|
+
module load gcc/14.2.0
|
|
64
|
+
conda activate my-env
|
|
65
|
+
```
|
|
66
|
+
|
|
67
|
+
Then run tasks against a pilot pool:
|
|
68
|
+
|
|
69
|
+
```python
|
|
70
|
+
from slurm_workflows import SlurmPilotExecutor, check_for_error
|
|
71
|
+
|
|
72
|
+
def square(x):
|
|
73
|
+
return x * x
|
|
74
|
+
|
|
75
|
+
DS_SERVICE_ADDRESS = "HOST-IP:5051"
|
|
76
|
+
|
|
77
|
+
# server_address points at your running ds-service instance.
|
|
78
|
+
executor = SlurmPilotExecutor(server_address=DS_SERVICE_ADDRESS)
|
|
79
|
+
|
|
80
|
+
# 1. Describe a kind of worker (nothing is launched yet).
|
|
81
|
+
executor.define_worker(
|
|
82
|
+
name="cpu",
|
|
83
|
+
sbatch_args=["-A my_alloc", "-p standard", "--cpus-per-task=4", "-t 01:00:00"],
|
|
84
|
+
setup_script="setup.sh",
|
|
85
|
+
)
|
|
86
|
+
|
|
87
|
+
# 2. Launch 4 pilot jobs of that kind.
|
|
88
|
+
executor.scale_workers("cpu", 4)
|
|
89
|
+
|
|
90
|
+
# 3. Submit tasks to a named queue; workers of that group pull from it.
|
|
91
|
+
tasks = [executor.submit("cpu", square, i) for i in range(100)]
|
|
92
|
+
|
|
93
|
+
# 4. Collect results as they complete (tqdm progress bar included).
|
|
94
|
+
for task in executor.as_completed(tasks, desc="squaring"):
|
|
95
|
+
... # task.output holds the return value
|
|
96
|
+
|
|
97
|
+
# Surface any tasks that raised on the worker.
|
|
98
|
+
for task in check_for_error(tasks):
|
|
99
|
+
print(task.task_id, task.output.error_id)
|
|
100
|
+
|
|
101
|
+
# 5. Cancel all pilot jobs when done.
|
|
102
|
+
executor.close()
|
|
103
|
+
```
|
|
104
|
+
|
|
105
|
+
`sbatch_args` are passed straight through to `sbatch`,
|
|
106
|
+
so any Slurm option works.
|
|
107
|
+
`submit` returns immediately with a `Task` handle;
|
|
108
|
+
`as_completed(tasks)` (or `wait(tasks)`) blocks until results are ready.
|
|
109
|
+
|
|
110
|
+
You don't have to wait for workers before submitting ---
|
|
111
|
+
tasks queue up and are picked up as pilot jobs start running.
|
|
112
|
+
|
|
113
|
+
### Stateful actors
|
|
114
|
+
|
|
115
|
+
To keep per-worker state warm across tasks,
|
|
116
|
+
register an actor class by its importable name.
|
|
117
|
+
Each worker instantiates it once at startup,
|
|
118
|
+
and you dispatch **method names** (as strings) instead of functions:
|
|
119
|
+
|
|
120
|
+
```python
|
|
121
|
+
# my_pkg/model.py
|
|
122
|
+
class Model:
|
|
123
|
+
def __init__(self):
|
|
124
|
+
self.model = load_expensive_model() # runs once per worker
|
|
125
|
+
|
|
126
|
+
def predict(self, x):
|
|
127
|
+
return self.model(x)
|
|
128
|
+
|
|
129
|
+
def close(self): # optional cleanup hook
|
|
130
|
+
self.model.release()
|
|
131
|
+
```
|
|
132
|
+
|
|
133
|
+
```python
|
|
134
|
+
executor.define_worker(
|
|
135
|
+
name="gpu",
|
|
136
|
+
sbatch_args=["-A my_alloc", "-p gpu", "--gres=gpu:1", "-t 02:00:00"],
|
|
137
|
+
setup_script="setup.sh",
|
|
138
|
+
actor_class_name="my_pkg.model.Model",
|
|
139
|
+
)
|
|
140
|
+
executor.scale_workers("gpu", 2)
|
|
141
|
+
|
|
142
|
+
tasks = [executor.submit("gpu", "predict", item) for item in dataset]
|
|
143
|
+
executor.wait(tasks)
|
|
144
|
+
```
|
|
145
|
+
|
|
146
|
+
The class must be importable on the compute node.
|
|
147
|
+
By default the executors's current working directory
|
|
148
|
+
is added to the workers' `sys.path`; add more with `python_paths=[...]`.
|
|
149
|
+
|
|
150
|
+
### One worker per job, or one per task
|
|
151
|
+
|
|
152
|
+
`is_batch_worker` controls how many worker processes each Slurm job starts:
|
|
153
|
+
|
|
154
|
+
| Setting | Script is run with | Workers per job |
|
|
155
|
+
| --- | --- | --- |
|
|
156
|
+
| `is_batch_worker=False` (default) | `srun` | one per Slurm task in the allocation |
|
|
157
|
+
| `is_batch_worker=True` | sourced directly | one, on the batch node |
|
|
158
|
+
|
|
159
|
+
So with the default, `--nodes=4 --ntasks-per-node=2`
|
|
160
|
+
gives you 8 worker processes from a single `scale_workers(..., 1)` call.
|
|
161
|
+
Use `is_batch_worker=True` when you want a single process
|
|
162
|
+
that owns the whole allocation (e.g. an MPI-style or whole-node job).
|
|
163
|
+
|
|
164
|
+
### Running the task-queue server
|
|
165
|
+
|
|
166
|
+
The executor and workers communicate only through a `ds-service` server
|
|
167
|
+
— they never talk to each other directly.
|
|
168
|
+
You can start one on the login node:
|
|
169
|
+
|
|
170
|
+
```python
|
|
171
|
+
from slurm_workflows.ds_service import DsService
|
|
172
|
+
|
|
173
|
+
with DsService(host="0.0.0.0", port=5051) as ds:
|
|
174
|
+
executor = SlurmPilotExecutor(server_address=ds.address)
|
|
175
|
+
...
|
|
176
|
+
```
|
|
177
|
+
|
|
178
|
+
The server must be reachable from the compute nodes,
|
|
179
|
+
so bind it to an address the workers can route to (`0.0.0.0` above),
|
|
180
|
+
and pass workers a routable host
|
|
181
|
+
— a login node's cluster-internal IP, not `localhost`.
|
|
182
|
+
|
|
183
|
+
## API reference
|
|
184
|
+
|
|
185
|
+
Import from the package root:
|
|
186
|
+
`from slurm_workflows import SlurmPilotExecutor, check_for_error`.
|
|
187
|
+
|
|
188
|
+
### `SlurmPilotExecutor(server_address, work_dir=None)`
|
|
189
|
+
|
|
190
|
+
`server_address` is the `host:port` of the `ds-service` server.
|
|
191
|
+
`work_dir` defaults to a timestamped directory under the platform cache dir
|
|
192
|
+
(`XDG_CACHE_HOME`-driven on Linux); generated scripts and all logs land there.
|
|
193
|
+
|
|
194
|
+
| Method | Purpose |
|
|
195
|
+
| --- | --- |
|
|
196
|
+
| `define_worker(name, sbatch_args, setup_script, ...)` | Register a worker group. Idempotent — redefining a group identically is a no-op, redefining it differently asserts. |
|
|
197
|
+
| `scale_workers(name, count)` | Submit or cancel pilot jobs so the group has `count` jobs. |
|
|
198
|
+
| `submit(queue, fn, *args, **kwargs) -> Task` | Enqueue a task. `queue` is a group name or a list of them; `fn` is a callable, or a method name (`str`) for actor workers. |
|
|
199
|
+
| `as_completed(tasks, desc=None, unit="task")` | Yield tasks as their results arrive, wrapped in a tqdm bar. |
|
|
200
|
+
| `wait(tasks, desc=None, unit="task")` | Same, but discards the iterator — just block until all are done. |
|
|
201
|
+
| `num_groups()` / `num_workers(detail=False)` | Counts of defined groups and submitted workers; `detail=True` returns a per-group dict. |
|
|
202
|
+
| `stop()` | Cancel all pilot jobs, keep the executor usable. |
|
|
203
|
+
| `close()` | Cancel all pilot jobs and close the queue-server connection. |
|
|
204
|
+
|
|
205
|
+
Remaining `define_worker` options:
|
|
206
|
+
|
|
207
|
+
| Argument | Default | Meaning |
|
|
208
|
+
| --- | --- | --- |
|
|
209
|
+
| `is_batch_worker` | `False` | See [above](#one-worker-per-job-or-one-per-task). |
|
|
210
|
+
| `actor_class_name` | `None` | Fully qualified class name to instantiate once per worker. |
|
|
211
|
+
| `python_paths` | `None` | Extra paths prepended to the workers' `sys.path`. |
|
|
212
|
+
| `add_cwd_to_python_path` | `True` | Also add the coordinator's cwd. |
|
|
213
|
+
| `worker_exe` | `"slurm-pilot-worker"` | Worker entry point, if you've wrapped or renamed it. |
|
|
214
|
+
|
|
215
|
+
### `Task`
|
|
216
|
+
|
|
217
|
+
`submit` returns a `Task` with `task_id`, `queue`, `priority`, `function`,
|
|
218
|
+
`input`, and `output`. `output` is a sentinel until the task completes; after
|
|
219
|
+
that it holds the return value — or a `RemoteExecutionError(error, error_id)`
|
|
220
|
+
if the worker raised.
|
|
221
|
+
|
|
222
|
+
### `check_for_error(tasks, verbose=True)`
|
|
223
|
+
|
|
224
|
+
Returns the subset of `tasks` whose `output` is a `RemoteExecutionError`,
|
|
225
|
+
printing each one's `error` and `error_id` unless `verbose=False`.
|
|
226
|
+
|
|
227
|
+
### Worker environment
|
|
228
|
+
|
|
229
|
+
Inside a task, these environment variables are set:
|
|
230
|
+
|
|
231
|
+
- `PILOT_WORKER_NAME` — e.g. `slurm_pilot_worker.cpu.0`
|
|
232
|
+
- `PILOT_WORKER_GROUP` — the group name
|
|
233
|
+
- `DS_SERVER_ADDRESS` — the queue server address
|
|
234
|
+
- plus the usual Slurm variables (`SLURM_JOB_ID`, …)
|
|
235
|
+
|
|
236
|
+
| Process | Runs on | Role |
|
|
237
|
+
| --- | --- | --- |
|
|
238
|
+
| Coordinator (`SlurmPilotExecutor`) | login node | defines worker groups, scales pilot jobs, submits tasks |
|
|
239
|
+
| `ds-service` | login node (or elsewhere) | holds tasks on named queues |
|
|
240
|
+
| Pilot workers | compute nodes | pull tasks, execute them, return results |
|
|
241
|
+
|
|
242
|
+
`scale_workers` renders a shell script and an sbatch wrapper
|
|
243
|
+
from Jinja templates and submits them.
|
|
244
|
+
Each job sources your setup script and launches `slurm-pilot-worker`,
|
|
245
|
+
which loops forever: fetch a task from its group's queue,
|
|
246
|
+
cloudpickle-load the function, run it, post the cloudpickled result back.
|
|
247
|
+
|
|
248
|
+
Two details worth knowing:
|
|
249
|
+
|
|
250
|
+
- **Exceptions are values.** A task that raises on a worker
|
|
251
|
+
does not propagate to the coordinator.
|
|
252
|
+
The worker catches it, logs the traceback under a generated `error_id`,
|
|
253
|
+
and returns a `RemoteExecutionError` as the task's `output`.
|
|
254
|
+
Always run `check_for_error` over a completed batch.
|
|
255
|
+
- **Submitting from inside a job works.** `sbatch` is invoked
|
|
256
|
+
with all `SLURM_*` / `PMI_*` / `SRUN_*` variables
|
|
257
|
+
stripped from the environment,
|
|
258
|
+
so a coordinator running inside a Slurm allocation
|
|
259
|
+
can still submit pilot jobs.
|
|
260
|
+
|
|
261
|
+
## Logs and troubleshooting
|
|
262
|
+
|
|
263
|
+
Everything for a run lives under the executor's `work_dir`
|
|
264
|
+
(printed as `executor.work_dir`):
|
|
265
|
+
|
|
266
|
+
| File | Contents |
|
|
267
|
+
| --- | --- |
|
|
268
|
+
| `coordinator.log` | Worker submission and cancellation from the executor's side |
|
|
269
|
+
| `<worker-name>.sh`, `<worker-name>.sbatch` | The generated scripts — read these first when a job dies immediately |
|
|
270
|
+
| `<worker-name>-<jobid>.out` | Slurm's stdout/stderr for the job, including setup-script failures |
|
|
271
|
+
| `<worker-name>-<jobid>-<host>-<pid>.log` | The worker process's own log: task-by-task progress and full tracebacks |
|
|
272
|
+
|
|
273
|
+
The `error_id` inside a `RemoteExecutionError` appears verbatim in the worker
|
|
274
|
+
log next to the traceback — grep for it across the work dir to find the failing
|
|
275
|
+
task's stack.
|
|
276
|
+
|
|
277
|
+
Common failure modes:
|
|
278
|
+
|
|
279
|
+
- **Tasks never complete, jobs are running.**
|
|
280
|
+
The queue name doesn't match a worker group name,
|
|
281
|
+
or the workers can't reach `ds-service` from the compute nodes.
|
|
282
|
+
Check the worker's `.log` file.
|
|
283
|
+
- **Jobs start and exit within seconds.**
|
|
284
|
+
The setup script failed. Check the `.out` file.
|
|
285
|
+
- **`ModuleNotFoundError` on a worker.**
|
|
286
|
+
The module isn't importable on the compute node
|
|
287
|
+
— add `python_paths=[...]` or install it into the environment the setup script activates.
|
|
288
|
+
|
|
289
|
+
## License
|
|
290
|
+
|
|
291
|
+
MIT — see [LICENSE](LICENSE).
|
|
@@ -0,0 +1,49 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["setuptools", "setuptools_scm"]
|
|
3
|
+
build-backend = "setuptools.build_meta"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "slurm-workflows"
|
|
7
|
+
dynamic = ["version"]
|
|
8
|
+
authors = [
|
|
9
|
+
{ name="Parantapa Bhattacharya", email="parantapa@virginia.edu" },
|
|
10
|
+
]
|
|
11
|
+
description = "HPC workflow helpers for Slurm clusters."
|
|
12
|
+
readme = "README.md"
|
|
13
|
+
requires-python = ">=3.12"
|
|
14
|
+
classifiers = [
|
|
15
|
+
"Programming Language :: Python :: 3",
|
|
16
|
+
"License :: OSI Approved :: MIT License",
|
|
17
|
+
"Operating System :: OS Independent",
|
|
18
|
+
]
|
|
19
|
+
|
|
20
|
+
dependencies = [
|
|
21
|
+
"click",
|
|
22
|
+
"tqdm",
|
|
23
|
+
"jinja2",
|
|
24
|
+
"json5",
|
|
25
|
+
"platformdirs",
|
|
26
|
+
"cloudpickle",
|
|
27
|
+
"typeguard",
|
|
28
|
+
"ds-service-client"
|
|
29
|
+
]
|
|
30
|
+
|
|
31
|
+
[project.optional-dependencies]
|
|
32
|
+
test = ["pytest"]
|
|
33
|
+
|
|
34
|
+
[tool.setuptools.packages.find]
|
|
35
|
+
where = ["src"]
|
|
36
|
+
|
|
37
|
+
[tool.pytest.ini_options]
|
|
38
|
+
testpaths = ["tests"]
|
|
39
|
+
addopts = "-ra"
|
|
40
|
+
|
|
41
|
+
[tool.setuptools_scm]
|
|
42
|
+
fallback_version = "1.0.0-dev"
|
|
43
|
+
|
|
44
|
+
[project.urls]
|
|
45
|
+
"Homepage" = "https://github.com/parantapa/slurm-hpc-workflows"
|
|
46
|
+
|
|
47
|
+
[project.scripts]
|
|
48
|
+
"run-jupyter" = "slurm_workflows.run_jupyter:run_jupyter"
|
|
49
|
+
"slurm-pilot-worker" = "slurm_workflows.slurm_pilot_worker:slurm_pilot_worker"
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
from .slurm_pilot_executor import SlurmPilotExecutor, check_for_error
|