diff --git a/CHANGELOG.md b/CHANGELOG.md index fcffc1f4..dca6ecb9 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/simvue/config/parameters.py b/simvue/config/parameters.py index 9e0b38bc..23770082 100644 --- a/simvue/config/parameters.py +++ b/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): diff --git a/simvue/metadata.py b/simvue/metadata.py index 0f63981c..eaeaae69 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 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 diff --git a/simvue/run.py b/simvue/run.py index 355cd693..4772cc7f 100644 --- a/simvue/run.py +++ b/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() diff --git a/tests/conftest.py b/tests/conftest.py index 01c5cf4e..f0b37c10 100644 --- a/tests/conftest.py +++ b/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"], diff --git a/tests/functional/test_run_class.py b/tests/functional/test_run_class.py index e493c5db..b248087c 100644 --- a/tests/functional/test_run_class.py +++ b/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) diff --git a/tests/unit/test_metadata.py b/tests/unit/test_metadata.py index 5c454e14..4cae6072 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 @@ -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" \ No newline at end of file + 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())