|
17 | 17 | ) |
18 | 18 |
|
19 | 19 |
|
20 | | -def test_build_runtime_env_has_required_keys(): |
21 | | - """PYTHONPATH (for megatron) and CUDA_DEVICE_MAX_CONNECTIONS (for Megatron TP) |
22 | | - are required — they mirror train.sh's inline python snippet and breaking them |
23 | | - breaks every downstream Ray worker.""" |
| 20 | +def test_build_runtime_env_required_keys_and_wandb_opt_in(): |
| 21 | + """PYTHONPATH + CUDA_DEVICE_MAX_CONNECTIONS are always present; wandb keys |
| 22 | + are opt-in, forwarded only when set in the parent env.""" |
24 | 23 | runner = SlimeRunner(**REQUIRED_KWARGS, megatron_dir="/opt/megatron") |
25 | 24 |
|
26 | 25 | with patch.dict("os.environ", {}, clear=True): |
27 | | - env = runner._build_runtime_env() |
28 | | - |
29 | | - assert env == { |
| 26 | + env_no_wandb = runner._build_runtime_env() |
| 27 | + assert env_no_wandb == { |
30 | 28 | "env_vars": { |
31 | 29 | "PYTHONPATH": "/opt/megatron", |
32 | 30 | "CUDA_DEVICE_MAX_CONNECTIONS": "1", |
33 | 31 | } |
34 | 32 | } |
35 | 33 |
|
36 | | - |
37 | | -def test_build_runtime_env_forwards_wandb_when_set(): |
38 | | - """Wandb keys are opt-in: absent from env → not in runtime_env.""" |
39 | | - runner = SlimeRunner(**REQUIRED_KWARGS) |
40 | | - |
41 | 34 | with patch.dict("os.environ", {"WANDB_API_KEY": "abc", "WANDB_ENTITY": "me"}, clear=True): |
42 | | - env = runner._build_runtime_env() |
43 | | - |
44 | | - assert env["env_vars"]["WANDB_API_KEY"] == "abc" |
45 | | - assert env["env_vars"]["WANDB_ENTITY"] == "me" |
| 35 | + env_with_wandb = runner._build_runtime_env() |
| 36 | + assert env_with_wandb["env_vars"]["WANDB_API_KEY"] == "abc" |
| 37 | + assert env_with_wandb["env_vars"]["WANDB_ENTITY"] == "me" |
46 | 38 |
|
47 | 39 |
|
48 | 40 | def test_from_yaml_round_trips(tmp_path: Path): |
|
0 commit comments