From 92a492205c97a096e013107a950df6eb719cb9f1 Mon Sep 17 00:00:00 2001 From: Dmitry Meyer Date: Mon, 12 Oct 2026 08:24:14 +0000 Subject: [PATCH] Support SQLAlchemy 2.1 SQLAlchemy was capped below 2.1 because sqlalchemy-utils didn't support it (kvesteri/sqlalchemy-utils#800). sqlalchemy-utils 0.43.0 adds SQLAlchemy 2.1 support, so require it and drop the cap. Python 3.10, which SQLAlchemy 2.1 doesn't support, still gets 2.0. Also: * delete_metrics() no longer runs its two DELETE statements concurrently with asyncio.gather() in separate sessions. In tests, StaticPool makes both sessions share one SQLite connection. In 2.1, the aiosqlite adapter runs execute + fetchall under a per-connection mutex that commit() doesn't take, so one session's COMMIT could land between the other session's DELETE ... RETURNING and its fetch, and fail with "cannot commit transaction - SQL statements in progress". Under 2.0 the interleaving happened to be safe. Production uses a connection pool, so only tests were affected. Both deletes now run in one session. * noload() is deprecated in 2.1. Replace it with raiseload(), which also skips the mapped lazy="selectin" load, but raises on access instead of returning an empty collection. RunModel.jobs isn't read after list_projects_run_models(), so the swap is enough there. FleetModel.instances is read downstream for candidate fleets that have no matching instances, so set it to an empty list explicitly with set_committed_value(). Co-Authored-By: Claude Opus 5.5 (1M context) --- pyproject.toml | 5 ++--- .../background/scheduled_tasks/metrics.py | 19 ++++++------------- .../server/services/runs/__init__.py | 4 ++-- .../_internal/server/services/runs/plan.py | 8 ++++++-- 4 files changed, 16 insertions(+), 20 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index d6b0d5f07b..1ddb8ef87d 100644 --- a/pyproject.toml +++ b/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", diff --git a/src/dstack/_internal/server/background/scheduled_tasks/metrics.py b/src/dstack/_internal/server/background/scheduled_tasks/metrics.py index a083d5956e..3d6695352c 100644 --- a/src/dstack/_internal/server/background/scheduled_tasks/metrics.py +++ b/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]): diff --git a/src/dstack/_internal/server/services/runs/__init__.py b/src/dstack/_internal/server/services/runs/__init__.py index 4d9e4a2de6..b6ff199641 100644 --- a/src/dstack/_internal/server/services/runs/__init__.py +++ b/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) ) diff --git a/src/dstack/_internal/server/services/runs/plan.py b/src/dstack/_internal/server/services/runs/plan.py index ae626b1c30..c5c175c57b 100644 --- a/src/dstack/_internal/server/services/runs/plan.py +++ b/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