Remove worker registration to gateway in PD

Author: Bihan Rana

src/dstack/_internal/proxy/gateway/services/registry.py

@@ -293,13 +293,6 @@ async def apply_service(
             )
             for replica, conn in replica_conns.items()
         ]
-    router_replicas = [r for r in service.replicas if r.is_router_replica]
-    if router_replicas:
-        replica_configs_for_nginx = [c for c in replica_configs if c.id == router_replicas[0].id]
-        service_config = await get_nginx_service_config(service, replica_configs_for_nginx)
-        await nginx.register(service_config, (await repo.get_config()).acme_settings)
-        return replica_failures
-
     service_config = await get_nginx_service_config(service, replica_configs)
     await nginx.register(service_config, (await repo.get_config()).acme_settings)
     return replica_failures

src/dstack/_internal/server/background/pipeline_tasks/jobs_running.py

@@ -14,7 +14,10 @@
 from dstack._internal.core.errors import GatewayError, SSHError
 from dstack._internal.core.models.backends.base import BackendType
 from dstack._internal.core.models.common import NetworkMode, RegistryAuth
-from dstack._internal.core.models.configurations import DevEnvironmentConfiguration
+from dstack._internal.core.models.configurations import (
+    DevEnvironmentConfiguration,
+    ServiceConfiguration,
+)
 from dstack._internal.core.models.files import FileArchiveMapping
 from dstack._internal.core.models.instances import InstanceStatus, SSHConnectionParams
 from dstack._internal.core.models.metrics import Metric
@@ -991,6 +994,22 @@ async def _register_service_replica(
     if context.run_model.gateway_id is None:
         return None
 
+    job_spec = JobSpec.__response__.parse_raw(context.job_model.job_spec_data)
+
+    # For router-based services (e.g. PD disaggregation), only router replicas should be
+    # registered with the gateway. Worker replicas are discovered by the router-worker
+    # sync pipeline and should not be routed to directly by the gateway.
+    config = context.run.run_spec.configuration
+    assert isinstance(config, ServiceConfiguration)
+    router_group = next((g for g in config.replica_groups if g.router is not None), None)
+    is_router_replica = router_group is not None and job_spec.replica_group == router_group.name
+    if router_group is not None and not is_router_replica:
+        logger.debug(
+            "%s: skipping gateway replica registration (non-router replica)",
+            fmt(context.job_model),
+        )
+        return None
+
     async with get_session_ctx() as session:
         gateway_model, conn = await get_or_add_gateway_connection(
             session, context.run_model.gateway_id
@@ -1011,7 +1030,7 @@ async def _register_service_replica(
         async with conn.client() as gateway_client:
             await gateway_client.register_replica(
                 run=context.run,
-                job_spec=JobSpec.__response__.parse_raw(context.job_model.job_spec_data),
+                job_spec=job_spec,
                 job_submission=job_submission,
                 instance_project_ssh_private_key=instance_project_ssh_private_key,
                 ssh_head_proxy=ssh_head_proxy,

src/dstack/_internal/server/background/pipeline_tasks/service_router_worker_sync.py

@@ -4,7 +4,7 @@
 from datetime import timedelta
 from typing import Sequence
 
-from sqlalchemy import or_, select, true, update
+from sqlalchemy import or_, select, update
 from sqlalchemy.orm import joinedload, load_only, selectinload
 from sqlalchemy.sql import false
 
@@ -224,7 +224,6 @@ async def process(self, item: ServiceRouterWorkerSyncPipelineItem) -> None:
                     selectinload(
                         RunModel.jobs.and_(
                             JobModel.status == JobStatus.RUNNING,
-                            JobModel.registered == true(),
                         )
                     )
                     .load_only(

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

@@ -17,7 +17,6 @@
 from dstack._internal.server.services.logging import fmt
 from dstack._internal.server.services.runs import run_spec_has_router_replica_group
 from dstack._internal.server.services.runs.replicas import (
-    is_replica_registered,
     job_belongs_to_group,
 )
 from dstack._internal.utils.logging import get_logger
@@ -107,7 +106,7 @@ def _get_router_job(run_model: RunModel, router_group) -> Optional[JobModel]:
         for j in run_model.jobs
         if job_belongs_to_group(j, group_name) and j.status == JobStatus.RUNNING
     ]
-    if not router_jobs or not is_replica_registered(router_jobs):
+    if not router_jobs:
         return None
     # Router replica group is currently validated to have count=1, so we assume a single active
     # router job here. When we support multiple router replicas for HA, this should be updated
@@ -297,8 +296,6 @@ async def _build_target_workers(
                 continue
             if job.status != JobStatus.RUNNING:
                 continue
-            if not is_replica_registered([job]):
-                continue
             jpd = get_job_provisioning_data(job)
             if jpd is None:
                 continue