Add support for env files.

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

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',