laco-submitit 1.0.0__py3-none-any.whl

laco/integrations/submitit/__init__.py

@@ -0,0 +1,38 @@
+"""SLURM/submitit integration for laco.
+
+Submits a laco task as a SLURM array job, one job per override combination.
+
+Examples
+--------
+::
+
+    import laco
+    import laco.integrations.submitit as laco_submitit
+
+    cfg = laco.load("configs/train.py")
+
+    executor = laco_submitit.executor(
+        folder="logs/slurm/%j",
+        slurm_partition="gpu",
+        slurm_gres="gpu:1",
+        timeout_min=120,
+    )
+
+    # Submit a sweep over lr and epochs:
+    jobs = laco_submitit.sweep(
+        executor,
+        cfg,
+        task_fn=train,
+        overrides=[
+            ["lr=1e-3", "epochs=10"],
+            ["lr=1e-4", "epochs=20"],
+        ],
+    )
+    results = [j.result() for j in jobs]
+"""
+
+from laco.integrations.submitit._core import executor as executor
+from laco.integrations.submitit._core import submit as submit
+from laco.integrations.submitit._core import sweep as sweep
+
+__all__ = ["executor", "submit", "sweep"]

laco/integrations/submitit/_core.py

@@ -0,0 +1,104 @@
+"""Implementation for laco-submitit."""
+
+from __future__ import annotations
+
+import typing
+
+if typing.TYPE_CHECKING:
+    import submitit
+    from omegaconf import DictConfig
+
+
+def executor(
+    folder: str = "logs/slurm/%j",
+    **slurm_kwargs: typing.Any,
+) -> submitit.AutoExecutor:
+    """Create a configured submitit AutoExecutor.
+
+    Parameters
+    ----------
+    folder : str
+        Log directory pattern. ``%j`` is replaced by the job ID.
+    **slurm_kwargs
+        SLURM parameters forwarded to ``executor.update_parameters()``
+        (e.g. ``slurm_partition``, ``slurm_gres``, ``timeout_min``, ``mem_gb``).
+
+    Returns
+    -------
+    submitit.AutoExecutor
+        Configured executor ready for job submission.
+    """
+    import submitit as _submitit
+
+    ex = _submitit.AutoExecutor(folder=folder)
+    if slurm_kwargs:
+        ex.update_parameters(**slurm_kwargs)
+    return ex
+
+
+def submit(
+    ex: submitit.AutoExecutor,
+    cfg: DictConfig,
+    task_fn: typing.Callable[[DictConfig], typing.Any],
+    *,
+    overrides: list[str] | None = None,
+) -> submitit.Job[typing.Any]:
+    """Submit a single laco task as a SLURM job.
+
+    Parameters
+    ----------
+    ex : submitit.AutoExecutor
+        Configured executor from :func:`executor`.
+    cfg : DictConfig
+        Base laco DictConfig.
+    task_fn : Callable
+        Function that accepts the resolved DictConfig.
+    overrides : list[str] | None
+        Hydra-style override strings applied before submission.
+
+    Returns
+    -------
+    submitit.Job
+        Job handle.
+    """
+    from laco._overrides import apply_overrides
+
+    resolved = apply_overrides(cfg, overrides) if overrides else cfg
+    return ex.submit(task_fn, resolved)
+
+
+def sweep(
+    ex: submitit.AutoExecutor,
+    cfg: DictConfig,
+    task_fn: typing.Callable[[DictConfig], typing.Any],
+    overrides: list[list[str]],
+) -> list[submitit.Job[typing.Any]]:
+    """Submit a sweep of laco tasks as SLURM array jobs.
+
+    Each entry in *overrides* produces one job with those overrides applied on
+    top of the base config.
+
+    Parameters
+    ----------
+    ex : submitit.AutoExecutor
+        Configured executor from :func:`executor`.
+    cfg : DictConfig
+        Base config shared across all sweep jobs.
+    task_fn : Callable
+        Function that accepts a resolved DictConfig.
+    overrides : list[list[str]]
+        List of override lists, one per sweep point.
+        E.g. ``[["lr=1e-3"], ["lr=1e-4"]]``.
+
+    Returns
+    -------
+    list[submitit.Job]
+        Job handles in the same order as *overrides*.
+    """
+    from copy import deepcopy
+
+    from laco._overrides import apply_overrides
+
+    configs = [apply_overrides(deepcopy(cfg), ov) for ov in overrides]
+    jobs = ex.map_array(task_fn, configs)
+    return list(jobs)

laco_submitit-1.0.0.dist-info/METADATA

@@ -0,0 +1,29 @@
+Metadata-Version: 2.4
+Name: laco-submitit
+Version: 1.0.0
+Summary: SLURM/submitit cluster job submission for laco.
+Author-email: Kurt Stolle <kurt@khws.io>
+Requires-Python: >=3.13
+Description-Content-Type: text/markdown
+Requires-Dist: laco>=1.0.0
+Requires-Dist: submitit>=1.5
+
+# Laco-Submitit
+
+SLURM/submitit cluster job submission for laco.
+
+Part of the [laco](https://github.com/khwstolle/laco) project — see the [root README](../../README.md) for an overview.
+
+## Installation
+
+```bash
+pip install laco-submitit
+```
+
+## Features
+
+`executor()`, `submit()`, `sweep()`
+
+## Usage
+
+See [`docs/index.md`](docs/index.md) for the full guide.

laco_submitit-1.0.0.dist-info/RECORD

@@ -0,0 +1,6 @@
+laco/integrations/submitit/__init__.py,sha256=gv9gz980IEhsKPK5ph-ZIzpSsEGoVcDCf54rZOeDFIU,932
+laco/integrations/submitit/_core.py,sha256=Zl1mNsKBqpDpWGNbN-MjWf3aVrdQKCBDMf4HziTcFa4,2779
+laco_submitit-1.0.0.dist-info/METADATA,sha256=9v-AIXgQ715D2WM720T32a4yc1MSuDL0j1WQ3O_uoFI,638
+laco_submitit-1.0.0.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
+laco_submitit-1.0.0.dist-info/top_level.txt,sha256=G2kLu09Aje44OkSqu-Tae3mjmTYhyRc2VrTyh3OmxFw,5
+laco_submitit-1.0.0.dist-info/RECORD,,

laco_submitit-1.0.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (82.0.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

laco_submitit-1.0.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+laco