[Core] Allow override env for all task in dag

sky/cli.py

@@ -765,6 +765,8 @@ def _make_task_or_dag_from_entrypoint_with_overrides(
                     f'WARNING: override params {override_params} are ignored, '
                     'since the yaml file contains multiple tasks.',
                     fg='yellow')
+            for task in dag.tasks:
+                task.update_envs(env)
             return dag
         assert len(dag.tasks) == 1, (
             f'If you see this, please file an issue; tasks: {dag.tasks}')

sky/utils/dag_utils.py

@@ -90,12 +90,6 @@ def load_chain_dag_from_yaml(
         # YAML has only `name: xxx`. Still instantiate a task.
         configs = [{'name': dag_name}]
 
-    if len(configs) > 1:
-        # TODO(zongheng): in a chain DAG of N tasks, cli.py currently makes the
-        # decision to not apply overrides. Here we maintain this behavior. We
-        # can listen to user feedback to change this.
-        env_overrides = None
-
     current_task = None
     with dag_lib.Dag() as dag:
         for task_config in configs: