fix(orchestration): forward explicit runtime variables to remote Slurm

fern/versions/latest/pages/orchestration/configuration.mdx

@@ -105,9 +105,30 @@ A mapping with exactly one entry; its key is a name you choose (the examples use
 | `account` | required | Slurm account (`--account`). |
 | `hostname` | `null` | Login node to submit through over SSH. `null`, or the name of the machine you are on, means run `sbatch` locally with no SSH. |
 | `walltime` | `null` | Job time limit (`--time`). |
+| `forward_env` | `[]` | Environment variable names to forward from the submitting shell to `sbatch`. Missing variables fail submission before staging. Values are sent over SSH stdin and are not stored in job artifacts. |
 | `node_pools` | `{}` | Named groups of nodes. See below. |
 | `extra_args` | `{}` | Extra job-level `#SBATCH` directives, written without leading dashes: `{comment: my-run}` becomes `#SBATCH --comment=my-run`. |
 
+Use `forward_env` with `runtime:` for credentials needed by a remote job:
+
+```yaml
+compute:
+  cluster:
+    type: slurm
+    hostname: login.example.com
+    account: my-account
+    forward_env: [HF_TOKEN, INFERENCE_API_KEY]
+driver:
+  env:
+    HF_TOKEN: runtime:HF_TOKEN
+    INFERENCE_API_KEY: runtime:INFERENCE_API_KEY
+```
+
+Only listed names are forwarded. Leave scheduler-provided values such as
+`CUDA_VISIBLE_DEVICES` out of this list. `--resolve-only` and dry runs do not
+require forwarded values to be set. Slurm must export the submission environment
+to the job; an explicit `--export=NONE` prevents that inheritance.
+
 #### Node pools
 
 | Field | Default | Description |

nemo_gym/orchestration/api.py

@@ -345,6 +345,18 @@ class SlurmComputeConfig(BaseComputeConfig):
     type: Literal["slurm"]
     account: str
     hostname: str | None = None  # None means we're already on the login node; skip SSH.
+    # Submit-side variables explicitly exported before sbatch, without writing values
+    # into job.sh or the resolved config. Runtime-only Slurm variables stay untouched.
+    forward_env: list[str] = []
+
+    @field_validator("forward_env")
+    @classmethod
+    def _validate_forward_env(cls, names: list[str]) -> list[str]:
+        for name in names:
+            if not re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*", name):
+                raise ValueError(f"Invalid forward_env variable name: {name!r}")
+        return names
+
     walltime: str | None = None
     node_pools: dict[str, NodePool] = {}
     extra_args: dict[str, str] = {}  # Job-level #SBATCH directives (e.g. --comment, --mail-user).

nemo_gym/orchestration/executors/slurm.py

@@ -14,6 +14,7 @@
 # limitations under the License.
 
 import getpass
+import os
 import re
 import shlex
 import tempfile
@@ -148,6 +149,14 @@ def _validate_mounts(config: SubmitConfig, conn: Connection) -> None:
         raise ValueError(f"Mount src paths do not exist:\n{details}")
 
 
+def _submission_env_exports(compute: SlurmComputeConfig) -> list[str]:
+    """Forward explicitly selected values over SSH stdin, never into artifacts."""
+    missing = [name for name in compute.forward_env if name not in os.environ]
+    if missing:
+        raise ValueError(f"compute.forward_env variables are unset: {', '.join(missing)}")
+    return [f"export {name}={shlex.quote(os.environ[name])}" for name in sorted(set(compute.forward_env))]
+
+
 class SlurmExecutor(BaseExecutor):
     """Slurm executor for Pyxis-enabled clusters (https://github.com/NVIDIA/pyxis).
 
@@ -161,6 +170,7 @@ class SlurmExecutor(BaseExecutor):
     def run(self, config: SubmitConfig, *, dry_run: bool = False) -> SubmissionRecord | None:
         compute = next(iter(config.compute.values()))
         cluster = next(iter(config.compute))
+        submission_env = _submission_env_exports(compute) if not dry_run else []
         benchmark_names = list(config.driver.benchmarks)
         _validate_benchmark_names(benchmark_names)
         # The remote lookup overlaps the other validation, staging, the SSH connect and the mount
@@ -190,7 +200,8 @@ def run(self, config: SubmitConfig, *, dry_run: bool = False) -> SubmissionRecor
                 conn.copy(staging, remote_run_dir)
                 token_export = [f"export {config.otel.token_env}={shlex.quote(token)}"] if token is not None else []
                 output = conn.run(
-                    token_export
+                    submission_env
+                    + token_export
                     + [_sbatch_command(name, remote_run_dir / name / "job.sh") for name in benchmark_names]
                 )
                 record = self._build_record(cluster, compute, gym_job_id, now, remote_run_dir, benchmark_names, output)

tests/unit_tests/test_slurm_executor.py

@@ -422,3 +422,58 @@ def reject(config):
         SlurmExecutor().run(_submit_config(tmp_path, ["bench_a"]))
 
     assert conn.commands == []
+
+
+def test_runtime_credentials_reach_remote_sbatch_without_being_persisted(tmp_path, monkeypatch):
+    import subprocess
+
+    secret = "test secret with ' quotes and $(false)"
+    monkeypatch.setenv("SUBMIT_JUDGE_KEY", secret)
+    monkeypatch.setenv("UNREFERENCED_SECRET", "must stay local")
+    monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / "cache"))
+    config = _submit_config(tmp_path, ["bench_a"])
+    config.driver.env = {"JUDGE_KEY": "runtime:SUBMIT_JUDGE_KEY"}
+    config.compute["hsg"].forward_env = ["SUBMIT_JUDGE_KEY"]
+    conn = _FakeConnection(["__GYM_JOB:bench_a:0:111 "])
+    _install(monkeypatch, conn)
+
+    record = SlurmExecutor().run(config)
+
+    commands = conn.commands[0]
+    exports = [command for command in commands if command.startswith("export ")]
+    probe = subprocess.run(
+        ["bash", "-s"],
+        input="\n".join(exports + ['printf %s "$SUBMIT_JUDGE_KEY"']),
+        text=True,
+        capture_output=True,
+        check=True,
+        env={},
+    )
+    assert probe.stdout == secret
+    assert "UNREFERENCED_SECRET" not in "\n".join(commands)
+    for path in Path(record.run_dir).rglob("*"):
+        if path.is_file():
+            assert secret not in path.read_text()
+
+
+def test_runtime_variable_is_not_implicitly_forwarded(tmp_path, monkeypatch):
+    monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0,1")
+    config = _submit_config(tmp_path, ["bench_a"])
+    config.driver.env = {"CUDA_VISIBLE_DEVICES": "runtime:CUDA_VISIBLE_DEVICES"}
+    assert slurm_module._submission_env_exports(config.compute["hsg"]) == []
+
+
+def test_missing_forwarded_variable_fails_before_submission(tmp_path, monkeypatch):
+    monkeypatch.delenv("MISSING_JUDGE_KEY", raising=False)
+    config = _submit_config(tmp_path, ["bench_a"])
+    config.compute["hsg"].forward_env = ["MISSING_JUDGE_KEY"]
+    with pytest.raises(ValueError, match="MISSING_JUDGE_KEY"):
+        SlurmExecutor().run(config)
+
+
+@pytest.mark.parametrize("name", ["BAD;false", "X=Y", "", "1NAME"])
+def test_forwarded_variable_name_is_validated(tmp_path, name):
+    raw = _submit_config(tmp_path, ["bench_a"]).model_dump()
+    raw["compute"]["hsg"]["forward_env"] = [name]
+    with pytest.raises(ValueError, match="Invalid forward_env"):
+        SubmitConfig.model_validate(raw)