[Core] Admin policy enforcement plugin

docs/source/cloud-setup/policy.rst

@@ -0,0 +1,169 @@
+.. _advanced-policy-config:
+
+Admin Policy Enforcement
+========================
+
+
+SkyPilot allows admins to enforce policies on users' SkyPilot usage by applying
+custom validation and mutation logic on user's task and SkyPilot config.
+
+In short, admins offers a Python package with a customized inheritance of SkyPilot's
+``AdminPolicy`` interface, and a user just needs to set the ``admin_policy`` field in
+the SkyPilot config ``~/.sky/config.yaml`` to enforce the policy to all their
+tasks.
+
+Overview
+--------
+
+
+
+User-Side
+~~~~~~~~~~
+
+To apply the policy, a user needs to set the ``admin_policy`` field in the SkyPilot config
+``~/.sky/config.yaml`` to the path of the Python package that implements the policy.
+For example:
+
+.. code-block:: yaml
+
+ admin_policy: mypackage.subpackage.MyPolicy
+
+
+.. hint::
+
+ SkyPilot loads the policy from the given package in the same Python environment.
+ You can test the existance of the policy by running:
+
+ .. code-block:: bash
+
+ python -c "from mypackage.subpackage import MyPolicy"
+
+
+Admin-Side
+~~~~~~~~~~
+
+An admin can distribute the Python package to users with pre-defined policy. The
+policy should follow the following interface:
+
+.. code-block:: python
+
+ import sky
+
+ class MyPolicy(sky.AdminPolicy):
+ @classmethod
+ def validate_and_mutate(cls, user_request: sky.UserRequest) -> sky.MutatedUserRequest:
+ # Logics for validate and modify user requests.
+ ...
+ return sky.MutatedUserRequest(user_request.task,
+ user_request.skypilot_config)
+
+
+``UserRequest`` and ``MutatedUserRequest`` are defined as follows:
+
+.. code-block:: python
+
+ class UserRequest:
+ task: sky.Task
+ skypilot_config: sky.NestedConfig
+ operation_args: sky.OperationArgs
+
+ class MutatedUserRequest:
+ task: sky.Task
+ skypilot_config: sky.NestedConfig
+
+That said, an ``AdminPolicy`` can mutate any fields of a user request, including
+the :ref:`task <yaml-spec>` and the :ref:`global skypilot config <config-yaml>`,
+giving admins a lot of flexibility to control user's SkyPilot usage.
+
+An ``AdminPolicy`` is responsible to both validate and mutate user requests. If
+a request should be rejected, the policy should raise an exception.
+
+
+Example Policies 
+----------------
+
+Reject All
+~~~~~~~~~~
+
+.. code-block:: python
+
+ class RejectAllPolicy(sky.AdminPolicy):
+ @classmethod
+ def validate_and_mutate(cls, user_request: sky.UserRequest) -> sky.MutatedUserRequest:
+ raise RuntimeError("This policy rejects all user requests.")
+
+.. code-block:: yaml
+
+ admin_policy: examples.admin_policy.reject_all.RejectAllPolicy
+
+
+Add Kubernetes Labels for all Tasks
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. code-block:: python
+
+ class AddLabelsPolicy(sky.AdminPolicy):
+ @classmethod
+ def validate_and_mutate(cls, user_request: sky.UserRequest) -> sky.MutatedUserRequest: 
+ config = user_request.skypilot_config
+ labels = config.get_nested(('kubernetes', 'labels'), {})
+ labels['app'] = 'skypilot'
+ config.set_nested(('kubernetes', 'labels'), labels)
+ return sky.MutatedUserRequest(user_request.task, config)
+
+.. code-block:: yaml
+
+ admin_policy: examples.admin_policy.add_labels.AddLabelsPolicy
+
+
+Always Disable Public IP for AWS Tasks
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. code-block:: python
+
+ class DisablePublicIPPolicy(sky.AdminPolicy):
+ @classmethod
+ def validate_and_mutate(cls, user_request: sky.UserRequest) -> sky.MutatedUserRequest:
+ config = user_request.skypilot_config
+ config.set_nested(('aws', 'use_internal_ip'), True)
+ if config.get_nested(('aws', 'vpc_name'), None) is None:
+ # If no VPC name is specified, it is likely a mistake. We should
+ # reject the request
+ raise RuntimeError('VPC name should be set. Check organization '
+ 'wiki for more information.')
+ return sky.MutatedUserRequest(user_request.task, config)
+
+.. code-block:: yaml
+
+ admin_policy: examples.admin_policy.disable_public_ip.DisablePublicIPPolicy
+
+
+Enforce Autostop for all Tasks
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. code-block:: python
+
+ class EnforceAutostopPolicy(sky.AdminPolicy):
+ @classmethod
+ def validate_and_mutate(cls, user_request: sky.UserRequest) -> sky.MutatedUserRequest:
+ operation_args = user_request.operation_args
+ # Operation args can be None for jobs and services, for which we
+ # don't need to enforce autostop, as they are already managed.
+ if operation_args is None:
+ return sky.MutatedUserRequest(
+ task=user_request.task,
+ skypilot_config=user_request.skypilot_config)
+ idle_minutes_to_autostop = operation_args.idle_minutes_to_autostop
+ # Enforce autostop/down to be set for all tasks for new clusters.
+ if not operation_args.cluster_exists and (
+ idle_minutes_to_autostop is None or
+ idle_minutes_to_autostop < 0):
+ raise RuntimeError('Autostop/down must be set for all newly '
+ 'launched clusters.')
+ return sky.MutatedUserRequest(
+ task=user_request.task,
+ skypilot_config=user_request.skypilot_config)
+
+.. code-block:: yaml
+
+ admin_policy: examples.admin_policy.enforce_autostop.EnforceAutostopPolicy

docs/source/docs/index.rst

@@ -201,7 +201,8 @@ Read the research:
  ../cloud-setup/cloud-permissions/index
  ../cloud-setup/cloud-auth
  ../cloud-setup/quota
-
+ ../cloud-setup/policy
+
 .. toctree::
  :hidden:
  :maxdepth: 1

docs/source/reference/config.rst

@@ -89,11 +89,11 @@ Available fields and semantics:
 
  # Custom policy to be applied to all tasks.
  #
- # The policy function to be applied and mutate all tasks, which can be used to
+ # The policy class to be applied and mutate all tasks, which can be used to
  # enforce certain policies on all tasks.
  #
- # See details in: <TODO: add link to policy docs>
- policy: my_package.skypilot_policy_fn_v1
+ # The policy class should implement the sky.AdminPolicy interface.
+ admin_policy: my_package.SkyPilotPolicyV1
 
  # Advanced AWS configurations (optional).
  # Apply to all new instances but not existing ones.

examples/admin_policy/config_label_config.yaml

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

examples/admin_policy/enforce_autostop.yaml

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

examples/admin_policy/example_policy/example_policy/__init__.py

@@ -0,0 +1,6 @@
+"""Example module for SkyPilot admin policies."""
+
+from example_policy.skypilot_policy import ConfigLabelPolicy
+from example_policy.skypilot_policy import EnforceAutostopPolicy
+from example_policy.skypilot_policy import RejectAllPolicy
+from example_policy.skypilot_policy import TaskLabelPolicy

examples/admin_policy/example_policy/example_policy/skypilot_policy.py

@@ -0,0 +1,74 @@
+import copy
+import getpass
+
+import sky
+
+
+class TaskLabelPolicy(sky.AdminPolicy):
+ """Example policy: add label for task with the local user name."""
+
+ @classmethod
+ def validate_and_mutate(
+ cls, user_request: sky.UserRequest) -> sky.MutatedUserRequest:
+ """Add label for task with the local user name."""
+ local_user_name = getpass.getuser()
+
+ # Add label for task with the local user name
+ task = user_request.task
+ for r in task.resources:
+ r.labels['local_user'] = local_user_name
+
+ return sky.MutatedUserRequest(
+ task=task, skypilot_config=user_request.skypilot_config)
+
+
+class ConfigLabelPolicy(sky.AdminPolicy):
+ """Example policy: add label for skypilot_config with the local user name."""
+
+ @classmethod
+ def validate_and_mutate(
+ cls, user_request: sky.UserRequest) -> sky.MutatedUserRequest:
+ """Add label for skypilot_config with the local user name."""
+ local_user_name = getpass.getuser()
+
+ # Add label for skypilot_config with the local user name
+ skypilot_config = copy.deepcopy(user_request.skypilot_config)
+ skypilot_config.set_nested(('gcp', 'labels', 'local_user'),
+ local_user_name)
+ return sky.MutatedUserRequest(task=user_request.task,
+ skypilot_config=skypilot_config)
+
+
+class RejectAllPolicy(sky.AdminPolicy):
+ """Example policy: reject all user requests."""
+
+ @classmethod
+ def validate_and_mutate(
+ cls, user_request: sky.UserRequest) -> sky.MutatedUserRequest:
+ """Reject all user requests."""
+ del user_request
+ raise RuntimeError('Reject all policy')
+
+
+class EnforceAutostopPolicy(sky.AdminPolicy):
+ """Example policy: enforce autostop for all tasks."""
+
+ @classmethod
+ def validate_and_mutate(
+ cls, user_request: sky.UserRequest) -> sky.MutatedUserRequest:
+ """Enforce autostop for all tasks."""
+ operation_args = user_request.operation_args
+ if operation_args is None:
+ return sky.MutatedUserRequest(
+ task=user_request.task,
+ skypilot_config=user_request.skypilot_config)
+ idle_minutes_to_autostop = operation_args.idle_minutes_to_autostop
+ # Enforce autostop/down to be set for all tasks for new clusters.
+ if not operation_args.cluster_exists and (
+ idle_minutes_to_autostop is None or
+ idle_minutes_to_autostop < 0):
+ raise RuntimeError('Autostop/down must be set for all newly '
+ 'launched clusters.')
+ return sky.MutatedUserRequest(
+ task=user_request.task,
+ skypilot_config=user_request.skypilot_config)

examples/admin_policy/reject_all_config.yaml

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

examples/admin_policy/task_label_config.yaml

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

sky/__init__.py

@@ -110,10 +110,14 @@ def set_proxy_env_var(proxy_var: str, urllib_var: Optional[str]):
 from sky.jobs.core import spot_tail_logs
 from sky.optimizer import Optimizer
 from sky.optimizer import OptimizeTarget
-from sky.policy import MutatedUserTask
-from sky.policy import UserTask
+# Admin Policy interfaces
+from sky.policy import AdminPolicy
+from sky.policy import MutatedUserRequest
+from sky.policy import OperationArgs
+from sky.policy import UserRequest
 from sky.resources import Resources
 from sky.skylet.job_lib import JobStatus
+from sky.skypilot_config import NestedConfig
 from sky.status_lib import ClusterStatus
 from sky.task import Task
 
@@ -187,6 +191,10 @@ def set_proxy_env_var(proxy_var: str, urllib_var: Optional[str]):
  # core APIs Storage Management
  'storage_ls',
  'storage_delete',
- 'UserTask',
- 'MutatedUserTask',
+ # Admin Policy
+ 'UserRequest',
+ 'MutatedUserRequest',
+ 'AdminPolicy',
+ 'NestedConfig',
+ 'OperationArgs',
 ]

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 policy."""
+ pass