Skip to content

Commit fe69caf

Browse files
authored
Merge pull request #771 from simvue-io/feature/env-var-metadata
Add functionality for including environment variables as metadata
2 parents 831cc3c + 3d293c9 commit fe69caf

File tree

6 files changed

+75
-4
lines changed

6 files changed

+75
-4
lines changed

CHANGELOG.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44

55
- Improves handling of Conda based environments in metadata collection.
66
- Adds additional options to `Client.get_runs`.
7+
- Added ability to include environment variables within metadata for runs.
78

89
## [v2.1.2](https://github.com/simvue-io/client/releases/tag/v2.1.2) - 2025-06-25
910

simvue/config/parameters.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,7 @@ class DefaultRunSpecifications(pydantic.BaseModel):
7272
folder: str = pydantic.Field("/", pattern=sv_models.FOLDER_REGEX)
7373
metadata: dict[str, str | int | float | bool] | None = None
7474
mode: typing.Literal["offline", "disabled", "online"] = "online"
75+
record_shell_vars: list[str] | None = None
7576

7677

7778
class ClientGeneralOptions(pydantic.BaseModel):

simvue/metadata.py

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,8 @@
99
import contextlib
1010
import typing
1111
import json
12+
import os
13+
import fnmatch
1214
import toml
1315
import yaml
1416
import logging
@@ -257,7 +259,22 @@ def _node_js_env(repository: pathlib.Path) -> dict[str, typing.Any]:
257259
return js_meta
258260

259261

260-
def environment(repository: pathlib.Path = pathlib.Path.cwd()) -> dict[str, typing.Any]:
262+
def _environment_variables(glob_exprs: list[str]) -> dict[str, str]:
263+
"""Retrieve values for environment variables."""
264+
_env_vars: list[str] = list(os.environ.keys())
265+
_metadata: dict[str, str] = {}
266+
267+
for pattern in glob_exprs:
268+
for key in fnmatch.filter(_env_vars, pattern):
269+
_metadata[key] = os.environ[key]
270+
271+
return _metadata
272+
273+
274+
def environment(
275+
repository: pathlib.Path = pathlib.Path.cwd(),
276+
env_var_glob_exprs: set[str] | None = None,
277+
) -> dict[str, typing.Any]:
261278
"""Retrieve environment metadata"""
262279
_environment_meta = {}
263280
if _python_meta := _python_env(repository):
@@ -268,4 +285,6 @@ def environment(repository: pathlib.Path = pathlib.Path.cwd()) -> dict[str, typi
268285
_environment_meta["julia"] = _julia_meta
269286
if _js_meta := _node_js_env(repository):
270287
_environment_meta["javascript"] = _js_meta
288+
if env_var_glob_exprs:
289+
_environment_meta["shell"] = _environment_variables(env_var_glob_exprs)
271290
return _environment_meta

simvue/run.py

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -623,6 +623,7 @@ def init(
623623
timeout: int | None = 180,
624624
visibility: typing.Literal["public", "tenant"] | list[str] | None = None,
625625
no_color: bool = False,
626+
record_shell_vars: set[str] | None = None,
626627
) -> bool:
627628
"""Initialise a Simvue run
628629
@@ -660,6 +661,9 @@ def init(
660661
* A list of usernames with which to share this run
661662
no_color : bool, optional
662663
disable terminal colors. Default False.
664+
record_shell_vars : list[str] | None,
665+
list of environment variables to store as metadata, these can
666+
either be defined as literal strings or globular expressions
663667
664668
Returns
665669
-------
@@ -677,6 +681,7 @@ def init(
677681
folder = folder or self._user_config.run.folder
678682
name = name or self._user_config.run.name
679683
metadata = (metadata or {}) | (self._user_config.run.metadata or {})
684+
record_shell_vars = record_shell_vars or self._user_config.run.record_shell_vars
680685

681686
self._term_color = not no_color
682687

@@ -742,7 +747,11 @@ def init(
742747
self._sv_obj.ttl = self._retention
743748
self._sv_obj.status = self._status
744749
self._sv_obj.tags = tags
745-
self._sv_obj.metadata = (metadata or {}) | git_info(os.getcwd()) | environment()
750+
self._sv_obj.metadata = (
751+
(metadata or {})
752+
| git_info(os.getcwd())
753+
| environment(env_var_glob_exprs=record_shell_vars)
754+
)
746755
self._sv_obj.heartbeat_timeout = timeout
747756
self._sv_obj.alerts = []
748757
self._sv_obj.created = time.time()

tests/functional/test_run_class.py

Lines changed: 20 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1285,10 +1285,28 @@ def test_reconnect_functionality(mode, monkeypatch: pytest.MonkeyPatch) -> None:
12851285
assert dict(_reconnected_run.metrics)["test_metric"]["last"] == 1
12861286
assert client.get_events(run_id)[0]["message"] == "Testing!"
12871287

1288-
if temp_d:
1289-
temp_d.cleanup()
12901288

1289+
@pytest.mark.run
1290+
def test_env_var_metadata() -> None:
1291+
# Add some environment variables to glob
1292+
_recorded_env = {
1293+
"SIMVUE_RUN_TEST_VAR_1": "1",
1294+
"SIMVUE_RUN_TEST_VAR_2": "hello"
1295+
}
1296+
os.environ.update(_recorded_env)
1297+
with simvue.Run() as run:
1298+
run.init(
1299+
name="test_reconnect",
1300+
folder="/simvue_unit_testing",
1301+
retention_period="2 minutes",
1302+
timeout=None,
1303+
running=False,
1304+
record_shell_vars={"SIMVUE_RUN_TEST_VAR_*"}
1305+
)
1306+
_recorded_meta = RunObject(identifier=run.id).metadata
1307+
assert all(key in _recorded_meta.get("shell") for key in _recorded_env)
12911308

1309+
@pytest.mark.run
12921310
def test_reconnect_with_process() -> None:
12931311
_uuid = f"{uuid.uuid4()}".split("-")[0]
12941312
with simvue.Run() as run:

tests/unit/test_metadata.py

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import os
12
import pytest
23
import pathlib
34
import re
@@ -55,3 +56,25 @@ def test_environment() -> None:
5556
assert metadata["rust"]["project"]["name"] == "example_project"
5657
assert metadata["julia"]["project"]["name"] == "Julia Demo Project"
5758
assert metadata["javascript"]["project"]["name"] == "my-awesome-project"
59+
60+
@pytest.mark.metadata
61+
@pytest.mark.local
62+
def test_slurm_env_var_capture() -> None:
63+
_slurm_env = {
64+
"SLURM_CPUS_PER_TASK": "2",
65+
"SLURM_TASKS_PER_NODE": "1",
66+
"SLURM_NNODES": "1",
67+
"SLURM_NTASKS_PER_NODE": "1",
68+
"SLURM_NTASKS": "1",
69+
"SLURM_JOB_CPUS_PER_NODE": "2",
70+
"SLURM_CPUS_ON_NODE": "2",
71+
"SLURM_JOB_NUM_NODES": "1",
72+
"SLURM_MEM_PER_NODE": "2000",
73+
"SLURM_NPROCS": "1",
74+
"SLURM_TRES_PER_TASK": "cpu:2",
75+
}
76+
os.environ.update(_slurm_env)
77+
78+
sv_meta.metadata = sv_meta.environment(env_var_glob_exprs={"SLURM_*"})
79+
assert all((key, value) in sv_meta.metadata["shell"].items() for key, value in _slurm_env.items())
80+

0 commit comments

Comments
 (0)