From 9648b88fd339534836bb8071f39b319b98420b13 Mon Sep 17 00:00:00 2001 From: Victor Skvortsov Date: Thu, 23 Jul 2026 13:53:29 +0500 Subject: [PATCH] Implement within_fleets filter via entity_fleet_id --- ...75_add_eventtargetmodel_entity_fleet_id.py | 47 +++++++++++++++++++ ...backfill_eventtargetmodel_entity_fleet_.py | 46 ++++++++++++++++++ src/dstack/_internal/server/models.py | 5 ++ .../_internal/server/services/events.py | 30 +++++------- .../_internal/server/services/fleets.py | 6 +-- .../_internal/server/routers/test_events.py | 28 ++++++++++- 6 files changed, 141 insertions(+), 21 deletions(-) create mode 100644 src/dstack/_internal/server/migrations/versions/2026/07_23_0822_0ebd6564f375_add_eventtargetmodel_entity_fleet_id.py create mode 100644 src/dstack/_internal/server/migrations/versions/2026/07_23_0825_4d3cbb932bb2_backfill_eventtargetmodel_entity_fleet_.py diff --git a/src/dstack/_internal/server/migrations/versions/2026/07_23_0822_0ebd6564f375_add_eventtargetmodel_entity_fleet_id.py b/src/dstack/_internal/server/migrations/versions/2026/07_23_0822_0ebd6564f375_add_eventtargetmodel_entity_fleet_id.py new file mode 100644 index 000000000..ae49fb500 --- /dev/null +++ b/src/dstack/_internal/server/migrations/versions/2026/07_23_0822_0ebd6564f375_add_eventtargetmodel_entity_fleet_id.py @@ -0,0 +1,47 @@ +"""Add EventTargetModel.entity_fleet_id + +Revision ID: 0ebd6564f375 +Revises: 87d4312605e5 +Create Date: 2026-07-23 08:22:23.295309+00:00 + +""" + +import sqlalchemy as sa +import sqlalchemy_utils +from alembic import op + +# revision identifiers, used by Alembic. +revision = "0ebd6564f375" +down_revision = "87d4312605e5" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + with op.batch_alter_table("event_targets", schema=None) as batch_op: + batch_op.add_column( + sa.Column( + "entity_fleet_id", + sqlalchemy_utils.types.uuid.UUIDType(binary=False), + nullable=True, + ) + ) + batch_op.create_index( + batch_op.f("ix_event_targets_entity_fleet_id"), ["entity_fleet_id"], unique=False + ) + batch_op.create_foreign_key( + batch_op.f("fk_event_targets_entity_fleet_id_fleets"), + "fleets", + ["entity_fleet_id"], + ["id"], + ondelete="CASCADE", + ) + + +def downgrade() -> None: + with op.batch_alter_table("event_targets", schema=None) as batch_op: + batch_op.drop_constraint( + batch_op.f("fk_event_targets_entity_fleet_id_fleets"), type_="foreignkey" + ) + batch_op.drop_index(batch_op.f("ix_event_targets_entity_fleet_id")) + batch_op.drop_column("entity_fleet_id") diff --git a/src/dstack/_internal/server/migrations/versions/2026/07_23_0825_4d3cbb932bb2_backfill_eventtargetmodel_entity_fleet_.py b/src/dstack/_internal/server/migrations/versions/2026/07_23_0825_4d3cbb932bb2_backfill_eventtargetmodel_entity_fleet_.py new file mode 100644 index 000000000..1aa715171 --- /dev/null +++ b/src/dstack/_internal/server/migrations/versions/2026/07_23_0825_4d3cbb932bb2_backfill_eventtargetmodel_entity_fleet_.py @@ -0,0 +1,46 @@ +"""Backfill EventTargetModel.entity_fleet_id + +Revision ID: 4d3cbb932bb2 +Revises: 0ebd6564f375 +Create Date: 2026-07-23 08:25:22.615590+00:00 + +""" + +from alembic import op + +# revision identifiers, used by Alembic. +revision = "4d3cbb932bb2" +down_revision = "0ebd6564f375" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + # Events recorded before entity_fleet_id was introduced have it unset. + # The within_fleets events filter relies on entity_fleet_id, so backfill it. + # Instance targets are backfilled via the instances table, so this migration + # must run before the instance models the events reference are deleted + # (e.g. placeholder instances deleted on job termination). + # Old replicas can still record events without entity_fleet_id while + # this migration is being deployed. Such events won't be backfilled and + # won't be returned by the within_fleets events filter. + op.execute( + """ + UPDATE event_targets SET entity_fleet_id = entity_id + WHERE entity_type = 'FLEET' AND entity_fleet_id IS NULL + """ + ) + op.execute( + """ + UPDATE event_targets SET entity_fleet_id = instances.fleet_id + FROM instances + WHERE instances.id = event_targets.entity_id + AND instances.fleet_id IS NOT NULL + AND event_targets.entity_type = 'INSTANCE' + AND event_targets.entity_fleet_id IS NULL + """ + ) + + +def downgrade() -> None: + pass diff --git a/src/dstack/_internal/server/models.py b/src/dstack/_internal/server/models.py index d5452f059..e49726540 100644 --- a/src/dstack/_internal/server/models.py +++ b/src/dstack/_internal/server/models.py @@ -1188,6 +1188,11 @@ class EventTargetModel(BaseModel): ) entity_run: Mapped[Optional["RunModel"]] = relationship() + entity_fleet_id: Mapped[Optional[uuid.UUID]] = mapped_column( + ForeignKey("fleets.id", ondelete="CASCADE"), nullable=True, index=True + ) + entity_fleet: Mapped[Optional["FleetModel"]] = relationship() + entity_type: Mapped[EventTargetType] = mapped_column( EnumAsString(EventTargetType, 100), index=True ) diff --git a/src/dstack/_internal/server/services/events.py b/src/dstack/_internal/server/services/events.py index 5eaf73b06..65dd83ff2 100644 --- a/src/dstack/_internal/server/services/events.py +++ b/src/dstack/_internal/server/services/events.py @@ -77,6 +77,7 @@ class Target: id: uuid.UUID name: str run_id: Optional[uuid.UUID] = None + fleet_id: Optional[uuid.UUID] = None def __post_init__(self): if self.type == EventTargetType.USER and self.project_id is not None: @@ -89,6 +90,8 @@ def __post_init__(self): raise ValueError(f"{self.type} target must have run_id") if self.type == EventTargetType.RUN and self.id != self.run_id: raise ValueError("Run target id must be equal to run_id") + if self.type == EventTargetType.FLEET and self.id != self.fleet_id: + raise ValueError("Fleet target id must be equal to fleet_id") @staticmethod def from_model( @@ -110,6 +113,7 @@ def from_model( project_id=model.project_id or model.project.id, id=model.id, name=model.name, + fleet_id=model.id, ) if isinstance(model, GatewayModel): return Target( @@ -119,11 +123,18 @@ def from_model( name=model.name, ) if isinstance(model, InstanceModel): + fleet_id = model.fleet_id + if fleet_id is None: + # Not-yet-flushed models may only have the fleet relationship set. + fleet = model.__dict__.get("fleet") + if fleet is not None: + fleet_id = fleet.id return Target( type=EventTargetType.INSTANCE, project_id=model.project_id or model.project.id, id=model.id, name=model.name, + fleet_id=fleet_id, ) if isinstance(model, JobModel): return Target( @@ -233,6 +244,7 @@ def emit(session: AsyncSession, message: str, actor: AnyActor, targets: list[Tar entity_id=target.id, entity_name=target.name, entity_run_id=target.run_id, + entity_fleet_id=target.fleet_id, ) ) session.add(event) @@ -344,23 +356,7 @@ async def list_events( if within_projects is not None: target_filters.append(EventTargetModel.entity_project_id.in_(within_projects)) if within_fleets is not None: - query = select(InstanceModel.id).where(InstanceModel.fleet_id.in_(within_fleets)) - res = await session.execute(query) - # In Postgres, fetching instance IDs separately is orders of magnitude faster - # than using a subquery. - instance_ids = list(res.unique().scalars().all()) - target_filters.append( - or_( - and_( - EventTargetModel.entity_type == EventTargetType.FLEET, - EventTargetModel.entity_id.in_(within_fleets), - ), - and_( - EventTargetModel.entity_type == EventTargetType.INSTANCE, - EventTargetModel.entity_id.in_(instance_ids), - ), - ) - ) + target_filters.append(EventTargetModel.entity_fleet_id.in_(within_fleets)) if within_runs is not None: target_filters.append(EventTargetModel.entity_run_id.in_(within_runs)) if include_target_types is not None: diff --git a/src/dstack/_internal/server/services/fleets.py b/src/dstack/_internal/server/services/fleets.py index 547f91d52..2d73942c5 100644 --- a/src/dstack/_internal/server/services/fleets.py +++ b/src/dstack/_internal/server/services/fleets.py @@ -1074,6 +1074,7 @@ async def _create_fleet( instance_num=i, host=host, ) + fleet_model.instances.append(instance_model) events.emit( session, ( @@ -1083,7 +1084,6 @@ async def _create_fleet( actor=events.UserActor.from_user(user), targets=[events.Target.from_model(instance_model)], ) - fleet_model.instances.append(instance_model) else: for i in range(_get_fleet_nodes_to_provision(spec)): instance_model = create_fleet_instance_model( @@ -1093,6 +1093,7 @@ async def _create_fleet( spec=spec, instance_num=i, ) + fleet_model.instances.append(instance_model) events.emit( session, ( @@ -1106,7 +1107,6 @@ async def _create_fleet( actor=events.SystemActor(), targets=[events.Target.from_model(instance_model)], ) - fleet_model.instances.append(instance_model) await session.commit() if spec.configuration.ssh_config is None: pipeline_hinter.hint_fetch(FleetModel.__name__) @@ -1186,13 +1186,13 @@ async def _update_fleet( instance_num=instance_num, host=host, ) + fleet_model.instances.append(instance_model) events.emit( session, f"Instance created on fleet update. Status: {instance_model.status.upper()}", actor=events.UserActor.from_user(user), targets=[events.Target.from_model(instance_model)], ) - fleet_model.instances.append(instance_model) active_instance_nums.add(instance_num) if removed_instance_nums: _terminate_fleet_instances(session, fleet_model, removed_instance_nums, actor=user) diff --git a/src/tests/_internal/server/routers/test_events.py b/src/tests/_internal/server/routers/test_events.py index 7c47cc724..7dd6359af 100644 --- a/src/tests/_internal/server/routers/test_events.py +++ b/src/tests/_internal/server/routers/test_events.py @@ -10,7 +10,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from dstack._internal.core.models.users import GlobalRole, ProjectRole -from dstack._internal.server.models import JobModel +from dstack._internal.server.models import InstanceModel, JobModel from dstack._internal.server.services import events from dstack._internal.server.services.projects import add_project_member from dstack._internal.server.testing.common import ( @@ -859,6 +859,32 @@ async def test_within_fleets(self, session: AsyncSession, client: AsyncClient) - resp.raise_for_status() assert len(resp.json()) == 3 + async def test_within_fleets_finds_events_of_deleted_instances( + self, session: AsyncSession, client: AsyncClient + ) -> None: + user = await create_user(session=session) + project = await create_project(session=session, owner=user) + fleet = await create_fleet(session=session, project=project) + instance = await create_instance(session=session, project=project, fleet=fleet) + events.emit( + session, + "Instance created for job", + actor=events.SystemActor(), + targets=[events.Target.from_model(instance)], + ) + await session.commit() + # Placeholder instances that never provisioned are deleted on job termination + await session.execute(delete(InstanceModel).where(InstanceModel.id == instance.id)) + await session.commit() + + resp = await client.post( + "/api/events/list", + headers=get_auth_headers(user.token), + json={"within_fleets": [str(fleet.id)]}, + ) + resp.raise_for_status() + assert len(resp.json()) == 1 + async def test_within_runs(self, session: AsyncSession, client: AsyncClient) -> None: user = await create_user(session=session) project = await create_project(session=session, owner=user)