add request name to admin policy (#7840)

* add request name to admin policy

* format

* fix ut

* plumb to userrequest

* add ut

* fix restful test

* add example policy using new field

* make request names available on sky

* better example admin policies

* controller launch

* add reprs

* format

---------

Co-authored-by: tv <[email protected]>
Co-authored-by: tv <[email protected]>

Author: 3 people

examples/admin_policy/add_labels_conditional.yaml

@@ -0,0 +1 @@
+admin_policy: example_policy.AddLabelsConditionalPolicy

examples/admin_policy/example_policy/example_policy/__init__.py

@@ -1,5 +1,6 @@
 """Example admin policy module and prebuilt policies."""
 from example_policy.client_policy import UseLocalGcpCredentialsPolicy
+from example_policy.skypilot_policy import AddLabelsConditionalPolicy
 from example_policy.skypilot_policy import AddLabelsPolicy
 from example_policy.skypilot_policy import AddVolumesPolicy
 from example_policy.skypilot_policy import DisablePublicIpPolicy

examples/admin_policy/example_policy/example_policy/skypilot_policy.py

@@ -42,6 +42,27 @@ def validate_and_mutate(
         return sky.MutatedUserRequest(user_request.task, config)
 
 
+class AddLabelsConditionalPolicy(sky.AdminPolicy):
+    """Example policy: adds a kubernetes label for skypilot_config
+    if the request is a cluster launch request."""
+
+    @classmethod
+    def validate_and_mutate(
+            cls, user_request: sky.UserRequest) -> sky.MutatedUserRequest:
+        if user_request.request_name in [
+                sky.AdminPolicyRequestName.VALIDATE,
+                sky.AdminPolicyRequestName.OPTIMIZE
+        ]:
+            return sky.MutatedUserRequest(user_request.task,
+                                          user_request.skypilot_config)
+        config = user_request.skypilot_config
+        labels = config.get_nested(('kubernetes', 'custom_metadata', 'labels'),
+                                   {})
+        labels['app'] = 'skypilot'
+        config.set_nested(('kubernetes', 'custom_metadata', 'labels'), labels)
+        return sky.MutatedUserRequest(user_request.task, config)
+
+
 class DisablePublicIpPolicy(sky.AdminPolicy):
     """Example policy: disables public IP for all AWS tasks."""
 
@@ -94,15 +115,17 @@ def validate_and_mutate(
         policy is applied, we should expect a few seconds latency when a user
         run a request.
         """
-        request_options = user_request.request_options
-
-        # Request options is None when a task is executed with `jobs launch` or
-        # `sky serve up`.
-        if request_options is None:
+        if user_request.request_name not in [
+                sky.AdminPolicyRequestName.CLUSTER_LAUNCH,
+                sky.AdminPolicyRequestName.CLUSTER_EXEC,
+        ]:
             return sky.MutatedUserRequest(
                 task=user_request.task,
                 skypilot_config=user_request.skypilot_config)
 
+        request_options = user_request.request_options
+        # Request options is not None when a task is executed with `sky launch`.
+        assert request_options is not None
         # Get the cluster record to operate on.
         cluster_name = request_options.cluster_name
         cluster_records: List[responses.StatusResponse] = []

sky/__init__.py

@@ -122,6 +122,7 @@ def set_proxy_env_var(proxy_var: str, urllib_var: Optional[str]):
 from sky.jobs import ManagedJobStatus
 from sky.optimizer import Optimizer
 from sky.resources import Resources
+from sky.server.requests.request_names import AdminPolicyRequestName
 from sky.skylet.job_lib import JobStatus
 from sky.task import Task
 from sky.utils.common import OptimizeTarget
@@ -228,6 +229,7 @@ def set_proxy_env_var(proxy_var: str, urllib_var: Optional[str]):
     'MutatedUserRequest',
     'AdminPolicy',
     'Config',
+    'AdminPolicyRequestName',
     # Registry
     'CLOUD_REGISTRY',
     'JOBS_RECOVERY_STRATEGY_REGISTRY',

sky/admin_policy.py

@@ -11,6 +11,7 @@
 from sky import exceptions
 from sky import models
 from sky.adaptors import common as adaptors_common
+from sky.server.requests import request_names
 from sky.utils import common_utils
 from sky.utils import config_utils
 from sky.utils import ux_utils
@@ -50,6 +51,7 @@ class _UserRequestBody(pydantic.BaseModel):
     # will be converted to JSON string, which will lose the None key.
     task: str
     skypilot_config: str
+    request_name: str
     request_options: Optional[RequestOptions] = None
     at_client_side: bool = False
     user: str
@@ -81,6 +83,7 @@ class UserRequest:
     """
     task: 'sky.Task'
     skypilot_config: 'sky.Config'
+    request_name: request_names.AdminPolicyRequestName
     request_options: Optional['RequestOptions'] = None
     at_client_side: bool = False
     user: Optional['models.User'] = None
@@ -90,6 +93,7 @@ def encode(self) -> str:
             task=yaml_utils.dump_yaml_str(self.task.to_yaml_config()),
             skypilot_config=yaml_utils.dump_yaml_str(dict(
                 self.skypilot_config)),
+            request_name=self.request_name.value,
             request_options=self.request_options,
             at_client_side=self.at_client_side,
             user=(yaml_utils.dump_yaml_str(self.user.to_dict())
@@ -110,6 +114,8 @@ def decode(cls, body: str) -> 'UserRequest':
             skypilot_config=config_utils.Config.from_dict(
                 yaml_utils.read_yaml_all_str(
                     user_request_body.skypilot_config)[0]),
+            request_name=request_names.AdminPolicyRequestName(
+                user_request_body.request_name),
             request_options=user_request_body.request_options,
             at_client_side=user_request_body.at_client_side,
             user=user,

sky/client/sdk.py

@@ -37,6 +37,7 @@
 from sky.server import rest
 from sky.server import versions
 from sky.server.requests import payloads
+from sky.server.requests import request_names
 from sky.server.requests import requests as requests_lib
 from sky.skylet import autostop_lib
 from sky.skylet import constants
@@ -603,7 +604,10 @@ def launch(
         down=down,
         dryrun=dryrun)
     with admin_policy_utils.apply_and_use_config_in_current_request(
-            dag, request_options=request_options, at_client_side=True) as dag:
+            dag,
+            request_name=request_names.AdminPolicyRequestName.CLUSTER_LAUNCH,
+            request_options=request_options,
+            at_client_side=True) as dag:
         return _launch(
             dag,
             cluster_name,

sky/core.py

@@ -25,6 +25,7 @@
 from sky.provision.kubernetes import constants as kubernetes_constants
 from sky.provision.kubernetes import utils as kubernetes_utils
 from sky.schemas.api import responses
+from sky.server.requests import request_names
 from sky.skylet import autostop_lib
 from sky.skylet import constants
 from sky.skylet import job_lib
@@ -84,7 +85,9 @@ def optimize(
     # but we do not apply the admin policy there. We should apply the admin
     # policy in the optimizer, but that will require some refactoring.
     with admin_policy_utils.apply_and_use_config_in_current_request(
-            dag, request_options=request_options) as dag:
+            dag,
+            request_name=request_names.AdminPolicyRequestName.OPTIMIZE,
+            request_options=request_options) as dag:
         dag.resolve_and_validate_volumes()
         return optimizer.Optimizer.optimize(dag=dag,
                                             minimize=minimize,

sky/execution.py

@@ -16,6 +16,7 @@
 from sky import optimizer
 from sky import sky_logging
 from sky.backends import backend_utils
+from sky.server.requests import request_names
 from sky.skylet import autostop_lib
 from sky.usage import usage_lib
 from sky.utils import admin_policy_utils
@@ -116,8 +117,10 @@ def _execute(
     no_setup: bool = False,
     clone_disk_from: Optional[str] = None,
     skip_unnecessary_provisioning: bool = False,
+    *,  #keyword only separator
     # Internal only:
     # pylint: disable=invalid-name
+    _request_name: request_names.AdminPolicyRequestName,
     _quiet_optimizer: bool = False,
     _is_launched_by_jobs_controller: bool = False,
     _is_launched_by_sky_serve_controller: bool = False,
@@ -187,6 +190,7 @@ def _execute(
                 idle_minutes_to_autostop = resource.autostop_config.idle_minutes
     with admin_policy_utils.apply_and_use_config_in_current_request(
             dag,
+            request_name=_request_name,
             request_options=admin_policy.RequestOptions(
                 cluster_name=cluster_name,
                 idle_minutes_to_autostop=idle_minutes_to_autostop,
@@ -519,6 +523,51 @@ def _planner(_t: 'sky.Task'):
 
 @timeline.event
 @usage_lib.entrypoint
+def cluster_launch(
+    task: Union['sky.Task', 'sky.Dag'],
+    cluster_name: Optional[str] = None,
+    retry_until_up: bool = False,
+    idle_minutes_to_autostop: Optional[int] = None,
+    dryrun: bool = False,
+    down: bool = False,
+    stream_logs: bool = True,
+    backend: Optional[backends.Backend] = None,
+    optimize_target: common.OptimizeTarget = common.OptimizeTarget.COST,
+    no_setup: bool = False,
+    clone_disk_from: Optional[str] = None,
+    fast: bool = False,
+    *,  #keyword only separator
+    # Internal only:
+    # pylint: disable=invalid-name
+    _quiet_optimizer: bool = False,
+    _is_launched_by_jobs_controller: bool = False,
+    _is_launched_by_sky_serve_controller: bool = False,
+    _disable_controller_check: bool = False,
+    job_logger: logging.Logger = logger,
+) -> Tuple[Optional[int], Optional[backends.ResourceHandle]]:
+    return launch(
+        task=task,
+        cluster_name=cluster_name,
+        retry_until_up=retry_until_up,
+        idle_minutes_to_autostop=idle_minutes_to_autostop,
+        dryrun=dryrun,
+        down=down,
+        stream_logs=stream_logs,
+        backend=backend,
+        optimize_target=optimize_target,
+        no_setup=no_setup,
+        clone_disk_from=clone_disk_from,
+        fast=fast,
+        _quiet_optimizer=_quiet_optimizer,
+        _is_launched_by_jobs_controller=_is_launched_by_jobs_controller,
+        _is_launched_by_sky_serve_controller=
+        _is_launched_by_sky_serve_controller,
+        _disable_controller_check=_disable_controller_check,
+        _request_name=request_names.AdminPolicyRequestName.CLUSTER_LAUNCH,
+        job_logger=job_logger,
+    )
+
+
 # A launch routine will share tempfiles between steps, so we init a tempdir
 # for the launch routine and gc the entire dir after launch.
 @tempstore.with_tempdir
@@ -535,12 +584,14 @@ def launch(
     no_setup: bool = False,
     clone_disk_from: Optional[str] = None,
     fast: bool = False,
+    *,  #keyword only separator
     # Internal only:
     # pylint: disable=invalid-name
     _quiet_optimizer: bool = False,
     _is_launched_by_jobs_controller: bool = False,
     _is_launched_by_sky_serve_controller: bool = False,
     _disable_controller_check: bool = False,
+    _request_name: request_names.AdminPolicyRequestName,
     job_logger: logging.Logger = logger,
 ) -> Tuple[Optional[int], Optional[backends.ResourceHandle]]:
     # NOTE(dev): Keep the docstring consistent between the Python API and CLI.
@@ -707,6 +758,7 @@ def launch(
         _is_launched_by_jobs_controller=_is_launched_by_jobs_controller,
         _is_launched_by_sky_serve_controller=
         _is_launched_by_sky_serve_controller,
+        _request_name=_request_name,
         job_logger=job_logger)
 
 
@@ -794,4 +846,5 @@ def exec(  # pylint: disable=redefined-builtin
         ],
         cluster_name=cluster_name,
         job_logger=job_logger,
+        _request_name=request_names.AdminPolicyRequestName.CLUSTER_EXEC,
     )

sky/jobs/client/sdk.py

@@ -15,6 +15,7 @@
 from sky.server import rest
 from sky.server import versions
 from sky.server.requests import payloads
+from sky.server.requests import request_names
 from sky.skylet import constants
 from sky.usage import usage_lib
 from sky.utils import admin_policy_utils
@@ -84,7 +85,9 @@ def launch(
 
     dag = dag_utils.convert_entrypoint_to_dag(task)
     with admin_policy_utils.apply_and_use_config_in_current_request(
-            dag, at_client_side=True) as dag:
+            dag,
+            request_name=request_names.AdminPolicyRequestName.JOBS_LAUNCH,
+            at_client_side=True) as dag:
         sdk.validate(dag)
         if _need_confirmation:
             job_identity = 'a managed job'

sky/jobs/server/core.py

@@ -35,6 +35,7 @@
 from sky.serve import serve_state
 from sky.serve import serve_utils
 from sky.serve.server import impl
+from sky.server.requests import request_names
 from sky.skylet import constants as skylet_constants
 from sky.usage import usage_lib
 from sky.utils import admin_policy_utils
@@ -237,7 +238,8 @@ def launch(
     # Always apply the policy again here, even though it might have been applied
     # in the CLI. This is to ensure that we apply the policy to the final DAG
     # and get the mutated config.
-    dag, mutated_user_config = admin_policy_utils.apply(dag)
+    dag, mutated_user_config = admin_policy_utils.apply(
+        dag, request_name=request_names.AdminPolicyRequestName.JOBS_LAUNCH)
     dag.resolve_and_validate_volumes()
     if not dag.is_chain():
         with ux_utils.print_exception_no_traceback():
@@ -465,12 +467,15 @@ def _submit_one(
                     # intermediate bucket and newly created bucket should be in
                     # workspace A.
                     if consolidation_mode_job_id is None:
-                        return execution.launch(task=controller_task,
-                                                cluster_name=controller_name,
-                                                stream_logs=stream_logs,
-                                                retry_until_up=True,
-                                                fast=True,
-                                                _disable_controller_check=True)
+                        return execution.launch(
+                            task=controller_task,
+                            cluster_name=controller_name,
+                            stream_logs=stream_logs,
+                            retry_until_up=True,
+                            fast=True,
+                            _request_name=request_names.AdminPolicyRequestName.
+                            JOBS_LAUNCH_CONTROLLER,
+                            _disable_controller_check=True)
                     # Manually launch the scheduler in consolidation mode.
                     local_handle = backend_utils.is_controller_accessible(
                         controller=controller, stopped_message='')