Add functionality for including environment variables as metadata

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## Unreleased
 * Added sorting to server queries, users can now specify to sort by columns during data retrieval from the database.
+* Added ability to include environment variables within metadata for runs.
 ## [v2.0.1](https://github.com/simvue-io/client/releases/tag/v2.0.1) - 2025-03-24
 * Improvements to docstrings on methods, classes and functions.
 ## [v2.0.0](https://github.com/simvue-io/client/releases/tag/v2.0.0) - 2025-03-07

simvue/config/parameters.py

@@ -64,6 +64,7 @@ class DefaultRunSpecifications(pydantic.BaseModel):
     folder: str = pydantic.Field("/", pattern=sv_models.FOLDER_REGEX)
     metadata: dict[str, str | int | float | bool] | None = None
     mode: typing.Literal["offline", "disabled", "online"] = "online"
+    record_shell_vars: list[str] | None = None
 
 
 class ClientGeneralOptions(pydantic.BaseModel):

simvue/metadata.py

@@ -9,6 +9,8 @@
 import contextlib
 import typing
 import json
+import os
+import fnmatch
 import toml
 import logging
 import pathlib
@@ -179,7 +181,22 @@ def _node_js_env(repository: pathlib.Path) -> dict[str, typing.Any]:
     return js_meta
 
 
-def environment(repository: pathlib.Path = pathlib.Path.cwd()) -> dict[str, typing.Any]:
+def _environment_variables(glob_exprs: list[str]) -> dict[str, str]:
+    """Retrieve values for environment variables."""
+    _env_vars: list[str] = list(os.environ.keys())
+    _metadata: dict[str, str] = {}
+
+    for pattern in glob_exprs:
+        for key in fnmatch.filter(_env_vars, pattern):
+            _metadata[key] = os.environ[key]
+
+    return _metadata
+
+
+def environment(
+    repository: pathlib.Path = pathlib.Path.cwd(),
+    env_var_glob_exprs: set[str] | None = None,
+) -> dict[str, typing.Any]:
     """Retrieve environment metadata"""
     _environment_meta = {}
     if _python_meta := _python_env(repository):
@@ -190,4 +207,6 @@ def environment(repository: pathlib.Path = pathlib.Path.cwd()) -> dict[str, typi
         _environment_meta["julia"] = _julia_meta
     if _js_meta := _node_js_env(repository):
         _environment_meta["javascript"] = _js_meta
+    if env_var_glob_exprs:
+        _environment_meta["shell"] = _environment_variables(env_var_glob_exprs)
     return _environment_meta

simvue/run.py

@@ -613,6 +613,7 @@ def init(
         timeout: int | None = 180,
         visibility: typing.Literal["public", "tenant"] | list[str] | None = None,
         no_color: bool = False,
+        record_shell_vars: set[str] | None = None,
     ) -> bool:
         """Initialise a Simvue run
 
@@ -650,6 +651,9 @@ def init(
                 * A list of usernames with which to share this run
         no_color : bool, optional
             disable terminal colors. Default False.
+        record_shell_vars : list[str] | None,
+            list of environment variables to store as metadata, these can
+            either be defined as literal strings or globular expressions
 
         Returns
         -------
@@ -667,6 +671,7 @@ def init(
         folder = folder or self._user_config.run.folder
         name = name or self._user_config.run.name
         metadata = (metadata or {}) | (self._user_config.run.metadata or {})
+        record_shell_vars = record_shell_vars or self._user_config.run.record_shell_vars
 
         self._term_color = not no_color
 
@@ -734,7 +739,11 @@ def init(
         self._sv_obj.ttl = self._retention
         self._sv_obj.status = self._status
         self._sv_obj.tags = tags
-        self._sv_obj.metadata = (metadata or {}) | git_info(os.getcwd()) | environment()
+        self._sv_obj.metadata = (
+            (metadata or {})
+            | git_info(os.getcwd())
+            | environment(env_var_glob_exprs=record_shell_vars)
+        )
         self._sv_obj.heartbeat_timeout = timeout
         self._sv_obj.alerts = []
         self._sv_obj.created = time.time()

tests/conftest.py

@@ -139,6 +139,7 @@ def setup_test_run(run: sv_run.Run, create_objects: bool, request: pytest.Fixtur
 
     run.config(suppress_errors=False)
     run._heartbeat_interval = 1
+
     run.init(
         name=TEST_DATA['metadata']['test_identifier'],
         tags=TEST_DATA["tags"],

tests/functional/test_run_class.py

@@ -1,6 +1,6 @@
 import os
 from os.path import basename
-from numpy import identity
+from numpy import identity, rec
 import pytest
 import pytest_mock
 import time
@@ -1052,3 +1052,24 @@ def test_reconnect(mode, monkeypatch: pytest.MonkeyPatch) -> None:
     _reconnected_run = client.get_run(run_id)
     assert dict(_reconnected_run.metrics)["test_metric"]["last"] == 1
     assert client.get_events(run_id)[0]["message"] == "Testing!"
+
+
+@pytest.mark.run
+def test_env_var_metadata() -> None:
+    # Add some environment variables to glob
+    _recorded_env = {
+        "SIMVUE_RUN_TEST_VAR_1": "1",
+        "SIMVUE_RUN_TEST_VAR_2": "hello"
+    }
+    os.environ.update(_recorded_env)
+    with simvue.Run() as run:
+        run.init(
+            name="test_reconnect",
+            folder="/simvue_unit_testing",
+            retention_period="2 minutes",
+            timeout=None,
+            running=False,
+            record_shell_vars={"SIMVUE_RUN_TEST_VAR_*"}
+        )
+    _recorded_meta = RunObject(identifier=run._id).metadata
+    assert all(key in _recorded_meta.get("shell") for key in _recorded_env)

tests/unit/test_metadata.py

@@ -1,3 +1,4 @@
+import os
 import pytest
 import pathlib
 import re
@@ -51,4 +52,26 @@ def test_environment() -> None:
     assert metadata["python"]["project"]["name"] == "example-repo"
     assert metadata["rust"]["project"]["name"] == "example_project"
     assert metadata["julia"]["project"]["name"] == "Julia Demo Project"
-    assert metadata["javascript"]["project"]["name"] == "my-awesome-project"
+    assert metadata["javascript"]["project"]["name"] == "my-awesome-project"
+
+
+@pytest.mark.metadata
+@pytest.mark.local
+def test_slurm_env_var_capture() -> None:
+    _slurm_env = {
+        "SLURM_CPUS_PER_TASK": "2",
+        "SLURM_TASKS_PER_NODE": "1",
+        "SLURM_NNODES": "1",
+        "SLURM_NTASKS_PER_NODE": "1",
+        "SLURM_NTASKS": "1",
+        "SLURM_JOB_CPUS_PER_NODE": "2",
+        "SLURM_CPUS_ON_NODE": "2",
+        "SLURM_JOB_NUM_NODES": "1",
+        "SLURM_MEM_PER_NODE": "2000",
+        "SLURM_NPROCS": "1",
+        "SLURM_TRES_PER_TASK": "cpu:2",
+    }
+    os.environ.update(_slurm_env)
+
+    sv_meta.metadata = sv_meta.environment(env_var_glob_exprs={"SLURM_*"})
+    assert all((key, value) in sv_meta.metadata["shell"].items() for key, value in _slurm_env.items())