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