[API Server] Fix admin policy enforcement on `validate` and `optimize` #4820

romilbhardwaj · 2025-02-26T02:52:56Z

Closes #4818.

Our admin policies were not being enforced during the validate and optimize calls, which would break policies which dynamically fetch Kubernetes credentials.

This PR fixes that but with a caveat: admin policies are not applied in the optimizer during failover. This is consistent with our previous admin policy enforcement, but we should revisit that design.

Code formatting: bash format.sh
Manually tested with dynamic k8s fetching admin policy
All smoke tests: pytest tests/test_smoke.py

romilbhardwaj · 2025-02-26T06:51:00Z

/smoke-test

romilbhardwaj · 2025-02-26T06:51:13Z

/smoke-test --kubernetes

Michaelvll

Thanks for identifying and fixing this @romilbhardwaj! Left a comment : )

Michaelvll · 2025-02-26T07:00:19Z

sky/client/sdk.py

+    cluster_name: Optional[str] = None,
+    idle_minutes_to_autostop: Optional[int] = None,
+    down: bool = False,  # pylint: disable=redefined-outer-name
+    dryrun: bool = False


Should we emphasize that these are only required when admin_policy is specified, or more aggressively we just make this:

Suggested change

cluster_name: Optional[str] = None,

idle_minutes_to_autostop: Optional[int] = None,

down: bool = False, # pylint: disable=redefined-outer-name

dryrun: bool = False

admin_policy_request_options: Optional[admin_policy.RequestOptions] = None

Yep good idea, I was also thinking about it but wasn't sure if we want to expose admin_policy.RequestOptions in the SDK to end users. Should be ok.

Michaelvll · 2025-02-26T07:01:40Z

sky/server/requests/payloads.py

@@ -116,12 +117,14 @@ class CheckBody(RequestBody):
 class ValidateBody(RequestBody):
    """The request body for the validate endpoint."""
    dag: str
+    request_options: admin_policy.RequestOptions


Can this be made optional and we can error out on the server side if admin_policy is set but RequestOptions are not supplied?

Yep, made optional now. When None, we will now follow the default behavior in admin_policy_utils.apply().

(Current behavior is it will silently accept request_options as None, but we should update that behavior in a future PR. That's a deeper change since serve and jobs launch also invoke admin_policy_utils.apply() without request_options)

Michaelvll · 2025-02-26T07:15:00Z

sky/server/server.py

+
+    # We make the validation a separate request as it may require expensive
+    # network calls if an admin policy is applied.
+    # TODO: Our current launch process is broken down into three calls:
+    # validate, optimize, and launch. This requires us to apply the admin policy
+    # in each step, which may be an expensive operation. We should consolidate
+    # these into a single call or have a TTL cache for (task, admin_policy)
+    # pairs.
    logger.debug(f'Validating tasks: {validate_body.dag}')
-    try:
-        dag = dag_utils.load_chain_dag_from_yaml_str(validate_body.dag)
-        for task in dag.tasks:
-            # Will validate workdir and file_mounts in the backend, as those
-            # need to be validated after the files are uploaded to the SkyPilot
-            # API server with `upload_mounts_to_api_server`.
-            task.validate_name()
-            task.validate_run()
-            for r in task.resources:
-                r.validate()
-    except Exception as e:  # pylint: disable=broad-except
-        raise fastapi.HTTPException(
-            status_code=400, detail=exceptions.serialize_exception(e)) from e
+    executor.schedule_request(request_id=request.state.request_id,
+                              request_name='validate',
+                              request_body=validate_body,
+                              ignore_return_value=True,
+                              func=core.validate_dag,
+                              schedule_type=requests_lib.ScheduleType.SHORT)


We need to double check how much additional overhead we add to a normal launch with this modification. It may not worth it to slow down all the launch requests just for the admin policy, which is a less common codepath : )

Thanks for flagging, it does add a significant penalty:

[Master] validate without executor: 0.007498 seconds [This branch] validate with executor: 0.150482 seconds

I'm going to revert this change in favor of keeping the codepath fast for the common case when there's no admin policy.

Benchmark script:

from sky.client import sdk from sky import Task, Dag, Resources from sky.utils import dag_utils import timeit task = Task(name='test', run="echo 'Hello, World!'") # resources = Resources(labels={'invalid'*50: 'bar'}) # task.set_resources(resources) dag = dag_utils.convert_entrypoint_to_dag(task) # First run once to get the result and warm up any caches print("Running initial validation...") sdk.validate(dag) print("Initial validation complete.") # Now benchmark with timeit print("\nStarting benchmark...") number_of_runs = 10 setup = "from __main__ import sdk, dag" # Run the benchmark print(f"Running sdk.validate() {number_of_runs} times...") execution_time = timeit.timeit("sdk.validate(dag)", setup=setup, number=number_of_runs) average_time = execution_time / number_of_runs # Print results print("\nBenchmark Results:") print(f"Total time for {number_of_runs} runs: {execution_time:.6f} seconds") print(f"Average time per run: {average_time:.6f} seconds")

romilbhardwaj · 2025-02-27T00:08:30Z

/quicktest-core

romilbhardwaj · 2025-02-27T02:12:37Z

/quicktest-core

romilbhardwaj · 2025-02-27T19:16:02Z

/smoke-test

Michaelvll

Thanks @romilbhardwaj! Looks mostly good to me. Left some minor comments.

sky/client/sdk.py

sky/core.py

Michaelvll · 2025-02-28T00:42:04Z

sky/server/requests/payloads.py

+    """The request body for the optimize endpoint."""
+    dag: str
+    minimize: common_lib.OptimizeTarget = common_lib.OptimizeTarget.COST
+    request_options: Optional[admin_policy.RequestOptions]


Can we try if a API server exists, and a user upgrade to the current PR, will the CLIs show clear warning/error that the api server needs to be restarted? We may want to avoid pydantic errors.

The case you mentioned works (Old API server, new client), but the reverse (New API server, old client) did not. I've bumped up the API version to 3 to throw a clean error.

…o api_admin_policy_fix

Michaelvll

Thanks @romilbhardwaj for fixing this! It is a great catch.

Michaelvll · 2025-03-03T18:53:29Z

sky/client/sdk.py

@@ -212,13 +213,17 @@ def list_accelerator_counts(
 @annotations.client_api
 def optimize(
    dag: 'sky.Dag',
-    minimize: common.OptimizeTarget = common.OptimizeTarget.COST
+    minimize: common.OptimizeTarget = common.OptimizeTarget.COST,
+    admin_policy_request_options: Optional[admin_policy.RequestOptions] = None


We should probably fix the error message where Kubernetes raises a None context. We can mention the active/allowed context there?

Agreed, will do in another PR

#4822 should also help alleviate this.

sky/server/server.py

romilbhardwaj · 2025-03-03T20:32:28Z

/smoke-test

romilbhardwaj added 10 commits February 25, 2025 15:50

Add admin policy to validate

616bfbf

Add admin policy to validate

57a6ab2

Add admin policy to optimize

714e65c

docs

6423aad

imports

5010494

Move dag validation to core

0b65617

Fixes

d37644b

lint

c5164dc

Add comments

41f4d26

lint

aa33534

Michaelvll reviewed Feb 26, 2025

View reviewed changes

romilbhardwaj added 6 commits February 26, 2025 15:32

Fixed executor based validate implementation

20828be

Revert executor based validate implementation

6362328

lint

6f94e3b

lint

78169ce

Add validation during optimize

0d9da0d

lint

478b2a4

romilbhardwaj added 2 commits February 26, 2025 16:10

Remove validate from core

37cd372

Remove admin policy apply when validating dag for exec

0c7c078

romilbhardwaj requested a review from Michaelvll February 27, 2025 00:30

Michaelvll reviewed Feb 28, 2025

View reviewed changes

romilbhardwaj added 3 commits February 28, 2025 12:13

comments

07f2f3d

Merge branch 'master' of https://github.com/skypilot-org/skypilot int…

2382722

…o api_admin_policy_fix

Bump API version

ff27bcc

romilbhardwaj requested a review from Michaelvll March 2, 2025 22:18

Michaelvll approved these changes Mar 3, 2025

View reviewed changes

comments

6691dc1

romilbhardwaj merged commit 87a6b59 into master Mar 3, 2025
19 checks passed

romilbhardwaj deleted the api_admin_policy_fix branch March 3, 2025 21:23

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[API Server] Fix admin policy enforcement on `validate` and `optimize` #4820

[API Server] Fix admin policy enforcement on `validate` and `optimize` #4820

romilbhardwaj commented Feb 26, 2025 •

edited

Loading

romilbhardwaj commented Feb 26, 2025

romilbhardwaj commented Feb 26, 2025

Michaelvll left a comment

Michaelvll Feb 26, 2025

romilbhardwaj Feb 26, 2025

Michaelvll Feb 26, 2025

romilbhardwaj Feb 26, 2025

Michaelvll Feb 26, 2025

romilbhardwaj Feb 26, 2025

romilbhardwaj Feb 26, 2025

romilbhardwaj commented Feb 27, 2025

romilbhardwaj commented Feb 27, 2025

romilbhardwaj commented Feb 27, 2025

Michaelvll left a comment

Michaelvll Feb 28, 2025

romilbhardwaj Feb 28, 2025

Michaelvll left a comment

Michaelvll Mar 3, 2025

romilbhardwaj Mar 3, 2025

romilbhardwaj Mar 3, 2025

romilbhardwaj commented Mar 3, 2025

[API Server] Fix admin policy enforcement on validate and optimize #4820

[API Server] Fix admin policy enforcement on validate and optimize #4820

Conversation

romilbhardwaj commented Feb 26, 2025 • edited Loading

romilbhardwaj commented Feb 26, 2025

romilbhardwaj commented Feb 26, 2025

Michaelvll left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

romilbhardwaj commented Feb 27, 2025

romilbhardwaj commented Feb 27, 2025

romilbhardwaj commented Feb 27, 2025

Michaelvll left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Michaelvll left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

romilbhardwaj commented Mar 3, 2025

[API Server] Fix admin policy enforcement on `validate` and `optimize` #4820

[API Server] Fix admin policy enforcement on `validate` and `optimize` #4820

romilbhardwaj commented Feb 26, 2025 •

edited

Loading