Metadata-Version: 2.4
Name: prefect-submitit
Version: 0.1.2
Summary: A Prefect 3 TaskRunner that submits tasks to SLURM clusters via submitit
Project-URL: Homepage, https://github.com/dexterity-systems/prefect-submitit
Project-URL: Repository, https://github.com/dexterity-systems/prefect-submitit
Project-URL: Issues, https://github.com/dexterity-systems/prefect-submitit/issues
Maintainer: Dexterity Systems
License-Expression: BSD-3-Clause
License-File: LICENSE
Keywords: hpc,prefect,slurm,submitit,task-runner,workflow
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: BSD License
Classifier: Operating System :: POSIX :: Linux
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Topic :: Scientific/Engineering
Classifier: Topic :: System :: Distributed Computing
Requires-Python: >=3.12
Description-Content-Type: text/markdown

# prefect-submitit

A [Prefect 3](https://docs.prefect.io/) TaskRunner that submits tasks to
[SLURM](https://slurm.schedmd.com/) clusters via
[submitit](https://github.com/facebookincubator/submitit).

## Features

- **Single task submission** -- submit individual Prefect tasks as SLURM jobs
- **Job arrays** -- submit `task.map()` calls as SLURM job arrays with automatic
  chunking when the array exceeds cluster limits
- **Batched execution** -- group multiple items per SLURM job with
  `units_per_worker` to reduce scheduling overhead
- **Local mode** -- swap to local execution for testing without changing your flow
  code (`execution_mode="local"` or `SLURM_TASKRUNNER_BACKEND=local`)
- **Prefect UI integration** -- task run names include SLURM job IDs for easy
  cross-referencing with `squeue`/`sacct`

## Installation

```bash
pip install prefect-submitit
```

Or with conda (after the first conda-forge release):

```bash
conda install -c conda-forge prefect-submitit
```

## Quick Start

```python
from prefect import flow, task
from prefect_submitit import SlurmTaskRunner

@task
def add(x: int, y: int) -> int:
    return x + y

@flow(task_runner=SlurmTaskRunner(partition="cpu", time_limit="00:10:00"))
def my_flow():
    # Single task
    future = add.submit(1, 2)
    print(future.result())  # 3

    # Map over inputs (submitted as a SLURM job array)
    futures = add.map([1, 2, 3], [4, 5, 6])
    print([f.result() for f in futures])  # [5, 7, 9]

if __name__ == "__main__":
    my_flow()
```

## Configuration

```python
SlurmTaskRunner(
    partition="gpu",          # SLURM partition
    time_limit="04:00:00",    # Wall time (HH:MM:SS)
    mem_gb=16,                # Memory per job
    gpus_per_node=1,          # GPUs per job
    units_per_worker=10,      # Items per SLURM job (batched mode)
    execution_mode="slurm",   # "slurm" or "local"
    slurm_array_parallelism=100,  # Max concurrent array tasks
    log_folder="slurm_logs",  # Where submitit writes logs
    fail_on_error=True,       # Raise on SLURM job failure
)
```

Any additional keyword arguments are passed through to submitit (e.g.
`slurm_gres="gpu:a100:1"`).

## Local Testing

Set the environment variable to skip SLURM entirely:

```bash
export SLURM_TASKRUNNER_BACKEND=local
```

Or pass it directly:

```python
SlurmTaskRunner(execution_mode="local")
```

## Development

Requires [Pixi](https://pixi.sh):

```bash
pixi install
pixi run -e dev test
pixi run -e dev fmt
```

## License

BSD 3-Clause. See [LICENSE](LICENSE) for details.
