vmn-exp-sdk 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.
Files changed (90) hide show
  1. vmn_exp_sdk-0.1.0/LICENSE.txt +21 -0
  2. vmn_exp_sdk-0.1.0/PKG-INFO +38 -0
  3. vmn_exp_sdk-0.1.0/README.md +20 -0
  4. vmn_exp_sdk-0.1.0/pyproject.toml +29 -0
  5. vmn_exp_sdk-0.1.0/setup.cfg +4 -0
  6. vmn_exp_sdk-0.1.0/src/vmn_exp/_base.py +101 -0
  7. vmn_exp_sdk-0.1.0/src/vmn_exp/core/__init__.py +0 -0
  8. vmn_exp_sdk-0.1.0/src/vmn_exp/core/background.py +65 -0
  9. vmn_exp_sdk-0.1.0/src/vmn_exp/core/best_effort.py +49 -0
  10. vmn_exp_sdk-0.1.0/src/vmn_exp/core/env.py +331 -0
  11. vmn_exp_sdk-0.1.0/src/vmn_exp/core/fold.py +211 -0
  12. vmn_exp_sdk-0.1.0/src/vmn_exp/core/from_snapshot.py +102 -0
  13. vmn_exp_sdk-0.1.0/src/vmn_exp/core/index.py +420 -0
  14. vmn_exp_sdk-0.1.0/src/vmn_exp/core/index_io_process.py +249 -0
  15. vmn_exp_sdk-0.1.0/src/vmn_exp/core/index_listing.py +86 -0
  16. vmn_exp_sdk-0.1.0/src/vmn_exp/core/index_logs.py +123 -0
  17. vmn_exp_sdk-0.1.0/src/vmn_exp/core/index_record.py +79 -0
  18. vmn_exp_sdk-0.1.0/src/vmn_exp/core/index_snapshot.py +195 -0
  19. vmn_exp_sdk-0.1.0/src/vmn_exp/core/index_store.py +181 -0
  20. vmn_exp_sdk-0.1.0/src/vmn_exp/core/index_sweep.py +216 -0
  21. vmn_exp_sdk-0.1.0/src/vmn_exp/core/index_workers.py +121 -0
  22. vmn_exp_sdk-0.1.0/src/vmn_exp/core/inputs.py +87 -0
  23. vmn_exp_sdk-0.1.0/src/vmn_exp/core/jsonl_tail.py +43 -0
  24. vmn_exp_sdk-0.1.0/src/vmn_exp/core/log.py +195 -0
  25. vmn_exp_sdk-0.1.0/src/vmn_exp/core/logfiles.py +98 -0
  26. vmn_exp_sdk-0.1.0/src/vmn_exp/core/manage.py +42 -0
  27. vmn_exp_sdk-0.1.0/src/vmn_exp/core/provenance.py +147 -0
  28. vmn_exp_sdk-0.1.0/src/vmn_exp/core/query.py +427 -0
  29. vmn_exp_sdk-0.1.0/src/vmn_exp/core/record_files.py +37 -0
  30. vmn_exp_sdk-0.1.0/src/vmn_exp/core/refs.py +106 -0
  31. vmn_exp_sdk-0.1.0/src/vmn_exp/core/reserved.py +18 -0
  32. vmn_exp_sdk-0.1.0/src/vmn_exp/core/resolve_ref.py +50 -0
  33. vmn_exp_sdk-0.1.0/src/vmn_exp/core/status.py +230 -0
  34. vmn_exp_sdk-0.1.0/src/vmn_exp/core/storage_resolve.py +113 -0
  35. vmn_exp_sdk-0.1.0/src/vmn_exp/core/tree.py +285 -0
  36. vmn_exp_sdk-0.1.0/src/vmn_exp/core/values.py +82 -0
  37. vmn_exp_sdk-0.1.0/src/vmn_exp/core/writer.py +370 -0
  38. vmn_exp_sdk-0.1.0/src/vmn_exp/integrations/__init__.py +0 -0
  39. vmn_exp_sdk-0.1.0/src/vmn_exp/integrations/hf.py +146 -0
  40. vmn_exp_sdk-0.1.0/src/vmn_exp/integrations/hf_core.py +170 -0
  41. vmn_exp_sdk-0.1.0/src/vmn_exp/integrations/optuna_study.py +207 -0
  42. vmn_exp_sdk-0.1.0/src/vmn_exp/integrations/ray_tune.py +271 -0
  43. vmn_exp_sdk-0.1.0/src/vmn_exp/registry/__init__.py +0 -0
  44. vmn_exp_sdk-0.1.0/src/vmn_exp/registry/cli.py +399 -0
  45. vmn_exp_sdk-0.1.0/src/vmn_exp/registry/fold.py +137 -0
  46. vmn_exp_sdk-0.1.0/src/vmn_exp/registry/log.py +194 -0
  47. vmn_exp_sdk-0.1.0/src/vmn_exp/registry/names.py +101 -0
  48. vmn_exp_sdk-0.1.0/src/vmn_exp/registry/store.py +143 -0
  49. vmn_exp_sdk-0.1.0/src/vmn_exp/registry/view.py +149 -0
  50. vmn_exp_sdk-0.1.0/src/vmn_exp/sdk/__init__.py +66 -0
  51. vmn_exp_sdk-0.1.0/src/vmn_exp/sdk/autolog.py +360 -0
  52. vmn_exp_sdk-0.1.0/src/vmn_exp/sdk/autolog_adapter.py +257 -0
  53. vmn_exp_sdk-0.1.0/src/vmn_exp/sdk/autolog_hf.py +87 -0
  54. vmn_exp_sdk-0.1.0/src/vmn_exp/sdk/autolog_keras.py +155 -0
  55. vmn_exp_sdk-0.1.0/src/vmn_exp/sdk/autolog_lightning.py +87 -0
  56. vmn_exp_sdk-0.1.0/src/vmn_exp/sdk/autolog_sklearn.py +50 -0
  57. vmn_exp_sdk-0.1.0/src/vmn_exp/sdk/context.py +117 -0
  58. vmn_exp_sdk-0.1.0/src/vmn_exp/sdk/create.py +186 -0
  59. vmn_exp_sdk-0.1.0/src/vmn_exp/sdk/heartbeat.py +74 -0
  60. vmn_exp_sdk-0.1.0/src/vmn_exp/sdk/import_hooks.py +115 -0
  61. vmn_exp_sdk-0.1.0/src/vmn_exp/sdk/log_buffer.py +99 -0
  62. vmn_exp_sdk-0.1.0/src/vmn_exp/sdk/manage.py +50 -0
  63. vmn_exp_sdk-0.1.0/src/vmn_exp/sdk/models.py +231 -0
  64. vmn_exp_sdk-0.1.0/src/vmn_exp/sdk/ranks.py +89 -0
  65. vmn_exp_sdk-0.1.0/src/vmn_exp/sdk/reader.py +229 -0
  66. vmn_exp_sdk-0.1.0/src/vmn_exp/sdk/resume.py +75 -0
  67. vmn_exp_sdk-0.1.0/src/vmn_exp/sdk/run.py +535 -0
  68. vmn_exp_sdk-0.1.0/src/vmn_exp/sdk/run_artifacts.py +88 -0
  69. vmn_exp_sdk-0.1.0/src/vmn_exp/sdk/signals.py +75 -0
  70. vmn_exp_sdk-0.1.0/src/vmn_exp/sdk/state_publisher.py +65 -0
  71. vmn_exp_sdk-0.1.0/src/vmn_exp/sdk/sysmetrics.py +349 -0
  72. vmn_exp_sdk-0.1.0/src/vmn_exp/storage/__init__.py +0 -0
  73. vmn_exp_sdk-0.1.0/src/vmn_exp/storage/base.py +154 -0
  74. vmn_exp_sdk-0.1.0/src/vmn_exp/storage/buffered.py +166 -0
  75. vmn_exp_sdk-0.1.0/src/vmn_exp/storage/cached.py +326 -0
  76. vmn_exp_sdk-0.1.0/src/vmn_exp/storage/cached_logs.py +226 -0
  77. vmn_exp_sdk-0.1.0/src/vmn_exp/storage/files.py +225 -0
  78. vmn_exp_sdk-0.1.0/src/vmn_exp/storage/index_cache_dir.py +87 -0
  79. vmn_exp_sdk-0.1.0/src/vmn_exp/storage/listing.py +105 -0
  80. vmn_exp_sdk-0.1.0/src/vmn_exp/storage/local.py +404 -0
  81. vmn_exp_sdk-0.1.0/src/vmn_exp/storage/s3.py +174 -0
  82. vmn_exp_sdk-0.1.0/src/vmn_exp/storage/s3_base.py +147 -0
  83. vmn_exp_sdk-0.1.0/src/vmn_exp/storage/s3_listing.py +245 -0
  84. vmn_exp_sdk-0.1.0/src/vmn_exp/storage/s3_logs.py +174 -0
  85. vmn_exp_sdk-0.1.0/src/vmn_exp/storage/s3_records.py +121 -0
  86. vmn_exp_sdk-0.1.0/src/vmn_exp_sdk.egg-info/PKG-INFO +38 -0
  87. vmn_exp_sdk-0.1.0/src/vmn_exp_sdk.egg-info/SOURCES.txt +88 -0
  88. vmn_exp_sdk-0.1.0/src/vmn_exp_sdk.egg-info/dependency_links.txt +1 -0
  89. vmn_exp_sdk-0.1.0/src/vmn_exp_sdk.egg-info/requires.txt +8 -0
  90. vmn_exp_sdk-0.1.0/src/vmn_exp_sdk.egg-info/top_level.txt +1 -0
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2019-2026 Pavel Rogovoy
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,38 @@
1
+ Metadata-Version: 2.4
2
+ Name: vmn-exp-sdk
3
+ Version: 0.1.0
4
+ Summary: vmn experiment SDK: record runs and metrics from jobs, git-free
5
+ Author-email: Pavel Rogovoy <p.rogovoy@gmail.com>
6
+ License: MIT
7
+ Project-URL: Homepage, https://github.com/progovoy/vmn
8
+ Requires-Python: >=3.8
9
+ Description-Content-Type: text/markdown
10
+ License-File: LICENSE.txt
11
+ Requires-Dist: PyYAML>=5.4.1
12
+ Requires-Dist: filelock>=3.2.0
13
+ Provides-Extra: s3
14
+ Requires-Dist: boto3; extra == "s3"
15
+ Provides-Extra: sysmetrics
16
+ Requires-Dist: psutil>=5.9; extra == "sysmetrics"
17
+ Dynamic: license-file
18
+
19
+ # vmn-exp-sdk
20
+
21
+ The metrics writer of the [vmn](https://github.com/progovoy/vmn) experiment
22
+ platform, for jobs: `start_run()`, `log_metric`/`log_params`/artifacts,
23
+ autologging and framework integrations — with no git and no vmn installed.
24
+
25
+ ```sh
26
+ pip install vmn-exp-sdk # add [s3] to record to a bucket
27
+ ```
28
+
29
+ ```python
30
+ from vmn_exp.sdk import start_run
31
+
32
+ with start_run("my_app") as run: # VMN_SNAPSHOT_METADATA + VMN_EXPERIMENT_DIR/BUCKET
33
+ run.log_metric("loss", 0.12)
34
+ ```
35
+
36
+ Creating runs from a git checkout, the `vmn-exp` CLI and the dashboard live in
37
+ [`vmn-exp`](https://pypi.org/project/vmn-exp/). See
38
+ [docs/sdk.md](https://github.com/progovoy/vmn/blob/master/docs/sdk.md).
@@ -0,0 +1,20 @@
1
+ # vmn-exp-sdk
2
+
3
+ The metrics writer of the [vmn](https://github.com/progovoy/vmn) experiment
4
+ platform, for jobs: `start_run()`, `log_metric`/`log_params`/artifacts,
5
+ autologging and framework integrations — with no git and no vmn installed.
6
+
7
+ ```sh
8
+ pip install vmn-exp-sdk # add [s3] to record to a bucket
9
+ ```
10
+
11
+ ```python
12
+ from vmn_exp.sdk import start_run
13
+
14
+ with start_run("my_app") as run: # VMN_SNAPSHOT_METADATA + VMN_EXPERIMENT_DIR/BUCKET
15
+ run.log_metric("loss", 0.12)
16
+ ```
17
+
18
+ Creating runs from a git checkout, the `vmn-exp` CLI and the dashboard live in
19
+ [`vmn-exp`](https://pypi.org/project/vmn-exp/). See
20
+ [docs/sdk.md](https://github.com/progovoy/vmn/blob/master/docs/sdk.md).
@@ -0,0 +1,29 @@
1
+ [build-system]
2
+ requires = ["setuptools>=64"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [project]
6
+ name = "vmn-exp-sdk"
7
+ version = "0.1.0" # written by `vmn stamp vmn_exp`, shared with vmn-exp
8
+ description = "vmn experiment SDK: record runs and metrics from jobs, git-free"
9
+ readme = "README.md"
10
+ license = { text = "MIT" }
11
+ authors = [{ name = "Pavel Rogovoy", email = "p.rogovoy@gmail.com" }]
12
+ requires-python = ">=3.8"
13
+ dependencies = [
14
+ "PyYAML>=5.4.1",
15
+ "filelock>=3.2.0",
16
+ ]
17
+
18
+ [project.optional-dependencies]
19
+ s3 = ["boto3"]
20
+ # GPU metrics also need pynvml, which stays out of every extra: it is useless
21
+ # without a driver and would break installs on CPU-only hosts.
22
+ sysmetrics = ["psutil>=5.9"]
23
+
24
+ [project.urls]
25
+ Homepage = "https://github.com/progovoy/vmn"
26
+
27
+ [tool.setuptools.packages.find]
28
+ where = ["src"]
29
+ namespaces = true # vmn_exp is shared with vmn-exp (PEP 420)
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -0,0 +1,101 @@
1
+ """Helpers the job SDK needs from vmn, copied so vmn-exp-sdk does not depend on
2
+ vmn. tests/test_exp_base.py keeps each copy in step with version_stamp.core."""
3
+ import datetime
4
+ import hashlib
5
+ import logging
6
+ import os
7
+
8
+ import yaml
9
+ from filelock import FileLock
10
+
11
+ # The same logger the vmn CLI configures, so both sides share its handlers.
12
+ VMN_LOGGER = logging.getLogger("vmn")
13
+
14
+ # libyaml's loader is ~10x faster than the pure-Python one, and experiment
15
+ # listings parse one metadata.yml (and one run_state.yml) per run.
16
+ _FAST_SAFE_LOADER = getattr(yaml, "CSafeLoader", yaml.SafeLoader)
17
+
18
+ _UNSAFE_IN_COMPONENT = ("/", "\\", os.sep, "\0", "..")
19
+
20
+
21
+ def ensure_logger():
22
+ return VMN_LOGGER
23
+
24
+
25
+ def yaml_safe_load(stream_or_text):
26
+ """``yaml.safe_load`` semantics, on the C loader when libyaml is present."""
27
+ return yaml.load(stream_or_text, Loader=_FAST_SAFE_LOADER)
28
+
29
+
30
+ def parse_record_metadata(raw):
31
+ """A snapshot/experiment record's ``metadata.yml`` dict, or None.
32
+
33
+ None for a missing or unparsable file and for legacy ``create_snapshots``
34
+ verinfo files, which share the tree but carry no ``verstr``.
35
+ """
36
+ if not raw:
37
+ return None
38
+ try:
39
+ meta = yaml_safe_load(raw)
40
+ except yaml.YAMLError:
41
+ return None
42
+ return meta if isinstance(meta, dict) and "verstr" in meta else None
43
+
44
+
45
+ def valid_path_component(name):
46
+ """Whether *name* (a verstr, artifact name, ...) stays one path component."""
47
+ return (
48
+ bool(name)
49
+ and name != "."
50
+ and not any(bad in name for bad in _UNSAFE_IN_COMPONENT)
51
+ )
52
+
53
+
54
+ def valid_app_path(app_name):
55
+ """Whether *app_name* is a relative path of real components (``root/svc``)."""
56
+ return bool(app_name) and all(
57
+ valid_path_component(part) for part in app_name.split("/")
58
+ )
59
+
60
+
61
+ def now_iso():
62
+ """UTC now as the ``...Z`` ISO string every vmn record is timestamped with."""
63
+ return (
64
+ datetime.datetime.now(datetime.timezone.utc).isoformat().replace("+00:00", "Z")
65
+ )
66
+
67
+
68
+ def sha256_file(abs_path):
69
+ h = hashlib.sha256()
70
+ with open(abs_path, "rb") as f:
71
+ for chunk in iter(lambda: f.read(1 << 16), b""):
72
+ h.update(chunk)
73
+ return h.hexdigest()
74
+
75
+
76
+ def _is_root(path):
77
+ return os.path.exists(os.path.join(path, ".git")) or os.path.exists(
78
+ os.path.join(path, ".vmn")
79
+ )
80
+
81
+
82
+ def resolve_root_path():
83
+ """The nearest directory at or above the working dir holding .git or .vmn."""
84
+ path = os.path.realpath(
85
+ os.path.expanduser(os.environ.get("VMN_WORKING_DIR", os.getcwd()))
86
+ )
87
+ while not _is_root(path):
88
+ parent = os.path.realpath(os.path.join(path, ".."))
89
+ if parent == path:
90
+ raise RuntimeError("Running from an unmanaged directory")
91
+ path = parent
92
+ return path
93
+
94
+
95
+ def get_repo_lock(vmn_root_path):
96
+ """The per-repo FileLock every vmn command takes; ``$VMN_LOCK_FILE_PATH``
97
+ overrides its path for the whole process."""
98
+ return FileLock(
99
+ os.environ.get("VMN_LOCK_FILE_PATH")
100
+ or os.path.join(vmn_root_path, ".vmn", "vmn.lock")
101
+ )
File without changes
@@ -0,0 +1,65 @@
1
+ #!/usr/bin/env python3
2
+ """One background worker that only ever does the newest pending job.
3
+
4
+ Slow remote writes — the SDK's run-state copy and log sync, ``vmn-exp run``'s
5
+ log sync — must never run on the heartbeat thread: a hung S3 PUT there is a heartbeat that did
6
+ not happen. Both are "catch the remote up" jobs, where a job overtaken by a
7
+ newer one is not worth doing, so pending jobs coalesce into the newest and run
8
+ one at a time, in order.
9
+ """
10
+ import logging
11
+ import threading
12
+
13
+ _LOGGER = logging.getLogger(__name__)
14
+ _NOTHING = object()
15
+
16
+
17
+ class Coalescing:
18
+ """Run ``job(item)`` off-thread for the newest submitted *item*."""
19
+
20
+ def __init__(self, job, name):
21
+ self._job = job
22
+ self._name = name
23
+ self._cond = threading.Condition()
24
+ self._pending = _NOTHING
25
+ self._closed = False
26
+ self._thread = None
27
+
28
+ def submit(self, item=None):
29
+ with self._cond:
30
+ self._pending = item
31
+ self._cond.notify()
32
+ if self._thread is None:
33
+ self._thread = threading.Thread(
34
+ target=self._loop, name=self._name, daemon=True
35
+ )
36
+ self._thread.start()
37
+
38
+ def close(self, timeout):
39
+ """Finish what is pending, then stop. False if that took over *timeout*."""
40
+ with self._cond:
41
+ self._closed = True
42
+ self._cond.notify()
43
+ thread = self._thread
44
+ if thread is None:
45
+ return True
46
+ thread.join(timeout)
47
+ return not thread.is_alive()
48
+
49
+ def _next(self):
50
+ """The next item, or ``_NOTHING`` once closed with nothing pending."""
51
+ with self._cond:
52
+ while self._pending is _NOTHING and not self._closed:
53
+ self._cond.wait()
54
+ item, self._pending = self._pending, _NOTHING
55
+ return item
56
+
57
+ def _loop(self):
58
+ while True:
59
+ item = self._next()
60
+ if item is _NOTHING:
61
+ return
62
+ try:
63
+ self._job(item)
64
+ except Exception:
65
+ _LOGGER.debug(f"{self._name} job failed", exc_info=True)
@@ -0,0 +1,49 @@
1
+ #!/usr/bin/env python3
2
+ """Running the chores a run's liveness depends on, without ever raising.
3
+
4
+ Both run recorders — ``vmn-exp run`` supervising a child and the SDK's
5
+ in-process ``Run`` — have steps that must not end the run when they fail: a
6
+ heartbeat write, a remote sync, a final state. A failure there is reported,
7
+ never thrown at the workload.
8
+ """
9
+ import threading
10
+
11
+
12
+ class BestEffort:
13
+ """Call a step; on failure report it and return None instead of raising.
14
+
15
+ *describe(what, exc)* words the report. With *warn_first* the first failure
16
+ of each step is a warning — the user should learn their heartbeat or sync
17
+ is failing — and repeats go to debug, so a flapping backend cannot flood
18
+ the output. Without it every failure is debug-only.
19
+ """
20
+
21
+ def __init__(self, logger, describe, warn_first=True):
22
+ self._logger = logger
23
+ self._describe = describe
24
+ self._warn_first = warn_first
25
+ self._warned = set()
26
+ self._lock = threading.Lock()
27
+
28
+ def __call__(self, what, fn, *args, **kwargs):
29
+ try:
30
+ return fn(*args, **kwargs)
31
+ except Exception as exc:
32
+ if self._first_failure(what):
33
+ self._logger.warning(self._describe(what, exc))
34
+ self._logger.debug(f"{what} failed", exc_info=True)
35
+ return None
36
+
37
+ def _first_failure(self, what):
38
+ if not self._warn_first:
39
+ return False
40
+ with self._lock:
41
+ first = what not in self._warned
42
+ self._warned.add(what)
43
+ return first
44
+
45
+
46
+ def quiet(logger):
47
+ """A :class:`BestEffort` that only ever logs at debug."""
48
+ return BestEffort(logger, describe=None, warn_first=False)
49
+
@@ -0,0 +1,331 @@
1
+ """Stdlib-only environment capture for experiment provenance.
2
+
3
+ ``capture_env()`` -> full dict (python/platform/hostname/packages/cuda/container)
4
+ ``scan_packages()`` -> (dict name->version, truncated_flag) scanning sys.path
5
+ ``env_summary()`` -> <= 2 KB dict suitable for metadata.yml
6
+ ``packages_sha()`` -> stable sha256 over sorted name==version pairs
7
+ ``_cuda_info()`` -> reads torch.version.cuda ONLY if torch already loaded
8
+ ``_container_info()`` -> image digest from env vars / /.dockerenv presence
9
+ ``env_from_interpreter()`` -> probe an external Python exe, parse JSON from stdout
10
+
11
+ No third-party imports; safe to run as a subprocess script.
12
+ """
13
+ import hashlib
14
+ import importlib
15
+ import json
16
+ import os
17
+ import platform
18
+ import re
19
+ import socket
20
+ import subprocess
21
+ import sys
22
+
23
+ _PACKAGE_CAP = 2000
24
+ _PEP503_RE = re.compile(r"[-_.]+")
25
+
26
+ # Key ML packages shown verbatim in the summary.
27
+ _KEY_PACKAGES = {"torch", "tensorflow", "jax", "numpy", "transformers"}
28
+
29
+
30
+ # ---------------------------------------------------------------------------
31
+ # Name normalisation (PEP 503)
32
+ # ---------------------------------------------------------------------------
33
+
34
+ def _normalize_name(name: str) -> str:
35
+ return _PEP503_RE.sub("-", name).lower()
36
+
37
+
38
+ # ---------------------------------------------------------------------------
39
+ # Metadata file parser
40
+ # ---------------------------------------------------------------------------
41
+
42
+ def _parse_metadata_file(path: str):
43
+ """Return (name, version) from a METADATA or PKG-INFO file header section."""
44
+ name = version = None
45
+ try:
46
+ with open(path, encoding="utf-8", errors="replace") as fh:
47
+ for line in fh:
48
+ line = line.rstrip("\r\n")
49
+ if not line:
50
+ break
51
+ if line.startswith("Name:") and name is None:
52
+ name = line[5:].strip()
53
+ elif line.startswith("Version:") and version is None:
54
+ version = line[8:].strip()
55
+ if name and version:
56
+ break
57
+ except OSError:
58
+ pass
59
+ return name, version
60
+
61
+
62
+ def _parse_dir_stem(stem: str):
63
+ """Fallback: parse name and version from a dist/egg dir stem like 'numpy-1.24.0'."""
64
+ m = re.match(r"^(.+?)-(\d[\w.]*)(.*)$", stem)
65
+ if m:
66
+ tail = (m.group(2) + m.group(3)).rstrip("-")
67
+ return m.group(1), tail
68
+ return stem, "unknown"
69
+
70
+
71
+ # ---------------------------------------------------------------------------
72
+ # Package scanning
73
+ # ---------------------------------------------------------------------------
74
+
75
+ def scan_packages(paths=None):
76
+ """Scan *paths* (default: sys.path) for dist-info / egg-info records.
77
+
78
+ Returns ``(packages, truncated)`` where *packages* maps normalised package
79
+ names to version strings and *truncated* is ``True`` when the cap was hit.
80
+ """
81
+ if paths is None:
82
+ paths = sys.path
83
+ packages: dict = {}
84
+ for base in paths:
85
+ if not base or not os.path.isdir(base):
86
+ continue
87
+ try:
88
+ entries = os.listdir(base)
89
+ except OSError:
90
+ continue
91
+ for entry in entries:
92
+ if len(packages) >= _PACKAGE_CAP:
93
+ return packages, True
94
+ name = version = None
95
+ full = os.path.join(base, entry)
96
+ if entry.endswith(".dist-info"):
97
+ meta = os.path.join(full, "METADATA")
98
+ if os.path.isfile(meta):
99
+ name, version = _parse_metadata_file(meta)
100
+ if not name:
101
+ name, version = _parse_dir_stem(entry[: -len(".dist-info")])
102
+ elif entry.endswith(".egg-info"):
103
+ if os.path.isdir(full):
104
+ meta = os.path.join(full, "PKG-INFO")
105
+ else:
106
+ meta = full # single-file .egg-info
107
+ if os.path.isfile(meta):
108
+ name, version = _parse_metadata_file(meta)
109
+ if not name:
110
+ name, version = _parse_dir_stem(entry[: -len(".egg-info")])
111
+ if name:
112
+ norm = _normalize_name(name)
113
+ packages.setdefault(norm, version or "unknown")
114
+ return packages, False
115
+
116
+
117
+ # ---------------------------------------------------------------------------
118
+ # SHA over package list
119
+ # ---------------------------------------------------------------------------
120
+
121
+ def packages_sha(packages: dict) -> str:
122
+ """Return a stable sha256 hex digest over sorted ``name==version`` pairs."""
123
+ items = sorted(f"{k}=={v}" for k, v in packages.items())
124
+ return hashlib.sha256("\n".join(items).encode()).hexdigest()
125
+
126
+
127
+ # ---------------------------------------------------------------------------
128
+ # CUDA info (never imports torch; reads it only if already loaded)
129
+ # ---------------------------------------------------------------------------
130
+
131
+ def _cuda_info():
132
+ """Return cuda/driver info dict, or ``None`` if nothing is available."""
133
+ result: dict = {}
134
+ torch = sys.modules.get("torch")
135
+ if torch is not None:
136
+ try:
137
+ cuda_ver = torch.version.cuda
138
+ if cuda_ver is not None:
139
+ result["torch_cuda"] = cuda_ver
140
+ except Exception:
141
+ pass
142
+
143
+ pynvml = sys.modules.get("pynvml")
144
+ if pynvml is None:
145
+ try:
146
+ pynvml = importlib.import_module("pynvml")
147
+ except Exception:
148
+ pass
149
+ if pynvml is not None:
150
+ try:
151
+ pynvml.nvmlInit()
152
+ result["driver"] = pynvml.nvmlSystemGetDriverVersion()
153
+ except Exception:
154
+ pass
155
+ return result or None
156
+
157
+
158
+ # ---------------------------------------------------------------------------
159
+ # Container / image info (env vars only — never dumps full environ)
160
+ # ---------------------------------------------------------------------------
161
+
162
+ _DIGEST_ENV_VARS = (
163
+ "VMN_IMAGE_DIGEST",
164
+ "IMAGE_DIGEST",
165
+ "DOCKER_IMAGE_DIGEST",
166
+ "CONTAINER_IMAGE_DIGEST",
167
+ )
168
+
169
+
170
+ def _container_info():
171
+ """Return container info dict, or ``None`` if no evidence of a container."""
172
+ result: dict = {}
173
+ for var in _DIGEST_ENV_VARS:
174
+ val = os.environ.get(var)
175
+ if val:
176
+ result["image_digest"] = val
177
+ break
178
+ if os.path.exists("/.dockerenv"):
179
+ result["in_docker"] = True
180
+ return result or None
181
+
182
+
183
+ # ---------------------------------------------------------------------------
184
+ # Full env capture
185
+ # ---------------------------------------------------------------------------
186
+
187
+ def capture_env() -> dict:
188
+ """Return a full environment dict (python/platform/hostname/packages/cuda/container)."""
189
+ packages, truncated = scan_packages()
190
+ env: dict = {
191
+ "python": {
192
+ "version": platform.python_version(),
193
+ "implementation": platform.python_implementation(),
194
+ "executable": sys.executable,
195
+ },
196
+ "platform": {
197
+ "system": platform.system(),
198
+ "machine": platform.machine(),
199
+ "release": platform.release(),
200
+ },
201
+ "hostname": socket.gethostname(),
202
+ "packages": packages,
203
+ }
204
+ if truncated:
205
+ env["packages_truncated"] = True
206
+ cuda = _cuda_info()
207
+ if cuda:
208
+ env["cuda"] = cuda
209
+ container = _container_info()
210
+ if container:
211
+ env["container"] = container
212
+ return env
213
+
214
+
215
+ # ---------------------------------------------------------------------------
216
+ # Summary (<= 2 KB)
217
+ # ---------------------------------------------------------------------------
218
+
219
+ def env_summary(env: dict) -> dict:
220
+ """Return a compact (<= 2 KB) summary dict suitable for metadata.yml."""
221
+ packages = env.get("packages", {})
222
+ key_pkgs = {k: packages[k] for k in _KEY_PACKAGES if k in packages}
223
+ summary: dict = {
224
+ "python": env.get("python", {}).get("version"),
225
+ "platform": "{}/{}".format(
226
+ env.get("platform", {}).get("system", ""),
227
+ env.get("platform", {}).get("machine", ""),
228
+ ),
229
+ "packages_count": len(packages),
230
+ "packages_sha": packages_sha(packages),
231
+ }
232
+ if key_pkgs:
233
+ summary["key_packages"] = key_pkgs
234
+ cuda = env.get("cuda")
235
+ if cuda:
236
+ summary["cuda"] = cuda
237
+ container = env.get("container")
238
+ if container and container.get("image_digest"):
239
+ summary["image_digest"] = container["image_digest"]
240
+ return summary
241
+
242
+
243
+ # ---------------------------------------------------------------------------
244
+ # Probe an external interpreter
245
+ # ---------------------------------------------------------------------------
246
+
247
+ def env_from_interpreter(python_exe: str, timeout: int = 5) -> dict:
248
+ """Run this module file under *python_exe* with ``-I`` and return the JSON output.
249
+
250
+ The child runs in isolated mode so it does not inherit PYTHONPATH or
251
+ site-packages from the caller. If the child does not have vmn installed,
252
+ we still run the raw file (it is stdlib-only and has no imports from vmn).
253
+ """
254
+ module_file = os.path.abspath(__file__)
255
+ result = subprocess.run(
256
+ [python_exe, "-I", module_file],
257
+ capture_output=True,
258
+ text=True,
259
+ timeout=timeout,
260
+ )
261
+ if result.returncode != 0:
262
+ raise RuntimeError(
263
+ f"env_from_interpreter failed (exit {result.returncode}): {result.stderr}"
264
+ )
265
+ return json.loads(result.stdout)
266
+
267
+
268
+ # ---------------------------------------------------------------------------
269
+ # Capture-decision helpers
270
+ # ---------------------------------------------------------------------------
271
+
272
+ #: Environment variable that opts out of env capture when set to ``0`` or
273
+ #: a falsy string (``false``, ``no``, ``off``).
274
+ CAPTURE_ENV_ENV = "VMN_CAPTURE_ENV"
275
+
276
+ _OPT_OUT_VALS = frozenset(("0", "false", "False", "no", "off"))
277
+
278
+
279
+ def should_capture(explicit, exp_conf=None):
280
+ """Return True if the environment should be captured.
281
+
282
+ Precedence: explicit arg > ``VMN_CAPTURE_ENV`` > conf > True.
283
+
284
+ *explicit* controls the override:
285
+
286
+ * ``False`` — always skip (overrides env-var and conf).
287
+ * ``True`` — override conf, but still respect ``VMN_CAPTURE_ENV``.
288
+ * ``None`` — full opt-out chain (env-var → conf → default True).
289
+
290
+ *exp_conf* is the ``experiment:`` section of the app's conf.yml, as a dict.
291
+ """
292
+ if explicit is False:
293
+ return False
294
+ env_val = os.environ.get(CAPTURE_ENV_ENV)
295
+ if env_val in _OPT_OUT_VALS:
296
+ return False
297
+ if explicit is True:
298
+ return True # explicit True overrides conf but env-var already checked
299
+ # explicit is None: also check conf
300
+ if exp_conf:
301
+ conf_val = exp_conf.get("capture_env")
302
+ if conf_val is not None:
303
+ return bool(conf_val)
304
+ return True
305
+
306
+
307
+ def capture_env_safe(python_exe=None):
308
+ """Capture the environment dict; return ``None`` on any failure (never raises).
309
+
310
+ When *python_exe* is given the child interpreter is probed via
311
+ ``env_from_interpreter``; otherwise the current interpreter is captured.
312
+ """
313
+ try:
314
+ if python_exe:
315
+ return env_from_interpreter(python_exe, timeout=5)
316
+ return capture_env()
317
+ except Exception:
318
+ import logging
319
+
320
+ logging.getLogger(__name__).warning(
321
+ "vmn: could not capture environment for this run"
322
+ )
323
+ return None
324
+
325
+
326
+ # ---------------------------------------------------------------------------
327
+ # __main__: print JSON for env_from_interpreter / manual use
328
+ # ---------------------------------------------------------------------------
329
+
330
+ if __name__ == "__main__":
331
+ print(json.dumps(capture_env()))