skypilot-nightly 1.0.0.dev20241208__tar.gz → 1.0.0.dev20241210__tar.gz

{skypilot_nightly-1.0.0.dev20241208/skypilot_nightly.egg-info → skypilot_nightly-1.0.0.dev20241210}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: skypilot-nightly
-Version: 1.0.0.dev20241208
+Version: 1.0.0.dev20241210
 Summary: SkyPilot: An intercloud broker for the clouds
 Author: SkyPilot Team
 License: Apache 2.0

{skypilot_nightly-1.0.0.dev20241208 → skypilot_nightly-1.0.0.dev20241210}/sky/__init__.py

@@ -5,7 +5,7 @@ from typing import Optional
 import urllib.request
 
 # Replaced with the current commit when building the wheels.
-_SKYPILOT_COMMIT_SHA = '84f09efc3ee86089b45e19a6c668b5d7161ef0f6'
+_SKYPILOT_COMMIT_SHA = 'e0368888b6e64fd42cbafc698835b7aef0ce1a4c'
 
 
 def _get_git_commit():
@@ -35,7 +35,7 @@ def _get_git_commit():
 
 
 __commit__ = _get_git_commit()
-__version__ = '1.0.0.dev20241208'
+__version__ = '1.0.0.dev20241210'
 __root_dir__ = os.path.dirname(os.path.abspath(__file__))
 
 

{skypilot_nightly-1.0.0.dev20241208 → skypilot_nightly-1.0.0.dev20241210}/sky/clouds/runpod.py

@@ -24,8 +24,6 @@ class RunPod(clouds.Cloud):
     _REPR = 'RunPod'
     _CLOUD_UNSUPPORTED_FEATURES = {
         clouds.CloudImplementationFeatures.STOP: 'Stopping not supported.',
-        clouds.CloudImplementationFeatures.SPOT_INSTANCE:
-            ('Spot is not supported, as runpod API does not implement spot.'),
         clouds.CloudImplementationFeatures.MULTI_NODE:
             ('Multi-node not supported yet, as the interconnection among nodes '
              'are non-trivial on RunPod.'),
@@ -70,11 +68,8 @@ class RunPod(clouds.Cloud):
                               zone: Optional[str]) -> List[clouds.Region]:
         assert zone is None, 'RunPod does not support zones.'
         del accelerators, zone  # unused
-        if use_spot:
-            return []
-        else:
-            regions = service_catalog.get_region_zones_for_instance_type(
-                instance_type, use_spot, 'runpod')
+        regions = service_catalog.get_region_zones_for_instance_type(
+            instance_type, use_spot, 'runpod')
 
         if region is not None:
             regions = [r for r in regions if r.name == region]
@@ -176,11 +171,19 @@ class RunPod(clouds.Cloud):
         else:
             image_id = r.image_id[r.region]
 
+        instance_type = resources.instance_type
+        use_spot = resources.use_spot
+
+        hourly_cost = self.instance_type_to_hourly_cost(
+            instance_type=instance_type, use_spot=use_spot)
+
         return {
-            'instance_type': resources.instance_type,
+            'instance_type': instance_type,
             'custom_resources': custom_resources,
             'region': region.name,
             'image_id': image_id,
+            'use_spot': use_spot,
+            'bid_per_gpu': str(hourly_cost),
         }
 
     def _get_feasible_launchable_resources(

skypilot_nightly-1.0.0.dev20241210/sky/provision/runpod/api/__init__.py

@@ -0,0 +1,3 @@
+"""RunPod low level API support for spot pod."""
+
+from sky.provision.runpod.api.commands import create_spot_pod

skypilot_nightly-1.0.0.dev20241210/sky/provision/runpod/api/commands.py

@@ -0,0 +1,119 @@
+"""This module provides functions to generate GraphQL mutations for deploying
+spot instance Pods on RunPod.
+
+Reference:
+    https://github.com/runpod/runpod-python/blob/main/runpod/api/ctl_commands.py
+
+Functions:
+    generate_spot_pod_deployment_mutation: Generates a GraphQL mutation string
+        for deploying a spot instance Pod on RunPod.
+
+Example:
+    >>> mutation = generate_spot_pod_deployment_mutation(
+            name='test',
+            image_name='runpod/stack',
+            gpu_type_id='NVIDIA GeForce RTX 3070',
+            bid_per_gpu=0.3
+        )
+"""
+from typing import List, Optional
+
+from sky.adaptors import runpod
+from sky.provision.runpod.api.pods import generate_spot_pod_deployment_mutation
+
+_INTERRUPTABLE_POD_FIELD: str = 'podRentInterruptable'
+_RESPONSE_DATA_FIELD: str = 'data'
+
+
+def create_spot_pod(
+    name: str,
+    image_name: str,
+    gpu_type_id: str,
+    bid_per_gpu: float,
+    cloud_type: str = 'ALL',
+    volume_mount_path: str = '/runpod-volume',
+    gpu_count: Optional[int] = 1,
+    min_memory_in_gb: Optional[int] = 1,
+    min_vcpu_count: Optional[int] = 1,
+    container_disk_in_gb: Optional[int] = None,
+    volume_in_gb: Optional[int] = 0,
+    ports: Optional[str] = None,
+    start_ssh: Optional[bool] = True,
+    start_jupyter: Optional[bool] = False,
+    env: Optional[dict] = None,
+    docker_args: Optional[str] = '',
+    support_public_ip: Optional[bool] = True,
+    terminate_after: Optional[str] = None,
+    stop_after: Optional[str] = None,
+    data_center_id: Optional[str] = None,
+    country_code: Optional[str] = None,
+    network_volume_id: Optional[str] = None,
+    allowed_cuda_versions: Optional[List[str]] = None,
+    min_download: Optional[int] = None,
+    min_upload: Optional[int] = None,
+    cuda_version: Optional[str] = None,
+    template_id: Optional[str] = None,
+    volume_key: Optional[str] = None,
+) -> dict:
+    """This module provides functions to generate GraphQL mutations for
+    deploying spot instance Pods on RunPod.
+
+    Functions:
+        generate_spot_pod_deployment_mutation: Generates a GraphQL mutation
+            string for deploying a spot instance Pod on RunPod.
+
+    Example:
+        >>> mutation = generate_spot_pod_deployment_mutation(
+                name='test',
+                image_name='runpod/stack',
+                gpu_type_id='NVIDIA GeForce RTX 3070',
+                bid_per_gpu=0.3
+            )
+    """
+    runpod.runpod.get_gpu(gpu_type_id)
+    # refer to https://graphql-spec.runpod.io/#definition-CloudTypeEnum
+    if cloud_type not in ['ALL', 'COMMUNITY', 'SECURE']:
+        raise ValueError('cloud_type must be one of ALL, COMMUNITY or SECURE')
+
+    if network_volume_id and data_center_id is None:
+        user_info = runpod.runpod.get_user()
+        for network_volume in user_info['networkVolumes']:
+            if network_volume['id'] == network_volume_id:
+                data_center_id = network_volume['dataCenterId']
+                break
+
+    if container_disk_in_gb is None and template_id is None:
+        container_disk_in_gb = 10
+
+    mutation = generate_spot_pod_deployment_mutation(
+        name=name,
+        image_name=image_name,
+        gpu_type_id=gpu_type_id,
+        bid_per_gpu=bid_per_gpu,
+        cloud_type=cloud_type,
+        gpu_count=gpu_count,
+        min_memory_in_gb=min_memory_in_gb,
+        min_vcpu_count=min_vcpu_count,
+        container_disk_in_gb=container_disk_in_gb,
+        volume_in_gb=volume_in_gb,
+        volume_mount_path=volume_mount_path,
+        ports=ports,
+        start_ssh=start_ssh,
+        start_jupyter=start_jupyter,
+        env=env,
+        docker_args=docker_args,
+        support_public_ip=support_public_ip,
+        terminate_after=terminate_after,
+        stop_after=stop_after,
+        data_center_id=data_center_id,
+        country_code=country_code,
+        network_volume_id=network_volume_id,
+        allowed_cuda_versions=allowed_cuda_versions,
+        min_download=min_download,
+        min_upload=min_upload,
+        cuda_version=cuda_version,
+        template_id=template_id,
+        volume_key=volume_key,
+    )
+    response = runpod.runpod.api.graphql.run_graphql_query(mutation)
+    return response[_RESPONSE_DATA_FIELD][_INTERRUPTABLE_POD_FIELD]

skypilot_nightly-1.0.0.dev20241210/sky/provision/runpod/api/pods.py

@@ -0,0 +1,142 @@
+"""This module provides functions to generate GraphQL mutations for deploying
+spot instance Pods on RunPod.
+
+Reference:
+    https://github.com/runpod/runpod-python/blob/main/runpod/api/mutations/pods.py
+
+Functions:
+    generate_spot_pod_deployment_mutation: Generates a GraphQL mutation string
+        for deploying a spot instance Pod on RunPod.
+Example:
+    >>> mutation = generate_spot_pod_deployment_mutation(
+            name='test',
+            image_name='runpod/stack',
+            gpu_type_id='NVIDIA GeForce RTX 3070',
+            bid_per_gpu=0.3
+        )
+"""
+
+from typing import List, Optional
+
+
+# refer to https://graphql-spec.runpod.io/#definition-PodRentInterruptableInput
+def generate_spot_pod_deployment_mutation(
+    name: str,
+    image_name: str,
+    gpu_type_id: str,
+    bid_per_gpu: float,
+    volume_mount_path: str,
+    cloud_type: str = 'ALL',
+    gpu_count: Optional[int] = None,
+    min_memory_in_gb: Optional[int] = None,
+    min_vcpu_count: Optional[int] = None,
+    container_disk_in_gb: Optional[int] = None,
+    volume_in_gb: Optional[int] = None,
+    ports: Optional[str] = None,
+    start_ssh: Optional[bool] = True,
+    start_jupyter: Optional[bool] = False,
+    env: Optional[dict] = None,
+    docker_args: Optional[str] = None,
+    support_public_ip: Optional[bool] = True,
+    terminate_after: Optional[str] = None,
+    stop_after: Optional[str] = None,
+    data_center_id: Optional[str] = None,
+    country_code: Optional[str] = None,
+    network_volume_id: Optional[str] = None,
+    allowed_cuda_versions: Optional[List[str]] = None,
+    min_download: Optional[int] = None,
+    min_upload: Optional[int] = None,
+    cuda_version: Optional[str] = None,
+    template_id: Optional[str] = None,
+    volume_key: Optional[str] = None,
+) -> str:
+    input_fields = []
+
+    # Required Fields
+    input_fields.append(f'name: "{name}"')
+    input_fields.append(f'imageName: "{image_name}"')
+    input_fields.append(f'gpuTypeId: "{gpu_type_id}"')
+    input_fields.append(f'bidPerGpu: {bid_per_gpu}')
+    input_fields.append(f'volumeMountPath: "{volume_mount_path}"')
+
+    # Default Fields
+    input_fields.append(f'cloudType: {cloud_type}')
+
+    if start_ssh:
+        input_fields.append('startSsh: true')
+    if start_jupyter:
+        input_fields.append('startJupyter: true')
+    if support_public_ip:
+        input_fields.append('supportPublicIp: true')
+    else:
+        input_fields.append('supportPublicIp: false')
+
+    # Optional Fields
+    if gpu_count is not None:
+        input_fields.append(f'gpuCount: {gpu_count}')
+    if min_memory_in_gb is not None:
+        input_fields.append(f'minMemoryInGb: {min_memory_in_gb}')
+    if min_vcpu_count is not None:
+        input_fields.append(f'minVcpuCount: {min_vcpu_count}')
+    if container_disk_in_gb is not None:
+        input_fields.append(f'containerDiskInGb: {container_disk_in_gb}')
+    if volume_in_gb is not None:
+        input_fields.append(f'volumeInGb: {volume_in_gb}')
+    if ports is not None:
+        ports = ports.replace(' ', '')
+        input_fields.append(f'ports: "{ports}"')
+    if docker_args is not None:
+        input_fields.append(f'dockerArgs: "{docker_args}"')
+    if terminate_after is not None:
+        input_fields.append(f'terminateAfter: "{terminate_after}"')
+    if stop_after is not None:
+        input_fields.append(f'stopAfter: "{stop_after}"')
+    if data_center_id is not None:
+        input_fields.append(f'dataCenterId: "{data_center_id}"')
+    if country_code is not None:
+        input_fields.append(f'countryCode: "{country_code}"')
+    if network_volume_id is not None:
+        input_fields.append(f'networkVolumeId: "{network_volume_id}"')
+    if allowed_cuda_versions is not None:
+        allowed_cuda_versions_string = ', '.join(
+            [f'"{version}"' for version in allowed_cuda_versions])
+        input_fields.append(
+            f'allowedCudaVersions: [{allowed_cuda_versions_string}]')
+    if min_download is not None:
+        input_fields.append(f'minDownload: {min_download}')
+    if min_upload is not None:
+        input_fields.append(f'minUpload: {min_upload}')
+    if cuda_version is not None:
+        input_fields.append(f'cudaVersion: "{cuda_version}"')
+    if template_id is not None:
+        input_fields.append(f'templateId: "{template_id}"')
+    if volume_key is not None:
+        input_fields.append(f'volumeKey: "{volume_key}"')
+
+    if env is not None:
+        env_string = ', '.join([
+            f'{{ key: "{key}", value: "{value}" }}'
+            for key, value in env.items()
+        ])
+        input_fields.append(f'env: [{env_string}]')
+
+    # Format input fields
+    input_string = ', '.join(input_fields)
+    return f"""
+    mutation {{
+      podRentInterruptable(
+        input: {{
+          {input_string}
+        }}
+      ) {{
+        id
+        desiredStatus
+        imageName
+        env
+        machineId
+        machine {{
+          podHostId
+        }}
+      }}
+    }}
+    """

{skypilot_nightly-1.0.0.dev20241208 → skypilot_nightly-1.0.0.dev20241210}/sky/provision/runpod/instance.py

@@ -89,7 +89,10 @@ def run_instances(region: str, cluster_name_on_cloud: str,
                 disk_size=config.node_config['DiskSize'],
                 image_name=config.node_config['ImageId'],
                 ports=config.ports_to_open_on_launch,
-                public_key=config.node_config['PublicKey'])
+                public_key=config.node_config['PublicKey'],
+                preemptible=config.node_config['Preemptible'],
+                bid_per_gpu=config.node_config['BidPerGPU'],
+            )
         except Exception as e:  # pylint: disable=broad-except
             logger.warning(f'run_instances error: {e}')
             raise

{skypilot_nightly-1.0.0.dev20241208 → skypilot_nightly-1.0.0.dev20241210}/sky/provision/runpod/utils.py

@@ -6,6 +6,7 @@ from typing import Any, Dict, List, Optional
 
 from sky import sky_logging
 from sky.adaptors import runpod
+import sky.provision.runpod.api.commands as runpod_commands
 from sky.skylet import constants
 from sky.utils import common_utils
 
@@ -100,7 +101,8 @@ def list_instances() -> Dict[str, Dict[str, Any]]:
 
 
 def launch(name: str, instance_type: str, region: str, disk_size: int,
-           image_name: str, ports: Optional[List[int]], public_key: str) -> str:
+           image_name: str, ports: Optional[List[int]], public_key: str,
+           preemptible: Optional[bool], bid_per_gpu: float) -> str:
     """Launches an instance with the given parameters.
 
     Converts the instance_type to the RunPod GPU name, finds the specs for the
@@ -142,23 +144,35 @@ def launch(name: str, instance_type: str, region: str, disk_size: int,
     if ports is not None:
         custom_ports_str = ''.join([f'{p}/tcp,' for p in ports])
 
-    new_instance = runpod.runpod.create_pod(
-        name=name,
-        image_name=image_name,
-        gpu_type_id=gpu_type,
-        cloud_type=cloud_type,
-        container_disk_in_gb=disk_size,
-        min_vcpu_count=4 * gpu_quantity,
-        min_memory_in_gb=gpu_specs['memoryInGb'] * gpu_quantity,
-        gpu_count=gpu_quantity,
-        country_code=region,
-        ports=(f'22/tcp,'
-               f'{custom_ports_str}'
-               f'{constants.SKY_REMOTE_RAY_DASHBOARD_PORT}/http,'
-               f'{constants.SKY_REMOTE_RAY_PORT}/http'),
-        support_public_ip=True,
-        docker_args=
-        f'bash -c \'echo {encoded} | base64 --decode > init.sh; bash init.sh\'')
+    docker_args = (f'bash -c \'echo {encoded} | base64 --decode > init.sh; '
+                   f'bash init.sh\'')
+    ports = (f'22/tcp,'
+             f'{custom_ports_str}'
+             f'{constants.SKY_REMOTE_RAY_DASHBOARD_PORT}/http,'
+             f'{constants.SKY_REMOTE_RAY_PORT}/http')
+
+    params = {
+        'name': name,
+        'image_name': image_name,
+        'gpu_type_id': gpu_type,
+        'cloud_type': cloud_type,
+        'container_disk_in_gb': disk_size,
+        'min_vcpu_count': 4 * gpu_quantity,
+        'min_memory_in_gb': gpu_specs['memoryInGb'] * gpu_quantity,
+        'gpu_count': gpu_quantity,
+        'country_code': region,
+        'ports': ports,
+        'support_public_ip': True,
+        'docker_args': docker_args,
+    }
+
+    if preemptible is None or not preemptible:
+        new_instance = runpod.runpod.create_pod(**params)
+    else:
+        new_instance = runpod_commands.create_spot_pod(
+            bid_per_gpu=bid_per_gpu,
+            **params,
+        )
 
     return new_instance['id']
 

{skypilot_nightly-1.0.0.dev20241208 → skypilot_nightly-1.0.0.dev20241210}/sky/templates/runpod-ray.yml.j2

@@ -24,6 +24,8 @@ available_node_types:
       ImageId: {{image_id}}
       PublicKey: |-
         skypilot:ssh_public_key_content
+      Preemptible: {{use_spot}}
+      BidPerGPU: {{bid_per_gpu}}
 
 head_node_type: ray_head_default
 

{skypilot_nightly-1.0.0.dev20241208 → skypilot_nightly-1.0.0.dev20241210/skypilot_nightly.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: skypilot-nightly
-Version: 1.0.0.dev20241208
+Version: 1.0.0.dev20241210
 Summary: SkyPilot: An intercloud broker for the clouds
 Author: SkyPilot Team
 License: Apache 2.0

{skypilot_nightly-1.0.0.dev20241208 → skypilot_nightly-1.0.0.dev20241210}/skypilot_nightly.egg-info/SOURCES.txt

@@ -164,6 +164,9 @@ sky/provision/runpod/__init__.py
 sky/provision/runpod/config.py
 sky/provision/runpod/instance.py
 sky/provision/runpod/utils.py
+sky/provision/runpod/api/__init__.py
+sky/provision/runpod/api/commands.py
+sky/provision/runpod/api/pods.py
 sky/provision/vsphere/__init__.py
 sky/provision/vsphere/config.py
 sky/provision/vsphere/instance.py

{skypilot_nightly-1.0.0.dev20241208 → skypilot_nightly-1.0.0.dev20241210}/tests/test_smoke.py

@@ -37,7 +37,7 @@ import sys
 import tempfile
 import textwrap
 import time
-from typing import Dict, List, NamedTuple, Optional, Tuple
+from typing import Dict, List, NamedTuple, Optional, TextIO, Tuple
 import urllib.parse
 import uuid
 
@@ -1333,34 +1333,68 @@ def test_using_file_mounts_with_env_vars(generic_cloud: str):
 
 
 # ---------- storage ----------
+
+
+def _storage_mounts_commands_generator(f: TextIO, cluster_name: str,
+                                       storage_name: str, ls_hello_command: str,
+                                       cloud: str, only_mount: bool):
+    template_str = pathlib.Path(
+        'tests/test_yamls/test_storage_mounting.yaml.j2').read_text()
+    template = jinja2.Template(template_str)
+    content = template.render(
+        storage_name=storage_name,
+        cloud=cloud,
+        only_mount=only_mount,
+    )
+    f.write(content)
+    f.flush()
+    file_path = f.name
+    test_commands = [
+        *STORAGE_SETUP_COMMANDS,
+        f'sky launch -y -c {cluster_name} --cloud {cloud} {file_path}',
+        f'sky logs {cluster_name} 1 --status',  # Ensure job succeeded.
+        ls_hello_command,
+        f'sky stop -y {cluster_name}',
+        f'sky start -y {cluster_name}',
+        # Check if hello.txt from mounting bucket exists after restart in
+        # the mounted directory
+        f'sky exec {cluster_name} -- "set -ex; ls /mount_private_mount/hello.txt"',
+    ]
+    clean_command = f'sky down -y {cluster_name}; sky storage delete -y {storage_name}'
+    return test_commands, clean_command
+
+
 @pytest.mark.aws
 def test_aws_storage_mounts_with_stop():
     name = _get_cluster_name()
     cloud = 'aws'
     storage_name = f'sky-test-{int(time.time())}'
-    template_str = pathlib.Path(
-        'tests/test_yamls/test_storage_mounting.yaml.j2').read_text()
-    template = jinja2.Template(template_str)
-    content = template.render(storage_name=storage_name, cloud=cloud)
+    ls_hello_command = f'aws s3 ls {storage_name}/hello.txt'
     with tempfile.NamedTemporaryFile(suffix='.yaml', mode='w') as f:
-        f.write(content)
-        f.flush()
-        file_path = f.name
-        test_commands = [
-            *STORAGE_SETUP_COMMANDS,
-            f'sky launch -y -c {name} --cloud {cloud} {file_path}',
-            f'sky logs {name} 1 --status',  # Ensure job succeeded.
-            f'aws s3 ls {storage_name}/hello.txt',
-            f'sky stop -y {name}',
-            f'sky start -y {name}',
-            # Check if hello.txt from mounting bucket exists after restart in
-            # the mounted directory
-            f'sky exec {name} -- "set -ex; ls /mount_private_mount/hello.txt"'
-        ]
+        test_commands, clean_command = _storage_mounts_commands_generator(
+            f, name, storage_name, ls_hello_command, cloud, False)
         test = Test(
             'aws_storage_mounts',
             test_commands,
-            f'sky down -y {name}; sky storage delete -y {storage_name}',
+            clean_command,
+            timeout=20 * 60,  # 20 mins
+        )
+        run_one_test(test)
+
+
+@pytest.mark.aws
+def test_aws_storage_mounts_with_stop_only_mount():
+    name = _get_cluster_name()
+    cloud = 'aws'
+    storage_name = f'sky-test-{int(time.time())}'
+    ls_hello_command = f'aws s3 ls {storage_name}/hello.txt'
+    with tempfile.NamedTemporaryFile(suffix='.yaml', mode='w') as f:
+        test_commands, clean_command = _storage_mounts_commands_generator(
+            f, name, storage_name, ls_hello_command, cloud, True)
+        test = Test(
+            'aws_storage_mounts_only_mount',
+            test_commands,
+            clean_command,
             timeout=20 * 60,  # 20 mins
         )
         run_one_test(test)
@@ -1371,29 +1405,14 @@ def test_gcp_storage_mounts_with_stop():
     name = _get_cluster_name()
     cloud = 'gcp'
     storage_name = f'sky-test-{int(time.time())}'
-    template_str = pathlib.Path(
-        'tests/test_yamls/test_storage_mounting.yaml.j2').read_text()
-    template = jinja2.Template(template_str)
-    content = template.render(storage_name=storage_name, cloud=cloud)
+    ls_hello_command = f'gsutil ls gs://{storage_name}/hello.txt'
     with tempfile.NamedTemporaryFile(suffix='.yaml', mode='w') as f:
-        f.write(content)
-        f.flush()
-        file_path = f.name
-        test_commands = [
-            *STORAGE_SETUP_COMMANDS,
-            f'sky launch -y -c {name} --cloud {cloud} {file_path}',
-            f'sky logs {name} 1 --status',  # Ensure job succeeded.
-            f'gsutil ls gs://{storage_name}/hello.txt',
-            f'sky stop -y {name}',
-            f'sky start -y {name}',
-            # Check if hello.txt from mounting bucket exists after restart in
-            # the mounted directory
-            f'sky exec {name} -- "set -ex; ls /mount_private_mount/hello.txt"'
-        ]
+        test_commands, clean_command = _storage_mounts_commands_generator(
+            f, name, storage_name, ls_hello_command, cloud, False)
         test = Test(
             'gcp_storage_mounts',
             test_commands,
-            f'sky down -y {name}; sky storage delete -y {storage_name}',
+            clean_command,
             timeout=20 * 60,  # 20 mins
         )
         run_one_test(test)
@@ -1409,31 +1428,19 @@ def test_azure_storage_mounts_with_stop():
                             get_default_storage_account_name(default_region))
     storage_account_key = data_utils.get_az_storage_account_key(
         storage_account_name)
-    template_str = pathlib.Path(
-        'tests/test_yamls/test_storage_mounting.yaml.j2').read_text()
-    template = jinja2.Template(template_str)
-    content = template.render(storage_name=storage_name, cloud=cloud)
+    # if the file does not exist, az storage blob list returns '[]'
+    ls_hello_command = (f'output=$(az storage blob list -c {storage_name} '
+                        f'--account-name {storage_account_name} '
+                        f'--account-key {storage_account_key} '
+                        f'--prefix hello.txt) '
+                        f'[ "$output" = "[]" ] && exit 1 || exit 0')
     with tempfile.NamedTemporaryFile(suffix='.yaml', mode='w') as f:
-        f.write(content)
-        f.flush()
-        file_path = f.name
-        test_commands = [
-            *STORAGE_SETUP_COMMANDS,
-            f'sky launch -y -c {name} --cloud {cloud} {file_path}',
-            f'sky logs {name} 1 --status',  # Ensure job succeeded.
-            f'output=$(az storage blob list -c {storage_name} --account-name {storage_account_name} --account-key {storage_account_key} --prefix hello.txt)'
-            # if the file does not exist, az storage blob list returns '[]'
-            f'[ "$output" = "[]" ] && exit 1;'
-            f'sky stop -y {name}',
-            f'sky start -y {name}',
-            # Check if hello.txt from mounting bucket exists after restart in
-            # the mounted directory
-            f'sky exec {name} -- "set -ex; ls /mount_private_mount/hello.txt"'
-        ]
+        test_commands, clean_command = _storage_mounts_commands_generator(
+            f, name, storage_name, ls_hello_command, cloud, False)
         test = Test(
             'azure_storage_mounts',
             test_commands,
-            f'sky down -y {name}; sky storage delete -y {storage_name}',
+            clean_command,
             timeout=20 * 60,  # 20 mins
         )
         run_one_test(test)
@@ -1446,25 +1453,15 @@ def test_kubernetes_storage_mounts():
     # built for x86_64 only.
     name = _get_cluster_name()
     storage_name = f'sky-test-{int(time.time())}'
-    template_str = pathlib.Path(
-        'tests/test_yamls/test_storage_mounting.yaml.j2').read_text()
-    template = jinja2.Template(template_str)
-    content = template.render(storage_name=storage_name)
+    ls_hello_command = (f'aws s3 ls {storage_name}/hello.txt || '
+                        f'gsutil ls gs://{storage_name}/hello.txt')
     with tempfile.NamedTemporaryFile(suffix='.yaml', mode='w') as f:
-        f.write(content)
-        f.flush()
-        file_path = f.name
-        test_commands = [
-            *STORAGE_SETUP_COMMANDS,
-            f'sky launch -y -c {name} --cloud kubernetes {file_path}',
-            f'sky logs {name} 1 --status',  # Ensure job succeeded.
-            f'aws s3 ls {storage_name}/hello.txt || '
-            f'gsutil ls gs://{storage_name}/hello.txt',
-        ]
+        test_commands, clean_command = _storage_mounts_commands_generator(
+            f, name, storage_name, ls_hello_command, 'kubernetes', False)
         test = Test(
             'kubernetes_storage_mounts',
             test_commands,
-            f'sky down -y {name}; sky storage delete -y {storage_name}',
+            clean_command,
             timeout=20 * 60,  # 20 mins
         )
         run_one_test(test)