Support SQLAlchemy 2.1

pyproject.toml

@@ -235,9 +235,8 @@ server = [
     "httpx>=0.28.0",
     "jinja2",
     "watchfiles",
-    # sqlalchemy_utils is not compatible with SQLAlchemy >= 2.1.0
-    "sqlalchemy[asyncio]>=2.0.0,<2.1",
-    "sqlalchemy_utils>=0.40.0",
+    "sqlalchemy[asyncio]>=2.0.0",
+    "sqlalchemy_utils>=0.43.0",
     "alembic>=1.16.0",
     "aiosqlite",
     "docker>=6.0.0",

src/dstack/_internal/server/background/scheduled_tasks/metrics.py

@@ -3,7 +3,7 @@
 from collections.abc import Mapping
 from typing import List, Optional
 
-from sqlalchemy import Delete, delete, select
+from sqlalchemy import delete, select
 from sqlalchemy.orm import joinedload
 
 from dstack._internal.core.consts import DSTACK_RUNNER_HTTP_PORT
@@ -57,30 +57,23 @@ async def delete_metrics():
     finished_timestamp_micro_cutoff = (
         now_timestamp_micro - settings.SERVER_METRICS_FINISHED_TTL_SECONDS * 1_000_000
     )
-    await asyncio.gather(
-        _execute_delete_statement(
+    async with get_session_ctx() as session:
+        await session.execute(
             delete(JobMetricsPoint).where(
                 JobMetricsPoint.job_id.in_(
                     select(JobModel.id).where(JobModel.status.in_([JobStatus.RUNNING]))
                 ),
                 JobMetricsPoint.timestamp_micro < running_timestamp_micro_cutoff,
             )
-        ),
-        _execute_delete_statement(
+        )
+        await session.execute(
             delete(JobMetricsPoint).where(
                 JobMetricsPoint.job_id.in_(
                     select(JobModel.id).where(JobModel.status.in_(JobStatus.finished_statuses()))
                 ),
                 JobMetricsPoint.timestamp_micro < finished_timestamp_micro_cutoff,
             )
-        ),
-    )
-
-
-async def _execute_delete_statement(stmt: Delete) -> None:
-    async with get_session_ctx() as session:
-        await session.execute(stmt)
-        await session.commit()
+        )
 
 
 async def _collect_jobs_metrics(job_models: List[JobModel]):

src/dstack/_internal/server/services/runs/__init__.py

@@ -10,7 +10,7 @@
 from apscheduler.triggers.cron import CronTrigger
 from sqlalchemy import and_, func, or_, select, update
 from sqlalchemy.ext.asyncio import AsyncSession
-from sqlalchemy.orm import aliased, joinedload, noload, selectinload
+from sqlalchemy.orm import aliased, joinedload, raiseload, selectinload
 
 import dstack._internal.utils.common as common_utils
 from dstack._internal.core.errors import (
@@ -332,7 +332,7 @@ async def list_projects_run_models(
         .options(joinedload(RunModel.project).load_only(ProjectModel.id, ProjectModel.name))
         .options(joinedload(RunModel.user).load_only(UserModel.name))
         .options(joinedload(RunModel.fleet).load_only(FleetModel.id, FleetModel.name))
-        .options(noload(RunModel.jobs))
+        .options(raiseload(RunModel.jobs))
         .order_by(*order_by)
         .limit(limit)
     )

src/dstack/_internal/server/services/runs/plan.py

@@ -8,7 +8,8 @@
 
 from sqlalchemy import and_, exists, not_, or_, select
 from sqlalchemy.ext.asyncio import AsyncSession
-from sqlalchemy.orm import contains_eager, noload
+from sqlalchemy.orm import contains_eager, raiseload
+from sqlalchemy.orm.attributes import set_committed_value
 
 from dstack._internal.core.backends.base.backend import Backend
 from dstack._internal.core.models.common import EntityReference
@@ -336,10 +337,13 @@ async def select_run_candidate_fleet_models_with_filters(
                 not_(and_(*instance_filters)),
             )
         )
-        .options(noload(FleetModel.instances))
+        .options(raiseload(FleetModel.instances))
         .execution_options(populate_existing=True)
     )
     fleet_models_without_instances = list(res.unique().scalars().all())
+    for fleet_model in fleet_models_without_instances:
+        # No instances match the filters, consistent with contains_eager() in the first query
+        set_committed_value(fleet_model, "instances", [])
     return fleet_models_with_instances, fleet_models_without_instances