Allow setting the workdir using environment variables (#3792)

* Init * Add a test for env override in yaml * Format * Update tests/test_yaml_parser.py Co-authored-by: Romil Bhardwaj <[email protected]> --------- Co-authored-by: Romil Bhardwaj <[email protected]>
skypilot-org · Jul 28, 2024 · 92727c7 · 92727c7
1 parent e794971
commit 92727c7
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 0 deletions.
diff --git a/sky/task.py b/sky/task.py
@@ -393,6 +393,11 @@ def from_yaml_config(
             config['service'] = _fill_in_env_vars(config['service'],
                                                   config.get('envs', {}))
 
+        # Fill in any Task.envs into workdir
+        if config.get('workdir') is not None:
+            config['workdir'] = _fill_in_env_vars(config['workdir'],
+                                                  config.get('envs', {}))
+
         task = Task(
             config.pop('name', None),
             run=config.pop('run', None),

diff --git a/tests/test_yaml_parser.py b/tests/test_yaml_parser.py
@@ -146,3 +146,14 @@ def test_invalid_empty_envs(tmp_path):
     with pytest.raises(ValueError) as e:
         Task.from_yaml(config_path)
     assert 'Environment variable \'env_key2\' is None.' in e.value.args[0]
+
+
+def test_replace_envs_in_workdir(tmpdir, tmp_path):
+    config_path = _create_config_file(
+        textwrap.dedent(f"""\
+            envs:
+                env_key1: {tmpdir}
+            workdir: $env_key1
+            """), tmp_path)
+    task = Task.from_yaml(config_path)
+    assert task.workdir == tmpdir