Skip to content

Add functionality for including environment variables as metadata #771

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 8 commits into from
Aug 14, 2025
Merged
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
1 change: 1 addition & 0 deletions simvue/config/parameters.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
21 changes: 20 additions & 1 deletion simvue/metadata.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,8 @@
import contextlib
import typing
import json
import os
import fnmatch
import toml
import yaml
import logging
Expand Down Expand Up @@ -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):
Expand All @@ -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
11 changes: 10 additions & 1 deletion simvue/run.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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
-------
Expand All @@ -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

Expand Down Expand Up @@ -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()
Expand Down
22 changes: 20 additions & 2 deletions tests/functional/test_run_class.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
23 changes: 23 additions & 0 deletions tests/unit/test_metadata.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import os
import pytest
import pathlib
import re
Expand Down Expand Up @@ -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())