shardguard 1.0.0__tar.gz

shardguard-1.0.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Jonas Reiter
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

shardguard-1.0.0/PKG-INFO

@@ -0,0 +1,36 @@
+Metadata-Version: 2.4
+Name: shardguard
+Version: 1.0.0
+Summary: Cap the GPU memory of PyTorch at the GPU shards of a Slurm job
+Author: Jonas Reiter
+License-Expression: MIT
+Project-URL: Repository, https://github.com/jonaslreiter/shardguard
+Classifier: Programming Language :: Python :: 3
+Classifier: Operating System :: POSIX :: Linux
+Classifier: Topic :: System :: Distributed Computing
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Dynamic: license-file
+
+# shardguard
+
+Slurm can split a GPU into shards, but it does not limit the GPU memory of a job that holds shards, so the job can take memory that its co-tenants on the card were promised. shardguard caps the CUDA memory of PyTorch at the share of the GPU that the shards of the job stand for. Install it with
+
+```bash
+pip install shardguard
+```
+
+and import it in the entry point of your program, before anything initialises CUDA:
+
+```python
+import shardguard
+```
+
+Past the cap, PyTorch raises an out-of-memory error. Outside a Slurm job that holds shards, nothing happens.
+
+## How it works
+
+The share is the number of shards of the job over the number of shards one GPU of the node holds, which shardguard reads from `scontrol show node`. Setting `SHARDGUARD_SHARDS_PER_GPU` gives that number instead. The cap is `per_process_memory_fraction` in `PYTORCH_ALLOC_CONF` and `PYTORCH_CUDA_ALLOC_CONF`, together with `expandable_segments`, which keeps the memory PyTorch reserves close to what it allocates.
+
+The cap needs PyTorch 2.10 or later with its native allocator. PyTorch reads the setting when it initialises CUDA, so shardguard raises an error when it is imported after that. The cap applies to each process separately and does not cover the CUDA context, which takes a few hundred MB. shardguard also raises an error when the shards of the job are spread over several GPUs.

shardguard-1.0.0/README.md

@@ -0,0 +1,21 @@
+# shardguard
+
+Slurm can split a GPU into shards, but it does not limit the GPU memory of a job that holds shards, so the job can take memory that its co-tenants on the card were promised. shardguard caps the CUDA memory of PyTorch at the share of the GPU that the shards of the job stand for. Install it with
+
+```bash
+pip install shardguard
+```
+
+and import it in the entry point of your program, before anything initialises CUDA:
+
+```python
+import shardguard
+```
+
+Past the cap, PyTorch raises an out-of-memory error. Outside a Slurm job that holds shards, nothing happens.
+
+## How it works
+
+The share is the number of shards of the job over the number of shards one GPU of the node holds, which shardguard reads from `scontrol show node`. Setting `SHARDGUARD_SHARDS_PER_GPU` gives that number instead. The cap is `per_process_memory_fraction` in `PYTORCH_ALLOC_CONF` and `PYTORCH_CUDA_ALLOC_CONF`, together with `expandable_segments`, which keeps the memory PyTorch reserves close to what it allocates.
+
+The cap needs PyTorch 2.10 or later with its native allocator. PyTorch reads the setting when it initialises CUDA, so shardguard raises an error when it is imported after that. The cap applies to each process separately and does not cover the CUDA context, which takes a few hundred MB. shardguard also raises an error when the shards of the job are spread over several GPUs.

shardguard-1.0.0/pyproject.toml

@@ -0,0 +1,37 @@
+[build-system]
+requires = ["setuptools>=77.0.3"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "shardguard"
+version = "1.0.0"
+description = "Cap the GPU memory of PyTorch at the GPU shards of a Slurm job"
+readme = "README.md"
+requires-python = ">=3.10"
+license = "MIT"
+license-files = ["LICENSE"]
+authors = [{ name = "Jonas Reiter" }]
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "Operating System :: POSIX :: Linux",
+    "Topic :: System :: Distributed Computing",
+]
+
+[project.urls]
+Repository = "https://github.com/jonaslreiter/shardguard"
+
+[dependency-groups]
+dev = ["pytest>=8"]
+
+[tool.setuptools]
+package-dir = { "" = "src" }
+py-modules = ["shardguard"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+
+[tool.black]
+target-version = ["py310"]
+
+[tool.isort]
+profile = "black"

shardguard-1.0.0/setup.cfg

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

shardguard-1.0.0/src/shardguard.egg-info/PKG-INFO

@@ -0,0 +1,36 @@
+Metadata-Version: 2.4
+Name: shardguard
+Version: 1.0.0
+Summary: Cap the GPU memory of PyTorch at the GPU shards of a Slurm job
+Author: Jonas Reiter
+License-Expression: MIT
+Project-URL: Repository, https://github.com/jonaslreiter/shardguard
+Classifier: Programming Language :: Python :: 3
+Classifier: Operating System :: POSIX :: Linux
+Classifier: Topic :: System :: Distributed Computing
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Dynamic: license-file
+
+# shardguard
+
+Slurm can split a GPU into shards, but it does not limit the GPU memory of a job that holds shards, so the job can take memory that its co-tenants on the card were promised. shardguard caps the CUDA memory of PyTorch at the share of the GPU that the shards of the job stand for. Install it with
+
+```bash
+pip install shardguard
+```
+
+and import it in the entry point of your program, before anything initialises CUDA:
+
+```python
+import shardguard
+```
+
+Past the cap, PyTorch raises an out-of-memory error. Outside a Slurm job that holds shards, nothing happens.
+
+## How it works
+
+The share is the number of shards of the job over the number of shards one GPU of the node holds, which shardguard reads from `scontrol show node`. Setting `SHARDGUARD_SHARDS_PER_GPU` gives that number instead. The cap is `per_process_memory_fraction` in `PYTORCH_ALLOC_CONF` and `PYTORCH_CUDA_ALLOC_CONF`, together with `expandable_segments`, which keeps the memory PyTorch reserves close to what it allocates.
+
+The cap needs PyTorch 2.10 or later with its native allocator. PyTorch reads the setting when it initialises CUDA, so shardguard raises an error when it is imported after that. The cap applies to each process separately and does not cover the CUDA context, which takes a few hundred MB. shardguard also raises an error when the shards of the job are spread over several GPUs.

shardguard-1.0.0/src/shardguard.egg-info/SOURCES.txt

@@ -0,0 +1,9 @@
+LICENSE
+README.md
+pyproject.toml
+src/shardguard.py
+src/shardguard.egg-info/PKG-INFO
+src/shardguard.egg-info/SOURCES.txt
+src/shardguard.egg-info/dependency_links.txt
+src/shardguard.egg-info/top_level.txt
+tests/test_shardguard.py

shardguard-1.0.0/src/shardguard.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

shardguard-1.0.0/src/shardguard.egg-info/top_level.txt

@@ -0,0 +1 @@
+shardguard

shardguard-1.0.0/src/shardguard.py

@@ -0,0 +1,131 @@
+"""Importing shardguard caps the CUDA memory of PyTorch at the share of a GPU that the GPU shards of the Slurm job stand for.
+
+The cap is ``per_process_memory_fraction`` in the allocator configuration of PyTorch, set to the shards of the job over the shards one GPU of the node holds. ``expandable_segments`` keeps what PyTorch reserves close to what it allocates. PyTorch reads the configuration when it initialises CUDA, so shardguard must be imported before that, which it checks. Outside a job holding shards, it does nothing.
+"""
+
+from __future__ import annotations
+
+import math
+import os
+import re
+import subprocess
+import sys
+from importlib import metadata
+from typing import MutableMapping
+
+SHARDS_PER_GPU = "SHARDGUARD_SHARDS_PER_GPU"
+# PyTorch reads the first of these that is set, and releases before 2.9 only the second.
+ALLOC_CONF = ("PYTORCH_ALLOC_CONF", "PYTORCH_CUDA_ALLOC_CONF")
+FRACTION = re.compile(r"per_process_memory_fraction:([0-9.]+)")
+# The first release whose allocator configuration takes per_process_memory_fraction.
+OLDEST_TORCH = (2, 10)
+
+
+def shards_per_gpu(gres: str) -> float | None:
+    """Shards per GPU from the ``Gres`` of a node, e.g. 24 for ``gpu:A10:4,shard:A10:96``, or None if it lists no GPUs or no shards."""
+    counts = {"gpu": 0, "shard": 0}
+    for entry in re.sub(r"\([^)]*\)", "", gres).split(","):
+        fields = entry.split(":")
+        if fields[0] in counts and fields[-1].isdigit():
+            counts[fields[0]] += int(fields[-1])
+    if not counts["gpu"] or not counts["shard"]:
+        return None
+    return counts["shard"] / counts["gpu"]
+
+
+def node_gres(node: str) -> str:
+    """The ``Gres`` that ``scontrol show node`` reports for `node`."""
+    output = subprocess.run(
+        ["scontrol", "show", "node", node],
+        capture_output=True,
+        text=True,
+        check=True,
+        timeout=60,
+    ).stdout
+    match = re.search(r"\bGres=(\S*)", output)
+    return match.group(1) if match else ""
+
+
+def fraction(environ: MutableMapping[str, str] = os.environ) -> float | None:
+    """The share of one GPU that the shards of the job stand for, or None outside a job holding shards."""
+    shards = environ.get("SLURM_SHARDS_ON_NODE")
+    if not shards:
+        return None
+    if "," in environ.get("CUDA_VISIBLE_DEVICES", ""):
+        raise RuntimeError(
+            "shardguard: the shards of this job are on several GPUs, which one fraction per GPU cannot cap"
+        )
+    per_gpu = environ.get(SHARDS_PER_GPU)
+    if not per_gpu:
+        try:
+            per_gpu = shards_per_gpu(node_gres(environ["SLURMD_NODENAME"]))
+        except (KeyError, OSError, subprocess.SubprocessError):
+            per_gpu = None
+        if not per_gpu:
+            raise RuntimeError(
+                f"shardguard: cannot tell how many shards one GPU of this node holds, so set {SHARDS_PER_GPU}"
+            )
+        environ[SHARDS_PER_GPU] = f"{per_gpu:g}"
+    return min(1.0, int(shards) / float(per_gpu))
+
+
+def _rounded(share: float) -> float:
+    """`share` rounded down to four decimals."""
+    return math.floor(share * 1e4) / 1e4
+
+
+def with_cap(conf: str, share: float) -> str:
+    """An allocator configuration capped at `share`, with expandable segments unless it sets them itself. A lower cap it already has is kept."""
+    settings = [item for item in conf.split(",") if item]
+    capped = f"per_process_memory_fraction:{_rounded(share)}"
+    match = FRACTION.search(conf)
+    if match is None:
+        settings.append(capped)
+    elif float(match.group(1)) > share:
+        settings = [capped if FRACTION.fullmatch(item) else item for item in settings]
+    if not any(item.startswith("expandable_segments:") for item in settings):
+        settings.append("expandable_segments:True")
+    return ",".join(settings)
+
+
+def torch_version() -> tuple[int, int] | None:
+    """The major and minor version of the installed PyTorch, or None if it is not installed."""
+    try:
+        version = metadata.version("torch")
+    except metadata.PackageNotFoundError:
+        return None
+    major, minor = re.match(r"(\d+)\.(\d+)", version).groups()
+    return int(major), int(minor)
+
+
+def cap(environ: MutableMapping[str, str] = os.environ) -> float | None:
+    """Cap the CUDA memory of PyTorch at the shards of the job and return the share, or None outside a job holding shards."""
+    share = fraction(environ)
+    if share is None:
+        return None
+    version = torch_version()
+    if version is None or version < OLDEST_TORCH:
+        raise RuntimeError(
+            "shardguard: the cap needs PyTorch 2.10 or later, whose allocator configuration takes per_process_memory_fraction"
+        )
+    torch = sys.modules.get("torch")
+    if torch is not None and torch.cuda.is_initialized():
+        raise RuntimeError(
+            "shardguard: import it before CUDA is initialised, since PyTorch ignores the cap after that"
+        )
+    current = next((environ[name] for name in ALLOC_CONF if environ.get(name)), "")
+    if "cudaMallocAsync" in current:
+        raise RuntimeError(
+            "shardguard: the cudaMallocAsync backend ignores per_process_memory_fraction"
+        )
+    capped = with_cap(current, share)
+    for name in ALLOC_CONF:
+        environ[name] = capped
+    print(
+        f"shardguard: CUDA memory capped at {_rounded(share)} of a GPU, the share of {environ['SLURM_SHARDS_ON_NODE']} of its {environ[SHARDS_PER_GPU]} shards",
+        file=sys.stderr,
+    )
+    return share
+
+
+cap()

shardguard-1.0.0/tests/test_shardguard.py

@@ -0,0 +1,129 @@
+import sys
+import types
+
+import pytest
+
+import shardguard
+
+
+@pytest.fixture(autouse=True)
+def torch_2_14(monkeypatch):
+    monkeypatch.setattr(shardguard, "torch_version", lambda: (2, 14))
+
+
+@pytest.mark.parametrize(
+    "gres, expected",
+    [
+        ("gpu:A10:4,shard:A10:96", 24),
+        ("gpu:l40s:2(S:0-1),shard:l40s:96(S:0-1)", 48),
+        ("gpu:4", None),
+        ("(null)", None),
+    ],
+)
+def test_shards_per_gpu_is_read_off_the_gres(gres, expected):
+    assert shardguard.shards_per_gpu(gres) == expected
+
+
+def test_outside_a_shard_job_nothing_is_capped():
+    environ = {"SLURM_JOB_ID": "1"}
+    assert shardguard.cap(environ) is None
+    assert environ == {"SLURM_JOB_ID": "1"}
+
+
+def test_the_cap_is_the_jobs_share_of_a_gpu():
+    environ = {"SLURM_SHARDS_ON_NODE": "6", shardguard.SHARDS_PER_GPU: "24"}
+    assert shardguard.cap(environ) == 0.25
+    for name in shardguard.ALLOC_CONF:
+        assert (
+            environ[name] == "per_process_memory_fraction:0.25,expandable_segments:True"
+        )
+
+
+def test_the_shards_per_gpu_come_from_the_node_once(monkeypatch):
+    calls = []
+    monkeypatch.setattr(
+        shardguard,
+        "node_gres",
+        lambda node: calls.append(node) or "gpu:A10:4,shard:A10:96",
+    )
+    environ = {"SLURM_SHARDS_ON_NODE": "12", "SLURMD_NODENAME": "cn-1"}
+    assert shardguard.fraction(environ) == 0.5
+    assert shardguard.fraction(environ) == 0.5
+    assert calls == ["cn-1"]
+
+
+@pytest.mark.parametrize("gres", ["gpu:A10:4", OSError, KeyError])
+def test_a_node_that_cannot_say_its_shards_is_refused(monkeypatch, gres):
+    def node_gres(node):
+        if isinstance(gres, str):
+            return gres
+        raise gres
+
+    monkeypatch.setattr(shardguard, "node_gres", node_gres)
+    with pytest.raises(RuntimeError, match=shardguard.SHARDS_PER_GPU):
+        shardguard.fraction({"SLURM_SHARDS_ON_NODE": "4", "SLURMD_NODENAME": "cn-1"})
+
+
+def test_shards_on_several_gpus_are_refused():
+    with pytest.raises(RuntimeError, match="several GPUs"):
+        shardguard.fraction(
+            {"SLURM_SHARDS_ON_NODE": "4", "CUDA_VISIBLE_DEVICES": "0,1"}
+        )
+
+
+def settings(conf):
+    return dict(item.split(":", 1) for item in conf.split(","))
+
+
+def test_existing_settings_are_kept_and_a_lower_cap_wins():
+    assert settings(shardguard.with_cap("expandable_segments:False", 0.5)) == {
+        "expandable_segments": "False",
+        "per_process_memory_fraction": "0.5",
+    }
+    assert settings(shardguard.with_cap("per_process_memory_fraction:0.1", 0.5)) == {
+        "expandable_segments": "True",
+        "per_process_memory_fraction": "0.1",
+    }
+    assert (
+        settings(shardguard.with_cap("per_process_memory_fraction:0.9", 0.5))[
+            "per_process_memory_fraction"
+        ]
+        == "0.5"
+    )
+
+
+def test_the_cap_is_rounded_down_and_other_settings_are_left_alone():
+    conf = shardguard.with_cap("roundup_power2_divisions:[256:1,512:2]", 1 / 6)
+    assert (
+        conf
+        == "roundup_power2_divisions:[256:1,512:2],per_process_memory_fraction:0.1666,expandable_segments:True"
+    )
+
+
+def test_old_or_missing_pytorch_is_refused(monkeypatch):
+    for version in ((2, 9), None):
+        monkeypatch.setattr(shardguard, "torch_version", lambda: version)
+        with pytest.raises(RuntimeError, match="2.10"):
+            shardguard.cap(
+                {"SLURM_SHARDS_ON_NODE": "4", shardguard.SHARDS_PER_GPU: "24"}
+            )
+
+
+def test_the_async_allocator_is_refused():
+    with pytest.raises(RuntimeError, match="cudaMallocAsync"):
+        shardguard.cap(
+            {
+                "SLURM_SHARDS_ON_NODE": "4",
+                shardguard.SHARDS_PER_GPU: "24",
+                "PYTORCH_ALLOC_CONF": "backend:cudaMallocAsync",
+            }
+        )
+
+
+def test_importing_after_cuda_is_initialised_is_refused(monkeypatch):
+    torch = types.SimpleNamespace(
+        cuda=types.SimpleNamespace(is_initialized=lambda: True)
+    )
+    monkeypatch.setitem(sys.modules, "torch", torch)
+    with pytest.raises(RuntimeError, match="before CUDA"):
+        shardguard.cap({"SLURM_SHARDS_ON_NODE": "4", shardguard.SHARDS_PER_GPU: "24"})