skypilot-nightly 1.0.0.dev20250819__py3-none-any.whl → 1.0.0.dev20250820__py3-none-any.whl

sky/global_user_state.py

@@ -55,6 +55,13 @@ _SQLALCHEMY_ENGINE_LOCK = threading.Lock()
 DEFAULT_CLUSTER_EVENT_RETENTION_HOURS = 24.0
 MIN_CLUSTER_EVENT_DAEMON_INTERVAL_SECONDS = 3600
 
+_UNIQUE_CONSTRAINT_FAILED_ERROR_MSGS = [
+    # sqlite
+    'UNIQUE constraint failed',
+    # postgres
+    'duplicate key value violates unique constraint',
+]
+
 Base = declarative.declarative_base()
 
 config_table = sqlalchemy.Table(
@@ -735,17 +742,17 @@ def add_cluster_event(cluster_name: str,
                 ))
             session.commit()
         except sqlalchemy.exc.IntegrityError as e:
-            if 'UNIQUE constraint failed' in str(e):
-                # This can happen if the cluster event is added twice.
-                # We can ignore this error unless the caller requests
-                # to expose the error.
-                if expose_duplicate_error:
-                    raise db_utils.UniqueConstraintViolationError(
-                        value=reason, message=str(e))
-                else:
-                    pass
-            else:
-                raise e
+            for msg in _UNIQUE_CONSTRAINT_FAILED_ERROR_MSGS:
+                if msg in str(e):
+                    # This can happen if the cluster event is added twice.
+                    # We can ignore this error unless the caller requests
+                    # to expose the error.
+                    if expose_duplicate_error:
+                        raise db_utils.UniqueConstraintViolationError(
+                            value=reason, message=str(e))
+                    else:
+                        return
+            raise e
 
 
 def get_last_cluster_event(cluster_hash: str,

sky/jobs/server/core.py

@@ -188,11 +188,11 @@ def launch(
 
     dag_uuid = str(uuid.uuid4().hex[:4])
     dag = dag_utils.convert_entrypoint_to_dag(entrypoint)
-    dag.resolve_and_validate_volumes()
     # 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.resolve_and_validate_volumes()
     if not dag.is_chain():
         with ux_utils.print_exception_no_traceback():
             raise ValueError('Only single-task or chain DAG is '

sky/models.py

@@ -108,3 +108,4 @@ class VolumeConfig(pydantic.BaseModel):
     name_on_cloud: str
     size: Optional[str]
     config: Dict[str, Any] = {}
+    labels: Optional[Dict[str, str]] = None

sky/provision/aws/config.py

@@ -498,8 +498,8 @@ def _vpc_id_from_security_group_ids(ec2: 'mypy_boto3_ec2.ServiceResource',
     return vpc_ids[0]
 
 
-def _get_vpc_id_by_name(ec2: 'mypy_boto3_ec2.ServiceResource', vpc_name: str,
-                        region: str) -> str:
+def get_vpc_id_by_name(ec2: 'mypy_boto3_ec2.ServiceResource', vpc_name: str,
+                       region: str) -> str:
     """Returns the VPC ID of the unique VPC with a given name.
 
     Exits with code 1 if:
@@ -532,7 +532,7 @@ def _get_subnet_and_vpc_id(ec2: 'mypy_boto3_ec2.ServiceResource',
                            use_internal_ips: bool,
                            vpc_name: Optional[str]) -> Tuple[Any, str]:
     if vpc_name is not None:
-        vpc_id_of_sg = _get_vpc_id_by_name(ec2, vpc_name, region)
+        vpc_id_of_sg = get_vpc_id_by_name(ec2, vpc_name, region)
     elif security_group_ids:
         vpc_id_of_sg = _vpc_id_from_security_group_ids(ec2, security_group_ids)
     else:
@@ -614,8 +614,8 @@ def _get_or_create_vpc_security_group(ec2: 'mypy_boto3_ec2.ServiceResource',
             due to AWS service issues.
     """
     # Figure out which security groups with this name exist for each VPC...
-    security_group = _get_security_group_from_vpc_id(ec2, vpc_id,
-                                                     expected_sg_name)
+    security_group = get_security_group_from_vpc_id(ec2, vpc_id,
+                                                    expected_sg_name)
     if security_group is not None:
         return security_group
 
@@ -631,7 +631,7 @@ def _get_or_create_vpc_security_group(ec2: 'mypy_boto3_ec2.ServiceResource',
             # The security group already exists, but we didn't see it
             # because of eventual consistency.
             logger.warning(f'{expected_sg_name} already exists when creating.')
-            security_group = _get_security_group_from_vpc_id(
+            security_group = get_security_group_from_vpc_id(
                 ec2, vpc_id, expected_sg_name)
             assert (security_group is not None and
                     security_group.group_name == expected_sg_name), (
@@ -646,8 +646,8 @@ def _get_or_create_vpc_security_group(ec2: 'mypy_boto3_ec2.ServiceResource',
         logger.warning(message)
         raise exceptions.NoClusterLaunchedError(message) from e
 
-    security_group = _get_security_group_from_vpc_id(ec2, vpc_id,
-                                                     expected_sg_name)
+    security_group = get_security_group_from_vpc_id(ec2, vpc_id,
+                                                    expected_sg_name)
     assert security_group is not None, 'Failed to create security group'
     logger.info(f'Created new security group {colorama.Style.BRIGHT}'
                 f'{security_group.group_name}{colorama.Style.RESET_ALL} '
@@ -655,9 +655,9 @@ def _get_or_create_vpc_security_group(ec2: 'mypy_boto3_ec2.ServiceResource',
     return security_group
 
 
-def _get_security_group_from_vpc_id(ec2: 'mypy_boto3_ec2.ServiceResource',
-                                    vpc_id: str,
-                                    group_name: str) -> Optional[Any]:
+def get_security_group_from_vpc_id(ec2: 'mypy_boto3_ec2.ServiceResource',
+                                   vpc_id: str,
+                                   group_name: str) -> Optional[Any]:
     """Get security group by VPC ID and group name."""
     existing_groups = list(
         ec2.security_groups.filter(Filters=[{

sky/provision/aws/instance.py

@@ -18,6 +18,7 @@ from sky.clouds import aws as aws_cloud
 from sky.clouds.utils import aws_utils
 from sky.provision import common
 from sky.provision import constants
+from sky.provision.aws import config as aws_config
 from sky.provision.aws import utils
 from sky.utils import common_utils
 from sky.utils import resources_utils
@@ -685,7 +686,9 @@ def terminate_instances(
                                   filters,
                                   included_instances=None,
                                   excluded_instances=None)
-    default_sg = _get_sg_from_name(ec2, aws_cloud.DEFAULT_SECURITY_GROUP_NAME)
+    default_sg = aws_config.get_security_group_from_vpc_id(
+        ec2, _get_vpc_id(provider_config),
+        aws_cloud.DEFAULT_SECURITY_GROUP_NAME)
     if sg_name == aws_cloud.DEFAULT_SECURITY_GROUP_NAME:
         # Case 1: The default SG is used, we don't need to ensure instance are
         # terminated.
@@ -727,30 +730,6 @@ def terminate_instances(
     #  of most cloud implementations (including AWS).
 
 
-def _get_sg_from_name(
-    ec2: Any,
-    sg_name: str,
-) -> Any:
-    # GroupNames will only filter SGs in the default VPC, so we need to use
-    # Filters here. Ref:
-    # https://boto3.amazonaws.com/v1/documentation/api/1.26.112/reference/services/ec2/service-resource/security_groups.html  # pylint: disable=line-too-long
-    sgs = ec2.security_groups.filter(Filters=[{
-        'Name': 'group-name',
-        'Values': [sg_name]
-    }])
-    num_sg = len(list(sgs))
-    if num_sg == 0:
-        logger.warning(f'Expected security group {sg_name} not found. ')
-        return None
-    if num_sg > 1:
-        # TODO(tian): Better handle this case. Maybe we can check when creating
-        # the SG and throw an error if there is already an existing SG with the
-        # same name.
-        logger.warning(f'Found {num_sg} security groups with name {sg_name}. ')
-        return None
-    return list(sgs)[0]
-
-
 def _maybe_move_to_new_sg(
     instance: Any,
     expected_sg: Any,
@@ -803,7 +782,9 @@ def open_ports(
         with ux_utils.print_exception_no_traceback():
             raise ValueError('Instance with cluster name '
                              f'{cluster_name_on_cloud} not found.')
-    sg = _get_sg_from_name(ec2, sg_name)
+    sg = aws_config.get_security_group_from_vpc_id(ec2,
+                                                   _get_vpc_id(provider_config),
+                                                   sg_name)
     if sg is None:
         with ux_utils.print_exception_no_traceback():
             raise ValueError('Cannot find new security group '
@@ -899,7 +880,9 @@ def cleanup_ports(
         # We only want to delete the SG that is dedicated to this cluster (i.e.,
         # this cluster have opened some ports).
         return
-    sg = _get_sg_from_name(ec2, sg_name)
+    sg = aws_config.get_security_group_from_vpc_id(ec2,
+                                                   _get_vpc_id(provider_config),
+                                                   sg_name)
     if sg is None:
         logger.warning(
             'Find security group failed. Skip cleanup security group.')
@@ -1010,3 +993,23 @@ def get_cluster_info(
         provider_name='aws',
         provider_config=provider_config,
     )
+
+
+def _get_vpc_id(provider_config: Dict[str, Any]) -> str:
+    region = provider_config['region']
+    ec2 = _default_ec2_resource(provider_config['region'])
+    if 'vpc_name' in provider_config:
+        return aws_config.get_vpc_id_by_name(ec2, provider_config['vpc_name'],
+                                             region)
+    else:
+        # Retrieve the default VPC name from the region.
+        response = ec2.meta.client.describe_vpcs(Filters=[{
+            'Name': 'isDefault',
+            'Values': ['true']
+        }])
+        if len(response['Vpcs']) == 0:
+            raise ValueError(f'No default VPC found in region {region}')
+        elif len(response['Vpcs']) > 1:
+            raise ValueError(f'Multiple default VPCs found in region {region}')
+        else:
+            return response['Vpcs'][0]['VpcId']

sky/provision/do/utils.py

@@ -30,7 +30,7 @@ POSSIBLE_CREDENTIALS_PATHS = [
 INITIAL_BACKOFF_SECONDS = 10
 MAX_BACKOFF_FACTOR = 10
 MAX_ATTEMPTS = 6
-SSH_KEY_NAME_ON_DO = f'sky-key-{common_utils.get_user_hash()}'
+SSH_KEY_NAME_ON_DO_PREFIX = 'sky-key-'
 
 _client = None
 _ssh_key_id = None
@@ -125,7 +125,7 @@ def ssh_key_id(public_key: str):
 
         request = {
             'public_key': public_key,
-            'name': SSH_KEY_NAME_ON_DO,
+            'name': SSH_KEY_NAME_ON_DO_PREFIX + common_utils.get_user_hash(),
         }
         _ssh_key_id = client().ssh_keys.create(body=request)['ssh_key']
     return _ssh_key_id

sky/provision/kubernetes/network_utils.py

@@ -4,13 +4,13 @@ import time
 import typing
 from typing import Dict, List, Optional, Tuple, Union
 
-import sky
 from sky import exceptions
 from sky import sky_logging
 from sky import skypilot_config
 from sky.adaptors import common as adaptors_common
 from sky.adaptors import kubernetes
 from sky.provision.kubernetes import utils as kubernetes_utils
+from sky.utils import directory_utils
 from sky.utils import kubernetes_enums
 from sky.utils import ux_utils
 
@@ -80,7 +80,7 @@ def get_networking_mode(
 def fill_loadbalancer_template(namespace: str, context: Optional[str],
                                service_name: str, ports: List[int],
                                selector_key: str, selector_value: str) -> Dict:
-    template_path = os.path.join(sky.__root_dir__, 'templates',
+    template_path = os.path.join(directory_utils.get_sky_dir(), 'templates',
                                  _LOADBALANCER_TEMPLATE_NAME)
     if not os.path.exists(template_path):
         raise FileNotFoundError(
@@ -116,7 +116,7 @@ def fill_ingress_template(namespace: str, context: Optional[str],
                           service_details: List[Tuple[str, int,
                                                       str]], ingress_name: str,
                           selector_key: str, selector_value: str) -> Dict:
-    template_path = os.path.join(sky.__root_dir__, 'templates',
+    template_path = os.path.join(directory_utils.get_sky_dir(), 'templates',
                                  _INGRESS_TEMPLATE_NAME)
     if not os.path.exists(template_path):
         raise FileNotFoundError(

sky/provision/kubernetes/utils.py

@@ -14,7 +14,6 @@ import typing
 from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
 from urllib.parse import urlparse
 
-import sky
 from sky import clouds
 from sky import exceptions
 from sky import global_user_state
@@ -31,6 +30,7 @@ from sky.skylet import constants
 from sky.utils import annotations
 from sky.utils import common_utils
 from sky.utils import config_utils
+from sky.utils import directory_utils
 from sky.utils import env_options
 from sky.utils import kubernetes_enums
 from sky.utils import schemas
@@ -2444,7 +2444,7 @@ def clean_zombie_ssh_jump_pod(namespace: str, context: Optional[str],
 
 def fill_ssh_jump_template(ssh_key_secret: str, ssh_jump_image: str,
                            ssh_jump_name: str, service_type: str) -> Dict:
-    template_path = os.path.join(sky.__root_dir__, 'templates',
+    template_path = os.path.join(directory_utils.get_sky_dir(), 'templates',
                                  'kubernetes-ssh-jump.yml.j2')
     if not os.path.exists(template_path):
         raise FileNotFoundError(

sky/provision/kubernetes/volume.py

@@ -203,6 +203,8 @@ def _get_pvc_spec(namespace: str,
             },
         }
     }
+    if config.labels:
+        pvc_spec['metadata']['labels'].update(config.labels)
     storage_class = config.config.get('storage_class_name')
     if storage_class is not None:
         pvc_spec['spec']['storageClassName'] = storage_class

sky/serve/replica_managers.py

@@ -48,6 +48,13 @@ _PROCESS_POOL_REFRESH_INTERVAL = 20
 _RETRY_INIT_GAP_SECONDS = 60
 _DEFAULT_DRAIN_SECONDS = 120
 
+# TODO(tian): Backward compatibility. Remove this after 3 minor release, i.e.
+# 0.13.0. We move the ProcessStatus to common_utils.ProcessStatus in #6666, but
+# old ReplicaInfo in database will still tries to unpickle using ProcessStatus
+# in replica_managers. We set this alias to avoid breaking changes. See #6729
+# for more details.
+ProcessStatus = common_utils.ProcessStatus
+
 
 # TODO(tian): Combine this with
 # sky/spot/recovery_strategy.py::StrategyExecutor::launch

sky/serve/server/impl.py

@@ -129,11 +129,11 @@ def up(
                              f'{constants.CLUSTER_NAME_VALID_REGEX}')
 
     dag = dag_utils.convert_entrypoint_to_dag(task)
-    dag.resolve_and_validate_volumes()
     # 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.resolve_and_validate_volumes()
     dag.pre_mount_volumes()
     task = dag.tasks[0]
     assert task.service is not None

sky/server/requests/payloads.py

@@ -453,6 +453,7 @@ class VolumeApplyBody(RequestBody):
     zone: Optional[str] = None
     size: Optional[str] = None
     config: Optional[Dict[str, Any]] = None
+    labels: Optional[Dict[str, str]] = None
 
 
 class VolumeDeleteBody(RequestBody):

sky/server/requests/serializers/encoders.py

@@ -10,6 +10,7 @@ from typing import Any, Dict, List, Optional, Tuple
 
 from sky.schemas.api import responses
 from sky.server import constants as server_constants
+from sky.utils import serialize_utils
 
 if typing.TYPE_CHECKING:
     from sky import backends
@@ -22,6 +23,9 @@ handlers: Dict[str, Any] = {}
 
 def pickle_and_encode(obj: Any) -> str:
     try:
+        # Apply backwards compatibility processing at the lowest level
+        # to catch any handles that might have bypassed the encoders
+        obj = serialize_utils.prepare_handle_for_backwards_compatibility(obj)
         return base64.b64encode(pickle.dumps(obj)).decode('utf-8')
     except TypeError as e:
         raise ValueError(f'Failed to pickle object: {obj}') from e
@@ -58,7 +62,9 @@ def encode_status(
     for cluster in clusters:
         response_cluster = cluster.model_dump()
         response_cluster['status'] = cluster['status'].value
-        response_cluster['handle'] = pickle_and_encode(cluster['handle'])
+        handle = serialize_utils.prepare_handle_for_backwards_compatibility(
+            cluster['handle'])
+        response_cluster['handle'] = pickle_and_encode(handle)
         response_cluster['storage_mounts_metadata'] = pickle_and_encode(
             response_cluster['storage_mounts_metadata'])
         response.append(response_cluster)
@@ -70,6 +76,7 @@ def encode_launch(
     job_id_handle: Tuple[Optional[int], Optional['backends.ResourceHandle']]
 ) -> Dict[str, Any]:
     job_id, handle = job_id_handle
+    handle = serialize_utils.prepare_handle_for_backwards_compatibility(handle)
     return {
         'job_id': job_id,
         'handle': pickle_and_encode(handle),
@@ -78,6 +85,9 @@ def encode_launch(
 
 @register_encoder('start')
 def encode_start(resource_handle: 'backends.CloudVmRayResourceHandle') -> str:
+    resource_handle = (
+        serialize_utils.prepare_handle_for_backwards_compatibility(
+            resource_handle))
     return pickle_and_encode(resource_handle)
 
 
@@ -143,7 +153,9 @@ def _encode_serve_status(
         service_status['status'] = service_status['status'].value
         for replica_info in service_status.get('replica_info', []):
             replica_info['status'] = replica_info['status'].value
-            replica_info['handle'] = pickle_and_encode(replica_info['handle'])
+            handle = serialize_utils.prepare_handle_for_backwards_compatibility(
+                replica_info['handle'])
+            replica_info['handle'] = pickle_and_encode(handle)
     return service_statuses
 
 

sky/server/server.py

@@ -83,6 +83,8 @@ else:
 
 P = ParamSpec('P')
 
+_SERVER_USER_HASH_KEY = 'server_user_hash'
+
 
 def _add_timestamp_prefix_for_server_logs() -> None:
     server_logger = sky_logging.init_logger('sky.server')
@@ -1821,6 +1823,35 @@ async def root():
     return fastapi.responses.RedirectResponse(url='/dashboard/')
 
 
+def _init_or_restore_server_user_hash():
+    """Restores the server user hash from the global user state db.
+
+    The API server must have a stable user hash across restarts and potential
+    multiple replicas. Thus we persist the user hash in db and restore it on
+    startup. When upgrading from old version, the user hash will be read from
+    the local file (if any) to keep the user hash consistent.
+    """
+
+    def apply_user_hash(user_hash: str) -> None:
+        # For local API server, the user hash in db and local file should be
+        # same so there is no harm to override here.
+        common_utils.set_user_hash_locally(user_hash)
+        # Refresh the server user hash for current process after restore or
+        # initialize the user hash in db, child processes will get the correct
+        # server id from the local cache file.
+        common_lib.refresh_server_id()
+
+    user_hash = global_user_state.get_system_config(_SERVER_USER_HASH_KEY)
+    if user_hash is not None:
+        apply_user_hash(user_hash)
+        return
+
+    # Initial deployment, generate a user hash and save it to the db.
+    user_hash = common_utils.get_user_hash()
+    global_user_state.set_system_config(_SERVER_USER_HASH_KEY, user_hash)
+    apply_user_hash(user_hash)
+
+
 if __name__ == '__main__':
     import uvicorn
 
@@ -1830,6 +1861,8 @@ if __name__ == '__main__':
     global_user_state.initialize_and_get_db()
     # Initialize request db
     requests_lib.reset_db_and_logs()
+    # Restore the server user hash
+    _init_or_restore_server_user_hash()
 
     parser = argparse.ArgumentParser()
     parser.add_argument('--host', default='127.0.0.1')

sky/setup_files/dependencies.py

@@ -72,12 +72,27 @@ install_requires = [
     'aiohttp',
 ]
 
+# See requirements-dev.txt for the version of grpc and protobuf
+# used to generate the code during development.
+
+# The grpc version at runtime has to be newer than the version
+# used to generate the code.
+GRPC = 'grpcio>=1.63.0'
+# >= 5.26.1 because the runtime version can't be older than the version
+# used to generate the code.
+# < 7.0.0 because code generated for a major version V will be supported by
+# protobuf runtimes of version V and V+1.
+# https://protobuf.dev/support/cross-version-runtime-guarantee
+PROTOBUF = 'protobuf>=5.26.1, < 7.0.0'
+
 server_dependencies = [
     'casbin',
     'sqlalchemy_adapter',
     'passlib',
     'pyjwt',
     'aiohttp',
+    GRPC,
+    PROTOBUF,
 ]
 
 local_ray = [
@@ -88,18 +103,9 @@ local_ray = [
     'ray[default] >= 2.2.0, != 2.6.0',
 ]
 
-# See requirements-dev.txt for the version of grpc and protobuf
-# used to generate the code during development.
 remote = [
-    # The grpc version at runtime has to be newer than the version
-    # used to generate the code.
-    'grpcio>=1.63.0',
-    # >= 5.26.1 because the runtime version can't be older than the version
-    # used to generate the code.
-    # < 7.0.0 because code generated for a major version V will be supported by
-    # protobuf runtimes of version V and V+1.
-    # https://protobuf.dev/support/cross-version-runtime-guarantee
-    'protobuf >= 5.26.1, < 7.0.0',
+    GRPC,
+    PROTOBUF,
 ]
 
 # NOTE: Change the templates/jobs-controller.yaml.j2 file if any of the

sky/utils/common.py

@@ -11,18 +11,38 @@ from sky.utils import common_utils
 
 SKY_SERVE_CONTROLLER_PREFIX: str = 'sky-serve-controller-'
 JOB_CONTROLLER_PREFIX: str = 'sky-jobs-controller-'
+
 # We use the user hash (machine-specific) for the controller name. It will be
 # the same across the whole lifecycle of the server, including:
-# 1. all requests, because this global variable is set once during server
-#    starts.
-# 2. SkyPilot API server restarts, as long as the `~/.sky` folder is persisted
-#    and the env var set during starting the server is the same.
+# 1. all requests, because all the server processes share the same user hash
+#    cache file.
+# 2. SkyPilot API server restarts, because the API server will restore the
+#    user hash from the global user state db on startup.
+# 3. Potential multiple server replicas, because multiple server replicas of
+#    a same deployment will share the same global user state db.
 # This behavior is the same for the local API server (where SERVER_ID is the
 # same as the normal user hash). This ensures backwards-compatibility with jobs
 # controllers from before #4660.
-SERVER_ID = common_utils.get_user_hash()
-SKY_SERVE_CONTROLLER_NAME: str = f'{SKY_SERVE_CONTROLLER_PREFIX}{SERVER_ID}'
-JOB_CONTROLLER_NAME: str = f'{JOB_CONTROLLER_PREFIX}{SERVER_ID}'
+SERVER_ID: str
+SKY_SERVE_CONTROLLER_NAME: str
+JOB_CONTROLLER_NAME: str
+
+
+def refresh_server_id() -> None:
+    """Refresh the server id.
+
+    This function is used to ensure the server id is read from the authorative
+    source.
+    """
+    global SERVER_ID
+    global SKY_SERVE_CONTROLLER_NAME
+    global JOB_CONTROLLER_NAME
+    SERVER_ID = common_utils.get_user_hash()
+    SKY_SERVE_CONTROLLER_NAME = f'{SKY_SERVE_CONTROLLER_PREFIX}{SERVER_ID}'
+    JOB_CONTROLLER_NAME = f'{JOB_CONTROLLER_PREFIX}{SERVER_ID}'
+
+
+refresh_server_id()
 
 
 @contextlib.contextmanager

sky/utils/common_utils.py

@@ -28,7 +28,6 @@ from sky.adaptors import common as adaptors_common
 from sky.skylet import constants
 from sky.usage import constants as usage_constants
 from sky.utils import annotations
-from sky.utils import common_utils
 from sky.utils import ux_utils
 from sky.utils import validator
 
@@ -41,7 +40,7 @@ else:
     psutil = adaptors_common.LazyImport('psutil')
     yaml = adaptors_common.LazyImport('yaml')
 
-_USER_HASH_FILE = os.path.expanduser('~/.sky/user_hash')
+USER_HASH_FILE = os.path.expanduser('~/.sky/user_hash')
 USER_HASH_LENGTH = 8
 
 # We are using base36 to reduce the length of the hash. 2 chars -> 36^2 = 1296
@@ -131,21 +130,26 @@ def get_user_hash() -> str:
         assert user_hash is not None
         return user_hash
 
-    if os.path.exists(_USER_HASH_FILE):
+    if os.path.exists(USER_HASH_FILE):
         # Read from cached user hash file.
-        with open(_USER_HASH_FILE, 'r', encoding='utf-8') as f:
+        with open(USER_HASH_FILE, 'r', encoding='utf-8') as f:
             # Remove invalid characters.
             user_hash = f.read().strip()
         if is_valid_user_hash(user_hash):
             return user_hash
 
     user_hash = generate_user_hash()
-    os.makedirs(os.path.dirname(_USER_HASH_FILE), exist_ok=True)
-    with open(_USER_HASH_FILE, 'w', encoding='utf-8') as f:
-        f.write(user_hash)
+    set_user_hash_locally(user_hash)
     return user_hash
 
 
+def set_user_hash_locally(user_hash: str) -> None:
+    """Sets the user hash to local file."""
+    os.makedirs(os.path.dirname(USER_HASH_FILE), exist_ok=True)
+    with open(USER_HASH_FILE, 'w', encoding='utf-8') as f:
+        f.write(user_hash)
+
+
 def base36_encode(hex_str: str) -> str:
     """Converts a hex string to a base36 string."""
     int_value = int(hex_str, 16)
@@ -343,7 +347,7 @@ def get_current_user() -> 'models.User':
 
 def get_current_user_name() -> str:
     """Returns the current user name."""
-    name = common_utils.get_current_user().name
+    name = get_current_user().name
     assert name is not None
     return name
 
@@ -886,7 +890,7 @@ def get_cleaned_username(username: str = '') -> str:
     Returns:
       A cleaned username.
     """
-    username = username or common_utils.get_current_user_name()
+    username = username or get_current_user_name()
     username = username.lower()
     username = re.sub(r'[^a-z0-9-_]', '', username)
     username = re.sub(r'^[0-9-]+', '', username)

sky/utils/directory_utils.py

@@ -0,0 +1,12 @@
+"""Directory utilities."""
+
+import os
+
+# This file is in '<project_root>/sky/utils/directory_utils.py'
+# So we need to go up 2 levels to get to the '<project_root>/sky' directory
+SKY_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+
+
+def get_sky_dir():
+    """Get the sky root directory."""
+    return SKY_DIR

sky/utils/env_options.py

@@ -24,6 +24,9 @@ class Options(enum.Enum):
     # running in a Buildkite container environment, which requires special
     # handling for networking between containers.
     RUNNING_IN_BUILDKITE = ('BUILDKITE', False)
+    # Internal: This is used for testing to enable grpc for communication
+    # between the API server and the Skylet.
+    ENABLE_GRPC = ('SKYPILOT_ENABLE_GRPC', False)
 
     def __init__(self, env_var: str, default: bool) -> None:
         super().__init__()