flyteplugins-slurm 0.0.0a0__py3-none-any.whl
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.
- flyteplugins/slurm/__init__.py +50 -0
- flyteplugins/slurm/connector.py +773 -0
- flyteplugins/slurm/script.py +365 -0
- flyteplugins/slurm/task.py +360 -0
- flyteplugins/slurm/transport.py +373 -0
- flyteplugins_slurm-0.0.0a0.dist-info/METADATA +536 -0
- flyteplugins_slurm-0.0.0a0.dist-info/RECORD +10 -0
- flyteplugins_slurm-0.0.0a0.dist-info/WHEEL +5 -0
- flyteplugins_slurm-0.0.0a0.dist-info/entry_points.txt +2 -0
- flyteplugins_slurm-0.0.0a0.dist-info/top_level.txt +1 -0
|
@@ -0,0 +1,50 @@
|
|
|
1
|
+
"""Slurm connector plugin for Flyte.
|
|
2
|
+
|
|
3
|
+
Run Flyte 2 tasks on an existing Slurm cluster. The connector submits jobs over
|
|
4
|
+
SSH to a login node, so it works with any Slurm installation, including
|
|
5
|
+
Soperator-managed clusters.
|
|
6
|
+
|
|
7
|
+
Two task types are provided:
|
|
8
|
+
|
|
9
|
+
- `slurm`: a Python task run inside its own container image via Pyxis/Enroot.
|
|
10
|
+
Typed inputs and outputs, caching, retries and error reporting work as they
|
|
11
|
+
do for a Kubernetes task. Remove `plugin_config` and the same task runs on
|
|
12
|
+
Kubernetes.
|
|
13
|
+
- `slurm_script`: an existing `sbatch` script submitted as-is. Phase and
|
|
14
|
+
logs only.
|
|
15
|
+
|
|
16
|
+
```python
|
|
17
|
+
import flyte
|
|
18
|
+
from flyteplugins.slurm import Slurm
|
|
19
|
+
|
|
20
|
+
env = flyte.TaskEnvironment(
|
|
21
|
+
name="train",
|
|
22
|
+
plugin_config=Slurm(
|
|
23
|
+
partition="main",
|
|
24
|
+
nodes=1,
|
|
25
|
+
gres="gpu:8",
|
|
26
|
+
time_limit="4:00:00",
|
|
27
|
+
host="login.slurm.example.com",
|
|
28
|
+
username="flyte",
|
|
29
|
+
ssh_private_key="slurm-ssh-key",
|
|
30
|
+
),
|
|
31
|
+
image=flyte.Image.from_debian_base().with_pip_packages("flyteplugins-slurm"),
|
|
32
|
+
)
|
|
33
|
+
|
|
34
|
+
@env.task
|
|
35
|
+
async def train(steps: int) -> float:
|
|
36
|
+
...
|
|
37
|
+
```
|
|
38
|
+
"""
|
|
39
|
+
|
|
40
|
+
from flyteplugins.slurm.connector import SlurmConnector, SlurmJobMetadata, SlurmScriptConnector
|
|
41
|
+
from flyteplugins.slurm.task import Slurm, SlurmFunctionTask, SlurmScriptTask
|
|
42
|
+
|
|
43
|
+
__all__ = [
|
|
44
|
+
"Slurm",
|
|
45
|
+
"SlurmConnector",
|
|
46
|
+
"SlurmFunctionTask",
|
|
47
|
+
"SlurmJobMetadata",
|
|
48
|
+
"SlurmScriptConnector",
|
|
49
|
+
"SlurmScriptTask",
|
|
50
|
+
]
|