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.
@@ -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
+ ![Futuristic banner image.](extra/banner-image.png "Futuristic banner image.")
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
+ ![Futuristic banner image.](extra/banner-image.png "Futuristic banner image.")
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,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -0,0 +1 @@
1
+ from .slurm_pilot_executor import SlurmPilotExecutor, check_for_error