Repository: skypilot-org/skypilot
Pull request: Add support for env files. 
<!-- Describe the changes in this PR -->

Fixes #2295

A new argument is added `--env-file`.  This reads from a dotenv file to read multiple environment variables.  These are merged with any values passed via `--env` arguments with `--env` taking precedence in any conflicts.

<!-- Describe the tests ran -->

No unittests seem to exist for this.  Manually ran the command:

```sh
sky launch -c worker-unified-llama13b --env-file .env sky/fastchat_unified_llama13b.yaml
```

with a local config that depends on two env variables within `.env`.

<!-- 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)
- [ ] 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:
[Feature Request] Support passing env variables via a .env file
When launching or starting a new job, env variables currently need to be embedded in the config file or passed in through multiple `--env` arguments.  In some setups, its more convenient to have multiple env variables in a `.env` file.  Much like Docker's support, it'd be nice to be able to do 

```sh
sky launch -c some-cluster --env-file .env my-job.yaml
```

Instead of having to write all the env variables each time.

This is probably fairly simple with support for reading values from the file using `python-dotenv`

Hints from issue discussion:
Thanks for the great request @fozziethebeat! We got a similar request in #2151 as well. We will increase the priority for this. : )
Ah yeah that's fairly similar.  When I use workdir, copying over the env file works well.  But sometimes I want to skip workdir since my workdir can be a mess.
Good point! We are happy to support it. Would you like to contribute a PR? It would be very appreciated. : )
Yeah, i'll give it a shot!
Notes to self for where to implement the change:
  - [ ] `spot_launch`
  - [ ] `benchmark_launch`
  - [ ] `launch`
  - [ ] `exec`

In all cases we have to check the contents of an `env-file` argument and override them with `env` values if they conflict (`--env` has more precedence).  I'm pretty sure this is what Docker does so might as well be consistent with that.

Commit message: Add support for env files.  Env args will override any env file values and be passed on to override any yaml values

Patch to review:
diff --git a/sky/cli.py b/sky/cli.py
index 3136a16678a..5dd15e24c8d 100644
--- a/sky/cli.py
+++ b/sky/cli.py
@@ -25,6 +25,7 @@
 """
 import copy
 import datetime
+from dotenv import dotenv_values
 import functools
 import multiprocessing
 import os
@@ -330,6 +331,15 @@ def _parse_env_var(env_var: str) -> Tuple[str, str]:
             'or KEY.')
     return ret[0], ret[1]
 
+def _merge_env_vars(env_dict: Optional[Dict[str, str]],
+                    env_list: List[Tuple[str, str]]) -> List[Tuple[str, str]]:
+    """Merges all values from env_list into env_dict, overwriting any old values."""
+    if not env_dict:
+        return env_list
+    for (key, value) in env_list:
+        env_dict[key] = value
+    return list(env_dict.items())
+
 
 _TASK_OPTIONS = [
     click.option('--name',
@@ -382,6 +392,17 @@ def _parse_env_var(env_var: str) -> Tuple[str, str]:
                  default=None,
                  help=('Custom image id for launching the instances. '
                        'Passing "none" resets the config.')),
+    click.option(
+        '--env-file',
+        required=False,
+        type=dotenv_values,
+        help="""\
+        Path to a dotenv file with environment variables to set on the remote
+        node.
+
+        If any values from ``--env-file`` conflict with values set by
+        ``--env``, the ``--env`` value will be preferred."""
+    ),
     click.option(
         '--env',
         required=False,
@@ -1308,6 +1329,7 @@ def launch(
     use_spot: Optional[bool],
     image_id: Optional[str],
     env: List[Tuple[str, str]],
+    env_file: Optional[Dict[str, str]],
     disk_size: Optional[int],
     disk_tier: Optional[str],
     idle_minutes_to_autostop: Optional[int],
@@ -1326,6 +1348,7 @@ def launch(
     In both cases, the commands are run under the task's workdir (if specified)
     and they undergo job queue scheduling.
     """
+    env = _merge_env_vars(env_file, env)
     backend_utils.check_cluster_name_not_reserved(
         cluster, operation_str='Launching tasks on it')
     if backend_name is None:
@@ -1424,6 +1447,7 @@ def exec(
     use_spot: Optional[bool],
     image_id: Optional[str],
     env: List[Tuple[str, str]],
+    env_file: Optional[Dict[str, str]],
 ):
     # NOTE(dev): Keep the docstring consistent between the Python API and CLI.
     """Execute a task or a command on a cluster (skip setup).
@@ -1483,6 +1507,7 @@ def exec(
       sky exec mycluster --env WANDB_API_KEY python train_gpu.py
 
     """
+    env = _merge_env_vars(env_file, env)
     backend_utils.check_cluster_name_not_reserved(
         cluster, operation_str='Executing task on it')
     handle = global_user_state.get_handle_from_cluster_name(cluster)
@@ -3460,6 +3485,7 @@ def spot_launch(
     image_id: Optional[str],
     spot_recovery: Optional[str],
     env: List[Tuple[str, str]],
+    env_file: Optional[Dict[str, str]],
     disk_size: Optional[int],
     disk_tier: Optional[str],
     detach_run: bool,
@@ -3480,6 +3506,7 @@ def spot_launch(
 
       sky spot launch 'echo hello!'
     """
+    env = _merge_env_vars(env_file, env)
     task_or_dag = _make_task_or_dag_from_entrypoint_with_overrides(
         entrypoint,
         name=name,
@@ -3908,6 +3935,7 @@ def benchmark_launch(
     use_spot: Optional[bool],
     image_id: Optional[str],
     env: List[Tuple[str, str]],
+    env_file: Optional[Dict[str, str]],
     disk_size: Optional[int],
     disk_tier: Optional[str],
     idle_minutes_to_autostop: Optional[int],
@@ -3920,6 +3948,7 @@ def benchmark_launch(
     Alternatively, specify the benchmarking resources in your YAML (see doc),
     which allows benchmarking on many more resource fields.
     """
+    env = _merge_env_vars(env_file, env)
     record = benchmark_state.get_benchmark_from_name(benchmark)
     if record is not None:
         raise click.BadParameter(f'Benchmark {benchmark} already exists. '
diff --git a/sky/setup_files/setup.py b/sky/setup_files/setup.py
index 6a738bdad38..773d85deeff 100644
--- a/sky/setup_files/setup.py
+++ b/sky/setup_files/setup.py
@@ -86,6 +86,7 @@ def parse_readme(readme: str) -> str:
     # PrettyTable with version >=2.0.0 is required for the support of
     # `add_rows` method.
     'PrettyTable>=2.0.0',
+    'python-dotenv==1.0.0',
     # Lower version of ray will cause dependency conflict for
     # click/grpcio/protobuf.
     'ray[default]>=2.2.0,<=2.4.0',


Code review comments:
diff_hunk: @@ -86,6 +86,7 @@ def parse_readme(readme: str) -> str:
     # PrettyTable with version >=2.0.0 is required for the support of
     # `add_rows` method.
     'PrettyTable>=2.0.0',
+    'python-dotenv==1.0.0',
path: sky/setup_files/setup.py
text: @user1:
Can we relax this? This requires >= 3.8. Even though 3.7 is end of life now, it seems removing this will install 0.21.1 and it still works for the use case (tested).
```
» pip install python-dotenv==1.0.0      130 ↵
ERROR: Ignored the following versions that require a different python version: 1.0.0 Requires-Python >=3.8
ERROR: Could not find a version that satisfies the requirement python-dotenv==1.0.0 (from versions: 0.1.0, 0.1.2, 0.1.3, 0.1.5, 0.2.0, 0.3.0, 0.4.0, 0.5.0, 0.5.1, 0.6.0, 0.6.1, 0.6.2, 0.6.3, 0.6.4, 0.6.5, 0.7.0, 0.7.1, 0.8.0, 0.8.1, 0.8.2, 0.9.0, 0.9.1, 0.10.0, 0.10.1, 0.10.2, 0.10.3, 0.10.4, 0.10.5, 0.11.0, 0.12.0, 0.13.0, 0.14.0, 0.15.0, 0.16.0, 0.17.0, 0.17.1, 0.18.0, 0.19.0, 0.19.1, 0.19.2, 0.20.0, 0.21.0, 0.21.1)
ERROR: No matching distribution found for python-dotenv==1.0.0
```

@author:
done~

diff_hunk: @@ -25,6 +25,7 @@
 """
 import copy
 import datetime
+from dotenv import dotenv_values
path: sky/cli.py
text: @user1:
nit: we use google [styleguide](https://google.github.io/styleguide/pyguide.html#313-imports-formatting) where 3rd party libs are imported in 2nd block, L42-45. Also from styleguide: `import dotenv` and use `dotenv.dotenv_values` below?

@author:
done~

diff_hunk: @@ -382,6 +392,17 @@ def _parse_env_var(env_var: str) -> Tuple[str, str]:
                  default=None,
                  help=('Custom image id for launching the instances. '
                        'Passing "none" resets the config.')),
+    click.option(
+        '--env-file',
+        required=False,
+        type=dotenv_values,
path: sky/cli.py
text: @user1:
nit: `dotenv.dotenv_values`

@author:
done~