skypilot-nightly 1.0.0.dev20240923__tar.gz → 1.0.0.dev20240924__tar.gz

{skypilot_nightly-1.0.0.dev20240923/skypilot_nightly.egg-info → skypilot_nightly-1.0.0.dev20240924}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: skypilot-nightly
-Version: 1.0.0.dev20240923
+Version: 1.0.0.dev20240924
 Summary: SkyPilot: An intercloud broker for the clouds
 Author: SkyPilot Team
 License: Apache 2.0

{skypilot_nightly-1.0.0.dev20240923 → skypilot_nightly-1.0.0.dev20240924}/sky/__init__.py

@@ -5,7 +5,7 @@ from typing import Optional
 import urllib.request
 
 # Replaced with the current commit when building the wheels.
-_SKYPILOT_COMMIT_SHA = 'd602225c897f10ca67a2d5b5db21982c0dc8c1ec'
+_SKYPILOT_COMMIT_SHA = '800f7d6971bd604f266faebb33d044c7d5baca55'
 
 
 def _get_git_commit():
@@ -35,7 +35,7 @@ def _get_git_commit():
 
 
 __commit__ = _get_git_commit()
-__version__ = '1.0.0.dev20240923'
+__version__ = '1.0.0.dev20240924'
 __root_dir__ = os.path.dirname(os.path.abspath(__file__))
 
 
@@ -82,6 +82,9 @@ _set_http_proxy_env_vars()
 from sky import backends
 from sky import benchmark
 from sky import clouds
+from sky.admin_policy import AdminPolicy
+from sky.admin_policy import MutatedUserRequest
+from sky.admin_policy import UserRequest
 from sky.clouds.service_catalog import list_accelerators
 from sky.core import autostop
 from sky.core import cancel
@@ -112,6 +115,7 @@ from sky.optimizer import Optimizer
 from sky.optimizer import OptimizeTarget
 from sky.resources import Resources
 from sky.skylet.job_lib import JobStatus
+from sky.skypilot_config import Config
 from sky.status_lib import ClusterStatus
 from sky.task import Task
 
@@ -185,4 +189,9 @@ __all__ = [
     # core APIs Storage Management
     'storage_ls',
     'storage_delete',
+    # Admin Policy
+    'UserRequest',
+    'MutatedUserRequest',
+    'AdminPolicy',
+    'Config',
 ]

skypilot_nightly-1.0.0.dev20240924/sky/admin_policy.py

@@ -0,0 +1,101 @@
+"""Interface for admin-defined policy for user requests."""
+import abc
+import dataclasses
+import typing
+from typing import Optional
+
+if typing.TYPE_CHECKING:
+    import sky
+
+
+@dataclasses.dataclass
+class RequestOptions:
+    """Request options for admin policy.
+
+    Args:
+        cluster_name: Name of the cluster to create/reuse. It is None if not
+            specified by the user.
+        idle_minutes_to_autostop: Autostop setting requested by a user. The
+            cluster will be set to autostop after this many minutes of idleness.
+        down: If true, use autodown rather than autostop.
+        dryrun: Is the request a dryrun?
+    """
+    cluster_name: Optional[str]
+    idle_minutes_to_autostop: Optional[int]
+    down: bool
+    dryrun: bool
+
+
+@dataclasses.dataclass
+class UserRequest:
+    """A user request.
+
+    A "user request" is defined as a `sky launch / exec` command or its API
+    equivalent.
+
+    `sky jobs launch / serve up` involves multiple launch requests, including
+    the launch of controller and clusters for a job (which can have multiple
+    tasks if it is a pipeline) or service replicas. Each launch is a separate
+    request.
+
+    This class wraps the underlying task, the global skypilot config used to run
+    a task, and the request options.
+
+    Args:
+        task: User specified task.
+        skypilot_config: Global skypilot config to be used in this request.
+        request_options: Request options. It is None for jobs and services.
+    """
+    task: 'sky.Task'
+    skypilot_config: 'sky.Config'
+    request_options: Optional['RequestOptions'] = None
+
+
+@dataclasses.dataclass
+class MutatedUserRequest:
+    task: 'sky.Task'
+    skypilot_config: 'sky.Config'
+
+
+# pylint: disable=line-too-long
+class AdminPolicy:
+    """Abstract interface of an admin-defined policy for all user requests.
+
+    Admins can implement a subclass of AdminPolicy with the following signature:
+
+        import sky
+
+        class SkyPilotPolicyV1(sky.AdminPolicy):
+            def validate_and_mutate(user_request: UserRequest) -> MutatedUserRequest:
+                ...
+                return MutatedUserRequest(task=..., skypilot_config=...)
+
+    The policy can mutate both task and skypilot_config. Admins then distribute
+    a simple module that contains this implementation, installable in a way
+    that it can be imported by users from the same Python environment where
+    SkyPilot is running.
+
+    Users can register a subclass of AdminPolicy in the SkyPilot config file
+    under the key 'admin_policy', e.g.
+
+        admin_policy: my_package.SkyPilotPolicyV1
+    """
+
+    @classmethod
+    @abc.abstractmethod
+    def validate_and_mutate(cls,
+                            user_request: UserRequest) -> MutatedUserRequest:
+        """Validates and mutates the user request and returns mutated request.
+
+        Args:
+            user_request: The user request to validate and mutate.
+                UserRequest contains (sky.Task, sky.Config)
+
+        Returns:
+            MutatedUserRequest: The mutated user request.
+
+        Raises:
+            Exception to throw if the user request failed the validation.
+        """
+        raise NotImplementedError(
+            'Your policy must implement validate_and_mutate')

{skypilot_nightly-1.0.0.dev20240923 → skypilot_nightly-1.0.0.dev20240924}/sky/dag.py

@@ -1,8 +1,12 @@
 """DAGs: user applications to be run."""
 import pprint
 import threading
+import typing
 from typing import List, Optional
 
+if typing.TYPE_CHECKING:
+    from sky import task
+
 
 class Dag:
     """Dag: a user application, represented as a DAG of Tasks.
@@ -13,37 +17,37 @@ class Dag:
         >>>     task = sky.Task(...)
     """
 
-    def __init__(self):
-        self.tasks = []
+    def __init__(self) -> None:
+        self.tasks: List['task.Task'] = []
         import networkx as nx  # pylint: disable=import-outside-toplevel
 
         self.graph = nx.DiGraph()
-        self.name = None
+        self.name: Optional[str] = None
 
-    def add(self, task):
+    def add(self, task: 'task.Task') -> None:
         self.graph.add_node(task)
         self.tasks.append(task)
 
-    def remove(self, task):
+    def remove(self, task: 'task.Task') -> None:
         self.tasks.remove(task)
         self.graph.remove_node(task)
 
-    def add_edge(self, op1, op2):
+    def add_edge(self, op1: 'task.Task', op2: 'task.Task') -> None:
         assert op1 in self.graph.nodes
         assert op2 in self.graph.nodes
         self.graph.add_edge(op1, op2)
 
-    def __len__(self):
+    def __len__(self) -> int:
         return len(self.tasks)
 
-    def __enter__(self):
+    def __enter__(self) -> 'Dag':
         push_dag(self)
         return self
 
-    def __exit__(self, exc_type, exc_value, traceback):
+    def __exit__(self, exc_type, exc_value, traceback) -> None:
         pop_dag()
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         pformat = pprint.pformat(self.tasks)
         return f'DAG:\n{pformat}'
 
@@ -70,15 +74,15 @@ class Dag:
 
 class _DagContext(threading.local):
     """A thread-local stack of Dags."""
-    _current_dag = None
+    _current_dag: Optional[Dag] = None
     _previous_dags: List[Dag] = []
 
-    def push_dag(self, dag):
+    def push_dag(self, dag: Dag):
         if self._current_dag is not None:
             self._previous_dags.append(self._current_dag)
         self._current_dag = dag
 
-    def pop_dag(self):
+    def pop_dag(self) -> Optional[Dag]:
         old_dag = self._current_dag
         if self._previous_dags:
             self._current_dag = self._previous_dags.pop()

{skypilot_nightly-1.0.0.dev20240923 → skypilot_nightly-1.0.0.dev20240924}/sky/exceptions.py

@@ -286,3 +286,8 @@ class ServeUserTerminatedError(Exception):
 
 class PortDoesNotExistError(Exception):
     """Raised when the port does not exist."""
+
+
+class UserRequestRejectedByPolicy(Exception):
+    """Raised when a user request is rejected by an admin policy."""
+    pass

{skypilot_nightly-1.0.0.dev20240923 → skypilot_nightly-1.0.0.dev20240924}/sky/execution.py

@@ -9,6 +9,7 @@ from typing import List, Optional, Tuple, Union
 import colorama
 
 import sky
+from sky import admin_policy
 from sky import backends
 from sky import clouds
 from sky import global_user_state
@@ -16,6 +17,7 @@ from sky import optimizer
 from sky import sky_logging
 from sky.backends import backend_utils
 from sky.usage import usage_lib
+from sky.utils import admin_policy_utils
 from sky.utils import controller_utils
 from sky.utils import dag_utils
 from sky.utils import env_options
@@ -158,7 +160,16 @@ def _execute(
       handle: Optional[backends.ResourceHandle]; the handle to the cluster. None
         if dryrun.
     """
+
     dag = dag_utils.convert_entrypoint_to_dag(entrypoint)
+    dag, _ = admin_policy_utils.apply(
+        dag,
+        request_options=admin_policy.RequestOptions(
+            cluster_name=cluster_name,
+            idle_minutes_to_autostop=idle_minutes_to_autostop,
+            down=down,
+            dryrun=dryrun,
+        ))
     assert len(dag) == 1, f'We support 1 task for now. {dag}'
     task = dag.tasks[0]
 
@@ -170,9 +181,8 @@ def _execute(
 
     cluster_exists = False
     if cluster_name is not None:
-        existing_handle = global_user_state.get_handle_from_cluster_name(
-            cluster_name)
-        cluster_exists = existing_handle is not None
+        cluster_record = global_user_state.get_cluster_from_name(cluster_name)
+        cluster_exists = cluster_record is not None
         # TODO(woosuk): If the cluster exists, print a warning that
         # `cpus` and `memory` are not used as a job scheduling constraint,
         # unlike `gpus`.

{skypilot_nightly-1.0.0.dev20240923 → skypilot_nightly-1.0.0.dev20240924}/sky/jobs/controller.py

@@ -64,6 +64,7 @@ class JobsController:
             if len(self._dag.tasks) <= 1:
                 task_name = self._dag_name
             else:
+                assert task.name is not None, task
                 task_name = task.name
                 # This is guaranteed by the spot_launch API, where we fill in
                 # the task.name with
@@ -447,6 +448,7 @@ def _cleanup(job_id: int, dag_yaml: str):
     # controller, we should keep it in sync with JobsController.__init__()
     dag, _ = _get_dag_and_name(dag_yaml)
     for task in dag.tasks:
+        assert task.name is not None, task
         cluster_name = managed_job_utils.generate_managed_job_cluster_name(
             task.name, job_id)
         recovery_strategy.terminate_cluster(cluster_name)

{skypilot_nightly-1.0.0.dev20240923 → skypilot_nightly-1.0.0.dev20240924}/sky/jobs/core.py

@@ -18,6 +18,7 @@ from sky.jobs import constants as managed_job_constants
 from sky.jobs import utils as managed_job_utils
 from sky.skylet import constants as skylet_constants
 from sky.usage import usage_lib
+from sky.utils import admin_policy_utils
 from sky.utils import common_utils
 from sky.utils import controller_utils
 from sky.utils import dag_utils
@@ -54,6 +55,8 @@ def launch(
     dag_uuid = str(uuid.uuid4().hex[:4])
 
     dag = dag_utils.convert_entrypoint_to_dag(entrypoint)
+    dag, mutated_user_config = admin_policy_utils.apply(
+        dag, use_mutated_config_in_current_request=False)
     if not dag.is_chain():
         with ux_utils.print_exception_no_traceback():
             raise ValueError('Only single-task or chain DAG is '
@@ -103,6 +106,7 @@ def launch(
             **controller_utils.shared_controller_vars_to_fill(
                 controller_utils.Controllers.JOBS_CONTROLLER,
                 remote_user_config_path=remote_user_config_path,
+                local_user_config=mutated_user_config,
             ),
         }
 

{skypilot_nightly-1.0.0.dev20240923 → skypilot_nightly-1.0.0.dev20240924}/sky/serve/core.py

@@ -17,6 +17,7 @@ from sky.serve import serve_state
 from sky.serve import serve_utils
 from sky.skylet import constants
 from sky.usage import usage_lib
+from sky.utils import admin_policy_utils
 from sky.utils import common_utils
 from sky.utils import controller_utils
 from sky.utils import resources_utils
@@ -124,6 +125,10 @@ def up(
 
     _validate_service_task(task)
 
+    dag, mutated_user_config = admin_policy_utils.apply(
+        task, use_mutated_config_in_current_request=False)
+    task = dag.tasks[0]
+
     controller_utils.maybe_translate_local_file_mounts_and_sync_up(task,
                                                                    path='serve')
 
@@ -158,6 +163,7 @@ def up(
             **controller_utils.shared_controller_vars_to_fill(
                 controller=controller_utils.Controllers.SKY_SERVE_CONTROLLER,
                 remote_user_config_path=remote_config_yaml_path,
+                local_user_config=mutated_user_config,
             ),
         }
         common_utils.fill_template(serve_constants.CONTROLLER_TEMPLATE,

{skypilot_nightly-1.0.0.dev20240923 → skypilot_nightly-1.0.0.dev20240924}/sky/skypilot_config.py

@@ -61,6 +61,8 @@ from sky.utils import common_utils
 from sky.utils import schemas
 from sky.utils import ux_utils
 
+logger = sky_logging.init_logger(__name__)
+
 # The config path is discovered in this order:
 #
 # (1) (Used internally) If env var {ENV_VAR_SKYPILOT_CONFIG} exists, use its
@@ -78,11 +80,57 @@ ENV_VAR_SKYPILOT_CONFIG = 'SKYPILOT_CONFIG'
 # Path to the local config file.
 CONFIG_PATH = '~/.sky/config.yaml'
 
-logger = sky_logging.init_logger(__name__)
+
+class Config(Dict[str, Any]):
+    """SkyPilot config that supports setting/getting values with nested keys."""
+
+    def get_nested(self,
+                   keys: Tuple[str, ...],
+                   default_value: Any,
+                   override_configs: Optional[Dict[str, Any]] = None) -> Any:
+        """Gets a nested key.
+
+        If any key is not found, or any intermediate key does not point to a
+        dict value, returns 'default_value'.
+
+        Args:
+            keys: A tuple of strings representing the nested keys.
+            default_value: The default value to return if the key is not found.
+            override_configs: A dict of override configs with the same schema as
+                the config file, but only containing the keys to override.
+
+        Returns:
+            The value of the nested key, or 'default_value' if not found.
+        """
+        config = copy.deepcopy(self)
+        if override_configs is not None:
+            config = _recursive_update(config, override_configs)
+        return _get_nested(config, keys, default_value)
+
+    def set_nested(self, keys: Tuple[str, ...], value: Any) -> None:
+        """In-place sets a nested key to value.
+
+        Like get_nested(), if any key is not found, this will not raise an
+        error.
+        """
+        override = {}
+        for i, key in enumerate(reversed(keys)):
+            if i == 0:
+                override = {key: value}
+            else:
+                override = {key: override}
+        _recursive_update(self, override)
+
+    @classmethod
+    def from_dict(cls, config: Optional[Dict[str, Any]]) -> 'Config':
+        if config is None:
+            return cls()
+        return cls(**config)
+
 
 # The loaded config.
-_dict: Optional[Dict[str, Any]] = None
-_loaded_config_path = None
+_dict = Config()
+_loaded_config_path: Optional[str] = None
 
 
 def _get_nested(configs: Optional[Dict[str, Any]], keys: Iterable[str],
@@ -131,17 +179,11 @@ def get_nested(keys: Tuple[str, ...],
     ), (f'Override configs must not be provided when keys {keys} is not within '
         'constants.OVERRIDEABLE_CONFIG_KEYS: '
         f'{constants.OVERRIDEABLE_CONFIG_KEYS}')
-    config: Dict[str, Any] = {}
-    if _dict is not None:
-        config = copy.deepcopy(_dict)
-    if override_configs is None:
-        override_configs = {}
-    config = _recursive_update(config, override_configs)
-    return _get_nested(config, keys, default_value)
+    return _dict.get_nested(keys, default_value, override_configs)
 
 
-def _recursive_update(base_config: Dict[str, Any],
-                      override_config: Dict[str, Any]) -> Dict[str, Any]:
+def _recursive_update(base_config: Config,
+                      override_config: Dict[str, Any]) -> Config:
     """Recursively updates base configuration with override configuration"""
     for key, value in override_config.items():
         if (isinstance(value, dict) and key in base_config and
@@ -157,22 +199,14 @@ def set_nested(keys: Tuple[str, ...], value: Any) -> Dict[str, Any]:
 
     Like get_nested(), if any key is not found, this will not raise an error.
     """
-    _check_loaded_or_die()
-    assert _dict is not None
-    override = {}
-    for i, key in enumerate(reversed(keys)):
-        if i == 0:
-            override = {key: value}
-        else:
-            override = {key: override}
-    return _recursive_update(copy.deepcopy(_dict), override)
+    copied_dict = copy.deepcopy(_dict)
+    copied_dict.set_nested(keys, value)
+    return dict(**copied_dict)
 
 
-def to_dict() -> Dict[str, Any]:
+def to_dict() -> Config:
     """Returns a deep-copied version of the current config."""
-    if _dict is not None:
-        return copy.deepcopy(_dict)
-    return {}
+    return copy.deepcopy(_dict)
 
 
 def _try_load_config() -> None:
@@ -192,13 +226,14 @@ def _try_load_config() -> None:
     config_path = os.path.expanduser(config_path)
     if os.path.exists(config_path):
         logger.debug(f'Using config path: {config_path}')
-        _loaded_config_path = config_path
         try:
-            _dict = common_utils.read_yaml(config_path)
+            config = common_utils.read_yaml(config_path)
+            _dict = Config.from_dict(config)
+            _loaded_config_path = config_path
             logger.debug(f'Config loaded:\n{pprint.pformat(_dict)}')
         except yaml.YAMLError as e:
             logger.error(f'Error in loading config file ({config_path}):', e)
-        if _dict is not None:
+        if _dict:
             common_utils.validate_schema(
                 _dict,
                 schemas.get_config_schema(),
@@ -219,14 +254,6 @@ def loaded_config_path() -> Optional[str]:
 _try_load_config()
 
 
-def _check_loaded_or_die():
-    """Checks loaded() is true; otherwise raises RuntimeError."""
-    if _dict is None:
-        raise RuntimeError(
-            f'No user configs loaded. Check {CONFIG_PATH} exists and '
-            'can be loaded.')
-
-
 def loaded() -> bool:
     """Returns if the user configurations are loaded."""
-    return _dict is not None
+    return bool(_dict)

{skypilot_nightly-1.0.0.dev20240923 → skypilot_nightly-1.0.0.dev20240924}/sky/templates/jobs-controller.yaml.j2

@@ -4,7 +4,9 @@ name: {{dag_name}}
 
 file_mounts:
   {{remote_user_yaml_path}}: {{user_yaml_path}}
-  {{remote_user_config_path}}: skypilot:local_skypilot_config_path
+  {%- if local_user_config_path is not none %}
+  {{remote_user_config_path}}: {{local_user_config_path}}
+  {%- endif %}
   {%- for remote_catalog_path, local_catalog_path in modified_catalogs.items() %}
   {{remote_catalog_path}}: {{local_catalog_path}}
   {%- endfor %}

{skypilot_nightly-1.0.0.dev20240923 → skypilot_nightly-1.0.0.dev20240924}/sky/templates/sky-serve-controller.yaml.j2

@@ -23,7 +23,9 @@ setup: |
 
 file_mounts:
   {{remote_task_yaml_path}}: {{local_task_yaml_path}}
-  {{remote_user_config_path}}: skypilot:local_skypilot_config_path
+  {%- if local_user_config_path is not none %}
+  {{remote_user_config_path}}: {{local_user_config_path}}
+  {%- endif %}
   {%- for remote_catalog_path, local_catalog_path in modified_catalogs.items() %}
   {{remote_catalog_path}}: {{local_catalog_path}}
   {%- endfor %}

skypilot_nightly-1.0.0.dev20240924/sky/utils/admin_policy_utils.py

@@ -0,0 +1,145 @@
+"""Admin policy utils."""
+import copy
+import importlib
+import os
+import tempfile
+from typing import Optional, Tuple, Union
+
+import colorama
+
+from sky import admin_policy
+from sky import dag as dag_lib
+from sky import exceptions
+from sky import sky_logging
+from sky import skypilot_config
+from sky import task as task_lib
+from sky.utils import common_utils
+from sky.utils import ux_utils
+
+logger = sky_logging.init_logger(__name__)
+
+
+def _get_policy_cls(
+        policy: Optional[str]) -> Optional[admin_policy.AdminPolicy]:
+    """Gets admin-defined policy."""
+    if policy is None:
+        return None
+    try:
+        module_path, class_name = policy.rsplit('.', 1)
+        module = importlib.import_module(module_path)
+    except ImportError as e:
+        with ux_utils.print_exception_no_traceback():
+            raise ImportError(
+                f'Failed to import policy module: {policy}. '
+                'Please check if the module is installed in your Python '
+                'environment.') from e
+
+    try:
+        policy_cls = getattr(module, class_name)
+    except AttributeError as e:
+        with ux_utils.print_exception_no_traceback():
+            raise AttributeError(
+                f'Could not find {class_name} class in module {module_path}. '
+                'Please check with your policy admin for details.') from e
+
+    # Check if the module implements the AdminPolicy interface.
+    if not issubclass(policy_cls, admin_policy.AdminPolicy):
+        with ux_utils.print_exception_no_traceback():
+            raise ValueError(
+                f'Policy class {policy!r} does not implement the AdminPolicy '
+                'interface. Please check with your policy admin for details.')
+    return policy_cls
+
+
+def apply(
+    entrypoint: Union['dag_lib.Dag', 'task_lib.Task'],
+    use_mutated_config_in_current_request: bool = True,
+    request_options: Optional[admin_policy.RequestOptions] = None,
+) -> Tuple['dag_lib.Dag', skypilot_config.Config]:
+    """Applies an admin policy (if registered) to a DAG or a task.
+
+    It mutates a Dag by applying any registered admin policy and also
+    potentially updates (controlled by `use_mutated_config_in_current_request`)
+    the global SkyPilot config if there is any changes made by the policy.
+
+    Args:
+        dag: The dag to be mutated by the policy.
+        use_mutated_config_in_current_request: Whether to use the mutated
+            config in the current request.
+        request_options: Additional options user passed for the current request.
+
+    Returns:
+        - The new copy of dag after applying the policy
+        - The new copy of skypilot config after applying the policy.
+    """
+    if isinstance(entrypoint, task_lib.Task):
+        dag = dag_lib.Dag()
+        dag.add(entrypoint)
+    else:
+        dag = entrypoint
+
+    policy = skypilot_config.get_nested(('admin_policy',), None)
+    policy_cls = _get_policy_cls(policy)
+    if policy_cls is None:
+        return dag, skypilot_config.to_dict()
+
+    logger.info(f'Applying policy: {policy}')
+    original_config = skypilot_config.to_dict()
+    config = copy.deepcopy(original_config)
+    mutated_dag = dag_lib.Dag()
+    mutated_dag.name = dag.name
+
+    mutated_config = None
+    for task in dag.tasks:
+        user_request = admin_policy.UserRequest(task, config, request_options)
+        try:
+            mutated_user_request = policy_cls.validate_and_mutate(user_request)
+        except Exception as e:  # pylint: disable=broad-except
+            with ux_utils.print_exception_no_traceback():
+                raise exceptions.UserRequestRejectedByPolicy(
+                    f'{colorama.Fore.RED}User request rejected by policy '
+                    f'{policy!r}{colorama.Fore.RESET}: '
+                    f'{common_utils.format_exception(e, use_bracket=True)}'
+                ) from e
+        if mutated_config is None:
+            mutated_config = mutated_user_request.skypilot_config
+        else:
+            if mutated_config != mutated_user_request.skypilot_config:
+                # In the case of a pipeline of tasks, the mutated config
+                # generated should remain the same for all tasks for now for
+                # simplicity.
+                # TODO(zhwu): We should support per-task mutated config or
+                # allowing overriding required global config in task YAML.
+                with ux_utils.print_exception_no_traceback():
+                    raise exceptions.UserRequestRejectedByPolicy(
+                        'All tasks must have the same SkyPilot config after '
+                        'applying the policy. Please check with your policy '
+                        'admin for details.')
+        mutated_dag.add(mutated_user_request.task)
+    assert mutated_config is not None, dag
+
+    # Update the new_dag's graph with the old dag's graph
+    for u, v in dag.graph.edges:
+        u_idx = dag.tasks.index(u)
+        v_idx = dag.tasks.index(v)
+        mutated_dag.graph.add_edge(mutated_dag.tasks[u_idx],
+                                   mutated_dag.tasks[v_idx])
+
+    if (use_mutated_config_in_current_request and
+            original_config != mutated_config):
+        with tempfile.NamedTemporaryFile(
+                delete=False,
+                mode='w',
+                prefix='policy-mutated-skypilot-config-',
+                suffix='.yaml') as temp_file:
+
+            common_utils.dump_yaml(temp_file.name, dict(**mutated_config))
+            os.environ[skypilot_config.ENV_VAR_SKYPILOT_CONFIG] = temp_file.name
+            logger.debug(f'Updated SkyPilot config: {temp_file.name}')
+            # TODO(zhwu): This is not a clean way to update the SkyPilot config,
+            # because we are resetting the global context for a single DAG,
+            # which is conceptually weird.
+            importlib.reload(skypilot_config)
+
+    logger.debug(f'Mutated user request: {mutated_user_request}')
+    return mutated_dag, mutated_config