dstack 0.19.17__py3-none-any.whl → 0.19.19__py3-none-any.whl

dstack/_internal/cli/services/configurators/fleet.py

@@ -25,6 +25,7 @@ from dstack._internal.core.errors import (
     ServerClientError,
     URLNotFoundError,
 )
+from dstack._internal.core.models.common import ApplyAction
 from dstack._internal.core.models.configurations import ApplyConfigurationType
 from dstack._internal.core.models.fleets import (
     Fleet,
@@ -35,6 +36,7 @@ from dstack._internal.core.models.fleets import (
 )
 from dstack._internal.core.models.instances import InstanceAvailability, InstanceStatus, SSHKey
 from dstack._internal.core.models.repos.base import Repo
+from dstack._internal.core.services.diff import diff_models
 from dstack._internal.utils.common import local_time
 from dstack._internal.utils.logging import get_logger
 from dstack._internal.utils.ssh import convert_ssh_key_to_pem, generate_public_key, pkey_from_str
@@ -71,7 +73,14 @@ class FleetConfigurator(ApplyEnvVarsConfiguratorMixin, BaseApplyConfigurator):
                 spec=spec,
             )
         _print_plan_header(plan)
+        if plan.action is not None:
+            self._apply_plan(plan, command_args)
+        else:
+            # Old servers don't support spec update
+            self._apply_plan_on_old_server(plan, command_args)
 
+    def _apply_plan(self, plan: FleetPlan, command_args: argparse.Namespace):
+        delete_fleet_name: Optional[str] = None
         action_message = ""
         confirm_message = ""
         if plan.current_resource is None:
@@ -82,7 +91,108 @@ class FleetConfigurator(ApplyEnvVarsConfiguratorMixin, BaseApplyConfigurator):
             confirm_message += "Create the fleet?"
         else:
             action_message += f"Found fleet [code]{plan.spec.configuration.name}[/]."
-            if plan.current_resource.spec.configuration == plan.spec.configuration:
+            if plan.action == ApplyAction.CREATE:
+                delete_fleet_name = plan.current_resource.name
+                action_message += (
+                    " Configuration changes detected. Cannot update the fleet in-place"
+                )
+                confirm_message += "Re-create the fleet?"
+            elif plan.current_resource.spec == plan.effective_spec:
+                if command_args.yes and not command_args.force:
+                    # --force is required only with --yes,
+                    # otherwise we may ask for force apply interactively.
+                    console.print(
+                        "No configuration changes detected. Use --force to apply anyway."
+                    )
+                    return
+                delete_fleet_name = plan.current_resource.name
+                action_message += " No configuration changes detected."
+                confirm_message += "Re-create the fleet?"
+            else:
+                action_message += " Configuration changes detected."
+                confirm_message += "Update the fleet in-place?"
+
+        console.print(action_message)
+        if not command_args.yes and not confirm_ask(confirm_message):
+            console.print("\nExiting...")
+            return
+
+        if delete_fleet_name is not None:
+            with console.status("Deleting existing fleet..."):
+                self.api.client.fleets.delete(
+                    project_name=self.api.project, names=[delete_fleet_name]
+                )
+                # Fleet deletion is async. Wait for fleet to be deleted.
+                while True:
+                    try:
+                        self.api.client.fleets.get(
+                            project_name=self.api.project, name=delete_fleet_name
+                        )
+                    except ResourceNotExistsError:
+                        break
+                    else:
+                        time.sleep(1)
+
+        try:
+            with console.status("Applying plan..."):
+                fleet = self.api.client.fleets.apply_plan(project_name=self.api.project, plan=plan)
+        except ServerClientError as e:
+            raise CLIError(e.msg)
+        if command_args.detach:
+            console.print("Fleet configuration submitted. Exiting...")
+            return
+        try:
+            with MultiItemStatus(
+                f"Provisioning [code]{fleet.name}[/]...", console=console
+            ) as live:
+                while not _finished_provisioning(fleet):
+                    table = get_fleets_table([fleet])
+                    live.update(table)
+                    time.sleep(LIVE_TABLE_PROVISION_INTERVAL_SECS)
+                    fleet = self.api.client.fleets.get(self.api.project, fleet.name)
+        except KeyboardInterrupt:
+            if confirm_ask("Delete the fleet before exiting?"):
+                with console.status("Deleting fleet..."):
+                    self.api.client.fleets.delete(
+                        project_name=self.api.project, names=[fleet.name]
+                    )
+            else:
+                console.print("Exiting... Fleet provisioning will continue in the background.")
+            return
+        console.print(
+            get_fleets_table(
+                [fleet],
+                verbose=_failed_provisioning(fleet),
+                format_date=local_time,
+            )
+        )
+        if _failed_provisioning(fleet):
+            console.print("\n[error]Some instances failed. Check the table above for errors.[/]")
+            exit(1)
+
+    def _apply_plan_on_old_server(self, plan: FleetPlan, command_args: argparse.Namespace):
+        action_message = ""
+        confirm_message = ""
+        if plan.current_resource is None:
+            if plan.spec.configuration.name is not None:
+                action_message += (
+                    f"Fleet [code]{plan.spec.configuration.name}[/] does not exist yet."
+                )
+            confirm_message += "Create the fleet?"
+        else:
+            action_message += f"Found fleet [code]{plan.spec.configuration.name}[/]."
+            diff = diff_models(
+                old=plan.current_resource.spec.configuration,
+                new=plan.spec.configuration,
+                reset={
+                    "ssh_config": {
+                        "ssh_key": True,
+                        "proxy_jump": {"ssh_key"},
+                        "hosts": {"__all__": {"ssh_key": True, "proxy_jump": {"ssh_key"}}},
+                    }
+                },
+            )
+            if not diff:
                 if command_args.yes and not command_args.force:
                     # --force is required only with --yes,
                     # otherwise we may ask for force apply interactively.

dstack/_internal/cli/services/profile.py

@@ -159,7 +159,7 @@ def apply_profile_args(
     if args.idle_duration is not None:
         profile_settings.idle_duration = args.idle_duration
     elif args.dont_destroy:
-        profile_settings.idle_duration = False
+        profile_settings.idle_duration = "off"
     if args.creation_policy_reuse:
         profile_settings.creation_policy = CreationPolicy.REUSE
 

dstack/_internal/core/backends/aws/compute.py

@@ -1,14 +1,21 @@
+import threading
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from typing import Any, Dict, List, Optional, Tuple
 
 import boto3
 import botocore.client
 import botocore.exceptions
+from cachetools import Cache, TTLCache, cachedmethod
+from cachetools.keys import hashkey
 from pydantic import ValidationError
 
 import dstack._internal.core.backends.aws.resources as aws_resources
 from dstack._internal import settings
-from dstack._internal.core.backends.aws.models import AWSAccessKeyCreds, AWSConfig
+from dstack._internal.core.backends.aws.models import (
+    AWSAccessKeyCreds,
+    AWSConfig,
+    AWSOSImageConfig,
+)
 from dstack._internal.core.backends.base.compute import (
     Compute,
     ComputeWithCreateInstanceSupport,
@@ -26,7 +33,12 @@ from dstack._internal.core.backends.base.compute import (
     merge_tags,
 )
 from dstack._internal.core.backends.base.offers import get_catalog_offers
-from dstack._internal.core.errors import ComputeError, NoCapacityError, PlacementGroupInUseError
+from dstack._internal.core.errors import (
+    ComputeError,
+    NoCapacityError,
+    PlacementGroupInUseError,
+    PlacementGroupNotSupportedError,
+)
 from dstack._internal.core.models.backends.base import BackendType
 from dstack._internal.core.models.common import CoreModel
 from dstack._internal.core.models.gateways import (
@@ -39,7 +51,11 @@ from dstack._internal.core.models.instances import (
     InstanceOffer,
     InstanceOfferWithAvailability,
 )
-from dstack._internal.core.models.placement import PlacementGroup, PlacementGroupProvisioningData
+from dstack._internal.core.models.placement import (
+    PlacementGroup,
+    PlacementGroupProvisioningData,
+    PlacementStrategy,
+)
 from dstack._internal.core.models.resources import Memory, Range
 from dstack._internal.core.models.runs import JobProvisioningData, Requirements
 from dstack._internal.core.models.volumes import (
@@ -66,6 +82,10 @@ class AWSVolumeBackendData(CoreModel):
     iops: int
 
 
+def _ec2client_cache_methodkey(self, ec2_client, *args, **kwargs):
+    return hashkey(*args, **kwargs)
+
+
 class AWSCompute(
     ComputeWithCreateInstanceSupport,
     ComputeWithMultinodeSupport,
@@ -86,6 +106,24 @@ class AWSCompute(
             )
         else:  # default creds
             self.session = boto3.Session()
+        # Caches to avoid redundant API calls when provisioning many instances
+        # get_offers is already cached but we still cache its sub-functions
+        # with more aggressive/longer caches.
+        self._get_regions_to_quotas_cache_lock = threading.Lock()
+        self._get_regions_to_quotas_execution_lock = threading.Lock()
+        self._get_regions_to_quotas_cache = TTLCache(maxsize=10, ttl=300)
+        self._get_regions_to_zones_cache_lock = threading.Lock()
+        self._get_regions_to_zones_cache = Cache(maxsize=10)
+        self._get_vpc_id_subnet_id_or_error_cache_lock = threading.Lock()
+        self._get_vpc_id_subnet_id_or_error_cache = TTLCache(maxsize=100, ttl=600)
+        self._get_maximum_efa_interfaces_cache_lock = threading.Lock()
+        self._get_maximum_efa_interfaces_cache = Cache(maxsize=100)
+        self._get_subnets_availability_zones_cache_lock = threading.Lock()
+        self._get_subnets_availability_zones_cache = Cache(maxsize=100)
+        self._create_security_group_cache_lock = threading.Lock()
+        self._create_security_group_cache = TTLCache(maxsize=100, ttl=600)
+        self._get_image_id_and_username_cache_lock = threading.Lock()
+        self._get_image_id_and_username_cache = TTLCache(maxsize=100, ttl=600)
 
     def get_offers(
         self, requirements: Optional[Requirements] = None
@@ -126,8 +164,11 @@ class AWSCompute(
             extra_filter=filter,
         )
         regions = list(set(i.region for i in offers))
-        regions_to_quotas = _get_regions_to_quotas(self.session, regions)
-        regions_to_zones = _get_regions_to_zones(self.session, regions)
+        with self._get_regions_to_quotas_execution_lock:
+            # Cache lock does not prevent concurrent execution.
+            # We use a separate lock to avoid requesting quotas in parallel and hitting rate limits.
+            regions_to_quotas = self._get_regions_to_quotas(self.session, regions)
+        regions_to_zones = self._get_regions_to_zones(self.session, regions)
 
         availability_offers = []
         for offer in offers:
@@ -186,21 +227,24 @@ class AWSCompute(
         tags = aws_resources.filter_invalid_tags(tags)
 
         disk_size = round(instance_offer.instance.resources.disk.size_mib / 1024)
-        max_efa_interfaces = _get_maximum_efa_interfaces(
-            ec2_client=ec2_client, instance_type=instance_offer.instance.name
+        max_efa_interfaces = self._get_maximum_efa_interfaces(
+            ec2_client=ec2_client,
+            region=instance_offer.region,
+            instance_type=instance_offer.instance.name,
         )
         enable_efa = max_efa_interfaces > 0
         is_capacity_block = False
         try:
-            vpc_id, subnet_ids = get_vpc_id_subnet_id_or_error(
+            vpc_id, subnet_ids = self._get_vpc_id_subnet_id_or_error(
                 ec2_client=ec2_client,
                 config=self.config,
                 region=instance_offer.region,
                 allocate_public_ip=allocate_public_ip,
                 availability_zones=zones,
             )
-            subnet_id_to_az_map = aws_resources.get_subnets_availability_zones(
+            subnet_id_to_az_map = self._get_subnets_availability_zones(
                 ec2_client=ec2_client,
+                region=instance_offer.region,
                 subnet_ids=subnet_ids,
             )
             if instance_config.reservation:
@@ -229,12 +273,19 @@ class AWSCompute(
             tried_zones.add(az)
             try:
                 logger.debug("Trying provisioning %s in %s", instance_offer.instance.name, az)
-                image_id, username = aws_resources.get_image_id_and_username(
+                image_id, username = self._get_image_id_and_username(
                     ec2_client=ec2_client,
+                    region=instance_offer.region,
                     cuda=len(instance_offer.instance.resources.gpus) > 0,
                     instance_type=instance_offer.instance.name,
                     image_config=self.config.os_images,
                 )
+                security_group_id = self._create_security_group(
+                    ec2_client=ec2_client,
+                    region=instance_offer.region,
+                    project_id=project_name,
+                    vpc_id=vpc_id,
+                )
                 response = ec2_resource.create_instances(
                     **aws_resources.create_instances_struct(
                         disk_size=disk_size,
@@ -243,11 +294,7 @@ class AWSCompute(
                         iam_instance_profile=self.config.iam_instance_profile,
                         user_data=get_user_data(authorized_keys=instance_config.get_public_keys()),
                         tags=aws_resources.make_tags(tags),
-                        security_group_id=aws_resources.create_security_group(
-                            ec2_client=ec2_client,
-                            project_id=project_name,
-                            vpc_id=vpc_id,
-                        ),
+                        security_group_id=security_group_id,
                         spot=instance_offer.instance.resources.spot,
                         subnet_id=subnet_id,
                         allocate_public_ip=allocate_public_ip,
@@ -296,6 +343,8 @@ class AWSCompute(
         placement_group: PlacementGroup,
         master_instance_offer: InstanceOffer,
     ) -> PlacementGroupProvisioningData:
+        if not _offer_supports_placement_group(master_instance_offer, placement_group):
+            raise PlacementGroupNotSupportedError()
         ec2_client = self.session.client("ec2", region_name=placement_group.configuration.region)
         logger.debug("Creating placement group %s...", placement_group.name)
         ec2_client.create_placement_group(
@@ -332,6 +381,8 @@ class AWSCompute(
         placement_group: PlacementGroup,
         instance_offer: InstanceOffer,
     ) -> bool:
+        if not _offer_supports_placement_group(instance_offer, placement_group):
+            return False
         return (
             placement_group.configuration.backend == BackendType.AWS
             and placement_group.configuration.region == instance_offer.region
@@ -361,7 +412,7 @@ class AWSCompute(
         tags = aws_resources.filter_invalid_tags(tags)
         tags = aws_resources.make_tags(tags)
 
-        vpc_id, subnets_ids = get_vpc_id_subnet_id_or_error(
+        vpc_id, subnets_ids = self._get_vpc_id_subnet_id_or_error(
             ec2_client=ec2_client,
             config=self.config,
             region=configuration.region,
@@ -696,6 +747,165 @@ class AWSCompute(
             return True
         return True
 
+    def _get_regions_to_quotas_key(
+        self,
+        session: boto3.Session,
+        regions: List[str],
+    ) -> tuple:
+        return hashkey(tuple(regions))
+
+    @cachedmethod(
+        cache=lambda self: self._get_regions_to_quotas_cache,
+        key=_get_regions_to_quotas_key,
+        lock=lambda self: self._get_regions_to_quotas_cache_lock,
+    )
+    def _get_regions_to_quotas(
+        self,
+        session: boto3.Session,
+        regions: List[str],
+    ) -> Dict[str, Dict[str, int]]:
+        return _get_regions_to_quotas(session=session, regions=regions)
+
+    def _get_regions_to_zones_key(
+        self,
+        session: boto3.Session,
+        regions: List[str],
+    ) -> tuple:
+        return hashkey(tuple(regions))
+
+    @cachedmethod(
+        cache=lambda self: self._get_regions_to_zones_cache,
+        key=_get_regions_to_zones_key,
+        lock=lambda self: self._get_regions_to_zones_cache_lock,
+    )
+    def _get_regions_to_zones(
+        self,
+        session: boto3.Session,
+        regions: List[str],
+    ) -> Dict[str, List[str]]:
+        return _get_regions_to_zones(session=session, regions=regions)
+
+    def _get_vpc_id_subnet_id_or_error_cache_key(
+        self,
+        ec2_client: botocore.client.BaseClient,
+        config: AWSConfig,
+        region: str,
+        allocate_public_ip: bool,
+        availability_zones: Optional[List[str]] = None,
+    ) -> tuple:
+        return hashkey(
+            region, allocate_public_ip, tuple(availability_zones) if availability_zones else None
+        )
+
+    @cachedmethod(
+        cache=lambda self: self._get_vpc_id_subnet_id_or_error_cache,
+        key=_get_vpc_id_subnet_id_or_error_cache_key,
+        lock=lambda self: self._get_vpc_id_subnet_id_or_error_cache_lock,
+    )
+    def _get_vpc_id_subnet_id_or_error(
+        self,
+        ec2_client: botocore.client.BaseClient,
+        config: AWSConfig,
+        region: str,
+        allocate_public_ip: bool,
+        availability_zones: Optional[List[str]] = None,
+    ) -> Tuple[str, List[str]]:
+        return get_vpc_id_subnet_id_or_error(
+            ec2_client=ec2_client,
+            config=config,
+            region=region,
+            allocate_public_ip=allocate_public_ip,
+            availability_zones=availability_zones,
+        )
+
+    @cachedmethod(
+        cache=lambda self: self._get_maximum_efa_interfaces_cache,
+        key=_ec2client_cache_methodkey,
+        lock=lambda self: self._get_maximum_efa_interfaces_cache_lock,
+    )
+    def _get_maximum_efa_interfaces(
+        self,
+        ec2_client: botocore.client.BaseClient,
+        region: str,
+        instance_type: str,
+    ) -> int:
+        return _get_maximum_efa_interfaces(
+            ec2_client=ec2_client,
+            instance_type=instance_type,
+        )
+
+    def _get_subnets_availability_zones_key(
+        self,
+        ec2_client: botocore.client.BaseClient,
+        region: str,
+        subnet_ids: List[str],
+    ) -> tuple:
+        return hashkey(region, tuple(subnet_ids))
+
+    @cachedmethod(
+        cache=lambda self: self._get_subnets_availability_zones_cache,
+        key=_get_subnets_availability_zones_key,
+        lock=lambda self: self._get_subnets_availability_zones_cache_lock,
+    )
+    def _get_subnets_availability_zones(
+        self,
+        ec2_client: botocore.client.BaseClient,
+        region: str,
+        subnet_ids: List[str],
+    ) -> Dict[str, str]:
+        return aws_resources.get_subnets_availability_zones(
+            ec2_client=ec2_client,
+            subnet_ids=subnet_ids,
+        )
+
+    @cachedmethod(
+        cache=lambda self: self._create_security_group_cache,
+        key=_ec2client_cache_methodkey,
+        lock=lambda self: self._create_security_group_cache_lock,
+    )
+    def _create_security_group(
+        self,
+        ec2_client: botocore.client.BaseClient,
+        region: str,
+        project_id: str,
+        vpc_id: Optional[str],
+    ) -> str:
+        return aws_resources.create_security_group(
+            ec2_client=ec2_client,
+            project_id=project_id,
+            vpc_id=vpc_id,
+        )
+
+    def _get_image_id_and_username_cache_key(
+        self,
+        ec2_client: botocore.client.BaseClient,
+        region: str,
+        cuda: bool,
+        instance_type: str,
+        image_config: Optional[AWSOSImageConfig] = None,
+    ) -> tuple:
+        return hashkey(region, cuda, instance_type, image_config.json() if image_config else None)
+
+    @cachedmethod(
+        cache=lambda self: self._get_image_id_and_username_cache,
+        key=_get_image_id_and_username_cache_key,
+        lock=lambda self: self._get_image_id_and_username_cache_lock,
+    )
+    def _get_image_id_and_username(
+        self,
+        ec2_client: botocore.client.BaseClient,
+        region: str,
+        cuda: bool,
+        instance_type: str,
+        image_config: Optional[AWSOSImageConfig] = None,
+    ) -> tuple[str, str]:
+        return aws_resources.get_image_id_and_username(
+            ec2_client=ec2_client,
+            cuda=cuda,
+            instance_type=instance_type,
+            image_config=image_config,
+        )
+
 
 def get_vpc_id_subnet_id_or_error(
     ec2_client: botocore.client.BaseClient,
@@ -798,7 +1008,7 @@ def _get_regions_to_quotas(
         return region_quotas
 
     regions_to_quotas = {}
-    with ThreadPoolExecutor(max_workers=8) as executor:
+    with ThreadPoolExecutor(max_workers=12) as executor:
         future_to_region = {}
         for region in regions:
             future = executor.submit(
@@ -823,7 +1033,7 @@ def _has_quota(quotas: Dict[str, int], instance_name: str) -> Optional[bool]:
 
 def _get_regions_to_zones(session: boto3.Session, regions: List[str]) -> Dict[str, List[str]]:
     regions_to_zones = {}
-    with ThreadPoolExecutor(max_workers=8) as executor:
+    with ThreadPoolExecutor(max_workers=12) as executor:
         future_to_region = {}
         for region in regions:
             future = executor.submit(
@@ -862,6 +1072,15 @@ def _supported_instances(offer: InstanceOffer) -> bool:
     return False
 
 
+def _offer_supports_placement_group(offer: InstanceOffer, placement_group: PlacementGroup) -> bool:
+    if placement_group.configuration.placement_strategy != PlacementStrategy.CLUSTER:
+        return True
+    for family in ["t3.", "t2."]:
+        if offer.instance.name.startswith(family):
+            return False
+    return True
+
+
 def _get_maximum_efa_interfaces(ec2_client: botocore.client.BaseClient, instance_type: str) -> int:
     try:
         response = ec2_client.describe_instance_types(

dstack/_internal/core/backends/base/compute.py

@@ -57,7 +57,7 @@ class Compute(ABC):
 
     def __init__(self):
         self._offers_cache_lock = threading.Lock()
-        self._offers_cache = TTLCache(maxsize=5, ttl=30)
+        self._offers_cache = TTLCache(maxsize=10, ttl=180)
 
     @abstractmethod
     def get_offers(
@@ -559,7 +559,8 @@ def get_shim_commands(
     backend_shim_env: Optional[Dict[str, str]] = None,
     arch: Optional[str] = None,
 ) -> List[str]:
-    commands = get_shim_pre_start_commands(
+    commands = get_setup_cloud_instance_commands()
+    commands += get_shim_pre_start_commands(
         base_path=base_path,
         bin_path=bin_path,
         arch=arch,
@@ -641,6 +642,23 @@ def get_dstack_shim_download_url(arch: Optional[str] = None) -> str:
     return url_template.format(version=version, arch=arch)
 
 
+def get_setup_cloud_instance_commands() -> list[str]:
+    return [
+        # Workaround for https://github.com/NVIDIA/nvidia-container-toolkit/issues/48
+        # Attempts to patch /etc/docker/daemon.json while keeping any custom settings it may have.
+        (
+            "/bin/sh -c '"  # wrap in /bin/sh to avoid interfering with other cloud init commands
+            " grep -q nvidia /etc/docker/daemon.json"
+            " && ! grep -q native.cgroupdriver /etc/docker/daemon.json"
+            " && jq '\\''.\"exec-opts\" = ((.\"exec-opts\" // []) + [\"native.cgroupdriver=cgroupfs\"])'\\'' /etc/docker/daemon.json > /tmp/daemon.json"
+            " && sudo mv /tmp/daemon.json /etc/docker/daemon.json"
+            " && sudo service docker restart"
+            " || true"
+            "'"
+        ),
+    ]
+
+
 def get_shim_pre_start_commands(
     base_path: Optional[PathLike] = None,
     bin_path: Optional[PathLike] = None,

dstack/_internal/core/backends/cudo/compute.py

@@ -65,12 +65,13 @@ class CudoCompute(
         public_keys = instance_config.get_public_keys()
         memory_size = round(instance_offer.instance.resources.memory_mib / 1024)
         disk_size = round(instance_offer.instance.resources.disk.size_mib / 1024)
-        commands = get_shim_commands(authorized_keys=public_keys)
         gpus_no = len(instance_offer.instance.resources.gpus)
-        shim_commands = " ".join([" && ".join(commands)])
-        startup_script = (
-            shim_commands if gpus_no > 0 else f"{install_docker_script()} && {shim_commands}"
-        )
+        if gpus_no > 0:
+            # we'll need jq for patching /etc/docker/daemon.json, see get_shim_commands()
+            commands = install_jq_commands()
+        else:
+            commands = install_docker_commands()
+        commands += get_shim_commands(authorized_keys=public_keys)
 
         try:
             resp_data = self.api_client.create_virtual_machine(
@@ -85,7 +86,7 @@ class CudoCompute(
                 memory_gib=memory_size,
                 vcpus=instance_offer.instance.resources.cpus,
                 vm_id=vm_id,
-                start_script=startup_script,
+                start_script=" && ".join(commands),
                 password=None,
                 customSshKeys=public_keys,
             )
@@ -151,6 +152,19 @@ def _get_image_id(cuda: bool) -> str:
     return image_name
 
 
-def install_docker_script():
-    commands = 'export DEBIAN_FRONTEND="noninteractive" && mkdir -p /etc/apt/keyrings && curl --max-time 60 -fsSL https://download.docker.com/linux/ubuntu/gpg | gpg --dearmor -o /etc/apt/keyrings/docker.gpg && echo "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.gpg] https://download.docker.com/linux/ubuntu $(lsb_release -cs) stable" | tee /etc/apt/sources.list.d/docker.list > /dev/null && apt-get update && apt-get --assume-yes install docker-ce docker-ce-cli containerd.io docker-compose-plugin'
-    return commands
+def install_jq_commands():
+    return [
+        "export DEBIAN_FRONTEND=noninteractive",
+        "apt-get --assume-yes install jq",
+    ]
+
+
+def install_docker_commands():
+    return [
+        "export DEBIAN_FRONTEND=noninteractive",
+        "mkdir -p /etc/apt/keyrings",
+        "curl --max-time 60 -fsSL https://download.docker.com/linux/ubuntu/gpg | gpg --dearmor -o /etc/apt/keyrings/docker.gpg",
+        'echo "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.gpg] https://download.docker.com/linux/ubuntu $(lsb_release -cs) stable" | tee /etc/apt/sources.list.d/docker.list > /dev/null',
+        "apt-get update",
+        "apt-get --assume-yes install docker-ce docker-ce-cli containerd.io docker-compose-plugin",
+    ]