Repository: skypilot-org/skypilot
Pull request #6125 (2025-07-01 13:35:27)

[Core] Support GPU provisioning with flex-start in GKE
<!-- Describe the changes in this PR -->

Support GPU provisioning with flex-start in GKE

<!-- Describe the tests ran -->
<!-- Unit tests (tests/test_*.py) are part of GitHub CI; below are tests that launch on the cloud. -->

Tested (run the relevant ones):

- [x] Code formatting: install pre-commit (auto-check on commit) or `bash format.sh`
- [x] Any manual or new tests for this PR (please specify below)
  - Launch cluster with `accelerators: L4:1` in the flex-start case
  - Launch cluster with `accelerators: L4:1` in the flex-start case with config in task.yaml
  - Launch cluster with `accelerators: L4:1` in the flex-start with queued provisioning case with one Kueue resource flavor
  - Launch cluster with `accelerators: L4:1` in the flex-start with queued provisioning case with one Kueue resource flavor with config in task.yaml
  - Launch cluster with `accelerators: L4:1` in the flex-start with queued provisioning case with two Kueue resource flavors
- [ ] All smoke tests: `/smoke-test` (CI) or `pytest tests/test_smoke.py` (local)
- [ ] Relevant individual tests: `/smoke-test -k test_name` (CI) or `pytest tests/test_smoke.py::test_name` (local)
- [ ] Backward compatibility: `/quicktest-core` (CI) or `pytest tests/smoke_tests/test_backward_compat.py` (local)

<!-- CI commands (/-prefixed) can only be triggered by repo members -->


DanielZhangQD: Hi @Michaelvll, comments addressed. PTAL again. Thanks!
DanielZhangQD: /smoke-test --kubernetes

Patch:
diff --git a/.github/workflows/nightly-build.yml b/.github/workflows/nightly-build.yml
index 656f246ea..fbf13f674 100644
--- a/.github/workflows/nightly-build.yml
+++ b/.github/workflows/nightly-build.yml
@@ -1,7 +1,7 @@
 name: nightly-build
 on:
   schedule:
-    - cron: '35 10 * * *' # 10:35am UTC, 2:35am PST, 5:35am EST
+    - cron: '35 8 * * *' # 8:35am UTC, 1:35am PST, 4:35am EST
   workflow_dispatch:
     inputs:
       skip_buildkite:
diff --git a/README.md b/README.md
index f7e506f1d..012b0318b 100644
--- a/README.md
+++ b/README.md
@@ -117,7 +117,8 @@ workdir: ~/torch_examples
 # Commands to be run before executing the job.
 # Typical use: pip install -r requirements.txt, git clone, etc.
 setup: |
-  pip install "torch<2.2" torchvision --index-url https://download.pytorch.org/whl/cu121
+  cd mnist
+  pip install -r requirements.txt
 
 # Commands to run as a job.
 # Typical use: launch the main program.
diff --git a/docs/source/examples/frameworks/index.rst b/docs/source/examples/frameworks/index.rst
index e39201806..43c385782 100644
--- a/docs/source/examples/frameworks/index.rst
+++ b/docs/source/examples/frameworks/index.rst
@@ -6,7 +6,7 @@ Frameworks
 
    Cross-cloud data transfer <https://nebius.com/blog/posts/bulk-object-storage-s3-data-migration-with-skypilot>
    DVC <dvc>
-   GCP DWS/Kueue <https://gke-ai-labs.dev/docs/tutorials/skypilot/resource-management-using-kueue/>
+   GCP DWS <https://docs.skypilot.co/en/latest/reservations/reservations.html#gcp-dynamic-workload-scheduler-dws>
    Jupyter <jupyter>
    MLFlow <https://nebius.com/blog/posts/orchestrating-llm-fine-tuning-k8s-skypilot-mlflow>
    MPI <mpi>
diff --git a/docs/source/reference/config-sources.rst b/docs/source/reference/config-sources.rst
index e03a6bed1..0f657dc8e 100644
--- a/docs/source/reference/config-sources.rst
+++ b/docs/source/reference/config-sources.rst
@@ -109,6 +109,8 @@ The following fields are supported in SkyPilot YAML inline configuration:
 * :ref:`nvidia_gpus.disable_ecc <config-yaml-nvidia-gpus-disable-ecc>`
 * :ref:`kubernetes.pod_config <config-yaml-kubernetes-pod-config>`
 * :ref:`kubernetes.provision_timeout <config-yaml-kubernetes-provision-timeout>`
+* :ref:`kubernetes.dws <config-yaml-kubernetes-dws>`
+* :ref:`kubernetes.kueue <config-yaml-kubernetes-kueue>`
 * :ref:`gcp.managed_instance_group <config-yaml-gcp-managed-instance-group>`
 
 Example:
@@ -122,6 +124,8 @@ Example:
     kubernetes:
       pod_config: ...
       provision_timeout: ...
+      dws: ...
+      kueue: ...
     gcp:
       managed_instance_group: ...
     nvidia_gpus:
diff --git a/docs/source/reference/config.rst b/docs/source/reference/config.rst
index 1bf88a6b7..2d1c811fe 100644
--- a/docs/source/reference/config.rst
+++ b/docs/source/reference/config.rst
@@ -79,6 +79,9 @@ Below is the configuration syntax and some example values. See detailed explanat
         runtimeClassName: nvidia
     :ref:`kueue <config-yaml-kubernetes-kueue>`:
       :ref:`local_queue_name <config-yaml-kubernetes-kueue-local-queue-name>`: skypilot-local-queue
+    :ref:`dws <config-yaml-kubernetes-dws>`:
+      enabled: true
+      max_run_duration: 10m
     :ref:`context_configs <config-yaml-kubernetes-context-configs>`:
       context1:
         pod_config:
@@ -1111,6 +1114,39 @@ Kueue configuration (optional).
 
 Name of the `local queue <https://kueue.sigs.k8s.io/docs/concepts/local_queue/>`_ to use for SkyPilot jobs.
 
+.. _config-yaml-kubernetes-dws:
+
+``kubernetes.dws``
+~~~~~~~~~~~~~~~~~~
+
+GKE DWS configuration (optional).
+
+Refer to :ref:`Using DWS on GKE <dws-on-gke>` for more details.
+
+``kubernetes.dws.enabled``
+~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Whether to enable DWS (optional).
+
+When ``enabled: true``, SkyPilot will automatically use DWS with flex-start mode. If ``kubernetes.kueue.local_queue_name`` is set, it will use flex-start with queued provisioning mode.
+
+Default: ``false``.
+
+``kubernetes.dws.max_run_duration``
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Maximum runtime of a node (optional), only used in ``flex-start-queued-provisioning`` mode.
+
+Default: ``null``.
+
+Example:
+
+.. code-block:: yaml
+
+  kubernetes:
+    dws:
+      enabled: true
+      max_run_duration: 10m
 
 .. _config-yaml-kubernetes-context-configs:
 
diff --git a/docs/source/reservations/reservations.rst b/docs/source/reservations/reservations.rst
index 9dbf37a02..b160a88be 100644
--- a/docs/source/reservations/reservations.rst
+++ b/docs/source/reservations/reservations.rst
@@ -123,6 +123,7 @@ Example:
 
 SkyPilot will utilize the reservations similar to AWS reservations as described in :ref:`utilizing-reservations`.
 
+.. _gcp-dws:
 
 GCP Dynamic Workload Scheduler (DWS)
 -------------------------------------
@@ -170,32 +171,141 @@ In case you want to specify the DWS configuration for each job/cluster, you can
 
 .. _dws-on-gke:
 
-Using DWS on GKE with Kueue
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Using DWS on GKE
+~~~~~~~~~~~~~~~~~
+
+.. dropdown:: DWS flex-start modes on GKE
+
+  Flex-start, powered by `Dynamic Workload Scheduler <https://cloud.google.com/blog/products/compute/introducing-dynamic-workload-scheduler>`_, provides a flexible and cost-effective technique to obtain GPUs when you need to run AI/ML workloads.
+
+  GKE supports the two kinds of flex-start configurations:
+
+  * Flex-start, where GKE allocates resources node by node.
+
+    * This configuration only requires you to set the ``flex-start`` flag during node creation. 
+
+    * Flex-start is recommended for small to medium (single-node) workloads such as small training jobs, offline inference, or batch jobs.
+
+  * Flex-start with queued provisioning, where GKE allocates all requested resources at the same time.
+
+    * To use this configuration, add the ``--flex-start`` and ``--enable-queued-provisioning`` flags when you create the node pool.
+
+    * Flex-start with queued provisioning is recommended for multi-node workloads, which require multiple resources and can't start running until all nodes are provisioned and ready at the same time. For example, distributed machine learning training workloads.
+
+  See `GKE DWS documentation <https://cloud.google.com/kubernetes-engine/docs/concepts/dws>`_ for more details.
+
+  .. note::
+
+      If you are using Kueue and ``kubernetes.kueue.local_queue_name`` is specified, SkyPilot will automatically enable flex-start with queued provisioning mode, otherwise it will use the regular flex-start mode.
+
+
+Using DWS without Kueue
+^^^^^^^^^^^^^^^^^^^^^^^
+
+To launch clusters or managed jobs using DWS without `Kueue <https://kueue.sigs.k8s.io/>`_:
 
-DWS is also supported on Google Kubernetes Engine (GKE) with Kueue. To enable DWS on GKE, you need to set up your GKE cluster with Kueue and DWS; see the `GCP official docs <https://cloud.google.com/kubernetes-engine/docs/how-to/provisioningrequest>`__.
+1. Follow the `official documentation <https://cloud.google.com/kubernetes-engine/docs/how-to/dws-flex-start-training#node-pool-flex>`_ to create a node pool.
 
-To launch a SkyPilot cluster or job on GKE with DWS, you can specify the DWS configuration in the SkyPilot task YAML:
+2. Configure the following fields in ``~/.sky/config.yaml``:
 
 .. code-block:: yaml
 
+    kubernetes:
+      # provision_timeout: 1200
+      autoscaler: gke
+      dws:
+        enabled: true
+
+When DWS is enabled, the default ``provision_timeout`` is set to ``600`` seconds (10 minutes). If you encounter provisioning timeout issues, you can increase this value in your configuration. For example, set it to ``1200`` seconds (20 minutes) or higher depending on your workload and cluster size.
+
+3. Launch your clusters or managed jobs.
+
+.. code-block:: yaml
+
+    name: dws
+
+    resources:
+      infra: k8s
+      accelerators: L4:1
+
+    num_nodes: 1
+
+In case you want to specify the DWS configuration for each job/cluster, you can also specify the configuration in the SkyPilot task YAML (see :ref:`here <config-client-job-task-yaml>`):
+
+.. code-block:: yaml
+
+    name: dws
+
+    resources:
+      infra: k8s
+      accelerators: L4:1
+
+    num_nodes: 1
+
     config:
       kubernetes:
-        pod_config:
-          metadata:
-            annotations:
-              provreq.kueue.x-k8s.io/maxRunDurationSeconds: "3600"
-        provision_timeout: 900
+        dws:
+          enabled: true
+
+Using DWS with Kueue
+^^^^^^^^^^^^^^^^^^^^
+
+To launch clusters or managed jobs with DWS with Kueue:
+
+1. Follow the `GKE documentation <https://cloud.google.com/kubernetes-engine/docs/how-to/provisioningrequest#create-node-pool>`_ to create a node pool.
+
+2. Follow the :ref:`Kueue example <kubernetes-example-kueue>` to install Kueue, update configuration of Kueue to support plain Pods, create Kueue resource flavor, cluster queue and local queue.
+
+3. Configure the following fields in ``~/.sky/config.yaml``:
+
+.. code-block:: yaml
+
+    kubernetes:
+      # provision_timeout: 1200
+      autoscaler: gke
+      dws:
+        enabled: true
+        # Optional, the maximum runtime of a node,
+        # up to the default of seven days
+        max_run_duration: 10m
+      kueue:
+        local_queue_name: skypilot-local-queue
+
+When DWS is enabled, the default ``provision_timeout`` is set to ``600`` seconds (10 minutes). If you encounter provisioning timeout issues, you can increase this value in your configuration. For example, set it to ``1200`` seconds (20 minutes) or higher depending on your workload and cluster size.
+
+4. Launch your clusters or managed jobs.
+
+.. code-block:: yaml
+
+    name: dws-kueue
 
     resources:
-      infra: kubernetes
-      accelerators: A100:8
-      labels:
-        kueue.x-k8s.io/queue-name: dws-local-queue
+      infra: k8s
+      accelerators: L4:1
+
+    num_nodes: 2
+
+In case you want to specify the DWS configuration for each job/cluster, you can also specify the configuration in the SkyPilot task YAML (see :ref:`here <config-client-job-task-yaml>`):
+
+.. code-block:: yaml
 
-1. ``kueue.x-k8s.io/queue-name``: name of the Kueue queue to submit your resource request to.
-2. ``provreq.kueue.x-k8s.io/maxRunDurationSeconds``: maximum duration for a created instance to be kept alive (in seconds, required).
-3. ``provision_timeout``: timeout for provisioning an instance with DWS (in seconds, optional). If the timeout is reached without getting the requested resources, SkyPilot will automatically :ref:`failover <auto-failover>` to other clouds/regions/zones to get the resources.
+    name: dws-kueue
+
+    resources:
+      infra: k8s
+      accelerators: L4:1
+
+    num_nodes: 2
+
+    config:
+      kubernetes:
+        dws:
+          enabled: true
+          # Optional, the maximum runtime of a node,
+          # up to the default of seven days
+          max_run_duration: 10m
+        kueue:
+          local_queue_name: skypilot-local-queue
 
 Long-term reservations
 ----------------------
diff --git a/sky/backends/backend_utils.py b/sky/backends/backend_utils.py
index 9e2d584c7..5987efcd0 100644
--- a/sky/backends/backend_utils.py
+++ b/sky/backends/backend_utils.py
@@ -51,6 +51,7 @@ from sky.utils import rich_utils
 from sky.utils import schemas
 from sky.utils import status_lib
 from sky.utils import subprocess_utils
+from sky.utils import tempstore
 from sky.utils import timeline
 from sky.utils import ux_utils
 from sky.workspaces import core as workspaces_core
@@ -247,7 +248,7 @@ def _optimize_file_mounts(tmp_yaml_path: str) -> None:
     #  - use a remote command to move all runtime files to their right places.
 
     # Local tmp dir holding runtime files.
-    local_runtime_files_dir = tempfile.mkdtemp()
+    local_runtime_files_dir = tempstore.mkdtemp()
     new_file_mounts = {_REMOTE_RUNTIME_FILES_DIR: local_runtime_files_dir}
 
     # Generate local_src -> unique_name.
diff --git a/sky/catalog/cudo_catalog.py b/sky/catalog/cudo_catalog.py
index adccccd0e..66cd8b73a 100644
--- a/sky/catalog/cudo_catalog.py
+++ b/sky/catalog/cudo_catalog.py
@@ -10,7 +10,7 @@ from sky.utils import ux_utils
 if typing.TYPE_CHECKING:
     from sky.clouds import cloud
 
-_PULL_FREQUENCY_HOURS = 1
+_PULL_FREQUENCY_HOURS = 7
 _df = common.read_catalog(cudo_mt.VMS_CSV,
                           pull_frequency_hours=_PULL_FREQUENCY_HOURS)
 
diff --git a/sky/clouds/kubernetes.py b/sky/clouds/kubernetes.py
index d3575e2aa..f280136e1 100644
--- a/sky/clouds/kubernetes.py
+++ b/sky/clouds/kubernetes.py
@@ -14,6 +14,7 @@ from sky import exceptions
 from sky import sky_logging
 from sky import skypilot_config
 from sky.adaptors import kubernetes
+from sky.clouds.utils import gcp_utils
 from sky.provision import instance_setup
 from sky.provision.kubernetes import network_utils
 from sky.provision.kubernetes import utils as kubernetes_utils
@@ -22,6 +23,7 @@ from sky.provision.kubernetes.utils import normalize_tpu_accelerator_name
 from sky.skylet import constants
 from sky.utils import annotations
 from sky.utils import common_utils
+from sky.utils import kubernetes_enums
 from sky.utils import registry
 from sky.utils import resources_utils
 from sky.utils import schemas
@@ -281,7 +283,8 @@ class Kubernetes(clouds.Cloud):
                 default_value=None)
             if (autoscaler_type is not None and
                     not kubernetes_utils.get_autoscaler(
-                        autoscaler_type).can_query_backend):
+                        kubernetes_enums.KubernetesAutoscalerType(
+                            autoscaler_type)).can_query_backend):
                 # Unsupported autoscaler type. Rely on the autoscaler to
                 # provision the right instance type without running checks.
                 # Worst case, if autoscaling fails, the pod will be stuck in
@@ -298,7 +301,8 @@ class Kubernetes(clouds.Cloud):
 
             if autoscaler_type is None:
                 continue
-            autoscaler = kubernetes_utils.get_autoscaler(autoscaler_type)
+            autoscaler = kubernetes_utils.get_autoscaler(
+                kubernetes_enums.KubernetesAutoscalerType(autoscaler_type))
             logger.debug(f'{context} has autoscaler of type: {autoscaler_type}')
             if autoscaler.can_create_new_instance_of_type(
                     context, instance_type):
@@ -406,8 +410,10 @@ class Kubernetes(clouds.Cloud):
 
     @staticmethod
     def _calculate_provision_timeout(
-            num_nodes: int,
-            volume_mounts: Optional[List['volume_lib.VolumeMount']]) -> int:
+        num_nodes: int,
+        volume_mounts: Optional[List['volume_lib.VolumeMount']],
+        enable_flex_start: bool,
+    ) -> int:
         """Calculate provision timeout based on number of nodes.
 
         The timeout scales linearly with the number of nodes to account for
@@ -415,6 +421,8 @@ class Kubernetes(clouds.Cloud):
 
         Args:
             num_nodes: Number of nodes being provisioned
+            volume_mounts: Volume mounts for the pod
+            enable_flex_start: Whether flex start is enabled
 
         Returns:
             Timeout in seconds
@@ -422,7 +430,12 @@ class Kubernetes(clouds.Cloud):
         base_timeout = 10  # Base timeout for single node
         per_node_timeout = 0.2  # Additional seconds per node
         max_timeout = 60  # Cap at 1 minute
-        if volume_mounts is not None:
+        if enable_flex_start:
+            # Flex start takes longer to provision.
+            base_timeout = 600
+            per_node_timeout = 10
+            max_timeout = 900
+        elif volume_mounts is not None:
             for volume_mount in volume_mounts:
                 if (volume_mount.volume_config.type ==
                         volume_lib.VolumeType.PVC.value):
@@ -581,24 +594,6 @@ class Kubernetes(clouds.Cloud):
         if resources.use_spot:
             spot_label_key, spot_label_value = kubernetes_utils.get_spot_label()
 
-        # Timeout for resource provisioning. This timeout determines how long to
-        # wait for pod to be in pending status before giving up.
-        # Larger timeout may be required for autoscaling clusters, since
-        # autoscaler may take some time to provision new nodes.
-        # Note that this timeout includes time taken by the Kubernetes scheduler
-        # itself, which can be upto 2-3 seconds, and up to 10-15 seconds when
-        # scheduling 100s of pods.
-        # We use a linear scaling formula to determine the timeout based on the
-        # number of nodes.
-
-        timeout = self._calculate_provision_timeout(num_nodes, volume_mounts)
-        timeout = skypilot_config.get_effective_region_config(
-            cloud='kubernetes',
-            region=context,
-            keys=('provision_timeout',),
-            default_value=timeout,
-            override_configs=resources.cluster_config_overrides)
-
         # Check if this cluster supports high performance networking and
         # configure IPC_LOCK capability for clusters like Nebius that support it
         k8s_ipc_lock_capability = False
@@ -640,7 +635,46 @@ class Kubernetes(clouds.Cloud):
                 cloud='kubernetes',
                 region=context,
                 keys=('kueue', 'local_queue_name'),
-                default_value=None))
+                default_value=None,
+                override_configs=resources.cluster_config_overrides))
+
+        # Check DWS configuration for GKE.
+        (enable_flex_start, enable_flex_start_queued_provisioning,
+         max_run_duration_seconds) = gcp_utils.get_dws_config(
+             context, k8s_kueue_local_queue_name,
+             resources.cluster_config_overrides)
+        if enable_flex_start_queued_provisioning or enable_flex_start:
+            # DWS is only supported in GKE, check the autoscaler type.
+            autoscaler_type = skypilot_config.get_effective_region_config(
+                cloud='kubernetes',
+                region=context,
+                keys=('autoscaler',),
+                default_value=None)
+            if (autoscaler_type !=
+                    kubernetes_enums.KubernetesAutoscalerType.GKE.value):
+                raise ValueError(
+                    f'DWS is only supported in GKE, but the autoscaler type '
+                    f'for context {context} is {autoscaler_type}')
+
+        # Timeout for resource provisioning. This timeout determines how long to
+        # wait for pod to be in pending status before giving up.
+        # Larger timeout may be required for autoscaling clusters, since
+        # autoscaler may take some time to provision new nodes.
+        # Note that this timeout includes time taken by the Kubernetes scheduler
+        # itself, which can be upto 2-3 seconds, and up to 10-15 seconds when
+        # scheduling 100s of pods.
+        # We use a linear scaling formula to determine the timeout based on the
+        # number of nodes.
+
+        timeout = self._calculate_provision_timeout(
+            num_nodes, volume_mounts, enable_flex_start or
+            enable_flex_start_queued_provisioning)
+        timeout = skypilot_config.get_effective_region_config(
+            cloud='kubernetes',
+            region=context,
+            keys=('provision_timeout',),
+            default_value=timeout,
+            override_configs=resources.cluster_config_overrides)
 
         deploy_vars = {
             'instance_type': resources.instance_type,
@@ -696,6 +730,8 @@ class Kubernetes(clouds.Cloud):
                 (k8s_ha_storage_class_name),
             'avoid_label_keys': avoid_label_keys,
             'k8s_ipc_lock_capability': k8s_ipc_lock_capability,
+            'k8s_enable_flex_start': enable_flex_start,
+            'k8s_max_run_duration_seconds': max_run_duration_seconds,
         }
 
         # Add kubecontext if it is set. It may be None if SkyPilot is running
diff --git a/sky/clouds/utils/gcp_utils.py b/sky/clouds/utils/gcp_utils.py
index 417589d10..4eadc2b17 100644
--- a/sky/clouds/utils/gcp_utils.py
+++ b/sky/clouds/utils/gcp_utils.py
@@ -10,7 +10,7 @@ import dataclasses
 import json
 import time
 import typing
-from typing import List, Optional, Set
+from typing import Any, Dict, List, Optional, Set, Tuple
 
 import cachetools
 
@@ -18,6 +18,7 @@ from sky import sky_logging
 from sky import skypilot_config
 from sky.provision.gcp import constants
 from sky.provision.kubernetes import utils as kubernetes_utils
+from sky.utils import resources_utils
 from sky.utils import subprocess_utils
 
 if typing.TYPE_CHECKING:
@@ -206,3 +207,62 @@ def get_minimal_storage_permissions() -> List[str]:
     permissions += constants.GCP_MINIMAL_PERMISSIONS
 
     return permissions
+
+
+# Get the DWS configuration for the given context in GKE.
+def get_dws_config(
+    context: str,
+    k8s_kueue_local_queue_name: Optional[str],
+    cluster_config_overrides: Optional[Dict[str, Any]] = None,
+) -> Tuple[bool, bool, Optional[int]]:
+    """Get the DWS configuration for the given context.
+
+        Args:
+            context: The context to get the DWS configuration for.
+            k8s_kueue_local_queue_name: The name of the Kueue local queue.
+            cluster_config_overrides: The cluster config overrides.
+
+        Returns:
+            A tuple of (enable_flex_start,
+                        enable_flex_start_queued_provisioning,
+                        max_run_duration_seconds).
+
+        Raises:
+            ValueError: If k8s_kueue_local_queue_name is missing to enable
+                        flex start queued provisioning for the given context.
+        """
+    dws_config = skypilot_config.get_effective_region_config(
+        cloud='kubernetes',
+        region=context,
+        keys=('dws',),
+        default_value={},
+        override_configs=cluster_config_overrides)
+    if not dws_config:
+        return False, False, None
+
+    enabled = dws_config.get('enabled', False)
+    if not enabled:
+        return False, False, None
+
+    enable_flex_start = False
+    enable_flex_start_queued_provisioning = False
+    max_run_duration_seconds = None
+    # If users already use Kueue, use the flex start with queued
+    # provisioning mode.
+    if k8s_kueue_local_queue_name:
+        enable_flex_start_queued_provisioning = True
+    else:
+        enable_flex_start = True
+
+    if not enable_flex_start_queued_provisioning:
+        return (enable_flex_start, enable_flex_start_queued_provisioning,
+                max_run_duration_seconds)
+
+    # Max run duration is only used in the flex start with queued
+    # provisioning mode.
+    max_run_duration = dws_config.get('max_run_duration', None)
+    if max_run_duration:
+        max_run_duration_seconds = resources_utils.parse_time_minutes(
+            max_run_duration) * 60
+    return (enable_flex_start, enable_flex_start_queued_provisioning,
+            max_run_duration_seconds)
diff --git a/sky/provision/kubernetes/instance.py b/sky/provision/kubernetes/instance.py
index 58a3a8d82..7ddd3e7a9 100644
--- a/sky/provision/kubernetes/instance.py
+++ b/sky/provision/kubernetes/instance.py
@@ -876,6 +876,22 @@ def _create_pods(region: str, cluster_name_on_cloud: str,
             pod_spec_copy['spec']['tolerations'] = existing_tolerations + [
                 tpu_toleration
             ]
+        # Add GPU toleration if GPU is requested.
+        # The nodes provisioned by DWS with flex start with queued provisioning
+        # mode have the GPU taint, so we have to add the GPU toleration.
+        # No need to check if DWS is enabled here since this has no side effect
+        # to the non-DWS case.
+        if needs_gpus:
+            gpu_toleration = {
+                'key': kubernetes_utils.get_gpu_resource_key(),
+                'operator': 'Exists',
+                'effect': 'NoSchedule'
+            }
+            # Preserve existing tolerations if any
+            existing_tolerations = pod_spec_copy['spec'].get('tolerations', [])
+            pod_spec_copy['spec']['tolerations'] = existing_tolerations + [
+                gpu_toleration
+            ]
 
         if to_create_deployment:
             volume.create_persistent_volume_claim(namespace, context, pvc_spec)
diff --git a/sky/provision/kubernetes/utils.py b/sky/provision/kubernetes/utils.py
index 9abfda4be..346fc4a66 100644
--- a/sky/provision/kubernetes/utils.py
+++ b/sky/provision/kubernetes/utils.py
@@ -451,7 +451,10 @@ class GKELabelFormatter(GPULabelFormatter):
 
         e.g. tpu-v5-lite-podslice:8 -> '2x4'
         """
-        acc_type, acc_count = normalize_tpu_accelerator_name(acc_type)
+        # If the TPU type is in the GKE_TPU_ACCELERATOR_TO_GENERATION, it means
+        # that it has been normalized before, no need to normalize again.
+        if acc_type not in GKE_TPU_ACCELERATOR_TO_GENERATION:
+            acc_type, acc_count = normalize_tpu_accelerator_name(acc_type)
         count_to_topology = cls.GKE_TPU_TOPOLOGIES.get(acc_type,
                                                        {}).get(acc_count, None)
         if count_to_topology is None:
@@ -1203,7 +1206,8 @@ def get_accelerator_label_key_values(
             # early since we assume the cluster autoscaler will handle GPU
             # node provisioning.
             return None, None, None, None
-        autoscaler = AUTOSCALER_TYPE_TO_AUTOSCALER.get(autoscaler_type)
+        autoscaler = AUTOSCALER_TYPE_TO_AUTOSCALER.get(
+            kubernetes_enums.KubernetesAutoscalerType(autoscaler_type))
         assert autoscaler is not None, ('Unsupported autoscaler type:'
                                         f' {autoscaler_type}')
         formatter = autoscaler.label_formatter
diff --git a/sky/resources.py b/sky/resources.py
index 3041125bd..69f3f6280 100644
--- a/sky/resources.py
+++ b/sky/resources.py
@@ -1,7 +1,6 @@
 """Resources: compute requirements of Tasks."""
 import collections
 import dataclasses
-import math
 import re
 import textwrap
 import typing
@@ -92,7 +91,7 @@ class AutostopConfig:
             return cls(idle_minutes=config, down=False, enabled=True)
 
         if isinstance(config, str):
-            return cls(idle_minutes=parse_time_minutes(config),
+            return cls(idle_minutes=resources_utils.parse_time_minutes(config),
                        down=False,
                        enabled=True)
 
@@ -2415,31 +2414,3 @@ def _maybe_add_docker_prefix_to_image_id(
     for k, v in image_id_dict.items():
         if not v.startswith('docker:'):
             image_id_dict[k] = f'docker:{v}'
-
-
-def parse_time_minutes(time: str) -> int:
-    """Convert a time string to minutes.
-
-    Args:
-        time: Time string with optional unit suffix (e.g., '30m', '2h', '1d')
-
-    Returns:
-        Time in minutes as an integer
-    """
-    time_str = str(time)
-
-    if time_str.isdecimal():
-        # We assume it is already in minutes to maintain backwards
-        # compatibility
-        return int(time_str)
-
-    time_str = time_str.lower()
-    for unit, multiplier in constants.TIME_UNITS.items():
-        if time_str.endswith(unit):
-            try:
-                value = int(time_str[:-len(unit)])
-                return math.ceil(value * multiplier)
-            except ValueError:
-                continue
-
-    raise ValueError(f'Invalid time format: {time}')
diff --git a/sky/server/requests/executor.py b/sky/server/requests/executor.py
index 5bb40a342..1066dcc75 100644
--- a/sky/server/requests/executor.py
+++ b/sky/server/requests/executor.py
@@ -53,6 +53,7 @@ from sky.utils import common_utils
 from sky.utils import context
 from sky.utils import context_utils
 from sky.utils import subprocess_utils
+from sky.utils import tempstore
 from sky.utils import timeline
 from sky.workspaces import core as workspaces_core
 
@@ -338,6 +339,7 @@ def _request_execution_wrapper(request_id: str,
     2. Update the request status based on the execution result;
     3. Redirect the stdout and stderr of the execution to log file;
     4. Handle the SIGTERM signal to abort the request gracefully.
+    5. Maintain the lifecycle of the temp dir used by the request.
     """
     # Handle the SIGTERM signal to abort the request processing gracefully.
     signal.signal(signal.SIGTERM, _sigterm_handler)
@@ -361,7 +363,8 @@ def _request_execution_wrapper(request_id: str,
         # config, as there can be some logs during override that needs to be
         # captured in the log file.
         try:
-            with override_request_env_and_config(request_body):
+            with override_request_env_and_config(request_body), \
+                tempstore.tempdir():
                 if sky_logging.logging_enabled(logger, sky_logging.DEBUG):
                     config = skypilot_config.to_dict()
                     logger.debug(f'request config: \n'
diff --git a/sky/server/uvicorn.py b/sky/server/uvicorn.py
index 6415aa967..2161fa69f 100644
--- a/sky/server/uvicorn.py
+++ b/sky/server/uvicorn.py
@@ -159,6 +159,8 @@ class Server(uvicorn.Server):
     def run(self, *args, **kwargs):
         """Run the server process."""
         context_utils.hijack_sys_attrs()
+        # Use default loop policy of uvicorn (use uvloop if available).
+        self.config.setup_event_loop()
         with self.capture_signals():
             asyncio.run(self.serve(*args, **kwargs))
 
diff --git a/sky/skylet/constants.py b/sky/skylet/constants.py
index ef02b77f8..a7f0e473d 100644
--- a/sky/skylet/constants.py
+++ b/sky/skylet/constants.py
@@ -375,6 +375,8 @@ OVERRIDEABLE_CONFIG_KEYS_IN_TASK: List[Tuple[str, ...]] = [
     ('ssh', 'pod_config'),
     ('kubernetes', 'pod_config'),
     ('kubernetes', 'provision_timeout'),
+    ('kubernetes', 'dws'),
+    ('kubernetes', 'kueue'),
     ('gcp', 'managed_instance_group'),
     ('gcp', 'enable_gvnic'),
     ('gcp', 'enable_gpu_direct'),
@@ -460,8 +462,10 @@ TIME_UNITS = {
     'w': 7 * 24 * 60,
 }
 
-TIME_PATTERN: str = (
-    f'^[0-9]+({"|".join([unit.lower() for unit in TIME_UNITS])})?$/i')
+TIME_PATTERN: str = ('^[0-9]+('
+                     f'{"|".join([unit.lower() for unit in TIME_UNITS])}|'
+                     f'{"|".join([unit.upper() for unit in TIME_UNITS])}|'
+                     ')?$')
 
 MEMORY_SIZE_UNITS = {
     'kb': 2**10,
diff --git a/sky/templates/kubernetes-ray.yml.j2 b/sky/templates/kubernetes-ray.yml.j2
index 32d02f05d..292701a95 100644
--- a/sky/templates/kubernetes-ray.yml.j2
+++ b/sky/templates/kubernetes-ray.yml.j2
@@ -297,6 +297,9 @@ available_node_types:
         annotations:
             kueue.x-k8s.io/retriable-in-group: "false"
             kueue.x-k8s.io/pod-group-total-count: "{{ num_nodes|string }}"
+            {% if k8s_max_run_duration_seconds %}
+            provreq.kueue.x-k8s.io/maxRunDurationSeconds: "{{k8s_max_run_duration_seconds|string}}"
+            {% endif %}
         {% endif %}
       spec:
         # serviceAccountName: skypilot-service-account
@@ -309,7 +312,7 @@ available_node_types:
         {% endif %}
 
         # Add node selector if GPU/TPUs are requested:
-        {% if (k8s_topology_label_key is not none and k8s_topology_label_value is not none) or (k8s_spot_label_key is not none) %}
+        {% if (k8s_topology_label_key is not none and k8s_topology_label_value is not none) or (k8s_spot_label_key is not none) or (k8s_enable_flex_start) %}
         nodeSelector:
             {% if k8s_topology_label_key is not none and k8s_topology_label_value is not none %}
             {{k8s_topology_label_key}}: {{k8s_topology_label_value}}
@@ -317,6 +320,9 @@ available_node_types:
             {% if k8s_spot_label_key is not none %}
             {{k8s_spot_label_key}}: {{k8s_spot_label_value|tojson}}
             {% endif %}
+            {% if k8s_enable_flex_start %}
+            cloud.google.com/gke-flex-start: "true"
+            {% endif %}
         {% endif %}
         {% if (k8s_acc_label_key is not none and k8s_acc_label_values is not none) or (avoid_label_keys is not none) %}
         affinity:
@@ -392,7 +398,7 @@ available_node_types:
         {% endfor %}
         containers:
         - name: ray-node
-          imagePullPolicy: IfNotPresent
+          imagePullPolicy: Always
           image: {{image_id}}
           env:
             - name: SKYPILOT_POD_NODE_TYPE
diff --git a/sky/utils/resources_utils.py b/sky/utils/resources_utils.py
index 31b9756e2..7b70c2b49 100644
--- a/sky/utils/resources_utils.py
+++ b/sky/utils/resources_utils.py
@@ -405,3 +405,33 @@ def parse_memory_resource(resource_qty_str: Union[str, int, float],
                 continue
 
     raise ValueError(error_msg)
+
+
+def parse_time_minutes(time: str) -> int:
+    """Convert a time string to minutes.
+
+    Args:
+        time: Time string with optional unit suffix (e.g., '30m', '2h', '1d')
+
+    Returns:
+        Time in minutes as an integer
+    """
+    time_str = str(time)
+
+    if time_str.isdecimal():
+        # We assume it is already in minutes to maintain backwards
+        # compatibility
+        return int(time_str)
+
+    time_str = time_str.lower()
+    for unit, multiplier in constants.TIME_UNITS.items():
+        if time_str.endswith(unit):
+            try:
+                value = float(time_str[:-len(unit)])
+                final_value = math.ceil(value * multiplier)
+                if final_value >= 0:
+                    return final_value
+            except ValueError:
+                continue
+
+    raise ValueError(f'Invalid time format: {time}')
diff --git a/sky/utils/schemas.py b/sky/utils/schemas.py
index 3450b9d41..bd6f2755c 100644
--- a/sky/utils/schemas.py
+++ b/sky/utils/schemas.py
@@ -1084,6 +1084,25 @@ _CONTEXT_CONFIG_SCHEMA_KUBERNETES = {
             },
         },
     },
+    'dws': {
+        'type': 'object',
+        'required': [],
+        'additionalProperties': False,
+        'properties': {
+            'enabled': {
+                'type': 'boolean',
+            },
+            # Only used when Kueue is enabled.
+            'max_run_duration': {
+                'anyOf': [{
+                    'type': 'string',
+                    'pattern': constants.TIME_PATTERN,
+                }, {
+                    'type': 'integer',
+                }]
+            },
+        },
+    },
 }
 
 
diff --git a/sky/utils/tempstore.py b/sky/utils/tempstore.py
new file mode 100644
index 000000000..b47be4926
--- /dev/null
+++ b/sky/utils/tempstore.py
@@ -0,0 +1,51 @@
+"""Temporary storage context manager."""
+
+import contextlib
+import contextvars
+import os
+import tempfile
+from typing import Iterator, Optional
+
+_TEMP_DIR: contextvars.ContextVar[Optional[str]] = contextvars.ContextVar(
+    'temp_store_dir', default=None)
+
+
+@contextlib.contextmanager
+def tempdir() -> Iterator[str]:
+    """Context manager for temporary directory of current context.
+
+    This wraps tempfile.TemporaryDirectory and makes the temp dir available
+    throughout the context, eliminating the need to pass the temp dir to
+    the nested functions that need it.
+
+    This context manager is nestable - nested calls will create new temp dirs
+    and restore the previous temp dir when exiting.
+    """
+    with tempfile.TemporaryDirectory(prefix='sky-tmp') as temp_dir:
+        token = _TEMP_DIR.set(temp_dir)
+        try:
+            yield temp_dir
+        finally:
+            _TEMP_DIR.reset(token)
+
+
+# Keep the function signature same as tempfile.mkdtemp.
+# pylint: disable=redefined-builtin
+def mkdtemp(suffix: Optional[str] = None,
+            prefix: Optional[str] = None,
+            dir: Optional[str] = None) -> str:
+    """Create a temporary directory in the temp dir of current context.
+
+    The directory will be cleaned when the current context exits.
+    If there is no temp dir in current context, this function is equivalent to
+    tempfile.mkdtemp.
+    """
+    context_temp_dir = _TEMP_DIR.get()
+
+    if context_temp_dir is not None and dir is None:
+        dir = context_temp_dir
+    elif context_temp_dir is not None and dir is not None:
+        dir = os.path.join(context_temp_dir, dir)
+        os.makedirs(dir, exist_ok=True)
+
+    return tempfile.mkdtemp(suffix=suffix, prefix=prefix, dir=dir)


Test patch:
diff --git a/tests/load_tests/test_load_on_server.py b/tests/load_tests/test_load_on_server.py
index 11e532319..8497e231a 100644
--- a/tests/load_tests/test_load_on_server.py
+++ b/tests/load_tests/test_load_on_server.py
@@ -233,8 +233,13 @@ def test_validate_api(num_requests):
     run_concurrent_api_requests(num_requests, validate, 'API /validate')
 
 
+def test_api_status(num_requests):
+    print(f"Testing {num_requests} API status requests")
+    run_concurrent_api_requests(num_requests, sdk.api_status, 'API /status')
+
+
 all_requests = ['launch', 'status', 'logs', 'jobs', 'serve']
-all_apis = ['status', 'cli_status', 'tail_logs', 'validate']
+all_apis = ['status', 'cli_status', 'tail_logs', 'validate', 'api_status']
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
@@ -319,6 +324,9 @@ if __name__ == '__main__':
             elif api == 'validate':
                 thread = threading.Thread(target=test_validate_api,
                                           args=(args.n,))
+            elif api == 'api_status':
+                thread = threading.Thread(target=test_api_status,
+                                          args=(args.n,))
             if thread:
                 test_threads.append(thread)
                 thread.start()
diff --git a/tests/unit_tests/test_sky/clouds/test_kubernetes.py b/tests/unit_tests/test_sky/clouds/test_kubernetes.py
index 68e11ceb2..2e9653e91 100644
--- a/tests/unit_tests/test_sky/clouds/test_kubernetes.py
+++ b/tests/unit_tests/test_sky/clouds/test_kubernetes.py
@@ -8,6 +8,7 @@ from unittest.mock import patch
 import pytest
 
 from sky.clouds import kubernetes
+from sky.clouds.utils import gcp_utils
 from sky.provision.kubernetes import utils as kubernetes_utils
 
 
@@ -782,5 +783,87 @@ class TestKubernetesSecurityContextMerging(unittest.TestCase):
                 os.unlink(tmp_path)
 
 
+class TestGKEDWSConfig(unittest.TestCase):
+    """Test cases for gcp_utils.get_dws_config method."""
+
+    @patch('sky.skypilot_config.get_effective_region_config')
+    def test_no_dws_config(self, mock_get_config):
+        """Test when no DWS config is provided."""
+        mock_get_config.return_value = {}
+        enable_flex_start, enable_flex_start_queued_provisioning, max_run_duration = (
+            gcp_utils.get_dws_config('test-context', None))
+
+        self.assertFalse(enable_flex_start)
+        self.assertFalse(enable_flex_start_queued_provisioning)
+        self.assertIsNone(max_run_duration)
+
+    @patch('sky.skypilot_config.get_effective_region_config')
+    def test_dws_not_enabled(self, mock_get_config):
+        """Test when DWS is configured but not enabled."""
+        mock_get_config.return_value = {'enabled': False}
+        enable_flex_start, enable_flex_start_queued_provisioning, max_run_duration = (
+            gcp_utils.get_dws_config('test-context', None))
+
+        self.assertFalse(enable_flex_start)
+        self.assertFalse(enable_flex_start_queued_provisioning)
+        self.assertIsNone(max_run_duration)
+
+    @patch('sky.skypilot_config.get_effective_region_config')
+    def test_dws_enabled_without_kueue(self, mock_get_config):
+        """Test when DWS is enabled without Kueue queue name."""
+        mock_get_config.return_value = {'enabled': True}
+        enable_flex_start, enable_flex_start_queued_provisioning, max_run_duration = (
+            gcp_utils.get_dws_config('test-context', None))
+
+        self.assertTrue(enable_flex_start)
+        self.assertFalse(enable_flex_start_queued_provisioning)
+        self.assertIsNone(max_run_duration)
+
+    @patch('sky.skypilot_config.get_effective_region_config')
+    def test_dws_enabled_with_kueue(self, mock_get_config):
+        """Test when DWS is enabled with Kueue queue name."""
+        mock_get_config.return_value = {'enabled': True}
+        enable_flex_start, enable_flex_start_queued_provisioning, max_run_duration = (
+            gcp_utils.get_dws_config('test-context', 'test-queue'))
+
+        self.assertFalse(enable_flex_start)
+        self.assertTrue(enable_flex_start_queued_provisioning)
+        self.assertIsNone(max_run_duration)
+
+    @patch('sky.skypilot_config.get_effective_region_config')
+    def test_dws_enabled_with_max_duration(self, mock_get_config):
+        """Test when DWS is enabled with max run duration."""
+        mock_get_config.return_value = {
+            'enabled': True,
+            'max_run_duration': '2h'
+        }
+        enable_flex_start, enable_flex_start_queued_provisioning, max_run_duration = (
+            gcp_utils.get_dws_config('test-context', 'test-queue'))
+
+        self.assertFalse(enable_flex_start)
+        self.assertTrue(enable_flex_start_queued_provisioning)
+        # 2h = 120 minutes = 7200 seconds
+        self.assertEqual(max_run_duration, 7200)
+
+    @patch('sky.skypilot_config.get_effective_region_config')
+    def test_dws_with_cluster_overrides(self, mock_get_config):
+        """Test DWS config with cluster overrides."""
+        mock_get_config.return_value = {'enabled': False}
+        cluster_overrides = {'dws': {'enabled': False}}
+        enable_flex_start, enable_flex_start_queued_provisioning, max_run_duration = (
+            gcp_utils.get_dws_config('test-context', None, cluster_overrides))
+
+        self.assertFalse(enable_flex_start)
+        self.assertFalse(enable_flex_start_queued_provisioning)
+        self.assertIsNone(max_run_duration)
+        # Verify the override_configs parameter was passed
+        mock_get_config.assert_called_once_with(
+            cloud='kubernetes',
+            region='test-context',
+            keys=('dws',),
+            default_value={},
+            override_configs=cluster_overrides)
+
+
 if __name__ == '__main__':
     unittest.main()
diff --git a/tests/unit_tests/test_sky/server/test_server.py b/tests/unit_tests/test_sky/server/test_server.py
index d4a9ce104..e28bee88c 100644
--- a/tests/unit_tests/test_sky/server/test_server.py
+++ b/tests/unit_tests/test_sky/server/test_server.py
@@ -212,3 +212,49 @@ async def test_logs():
             request_id=mock.ANY,
             logs_path=mock_request_task.log_path,
             background_tasks=mock_background_tasks)
+
+
+@mock.patch('sky.utils.context_utils.hijack_sys_attrs')
+@mock.patch('asyncio.run')
+def test_server_run_uses_uvloop(mock_asyncio_run, mock_hijack_sys_attrs):
+    """Test that Server.run uses uvloop event loop policy."""
+    from sky.server.uvicorn import Server
+
+    config = uvicorn.Config(app='sky.server.server:app',
+                            host='127.0.0.1',
+                            port=8000)
+    server_instance = Server(config)
+    original_setup = config.setup_event_loop
+
+    uvloop_policy_set = False
+    uvloop_available = True
+
+    def setup_and_check():
+        # Call original setup to configure event loop
+        original_setup()
+        # Check if uvloop policy is now set
+        nonlocal uvloop_policy_set, uvloop_available
+        import asyncio
+        try:
+            import uvloop
+            policy = asyncio.get_event_loop_policy()
+            uvloop_policy_set = isinstance(policy, uvloop.EventLoopPolicy)
+        except ImportError:
+            # uvloop not available
+            uvloop_available = False
+
+    with mock.patch.object(config,
+                           'setup_event_loop',
+                           side_effect=setup_and_check):
+        # Call server.run
+        server_instance.run()
+
+    mock_asyncio_run.assert_called_once()
+
+    # Check uvloop policy was set (if uvloop is available)
+    if uvloop_available:
+        assert uvloop_policy_set, (
+            "Expected uvloop event loop policy to be set when uvloop "
+            "is available")
+    else:
+        pytest.skip("uvloop not available, skipping uvloop policy check")
diff --git a/tests/unit_tests/test_sky/utils/text_context.py b/tests/unit_tests/test_sky/utils/test_context.py
similarity index 100%
rename from tests/unit_tests/test_sky/utils/text_context.py
rename to tests/unit_tests/test_sky/utils/test_context.py
diff --git a/tests/unit_tests/test_sky/utils/test_resources_utils.py b/tests/unit_tests/test_sky/utils/test_resources_utils.py
new file mode 100644
index 000000000..742d454e1
--- /dev/null
+++ b/tests/unit_tests/test_sky/utils/test_resources_utils.py
@@ -0,0 +1,87 @@
+"""Tests for resources_utils.py."""
+import unittest
+
+from sky.utils import resources_utils
+
+
+class TestParseTimeMinutes(unittest.TestCase):
+    """Test cases for parse_time_minutes function."""
+
+    def test_pure_number(self):
+        """Test parsing pure numbers (assumed to be in minutes)."""
+        self.assertEqual(resources_utils.parse_time_minutes('30'), 30)
+        self.assertEqual(resources_utils.parse_time_minutes('0'), 0)
+        self.assertEqual(resources_utils.parse_time_minutes('120'), 120)
+
+    def test_minutes(self):
+        """Test parsing time strings with minute units."""
+        self.assertEqual(resources_utils.parse_time_minutes('30m'), 30)
+        self.assertEqual(resources_utils.parse_time_minutes('45M'), 45)
+        self.assertEqual(resources_utils.parse_time_minutes('0m'), 0)
+
+    def test_hours(self):
+        """Test parsing time strings with hour units."""
+        self.assertEqual(resources_utils.parse_time_minutes('1h'), 60)
+        self.assertEqual(resources_utils.parse_time_minutes('2H'), 120)
+        self.assertEqual(resources_utils.parse_time_minutes('0.5h'), 30)
+        # Test rounding up
+        self.assertEqual(resources_utils.parse_time_minutes('1.7h'), 102)
+
+    def test_days(self):
+        """Test parsing time strings with day units."""
+        self.assertEqual(resources_utils.parse_time_minutes('1d'), 1440)
+        self.assertEqual(resources_utils.parse_time_minutes('2D'), 2880)
+        self.assertEqual(resources_utils.parse_time_minutes('0.5d'), 720)
+        # Test rounding up
+        self.assertEqual(resources_utils.parse_time_minutes('1.2d'), 1728)
+
+    def test_invalid_format(self):
+        """Test invalid time format handling."""
+        invalid_inputs = [
+            '',  # Empty string
+            'abc',  # Non-numeric string
+            '30x',  # Invalid unit
+            'm30',  # Unit before number
+            '1.2.3h',  # Multiple decimal points
+            'h',  # Missing number
+            '30mm',  # Duplicate unit
+            '-30m',  # Negative time
+            '30m2h',  # Multiple units
+        ]
+        for invalid_input in invalid_inputs:
+            with self.assertRaises(
+                    ValueError, msg=f'Should fail for input: {invalid_input}'):
+                resources_utils.parse_time_minutes(invalid_input)
+
+    def test_case_insensitive(self):
+        """Test that unit parsing is case-insensitive."""
+        self.assertEqual(resources_utils.parse_time_minutes('30m'), 30)
+        self.assertEqual(resources_utils.parse_time_minutes('30M'), 30)
+        self.assertEqual(resources_utils.parse_time_minutes('1h'), 60)
+        self.assertEqual(resources_utils.parse_time_minutes('1H'), 60)
+        self.assertEqual(resources_utils.parse_time_minutes('1d'), 1440)
+        self.assertEqual(resources_utils.parse_time_minutes('1D'), 1440)
+
+    def test_decimal_values(self):
+        """Test handling of decimal values."""
+        test_cases = [
+            ('1.5h', 90),  # 1.5 hours = 90 minutes
+            ('0.5d', 720),  # 0.5 days = 12 hours = 720 minutes
+            ('2.7h', 162),  # 2.7 hours = 162 minutes
+            ('0.1d', 144),  # 0.1 days = 144 minutes
+        ]
+        for input_str, expected in test_cases:
+            self.assertEqual(resources_utils.parse_time_minutes(input_str),
+                             expected, f'Failed for input: {input_str}')
+
+    def test_rounding_behavior(self):
+        """Test that results are properly rounded up."""
+        test_cases = [
+            ('1.01h', 61),  # 60.6 minutes rounds to 61
+            ('0.01d', 15),  # 14.4 minutes rounds to 15
+            ('1.99h', 120),  # 119.4 minutes rounds to 120
+            ('0.99d', 1426),  # 1425.6 minutes rounds to 1426
+        ]
+        for input_str, expected in test_cases:
+            self.assertEqual(resources_utils.parse_time_minutes(input_str),
+                             expected, f'Failed for input: {input_str}')
diff --git a/tests/unit_tests/test_sky/utils/test_tempstore.py b/tests/unit_tests/test_sky/utils/test_tempstore.py
new file mode 100644
index 000000000..9a0b450d2
--- /dev/null
+++ b/tests/unit_tests/test_sky/utils/test_tempstore.py
@@ -0,0 +1,282 @@
+"""Unit tests for sky.utils.tempstore module."""
+
+import os
+import pathlib
+import tempfile
+from unittest import mock
+
+import pytest
+
+from sky.utils import tempstore
+
+
+def test_tempdir_context_manager():
+    """Test tempdir context manager creates and cleans up temp directory."""
+    temp_path = None
+
+    with tempstore.tempdir() as temp_dir:
+        temp_path = pathlib.Path(temp_dir)
+        assert temp_path.exists()
+        assert temp_path.is_dir()
+        assert temp_path.name.startswith('sky-tmp')
+
+    # Directory should be cleaned up after context exits
+    assert not temp_path.exists()
+
+
+def test_tempdir_context_isolation():
+    """Test tempdir context manager provides isolated temp directories."""
+    path1 = None
+    path2 = None
+
+    with tempstore.tempdir() as temp_dir1:
+        path1 = temp_dir1
+        with tempstore.tempdir() as temp_dir2:
+            path2 = temp_dir2
+            # Should be different directories
+            assert path1 != path2
+            assert pathlib.Path(path1).exists()
+            assert pathlib.Path(path2).exists()
+
+        # Inner temp dir should be cleaned up
+        assert not pathlib.Path(path2).exists()
+        # Outer temp dir should still exist
+        assert pathlib.Path(path1).exists()
+
+    # Both should be cleaned up now
+    assert not pathlib.Path(path1).exists()
+    assert not pathlib.Path(path2).exists()
+
+
+def test_tempdir_nested_contexts():
+    """Test nested tempdir contexts work correctly."""
+    outer_dir = None
+    inner_dir = None
+
+    with tempstore.tempdir() as temp_dir1:
+        outer_dir = temp_dir1
+        # Create a temp dir inside the context
+        inner_temp = tempstore.mkdtemp()
+        assert inner_temp.startswith(outer_dir)
+
+        with tempstore.tempdir() as temp_dir2:
+            inner_dir = temp_dir2
+            # This should create temp dir in the new context
+            nested_temp = tempstore.mkdtemp()
+            assert nested_temp.startswith(inner_dir)
+            assert not nested_temp.startswith(outer_dir)
+
+        # Back to outer context
+        outer_temp = tempstore.mkdtemp()
+        assert outer_temp.startswith(outer_dir)
+
+
+def test_mkdtemp_no_context():
+    """Test mkdtemp without tempdir context."""
+    # Should work like regular tempfile.mkdtemp
+    temp_dir = tempstore.mkdtemp()
+    try:
+        assert os.path.exists(temp_dir)
+        assert os.path.isdir(temp_dir)
+        # Should not have sky-tmp prefix since no context
+        assert not os.path.basename(temp_dir).startswith('sky-tmp')
+    finally:
+        # Manual cleanup since no context
+        if os.path.exists(temp_dir):
+            os.rmdir(temp_dir)
+
+
+def test_mkdtemp_with_context():
+    """Test mkdtemp with tempdir context."""
+    with tempstore.tempdir() as context_dir:
+        temp_dir = tempstore.mkdtemp()
+
+        # Should be created inside context directory
+        assert temp_dir.startswith(context_dir)
+        assert os.path.exists(temp_dir)
+        assert os.path.isdir(temp_dir)
+
+        # Should be automatically cleaned up when context exits
+
+
+def test_mkdtemp_with_parameters():
+    """Test mkdtemp with suffix, prefix, and dir parameters."""
+    with tempstore.tempdir() as context_dir:
+        # Test with suffix and prefix
+        temp_dir = tempstore.mkdtemp(suffix='_test', prefix='test_')
+        assert os.path.basename(temp_dir).startswith('test_')
+        assert os.path.basename(temp_dir).endswith('_test')
+        assert temp_dir.startswith(context_dir)
+
+        # Test with explicit dir parameter
+        subdir = 'subdir'
+        temp_dir2 = tempstore.mkdtemp(dir=subdir)
+        expected_parent = os.path.join(context_dir, subdir)
+        assert temp_dir2.startswith(expected_parent)
+
+
+def test_mkdtemp_explicit_dir_no_context():
+    """Test mkdtemp with explicit dir but no context."""
+    with tempfile.TemporaryDirectory() as temp_base:
+        subdir = 'testsubdir'
+        full_subdir = os.path.join(temp_base, subdir)
+        os.makedirs(full_subdir)
+
+        temp_dir = tempstore.mkdtemp(dir=full_subdir)
+        try:
+            assert temp_dir.startswith(full_subdir)
+            assert os.path.exists(temp_dir)
+        finally:
+            if os.path.exists(temp_dir):
+                os.rmdir(temp_dir)
+
+
+def test_mkdtemp_explicit_dir_with_context():
+    """Test mkdtemp with explicit dir and context."""
+    with tempstore.tempdir() as context_dir:
+        subdir = 'mysubdir'
+        temp_dir = tempstore.mkdtemp(dir=subdir)
+
+        # Should be created in context_dir/subdir
+        expected_parent = os.path.join(context_dir, subdir)
+        assert temp_dir.startswith(expected_parent)
+        assert os.path.exists(temp_dir)
+
+        # Test nested directory creation
+        nested_dir = 'level1/level2/level3'
+        nested_temp = tempstore.mkdtemp(dir=nested_dir)
+        expected_nested = os.path.join(context_dir, nested_dir)
+        assert nested_temp.startswith(expected_nested)
+        assert os.path.exists(nested_temp)
+
+
+def test_mkdtemp_multiple_calls():
+    """Test multiple mkdtemp calls in same context create different dirs."""
+    with tempstore.tempdir() as context_dir:
+        temp_dir1 = tempstore.mkdtemp()
+        temp_dir2 = tempstore.mkdtemp()
+        temp_dir3 = tempstore.mkdtemp(prefix='special_')
+
+        # All should be different
+        assert temp_dir1 != temp_dir2 != temp_dir3
+
+        # All should exist
+        assert os.path.exists(temp_dir1)
+        assert os.path.exists(temp_dir2)
+        assert os.path.exists(temp_dir3)
+
+        # All should be in context directory
+        assert temp_dir1.startswith(context_dir)
+        assert temp_dir2.startswith(context_dir)
+        assert temp_dir3.startswith(context_dir)
+
+        # Special prefix should be honored
+        assert os.path.basename(temp_dir3).startswith('special_')
+
+
+def test_tempdir_exception_cleanup():
+    """Test tempdir cleans up even when exception occurs."""
+    temp_path = None
+
+    with pytest.raises(ValueError):
+        with tempstore.tempdir() as temp_dir:
+            temp_path = pathlib.Path(temp_dir)
+            assert temp_path.exists()
+            raise ValueError("Test exception")
+
+    # Directory should still be cleaned up despite exception
+    assert not temp_path.exists()
+
+
+def test_contextvar_isolation():
+    """Test that context variable is properly isolated."""
+    import threading
+
+    results = {}
+
+    def thread_func(thread_id):
+        with tempstore.tempdir() as temp_dir:
+            results[thread_id] = temp_dir
+            # Create a temp directory in this thread's context
+            thread_temp = tempstore.mkdtemp()
+            results[f'{thread_id}_temp'] = thread_temp
+
+    # Run in multiple threads
+    threads = []
+    for i in range(3):
+        thread = threading.Thread(target=thread_func, args=(i,))
+        threads.append(thread)
+        thread.start()
+
+    for thread in threads:
+        thread.join()
+
+    # Each thread should have different temp directories
+    assert len(set(results[i] for i in range(3))) == 3
+
+    # Each thread's mkdtemp should be in its own context
+    for i in range(3):
+        assert results[f'{i}_temp'].startswith(results[i])
+
+
+def test_tempdir_prefix():
+    """Test that tempdir uses correct prefix."""
+    with tempstore.tempdir() as temp_dir:
+        dir_name = os.path.basename(temp_dir)
+        assert dir_name.startswith('sky-tmp')
+
+
+def test_mkdtemp_with_none_parameters():
+    """Test mkdtemp handles None parameters correctly."""
+    with tempstore.tempdir() as context_dir:
+        # All None parameters should work
+        temp_dir = tempstore.mkdtemp(suffix=None, prefix=None, dir=None)
+        assert temp_dir.startswith(context_dir)
+        assert os.path.exists(temp_dir)
+
+
+def test_tempdir_returns_string():
+    """Test that tempdir yields a string path."""
+    with tempstore.tempdir() as temp_dir:
+        assert isinstance(temp_dir, str)
+        assert os.path.isabs(temp_dir)
+
+
+def test_mkdtemp_returns_string():
+    """Test that mkdtemp returns a string path."""
+    with tempstore.tempdir():
+        temp_dir = tempstore.mkdtemp()
+        assert isinstance(temp_dir, str)
+        assert os.path.isabs(temp_dir)
+
+
+def test_mkdtemp_creates_intermediate_dirs():
+    """Test that mkdtemp creates intermediate directories when needed."""
+    with tempstore.tempdir() as context_dir:
+        # Test deep nested directory creation
+        deep_dir = 'a/b/c/d/e'
+        temp_dir = tempstore.mkdtemp(dir=deep_dir)
+
+        # All intermediate directories should be created
+        intermediate_path = os.path.join(context_dir, deep_dir)
+        assert os.path.exists(intermediate_path)
+        assert temp_dir.startswith(intermediate_path)
+
+        # Verify we can create another temp dir in the same deep structure
+        temp_dir2 = tempstore.mkdtemp(dir=deep_dir)
+        assert temp_dir2.startswith(intermediate_path)
+        assert temp_dir != temp_dir2
+
+
+def test_mkdtemp_dir_with_existing_path():
+    """Test mkdtemp with dir parameter when path already exists."""
+    with tempstore.tempdir() as context_dir:
+        # Create a subdirectory first
+        subdir = 'existing_subdir'
+        full_subdir = os.path.join(context_dir, subdir)
+        os.makedirs(full_subdir)
+
+        # Should work fine even if directory already exists
+        temp_dir = tempstore.mkdtemp(dir=subdir)
+        assert temp_dir.startswith(full_subdir)
+        assert os.path.exists(temp_dir)