add user data to admin policy (#7800)

* add user info

* encode/decode test

* another ut

* add comment

* fix syntax warning

---------

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

Author: SeungjinYang

sky/admin_policy.py

@@ -9,6 +9,7 @@
 
 import sky
 from sky import exceptions
+from sky import models
 from sky.adaptors import common as adaptors_common
 from sky.utils import common_utils
 from sky.utils import config_utils
@@ -51,6 +52,7 @@ class _UserRequestBody(pydantic.BaseModel):
     skypilot_config: str
     request_options: Optional[RequestOptions] = None
     at_client_side: bool = False
+    user: str
 
 
 @dataclasses.dataclass
@@ -73,11 +75,15 @@ class UserRequest:
         skypilot_config: Global skypilot config to be used in this request.
         request_options: Request options. It is None for jobs and services.
         at_client_side: Is the request intercepted by the policy at client-side?
+        user: User who made the request.
+              Only available on the server side.
+              This value is None if at_client_side is True.
     """
     task: 'sky.Task'
     skypilot_config: 'sky.Config'
     request_options: Optional['RequestOptions'] = None
     at_client_side: bool = False
+    user: Optional['models.User'] = None
 
     def encode(self) -> str:
         return _UserRequestBody(
@@ -86,11 +92,18 @@ def encode(self) -> str:
                 self.skypilot_config)),
             request_options=self.request_options,
             at_client_side=self.at_client_side,
+            user=(yaml_utils.dump_yaml_str(self.user.to_dict())
+                  if self.user is not None else ''),
         ).model_dump_json()
 
     @classmethod
     def decode(cls, body: str) -> 'UserRequest':
         user_request_body = _UserRequestBody.model_validate_json(body)
+        user_dict = yaml_utils.read_yaml_str(
+            user_request_body.user) if user_request_body.user != '' else None
+        user = models.User(
+            id=user_dict['id'],
+            name=user_dict['name']) if user_dict is not None else None
         return cls(
             task=sky.Task.from_yaml_config(
                 yaml_utils.read_yaml_all_str(user_request_body.task)[0]),
@@ -99,6 +112,7 @@ def decode(cls, body: str) -> 'UserRequest':
                     user_request_body.skypilot_config)[0]),
             request_options=user_request_body.request_options,
             at_client_side=user_request_body.at_client_side,
+            user=user,
         )
 
 

sky/utils/admin_policy_utils.py

@@ -2,6 +2,7 @@
 import contextlib
 import copy
 import importlib
+import typing
 from typing import Iterator, Optional, Tuple, Union
 import urllib.parse
 
@@ -19,6 +20,9 @@
 
 logger = sky_logging.init_logger(__name__)
 
+if typing.TYPE_CHECKING:
+    from sky import models
+
 
 def _is_url(policy_string: str) -> bool:
     """Check if the policy string is a URL."""
@@ -126,9 +130,13 @@ def apply(
     if policy is None:
         return dag, skypilot_config.to_dict()
 
+    user = None
     if at_client_side:
         logger.info(f'Applying client admin policy: {policy}')
     else:
+        # When being called by the server, the middleware has set the
+        # current user and this information is available at this point.
+        user = common_utils.get_current_user()
         logger.info(f'Applying server admin policy: {policy}')
     config = copy.deepcopy(skypilot_config.to_dict())
     mutated_dag = dag_lib.Dag()
@@ -137,7 +145,7 @@ def apply(
     mutated_config = None
     for task in dag.tasks:
         user_request = admin_policy.UserRequest(task, config, request_options,
-                                                at_client_side)
+                                                at_client_side, user)
         try:
             mutated_user_request = policy.apply(user_request)
         # Avoid duplicate exception wrapping.

tests/unit_tests/test_admin_policy.py

@@ -14,6 +14,7 @@
 
 import sky
 from sky import exceptions
+from sky import models
 from sky import sky_logging
 from sky import skypilot_config
 from sky.utils import admin_policy_utils
@@ -379,6 +380,21 @@ def test_use_local_gcp_credentials_policy(add_example_policy_paths, task):
             policy.apply(user_request)
 
 
+def test_user_request_encode_decode(task):
+    with mock.patch('sky.utils.common_utils.get_current_user',
+                    return_value=models.User(id='123', name='test')):
+        user_request = sky.UserRequest(task=task,
+                                       skypilot_config=sky.Config(),
+                                       at_client_side=False,
+                                       user=models.User(id='123', name='test'))
+        encoded_request = user_request.encode()
+        decoded_request = sky.UserRequest.decode(encoded_request)
+        assert repr(decoded_request.task) == repr(task)
+        assert decoded_request.skypilot_config == sky.Config()
+        assert decoded_request.at_client_side == False
+        assert decoded_request.user == models.User(id='123', name='test')
+
+
 def test_restful_policy(add_example_policy_paths, task):
     """Test RestfulAdminPolicy for various scenarios."""
 

tests/unit_tests/test_admin_policy_restful.py

@@ -19,6 +19,7 @@
 import threading
 import time
 from typing import Optional, Tuple
+from unittest import mock
 
 from fastapi import FastAPI
 from fastapi import Request
@@ -28,6 +29,7 @@
 
 import sky
 from sky import admin_policy
+from sky import models
 from sky import skypilot_config
 from sky.utils import admin_policy_utils
 from sky.utils import common_utils
@@ -371,6 +373,44 @@ def test_restful_policy_with_request_options(monkeypatch):
             os.unlink(config_path)
 
 
+def test_restful_policy_with_user(monkeypatch):
+    """Test RESTful admin policy receiving user information."""
+    with mock.patch('sky.utils.common_utils.get_current_user',
+                    return_value=models.User(id='123', name='test')):
+        with PolicyServer() as server:
+            ImageIdInspectorPolicy.received_requests.clear()
+
+            # Create a test task
+            task = create_test_task()
+
+            # Create temporary config and apply policy using existing function
+            with tempfile.NamedTemporaryFile(mode='w',
+                                             suffix='.yaml',
+                                             delete=False) as f:
+                f.write(f'admin_policy: http://127.0.0.1:{server.port}\n')
+                config_path = f.name
+
+            try:
+                dag, config = _load_task_and_apply_policy(
+                    task, config_path, monkeypatch)
+
+                # Verify the policy was called with proper request structure
+                assert len(ImageIdInspectorPolicy.received_requests) == 1
+                request = ImageIdInspectorPolicy.received_requests[0]
+
+                # Check that user information was properly included
+                assert request.user is not None
+                assert request.user.id == '123'
+                assert request.user.name == 'test'
+
+                # Check that we got valid results back
+                assert dag is not None
+                assert len(dag.tasks) == 1
+                assert config is not None
+            finally:
+                os.unlink(config_path)
+
+
 def test_restful_policy_basic_functionality(monkeypatch):
     """Test basic RESTful admin policy functionality using real patterns."""
     with PolicyServer() as server: