src-snapshot 0.1.0__tar.gz

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.
@@ -0,0 +1,74 @@
1
+ name: CI
2
+
3
+ on:
4
+ push:
5
+ branches: [main]
6
+ tags: ["v*"]
7
+ pull_request:
8
+
9
+ jobs:
10
+ test:
11
+ name: Smoke test (Python ${{ matrix.python-version }}, ${{ matrix.os }})
12
+ runs-on: ${{ matrix.os }}
13
+ strategy:
14
+ fail-fast: false
15
+ matrix:
16
+ # Oldest and newest Python on Linux, newest only on Windows and macOS.
17
+ include:
18
+ - { os: ubuntu-latest, python-version: "3.10" }
19
+ - { os: ubuntu-latest, python-version: "3.13" }
20
+ - { os: windows-latest, python-version: "3.13" }
21
+ - { os: macos-latest, python-version: "3.13" }
22
+ steps:
23
+ - uses: actions/checkout@v4
24
+ - uses: actions/setup-python@v5
25
+ with:
26
+ python-version: ${{ matrix.python-version }}
27
+ - uses: astral-sh/setup-uv@v6
28
+
29
+ - name: Build
30
+ run: uv build
31
+
32
+ - name: Install wheel
33
+ shell: bash
34
+ run: python -m pip install dist/*.whl
35
+
36
+ - name: Snapshot from a script
37
+ shell: bash
38
+ run: |
39
+ python examples/create_zip.py
40
+ python -c "import zipfile; names = zipfile.ZipFile('code.zip').namelist(); print(names); assert '__main__.py' in names and 'metadata.json' in names"
41
+
42
+ - name: Snapshot without a main script (python -c)
43
+ shell: bash
44
+ run: |
45
+ python -c "from src_snapshot import create_snapshot; create_snapshot('interactive.zip', source_folders='src')"
46
+ python -c "import zipfile; names = zipfile.ZipFile('interactive.zip').namelist(); print(names); assert '__main__.py' not in names"
47
+
48
+ publish:
49
+ name: Publish to PyPI
50
+ needs: test
51
+ if: startsWith(github.ref, 'refs/tags/v')
52
+ runs-on: ubuntu-latest
53
+ environment:
54
+ name: pypi
55
+ url: https://pypi.org/p/src-snapshot
56
+ permissions:
57
+ id-token: write # required for PyPI trusted publishing
58
+ steps:
59
+ - uses: actions/checkout@v4
60
+ - uses: astral-sh/setup-uv@v6
61
+
62
+ - name: Check tag matches package version
63
+ run: |
64
+ version=$(uv version --short)
65
+ if [ "v$version" != "$GITHUB_REF_NAME" ]; then
66
+ echo "Tag $GITHUB_REF_NAME does not match pyproject version $version"
67
+ exit 1
68
+ fi
69
+
70
+ - name: Build
71
+ run: uv build
72
+
73
+ - name: Publish
74
+ run: uv publish --trusted-publishing always
@@ -0,0 +1,10 @@
1
+ # Python-generated files
2
+ __pycache__/
3
+ *.py[oc]
4
+ build/
5
+ dist/
6
+ wheels/
7
+ *.egg-info
8
+
9
+ # Virtual environments
10
+ .venv
@@ -0,0 +1 @@
1
+ 3.13
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 Dominik Woiwode
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,107 @@
1
+ Metadata-Version: 2.5
2
+ Name: src-snapshot
3
+ Version: 0.1.0
4
+ Summary: Lightweight utility to snapshot your source code into a zip archive for reproducible experiments.
5
+ Project-URL: Homepage, https://github.com/dwoiwode/src-snapshot
6
+ Project-URL: Repository, https://github.com/dwoiwode/src-snapshot
7
+ Project-URL: Issues, https://github.com/dwoiwode/src-snapshot/issues
8
+ Author: Dominik Woiwode
9
+ License-Expression: MIT
10
+ License-File: LICENSE
11
+ Keywords: experiments,provenance,reproducibility,research,slurm,snapshot
12
+ Classifier: Development Status :: 3 - Alpha
13
+ Classifier: Intended Audience :: Developers
14
+ Classifier: Intended Audience :: Science/Research
15
+ Classifier: Operating System :: OS Independent
16
+ Classifier: Programming Language :: Python :: 3
17
+ Classifier: Programming Language :: Python :: 3 :: Only
18
+ Classifier: Programming Language :: Python :: 3.10
19
+ Classifier: Programming Language :: Python :: 3.11
20
+ Classifier: Programming Language :: Python :: 3.12
21
+ Classifier: Programming Language :: Python :: 3.13
22
+ Classifier: Topic :: Scientific/Engineering
23
+ Classifier: Topic :: System :: Archiving
24
+ Requires-Python: >=3.10
25
+ Requires-Dist: psutil
26
+ Description-Content-Type: text/markdown
27
+
28
+ # src-snapshot
29
+
30
+ [![CI](https://github.com/dwoiwode/src-snapshot/actions/workflows/ci.yml/badge.svg)](https://github.com/dwoiwode/src-snapshot/actions/workflows/ci.yml)
31
+ [![PyPI](https://img.shields.io/pypi/v/src-snapshot)](https://pypi.org/project/src-snapshot/)
32
+ [![Python](https://img.shields.io/pypi/pyversions/src-snapshot)](https://pypi.org/project/src-snapshot/)
33
+ [![License: MIT](https://img.shields.io/badge/License-MIT-blue.svg)](LICENSE)
34
+
35
+ Lightweight utility to snapshot your source code into a zip archive for reproducible experiments.
36
+
37
+ Designed for researchers who want a simple, reliable way to store the exact code used for a run (and some metadata).
38
+
39
+ ## Installation
40
+
41
+ ```bash
42
+ pip install src-snapshot
43
+ ```
44
+
45
+ The package is installed as `src-snapshot` and imported as `src_snapshot`.
46
+
47
+ ## Usage
48
+ ```python
49
+ from src_snapshot import create_snapshot
50
+
51
+ create_snapshot("code.zip", source_folders="src")
52
+ ```
53
+
54
+ That’s it. A zip archive will be created containing your source files and optional metadata.
55
+
56
+ ## Features
57
+ - Copy selected source folders into a zip archive
58
+ - Filter files by suffix (default: .py, .yaml)
59
+ - Include the main entrypoint script (automatically or manual path)
60
+ - Store current git commit and diff
61
+ - Generate a metadata file with environment, hardware, SLURM and run information
62
+
63
+ ## API
64
+ ```python
65
+ def create_snapshot(
66
+ filepath: Path | str,
67
+ *,
68
+ source_folders: Optional[Path | str | Iterable[Path | str]] = None,
69
+ suffixes: Optional[str | tuple[str]] = (".py", ".yaml"),
70
+ add_main_file: bool | Path | str = True,
71
+ add_git_diff: bool = True,
72
+ add_metadata: bool = True,
73
+ )
74
+ ```
75
+
76
+ ### Arguments
77
+ - ``filepath``: Output zip file (must not exist yet)
78
+ - ``source_folders``: Folder(s) to include (default: current directory)
79
+ - ``suffixes``: File extensions to include
80
+ - ``add_main_file``: Include the main script as ``__main__.py``. ``True`` detects it automatically; pass a path to choose it explicitly.
81
+ In interactive sessions (Jupyter, REPL, ``python -c``) there is no main script, so it is skipped with a warning.
82
+ - ``add_git_diff``: Store uncommitted changes as ``git_diff.patch``
83
+ - ``add_metadata``: Write ``metadata.json`` (timestamp, environment, git info, installed packages)
84
+
85
+ ## Archive contents
86
+ ```
87
+ code.zip
88
+ ├── <source_folder>/... # your source files
89
+ ├── __main__.py # the entrypoint script
90
+ ├── git_diff.patch # uncommitted changes
91
+ └── metadata.json # run metadata
92
+ ```
93
+
94
+ ## Privacy
95
+ ``metadata.json`` contains the username, hostname, working directory, command line and the list of installed packages.
96
+ Keep this in mind before sharing snapshots publicly, or set ``add_metadata=False``.
97
+
98
+ ## Releasing
99
+ Releases are published to PyPI by GitHub Actions using [trusted publishing](https://docs.pypi.org/trusted-publishers/):
100
+
101
+ 1. Bump ``version`` in ``pyproject.toml``.
102
+ 2. Commit, then tag and push: ``git tag v0.1.0 && git push origin v0.1.0``.
103
+
104
+ The workflow runs the tests, checks that the tag matches the version, builds and uploads the package.
105
+
106
+ ## License
107
+ [MIT](LICENSE)
@@ -0,0 +1,80 @@
1
+ # src-snapshot
2
+
3
+ [![CI](https://github.com/dwoiwode/src-snapshot/actions/workflows/ci.yml/badge.svg)](https://github.com/dwoiwode/src-snapshot/actions/workflows/ci.yml)
4
+ [![PyPI](https://img.shields.io/pypi/v/src-snapshot)](https://pypi.org/project/src-snapshot/)
5
+ [![Python](https://img.shields.io/pypi/pyversions/src-snapshot)](https://pypi.org/project/src-snapshot/)
6
+ [![License: MIT](https://img.shields.io/badge/License-MIT-blue.svg)](LICENSE)
7
+
8
+ Lightweight utility to snapshot your source code into a zip archive for reproducible experiments.
9
+
10
+ Designed for researchers who want a simple, reliable way to store the exact code used for a run (and some metadata).
11
+
12
+ ## Installation
13
+
14
+ ```bash
15
+ pip install src-snapshot
16
+ ```
17
+
18
+ The package is installed as `src-snapshot` and imported as `src_snapshot`.
19
+
20
+ ## Usage
21
+ ```python
22
+ from src_snapshot import create_snapshot
23
+
24
+ create_snapshot("code.zip", source_folders="src")
25
+ ```
26
+
27
+ That’s it. A zip archive will be created containing your source files and optional metadata.
28
+
29
+ ## Features
30
+ - Copy selected source folders into a zip archive
31
+ - Filter files by suffix (default: .py, .yaml)
32
+ - Include the main entrypoint script (automatically or manual path)
33
+ - Store current git commit and diff
34
+ - Generate a metadata file with environment, hardware, SLURM and run information
35
+
36
+ ## API
37
+ ```python
38
+ def create_snapshot(
39
+ filepath: Path | str,
40
+ *,
41
+ source_folders: Optional[Path | str | Iterable[Path | str]] = None,
42
+ suffixes: Optional[str | tuple[str]] = (".py", ".yaml"),
43
+ add_main_file: bool | Path | str = True,
44
+ add_git_diff: bool = True,
45
+ add_metadata: bool = True,
46
+ )
47
+ ```
48
+
49
+ ### Arguments
50
+ - ``filepath``: Output zip file (must not exist yet)
51
+ - ``source_folders``: Folder(s) to include (default: current directory)
52
+ - ``suffixes``: File extensions to include
53
+ - ``add_main_file``: Include the main script as ``__main__.py``. ``True`` detects it automatically; pass a path to choose it explicitly.
54
+ In interactive sessions (Jupyter, REPL, ``python -c``) there is no main script, so it is skipped with a warning.
55
+ - ``add_git_diff``: Store uncommitted changes as ``git_diff.patch``
56
+ - ``add_metadata``: Write ``metadata.json`` (timestamp, environment, git info, installed packages)
57
+
58
+ ## Archive contents
59
+ ```
60
+ code.zip
61
+ ├── <source_folder>/... # your source files
62
+ ├── __main__.py # the entrypoint script
63
+ ├── git_diff.patch # uncommitted changes
64
+ └── metadata.json # run metadata
65
+ ```
66
+
67
+ ## Privacy
68
+ ``metadata.json`` contains the username, hostname, working directory, command line and the list of installed packages.
69
+ Keep this in mind before sharing snapshots publicly, or set ``add_metadata=False``.
70
+
71
+ ## Releasing
72
+ Releases are published to PyPI by GitHub Actions using [trusted publishing](https://docs.pypi.org/trusted-publishers/):
73
+
74
+ 1. Bump ``version`` in ``pyproject.toml``.
75
+ 2. Commit, then tag and push: ``git tag v0.1.0 && git push origin v0.1.0``.
76
+
77
+ The workflow runs the tests, checks that the tag matches the version, builds and uploads the package.
78
+
79
+ ## License
80
+ [MIT](LICENSE)
@@ -0,0 +1,6 @@
1
+ from pathlib import Path
2
+
3
+ from src_snapshot import create_snapshot
4
+
5
+ # Snapshot the package source next to this repository's examples folder
6
+ create_snapshot("code.zip", source_folders=Path(__file__).parent.parent / "src")
@@ -0,0 +1,39 @@
1
+ [build-system]
2
+ requires = ["hatchling"]
3
+ build-backend = "hatchling.build"
4
+
5
+ [project]
6
+ name = "src-snapshot"
7
+ version = "0.1.0"
8
+ description = "Lightweight utility to snapshot your source code into a zip archive for reproducible experiments."
9
+ readme = "README.md"
10
+ license = "MIT"
11
+ license-files = ["LICENSE"]
12
+ authors = [{ name = "Dominik Woiwode" }]
13
+ requires-python = ">=3.10"
14
+ dependencies = [
15
+ "psutil",
16
+ ]
17
+ keywords = ["reproducibility", "snapshot", "experiments", "research", "provenance", "slurm"]
18
+ classifiers = [
19
+ "Development Status :: 3 - Alpha",
20
+ "Intended Audience :: Developers",
21
+ "Intended Audience :: Science/Research",
22
+ "Operating System :: OS Independent",
23
+ "Programming Language :: Python :: 3",
24
+ "Programming Language :: Python :: 3 :: Only",
25
+ "Programming Language :: Python :: 3.10",
26
+ "Programming Language :: Python :: 3.11",
27
+ "Programming Language :: Python :: 3.12",
28
+ "Programming Language :: Python :: 3.13",
29
+ "Topic :: Scientific/Engineering",
30
+ "Topic :: System :: Archiving",
31
+ ]
32
+
33
+ [project.urls]
34
+ Homepage = "https://github.com/dwoiwode/src-snapshot"
35
+ Repository = "https://github.com/dwoiwode/src-snapshot"
36
+ Issues = "https://github.com/dwoiwode/src-snapshot/issues"
37
+
38
+ [tool.hatch.build.targets.wheel]
39
+ packages = ["src/src_snapshot"]
@@ -0,0 +1 @@
1
+ from .snapshot import create_snapshot
@@ -0,0 +1,90 @@
1
+ import datetime
2
+ import json
3
+ import os
4
+ import tempfile
5
+ import warnings
6
+ from pathlib import Path
7
+ from typing import Optional, Iterable
8
+ from zipfile import ZipFile, ZIP_DEFLATED
9
+
10
+ from src_snapshot.utils import git_diff, create_metadata, zip_file_folders, get_main_script_path
11
+
12
+
13
+ def create_snapshot(
14
+ filepath: Path | str,
15
+ *,
16
+ source_folders: Optional[Path | str | Iterable[Path | str]] = None,
17
+ suffixes: Optional[str | tuple[str]] = (".py", ".yaml"),
18
+ add_main_file: bool | Path | str = True,
19
+ add_git_diff: bool = True,
20
+ add_metadata: bool = True
21
+ ):
22
+ # Prepare source folders (default: current directory)
23
+ if source_folders is None:
24
+ source_folders = (Path.cwd(),)
25
+ elif isinstance(source_folders, (str, Path)):
26
+ source_folders = (source_folders,)
27
+ else:
28
+ source_folders = tuple(source_folders)
29
+
30
+ # Prepare output folder
31
+ filepath = Path(filepath)
32
+ if filepath.exists():
33
+ raise FileExistsError(f"Code already exists at {filepath.absolute()}!")
34
+
35
+ filepath.parent.mkdir(exist_ok=True, parents=True)
36
+
37
+ # Write to a temp file next to the target and rename at the end,
38
+ # so a failed run never leaves a half-written zip behind.
39
+ fd, tmp_name = tempfile.mkstemp(dir=filepath.parent, prefix=f".{filepath.name}.", suffix=".tmp")
40
+ os.close(fd)
41
+ tmp_path = Path(tmp_name)
42
+ try:
43
+ _write_zip(tmp_path, source_folders, suffixes, add_main_file, add_git_diff, add_metadata)
44
+ os.replace(tmp_path, filepath)
45
+ finally:
46
+ tmp_path.unlink(missing_ok=True)
47
+
48
+
49
+ def _write_zip(zip_path, source_folders, suffixes, add_main_file, add_git_diff, add_metadata):
50
+ with ZipFile(zip_path, 'w', compression=ZIP_DEFLATED) as zip_file:
51
+ n = 0
52
+ for source_folder in source_folders:
53
+ source_folder = Path(source_folder).resolve()
54
+ n += zip_file_folders(zip_file, source_folder, source_folder.name, suffixes=suffixes)
55
+
56
+ # Add main file
57
+ if add_main_file:
58
+ if isinstance(add_main_file, (Path, str)):
59
+ main_file = Path(add_main_file)
60
+ if not main_file.is_file():
61
+ raise FileNotFoundError(f"Main file not found: {main_file.absolute()}")
62
+ else:
63
+ main_file = get_main_script_path()
64
+
65
+ if main_file is None:
66
+ # e.g. Jupyter, REPL or `python -c`: nothing to copy
67
+ warnings.warn(
68
+ "Could not determine the main script (interactive session?); "
69
+ "skipping __main__.py. Pass add_main_file=<path> to set it explicitly.",
70
+ stacklevel=3,
71
+ )
72
+ else:
73
+ main_file_content = (
74
+ f"# Originally copied on {datetime.datetime.now():%Y-%m-%d %H:%M:%S} from {main_file.absolute()}\n\n"
75
+ f"{main_file.read_text(encoding='utf-8')}"
76
+ )
77
+ zip_file.writestr('__main__.py', main_file_content)
78
+ n += 1
79
+
80
+ # Add git diff
81
+ if add_git_diff:
82
+ n += 1
83
+ zip_file.writestr('git_diff.patch', git_diff())
84
+
85
+ # Add metadata
86
+ if add_metadata:
87
+ n += 1
88
+ meta = create_metadata()
89
+ meta["files_copied"] = n
90
+ zip_file.writestr('metadata.json', json.dumps(meta, indent=2))
@@ -0,0 +1,195 @@
1
+ import datetime
2
+ import getpass
3
+ import importlib.metadata
4
+ import inspect
5
+ import os
6
+ import platform
7
+ import socket
8
+ import subprocess
9
+ import sys
10
+ from pathlib import Path
11
+ from typing import Optional
12
+
13
+ import psutil
14
+
15
+
16
+ def dependencies():
17
+ """Installed distributions as sorted ``name==version`` strings (like ``pip freeze``)."""
18
+ packages = {}
19
+ for dist in importlib.metadata.distributions():
20
+ name = dist.metadata["Name"]
21
+ if name: # skip broken installs without metadata
22
+ packages.setdefault(name, dist.version)
23
+ return [f"{name}=={version}" for name, version in sorted(packages.items(), key=lambda kv: kv[0].lower())]
24
+
25
+
26
+ def environment_info():
27
+ info = {
28
+ "python": sys.version.replace("\n", " "),
29
+ "executable": sys.executable,
30
+ "platform": platform.platform(),
31
+ "machine": platform.machine(),
32
+ "processor": platform.processor(),
33
+ "hardware": {
34
+ "GPU": gpu_info(),
35
+ "RAM": ram_info()
36
+ }
37
+ }
38
+ return info
39
+
40
+
41
+ def git_info():
42
+ def _run(cmd):
43
+ try:
44
+ return subprocess.check_output(cmd, stderr=subprocess.DEVNULL).decode().strip()
45
+ except Exception:
46
+ return "unknown"
47
+
48
+ return {
49
+ "repo_url": _run(["git", "config", "--get", "remote.origin.url"]),
50
+ "commit": _run(["git", "rev-parse", "HEAD"]),
51
+ "short_commit": _run(["git", "rev-parse", "--short", "HEAD"]),
52
+ "branch": _run(["git", "rev-parse", "--abbrev-ref", "HEAD"]),
53
+ "dirty": bool(_run(["git", "status", "--porcelain"])),
54
+ "tag": _run(["git", "describe", "--tags", "--abbrev=0"]),
55
+ }
56
+
57
+ def git_diff():
58
+ try:
59
+ return subprocess.check_output(["git", "diff"], stderr=subprocess.DEVNULL).decode().strip()
60
+ except (OSError, subprocess.CalledProcessError):
61
+ # git missing or not inside a repository
62
+ return ""
63
+
64
+ def create_metadata():
65
+ now = datetime.datetime.now(datetime.timezone.utc).isoformat()
66
+ main_file = get_main_script_path()
67
+ meta = {
68
+ "timestamp_utc": now,
69
+ "provenance": git_info(),
70
+ "user": user_info(),
71
+ "execution": {
72
+ "entrypoint": f"{main_file.absolute() if main_file else 'unknown'}",
73
+ "sys.argv": " ".join(sys.orig_argv),
74
+ "cwd": str(Path.cwd()),
75
+ "hostname": socket.gethostname(),
76
+ "pid": os.getpid(),
77
+ },
78
+ "environment": environment_info(),
79
+ "slurm": slurm_info(),
80
+ "dependencies": dependencies(),
81
+ }
82
+ return meta
83
+
84
+ def zip_file_folders(
85
+ zip_file,
86
+ in_dir: Path,
87
+ archive_path: str,
88
+ *,
89
+ suffixes: Optional[str | tuple[str]] = None
90
+ ) -> int:
91
+ if in_dir.is_dir():
92
+ n = 0
93
+ for file_folder in in_dir.iterdir():
94
+ n += zip_file_folders(zip_file, file_folder, f"{archive_path}/{file_folder.name}", suffixes=suffixes)
95
+ return n
96
+ elif in_dir.is_file():
97
+ if isinstance(suffixes, str) and not in_dir.suffix == suffixes:
98
+ return 0
99
+
100
+ if isinstance(suffixes, tuple | list) and in_dir.suffix not in suffixes:
101
+ return 0
102
+
103
+ # Add file to zip
104
+ zip_file.write(in_dir, archive_path)
105
+ return 1
106
+
107
+ else:
108
+ raise TypeError(f"What is {in_dir} if not file or folder?")
109
+
110
+
111
+ def get_main_script_path() -> Optional[Path]:
112
+ """
113
+ This function gets the current frame and then recursively goes back to the calling function until
114
+ - no more calling function exists (-> reached most top level file)
115
+ - the __name__ of the current file is "__main__"
116
+ :return:
117
+ """
118
+
119
+ current_frame = inspect.currentframe()
120
+ while current_frame is not None:
121
+ # Check if the module is the main module
122
+ if current_frame.f_globals.get("__name__") == "__main__":
123
+ main_file = current_frame.f_globals.get("__file__")
124
+ # No __file__ in REPLs, Jupyter or `python -c`
125
+ if main_file is None:
126
+ return None
127
+ main_file = Path(main_file)
128
+ return main_file if main_file.is_file() else None
129
+ current_frame = current_frame.f_back
130
+ return None
131
+
132
+ def slurm_info():
133
+ def env(key):
134
+ return os.environ.get(key)
135
+
136
+ info = {
137
+ "job_id": env("SLURM_JOB_ID"),
138
+ "job_name": env("SLURM_JOB_NAME"),
139
+ "array_job_id": env("SLURM_ARRAY_JOB_ID"),
140
+ "array_task_id": env("SLURM_ARRAY_TASK_ID"),
141
+ "partition": env("SLURM_JOB_PARTITION"),
142
+ "node_list": env("SLURM_NODELIST"),
143
+ "num_nodes": env("SLURM_JOB_NUM_NODES"),
144
+ "num_tasks": env("SLURM_NTASKS"),
145
+ "cpus_per_task": env("SLURM_CPUS_PER_TASK"),
146
+ "cpus_on_node": env("SLURM_CPUS_ON_NODE"),
147
+ "gpus": env("SLURM_GPUS"),
148
+ "gpus_per_node": env("SLURM_GPUS_ON_NODE"),
149
+ "ntasks_per_node": env("SLURM_NTASKS_PER_NODE"),
150
+ "local_id": env("SLURM_LOCALID"),
151
+ "global_id": env("SLURM_PROCID"),
152
+ }
153
+
154
+ # remove keys that are not set
155
+ info = {k: v for k, v in info.items() if v is not None}
156
+
157
+ return info if info else None
158
+
159
+ def gpu_info():
160
+ try:
161
+ out = subprocess.check_output([
162
+ "nvidia-smi",
163
+ "--query-gpu=index,name,memory.total",
164
+ "--format=csv,noheader,nounits"
165
+ ]).decode().strip().splitlines()
166
+ except (OSError, subprocess.CalledProcessError):
167
+ # nvidia-smi missing or failing (e.g. no driver) -> no GPU info
168
+ return []
169
+
170
+ gpus = []
171
+ for line in out:
172
+ try:
173
+ idx, name, mem = [x.strip() for x in line.split(",")]
174
+ except ValueError:
175
+ continue
176
+ gpus.append({
177
+ "index": int(idx) if idx.isdigit() else idx,
178
+ "name": name,
179
+ "memory_mb": int(mem) if mem.isdigit() else None,
180
+ })
181
+ return gpus
182
+
183
+ def ram_info():
184
+ vm = psutil.virtual_memory()
185
+ return {
186
+ "total_bytes": vm.total,
187
+ "available_bytes": vm.available,
188
+ "used_bytes": vm.used,
189
+ }
190
+
191
+ def user_info():
192
+ return {
193
+ "username": getpass.getuser(),
194
+ "uid": os.getuid() if hasattr(os, "getuid") else None,
195
+ }