skypilot-nightly 1.0.0.dev20250520__py3-none-any.whl → 1.0.0.dev20250522__py3-none-any.whl

sky/skypilot_config.py

@@ -52,6 +52,7 @@ import contextlib
 import copy
 import json
 import os
+import tempfile
 import threading
 import typing
 from typing import Any, Dict, Iterator, List, Optional, Tuple
@@ -62,6 +63,7 @@ from sky.adaptors import common as adaptors_common
 from sky.skylet import constants
 from sky.utils import common_utils
 from sky.utils import config_utils
+from sky.utils import context
 from sky.utils import schemas
 from sky.utils import ux_utils
 
@@ -105,13 +107,66 @@ ENV_VAR_PROJECT_CONFIG = f'{constants.SKYPILOT_ENV_VAR_PREFIX}PROJECT_CONFIG'
 _GLOBAL_CONFIG_PATH = '~/.sky/config.yaml'
 _PROJECT_CONFIG_PATH = '.sky.yaml'
 
-# The loaded config.
-_dict = config_utils.Config()
-_loaded_config_path: Optional[str] = None
-_config_overridden: bool = False
+
+class ConfigContext:
+
+    def __init__(self,
+                 config: config_utils.Config = config_utils.Config(),
+                 config_path: Optional[str] = None,
+                 config_overridden: bool = False):
+        self.config = config
+        self.config_path = config_path
+        self.config_overridden = config_overridden
+
+
+# The global loaded config.
+_global_config_context = ConfigContext()
 _reload_config_lock = threading.Lock()
 
 
+def _get_config_context() -> ConfigContext:
+    """Get config context for current context.
+
+    If no context is available, the global config context is returned.
+    """
+    ctx = context.get()
+    if not ctx:
+        return _global_config_context
+    if ctx.config_context is None:
+        # Config context for current context is not initialized, inherit from
+        # the global one.
+        ctx.config_context = ConfigContext(
+            config=copy.deepcopy(_global_config_context.config),
+            config_path=_global_config_context.config_path,
+            config_overridden=_global_config_context.config_overridden,
+        )
+    return ctx.config_context
+
+
+def _get_loaded_config() -> config_utils.Config:
+    return _get_config_context().config
+
+
+def _set_loaded_config(config: config_utils.Config) -> None:
+    _get_config_context().config = config
+
+
+def _get_loaded_config_path() -> Optional[str]:
+    return _get_config_context().config_path
+
+
+def _set_loaded_config_path(path: Optional[str]) -> None:
+    _get_config_context().config_path = path
+
+
+def _is_config_overridden() -> bool:
+    return _get_config_context().config_overridden
+
+
+def _set_config_overridden(config_overridden: bool) -> None:
+    _get_config_context().config_overridden = config_overridden
+
+
 def get_user_config_path() -> str:
     """Returns the path to the user config file."""
     return _GLOBAL_CONFIG_PATH
@@ -224,7 +279,7 @@ def get_nested(keys: Tuple[str, ...],
     Returns:
         The value of the nested key, or 'default_value' if not found.
     """
-    return _dict.get_nested(
+    return _get_loaded_config().get_nested(
         keys,
         default_value,
         override_configs,
@@ -237,14 +292,14 @@ def set_nested(keys: Tuple[str, ...], value: Any) -> Dict[str, Any]:
 
     Like get_nested(), if any key is not found, this will not raise an error.
     """
-    copied_dict = copy.deepcopy(_dict)
+    copied_dict = copy.deepcopy(_get_loaded_config())
     copied_dict.set_nested(keys, value)
     return dict(**copied_dict)
 
 
 def to_dict() -> config_utils.Config:
     """Returns a deep-copied version of the current config."""
-    return copy.deepcopy(_dict)
+    return copy.deepcopy(_get_loaded_config())
 
 
 def _get_config_file_path(envvar: str) -> Optional[str]:
@@ -345,10 +400,9 @@ def _parse_dotlist(dotlist: List[str]) -> config_utils.Config:
 
 
 def _reload_config_from_internal_file(internal_config_path: str) -> None:
-    global _dict, _loaded_config_path
     # Reset the global variables, to avoid using stale values.
-    _dict = config_utils.Config()
-    _loaded_config_path = None
+    _set_loaded_config(config_utils.Config())
+    _set_loaded_config_path(None)
 
     config_path = os.path.expanduser(internal_config_path)
     if not os.path.exists(config_path):
@@ -359,14 +413,13 @@ def _reload_config_from_internal_file(internal_config_path: str) -> None:
                 'exist. Please double check the path or unset the env var: '
                 f'unset {ENV_VAR_SKYPILOT_CONFIG}')
     logger.debug(f'Using config path: {config_path}')
-    _dict = parse_config_file(config_path)
-    _loaded_config_path = config_path
+    _set_loaded_config(parse_config_file(config_path))
+    _set_loaded_config_path(config_path)
 
 
 def _reload_config_as_server() -> None:
-    global _dict
     # Reset the global variables, to avoid using stale values.
-    _dict = config_utils.Config()
+    _set_loaded_config(config_utils.Config())
 
     overrides: List[config_utils.Config] = []
     server_config = get_server_config()
@@ -382,13 +435,12 @@ def _reload_config_as_server() -> None:
         logger.debug(
             f'server config: \n'
             f'{common_utils.dump_yaml_str(dict(overlaid_server_config))}')
-    _dict = overlaid_server_config
+    _set_loaded_config(overlaid_server_config)
 
 
 def _reload_config_as_client() -> None:
-    global _dict
     # Reset the global variables, to avoid using stale values.
-    _dict = config_utils.Config()
+    _set_loaded_config(config_utils.Config())
 
     overrides: List[config_utils.Config] = []
     user_config = get_user_config()
@@ -407,15 +459,15 @@ def _reload_config_as_client() -> None:
         logger.debug(
             f'client config (before task and CLI overrides): \n'
             f'{common_utils.dump_yaml_str(dict(overlaid_client_config))}')
-    _dict = overlaid_client_config
+    _set_loaded_config(overlaid_client_config)
 
 
 def loaded_config_path() -> Optional[str]:
     """Returns the path to the loaded config file, or
     '<overridden>' if the config is overridden."""
-    if _config_overridden:
+    if _is_config_overridden():
         return '<overridden>'
-    return _loaded_config_path
+    return _get_loaded_config_path()
 
 
 # Load on import, synchronization is guaranteed by python interpreter.
@@ -424,21 +476,20 @@ _reload_config()
 
 def loaded() -> bool:
     """Returns if the user configurations are loaded."""
-    return bool(_dict)
+    return bool(_get_loaded_config())
 
 
 @contextlib.contextmanager
 def override_skypilot_config(
         override_configs: Optional[Dict[str, Any]]) -> Iterator[None]:
     """Overrides the user configurations."""
-    global _dict, _config_overridden
     # TODO(SKY-1215): allow admin user to extend the disallowed keys or specify
     # allowed keys.
     if not override_configs:
         # If no override configs (None or empty dict), do nothing.
         yield
         return
-    original_config = _dict
+    original_config = _get_loaded_config()
     override_configs = config_utils.Config(override_configs)
     disallowed_diff_keys = []
     for key in constants.SKIPPED_CLIENT_OVERRIDE_KEYS:
@@ -455,7 +506,7 @@ def override_skypilot_config(
             'and will be ignored. Remove these keys to disable this warning. '
             'If you want to specify it, please modify it on server side or '
             'contact your administrator.')
-    config = _dict.get_nested(
+    config = original_config.get_nested(
         keys=tuple(),
         default_value=None,
         override_configs=dict(override_configs),
@@ -469,8 +520,8 @@ def override_skypilot_config(
             'https://docs.skypilot.co/en/latest/reference/config.html. '  # pylint: disable=line-too-long
             'Error: ',
             skip_none=False)
-        _config_overridden = True
-        _dict = config
+        _set_config_overridden(True)
+        _set_loaded_config(config)
         yield
     except exceptions.InvalidSkyPilotConfigError as e:
         with ux_utils.print_exception_no_traceback():
@@ -483,8 +534,43 @@ def override_skypilot_config(
                 f'{common_utils.dump_yaml_str(dict(override_configs))}\n'
                 f'Details: {e}') from e
     finally:
-        _dict = original_config
-        _config_overridden = False
+        _set_loaded_config(original_config)
+        _set_config_overridden(False)
+
+
+@contextlib.contextmanager
+def replace_skypilot_config(new_configs: config_utils.Config) -> Iterator[None]:
+    """Replaces the global config with the new configs.
+
+    This function is concurrent safe when it is:
+    1. called in different processes;
+    2. or called in a same process but with different context, refer to
+       sky_utils.context for more details.
+    """
+    original_config = _get_loaded_config()
+    original_env_var = os.environ.get(ENV_VAR_SKYPILOT_CONFIG)
+    if new_configs != original_config:
+        # Modify the global config of current process or context
+        _set_loaded_config(new_configs)
+        with tempfile.NamedTemporaryFile(delete=False,
+                                         mode='w',
+                                         prefix='mutated-skypilot-config-',
+                                         suffix='.yaml') as temp_file:
+            common_utils.dump_yaml(temp_file.name, dict(**new_configs))
+        # Modify the env var of current process or context so that the
+        # new config will be used by spawned sub-processes.
+        # Note that this code modifies os.environ directly because it
+        # will be hijacked to be context-aware if a context is active.
+        os.environ[ENV_VAR_SKYPILOT_CONFIG] = temp_file.name
+        yield
+        # Restore the original config and env var.
+        _set_loaded_config(original_config)
+        if original_env_var:
+            os.environ[ENV_VAR_SKYPILOT_CONFIG] = original_env_var
+        else:
+            os.environ.pop(ENV_VAR_SKYPILOT_CONFIG, None)
+    else:
+        yield
 
 
 def _compose_cli_config(cli_config: Optional[List[str]]) -> config_utils.Config:
@@ -529,11 +615,11 @@ def apply_cli_config(cli_config: Optional[List[str]]) -> Dict[str, Any]:
         cli_config: A path to a config file or a comma-separated
         list of key-value pairs.
     """
-    global _dict
     parsed_config = _compose_cli_config(cli_config)
     if sky_logging.logging_enabled(logger, sky_logging.DEBUG):
         logger.debug(f'applying following CLI overrides: \n'
                      f'{common_utils.dump_yaml_str(dict(parsed_config))}')
-    _dict = overlay_skypilot_config(original_config=_dict,
-                                    override_configs=parsed_config)
+    _set_loaded_config(
+        overlay_skypilot_config(original_config=_get_loaded_config(),
+                                override_configs=parsed_config))
     return parsed_config

sky/task.py

@@ -512,6 +512,7 @@ class Task:
         # storage objects with the storage/storage_mount objects.
         fm_storages = []
         file_mounts = config.pop('file_mounts', None)
+        volumes = []
         if file_mounts is not None:
             copy_mounts = {}
             for dst_path, src in file_mounts.items():
@@ -521,7 +522,27 @@ class Task:
                 # If the src is not a str path, it is likely a dict. Try to
                 # parse storage object.
                 elif isinstance(src, dict):
-                    fm_storages.append((dst_path, src))
+                    if (src.get('store') ==
+                            storage_lib.StoreType.VOLUME.value.lower()):
+                        # Build the volumes config for resources.
+                        volume_config = {
+                            'path': dst_path,
+                        }
+                        if src.get('name'):
+                            volume_config['name'] = src.get('name')
+                        persistent = src.get('persistent', False)
+                        volume_config['auto_delete'] = not persistent
+                        volume_config_detail = src.get('config', {})
+                        volume_config.update(volume_config_detail)
+                        volumes.append(volume_config)
+                        source_path = src.get('source')
+                        if source_path:
+                            # For volume, copy the source path to the
+                            # data directory of the volume mount point.
+                            copy_mounts[
+                                f'{dst_path.rstrip("/")}/data'] = source_path
+                    else:
+                        fm_storages.append((dst_path, src))
                 else:
                     with ux_utils.print_exception_no_traceback():
                         raise ValueError(f'Unable to parse file_mount '
@@ -599,6 +620,8 @@ class Task:
                 'experimental.config_overrides')
             resources_config[
                 '_cluster_config_overrides'] = cluster_config_override
+        if volumes:
+            resources_config['volumes'] = volumes
         task.set_resources(sky.Resources.from_yaml_config(resources_config))
 
         service = config.pop('service', None)

sky/templates/gcp-ray.yml.j2

@@ -109,12 +109,27 @@ available_node_types:
 {%- if tpu_vm %}
       acceleratorType: {{tpu_type}}
       runtimeVersion: {{runtime_version}}
+    {%- if volumes %}
+      dataDisks:
+      {%- for volume in volumes %}
+        {%- if volume.source %}
+        - sourceDisk: {{volume.source}}
+        {%- endif %}
+        {%- if volume.attach_mode %}
+          mode: {{volume.attach_mode}}
+        {%- endif %}
+      {%- endfor %}
+    {%- endif %}
       metadata:
         # TPU VM's metadata has different format than normal VMs.
         # After replacing the variables, this will become username:ssh_public_key_content.
         # This is a specific syntax required by GCP https://cloud.google.com/compute/docs/connect/add-ssh-keys
         ssh-keys: |-
               skypilot:ssh_user:skypilot:ssh_public_key_content
+      {%- if user_data is not none %}
+        startup-script: |-
+          {{ user_data | indent(10) }}
+      {%- endif %}
   {%- if use_spot %}
       schedulingConfig:
         preemptible: true
@@ -138,6 +153,34 @@ available_node_types:
           {%- if disk_iops %}
             provisionedIops: {{disk_iops}}
           {%- endif %}
+      {%- for volume in volumes %}
+        - boot: false
+          autoDelete: {{volume.auto_delete}}
+          type: {{volume.storage_type}}
+          deviceName: {{volume.device_name}}
+        {%- if volume.source %}
+          source: {{volume.source}}
+        {%- endif %}
+        {%- if volume.attach_mode %}
+          mode: {{volume.attach_mode}}
+        {%- endif %}
+        {%- if volume.interface_type %}
+          interface: {{volume.interface_type}}
+        {%- endif %}
+        {%- if volume.disk_tier %}
+          initializeParams:
+            diskType: zones/{{zones}}/diskTypes/{{volume.disk_tier}}
+        {%- endif %}
+          {%- if volume.disk_name %}
+            diskName: {{volume.disk_name}}
+          {%- endif %}
+          {%- if volume.disk_size %}
+            diskSizeGb: {{volume.disk_size}}
+          {%- endif %}
+          {%- if volume.iops %}
+            provisionedIops: {{volume.iops}}
+          {%- endif %}
+      {%- endfor %}
   {%- if gpu is not none %}
       guestAccelerators:
         - acceleratorType: projects/{{gcp_project_id}}/zones/{{zones}}/acceleratorTypes/{{gpu}}
@@ -157,7 +200,7 @@ available_node_types:
   {%- if user_data is not none %}
           - key: user-data
             value: |-
-              {{ user_data | indent(10) }}
+              {{ user_data | indent(14) }}
   {%- endif %}
   {%- if use_spot or gpu is not none %}
       scheduling:

sky/templates/nebius-ray.yml.j2

@@ -47,11 +47,9 @@ available_node_types:
       ImageId: {{image_id}}
       DiskSize: {{disk_size}}
       UserData: |
-        {%- if docker_image is not none %}
         runcmd:
           - sudo sed -i 's/^#\?AllowTcpForwarding.*/AllowTcpForwarding yes/' /etc/ssh/sshd_config
           - systemctl restart sshd
-        {%- endif %}
 
         {# Two available OS images:
                1. ubuntu22.04-driverless - requires Docker installation

sky/utils/admin_policy_utils.py

@@ -1,9 +1,8 @@
 """Admin policy utils."""
+import contextlib
 import copy
 import importlib
-import os
-import tempfile
-from typing import Optional, Tuple, Union
+from typing import Iterator, Optional, Tuple, Union
 
 import colorama
 
@@ -52,9 +51,31 @@ def _get_policy_cls(
     return policy_cls
 
 
+@contextlib.contextmanager
+def apply_and_use_config_in_current_request(
+    entrypoint: Union['dag_lib.Dag', 'task_lib.Task'],
+    request_options: Optional[admin_policy.RequestOptions] = None,
+) -> Iterator['dag_lib.Dag']:
+    """Applies an admin policy and override SkyPilot config for current request
+
+    This is a helper function of `apply()` that applies an admin policy and
+    overrides the SkyPilot config for the current request as a context manager.
+    The original SkyPilot config will be restored when the context manager is
+    exited.
+
+    Refer to `apply()` for more details.
+    """
+    original_config = skypilot_config.to_dict()
+    dag, mutated_config = apply(entrypoint, request_options)
+    if mutated_config != original_config:
+        with skypilot_config.replace_skypilot_config(mutated_config):
+            yield dag
+    else:
+        yield dag
+
+
 def apply(
     entrypoint: Union['dag_lib.Dag', 'task_lib.Task'],
-    use_mutated_config_in_current_request: bool = True,
     request_options: Optional[admin_policy.RequestOptions] = None,
 ) -> Tuple['dag_lib.Dag', config_utils.Config]:
     """Applies an admin policy (if registered) to a DAG or a task.
@@ -85,8 +106,7 @@ def apply(
         return dag, skypilot_config.to_dict()
 
     logger.info(f'Applying policy: {policy}')
-    original_config = skypilot_config.to_dict()
-    config = copy.deepcopy(original_config)
+    config = copy.deepcopy(skypilot_config.to_dict())
     mutated_dag = dag_lib.Dag()
     mutated_dag.name = dag.name
 
@@ -126,22 +146,6 @@ def apply(
         mutated_dag.graph.add_edge(mutated_dag.tasks[u_idx],
                                    mutated_dag.tasks[v_idx])
 
-    if (use_mutated_config_in_current_request and
-            original_config != mutated_config):
-        with tempfile.NamedTemporaryFile(
-                delete=False,
-                mode='w',
-                prefix='policy-mutated-skypilot-config-',
-                suffix='.yaml') as temp_file:
-
-            common_utils.dump_yaml(temp_file.name, dict(**mutated_config))
-            os.environ[skypilot_config.ENV_VAR_SKYPILOT_CONFIG] = temp_file.name
-            logger.debug(f'Updated SkyPilot config: {temp_file.name}')
-            # TODO(zhwu): This is not a clean way to update the SkyPilot config,
-            # because we are resetting the global context for a single DAG,
-            # which is conceptually weird.
-            importlib.reload(skypilot_config)
-
     logger.debug(f'Mutated user request: {mutated_user_request}')
     mutated_dag.policy_applied = True
     return mutated_dag, mutated_config