[Core] Admin policy enforcement plugin #3966

Michaelvll · 2024-09-20T01:00:41Z

This PR allows the admin to bring a customized policy enforcement script in Python, which can apply any mutation to all user tasks based on customized policy requirements.

Usage:

User-side

In ~/.sky/config.yaml:

policy: my_package.MyPolicy

Admin-side

Create the MyPolicy class with the following signature:

class MyPolicy(sky.AdminPolicy):
  @classmethod
  def validate_and_muteate(user_task: sky.UserRequest) -> sky.MutatedUserRequest:
      pass

See policy.py for the definition of these two types.

TODO:

Check when to apply policy, e.g. skip the policy apply for exec? (Let's apply for both for now)
Upload policy package to controllers (We now apply the policy locally before uploading to controller, so this should be fine)
Add to config.rst
Add to some admin doc page for how to set this up
Policy name/hash label (blocked by [UX] Make labels work in task.yaml without the requirement for specifying cloud #3965)

Tested (run the relevant ones):

Code formatting: bash format.sh
Any manual or new tests for this PR (please specify below)
- normal launch with customized labels
- normal launch/exec with autostop enforcement
- jobs launch with customized labels
- jobs launch with autostop enforcement
- service launch with customized labels

export SKYPILOT_CONFIG='/home/gcpuser/skypilot-dev/examples/admin_policy/config_label_config.yaml'
sky serve up -n test-global-label examples/serve/http_server/task.yaml --cloud gcp

All smoke tests: pytest tests/test_smoke.py
Relevant individual smoke tests: pytest tests/test_smoke.py::test_fill_in_the_name
Backward compatibility tests: conda deactivate; bash -i tests/backward_compatibility_tests.sh

docs/source/reference/config.rst

examples/policy/example_policy/example_policy/skypilot_policy.py

sky/policy.py

concretevitamin · 2024-09-20T20:28:24Z

sky/serve/core.py

@@ -124,6 +125,8 @@ def up(

 _validate_service_task(task)

+ task = policy.Policy().apply_to_task(task)


A bit unclear if it's in-place since it also returns the task. How about "apply_in_place"?

It is not in-place. Either way is fine, but I found not in-place is more commonly seen operations.

sky/utils/dag_utils.py

sky/policy.py

Co-authored-by: Zongheng Yang <zongheng.y@gmail.com>

…olicy-hook

Co-authored-by: Zongheng Yang <zongheng.y@gmail.com>

…olicy-hook

concretevitamin

LGTM @Michaelvll. I think two TODOs before merging:

(1) Smoke tests, since this touches config / backend
(2) Update PR description (policy)

docs/source/cloud-setup/policy.rst

examples/admin_policy/example_policy/example_policy/skypilot_policy.py

sky/execution.py

tests/unit_tests/test_admin_policy.py

examples/admin_policy/example_policy/example_policy/skypilot_policy.py

concretevitamin · 2024-09-23T14:50:01Z

sky/execution.py

- 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
+ cluster_running = (cluster_record is not None and


Q: what's the latency impact of this policy on sky launch/exec? Is it small? Do we expect users to use this example policy and be ok with the latency?

docs/source/cloud-setup/policy.rst

Co-authored-by: Zongheng Yang <zongheng.y@gmail.com>

…olicy-hook

…-hook

Michaelvll · 2024-09-24T01:42:33Z

Tested:

pytest tests/test_smoke.py (failed [Azure] test_cancel_azure failed due to error of resnet user program #3974 , [GCP] test_gcp_disk_tier for best tier failed due to best tier not mapping to pd-extreme #3973; fixed test_skyserve_fast_update)

…3976) Use wget instead of git clone for faster downloading

* [LLM] Update qwen examples (skypilot-org#3957) * update qwen examples * Fix misalign * Qwen 2.5 support (skypilot-org#3959) * Update qwen example for 2.5 release * Add support for qwen 2.5 example * Qwen 2.5 k8s (skypilot-org#3960) * Update qwen example for 2.5 release * Add support for qwen 2.5 example * add kubernetes * Integrating the Yi series models (skypilot-org#3958) * Add files via upload * Update and rename qwen2-7b.yaml to yi15-6b.yaml * Add files via upload * Update yi15-9b.yaml * Update yi15-34b.yaml * Update yi15-6b.yaml * Add files via upload * Update yicoder-1_5b.yaml * Update yicoder-9b.yaml * Add files via upload * Update yi15-34b.yaml * Update yi15-6b.yaml * Update yi15-9b.yaml * Update yicoder-1_5b.yaml * Update yicoder-9b.yaml * [Test] Fix Smoke Test `test-skyserve-fast-update` (skypilot-org#3956) * init * add newline * [LLM] Add Qwen2-VL multimodal example (skypilot-org#3961) Add multimodal example * Update README.md (skypilot-org#3969) * Add files via upload * Update and rename qwen2-7b.yaml to yi15-6b.yaml * Add files via upload * Update yi15-9b.yaml * Update yi15-34b.yaml * Update yi15-6b.yaml * Add files via upload * Update yicoder-1_5b.yaml * Update yicoder-9b.yaml * Add files via upload * Update yi15-34b.yaml * Update yi15-6b.yaml * Update yi15-9b.yaml * Update yicoder-1_5b.yaml * Update yicoder-9b.yaml * Update README.md * [Core] Admin policy enforcement plugin (skypilot-org#3966) * support policy hook * test task labels * Add test for policy that sets labels * Fix comment * format * use -e to make test related files visible * Add config.rst * Fix test * fix config rst * Apply policy to service * add policy for serving * Add docs * fix * format * Update interface * fix * Fix * fix * Fix test config * Fix mutated config * fix * Add policy doc * rename * minor * Add additional arguments for autostop * fix mypy * format * rejected message * format * Update sky/utils/policy_utils.py Co-authored-by: Zongheng Yang <zongheng.y@gmail.com> * Update sky/utils/policy_utils.py Co-authored-by: Zongheng Yang <zongheng.y@gmail.com> * Fix * Update examples/admin_policy/example_policy/example_policy/__init__.py Co-authored-by: Zongheng Yang <zongheng.y@gmail.com> * Update docs/source/reference/config.rst Co-authored-by: Zongheng Yang <zongheng.y@gmail.com> * Address comments * format * changes in examples * Fix enforce autostop * Fix autostop enforcement * fix test * Update docs/source/cloud-setup/policy.rst Co-authored-by: Zongheng Yang <zongheng.y@gmail.com> * Update sky/admin_policy.py Co-authored-by: Zongheng Yang <zongheng.y@gmail.com> * Update sky/admin_policy.py Co-authored-by: Zongheng Yang <zongheng.y@gmail.com> * wip * Update docs/source/cloud-setup/policy.rst Co-authored-by: Zongheng Yang <zongheng.y@gmail.com> * Update docs/source/cloud-setup/policy.rst Co-authored-by: Zongheng Yang <zongheng.y@gmail.com> * Update docs/source/cloud-setup/policy.rst Co-authored-by: Zongheng Yang <zongheng.y@gmail.com> * fix * fix * fix * Use sky.status for autostop * update policy * Update docs/source/cloud-setup/policy.rst Co-authored-by: Zongheng Yang <zongheng.y@gmail.com> * fix policy.rst * Add comment * Fix logging * fix CI * Update docs/source/cloud-setup/policy.rst Co-authored-by: Zongheng Yang <zongheng.y@gmail.com> * Use sphnix inline code * Add comment * fix skypilot config file mounts for jobs and serve --------- Co-authored-by: Zongheng Yang <zongheng.y@gmail.com> * [k8s] Autodown Serve controller on Kubernetes (skypilot-org#3984) * Add autodown for skyserve on k8s * lint * [Tests] Add missing changes from skypilot-org#3966 for fast service update test (skypilot-org#3976) Use wget instead of git clone for faster downloading * [Paperspace] add A4000, P4000, GPU+ (skypilot-org#3991) add A4000, P4000, GPU+ * [Docs] Fix highlighting in code block (skypilot-org#3994) Fix highlighting in code block Fixes skypilot-org#3993 * [LLM] Llama 3.2 guide (skypilot-org#3990) * Add llama 3.2 example * update * length * fix * update * update cpus limit * Use 11B instead for better performance * update * update * Add link * Fix reference * Fix vllm version * Update llm/llama-3_2/README.md Co-authored-by: Zongheng Yang <zongheng.y@gmail.com> * Update llm/llama-3_2/README.md Co-authored-by: Zongheng Yang <zongheng.y@gmail.com> * Update llm/llama-3_2/README.md Co-authored-by: Zongheng Yang <zongheng.y@gmail.com> * Update llm/llama-3_2/README.md Co-authored-by: Zongheng Yang <zongheng.y@gmail.com> * Fix title * news * no need to pin transformers * remove cover photo for now --------- Co-authored-by: Zongheng Yang <zongheng.y@gmail.com> * [k8s] Add cluster attributes(autodown, idle-minutes-to-autostop) as annotations to the pod (skypilot-org#3870) * add autodown annotations to the k8s pod * revert kubernetes ray template * revert backend_utils from invasive approach * nit * revert from invasive approaches * revert * updated approach * nit * nit * Use constant to represent idle_minutes_to_autostop for cancellation * revert using constants for cancel * nit * nit * add smoke tests * Update sky/provision/kubernetes/utils.py Co-authored-by: Romil Bhardwaj <romil.bhardwaj@gmail.com> * fix comments * nit * remove loops and annotate one by one * format * update with autodown annotation with context * format --------- Co-authored-by: Romil Bhardwaj <romil.bhardwaj@gmail.com> * [Examples] Add airflow example (skypilot-org#3982) * Airflow example * Airflow example * Airflow example * Airflow example * wip * Update airflow examples * Update airflow examples * Update airflow examples * Add to readme * Add to readme * Add to readme * lint * updates * less salesy * comments * comments * comments * [UX] default to minimal logging (no module/line number/timestamp). (skypilot-org#3980) * [UX] default to minimal logging (no module/line number/timestamp). * Fix mypy. * Fix typing * Update sky/utils/env_options.py Co-authored-by: Tian Xia <cblmemo@gmail.com> * Update sky/utils/env_options.py Co-authored-by: Tian Xia <cblmemo@gmail.com> * Account for debug flag. * Remove prefixes from docs. --------- Co-authored-by: Tian Xia <cblmemo@gmail.com> * Revert "[UX] default to minimal logging (no module/line number/timestamp)." (skypilot-org#4003) Revert "[UX] default to minimal logging (no module/line number/timestamp). (#…" This reverts commit b96a5b4. * [Docs] Clarify k8s private registry usage in docs (skypilot-org#3998) * Clarify k8s private registry auth in docs. * comments * [Docs] Various polishing. (skypilot-org#4002) * [Docs] Various polishing. * update * Reword. * lint --------- Co-authored-by: Zhanghao Wu <zhanghao.wu@outlook.com> Co-authored-by: Haijian Wang <130898843+Haijian06@users.noreply.github.com> Co-authored-by: Tian Xia <cblmemo@gmail.com> Co-authored-by: Zongheng Yang <zongheng.y@gmail.com> Co-authored-by: Romil Bhardwaj <romil.bhardwaj@berkeley.edu> Co-authored-by: Andy Lee <andylizf@outlook.com> Co-authored-by: landscapepainter <34902420+landscapepainter@users.noreply.github.com> Co-authored-by: Romil Bhardwaj <romil.bhardwaj@gmail.com>

Michaelvll added 4 commits September 19, 2024 22:01

support policy hook

cb28b8d

test task labels

b64efa0

Add test for policy that sets labels

cf89929

Fix comment

54c93ea

Michaelvll changed the title ~~[Core] Customized policy enforcement Plugin~~ [Core] Customized policy enforcement plugin Sep 20, 2024

Michaelvll added 10 commits September 20, 2024 01:14

format

1d1c500

use -e to make test related files visible

a0bdb2c

Add config.rst

543e66a

Fix test

520a2a1

fix config rst

b533351

Apply policy to service

466f7fe

add policy for serving

050dc7a

Add docs

31e0174

fix

0c74f2a

format

48a6cc9

concretevitamin reviewed Sep 20, 2024

View reviewed changes

Michaelvll added 14 commits September 20, 2024 22:33

Update interface

1ca5a8a

fix

14b2346

Fix

cb39c73

fix

1e3ddef

Fix test config

aa87df7

Fix mutated config

28487a4

fix

d1f0480

Add policy doc

f42ace5

rename

c04f3dc

minor

58f413c

Add additional arguments for autostop

52053bd

fix mypy

4a4f682

format

a8d1c44

rejected message

6c73d81

Michaelvll and others added 12 commits September 22, 2024 17:28

Update docs/source/cloud-setup/policy.rst

67552d7

Co-authored-by: Zongheng Yang <zongheng.y@gmail.com>

fix

7de757e

Merge branch 'policy-hook' of github.com:skypilot-org/skypilot into p…

8443ddc

…olicy-hook

fix

7fbc30d

fix

92b68fc

Use sky.status for autostop

7d8af9a

update policy

5b37f47

Update docs/source/cloud-setup/policy.rst

c7af310

Co-authored-by: Zongheng Yang <zongheng.y@gmail.com>

fix policy.rst

cb232a8

Merge branch 'policy-hook' of github.com:skypilot-org/skypilot into p…

5e9f544

…olicy-hook

Add comment

deb4c92

Fix logging

cbff59d

Michaelvll changed the title ~~[Core] Customized policy enforcement plugin~~ [Core] Admin policy enforcement plugin Sep 23, 2024

fix CI

1fe350a

Michaelvll requested a review from concretevitamin September 23, 2024 07:30

concretevitamin approved these changes Sep 23, 2024

View reviewed changes

Michaelvll and others added 6 commits September 23, 2024 10:07

Update docs/source/cloud-setup/policy.rst

2e8e41c

Co-authored-by: Zongheng Yang <zongheng.y@gmail.com>

Use sphnix inline code

aae42ce

Merge branch 'policy-hook' of github.com:skypilot-org/skypilot into p…

73c8fb7

…olicy-hook

Add comment

11bbd5e

fix skypilot config file mounts for jobs and serve

3630535

Merge branch 'master' of github.com:skypilot-org/skypilot into policy…

e020dea

…-hook

Michaelvll added this pull request to the merge queue Sep 24, 2024

Merged via the queue into master with commit 800f7d6 Sep 24, 2024
20 checks passed

Michaelvll deleted the policy-hook branch September 24, 2024 04:44

Michaelvll mentioned this pull request Sep 24, 2024

[Policy] Add label for policy applied to a user task #3983

Open

github-merge-queue bot pushed a commit that referenced this pull request Sep 25, 2024

[Tests] Add missing changes from #3966 for fast service update test (#…

be92944

…3976) Use wget instead of git clone for faster downloading

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Core] Admin policy enforcement plugin #3966

[Core] Admin policy enforcement plugin #3966

Michaelvll commented Sep 20, 2024 •

edited

Loading

concretevitamin Sep 20, 2024

Michaelvll Sep 20, 2024

concretevitamin left a comment

concretevitamin Sep 23, 2024

Michaelvll commented Sep 24, 2024

		@@ -124,6 +125,8 @@ def up(

		_validate_service_task(task)

		task = policy.Policy().apply_to_task(task)

[Core] Admin policy enforcement plugin #3966

[Core] Admin policy enforcement plugin #3966

Conversation

Michaelvll commented Sep 20, 2024 • edited Loading

Usage:

User-side

Admin-side

concretevitamin Sep 20, 2024

Choose a reason for hiding this comment

Michaelvll Sep 20, 2024

Choose a reason for hiding this comment

concretevitamin left a comment

Choose a reason for hiding this comment

concretevitamin Sep 23, 2024

Choose a reason for hiding this comment

Michaelvll commented Sep 24, 2024

Michaelvll commented Sep 20, 2024 •

edited

Loading