laco-submitit 1.0.0__tar.gz

laco_submitit-1.0.0/PKG-INFO

@@ -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/README.md

@@ -0,0 +1,19 @@
+# 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/pyproject.toml

@@ -0,0 +1,19 @@
+[build-system]
+requires = ["setuptools>=75", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "laco-submitit"
+version = "1.0.0"
+description = "SLURM/submitit cluster job submission for laco."
+readme = "README.md"
+requires-python = ">=3.13"
+authors = [{ name = "Kurt Stolle", email = "kurt@khws.io" }]
+dependencies = ["laco>=1.0.0", "submitit>=1.5"]
+
+[tool.setuptools.packages.find]
+where = ["sources"]
+namespaces = true
+
+[tool.uv.sources]
+laco = { workspace = true }

laco_submitit-1.0.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

laco_submitit-1.0.0/sources/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_submitit-1.0.0/sources/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/sources/laco_submitit.egg-info/PKG-INFO

@@ -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/sources/laco_submitit.egg-info/SOURCES.txt

@@ -0,0 +1,10 @@
+README.md
+pyproject.toml
+sources/laco/integrations/submitit/__init__.py
+sources/laco/integrations/submitit/_core.py
+sources/laco_submitit.egg-info/PKG-INFO
+sources/laco_submitit.egg-info/SOURCES.txt
+sources/laco_submitit.egg-info/dependency_links.txt
+sources/laco_submitit.egg-info/requires.txt
+sources/laco_submitit.egg-info/top_level.txt
+tests/test_laco_submitit.py

laco_submitit-1.0.0/sources/laco_submitit.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

laco_submitit-1.0.0/sources/laco_submitit.egg-info/requires.txt

@@ -0,0 +1,2 @@
+laco>=1.0.0
+submitit>=1.5

laco_submitit-1.0.0/sources/laco_submitit.egg-info/top_level.txt

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

laco_submitit-1.0.0/tests/test_laco_submitit.py

@@ -0,0 +1,89 @@
+"""Tests for laco-submitit."""
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+submitit = pytest.importorskip("submitit")
+
+
+def test_import():
+    import laco.integrations.submitit  # noqa: F401
+
+
+def test_public_api():
+    import laco.integrations.submitit as laco_submitit
+
+    assert set(laco_submitit.__all__) == {"executor", "submit", "sweep"}
+
+
+def test_executor_creates_auto_executor():
+    import laco.integrations.submitit as laco_submitit
+
+    ex = laco_submitit.executor(folder="/tmp/slurm/%j")
+
+    assert isinstance(ex, submitit.AutoExecutor)
+
+
+def test_executor_applies_slurm_kwargs():
+    import laco.integrations.submitit as laco_submitit
+
+    with patch("submitit.AutoExecutor") as mock_cls:
+        mock_ex = MagicMock()
+        mock_cls.return_value = mock_ex
+
+        laco_submitit.executor(
+            folder="/tmp/slurm/%j",
+            slurm_partition="gpu",
+            timeout_min=60,
+        )
+
+        mock_ex.update_parameters.assert_called_once_with(
+            slurm_partition="gpu", timeout_min=60
+        )
+
+
+def test_executor_no_kwargs_skips_update():
+    import laco.integrations.submitit as laco_submitit
+
+    with patch("submitit.AutoExecutor") as mock_cls:
+        mock_ex = MagicMock()
+        mock_cls.return_value = mock_ex
+
+        laco_submitit.executor(folder="/tmp/slurm/%j")
+
+        mock_ex.update_parameters.assert_not_called()
+
+
+def test_submit_without_overrides_passes_cfg_directly():
+    import laco.integrations.submitit as laco_submitit
+    from omegaconf import OmegaConf
+
+    cfg = OmegaConf.create({"lr": 0.001})
+    mock_ex = MagicMock()
+    task_fn = lambda c: c  # noqa: E731
+
+    laco_submitit.submit(mock_ex, cfg, task_fn=task_fn)
+
+    mock_ex.submit.assert_called_once_with(task_fn, cfg)
+
+
+def test_sweep_submits_one_job_per_override():
+    import laco.integrations.submitit as laco_submitit
+    from omegaconf import OmegaConf
+
+    cfg = OmegaConf.create({"lr": 0.001})
+    mock_ex = MagicMock()
+    mock_ex.map_array.return_value = [MagicMock(), MagicMock()]
+    task_fn = lambda c: c  # noqa: E731
+
+    overrides = [["lr=1e-3"], ["lr=1e-4"]]
+
+    with patch("laco._overrides.apply_overrides", side_effect=lambda c, _: c):
+        jobs = laco_submitit.sweep(mock_ex, cfg, task_fn=task_fn, overrides=overrides)
+
+    assert len(jobs) == 2
+    mock_ex.map_array.assert_called_once()
+    passed_task_fn, passed_configs = mock_ex.map_array.call_args[0]
+    assert passed_task_fn is task_fn
+    assert len(passed_configs) == 2