diff --git a/CHANGELOG.md b/CHANGELOG.md index 1dedd210..eddbea39 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,7 @@ - Improves handling of Conda based environments in metadata collection. - Adds additional options to `Client.get_runs`. +- Added ability to include environment variables within metadata for runs. ## [v2.1.2](https://github.com/simvue-io/client/releases/tag/v2.1.2) - 2025-06-25 diff --git a/simvue/config/parameters.py b/simvue/config/parameters.py index e4a5b526..28e343ec 100644 --- a/simvue/config/parameters.py +++ b/simvue/config/parameters.py @@ -72,6 +72,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): diff --git a/simvue/metadata.py b/simvue/metadata.py index 67b52af2..614de6e0 100644 --- a/simvue/metadata.py +++ b/simvue/metadata.py @@ -9,6 +9,8 @@ import contextlib import typing import json +import os +import fnmatch import toml import yaml import logging @@ -257,7 +259,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): @@ -268,4 +285,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 diff --git a/simvue/run.py b/simvue/run.py index 51419d7a..6854da9a 100644 --- a/simvue/run.py +++ b/simvue/run.py @@ -623,6 +623,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 @@ -660,6 +661,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 ------- @@ -677,6 +681,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 @@ -742,7 +747,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() diff --git a/tests/functional/test_run_class.py b/tests/functional/test_run_class.py index 595fb441..804f9bb6 100644 --- a/tests/functional/test_run_class.py +++ b/tests/functional/test_run_class.py @@ -1285,10 +1285,28 @@ def test_reconnect_functionality(mode, monkeypatch: pytest.MonkeyPatch) -> None: assert dict(_reconnected_run.metrics)["test_metric"]["last"] == 1 assert client.get_events(run_id)[0]["message"] == "Testing!" - if temp_d: - temp_d.cleanup() +@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) +@pytest.mark.run def test_reconnect_with_process() -> None: _uuid = f"{uuid.uuid4()}".split("-")[0] with simvue.Run() as run: diff --git a/tests/unit/test_metadata.py b/tests/unit/test_metadata.py index 5375e12b..296d1292 100644 --- a/tests/unit/test_metadata.py +++ b/tests/unit/test_metadata.py @@ -1,3 +1,4 @@ +import os import pytest import pathlib import re @@ -55,3 +56,25 @@ def test_environment() -> None: assert metadata["rust"]["project"]["name"] == "example_project" assert metadata["julia"]["project"]["name"] == "Julia Demo 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()) +