skypilot-nightly 1.0.0.dev20250919__py3-none-any.whl → 1.0.0.dev20250925__py3-none-any.whl

sky/utils/context.py

@@ -2,15 +2,21 @@
 
 import asyncio
 from collections.abc import Mapping
-from collections.abc import MutableMapping
 import contextvars
+import copy
 import functools
+import inspect
 import os
 import pathlib
 import subprocess
 import sys
-import typing
-from typing import Any, Callable, Dict, Optional, TextIO, TypeVar
+from typing import (Callable, Dict, Iterator, MutableMapping, Optional, TextIO,
+                    TYPE_CHECKING, TypeVar)
+
+from typing_extensions import ParamSpec
+
+if TYPE_CHECKING:
+    from sky.skypilot_config import ConfigContext
 
 
 class Context(object):
@@ -88,7 +94,7 @@ class Context(object):
         else:
             self._log_file_handle = open(log_file, 'a', encoding='utf-8')
         self._log_file = log_file
-        if original_log_file is not None:
+        if original_log_handle is not None:
             original_log_handle.close()
         return original_log_file
 
@@ -102,8 +108,30 @@ class Context(object):
         for k, v in envs.items():
             self.env_overrides[k] = v
 
+    def cleanup(self):
+        """Clean up the context."""
+        if self._log_file_handle is not None:
+            self._log_file_handle.close()
+            self._log_file_handle = None
+
+    def copy(self) -> 'Context':
+        """Create a copy of the context.
+
+        Changes to the current context after this call will not affect the copy.
+        The new context will get its own handle/fd for the log file.
+        The new context will get an independent copy of the env var overrides.
+        The new context will get an independent copy of the config context.
+        Cancellation of the current context will not be propagated to the copy.
+        """
+        new_context = Context()
+        new_context.redirect_log(self._log_file)
+        new_context.env_overrides = self.env_overrides.copy()
+        new_context.config_context = copy.deepcopy(self.config_context)
+        return new_context
 
-_CONTEXT = contextvars.ContextVar('sky_context', default=None)
+
+_CONTEXT = contextvars.ContextVar[Optional[Context]]('sky_context',
+                                                     default=None)
 
 
 def get() -> Optional[Context]:
@@ -116,7 +144,7 @@ def get() -> Optional[Context]:
     return _CONTEXT.get()
 
 
-class ContextualEnviron(MutableMapping):
+class ContextualEnviron(MutableMapping[str, str]):
     """Environment variables wrapper with contextual overrides.
 
     An instance of ContextualEnviron will typically be used to replace
@@ -155,10 +183,10 @@ class ContextualEnviron(MutableMapping):
        assert os.environ['FOO'] == 'BAR1'
     """
 
-    def __init__(self, environ):
+    def __init__(self, environ: 'os._Environ[str]') -> None:
         self._environ = environ
 
-    def __getitem__(self, key):
+    def __getitem__(self, key: str) -> str:
         ctx = get()
         if ctx is not None:
             if key in ctx.env_overrides:
@@ -170,10 +198,10 @@ class ContextualEnviron(MutableMapping):
                 return value
         return self._environ[key]
 
-    def __iter__(self):
-        ctx = get()
-        deleted_keys = set()
-        if ctx is not None:
+    def __iter__(self) -> Iterator[str]:
+
+        def iter_from_context(ctx: Context) -> Iterator[str]:
+            deleted_keys = set()
             for key, value in ctx.env_overrides.items():
                 if value is None:
                     deleted_keys.add(key)
@@ -182,20 +210,24 @@ class ContextualEnviron(MutableMapping):
                 # Deduplicate the keys
                 if key not in ctx.env_overrides and key not in deleted_keys:
                     yield key
+
+        ctx = get()
+        if ctx is not None:
+            return iter_from_context(ctx)
         else:
             return self._environ.__iter__()
 
-    def __len__(self):
+    def __len__(self) -> int:
         return len(dict(self))
 
-    def __setitem__(self, key, value):
+    def __setitem__(self, key: str, value: str) -> None:
         ctx = get()
         if ctx is not None:
             ctx.env_overrides[key] = value
         else:
             self._environ.__setitem__(key, value)
 
-    def __delitem__(self, key):
+    def __delitem__(self, key: str) -> None:
         ctx = get()
         if ctx is not None:
             if key in ctx.env_overrides:
@@ -211,10 +243,13 @@ class ContextualEnviron(MutableMapping):
         else:
             self._environ.__delitem__(key)
 
-    def __repr__(self):
-        return self._environ.__repr__()
+    def __repr__(self) -> str:
+        # Adapted from os._Environ.__repr__
+        formatted_items = ', '.join(
+            f'{key!r}: {value!r}' for key, value in self.items())
+        return f'ctx_environ({{{formatted_items}}})'
 
-    def copy(self):
+    def copy(self) -> Dict[str, str]:
         copied = self._environ.copy()
         ctx = get()
         if ctx is not None:
@@ -225,7 +260,7 @@ class ContextualEnviron(MutableMapping):
                     copied[key] = ctx.env_overrides[key]
         return copied
 
-    def setdefault(self, key, default=None):
+    def setdefault(self, key: str, default: str) -> str:
         return self._environ.setdefault(key, default)
 
     def __ior__(self, other):
@@ -260,27 +295,67 @@ class Popen(subprocess.Popen):
         super().__init__(*args, env=env, **kwargs)
 
 
-F = TypeVar('F', bound=Callable[..., Any])
+P = ParamSpec('P')
+T = TypeVar('T')
 
 
-def contextual(func: F) -> F:
+def contextual(func: Callable[P, T]) -> Callable[P, T]:
     """Decorator to initialize a context before executing the function.
 
-    If a context is already initialized, this decorator will reset the context,
-    i.e. all contextual variables set previously will be cleared.
+    If a context is already initialized, this decorator will create a new
+    context that inherits the values from the existing context.
     """
 
     @functools.wraps(func)
-    def wrapper(*args, **kwargs):
-        initialize()
-        return func(*args, **kwargs)
+    def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
+        original_ctx = get()
+        initialize(original_ctx)
+        ctx = get()
+        cleanup_after_await = False
+
+        def cleanup():
+            try:
+                if ctx is not None:
+                    ctx.cleanup()
+            finally:
+                # Note: _CONTEXT.reset() is not reliable - may fail with
+                # ValueError: <Token ... at ...> was created in a different
+                # Context
+                # We must make sure this happens because otherwise we may try to
+                # write to the wrong log.
+                _CONTEXT.set(original_ctx)
+
+        # There are two cases:
+        # 1. The function is synchronous (that is, return type is not awaitable)
+        #    In this case, we use a finally block to cleanup the context.
+        # 2. The function is asynchronous (that is, return type is awaitable)
+        #    In this case, we need to construct an async def wrapper and await
+        #    the value, then call the cleanup function in the finally block.
+
+        async def await_with_cleanup(awaitable):
+            try:
+                return await awaitable
+            finally:
+                cleanup()
+
+        try:
+            ret = func(*args, **kwargs)
+            if inspect.isawaitable(ret):
+                cleanup_after_await = True
+                return await_with_cleanup(ret)
+            else:
+                return ret
+        finally:
+            if not cleanup_after_await:
+                cleanup()
 
-    return typing.cast(F, wrapper)
+    return wrapper
 
 
-def initialize():
+def initialize(base_context: Optional[Context] = None) -> None:
     """Initialize the current SkyPilot context."""
-    _CONTEXT.set(Context())
+    new_context = base_context.copy() if base_context is not None else Context()
+    _CONTEXT.set(new_context)
 
 
 class _ContextualStream:

sky/utils/controller_utils.py

@@ -620,15 +620,16 @@ def get_controller_resources(
     controller_resources_to_use: resources.Resources = list(
         controller_resources)[0]
 
-    controller_record = global_user_state.get_cluster_from_name(
+    controller_handle = global_user_state.get_handle_from_cluster_name(
         controller.value.cluster_name)
-    if controller_record is not None:
-        handle = controller_record.get('handle', None)
-        if handle is not None:
+    if controller_handle is not None:
+        if controller_handle is not None:
             # Use the existing resources, but override the autostop config with
             # the one currently specified in the config.
-            controller_resources_to_use = handle.launched_resources.copy(
-                autostop=controller_resources_config_copied.get('autostop'))
+            controller_resources_to_use = (
+                controller_handle.launched_resources.copy(
+                    autostop=controller_resources_config_copied.get('autostop'))
+            )
 
     # If the controller and replicas are from the same cloud (and region/zone),
     # it should provide better connectivity. We will let the controller choose

sky/utils/kubernetes/create_cluster.sh

@@ -1,22 +1,19 @@
 #!/bin/bash
 # Creates a local Kubernetes cluster using kind with optional GPU support
-# Usage: ./create_cluster.sh [--gpus]
-# Invokes generate_kind_config.py to generate a kind-cluster.yaml with NodePort mappings
+# Usage: ./create_cluster.sh [name] [yaml_path] [--gpus]
 set -e
 
 # Images
 IMAGE="us-central1-docker.pkg.dev/skypilot-375900/skypilotk8s/skypilot:latest"
 IMAGE_GPU="us-central1-docker.pkg.dev/skypilot-375900/skypilotk8s/skypilot-gpu:latest"
 
-# Limit port range to speed up kind cluster creation
-PORT_RANGE_START=30000
-PORT_RANGE_END=30100
-
-USER_HASH=$1
+# Arguments
+NAME=$1
+YAML_PATH=$2
 
 # Check for GPU flag
 ENABLE_GPUS=false
-if [[ "$2" == "--gpus" ]]; then
+if [[ "$3" == "--gpus" ]]; then
     ENABLE_GPUS=true
 fi
 
@@ -82,28 +79,16 @@ fi
 # ====== End of dependency checks =======
 
 # Check if the local cluster already exists
-if kind get clusters | grep -q skypilot; then
-    echo "Local cluster already exists. Exiting."
+if kind get clusters | grep -q $NAME; then
+    echo "Local cluster $NAME already exists. Exiting."
     # Switch context to the local cluster
-    kind export kubeconfig --name skypilot
-    kubectl config use-context kind-skypilot
+    kind export kubeconfig --name $NAME
+    kubectl config use-context kind-$NAME
     exit 100
 fi
 
-# Generate cluster YAML
-YAML_PATH="/tmp/skypilot-kind-$USER_HASH.yaml"
-echo "Generating $YAML_PATH"
-
-# Add GPUs flag to the generate_kind_config.py command if GPUs are enabled
-if $ENABLE_GPUS; then
-    python -m sky.utils.kubernetes.generate_kind_config --path $YAML_PATH --port-start ${PORT_RANGE_START} --port-end ${PORT_RANGE_END} --gpus
-else
-  python -m sky.utils.kubernetes.generate_kind_config --path $YAML_PATH --port-start ${PORT_RANGE_START} --port-end ${PORT_RANGE_END}
-fi
-
-kind create cluster --config $YAML_PATH --name skypilot
-
-echo "Kind cluster created."
+kind create cluster --config $YAML_PATH --name $NAME
+echo "Kind cluster $NAME created."
 
 # Function to wait for GPU operator to be correctly installed
 wait_for_gpu_operator_installation() {
@@ -157,7 +142,7 @@ if $ENABLE_GPUS; then
     echo "Enabling GPU support..."
     # Run patch for missing ldconfig.real
     # https://github.com/NVIDIA/nvidia-docker/issues/614#issuecomment-423991632
-    docker exec -ti skypilot-control-plane /bin/bash -c '[ ! -f /sbin/ldconfig.real ] && ln -s /sbin/ldconfig /sbin/ldconfig.real || echo "/sbin/ldconfig.real already exists"'
+    docker exec -ti $NAME-control-plane /bin/bash -c '[ ! -f /sbin/ldconfig.real ] && ln -s /sbin/ldconfig /sbin/ldconfig.real || echo "/sbin/ldconfig.real already exists"'
 
     echo "Installing NVIDIA GPU operator..."
     # Install the NVIDIA GPU operator
@@ -185,4 +170,4 @@ if $ENABLE_GPUS; then
         echo "GPU support is enabled. Run 'sky show-gpus --cloud kubernetes' to see the GPUs available on the cluster."
     fi
 fi
-echo "Number of CPUs available on the local cluster: $NUM_CPUS"
+echo "Number of CPUs available on the local cluster $NAME: $NUM_CPUS"

sky/utils/kubernetes/delete_cluster.sh

@@ -1,9 +1,12 @@
 #!/bin/bash
-# Deletes the local kind cluster
-# Usage: ./delete_cluster.sh
-# Raises error code 100 if the local cluster does not exist
+# Deletes the local kind cluster of [name]
+# Usage: ./delete_cluster.sh [name]
+# Raises error code 100 if the specified local cluster does not exist
 
 set -e
+
+NAME="${1:-skypilot}"
+
 # Check if docker is running
 if ! docker info > /dev/null 2>&1; then
     >&2 echo "Docker is not running. Please start Docker and try again."
@@ -17,13 +20,13 @@ if ! kind version > /dev/null 2>&1; then
 fi
 
 # Check if the local cluster exists
-if ! kind get clusters | grep -q skypilot; then
-    echo "Local cluster does not exist. Exiting."
+if ! kind get clusters | grep -q $NAME; then
+    echo "Local cluster $NAME does not exist. Exiting."
     exit 100
 fi
 
-kind delete cluster --name skypilot
-echo "Local cluster deleted!"
+kind delete cluster --name $NAME
+echo "Local cluster $NAME deleted!"
 
 # Switch to the first available context
 AVAILABLE_CONTEXT=$(kubectl config get-contexts -o name | head -n 1)

sky/utils/kubernetes/generate_kind_config.py

@@ -3,67 +3,8 @@
 Maps specified ports from host to cluster container.
 """
 import argparse
-import textwrap
-
-
-def generate_kind_config(path: str,
-                         port_start: int = 30000,
-                         port_end: int = 32768,
-                         num_nodes: int = 1,
-                         gpus: bool = False) -> None:
-    """Generate a kind cluster config with ports mapped from host to container
-
-    Args:
-        path: Path to generate the config file at
-        port_start: Port range start
-        port_end: Port range end
-        num_nodes: Number of nodes in the cluster
-        gpus: If true, initialize kind cluster with GPU support
-    """
-
-    preamble = textwrap.dedent(f"""
-    apiVersion: kind.x-k8s.io/v1alpha4
-    kind: Cluster
-    kubeadmConfigPatches:
-    - |
-      kind: ClusterConfiguration
-      apiServer:
-        extraArgs:
-          "service-node-port-range": {port_start}-{port_end}
-    nodes:
-    - role: control-plane
-      kubeadmConfigPatches:
-      - |
-        kind: InitConfiguration
-        nodeRegistration:
-          kubeletExtraArgs:
-            node-labels: "ingress-ready=true"
-    """)
-    if gpus:
-        preamble += textwrap.indent(
-            textwrap.dedent("""
-        extraMounts:
-        - hostPath: /dev/null
-          containerPath: /var/run/nvidia-container-devices/all"""), ' ' * 2)
-    preamble += textwrap.indent(
-        textwrap.dedent("""
-      extraPortMappings:"""), ' ' * 2)
-    suffix = ''
-    if num_nodes > 1:
-        for _ in range(1, num_nodes):
-            suffix += """- role: worker\n"""
-    with open(path, 'w', encoding='utf-8') as f:
-        f.write(preamble)
-        for port in range(port_start, port_end + 1):
-            f.write(f"""
-      - containerPort: {port}
-        hostPort: {port}
-        listenAddress: "0.0.0.0"
-        protocol: tcp""")
-        f.write('\n')
-        if suffix:
-            f.write(suffix)
 
+from sky.utils.kubernetes import kubernetes_deploy_utils
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser(description='Generate a kind cluster '
@@ -77,10 +18,6 @@ if __name__ == '__main__':
                         type=int,
                         default=30000,
                         help='Port range start')
-    parser.add_argument('--port-end',
-                        type=int,
-                        default=32768,
-                        help='Port range end')
     parser.add_argument('--num-nodes',
                         type=int,
                         default=1,
@@ -90,5 +27,8 @@ if __name__ == '__main__':
                         action='store_true',
                         help='Initialize kind cluster with GPU support')
     args = parser.parse_args()
-    generate_kind_config(args.path, args.port_start, args.port_end,
-                         args.num_nodes, args.gpus)
+
+    with open(args.path, 'w', encoding='utf-8') as f:
+        f.write(
+            kubernetes_deploy_utils.generate_kind_config(
+                args.port_start, args.num_nodes, args.gpus))