dstack 0.19.25rc1__py3-none-any.whl → 0.19.26__py3-none-any.whl

dstack/_internal/server/background/tasks/process_submitted_jobs.py

@@ -5,9 +5,9 @@ import uuid
 from datetime import datetime, timedelta
 from typing import List, Optional, Tuple
 
-from sqlalchemy import and_, or_, select
+from sqlalchemy import and_, func, not_, or_, select
 from sqlalchemy.ext.asyncio import AsyncSession
-from sqlalchemy.orm import contains_eager, joinedload, load_only, selectinload
+from sqlalchemy.orm import contains_eager, joinedload, load_only, noload, selectinload
 
 from dstack._internal.core.backends.base.backend import Backend
 from dstack._internal.core.backends.base.compute import ComputeWithVolumeSupport
@@ -53,6 +53,7 @@ from dstack._internal.server.models import (
 from dstack._internal.server.services.backends import get_project_backend_by_type_or_error
 from dstack._internal.server.services.fleets import (
     fleet_model_to_fleet,
+    get_fleet_requirements,
 )
 from dstack._internal.server.services.instances import (
     filter_pool_instances,
@@ -71,6 +72,10 @@ from dstack._internal.server.services.jobs import (
 from dstack._internal.server.services.locking import get_locker
 from dstack._internal.server.services.logging import fmt
 from dstack._internal.server.services.offers import get_offers_by_requirements
+from dstack._internal.server.services.requirements.combine import (
+    combine_fleet_and_run_profiles,
+    combine_fleet_and_run_requirements,
+)
 from dstack._internal.server.services.runs import (
     check_run_spec_requires_instance_mounts,
     run_model_to_run,
@@ -148,8 +153,8 @@ async def _process_next_submitted_job():
             if job_model is None:
                 return
             lockset.add(job_model.id)
+        job_model_id = job_model.id
         try:
-            job_model_id = job_model.id
             await _process_submitted_job(session=session, job_model=job_model)
         finally:
             lockset.difference_update([job_model_id])
@@ -245,8 +250,8 @@ async def _process_submitted_job(session: AsyncSession, job_model: JobModel):
         ]
         if run_model.fleet is not None:
             fleet_filters.append(FleetModel.id == run_model.fleet_id)
-        if run_spec.configuration.fleets is not None:
-            fleet_filters.append(FleetModel.name.in_(run_spec.configuration.fleets))
+        if run_spec.merged_profile.fleets is not None:
+            fleet_filters.append(FleetModel.name.in_(run_spec.merged_profile.fleets))
 
         instance_filters = [
             InstanceModel.deleted == False,
@@ -264,9 +269,6 @@ async def _process_submitted_job(session: AsyncSession, job_model: JobModel):
                 [i.id for i in f.instances] for f in fleet_models_with_instances
             )
         )
-        fleet_models = fleet_models_with_instances + fleet_models_without_instances
-        fleets_ids = [f.id for f in fleet_models]
-
         if get_db().dialect_name == "sqlite":
             # Start new transaction to see committed changes after lock
             await session.commit()
@@ -275,13 +277,15 @@ async def _process_submitted_job(session: AsyncSession, job_model: JobModel):
             InstanceModel.__tablename__, instances_ids
         ):
             if get_db().dialect_name == "sqlite":
-                fleet_models = await _refetch_fleet_models(
+                fleets_with_instances_ids = [f.id for f in fleet_models_with_instances]
+                fleet_models_with_instances = await _refetch_fleet_models_with_instances(
                     session=session,
-                    fleets_ids=fleets_ids,
+                    fleets_ids=fleets_with_instances_ids,
                     instances_ids=instances_ids,
                     fleet_filters=fleet_filters,
                     instance_filters=instance_filters,
                 )
+            fleet_models = fleet_models_with_instances + fleet_models_without_instances
             fleet_model, fleet_instances_with_offers = _find_optimal_fleet_with_offers(
                 fleet_models=fleet_models,
                 run_model=run_model,
@@ -290,7 +294,7 @@ async def _process_submitted_job(session: AsyncSession, job_model: JobModel):
                 master_job_provisioning_data=master_job_provisioning_data,
                 volumes=volumes,
             )
-            if fleet_model is None and run_spec.configuration.fleets is not None:
+            if fleet_model is None and run_spec.merged_profile.fleets is not None:
                 # Run cannot create new fleets when fleets are specified
                 logger.debug("%s: failed to use specified fleets", fmt(job_model))
                 job_model.status = JobStatus.TERMINATING
@@ -361,6 +365,10 @@ async def _process_submitted_job(session: AsyncSession, job_model: JobModel):
                 project=project,
                 run=run,
             )
+        # FIXME: Fleet is not locked which may lead to duplicate instance_num.
+        # This is currently hard to fix without locking the fleet for entire provisioning duration.
+        # Processing should be done in multiple steps so that
+        # InstanceModel is created before provisioning.
         instance_num = await _get_next_instance_num(
             session=session,
             fleet_model=fleet_model,
@@ -438,14 +446,21 @@ async def _select_fleet_models(
             *fleet_filters,
             FleetModel.id.not_in(fleet_models_with_instances_ids),
         )
-        .where(InstanceModel.id.is_(None))
-        .options(contains_eager(FleetModel.instances))  # loading empty relation
+        .where(
+            or_(
+                InstanceModel.id.is_(None),
+                not_(and_(*instance_filters)),
+            )
+        )
+        # Load empty list of instances so that downstream code
+        # knows this fleet has no instances eligible for offers.
+        .options(noload(FleetModel.instances))
     )
     fleet_models_without_instances = list(res.unique().scalars().all())
     return fleet_models_with_instances, fleet_models_without_instances
 
 
-async def _refetch_fleet_models(
+async def _refetch_fleet_models_with_instances(
     session: AsyncSession,
     fleets_ids: list[uuid.UUID],
     instances_ids: list[uuid.UUID],
@@ -460,13 +475,8 @@ async def _refetch_fleet_models(
             *fleet_filters,
         )
         .where(
-            or_(
-                InstanceModel.id.is_(None),
-                and_(
-                    InstanceModel.id.in_(instances_ids),
-                    *instance_filters,
-                ),
-            )
+            InstanceModel.id.in_(instances_ids),
+            *instance_filters,
         )
         .options(contains_eager(FleetModel.instances))
         .execution_options(populate_existing=True)
@@ -533,7 +543,7 @@ def _find_optimal_fleet_with_offers(
                 fleet_priority,
             )
         )
-    if run_spec.configuration.fleets is None and all(
+    if run_spec.merged_profile.fleets is None and all(
         t[2] == 0 for t in candidate_fleets_with_offers
     ):
         # If fleets are not specified and no fleets have available offers, create a new fleet.
@@ -646,6 +656,8 @@ async def _run_job_on_new_instance(
 ) -> Optional[Tuple[JobProvisioningData, InstanceOfferWithAvailability]]:
     if volumes is None:
         volumes = []
+    profile = run.run_spec.merged_profile
+    requirements = job.job_spec.requirements
     fleet = None
     if fleet_model is not None:
         fleet = fleet_model_to_fleet(fleet_model)
@@ -654,13 +666,26 @@ async def _run_job_on_new_instance(
                 "%s: cannot fit new instance into fleet %s", fmt(job_model), fleet_model.name
             )
             return None
+        profile = combine_fleet_and_run_profiles(fleet.spec.merged_profile, profile)
+        if profile is None:
+            logger.debug("%s: cannot combine fleet %s profile", fmt(job_model), fleet_model.name)
+            return None
+        fleet_requirements = get_fleet_requirements(fleet.spec)
+        requirements = combine_fleet_and_run_requirements(fleet_requirements, requirements)
+        if requirements is None:
+            logger.debug(
+                "%s: cannot combine fleet %s requirements", fmt(job_model), fleet_model.name
+            )
+            return None
+        # TODO: Respect fleet provisioning properties such as tags
+
     multinode = job.job_spec.jobs_per_replica > 1 or (
         fleet is not None and fleet.spec.configuration.placement == InstanceGroupPlacement.CLUSTER
     )
     offers = await get_offers_by_requirements(
         project=project,
-        profile=run.run_spec.merged_profile,
-        requirements=job.job_spec.requirements,
+        profile=profile,
+        requirements=requirements,
         exclude_not_available=True,
         multinode=multinode,
         master_job_provisioning_data=master_job_provisioning_data,
@@ -752,25 +777,11 @@ def _create_fleet_model_for_job(
 
 
 async def _get_next_instance_num(session: AsyncSession, fleet_model: FleetModel) -> int:
-    if len(fleet_model.instances) == 0:
-        # No instances means the fleet is not in the db yet, so don't lock.
-        return 0
-    async with get_locker(get_db().dialect_name).lock_ctx(
-        FleetModel.__tablename__, [fleet_model.id]
-    ):
-        fleet_model = (
-            (
-                await session.execute(
-                    select(FleetModel)
-                    .where(FleetModel.id == fleet_model.id)
-                    .options(joinedload(FleetModel.instances))
-                    .execution_options(populate_existing=True)
-                )
-            )
-            .unique()
-            .scalar_one()
-        )
-        return len(fleet_model.instances)
+    res = await session.execute(
+        select(func.count(InstanceModel.id)).where(InstanceModel.fleet_id == fleet_model.id)
+    )
+    instance_count = res.scalar_one()
+    return instance_count
 
 
 def _create_instance_model_for_job(

dstack/_internal/server/background/tasks/process_terminating_jobs.py

@@ -75,9 +75,9 @@ async def _process_next_terminating_job():
                     return
                 instance_lockset.add(instance_model.id)
             job_lockset.add(job_model.id)
+        job_model_id = job_model.id
+        instance_model_id = job_model.used_instance_id
         try:
-            job_model_id = job_model.id
-            instance_model_id = job_model.used_instance_id
             await _process_job(
                 session=session,
                 job_model=job_model,

dstack/_internal/server/background/tasks/process_volumes.py

@@ -42,8 +42,8 @@ async def process_submitted_volumes():
             if volume_model is None:
                 return
             lockset.add(volume_model.id)
+        volume_model_id = volume_model.id
         try:
-            volume_model_id = volume_model.id
             await _process_submitted_volume(session=session, volume_model=volume_model)
         finally:
             lockset.difference_update([volume_model_id])

dstack/_internal/server/db.py

@@ -4,8 +4,12 @@ from typing import Optional
 from alembic import command, config
 from sqlalchemy import AsyncAdaptedQueuePool, event
 from sqlalchemy.engine.interfaces import DBAPIConnection
-from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, create_async_engine
-from sqlalchemy.orm import sessionmaker
+from sqlalchemy.ext.asyncio import (
+    AsyncEngine,
+    AsyncSession,
+    async_sessionmaker,
+    create_async_engine,
+)
 from sqlalchemy.pool import ConnectionPoolEntry
 
 from dstack._internal.server import settings
@@ -26,8 +30,8 @@ class Database:
                 pool_size=settings.DB_POOL_SIZE,
                 max_overflow=settings.DB_MAX_OVERFLOW,
             )
-        self.session_maker = sessionmaker(
-            bind=self.engine,
+        self.session_maker = async_sessionmaker(
+            bind=self.engine,  # type: ignore[assignment]
             expire_on_commit=False,
             class_=AsyncSession,
         )

dstack/_internal/server/models.py

@@ -622,6 +622,7 @@ class InstanceModel(BaseModel):
     backend: Mapped[Optional[BackendType]] = mapped_column(EnumAsString(BackendType, 100))
     backend_data: Mapped[Optional[str]] = mapped_column(Text)
 
+    # Not set for cloud fleets that haven't been provisioning
     offer: Mapped[Optional[str]] = mapped_column(Text)
     region: Mapped[Optional[str]] = mapped_column(String(2000))
     price: Mapped[Optional[float]] = mapped_column(Float)

dstack/_internal/server/routers/gpus.py

@@ -1,9 +1,7 @@
 from typing import Tuple
 
 from fastapi import APIRouter, Depends
-from sqlalchemy.ext.asyncio import AsyncSession
 
-from dstack._internal.server.db import get_session
 from dstack._internal.server.models import ProjectModel, UserModel
 from dstack._internal.server.schemas.gpus import ListGpusRequest, ListGpusResponse
 from dstack._internal.server.security.permissions import ProjectMember
@@ -20,10 +18,7 @@ project_router = APIRouter(
 @project_router.post("/list", response_model=ListGpusResponse, response_model_exclude_none=True)
 async def list_gpus(
     body: ListGpusRequest,
-    session: AsyncSession = Depends(get_session),
     user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectMember()),
 ) -> ListGpusResponse:
     _, project = user_project
-    return await list_gpus_grouped(
-        session=session, project=project, run_spec=body.run_spec, group_by=body.group_by
-    )
+    return await list_gpus_grouped(project=project, run_spec=body.run_spec, group_by=body.group_by)

dstack/_internal/server/schemas/runner.py

@@ -159,6 +159,16 @@ class GPUDevice(CoreModel):
     path_in_container: str
 
 
+class TaskListItem(CoreModel):
+    id: str
+    status: TaskStatus
+
+
+class TaskListResponse(CoreModel):
+    ids: Optional[list[str]] = None  # returned by pre-0.19.26 shim
+    tasks: Optional[list[TaskListItem]] = None  # returned by 0.19.26+ shim
+
+
 class TaskInfoResponse(CoreModel):
     id: str
     status: TaskStatus

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

@@ -17,8 +17,8 @@ from dstack._internal.core.backends.configurators import (
 )
 from dstack._internal.core.backends.local.backend import LocalBackend
 from dstack._internal.core.backends.models import (
-    AnyBackendConfig,
     AnyBackendConfigWithCreds,
+    AnyBackendConfigWithoutCreds,
 )
 from dstack._internal.core.errors import (
     BackendError,
@@ -126,19 +126,25 @@ async def get_backend_config(
             )
             continue
         if backend_model.type == backend_type:
-            return get_backend_config_from_backend_model(
-                configurator, backend_model, include_creds=True
-            )
+            return get_backend_config_with_creds_from_backend_model(configurator, backend_model)
     return None
 
 
-def get_backend_config_from_backend_model(
+def get_backend_config_with_creds_from_backend_model(
+    configurator: Configurator,
+    backend_model: BackendModel,
+) -> AnyBackendConfigWithCreds:
+    backend_record = get_stored_backend_record(backend_model)
+    backend_config = configurator.get_backend_config_with_creds(backend_record)
+    return backend_config
+
+
+def get_backend_config_without_creds_from_backend_model(
     configurator: Configurator,
     backend_model: BackendModel,
-    include_creds: bool,
-) -> AnyBackendConfig:
+) -> AnyBackendConfigWithoutCreds:
     backend_record = get_stored_backend_record(backend_model)
-    backend_config = configurator.get_backend_config(backend_record, include_creds=include_creds)
+    backend_config = configurator.get_backend_config_without_creds(backend_record)
     return backend_config
 
 

dstack/_internal/server/services/backends/handlers.py

@@ -55,7 +55,11 @@ async def _check_active_instances(
     )
     for fleet_model in fleet_models:
         for instance in fleet_model.instances:
-            if instance.status.is_active() and instance.backend in backends_types:
+            if (
+                instance.status.is_active()
+                and instance.backend is not None
+                and instance.backend in backends_types
+            ):
                 if error:
                     msg = (
                         f"Backend {instance.backend.value} has active instances."
@@ -83,6 +87,7 @@ async def _check_active_volumes(
         if (
             volume_model.status.is_active()
             and volume_model.provisioning_data is not None
+            and volume_model.provisioning_data.backend is not None
             and volume_model.provisioning_data.backend in backends_types
         ):
             if error:

dstack/_internal/server/services/docker.py

@@ -32,15 +32,15 @@ class DXFAuthAdapter:
 
 
 class DockerImage(CoreModel):
-    class Config(CoreModel.Config):
-        frozen = True
-
     image: str
     registry: Optional[str]
     repo: str
     tag: str
     digest: Optional[str]
 
+    class Config(CoreModel.Config):
+        frozen = True
+
 
 class ImageConfig(CoreModel):
     user: Annotated[Optional[str], Field(alias="User")] = None
@@ -77,7 +77,7 @@ def get_image_config(image_name: str, registry_auth: Optional[RegistryAuth]) ->
     registry_client = PatchedDXF(
         host=image.registry or DEFAULT_REGISTRY,
         repo=image.repo,
-        auth=DXFAuthAdapter(registry_auth),
+        auth=DXFAuthAdapter(registry_auth),  # type: ignore[assignment]
         timeout=REGISTRY_REQUEST_TIMEOUT,
     )
 
@@ -88,7 +88,7 @@ def get_image_config(image_name: str, registry_auth: Optional[RegistryAuth]) ->
             )
             manifest = ImageManifest.__response__.parse_raw(manifest_resp)
             config_stream = registry_client.pull_blob(manifest.config.digest)
-            config_resp = join_byte_stream_checked(config_stream, MAX_CONFIG_OBJECT_SIZE)
+            config_resp = join_byte_stream_checked(config_stream, MAX_CONFIG_OBJECT_SIZE)  # type: ignore[arg-type]
             if config_resp is None:
                 raise DockerRegistryError(
                     f"Image config object exceeds the size limit of {MAX_CONFIG_OBJECT_SIZE} bytes"

dstack/_internal/server/services/fleets.py

@@ -279,7 +279,7 @@ async def get_plan(
         offers_with_backends = await get_create_instance_offers(
             project=project,
             profile=effective_spec.merged_profile,
-            requirements=_get_fleet_requirements(effective_spec),
+            requirements=get_fleet_requirements(effective_spec),
             fleet_spec=effective_spec,
             blocks=effective_spec.configuration.blocks,
         )
@@ -458,7 +458,7 @@ async def create_fleet_instance_model(
     instance_num: int,
 ) -> InstanceModel:
     profile = spec.merged_profile
-    requirements = _get_fleet_requirements(spec)
+    requirements = get_fleet_requirements(spec)
     instance_model = await instances_services.create_instance_model(
         session=session,
         project=project,
@@ -504,6 +504,7 @@ async def create_fleet_ssh_instance_model(
         raise ServerClientError("ssh key or user not specified")
 
     if proxy_jump is not None:
+        assert proxy_jump.ssh_key is not None
         ssh_proxy = SSHConnectionParams(
             hostname=proxy_jump.hostname,
             port=proxy_jump.port or 22,
@@ -643,6 +644,17 @@ def is_fleet_empty(fleet_model: FleetModel) -> bool:
     return len(active_instances) == 0
 
 
+def get_fleet_requirements(fleet_spec: FleetSpec) -> Requirements:
+    profile = fleet_spec.merged_profile
+    requirements = Requirements(
+        resources=fleet_spec.configuration.resources or ResourcesSpec(),
+        max_price=profile.max_price,
+        spot=get_policy_map(profile.spot_policy, default=SpotPolicy.ONDEMAND),
+        reservation=fleet_spec.configuration.reservation,
+    )
+    return requirements
+
+
 async def _create_fleet(
     session: AsyncSession,
     project: ProjectModel,
@@ -1003,17 +1015,6 @@ def _terminate_fleet_instances(fleet_model: FleetModel, instance_nums: Optional[
             instance.status = InstanceStatus.TERMINATING
 
 
-def _get_fleet_requirements(fleet_spec: FleetSpec) -> Requirements:
-    profile = fleet_spec.merged_profile
-    requirements = Requirements(
-        resources=fleet_spec.configuration.resources or ResourcesSpec(),
-        max_price=profile.max_price,
-        spot=get_policy_map(profile.spot_policy, default=SpotPolicy.ONDEMAND),
-        reservation=fleet_spec.configuration.reservation,
-    )
-    return requirements
-
-
 def _get_next_instance_num(instance_nums: set[int]) -> int:
     if not instance_nums:
         return 0

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

@@ -93,6 +93,8 @@ async def create_gateway_compute(
     backend_id: Optional[uuid.UUID] = None,
 ) -> GatewayComputeModel:
     assert isinstance(backend_compute, ComputeWithGatewaySupport)
+    assert configuration.name is not None
+
     private_bytes, public_bytes = generate_rsa_key_pair_bytes()
     gateway_ssh_private_key = private_bytes.decode()
     gateway_ssh_public_key = public_bytes.decode()

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

@@ -7,7 +7,7 @@ from pydantic import parse_obj_as
 
 from dstack._internal.core.consts import DSTACK_RUNNER_SSH_PORT
 from dstack._internal.core.errors import GatewayError
-from dstack._internal.core.models.configurations import RateLimit, ServiceConfiguration
+from dstack._internal.core.models.configurations import RateLimit
 from dstack._internal.core.models.instances import SSHConnectionParams
 from dstack._internal.core.models.runs import JobSpec, JobSubmission, Run, get_service_port
 from dstack._internal.proxy.gateway.schemas.stats import ServiceStats
@@ -85,7 +85,7 @@ class GatewayClient:
         ssh_head_proxy: Optional[SSHConnectionParams],
         ssh_head_proxy_private_key: Optional[str],
     ):
-        assert isinstance(run.run_spec.configuration, ServiceConfiguration)
+        assert run.run_spec.configuration.type == "service"
         payload = {
             "job_id": job_submission.id.hex,
             "app_port": get_service_port(job_spec, run.run_spec.configuration),
@@ -93,6 +93,9 @@ class GatewayClient:
             "ssh_head_proxy_private_key": ssh_head_proxy_private_key,
         }
         jpd = job_submission.job_provisioning_data
+        assert jpd is not None
+        assert jpd.hostname is not None
+        assert jpd.ssh_port is not None
         if not jpd.dockerized:
             payload.update(
                 {

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

@@ -67,7 +67,7 @@ class GatewayConnection:
             # reverse_forwarded_sockets are added later in .open()
         )
         self.tunnel_id = uuid.uuid4()
-        self._client = GatewayClient(uds=self.gateway_socket_path)
+        self._client = GatewayClient(uds=str(self.gateway_socket_path))
 
     @staticmethod
     def _init_symlink_dir(connection_dir: Path) -> Tuple[TemporaryDirectory, Path]: