Remove implicit gateway_compute load

Author: r4victor

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

@@ -153,6 +153,7 @@ async def _process_provisioning_gateway(
     res = await session.execute(
         select(GatewayModel)
         .where(GatewayModel.id == gateway_model.id)
+        .options(joinedload(GatewayModel.gateway_compute))
         .execution_options(populate_existing=True)
     )
     gateway_model = res.unique().scalar_one()

src/dstack/_internal/server/models.py

@@ -518,7 +518,7 @@ class GatewayModel(PipelineModelMixin, BaseModel):
     gateway_compute_id: Mapped[Optional[uuid.UUID]] = mapped_column(
         ForeignKey("gateway_computes.id", ondelete="CASCADE")
     )
-    gateway_compute: Mapped[Optional["GatewayComputeModel"]] = relationship(lazy="joined")
+    gateway_compute: Mapped[Optional["GatewayComputeModel"]] = relationship()
 
     runs: Mapped[List["RunModel"]] = relationship(back_populates="gateway")
 

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

@@ -127,14 +127,23 @@ def get_gateway_status_change_message(
 
 
 async def list_project_gateways(session: AsyncSession, project: ProjectModel) -> List[Gateway]:
-    gateways = await list_project_gateway_models(session=session, project=project)
+    gateways = await list_project_gateway_models(
+        session=session,
+        project=project,
+        load_gateway_compute=True,
+    )
     return [gateway_model_to_gateway(g) for g in gateways]
 
 
 async def get_gateway_by_name(
     session: AsyncSession, project: ProjectModel, name: str
 ) -> Optional[Gateway]:
-    gateway = await get_project_gateway_model_by_name(session=session, project=project, name=name)
+    gateway = await get_project_gateway_model_by_name(
+        session=session,
+        project=project,
+        name=name,
+        load_gateway_compute=True,
+    )
     if gateway is None:
         return None
     return gateway_model_to_gateway(gateway)
@@ -494,20 +503,30 @@ async def set_default_gateway(
 
 
 async def list_project_gateway_models(
-    session: AsyncSession, project: ProjectModel
+    session: AsyncSession,
+    project: ProjectModel,
+    load_gateway_compute: bool = False,
 ) -> Sequence[GatewayModel]:
-    res = await session.execute(select(GatewayModel).where(GatewayModel.project_id == project.id))
+    stmt = select(GatewayModel).where(GatewayModel.project_id == project.id)
+    if load_gateway_compute:
+        stmt = stmt.options(joinedload(GatewayModel.gateway_compute))
+    res = await session.execute(stmt)
     return res.scalars().all()
 
 
 async def get_project_gateway_model_by_name(
-    session: AsyncSession, project: ProjectModel, name: str
+    session: AsyncSession,
+    project: ProjectModel,
+    name: str,
+    load_gateway_compute: bool = False,
 ) -> Optional[GatewayModel]:
-    res = await session.execute(
-        select(GatewayModel).where(
-            GatewayModel.project_id == project.id, GatewayModel.name == name
-        )
+    stmt = select(GatewayModel).where(
+        GatewayModel.project_id == project.id,
+        GatewayModel.name == name,
     )
+    if load_gateway_compute:
+        stmt = stmt.options(joinedload(GatewayModel.gateway_compute))
+    res = await session.execute(stmt)
     return res.scalar()
 
 
@@ -538,20 +557,24 @@ async def get_project_gateway_model_by_name_for_update(
             res = await session.execute(
                 select(GatewayModel)
                 .where(GatewayModel.id.in_([gateway_id]), *filters)
+                .options(joinedload(GatewayModel.gateway_compute))
                 .with_for_update(key_share=True, of=GatewayModel)
             )
             yield res.scalar_one_or_none()
 
 
 async def get_project_default_gateway_model(
-    session: AsyncSession, project: ProjectModel
+    session: AsyncSession,
+    project: ProjectModel,
+    load_gateway_compute: bool = False,
 ) -> Optional[GatewayModel]:
-    res = await session.execute(
-        select(GatewayModel).where(
-            GatewayModel.id == project.default_gateway_id,
-            GatewayModel.to_be_deleted == False,
-        )
+    stmt = select(GatewayModel).where(
+        GatewayModel.id == project.default_gateway_id,
+        GatewayModel.to_be_deleted == False,
     )
+    if load_gateway_compute:
+        stmt = stmt.options(joinedload(GatewayModel.gateway_compute))
+    res = await session.execute(stmt)
     return res.scalar_one_or_none()
 
 
@@ -567,7 +590,12 @@ async def generate_gateway_name(session: AsyncSession, project: ProjectModel) ->
 async def get_or_add_gateway_connection(
     session: AsyncSession, gateway_id: uuid.UUID
 ) -> tuple[GatewayModel, GatewayConnection]:
-    gateway = await session.get(GatewayModel, gateway_id)
+    gateway = await session.get(
+        GatewayModel,
+        gateway_id,
+        options=[joinedload(GatewayModel.gateway_compute)],
+        populate_existing=True,
+    )
     if gateway is None:
         raise GatewayError("Gateway not found")
     if gateway.gateway_compute is None:

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

@@ -90,7 +90,10 @@ async def register_service(session: AsyncSession, run_model: RunModel, run_spec:
 
     if isinstance(run_spec.configuration.gateway, str):
         gateway = await get_project_gateway_model_by_name(
-            session=session, project=run_model.project, name=run_spec.configuration.gateway
+            session=session,
+            project=run_model.project,
+            name=run_spec.configuration.gateway,
+            load_gateway_compute=True,
         )
         if gateway is None:
             raise ResourceNotExistsError(
@@ -104,7 +107,9 @@ async def register_service(session: AsyncSession, run_model: RunModel, run_spec:
         gateway = None
     else:
         gateway = await get_project_default_gateway_model(
-            session=session, project=run_model.project
+            session=session,
+            project=run_model.project,
+            load_gateway_compute=True,
         )
         if gateway is None and run_spec.configuration.gateway == True:
             raise ResourceNotExistsError(

src/tests/_internal/server/background/pipeline_tasks/test_gateways.py

@@ -5,6 +5,7 @@
 import pytest
 from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy.orm import joinedload
 
 from dstack._internal.core.errors import BackendError
 from dstack._internal.core.models.gateways import GatewayProvisioningData, GatewayStatus
@@ -77,6 +78,12 @@ async def test_submitted_to_provisioning(
             aws.compute.return_value.create_gateway.assert_called_once()
 
         await session.refresh(gateway)
+        res = await session.execute(
+            select(GatewayModel)
+            .where(GatewayModel.id == gateway.id)
+            .options(joinedload(GatewayModel.gateway_compute))
+        )
+        gateway = res.unique().scalar_one()
         assert gateway.status == GatewayStatus.PROVISIONING
         assert gateway.gateway_compute is not None
         assert gateway.gateway_compute.ip_address == "2.2.2.2"

src/tests/_internal/server/background/scheduled_tasks/test_gateways.py

@@ -1,11 +1,14 @@
 from unittest.mock import MagicMock, Mock, patch
 
 import pytest
+from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy.orm import joinedload
 
 from dstack._internal.core.errors import BackendError
 from dstack._internal.core.models.gateways import GatewayProvisioningData, GatewayStatus
 from dstack._internal.server.background.scheduled_tasks.gateways import process_gateways
+from dstack._internal.server.models import GatewayModel
 from dstack._internal.server.testing.common import (
     AsyncContextManager,
     ComputeMockSpec,
@@ -44,6 +47,12 @@ async def test_submitted_to_provisioning(self, test_db, session: AsyncSession):
             m.assert_called_once()
             aws.compute.return_value.create_gateway.assert_called_once()
         await session.refresh(gateway)
+        res = await session.execute(
+            select(GatewayModel)
+            .where(GatewayModel.id == gateway.id)
+            .options(joinedload(GatewayModel.gateway_compute))
+        )
+        gateway = res.unique().scalar_one()
         assert gateway.status == GatewayStatus.PROVISIONING
         assert gateway.gateway_compute is not None
         assert gateway.gateway_compute.ip_address == "2.2.2.2"