Repository: skypilot-org/skypilot
Pull request: API/YAML: Support using env vars to define file_mounts.
<!-- Describe the changes in this PR -->

```
# Use env vars in the "name" (bucket name) and "source" (local dir/file to
# upload) fields.
#
# Both syntaxes work: ${MY_BUCKET} and $MY_BUCKET.
file_mounts:
  /mydir:
    name: ${MY_BUCKET}  # Name of the bucket.
    store: gcs
    mode: MOUNT

  /another-dir:
    name: ${MY_BUCKET}-2
    source: ["~/${MY_LOCAL_PATH}"]
    store: gcs
    mode: MOUNT
```
Overriding with `sky launch --env` works too.

Also fixes `sky launch` parsing errors: 
- Do not load YAML twice (which creates storage twice)
- Properly load empty / name-only YAMLs

Fixes #2093.

<!-- Describe the tests ran -->
<!-- Unit tests (tests/test_*.py) are part of GitHub CI; below are tests that launch on the cloud. -->

Tested (run the relevant ones):

- [x] Code formatting: `bash format.sh`
- [x] Any manual or new tests for this PR (please specify below)
  - `pytest tests/test_smoke.py::test_using_file_mounts_with_env_vars`
  - `pytest tests/test_optimizer_dryruns.py`
- [x] All smoke tests: `pytest tests/test_smoke.py` 
- [ ] Relevant individual smoke tests: `pytest tests/test_smoke.py::test_fill_in_the_name` 
- [ ] Backward compatibility tests: `bash tests/backward_comaptibility_tests.sh`


Problem statement:
Use environment variable to define file mount
Hi and thanks for this great project !

I am in a situation where I would like to mount some directory in a bucket which changes based on the value of an environment variable. I am unsure how to do this with SkyPilot as defining a file mount would typically require that I specify the corresponding URL in my task YAML file beforehand.

Put simply I'd like to do
```yaml
file_mounts:
  /some/folder:
    source: $URL_IN_MY_ENVVAR
    mode: MOUNT
```

And be able to set `URL_IN_MY_ENVVAR ` at runtime.
Am I missing something or would it be a feature that may be of interest?

Hints from issue discussion:
Thanks for filing this report @QLutz and welcome to SkyPilot!

One solution to this could be to use `envsubst` to create a temporary yaml file that you can launch. Would that work?
```
export URL_IN_MY_ENVVAR=s3://mybucket/
envsubst < template.yaml > run.yaml       
sky launch run.yaml
```
Alternatively, you can also use the SkyPilot [python API](https://skypilot.readthedocs.io/en/latest/reference/api.html) if you find that easier to integrate into your workflow.

We would love to hear a bit more about your use case for this to see if we can add this as a feature. Can you share a little more about your workload and how you are using SkyPilot?
Thanks for your answer @romilbhardwaj 

My current workflow consists in setting some parameters for my deployment at launch, e.g.:
```bash
sky launch -c my_deployment --env ENV1=val1 --env ENV2=val2 --gpus=T4:2 my_task.yaml
```
with all parameters having default values (that I may or may not want to override) in `my_task.yaml` which would resemble:
```yaml
resources:
  accelerators: T4:1
  disk_size: 250

run: |
  ENV1=${ENV1:-default_val1} \
  ENV2=${ENV2:-default_val2} \
  ENV3=${ENV3:-default_val3} \
  sh ./scripts/run.sh
```
I find this to be very convenient as it allows for seamless customization of my deployment. Mounts however, cannot be set as easily in the CLI.
I understand that many solutions exist (and the one you offer is probably the simplest one at the moment) but I find that the current simplicity of Task YAMLs is very elegant and allows for highly maintainable workflows (using the Python API while more versatile is quite a bit more cumbersome).
It's definitely a convenience thing and I get that you may have more pressing issues to handle. Nevertheless, if you find that allowing for envvars to be used in the file mounts is not that hard, I think it has some value.
Thanks @QLutz! This certainly makes sense. We do want to get to it; before we have bandwidth any contribution from you guys is definitely very welcome.

Commit message: Fixes

Patch to review:
diff --git a/sky/cli.py b/sky/cli.py
index 1fcb9786f17..eeef88243f9 100644
--- a/sky/cli.py
+++ b/sky/cli.py
@@ -948,7 +948,11 @@ def _check_yaml(entrypoint: str) -> Tuple[bool, Optional[Dict[str, Any]]]:
     try:
         with open(entrypoint, 'r') as f:
             try:
-                config = list(yaml.safe_load_all(f))[0]
+                config = list(yaml.safe_load_all(f))
+                if config:
+                    config = config[0]
+                else:
+                    config = {}
                 if isinstance(config, str):
                     # 'sky exec cluster ./my_script.sh'
                     is_yaml = False
@@ -1041,7 +1045,7 @@ def _make_task_or_dag_from_entrypoint_with_overrides(
     if is_yaml:
         assert entrypoint is not None
         usage_lib.messages.usage.update_user_task_yaml(entrypoint)
-        dag = dag_utils.load_chain_dag_from_yaml(entrypoint)
+        dag = dag_utils.load_chain_dag_from_yaml(entrypoint, env_overrides=env)
         if len(dag.tasks) > 1:
             # When the dag has more than 1 task. It is unclear how to
             # override the params for the dag. So we just ignore the
@@ -1052,7 +1056,9 @@ def _make_task_or_dag_from_entrypoint_with_overrides(
                     'since the yaml file contains multiple tasks.',
                     fg='yellow')
             return dag
-        task = sky.Task.from_yaml(entrypoint)
+        assert len(dag.tasks) == 1, (
+            f'If you see this, please file an issue; tasks: {dag.tasks}')
+        task = dag.tasks[0]
     else:
         task = sky.Task(name='sky-cmd', run=entrypoint)
         task.set_resources({sky.Resources()})
diff --git a/sky/clouds/gcp.py b/sky/clouds/gcp.py
index 8f2ac579fce..c9d44f54cf8 100644
--- a/sky/clouds/gcp.py
+++ b/sky/clouds/gcp.py
@@ -622,7 +622,7 @@ def get_credential_file_mounts(self) -> Dict[str, str]:
         # be modified on the remote cluster by ray causing authentication
         # problems. The backup file will be updated to the remote cluster
         # whenever the original file is not empty and will be applied
-        # appropriately on the remote cluster when neccessary.
+        # appropriately on the remote cluster when necessary.
         if (os.path.exists(os.path.expanduser(GCP_CONFIG_PATH)) and
                 os.path.getsize(os.path.expanduser(GCP_CONFIG_PATH)) > 0):
             subprocess.run(f'cp {GCP_CONFIG_PATH} {GCP_CONFIG_SKY_BACKUP_PATH}',
diff --git a/sky/data/storage.py b/sky/data/storage.py
index 48dbe3134d3..bb2c9ab2ebb 100644
--- a/sky/data/storage.py
+++ b/sky/data/storage.py
@@ -1321,11 +1321,10 @@ def _validate(self):
         if not _is_storage_cloud_enabled(str(clouds.GCP())):
             with ux_utils.print_exception_no_traceback():
                 raise exceptions.ResourcesUnavailableError(
-                    'Storage \'store: gcs\' specified, but ' \
-                    'GCP access is disabled. To fix, enable '\
-                    'GCP by running `sky check`. '\
-                    'More info: https://skypilot.readthedocs.io/en/latest/getting-started/installation.html.' # pylint: disable=line-too-long
-                    )
+                    'Storage \'store: gcs\' specified, but '
+                    'GCP access is disabled. To fix, enable '
+                    'GCP by running `sky check`. '
+                    'More info: https://skypilot.readthedocs.io/en/latest/getting-started/installation.html.')  # pylint: disable=line-too-long
 
     @classmethod
     def validate_name(cls, name) -> str:
diff --git a/sky/execution.py b/sky/execution.py
index 810a4980a20..7b6beca7ff7 100644
--- a/sky/execution.py
+++ b/sky/execution.py
@@ -337,6 +337,7 @@ def _execute(
                 backend.sync_workdir(handle, task.workdir)
 
         if Stage.SYNC_FILE_MOUNTS in stages:
+            task.fill_env_vars_in_storage_mounts()
             backend.sync_file_mounts(handle, task.file_mounts,
                                      task.storage_mounts)
 
diff --git a/sky/task.py b/sky/task.py
index bc2ff18b66e..8c5e9275f7f 100644
--- a/sky/task.py
+++ b/sky/task.py
@@ -68,6 +68,45 @@ def _is_valid_env_var(name: str) -> bool:
     return bool(re.fullmatch(_VALID_ENV_VAR_REGEX, name))
 
 
+def _fill_in_env_vars_in_storage_config(
+        storage_config: Dict[str, Any], task_envs: Dict[str,
+                                                        str]) -> Dict[str, Any]:
+    """Detects env vars in storage_config and fills them with task.envs.
+
+    This func tries to replace two fields in storage_config: 'source' and
+    'name'.
+
+    Env vars of the following forms are detected:
+        - ${ENV}
+        - $ENV
+    where <ENV> must appear in task.envs.
+
+    Example:
+        {'mode': 'MOUNT', 'name': '${GSBUCKET}', 'store': 'gcs'}
+        -> {mode: 'MOUNT', name: 'mybucket', store: 'gcs'}
+    """
+
+    # TODO(zongheng): support ${ENV:-default}?
+    def _substitute_env_vars(target: Union[str, List[str]],
+                             envs: Dict[str, str]) -> Union[str, List[str]]:
+        for key, value in envs.items():
+            pattern = r'\$\{?\b' + key + r'\b\}?'
+            env_var_pattern = re.compile(pattern)
+            if isinstance(target, str):
+                target = env_var_pattern.sub(value, target)
+            else:
+                # source: [/local/path, /local/path2]
+                target = [env_var_pattern.sub(value, t) for t in target]
+        return target
+
+    fields_to_fill_in = ('name', 'source')
+    for field in fields_to_fill_in:
+        if field in storage_config:
+            storage_config[field] = _substitute_env_vars(
+                storage_config[field], task_envs)
+    return storage_config
+
+
 class Task:
     """Task: a computation to be run on the cloud."""
 
@@ -231,7 +270,19 @@ def _validate(self):
                         f'a symlink to a directory). {self.workdir} not found.')
 
     @staticmethod
-    def from_yaml_config(config: Dict[str, Any]) -> 'Task':
+    def from_yaml_config(
+        config: Dict[str, Any],
+        env_overrides: Optional[List[Tuple[str, str]]] = None,
+    ) -> 'Task':
+        if env_overrides is not None:
+            # We must override env vars before constructing the Task, because
+            # the Storage object creation is eager and it (its name/source
+            # fields) may depend on env vars. The eagerness / how we construct
+            # Task's from entrypoint (YAML, CLI args) should be fixed (FIXME).
+            new_envs = config.get('envs', {})
+            new_envs.update(env_overrides)
+            config['envs'] = new_envs
+
         backend_utils.validate_schema(config, schemas.get_task_schema(),
                                       'Invalid task YAML: ')
 
@@ -269,10 +320,12 @@ def from_yaml_config(config: Dict[str, Any]) -> 'Task':
         all_storages = fm_storages
         for storage in all_storages:
             mount_path = storage[0]
-            assert mount_path, \
-                'Storage mount path cannot be empty.'
+            assert mount_path, 'Storage mount path cannot be empty.'
             try:
-                storage_obj = storage_lib.Storage.from_yaml_config(storage[1])
+                storage_config = _fill_in_env_vars_in_storage_config(
+                    storage[1], task.envs)
+                storage_obj = storage_lib.Storage.from_yaml_config(
+                    storage_config)
             except exceptions.StorageSourceError as e:
                 # Patch the error message to include the mount path, if included
                 e.args = (e.args[0].replace('<destination_path>',
@@ -670,6 +723,15 @@ def update_storage_mounts(
         task_storage_mounts.update(storage_mounts)
         return self.set_storage_mounts(task_storage_mounts)
 
+    def fill_env_vars_in_storage_mounts(self) -> None:
+        """Updates self.storage_mounts with self.envs."""
+        new_storage_mounts = {}
+        for mnt_path, storage in self.storage_mounts.items():
+            new_storage_mounts[mnt_path] = storage_lib.Storage.from_yaml_config(
+                _fill_in_env_vars_in_storage_config(storage.to_yaml_config(),
+                                                    self.envs))
+        self.update_storage_mounts(new_storage_mounts)
+
     def get_preferred_store_type(self) -> storage_lib.StoreType:
         # TODO(zhwu, romilb): The optimizer should look at the source and
         #  destination to figure out the right stores to use. For now, we
diff --git a/sky/utils/common_utils.py b/sky/utils/common_utils.py
index 8f51c372034..1a281b07a06 100644
--- a/sky/utils/common_utils.py
+++ b/sky/utils/common_utils.py
@@ -182,7 +182,11 @@ def read_yaml(path) -> Dict[str, Any]:
 def read_yaml_all(path: str) -> List[Dict[str, Any]]:
     with open(path, 'r') as f:
         config = yaml.safe_load_all(f)
-        return list(config)
+        configs = list(config)
+        if not configs:
+            # Empty YAML file.
+            return [{}]
+        return configs
 
 
 def dump_yaml(path, config) -> None:
diff --git a/sky/utils/dag_utils.py b/sky/utils/dag_utils.py
index d8a8ed03d8c..446bfbd13e5 100644
--- a/sky/utils/dag_utils.py
+++ b/sky/utils/dag_utils.py
@@ -1,11 +1,24 @@
 """Utilities for loading and dumping DAGs from/to YAML files."""
+from typing import List, Optional, Tuple
+
 from sky import dag as dag_lib
 from sky import task as task_lib
 from sky.backends import backend_utils
 from sky.utils import common_utils
 
 
-def load_chain_dag_from_yaml(path: str) -> dag_lib.Dag:
+def load_chain_dag_from_yaml(
+    path: str,
+    env_overrides: Optional[List[Tuple[str, str]]] = None,
+) -> dag_lib.Dag:
+    """Loads a chain DAG from a YAML file.
+
+    Has special handling for an initial section in YAML that contains only the
+    'name' field, which is the DAG name.
+
+    Returns: a chain Dag with 1 or more tasks (an empty entrypoint would create
+      a trivial task).
+    """
     configs = common_utils.read_yaml_all(path)
     dag_name = None
     if set(configs[0].keys()) == {'name'}:
@@ -14,12 +27,16 @@ def load_chain_dag_from_yaml(path: str) -> dag_lib.Dag:
     elif len(configs) == 1:
         dag_name = configs[0].get('name')
 
+    if len(configs) == 0:
+        # YAML has only `name: xxx`. Still instantiate a task.
+        configs = [{'name': dag_name}]
+
     current_task = None
     with dag_lib.Dag() as dag:
         for task_config in configs:
             if task_config is None:
                 continue
-            task = task_lib.Task.from_yaml_config(task_config)
+            task = task_lib.Task.from_yaml_config(task_config, env_overrides)
             if current_task is not None:
                 current_task >> task  # pylint: disable=pointless-statement
             current_task = task
diff --git a/tests/test_optimizer_dryruns.py b/tests/test_optimizer_dryruns.py
index a3ca0798979..b66fc130dff 100644
--- a/tests/test_optimizer_dryruns.py
+++ b/tests/test_optimizer_dryruns.py
@@ -9,31 +9,38 @@
 from sky import exceptions
 
 
-def _test_parse_cpus(spec, expected_cpus):
+def _test_parse_task_yaml(spec: str, test_fn):
+    """Tests parsing a task from a YAML spec and running a test_fn."""
     with tempfile.NamedTemporaryFile('w') as f:
         f.write(spec)
         f.flush()
         with sky.Dag():
             task = sky.Task.from_yaml(f.name)
-            assert list(task.resources)[0].cpus == expected_cpus
+            test_fn(task)
+
+
+def _test_parse_cpus(spec, expected_cpus):
+
+    def test_fn(task):
+        assert list(task.resources)[0].cpus == expected_cpus
+
+    _test_parse_task_yaml(spec, test_fn)
 
 
 def _test_parse_memory(spec, expected_memory):
-    with tempfile.NamedTemporaryFile('w') as f:
-        f.write(spec)
-        f.flush()
-        with sky.Dag():
-            task = sky.Task.from_yaml(f.name)
-            assert list(task.resources)[0].memory == expected_memory
+
+    def test_fn(task):
+        assert list(task.resources)[0].memory == expected_memory
+
+    _test_parse_task_yaml(spec, test_fn)
 
 
 def _test_parse_accelerators(spec, expected_accelerators):
-    with tempfile.NamedTemporaryFile('w') as f:
-        f.write(spec)
-        f.flush()
-        with sky.Dag():
-            task = sky.Task.from_yaml(f.name)
-            assert list(task.resources)[0].accelerators == expected_accelerators
+
+    def test_fn(task):
+        assert list(task.resources)[0].accelerators == expected_accelerators
+
+    _test_parse_task_yaml(spec, test_fn)
 
 
 # Monkey-patching is required because in the test environment, no cloud is
@@ -547,3 +554,24 @@ def test_invalid_num_nodes():
                 task = sky.Task()
                 task.num_nodes = invalid_value
             assert 'num_nodes should be a positive int' in str(e.value)
+
+
+def test_parse_empty_yaml():
+    spec = textwrap.dedent("""\
+        """)
+
+    def test_fn(task):
+        assert task.num_nodes == 1
+
+    _test_parse_task_yaml(spec, test_fn)
+
+
+def test_parse_name_only_yaml():
+    spec = textwrap.dedent("""\
+        name: test_task
+        """)
+
+    def test_fn(task):
+        assert task.name == 'test_task'
+
+    _test_parse_task_yaml(spec, test_fn)
diff --git a/tests/test_smoke.py b/tests/test_smoke.py
index 27be66877cb..dcc159a9888 100644
--- a/tests/test_smoke.py
+++ b/tests/test_smoke.py
@@ -4,7 +4,7 @@
 # Run all tests except for AWS and Lambda Cloud
 # > pytest tests/test_smoke.py
 #
-# Terminate failed clusetrs after test finishes
+# Terminate failed clusters after test finishes
 # > pytest tests/test_smoke.py --terminate-on-failure
 #
 # Re-run last failed tests
@@ -738,6 +738,28 @@ def test_scp_file_mounts():
     run_one_test(test)
 
 
+def test_using_file_mounts_with_env_vars(generic_cloud: str):
+    name = _get_cluster_name()
+    test_commands = [
+        *storage_setup_commands,
+        (f'sky launch -y -c {name} --cpus 2+ --cloud {generic_cloud} '
+         'examples/using_file_mounts_with_env_vars.yaml'),
+        f'sky logs {name} 1 --status',  # Ensure the job succeeded.
+        # Override with --env:
+        (f'sky launch -y -c {name}-2 --cpus 2+ --cloud {generic_cloud} '
+         'examples/using_file_mounts_with_env_vars.yaml '
+         '--env MY_LOCAL_PATH=tmpfile'),
+        f'sky logs {name}-2 1 --status',  # Ensure the job succeeded.
+    ]
+    test = Test(
+        'using_file_mounts_with_env_vars',
+        test_commands,
+        f'sky down -y {name} {name}-2',
+        timeout=20 * 60,  # 20 mins
+    )
+    run_one_test(test)
+
+
 # ---------- storage ----------
 @pytest.mark.aws
 def test_aws_storage_mounts():


Code review comments:
diff_hunk: @@ -68,6 +68,45 @@ def _is_valid_env_var(name: str) -> bool:
     return bool(re.fullmatch(_VALID_ENV_VAR_REGEX, name))
 
 
+def _fill_in_env_vars_in_storage_config(
+        storage_config: Dict[str, Any], task_envs: Dict[str,
+                                                        str]) -> Dict[str, Any]:
+    """Detects env vars in storage_config and fills them with task.envs.
+
+    This func tries to replace two fields in storage_config: 'source' and
+    'name'.
+
+    Env vars of the following forms are detected:
+        - ${ENV}
+        - $ENV
+    where <ENV> must appear in task.envs.
+
+    Example:
+        {'mode': 'MOUNT', 'name': '${GSBUCKET}', 'store': 'gcs'}
+        -> {mode: 'MOUNT', name: 'mybucket', store: 'gcs'}
+    """
+
+    # TODO(zongheng): support ${ENV:-default}?
+    def _substitute_env_vars(target: Union[str, List[str]],
+                             envs: Dict[str, str]) -> Union[str, List[str]]:
+        for key, value in envs.items():
+            pattern = r'\$\{?\b' + key + r'\b\}?'
path: sky/task.py
text: @user1:
We may need to add a check for the key of the `task.envs` somewhere in our code, i.e. we need to make sure the `key ~ '[a-zA-Z][a-zA-Z0-9]*'`

@author:
This will get checked by storage.py's validate_name().

@user1:
Ahh, I was trying to say that if the user provides a envvar like the following:
```
envs:
  hello world: some_value

run: |
  echo ${hello world}
```
Here the run section will fail, as it is not a valid syntax in bash. However, if this env var is used in the `file_mounts` section, the envvar will be replaced correctly.
```
file_mounts:
  /dst:
    source: ${hello world}
```



@author:
Good call! Done.

diff_hunk: @@ -269,10 +320,12 @@ def from_yaml_config(config: Dict[str, Any]) -> 'Task':
         all_storages = fm_storages
         for storage in all_storages:
             mount_path = storage[0]
-            assert mount_path, \
-                'Storage mount path cannot be empty.'
+            assert mount_path, 'Storage mount path cannot be empty.'
             try:
-                storage_obj = storage_lib.Storage.from_yaml_config(storage[1])
+                storage_config = _fill_in_env_vars_in_storage_config(
+                    storage[1], task.envs)
+                storage_obj = storage_lib.Storage.from_yaml_config(
+                    storage_config)
path: sky/task.py
text: @user1:
Instead of overriding the `storage_config`, should we just replace all the env var in the `file_mounts` section? For example, in L285 we do:
```
if config.get('file_mounts') is not None:
    file_mounts_str = yaml.dumps(config['file_mounts'])
    new_file_mounts_str = _replace_env_var(file_mounts_str, config.get('envs', {}))
    config['file_mounts'] = yaml.loads(new_file_mounts_str)
```

With that, the user can do something like the following:
```
file_mounts:
  /model_path/llama-${SIZE}b: s3://llama-weights/llama-${SIZE}b
```

@author:
One worry with string replacement is that it'll make some weird things "work":
```
file_mounts:
  /path:
    mo${empty_envvar}de: xxxx
```
I think we should make replacement more structured.

That said, supporting env vars in src/dst seems a useful case. How about we support it when users ask?

@user1:
Sounds good, though I still think the replacement for `dst` and `src` might be quite important, as that was one thing we would like to have for the Vicuna for downloading the correct llama weight from the bucket based on the envvar.

For the problem you mentioned, we can do the schema validation before the replacement happens to avoid that issue.

@author:
Good call! Updated as such.