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

dstack/_internal/server/schemas/runs.py

@@ -9,12 +9,24 @@ from dstack._internal.core.models.runs import ApplyRunPlanInput, RunSpec
 
 
 class ListRunsRequest(CoreModel):
-    project_name: Optional[str]
-    repo_id: Optional[str]
-    username: Optional[str]
+    project_name: Optional[str] = None
+    repo_id: Optional[str] = None
+    username: Optional[str] = None
     only_active: bool = False
-    prev_submitted_at: Optional[datetime]
-    prev_run_id: Optional[UUID]
+    include_jobs: bool = Field(
+        True,
+        description=("Whether to include `jobs` in the response"),
+    )
+    job_submissions_limit: Optional[int] = Field(
+        None,
+        ge=0,
+        description=(
+            "Limit number of job submissions returned per job to avoid large responses."
+            "Drops older job submissions. No effect with `include_jobs: false`"
+        ),
+    )
+    prev_submitted_at: Optional[datetime] = None
+    prev_run_id: Optional[UUID] = None
     limit: int = Field(100, ge=0, le=100)
     ascending: bool = False
 

dstack/_internal/server/services/fleets.py

@@ -1,6 +1,8 @@
 import uuid
+from collections.abc import Callable
 from datetime import datetime, timezone
-from typing import List, Literal, Optional, Tuple, Union, cast
+from functools import wraps
+from typing import List, Literal, Optional, Tuple, TypeVar, Union, cast
 
 from sqlalchemy import and_, func, or_, select
 from sqlalchemy.ext.asyncio import AsyncSession
@@ -13,10 +15,12 @@ from dstack._internal.core.errors import (
     ResourceExistsError,
     ServerClientError,
 )
+from dstack._internal.core.models.common import ApplyAction, CoreModel
 from dstack._internal.core.models.envs import Env
 from dstack._internal.core.models.fleets import (
     ApplyFleetPlanInput,
     Fleet,
+    FleetConfiguration,
     FleetPlan,
     FleetSpec,
     FleetStatus,
@@ -40,6 +44,7 @@ from dstack._internal.core.models.resources import ResourcesSpec
 from dstack._internal.core.models.runs import Requirements, get_policy_map
 from dstack._internal.core.models.users import GlobalRole
 from dstack._internal.core.services import validate_dstack_resource_name
+from dstack._internal.core.services.diff import ModelDiff, copy_model, diff_models
 from dstack._internal.server.db import get_db
 from dstack._internal.server.models import (
     FleetModel,
@@ -49,7 +54,10 @@ from dstack._internal.server.models import (
 )
 from dstack._internal.server.services import instances as instances_services
 from dstack._internal.server.services import offers as offers_services
-from dstack._internal.server.services.instances import list_active_remote_instances
+from dstack._internal.server.services.instances import (
+    get_instance_remote_connection_info,
+    list_active_remote_instances,
+)
 from dstack._internal.server.services.locking import (
     get_locker,
     string_to_lock_id,
@@ -178,8 +186,9 @@ async def list_project_fleet_models(
 async def get_fleet(
     session: AsyncSession,
     project: ProjectModel,
-    name: Optional[str],
-    fleet_id: Optional[uuid.UUID],
+    name: Optional[str] = None,
+    fleet_id: Optional[uuid.UUID] = None,
+    include_sensitive: bool = False,
 ) -> Optional[Fleet]:
     if fleet_id is not None:
         fleet_model = await get_project_fleet_model_by_id(
@@ -193,7 +202,7 @@ async def get_fleet(
         raise ServerClientError("name or id must be specified")
     if fleet_model is None:
         return None
-    return fleet_model_to_fleet(fleet_model)
+    return fleet_model_to_fleet(fleet_model, include_sensitive=include_sensitive)
 
 
 async def get_project_fleet_model_by_id(
@@ -236,23 +245,32 @@ async def get_plan(
     spec: FleetSpec,
 ) -> FleetPlan:
     # Spec must be copied by parsing to calculate merged_profile
-    effective_spec = FleetSpec.parse_obj(spec.dict())
+    effective_spec = copy_model(spec)
     effective_spec = await apply_plugin_policies(
         user=user.name,
         project=project.name,
         spec=effective_spec,
     )
-    effective_spec = FleetSpec.parse_obj(effective_spec.dict())
-    _validate_fleet_spec_and_set_defaults(spec)
+    # Spec must be copied by parsing to calculate merged_profile
+    effective_spec = copy_model(effective_spec)
+    _validate_fleet_spec_and_set_defaults(effective_spec)
+
+    action = ApplyAction.CREATE
     current_fleet: Optional[Fleet] = None
     current_fleet_id: Optional[uuid.UUID] = None
+
     if effective_spec.configuration.name is not None:
-        current_fleet_model = await get_project_fleet_model_by_name(
-            session=session, project=project, name=effective_spec.configuration.name
+        current_fleet = await get_fleet(
+            session=session,
+            project=project,
+            name=effective_spec.configuration.name,
+            include_sensitive=True,
         )
-        if current_fleet_model is not None:
-            current_fleet = fleet_model_to_fleet(current_fleet_model)
-            current_fleet_id = current_fleet_model.id
+        if current_fleet is not None:
+            _set_fleet_spec_defaults(current_fleet.spec)
+            if _can_update_fleet_spec(current_fleet.spec, effective_spec):
+                action = ApplyAction.UPDATE
+            current_fleet_id = current_fleet.id
     await _check_ssh_hosts_not_yet_added(session, effective_spec, current_fleet_id)
 
     offers = []
@@ -265,7 +283,10 @@ async def get_plan(
             blocks=effective_spec.configuration.blocks,
         )
         offers = [offer for _, offer in offers_with_backends]
+
     _remove_fleet_spec_sensitive_info(effective_spec)
+    if current_fleet is not None:
+        _remove_fleet_spec_sensitive_info(current_fleet.spec)
     plan = FleetPlan(
         project_name=project.name,
         user=user.name,
@@ -275,6 +296,7 @@ async def get_plan(
         offers=offers[:50],
         total_offers=len(offers),
         max_offer_price=max((offer.price for offer in offers), default=None),
+        action=action,
     )
     return plan
 
@@ -327,11 +349,77 @@ async def apply_plan(
     plan: ApplyFleetPlanInput,
     force: bool,
 ) -> Fleet:
-    return await create_fleet(
+    spec = await apply_plugin_policies(
+        user=user.name,
+        project=project.name,
+        spec=plan.spec,
+    )
+    # Spec must be copied by parsing to calculate merged_profile
+    spec = copy_model(spec)
+    _validate_fleet_spec_and_set_defaults(spec)
+
+    if spec.configuration.ssh_config is not None:
+        _check_can_manage_ssh_fleets(user=user, project=project)
+
+    configuration = spec.configuration
+    if configuration.name is None:
+        return await _create_fleet(
+            session=session,
+            project=project,
+            user=user,
+            spec=spec,
+        )
+
+    fleet_model = await get_project_fleet_model_by_name(
+        session=session,
+        project=project,
+        name=configuration.name,
+    )
+    if fleet_model is None:
+        return await _create_fleet(
+            session=session,
+            project=project,
+            user=user,
+            spec=spec,
+        )
+
+    instances_ids = sorted(i.id for i in fleet_model.instances if not i.deleted)
+    await session.commit()
+    async with (
+        get_locker(get_db().dialect_name).lock_ctx(FleetModel.__tablename__, [fleet_model.id]),
+        get_locker(get_db().dialect_name).lock_ctx(InstanceModel.__tablename__, instances_ids),
+    ):
+        # Refetch after lock
+        # TODO: Lock instances with FOR UPDATE?
+        res = await session.execute(
+            select(FleetModel)
+            .where(
+                FleetModel.project_id == project.id,
+                FleetModel.id == fleet_model.id,
+                FleetModel.deleted == False,
+            )
+            .options(selectinload(FleetModel.instances))
+            .options(selectinload(FleetModel.runs))
+            .execution_options(populate_existing=True)
+            .order_by(FleetModel.id)  # take locks in order
+            .with_for_update(key_share=True)
+        )
+        fleet_model = res.scalars().unique().one_or_none()
+        if fleet_model is not None:
+            return await _update_fleet(
+                session=session,
+                project=project,
+                spec=spec,
+                current_resource=plan.current_resource,
+                force=force,
+                fleet_model=fleet_model,
+            )
+
+    return await _create_fleet(
         session=session,
         project=project,
         user=user,
-        spec=plan.spec,
+        spec=spec,
     )
 
 
@@ -341,73 +429,19 @@ async def create_fleet(
     user: UserModel,
     spec: FleetSpec,
 ) -> Fleet:
-    # Spec must be copied by parsing to calculate merged_profile
     spec = await apply_plugin_policies(
         user=user.name,
         project=project.name,
         spec=spec,
     )
-    spec = FleetSpec.parse_obj(spec.dict())
+    # Spec must be copied by parsing to calculate merged_profile
+    spec = copy_model(spec)
     _validate_fleet_spec_and_set_defaults(spec)
 
     if spec.configuration.ssh_config is not None:
         _check_can_manage_ssh_fleets(user=user, project=project)
 
-    lock_namespace = f"fleet_names_{project.name}"
-    if get_db().dialect_name == "sqlite":
-        # Start new transaction to see committed changes after lock
-        await session.commit()
-    elif get_db().dialect_name == "postgresql":
-        await session.execute(
-            select(func.pg_advisory_xact_lock(string_to_lock_id(lock_namespace)))
-        )
-
-    lock, _ = get_locker().get_lockset(lock_namespace)
-    async with lock:
-        if spec.configuration.name is not None:
-            fleet_model = await get_project_fleet_model_by_name(
-                session=session,
-                project=project,
-                name=spec.configuration.name,
-            )
-            if fleet_model is not None:
-                raise ResourceExistsError()
-        else:
-            spec.configuration.name = await generate_fleet_name(session=session, project=project)
-
-        fleet_model = FleetModel(
-            id=uuid.uuid4(),
-            name=spec.configuration.name,
-            project=project,
-            status=FleetStatus.ACTIVE,
-            spec=spec.json(),
-            instances=[],
-        )
-        session.add(fleet_model)
-        if spec.configuration.ssh_config is not None:
-            for i, host in enumerate(spec.configuration.ssh_config.hosts):
-                instances_model = await create_fleet_ssh_instance_model(
-                    project=project,
-                    spec=spec,
-                    ssh_params=spec.configuration.ssh_config,
-                    env=spec.configuration.env,
-                    instance_num=i,
-                    host=host,
-                )
-                fleet_model.instances.append(instances_model)
-        else:
-            for i in range(_get_fleet_nodes_to_provision(spec)):
-                instance_model = await create_fleet_instance_model(
-                    session=session,
-                    project=project,
-                    user=user,
-                    spec=spec,
-                    reservation=spec.configuration.reservation,
-                    instance_num=i,
-                )
-                fleet_model.instances.append(instance_model)
-        await session.commit()
-        return fleet_model_to_fleet(fleet_model)
+    return await _create_fleet(session=session, project=project, user=user, spec=spec)
 
 
 async def create_fleet_instance_model(
@@ -516,11 +550,12 @@ async def delete_fleets(
     await session.commit()
     logger.info("Deleting fleets: %s", [v.name for v in fleet_models])
     async with (
-        get_locker().lock_ctx(FleetModel.__tablename__, fleets_ids),
-        get_locker().lock_ctx(InstanceModel.__tablename__, instances_ids),
+        get_locker(get_db().dialect_name).lock_ctx(FleetModel.__tablename__, fleets_ids),
+        get_locker(get_db().dialect_name).lock_ctx(InstanceModel.__tablename__, instances_ids),
     ):
         # Refetch after lock
-        # TODO lock instances with FOR UPDATE?
+        # TODO: Lock instances with FOR UPDATE?
+        # TODO: Do not lock fleet when deleting only instances
         res = await session.execute(
             select(FleetModel)
             .where(
@@ -599,6 +634,235 @@ def is_fleet_empty(fleet_model: FleetModel) -> bool:
     return len(active_instances) == 0
 
 
+async def _create_fleet(
+    session: AsyncSession,
+    project: ProjectModel,
+    user: UserModel,
+    spec: FleetSpec,
+) -> Fleet:
+    lock_namespace = f"fleet_names_{project.name}"
+    if get_db().dialect_name == "sqlite":
+        # Start new transaction to see committed changes after lock
+        await session.commit()
+    elif get_db().dialect_name == "postgresql":
+        await session.execute(
+            select(func.pg_advisory_xact_lock(string_to_lock_id(lock_namespace)))
+        )
+
+    lock, _ = get_locker(get_db().dialect_name).get_lockset(lock_namespace)
+    async with lock:
+        if spec.configuration.name is not None:
+            fleet_model = await get_project_fleet_model_by_name(
+                session=session,
+                project=project,
+                name=spec.configuration.name,
+            )
+            if fleet_model is not None:
+                raise ResourceExistsError()
+        else:
+            spec.configuration.name = await generate_fleet_name(session=session, project=project)
+
+        fleet_model = FleetModel(
+            id=uuid.uuid4(),
+            name=spec.configuration.name,
+            project=project,
+            status=FleetStatus.ACTIVE,
+            spec=spec.json(),
+            instances=[],
+        )
+        session.add(fleet_model)
+        if spec.configuration.ssh_config is not None:
+            for i, host in enumerate(spec.configuration.ssh_config.hosts):
+                instances_model = await create_fleet_ssh_instance_model(
+                    project=project,
+                    spec=spec,
+                    ssh_params=spec.configuration.ssh_config,
+                    env=spec.configuration.env,
+                    instance_num=i,
+                    host=host,
+                )
+                fleet_model.instances.append(instances_model)
+        else:
+            for i in range(_get_fleet_nodes_to_provision(spec)):
+                instance_model = await create_fleet_instance_model(
+                    session=session,
+                    project=project,
+                    user=user,
+                    spec=spec,
+                    reservation=spec.configuration.reservation,
+                    instance_num=i,
+                )
+                fleet_model.instances.append(instance_model)
+        await session.commit()
+        return fleet_model_to_fleet(fleet_model)
+
+
+async def _update_fleet(
+    session: AsyncSession,
+    project: ProjectModel,
+    spec: FleetSpec,
+    current_resource: Optional[Fleet],
+    force: bool,
+    fleet_model: FleetModel,
+) -> Fleet:
+    fleet = fleet_model_to_fleet(fleet_model)
+    _set_fleet_spec_defaults(fleet.spec)
+    fleet_sensitive = fleet_model_to_fleet(fleet_model, include_sensitive=True)
+    _set_fleet_spec_defaults(fleet_sensitive.spec)
+
+    if not force:
+        if current_resource is not None:
+            _set_fleet_spec_defaults(current_resource.spec)
+        if (
+            current_resource is None
+            or current_resource.id != fleet.id
+            or current_resource.spec != fleet.spec
+        ):
+            raise ServerClientError(
+                "Failed to apply plan. Resource has been changed. Try again or use force apply."
+            )
+
+    _check_can_update_fleet_spec(fleet_sensitive.spec, spec)
+
+    spec_json = spec.json()
+    fleet_model.spec = spec_json
+
+    if (
+        fleet_sensitive.spec.configuration.ssh_config is not None
+        and spec.configuration.ssh_config is not None
+    ):
+        added_hosts, removed_hosts, changed_hosts = _calculate_ssh_hosts_changes(
+            current=fleet_sensitive.spec.configuration.ssh_config.hosts,
+            new=spec.configuration.ssh_config.hosts,
+        )
+        # `_check_can_update_fleet_spec` ensures hosts are not changed
+        assert not changed_hosts, changed_hosts
+        active_instance_nums: set[int] = set()
+        removed_instance_nums: list[int] = []
+        if removed_hosts or added_hosts:
+            for instance_model in fleet_model.instances:
+                if instance_model.deleted:
+                    continue
+                active_instance_nums.add(instance_model.instance_num)
+                rci = get_instance_remote_connection_info(instance_model)
+                if rci is None:
+                    logger.error(
+                        "Cloud instance %s in SSH fleet %s",
+                        instance_model.id,
+                        fleet_model.id,
+                    )
+                    continue
+                if rci.host in removed_hosts:
+                    removed_instance_nums.append(instance_model.instance_num)
+        if added_hosts:
+            await _check_ssh_hosts_not_yet_added(session, spec, fleet.id)
+            for host in added_hosts.values():
+                instance_num = _get_next_instance_num(active_instance_nums)
+                instance_model = await create_fleet_ssh_instance_model(
+                    project=project,
+                    spec=spec,
+                    ssh_params=spec.configuration.ssh_config,
+                    env=spec.configuration.env,
+                    instance_num=instance_num,
+                    host=host,
+                )
+                fleet_model.instances.append(instance_model)
+                active_instance_nums.add(instance_num)
+        if removed_instance_nums:
+            _terminate_fleet_instances(fleet_model, removed_instance_nums)
+
+    await session.commit()
+    return fleet_model_to_fleet(fleet_model)
+
+
+def _can_update_fleet_spec(current_fleet_spec: FleetSpec, new_fleet_spec: FleetSpec) -> bool:
+    try:
+        _check_can_update_fleet_spec(current_fleet_spec, new_fleet_spec)
+    except ServerClientError as e:
+        logger.debug("Run cannot be updated: %s", repr(e))
+        return False
+    return True
+
+
+M = TypeVar("M", bound=CoreModel)
+
+
+def _check_can_update(*updatable_fields: str):
+    def decorator(fn: Callable[[M, M, ModelDiff], None]) -> Callable[[M, M], None]:
+        @wraps(fn)
+        def inner(current: M, new: M):
+            diff = _check_can_update_inner(current, new, updatable_fields)
+            fn(current, new, diff)
+
+        return inner
+
+    return decorator
+
+
+def _check_can_update_inner(current: M, new: M, updatable_fields: tuple[str, ...]) -> ModelDiff:
+    diff = diff_models(current, new)
+    changed_fields = diff.keys()
+    if not (changed_fields <= set(updatable_fields)):
+        raise ServerClientError(
+            f"Failed to update fields {list(changed_fields)}."
+            f" Can only update {list(updatable_fields)}."
+        )
+    return diff
+
+
+@_check_can_update("configuration", "configuration_path")
+def _check_can_update_fleet_spec(current: FleetSpec, new: FleetSpec, diff: ModelDiff):
+    if "configuration" in diff:
+        _check_can_update_fleet_configuration(current.configuration, new.configuration)
+
+
+@_check_can_update("ssh_config")
+def _check_can_update_fleet_configuration(
+    current: FleetConfiguration, new: FleetConfiguration, diff: ModelDiff
+):
+    if "ssh_config" in diff:
+        current_ssh_config = current.ssh_config
+        new_ssh_config = new.ssh_config
+        if current_ssh_config is None:
+            if new_ssh_config is not None:
+                raise ServerClientError("Fleet type changed from Cloud to SSH, cannot update")
+        elif new_ssh_config is None:
+            raise ServerClientError("Fleet type changed from SSH to Cloud, cannot update")
+        else:
+            _check_can_update_ssh_config(current_ssh_config, new_ssh_config)
+
+
+@_check_can_update("hosts")
+def _check_can_update_ssh_config(current: SSHParams, new: SSHParams, diff: ModelDiff):
+    if "hosts" in diff:
+        _, _, changed_hosts = _calculate_ssh_hosts_changes(current.hosts, new.hosts)
+        if changed_hosts:
+            raise ServerClientError(
+                f"Hosts configuration changed, cannot update: {list(changed_hosts)}"
+            )
+
+
+def _calculate_ssh_hosts_changes(
+    current: list[Union[SSHHostParams, str]], new: list[Union[SSHHostParams, str]]
+) -> tuple[dict[str, Union[SSHHostParams, str]], set[str], set[str]]:
+    current_hosts = {h if isinstance(h, str) else h.hostname: h for h in current}
+    new_hosts = {h if isinstance(h, str) else h.hostname: h for h in new}
+    added_hosts = {h: new_hosts[h] for h in new_hosts.keys() - current_hosts}
+    removed_hosts = current_hosts.keys() - new_hosts
+    changed_hosts: set[str] = set()
+    for host in current_hosts.keys() & new_hosts:
+        current_host = current_hosts[host]
+        new_host = new_hosts[host]
+        if isinstance(current_host, str) or isinstance(new_host, str):
+            if current_host != new_host:
+                changed_hosts.add(host)
+        elif diff_models(
+            current_host, new_host, reset={"identity_file": True, "proxy_jump": {"identity_file"}}
+        ):
+            changed_hosts.add(host)
+    return added_hosts, removed_hosts, changed_hosts
+
+
 def _check_can_manage_ssh_fleets(user: UserModel, project: ProjectModel):
     if user.global_role == GlobalRole.ADMIN:
         return
@@ -653,6 +917,8 @@ def _validate_fleet_spec_and_set_defaults(spec: FleetSpec):
         validate_dstack_resource_name(spec.configuration.name)
     if spec.configuration.ssh_config is None and spec.configuration.nodes is None:
         raise ServerClientError("No ssh_config or nodes specified")
+    if spec.configuration.ssh_config is not None and spec.configuration.nodes is not None:
+        raise ServerClientError("ssh_config and nodes are mutually exclusive")
     if spec.configuration.ssh_config is not None:
         _validate_all_ssh_params_specified(spec.configuration.ssh_config)
         if spec.configuration.ssh_config.ssh_key is not None:
@@ -661,6 +927,10 @@ def _validate_fleet_spec_and_set_defaults(spec: FleetSpec):
             if isinstance(host, SSHHostParams) and host.ssh_key is not None:
                 _validate_ssh_key(host.ssh_key)
         _validate_internal_ips(spec.configuration.ssh_config)
+    _set_fleet_spec_defaults(spec)
+
+
+def _set_fleet_spec_defaults(spec: FleetSpec):
     if spec.configuration.resources is not None:
         set_resources_defaults(spec.configuration.resources)
 
@@ -733,3 +1003,16 @@ def _get_fleet_requirements(fleet_spec: FleetSpec) -> Requirements:
         reservation=fleet_spec.configuration.reservation,
     )
     return requirements
+
+
+def _get_next_instance_num(instance_nums: set[int]) -> int:
+    if not instance_nums:
+        return 0
+    min_instance_num = min(instance_nums)
+    if min_instance_num > 0:
+        return 0
+    instance_num = min_instance_num + 1
+    while True:
+        if instance_num not in instance_nums:
+            return instance_num
+        instance_num += 1

dstack/_internal/server/services/gateways/__init__.py

@@ -2,6 +2,7 @@ import asyncio
 import datetime
 import uuid
 from datetime import timedelta, timezone
+from functools import partial
 from typing import List, Optional, Sequence
 
 import httpx
@@ -162,7 +163,7 @@ async def create_gateway(
             select(func.pg_advisory_xact_lock(string_to_lock_id(lock_namespace)))
         )
 
-    lock, _ = get_locker().get_lockset(lock_namespace)
+    lock, _ = get_locker(get_db().dialect_name).get_lockset(lock_namespace)
     async with lock:
         if configuration.name is None:
             configuration.name = await generate_gateway_name(session=session, project=project)
@@ -186,6 +187,7 @@ async def create_gateway(
         return gateway_model_to_gateway(gateway)
 
 
+# NOTE: dstack Sky imports and uses this function
 async def connect_to_gateway_with_retry(
     gateway_compute: GatewayComputeModel,
 ) -> Optional[GatewayConnection]:
@@ -229,7 +231,9 @@ async def delete_gateways(
     gateways_ids = sorted([g.id for g in gateway_models])
     await session.commit()
     logger.info("Deleting gateways: %s", [g.name for g in gateway_models])
-    async with get_locker().lock_ctx(GatewayModel.__tablename__, gateways_ids):
+    async with get_locker(get_db().dialect_name).lock_ctx(
+        GatewayModel.__tablename__, gateways_ids
+    ):
         # Refetch after lock
         res = await session.execute(
             select(GatewayModel)
@@ -378,6 +382,8 @@ async def get_or_add_gateway_connection(
 async def init_gateways(session: AsyncSession):
     res = await session.execute(
         select(GatewayComputeModel).where(
+            # FIXME: should not include computes related to gateways in the `provisioning` status.
+            # Causes warnings and delays when restarting the server during gateway provisioning.
             GatewayComputeModel.active == True,
             GatewayComputeModel.deleted == False,
         )
@@ -419,7 +425,8 @@ async def init_gateways(session: AsyncSession):
 
         for gateway_compute, error in await gather_map_async(
             await gateway_connections_pool.all(),
-            configure_gateway,
+            # Need several attempts to handle short gateway downtime after update
+            partial(configure_gateway, attempts=7),
             return_exceptions=True,
         ):
             if isinstance(error, Exception):
@@ -459,7 +466,11 @@ def _recently_updated(gateway_compute_model: GatewayComputeModel) -> bool:
     ) > get_current_datetime() - timedelta(seconds=60)
 
 
-async def configure_gateway(connection: GatewayConnection) -> None:
+# NOTE: dstack Sky imports and uses this function
+async def configure_gateway(
+    connection: GatewayConnection,
+    attempts: int = GATEWAY_CONFIGURE_ATTEMPTS,
+) -> None:
     """
     Try submitting gateway config several times in case gateway's HTTP server is not
     running yet
@@ -467,7 +478,7 @@ async def configure_gateway(connection: GatewayConnection) -> None:
 
     logger.debug("Configuring gateway %s", connection.ip_address)
 
-    for attempt in range(GATEWAY_CONFIGURE_ATTEMPTS - 1):
+    for attempt in range(attempts - 1):
         try:
             async with connection.client() as client:
                 await client.submit_gateway_config()
@@ -476,7 +487,7 @@ async def configure_gateway(connection: GatewayConnection) -> None:
             logger.debug(
                 "Failed attempt %s/%s at configuring gateway %s: %r",
                 attempt + 1,
-                GATEWAY_CONFIGURE_ATTEMPTS,
+                attempts,
                 connection.ip_address,
                 e,
             )