Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@
InstanceModel,
JobModel,
PlacementGroupModel,
ProjectModel,
RunModel,
)
from dstack._internal.server.services import events
Expand Down Expand Up @@ -338,7 +339,7 @@ async def _refetch_locked_fleet_for_processing(
FleetModel.id == item.id,
FleetModel.lock_token == item.lock_token,
)
.options(joinedload(FleetModel.project))
.options(joinedload(FleetModel.project).joinedload(ProjectModel.owner))
.options(
selectinload(FleetModel.instances.and_(InstanceModel.deleted == False))
.joinedload(InstanceModel.jobs)
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,52 @@
"""Add missing foreign key indexes

Revision ID: a1c3f5e7b209
Revises: ecc9e8a0bfac
Create Date: 2026-08-11 09:40:00.000000+00:00

"""

from alembic import op

# revision identifiers, used by Alembic.
revision = "a1c3f5e7b209"
down_revision = "ecc9e8a0bfac"
branch_labels = None
depends_on = None


INDEXES = [
("ix_jobs_instance_id", "jobs", "instance_id"),
("ix_backends_project_id", "backends", "project_id"),
("ix_members_project_id", "members", "project_id"),
("ix_instances_compute_group_id", "instances", "compute_group_id"),
("ix_runs_fleet_id", "runs", "fleet_id"),
]


def _drop_indexes() -> None:
for index_name, table_name, _ in INDEXES:
op.drop_index(
index_name,
table_name=table_name,
if_exists=True,
postgresql_concurrently=True,
)


def upgrade() -> None:
with op.get_context().autocommit_block():
_drop_indexes()
for index_name, table_name, column_name in INDEXES:
op.create_index(
index_name,
table_name,
[column_name],
unique=False,
postgresql_concurrently=True,
)


def downgrade() -> None:
with op.get_context().autocommit_block():
_drop_indexes()
18 changes: 12 additions & 6 deletions src/dstack/_internal/server/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -260,7 +260,7 @@ class ProjectModel(BaseModel):
"""

owner_id: Mapped[uuid.UUID] = mapped_column(ForeignKey("users.id", ondelete="CASCADE"))
owner: Mapped[UserModel] = relationship(lazy="joined")
owner: Mapped[UserModel] = relationship()
members: Mapped[List["MemberModel"]] = relationship(
back_populates="project", order_by="MemberModel.member_num"
)
Expand Down Expand Up @@ -296,7 +296,9 @@ class MemberModel(BaseModel):
id: Mapped[uuid.UUID] = mapped_column(
UUIDType(binary=False), primary_key=True, default=uuid.uuid4
)
project_id: Mapped[uuid.UUID] = mapped_column(ForeignKey("projects.id", ondelete="CASCADE"))
project_id: Mapped[uuid.UUID] = mapped_column(
ForeignKey("projects.id", ondelete="CASCADE"), index=True
)
project: Mapped["ProjectModel"] = relationship()
user_id: Mapped[uuid.UUID] = mapped_column(ForeignKey("users.id", ondelete="CASCADE"))
user: Mapped[UserModel] = relationship(lazy="joined")
Expand All @@ -311,7 +313,9 @@ class BackendModel(BaseModel):
id: Mapped[uuid.UUID] = mapped_column(
UUIDType(binary=False), primary_key=True, default=uuid.uuid4
)
project_id: Mapped[uuid.UUID] = mapped_column(ForeignKey("projects.id", ondelete="CASCADE"))
project_id: Mapped[uuid.UUID] = mapped_column(
ForeignKey("projects.id", ondelete="CASCADE"), index=True
)
project: Mapped["ProjectModel"] = relationship()
type: Mapped[BackendType] = mapped_column(EnumAsString(BackendType, 100))

Expand Down Expand Up @@ -418,7 +422,7 @@ class RunModel(PipelineModelMixin, BaseModel):
repo_id: Mapped[uuid.UUID] = mapped_column(ForeignKey("repos.id", ondelete="CASCADE"))
repo: Mapped["RepoModel"] = relationship()

fleet_id: Mapped[Optional[uuid.UUID]] = mapped_column(ForeignKey("fleets.id"))
fleet_id: Mapped[Optional[uuid.UUID]] = mapped_column(ForeignKey("fleets.id"), index=True)
"""`fleet_id` keeps runs attached to fleets so the fleets cannot be deleted while they are used.
A fleet can have no busy instances but still be used by a run, for example a service with
zero replicas.
Expand Down Expand Up @@ -564,7 +568,7 @@ class JobModel(PipelineModelMixin, BaseModel):
If `instance_assigned` is `True` and `instance` is `None`, no instance was assigned.
"""
instance_id: Mapped[Optional[uuid.UUID]] = mapped_column(
ForeignKey("instances.id", ondelete="CASCADE")
ForeignKey("instances.id", ondelete="CASCADE"), index=True
)
instance: Mapped[Optional["InstanceModel"]] = relationship(back_populates="jobs")
used_instance_id: Mapped[Optional[uuid.UUID]] = mapped_column(UUIDType(binary=False))
Expand Down Expand Up @@ -853,7 +857,9 @@ class InstanceModel(PipelineModelMixin, BaseModel):
)
"""`fleet` can be `None` only for legacy instances created before fleets."""

compute_group_id: Mapped[Optional[uuid.UUID]] = mapped_column(ForeignKey("compute_groups.id"))
compute_group_id: Mapped[Optional[uuid.UUID]] = mapped_column(
ForeignKey("compute_groups.id"), index=True
)
compute_group: Mapped[Optional["ComputeGroupModel"]] = relationship(back_populates="instances")

status: Mapped[InstanceStatus] = mapped_column(EnumAsString(InstanceStatus, 100), index=True)
Expand Down
1 change: 0 additions & 1 deletion src/dstack/_internal/server/services/events.py
Original file line number Diff line number Diff line change
Expand Up @@ -405,7 +405,6 @@ async def list_events(
joinedload(EventModel.targets)
.joinedload(EventTargetModel.entity_project)
.load_only(ProjectModel.name, ProjectModel.original_name, ProjectModel.deleted)
.noload(ProjectModel.owner)
),
joinedload(EventModel.actor_user).load_only(
UserModel.name, UserModel.original_name, UserModel.deleted
Expand Down
1 change: 1 addition & 0 deletions src/dstack/_internal/server/services/fleets.py
Original file line number Diff line number Diff line change
Expand Up @@ -206,6 +206,7 @@ async def list_projects_with_no_active_fleets(
active_fleet_alias.id.is_(None),
)
.order_by(ProjectModel.created_at)
.options(joinedload(ProjectModel.owner))
)

res = await session.execute(query)
Expand Down
44 changes: 8 additions & 36 deletions src/dstack/_internal/server/services/projects.py
Original file line number Diff line number Diff line change
Expand Up @@ -147,7 +147,12 @@ async def list_user_accessible_projects(
if return_total_count:
res = await session.execute(stmt.with_only_columns(safunc.count(literal_column("1"))))
total_count = res.scalar_one()
res = await session.execute(stmt.where(*pagination_filters).order_by(*order_by).limit(limit))
res = await session.execute(
stmt.where(*pagination_filters)
.order_by(*order_by)
.limit(limit)
.options(joinedload(ProjectModel.owner))
)
project_models = res.unique().scalars().all()
projects = [
project_model_to_project(p, include_backends=False, include_members=False)
Expand Down Expand Up @@ -516,25 +521,6 @@ async def list_member_project_models(
return list(res.scalars().unique().all())


async def list_public_non_member_project_models(
session: AsyncSession,
user: UserModel,
) -> List[ProjectModel]:
"""
List public project models where user is NOT a member.
"""
res = await session.execute(
select(ProjectModel).where(
ProjectModel.deleted == False,
ProjectModel.is_public == True,
ProjectModel.id.notin_(
select(MemberModel.project_id).where(MemberModel.user_id == user.id)
),
)
)
return list(res.scalars().all())


async def list_user_owned_project_models(
session: AsyncSession, user: UserModel, include_deleted: bool = False
) -> List[ProjectModel]:
Expand Down Expand Up @@ -580,6 +566,7 @@ async def get_project_model_by_name(
res = await session.execute(
select(ProjectModel)
.where(*filters)
.options(joinedload(ProjectModel.owner))
.options(joinedload(ProjectModel.backends))
.options(joinedload(ProjectModel.members))
)
Expand All @@ -596,22 +583,7 @@ async def get_project_model_by_name_or_error(
ProjectModel.name == project_name,
ProjectModel.deleted == False,
)
.options(joinedload(ProjectModel.backends))
.options(joinedload(ProjectModel.members))
)
return res.unique().scalar_one()


async def get_project_model_by_id_or_error(
session: AsyncSession,
project_id: uuid.UUID,
) -> ProjectModel:
res = await session.execute(
select(ProjectModel)
.where(
ProjectModel.id == project_id,
ProjectModel.deleted == False,
)
.options(joinedload(ProjectModel.owner))
.options(joinedload(ProjectModel.backends))
.options(joinedload(ProjectModel.members))
)
Expand Down
14 changes: 11 additions & 3 deletions src/dstack/_internal/server/services/prometheus/custom_metrics.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,11 +10,12 @@
from prometheus_client.samples import Sample
from sqlalchemy import func, select
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import aliased, joinedload
from sqlalchemy.orm import aliased, joinedload, load_only

from dstack._internal.core.models.instances import InstanceStatus
from dstack._internal.core.models.runs import JobStatus, RunStatus
from dstack._internal.server.models import (
FleetModel,
InstanceModel,
JobMetricsPoint,
JobModel,
Expand Down Expand Up @@ -55,8 +56,15 @@ async def get_instance_metrics(session: AsyncSession) -> Iterable[Metric]:
)
.order_by(ProjectModel.name, InstanceModel.name)
.options(
joinedload(InstanceModel.project),
joinedload(InstanceModel.fleet),
load_only(
InstanceModel.name,
InstanceModel.created_at,
InstanceModel.price,
InstanceModel.backend,
InstanceModel.offer,
),
joinedload(InstanceModel.project).load_only(ProjectModel.name),
joinedload(InstanceModel.fleet).load_only(FleetModel.name),
)
)
instances = res.unique().scalars().all()
Expand Down
Loading