Merge pull request #6 from boeschf/eval/env_var

Environment variables

Author: theely

examples/alps-ml-vetting/config.yaml

@@ -3,6 +3,24 @@ scheduler: slurm
 pip:
    index_url: "https://jfrog.svc.cscs.ch/artifactory/api/pypi/pypi-remote/simple"
 evals:
+- name: Environment Variables
+  type: vetnode.evaluations.env_var_eval.EnvVarEval
+  expected:
+    CUDA_CACHE_DISABLE: "1"
+    NCCL_NET: "AWS Libfabric"
+    NCCL_CROSS_NIC: "1"
+    NCCL_NET_GDR_LEVEL: "PHB"
+    NCCL_PROTO: "^LL128"
+    FI_PROVIDER: "cxi"
+    FI_CXI_DEFAULT_CQ_SIZE: "131072"
+    FI_CXI_DEFAULT_TX_SIZE: "16384"
+    FI_CXI_DISABLE_HOST_REGISTER: "1"
+    FI_CXI_RDZV_PROTO: "alt_read"
+    FI_CXI_RDZV_EAGER_SIZE: "0"
+    FI_CXI_RDZV_GET_MIN: "0"
+    FI_CXI_RDZV_THRESHOLD: "0"
+    FI_CXI_RX_MATCH_MODE: "hybrid"
+    FI_MR_CACHE_MONITOR: "userfaultfd"
 - name: Check GPU
   type: vetnode.evaluations.gpu_eval.GPUEval
   max_temp:  30
@@ -56,4 +74,4 @@ evals:
     runs: 3
   requirements:
       - ['torch', '--index-url', 'https://download.pytorch.org/whl/cu130']
-      - numpy
+      - numpy

examples/image-vet/config.yaml

@@ -3,6 +3,24 @@ scheduler: slurm
 pip:
    index_url: "https://jfrog.svc.cscs.ch/artifactory/api/pypi/pypi-remote/simple"
 evals:
+- name: Environment Variables
+  type: vetnode.evaluations.env_var_eval.EnvVarEval
+  expected:
+    CUDA_CACHE_DISABLE: "1"
+    NCCL_NET: "AWS Libfabric"
+    NCCL_CROSS_NIC: "1"
+    NCCL_NET_GDR_LEVEL: "PHB"
+    NCCL_PROTO: "^LL128"
+    FI_PROVIDER: "cxi"
+    FI_CXI_DEFAULT_CQ_SIZE: "131072"
+    FI_CXI_DEFAULT_TX_SIZE: "16384"
+    FI_CXI_DISABLE_HOST_REGISTER: "1"
+    FI_CXI_RDZV_PROTO: "alt_read"
+    FI_CXI_RDZV_EAGER_SIZE: "0"
+    FI_CXI_RDZV_GET_MIN: "0"
+    FI_CXI_RDZV_THRESHOLD: "0"
+    FI_CXI_RX_MATCH_MODE: "hybrid"
+    FI_MR_CACHE_MONITOR: "userfaultfd"
 - name: CudaKernel
   type: vetnode.evaluations.cuda_eval.CUDAEval
   cuda_home: /usr/local/cuda 
@@ -52,4 +70,4 @@ evals:
     runs: 3
   requirements:
       - ['torch', '--index-url', 'https://download.pytorch.org/whl/cu130']
-      - numpy
+      - numpy

requirements-testing.txt

@@ -1,3 +1,4 @@
 pytest
+pytest-asyncio
 build
-twine
+twine

src/vetnode/cli.py

@@ -87,7 +87,7 @@ def diagnose(config,skip_install,verbose) -> None:
             click.secho(f"[{hostname}-{main_context.rank}: Vetted] ", fg='green',nl=False)
     else:
         if verbose:
-            click.secho(f"[{hostname}-{main_context.rank}:Cordon] ", fg='red',nl=False)
+            click.secho(f"[{hostname}-{main_context.rank}: Cordon] ", fg='red',nl=False)
         sys.exit(1)
 
 @click.command()
@@ -310,4 +310,4 @@ def load_requirements(requirements: List[str], index_url: str = None):
             cmd.append(package)
         else:
             cmd += package
-        subprocess.check_call(cmd)
+        subprocess.check_call(cmd)

src/vetnode/evaluations/env_var_eval.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+import os
+from typing import Literal
+from pydantic import Field
+
+from vetnode.evaluations.base_eval import BaseEval
+from vetnode.evaluations.models import EvalResultStatus
+
+
+class EnvVarEval(BaseEval):
+    """
+    Vetnode evaluation that verifies expected environment variables are present
+    and (optionally) match expected values.
+
+    Example config concept:
+      expected:
+        VAR: "value"   # exact string
+        SOME_VAR: null # must exist, any value
+    """
+
+    name: str
+    type: Literal["vetnode.evaluations.env_var_eval.EnvVarEval"]
+
+    # Map of env var -> expected value (str) OR None to mean "must exist"
+    expected: dict[str, str | None] = Field(default_factory=dict)
+
+    async def check(self, executor) -> tuple[EvalResultStatus, dict]:
+        found: dict[str, str | None] = {}
+        missing: list[str] = []
+        mismatched: dict[str, dict[str, str | None]] = {}
+
+        for key, value in (self.expected or {}).items():
+            actual = os.environ.get(key)
+            found[key] = actual
+
+            if actual is None:
+                missing.append(key)
+                continue
+
+            if value is not None and actual != value:
+                mismatched[key] = {"expected": value, "actual": actual}
+
+        ok = (not missing) and (not mismatched)
+        status = EvalResultStatus.SUCCESS if ok else EvalResultStatus.FAILED
+
+        metrics = {
+            "expected": self.expected,
+            "found": found,
+            "missing": missing,
+            "mismatched": mismatched,
+        }
+        return status, metrics

src/vetnode/evaluations/models.py

@@ -94,4 +94,4 @@ class BandwidthSize(ByteSize):
         'eb/s': 2**60,
     }
     byte_string_pattern = r'^\s*(\d*\.?\d+)\s*([\w\/]+)?'
-    byte_string_re = re.compile(byte_string_pattern, re.IGNORECASE)
+    byte_string_re = re.compile(byte_string_pattern, re.IGNORECASE)

tests/test_env_var_eval.py

@@ -0,0 +1,126 @@
+import pytest
+import textwrap
+from click.testing import CliRunner
+
+from vetnode.evaluations.models import EvalContext, EvalResultStatus
+from vetnode.evaluations.env_var_eval import EnvVarEval
+from vetnode.configuration import Configuration
+from vetnode.cli import build_context, load_evals, diagnose
+
+
+def mk_ctx() -> EvalContext:
+    return EvalContext(
+        scheduler="standalone",
+        rank=0,
+        local_rank=0,
+        eval_id=0,
+        world_size=1,
+        master_addr="localhost",
+        master_port=29500,
+    )
+
+
+@pytest.mark.asyncio
+async def test_must_exist_passes_when_set(monkeypatch):
+    monkeypatch.setenv("SOME_FLAG", "")
+    ev = EnvVarEval(mk_ctx(), name="env", type="vetnode.evaluations.env_var_eval.EnvVarEval",
+                    expected={"SOME_FLAG": None})
+
+    status, metrics = await ev.check(None)
+    assert status == EvalResultStatus.SUCCESS
+    assert metrics["missing"] == []
+    assert metrics["mismatched"] == {}
+
+
+@pytest.mark.asyncio
+async def test_must_exist_fails_when_missing(monkeypatch):
+    monkeypatch.delenv("SOME_FLAG", raising=False)
+    ev = EnvVarEval(mk_ctx(), name="env", type="vetnode.evaluations.env_var_eval.EnvVarEval",
+                    expected={"SOME_FLAG": None})
+
+    status, metrics = await ev.check(None)
+    assert status == EvalResultStatus.FAILED
+    assert metrics["missing"] == ["SOME_FLAG"]
+
+
+@pytest.mark.asyncio
+async def test_exact_match(monkeypatch):
+    monkeypatch.setenv("FI_CXI_RX_MATCH_MODE", "hybrid")
+    ev = EnvVarEval(mk_ctx(), name="env", type="vetnode.evaluations.env_var_eval.EnvVarEval",
+                    expected={"FI_CXI_RX_MATCH_MODE": "hybrid"})
+
+    status, _ = await ev.check(None)
+    assert status == EvalResultStatus.SUCCESS
+
+
+@pytest.mark.asyncio
+async def test_mismatch(monkeypatch):
+    monkeypatch.setenv("FI_CXI_RX_MATCH_MODE", "wrong")
+    ev = EnvVarEval(mk_ctx(), name="env", type="vetnode.evaluations.env_var_eval.EnvVarEval",
+                    expected={"FI_CXI_RX_MATCH_MODE": "hybrid"})
+
+    status, metrics = await ev.check(None)
+    assert status == EvalResultStatus.FAILED
+    assert "FI_CXI_RX_MATCH_MODE" in metrics["mismatched"]
+
+
+def test_yaml_populates_expected_env(tmp_path):
+    cfg = tmp_path / "cfg.yaml"
+    cfg.write_text(textwrap.dedent("""
+        name: test-config
+        scheduler: standalone
+        evals:
+          - name: env-check
+            type: vetnode.evaluations.env_var_eval.EnvVarEval
+            expected:
+              SOME_FLAG: null
+              FI_CXI_RX_MATCH_MODE: hybrid
+    """).lstrip())
+
+    # Make Configuration read this YAML
+    Configuration._yaml_file = str(cfg)
+    conf = Configuration()
+
+    ctx = build_context(conf)
+    evals = load_evals(ctx, conf.evals)
+
+    assert len(evals) == 1
+    ev = evals[0]
+
+    assert ev.expected["SOME_FLAG"] is None
+    assert ev.expected["FI_CXI_RX_MATCH_MODE"] == "hybrid"
+
+
+@pytest.mark.parametrize(
+    "env, expected_exit, expected_token",
+    [
+        ({"FI_CXI_RX_MATCH_MODE": "hybrid", "FI_MR_CACHE_MONITOR": "userfaultfd"}, 0, ": Vetted"),
+        ({"FI_CXI_RX_MATCH_MODE": "nope",   "FI_MR_CACHE_MONITOR": "userfaultfd"}, 1, ": Cordon"),
+        ({}, 1, ": Cordon"),
+    ],
+)
+def test_diagnose_env_eval(tmp_path, monkeypatch, env, expected_exit, expected_token):
+    # Set/clear relevant env vars
+    for k in ["FI_CXI_RX_MATCH_MODE", "FI_MR_CACHE_MONITOR"]:
+        monkeypatch.delenv(k, raising=False)
+    for k, v in env.items():
+        monkeypatch.setenv(k, v)
+
+    cfg = tmp_path / "cfg.yaml"
+    cfg.write_text(textwrap.dedent("""
+        name: test-config
+        scheduler: standalone
+        evals:
+          - name: env-check
+            type: vetnode.evaluations.env_var_eval.EnvVarEval
+            expected:
+              FI_CXI_RX_MATCH_MODE: hybrid
+              FI_MR_CACHE_MONITOR: userfaultfd
+    """).lstrip())
+
+    runner = CliRunner()
+    result = runner.invoke(diagnose, ["--skip-install", "--verbose", str(cfg)])
+
+    print(result.output)
+    assert result.exit_code == expected_exit
+    assert expected_token in result.output