API/YAML: Support using env vars to define file_mounts.

[concretevitamin]

# 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.

Tested (run the relevant ones):

• Code formatting: bash format.sh
• 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
• 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

[Michaelvll]

Thanks for adding this @concretevitamin! This should make our YAML interface more flexible. Left several questions.

[Michaelvll]

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]*'

[concretevitamin]

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

[Michaelvll]

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}

[concretevitamin]

Good call! Done.

[Michaelvll]

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

[concretevitamin]

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?

[Michaelvll]

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.

[concretevitamin]

Good call! Updated as such.

[Michaelvll]

Thanks for the fix @concretevitamin! The code looks good to me with a small corner case that we may need to be careful.

[Michaelvll]

nit: it seems bash replaces the envvar by looking for the envvar in the string first and match it with the existing envvar, while here we do it in a reversed way. That will cause a slight inconsistency in the behavior:

HELLO_WORLD=1
echo $HELLO_WORLD_1

The block above will output an empty string, but in our implementation, we will output 1_1.

Probably, an alternative way to do this is:

# Create a replacement function
def replace_var(match):
    var_name = match.group(1)  # Now directly accessing the variable name
    return task_envs.get(var_name, match.group(0))  # If the variable isn't in the dictionary, return it unchanged

pattern = r'\$\{?\b([a-zA-Z_][a-zA-Z0-9_]*)\b\}?'
# Use re.sub with the replacement function
result = re.sub(pattern, replace_var, file_mounts_str)

This is a very corner case, we can also leave it as is to see how people feel about it.

[concretevitamin]

Great catch! Used this & reran the smoke test.

[Michaelvll]

nice, we should have all our tests to be --cpus 2+ in the future to save cost. ; )

[concretevitamin]

• All smoke tests: pytest tests/test_smoke.py