diff --git a/src/dstack/_internal/server/background/pipeline_tasks/jobs_running.py b/src/dstack/_internal/server/background/pipeline_tasks/jobs_running.py index 720b0141b..99192af16 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/jobs_running.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/jobs_running.py @@ -20,7 +20,6 @@ ) 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 @@ -34,7 +33,6 @@ Job, JobProvisioningData, JobRuntimeData, - JobSpec, JobStatus, JobSubmission, JobTerminationReason, @@ -366,6 +364,7 @@ class _JobUpdateMap(ItemUpdateMap, total=False): disconnected_at: Optional[datetime] inactivity_secs: Optional[int] exit_status: Optional[int] + ready: bool registered: bool image_pull_progress: Optional[str] skip_min_processing_interval: bool @@ -1083,6 +1082,12 @@ def _emit_result_events( job_model.disconnected_at, ), ) + _emit_readiness_change_event( + session=session, + job_model=job_model, + old_ready=job_model.ready, + new_ready=result.job_update_map.get("ready", job_model.ready), + ) if result.replica_registration is not None: targets = [events.Target.from_model(job_model)] if result.replica_registration.gateway_target is not None: @@ -1175,13 +1180,32 @@ async def _maybe_register_replica( ) -> None: if ( context.run.run_spec.configuration.type != "service" - or _get_result_registered(context.job_model, result) or context.job_model.job_num != 0 or result.new_probe_models - or not is_job_ready(context.job_model.probes, context.job.job_spec.probes) ): return + is_ready = is_job_ready(context.job_model.probes, context.job.job_spec.probes) + if is_ready and not context.job_model.ready: + result.job_update_map["ready"] = True + + router_group = next( + (g for g in context.run.run_spec.configuration.replica_groups if g.router is not None), + None, + ) + is_router_replica = ( + router_group is not None and context.job.job_spec.replica_group == router_group.name + ) + # non-router replicas aren't registered if the service has a router + if router_group is not None and not is_router_replica: + if context.job_model.registered: + # migration edge case: a pre-0.21.1 server replica incorrectly set registered=True + result.job_update_map["registered"] = False + return + + if not is_ready or _get_result_registered(context.job_model, result): + return + ssh_head_proxy: Optional[SSHConnectionParams] = None ssh_head_proxy_private_key: Optional[str] = None instance = get_or_error(context.job_model.instance) @@ -1220,23 +1244,6 @@ async def _register_service_replica( ) -> Optional[events.Target]: if context.run_model.gateway_id is None: return None - - job_spec = validate_json_extra_ignore(JobSpec, 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, connections = await get_or_add_gateway_connections( session, context.run_model.gateway_id @@ -1261,7 +1268,7 @@ async def _register_service_replica( async with conn.client() as gateway_client: await gateway_client.register_replica( run=context.run, - job_spec=job_spec, + job_spec=context.job.job_spec, job_submission=job_submission, instance_project_ssh_private_key=instance_project_ssh_private_key, ssh_head_proxy=ssh_head_proxy, @@ -1877,6 +1884,23 @@ def _emit_reachability_change_event( ) +def _emit_readiness_change_event( + session: AsyncSession, + job_model: JobModel, + old_ready: bool, + new_ready: bool, +) -> None: + # ready: False -> True + if not old_ready and new_ready: + events.emit( + session, + "Service replica ready to receive requests", + actor=events.SystemActor(), + targets=[events.Target.from_model(job_model)], + ) + # ready: True -> False is not possible as of this writing + + def _terminate_job( job_model: JobModel, job_update_map: _JobUpdateMap, diff --git a/src/dstack/_internal/server/background/pipeline_tasks/runs/active.py b/src/dstack/_internal/server/background/pipeline_tasks/runs/active.py index aafe132ea..1ad52e05f 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/runs/active.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/runs/active.py @@ -686,11 +686,11 @@ async def _build_rolling_deployment_maps( max_new = max(j.replica_num for j in new_jobs) next_replica_num = max(next_replica_num, max_new + 1) - # Scale down: terminate unregistered out-of-date + excess registered replicas - replicas_to_stop = state.unregistered_out_of_date_replica_count + # Scale down: terminate unready out-of-date + excess ready replicas + replicas_to_stop = state.unready_out_of_date_replica_count replicas_to_stop += max( 0, - state.registered_non_terminating_replica_count - group_desired, + state.ready_non_terminating_replica_count - group_desired, ) if replicas_to_stop > 0: scale_down_maps = _build_scale_down_job_update_maps( diff --git a/src/dstack/_internal/server/background/pipeline_tasks/service_router_worker_sync.py b/src/dstack/_internal/server/background/pipeline_tasks/service_router_worker_sync.py index 5ded45afe..48fe52210 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/service_router_worker_sync.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/service_router_worker_sync.py @@ -233,13 +233,12 @@ async def process(self, item: ServiceRouterWorkerSyncPipelineItem) -> None: selectinload( RunModel.jobs.and_( JobModel.status == JobStatus.RUNNING, - JobModel.registered == True, + JobModel.ready == True, ) ) .load_only( JobModel.id, JobModel.status, - JobModel.registered, JobModel.job_spec_data, JobModel.job_provisioning_data, JobModel.job_runtime_data, diff --git a/src/dstack/_internal/server/migrations/versions/2026/08_04_0702_72cfa56364ad_add_jobmodel_ready.py b/src/dstack/_internal/server/migrations/versions/2026/08_04_0702_72cfa56364ad_add_jobmodel_ready.py new file mode 100644 index 000000000..61195dd93 --- /dev/null +++ b/src/dstack/_internal/server/migrations/versions/2026/08_04_0702_72cfa56364ad_add_jobmodel_ready.py @@ -0,0 +1,120 @@ +"""Add JobModel.ready + +Revision ID: 72cfa56364ad +Revises: a1c3f5e7b209 +Create Date: 2026-08-04 07:02:02.041000+00:00 + +""" + +import json +import uuid +from typing import Optional + +import sqlalchemy as sa +from alembic import op +from sqlalchemy_utils import UUIDType + +# revision identifiers, used by Alembic. +revision = "72cfa56364ad" +down_revision = "a1c3f5e7b209" +branch_labels = None +depends_on = None + +# Partial table descriptions - only columns needed for the data migration below. +jobs_table = sa.Table( + "jobs", + sa.MetaData(), + sa.Column("id", UUIDType(binary=False), primary_key=True, default=uuid.uuid4), + sa.Column("run_id", UUIDType(binary=False)), + sa.Column("status", sa.String(100)), + sa.Column("registered", sa.Boolean()), + sa.Column("ready", sa.Boolean()), + sa.Column("job_spec_data", sa.Text()), +) +runs_table = sa.Table( + "runs", + sa.MetaData(), + sa.Column("id", UUIDType(binary=False), primary_key=True, default=uuid.uuid4), + sa.Column("run_spec", sa.Text()), + sa.Column("service_spec", sa.Text(), nullable=True), +) + + +def _get_router_group_name(run_spec_data: str) -> Optional[str]: + configuration = json.loads(run_spec_data).get("configuration") or {} + if configuration.get("type") != "service": + return None + replica_groups = configuration.get("replicas") + if not isinstance(replica_groups, list): + return None + for group in replica_groups: + if isinstance(group, dict) and group.get("router") is not None: + return group.get("name") + return None + + +def _get_job_replica_group(job_spec_data: str) -> str: + return json.loads(job_spec_data).get("replica_group", "0") + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + with op.batch_alter_table("jobs", schema=None) as batch_op: + batch_op.add_column( + sa.Column("ready", sa.Boolean(), server_default=sa.false(), nullable=False) + ) + + # ### end Alembic commands ### + + bind = op.get_bind() + + # backfill ready=True for existing registered replicas + + bind.execute(jobs_table.update().where(jobs_table.c.registered == True).values(ready=True)) + + # set registered=False for non-router replicas in services with a router + + candidate_jobs = bind.execute( + sa.select(jobs_table.c.id, jobs_table.c.run_id, jobs_table.c.job_spec_data) + .select_from(jobs_table.join(runs_table, jobs_table.c.run_id == runs_table.c.id)) + .where( + jobs_table.c.registered == True, + jobs_table.c.status == "RUNNING", + runs_table.c.service_spec.is_not(None), + sa.or_( + runs_table.c.run_spec.like('%"sglang"%'), + runs_table.c.run_spec.like('%"dynamo"%'), + ), + ) + ).all() + + run_ids = {row.run_id for row in candidate_jobs} + router_group_name_by_run_id = {} + if run_ids: + for row in bind.execute( + sa.select(runs_table.c.id, runs_table.c.run_spec).where(runs_table.c.id.in_(run_ids)) + ).all(): + router_group_name_by_run_id[row.id] = _get_router_group_name(row.run_spec) + + non_router_job_ids = [] + for row in candidate_jobs: + router_group_name = router_group_name_by_run_id.get(row.run_id) + if router_group_name is None: + continue + if _get_job_replica_group(row.job_spec_data) != router_group_name: + non_router_job_ids.append(row.id) + + if non_router_job_ids: + bind.execute( + jobs_table.update() + .where(jobs_table.c.id.in_(non_router_job_ids)) + .values(registered=False) + ) + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + with op.batch_alter_table("jobs", schema=None) as batch_op: + batch_op.drop_column("ready") + + # ### end Alembic commands ### diff --git a/src/dstack/_internal/server/models.py b/src/dstack/_internal/server/models.py index fd0a7a155..23fb31ed7 100644 --- a/src/dstack/_internal/server/models.py +++ b/src/dstack/_internal/server/models.py @@ -578,9 +578,14 @@ class JobModel(PipelineModelMixin, BaseModel): probes: Mapped[list["ProbeModel"]] = relationship( back_populates="job", order_by="ProbeModel.probe_num" ) + ready: Mapped[bool] = mapped_column(Boolean, server_default=false()) + """Whether the replica is ready to receive service requests based on probe statuses. + Always `False` for non-service runs. + """ registered: Mapped[bool] = mapped_column(Boolean, server_default=false()) - """`registered` shows whether the replica is registered to receive service requests. - It is always `False` for non-service runs. + """Whether the replica is registered to receive service requests from dstack-proxy. + Always `False` for non-service runs or jobs that shouldn't be registered + (e.g., non-router replicas for services with routers). """ waiting_master_job: Mapped[Optional[bool]] = mapped_column(Boolean) """`waiting_master_job` is `True` for non-master jobs that have to wait for master processing before diff --git a/src/dstack/_internal/server/services/proxy/repo.py b/src/dstack/_internal/server/services/proxy/repo.py index 5a0caffd7..d0afd41de 100644 --- a/src/dstack/_internal/server/services/proxy/repo.py +++ b/src/dstack/_internal/server/services/proxy/repo.py @@ -113,10 +113,6 @@ async def get_service(self, project_name: str, run_name: str) -> Optional[Servic ssh_head_proxy = rci.ssh_proxy ssh_head_proxy_private_key = get_or_error(rci.ssh_proxy_keys)[0].private job_spec = get_job_spec(job) - if router_group is not None and job_spec.replica_group != router_group.name: - # Strict router-only: when a router is configured, the proxy should only be aware - # of router replicas. - continue replica = Replica( id=job.id.hex, app_port=get_service_port(job_spec, run_spec.configuration), diff --git a/src/dstack/_internal/server/services/runs/replicas.py b/src/dstack/_internal/server/services/runs/replicas.py index 932088152..633b95999 100644 --- a/src/dstack/_internal/server/services/runs/replicas.py +++ b/src/dstack/_internal/server/services/runs/replicas.py @@ -19,8 +19,8 @@ class GroupRolloutState: inactive_replicas: List[Tuple[int, bool, int, List[JobModel]]] has_out_of_date_replicas: bool non_terminated_replica_count: int - unregistered_out_of_date_replica_count: int - registered_non_terminating_replica_count: int + unready_out_of_date_replica_count: int + ready_non_terminating_replica_count: int class RouterEnvStatus(str, Enum): @@ -77,7 +77,7 @@ def build_replica_lists( elif {JobStatus.PROVISIONING, JobStatus.PULLING} & statuses: # if there are any provisioning or pulling jobs, the replica is active and has the importance of 1 active_replicas.append((1, is_out_of_date, replica_num, replica_jobs)) - elif not is_replica_registered(replica_jobs): + elif not is_replica_ready(replica_jobs): # all jobs are running, but not receiving traffic, the replica is active and has the importance of 2 active_replicas.append((2, is_out_of_date, replica_num, replica_jobs)) else: @@ -98,8 +98,8 @@ def get_group_rollout_state(run_model: RunModel, group: ReplicaGroup) -> GroupRo ) non_terminated_replica_nums = set() - unregistered_out_of_date_replica_count = 0 - registered_non_terminating_replica_count = 0 + unready_out_of_date_replica_count = 0 + ready_non_terminating_replica_count = 0 for _, jobs in group_jobs_by_replica_latest(run_model.jobs): if not job_belongs_to_group(jobs[0], group.name): @@ -114,20 +114,20 @@ def get_group_rollout_state(run_model: RunModel, group: ReplicaGroup) -> GroupRo j.status not in [JobStatus.TERMINATING] + JobStatus.finished_statuses() for j in jobs ) - and not is_replica_registered(jobs) + and not is_replica_ready(jobs) ): - unregistered_out_of_date_replica_count += 1 + unready_out_of_date_replica_count += 1 - if is_replica_registered(jobs) and all(j.status != JobStatus.TERMINATING for j in jobs): - registered_non_terminating_replica_count += 1 + if is_replica_ready(jobs) and all(j.status != JobStatus.TERMINATING for j in jobs): + ready_non_terminating_replica_count += 1 return GroupRolloutState( active_replicas=active_replicas, inactive_replicas=inactive_replicas, has_out_of_date_replicas=has_out_of_date_replicas(run_model, group_filter=group.name), non_terminated_replica_count=len(non_terminated_replica_nums), - unregistered_out_of_date_replica_count=unregistered_out_of_date_replica_count, - registered_non_terminating_replica_count=registered_non_terminating_replica_count, + unready_out_of_date_replica_count=unready_out_of_date_replica_count, + ready_non_terminating_replica_count=ready_non_terminating_replica_count, ) @@ -149,9 +149,9 @@ def has_out_of_date_replicas(run: RunModel, group_filter: Optional[str] = None) return False -def is_replica_registered(jobs: list[JobModel]) -> bool: +def is_replica_ready(jobs: list[JobModel]) -> bool: # Only job_num=0 is supposed to receive service requests - return jobs[0].registered + return jobs[0].ready def get_router_replica_group(run_spec: RunSpec) -> Optional[ReplicaGroup]: diff --git a/src/dstack/_internal/server/services/runs/router_worker_sync.py b/src/dstack/_internal/server/services/runs/router_worker_sync.py index 4b9b8af65..9876d2415 100644 --- a/src/dstack/_internal/server/services/runs/router_worker_sync.py +++ b/src/dstack/_internal/server/services/runs/router_worker_sync.py @@ -1,4 +1,4 @@ -"""Reconcile SGLang router /workers with dstack's registered worker replicas (async, SSH-tunneled).""" +"""Reconcile SGLang router /workers with dstack's ready worker replicas (async, SSH-tunneled).""" import json from typing import Any, List, Literal, Optional, TypedDict diff --git a/src/dstack/_internal/server/testing/common.py b/src/dstack/_internal/server/testing/common.py index 46969a8e0..599fbd561 100644 --- a/src/dstack/_internal/server/testing/common.py +++ b/src/dstack/_internal/server/testing/common.py @@ -442,13 +442,21 @@ async def create_job( instance_assigned: bool = False, disconnected_at: Optional[datetime] = None, registered: bool = False, + ready: bool = False, waiting_master_job: Optional[bool] = None, + replica_group_name: Optional[str] = None, ) -> JobModel: + assert not (registered and not ready), "registered=True with ready=False is invalid" if deployment_num is None: deployment_num = run.deployment_num run_spec = validate_json_extra_ignore(RunSpec, run.run_spec) job_spec = ( - await get_job_specs_from_run_spec(run_spec=run_spec, secrets={}, replica_num=replica_num) + await get_job_specs_from_run_spec( + run_spec=run_spec, + secrets={}, + replica_num=replica_num, + replica_group_name=replica_group_name, + ) )[0] job_spec.job_num = job_num job = JobModel( @@ -476,6 +484,7 @@ async def create_job( disconnected_at=disconnected_at, probes=[], registered=registered, + ready=ready, waiting_master_job=waiting_master_job, ) session.add(job) diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_running_jobs.py b/src/tests/_internal/server/background/pipeline_tasks/test_running_jobs.py index 034dbf7c2..676223793 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_running_jobs.py +++ b/src/tests/_internal/server/background/pipeline_tasks/test_running_jobs.py @@ -69,6 +69,7 @@ from dstack._internal.server.services.runs.replicas import RouterEnvStatus from dstack._internal.server.services.volumes import volume_model_to_volume from dstack._internal.server.testing.common import ( + clear_events, create_backend, create_code, create_export, @@ -1969,6 +1970,7 @@ async def test_registers_service_replica_immediately_if_no_probes( events = await list_events(session) assert {event.message for event in events} == { "Job status changed PULLING -> RUNNING", + "Service replica ready to receive requests", "Service replica registered to receive requests", } @@ -2058,8 +2060,11 @@ async def test_registers_service_replica_only_after_probes_pass( events = await list_events(session) if expect_to_register: assert job.registered - assert len(events) == 1 - assert events[0].message == "Service replica registered to receive requests" + assert len(events) == 2 + assert {event.message for event in events} == { + "Service replica ready to receive requests", + "Service replica registered to receive requests", + } else: assert not job.registered assert not events @@ -2130,6 +2135,7 @@ async def test_registers_service_replica_in_gateway( events = await list_events(session) assert {event.message for event in events} == { "Job status changed PULLING -> RUNNING", + "Service replica ready to receive requests", "Service replica registered to receive requests", } mock_gateway_connection.return_value.client.return_value.__aenter__.return_value.register_replica.assert_called_once_with( @@ -2216,6 +2222,7 @@ async def test_registers_service_replica_in_gateway_when_running_on_imported_ins events = await list_events(session) assert {event.message for event in events} == { "Job status changed PULLING -> RUNNING", + "Service replica ready to receive requests", "Service replica registered to receive requests", } mock_gateway_connection.return_value.client.return_value.__aenter__.return_value.register_replica.assert_called_once_with( @@ -2456,13 +2463,168 @@ async def test_provisioning_shim_uses_server_default_registry( assert call_kwargs["registry_username"] == "server-user" assert call_kwargs["registry_password"] == "server-pass" + async def test_registers_router_replica_but_not_worker_replica_in_gateway( + self, + test_db, + session: AsyncSession, + worker: JobRunningWorker, + ssh_tunnel_mock: Mock, + runner_client_mock: Mock, + mock_gateway_connection: AsyncMock, + ): + user = await create_user(session=session) + project = await create_project(session=session, owner=user) + repo = await create_repo(session=session, project_id=project.id) + backend = await create_backend(session=session, project_id=project.id) + gateway = await create_gateway( + session=session, + project_id=project.id, + backend_id=backend.id, + status=GatewayStatus.RUNNING, + name="test-gateway", + wildcard_domain="example.com", + ) + await create_gateway_compute( + session=session, + backend_id=backend.id, + gateway_id=gateway.id, + ) + run = await create_run( + session=session, + project=project, + repo=repo, + user=user, + run_spec=get_run_spec( + run_name="test", + repo_id=repo.name, + configuration=_router_service_configuration("sglang", gateway="test-gateway"), + ), + gateway=gateway, + ) + instance = await create_instance( + session=session, + project=project, + status=InstanceStatus.BUSY, + ) + router_job = await create_job( + session=session, + run=run, + replica_num=0, + replica_group_name="router", + status=JobStatus.RUNNING, + job_provisioning_data=get_job_provisioning_data(dockerized=True), + instance=instance, + instance_assigned=True, + ) + worker_job = await create_job( + session=session, + run=run, + replica_num=1, + replica_group_name="worker", + status=JobStatus.RUNNING, + job_provisioning_data=get_job_provisioning_data(dockerized=True), + instance=instance, + instance_assigned=True, + ) + runner_client_mock.pull.return_value = PullResponse( + job_states=[], job_logs=[], runner_logs=[], last_updated=0 + ) + + await _process_job(session, worker, router_job) + + await session.refresh(router_job) + assert router_job.registered + assert router_job.ready + events = await list_events(session) + assert {event.message for event in events} == { + "Service replica ready to receive requests", + "Service replica registered to receive requests", + } + + await clear_events(session) + + await _process_job(session, worker, worker_job) + + await session.refresh(worker_job) + assert not worker_job.registered + assert worker_job.ready + events = await list_events(session) + assert {event.message for event in events} == { + "Service replica ready to receive requests", + } + + async def test_resets_stale_registered_flag_for_non_router_replica( + self, + test_db, + session: AsyncSession, + worker: JobRunningWorker, + ssh_tunnel_mock: Mock, + runner_client_mock: Mock, + ): + """Migration edge case: a pre-0.21.1 server may have incorrectly marked + a non-router replica as registered. Should be corrected. + """ + project = await create_project(session=session) + user = await create_user(session=session) + repo = await create_repo(session=session, project_id=project.id) + run = await create_run( + session=session, + project=project, + repo=repo, + user=user, + run_spec=get_run_spec( + run_name="test", + repo_id=repo.name, + configuration=_router_service_configuration( + "sglang", + probes=[ProbeConfig(type="http", url="/health", ready_after=1)], + ), + ), + ) + instance = await create_instance( + session=session, + project=project, + status=InstanceStatus.BUSY, + ) + worker_job = await create_job( + session=session, + run=run, + replica_num=1, + replica_group_name="worker", + status=JobStatus.RUNNING, + job_provisioning_data=get_job_provisioning_data(dockerized=True), + instance=instance, + instance_assigned=True, + registered=True, + ready=True, + ) + await create_probe(session=session, job=worker_job, probe_num=0, success_streak=0) + runner_client_mock.pull.return_value = PullResponse( + job_states=[], job_logs=[], runner_logs=[], last_updated=0 + ) + + await _process_job(session, worker, worker_job) + + await session.refresh(worker_job) + assert worker_job.status == JobStatus.RUNNING + assert not worker_job.registered + events = await list_events(session) + assert events == [] + -def _router_service_configuration(router_type: str) -> ServiceConfiguration: +def _router_service_configuration( + router_type: str, + *, + gateway: Optional[str] = None, + probes: Optional[list[ProbeConfig]] = None, +) -> ServiceConfiguration: return ServiceConfiguration.model_validate( { "type": "service", "port": 8000, "image": "ubuntu", + "gateway": gateway, + "probes": probes, "replicas": [ {"name": "worker", "commands": ["echo worker"], "count": 1}, { diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_runs/test_active.py b/src/tests/_internal/server/background/pipeline_tasks/test_runs/test_active.py index 470b7824d..ed085e09d 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_runs/test_active.py +++ b/src/tests/_internal/server/background/pipeline_tasks/test_runs/test_active.py @@ -950,7 +950,7 @@ async def test_service_rolling_deployment_scale_up( self, test_db, session: AsyncSession, worker: RunWorker ) -> None: """Service with 1 out-of-date RUNNING replica whose spec differs from the new - deployment, desired=1 → creates 1 new replica (surge), old registered replica + deployment, desired=1 → creates 1 new replica (surge), old ready replica untouched.""" project = await create_project(session=session) user = await create_user(session=session) @@ -979,6 +979,7 @@ async def test_service_rolling_deployment_scale_up( status=JobStatus.RUNNING, deployment_num=0, registered=True, + ready=True, replica_num=0, ) # Make the old job's spec differ from the current run_spec so in-place bump @@ -1001,7 +1002,7 @@ async def test_service_rolling_deployment_scale_up( ) jobs = list(res.scalars().all()) assert len(jobs) == 2 - # Old replica still RUNNING (registered, not terminated during rolling) + # Old replica still RUNNING (ready, not terminated during rolling) assert jobs[0].status == JobStatus.RUNNING assert jobs[0].deployment_num == 0 # New surge replica created @@ -1041,6 +1042,7 @@ async def test_service_rolling_deployment_scale_down_old_unregistered( status=JobStatus.RUNNING, deployment_num=1, registered=True, + ready=True, replica_num=0, ) # Out-of-date unregistered replica with different spec @@ -1050,6 +1052,7 @@ async def test_service_rolling_deployment_scale_down_old_unregistered( status=JobStatus.RUNNING, deployment_num=0, registered=False, + ready=False, replica_num=1, ) old_spec = get_job_spec(old_job) @@ -1069,6 +1072,207 @@ async def test_service_rolling_deployment_scale_down_old_unregistered( assert old_job.status == JobStatus.TERMINATING assert old_job.termination_reason == JobTerminationReason.SCALED_DOWN + async def test_service_router_rolling_deployment_surges_ready_worker_replica( + self, test_db, session: AsyncSession, worker: RunWorker + ) -> None: + project = await create_project(session=session) + user = await create_user(session=session) + repo = await create_repo(session=session, project_id=project.id) + run_spec = get_run_spec( + repo_id=repo.name, + run_name="service-run", + configuration=_router_worker_service_configuration(), + ) + run = await create_run( + session=session, + project=project, + repo=repo, + user=user, + run_name="service-run", + run_spec=run_spec, + status=RunStatus.RUNNING, + deployment_num=1, + ) + # Up-to-date router replica + await create_job( + session=session, + run=run, + status=JobStatus.RUNNING, + deployment_num=1, + registered=True, + ready=True, + replica_num=0, + replica_group_name="router", + ) + # Out-of-date worker replica: ready to serve traffic but never registered. + old_worker_job = await create_job( + session=session, + run=run, + status=JobStatus.RUNNING, + deployment_num=0, + registered=False, + ready=True, + replica_num=1, + replica_group_name="worker", + ) + old_spec = get_job_spec(old_worker_job) + old_spec.commands = ["echo old worker!"] + old_worker_job.job_spec_data = old_spec.model_dump_json() + await session.commit() + + lock_run(run) + await session.commit() + + await worker.process(run_to_pipeline_item(run)) + + await session.refresh(run) + assert run.status == RunStatus.RUNNING + + await session.refresh(old_worker_job) + assert old_worker_job.status == JobStatus.RUNNING + assert old_worker_job.ready + assert not old_worker_job.registered + assert old_worker_job.deployment_num == 0 + + res = await session.execute( + select(JobModel).where( + JobModel.run_id == run.id, + JobModel.replica_num == 2, + ) + ) + new_worker_job = res.scalar_one() + assert new_worker_job.status == JobStatus.SUBMITTED + assert new_worker_job.deployment_num == 1 + + async def test_service_router_rolling_deployment_scales_down_unready_worker_replica( + self, test_db, session: AsyncSession, worker: RunWorker + ) -> None: + project = await create_project(session=session) + user = await create_user(session=session) + repo = await create_repo(session=session, project_id=project.id) + run_spec = get_run_spec( + repo_id=repo.name, + run_name="service-run", + configuration=_router_worker_service_configuration(), + ) + run = await create_run( + session=session, + project=project, + repo=repo, + user=user, + run_name="service-run", + run_spec=run_spec, + status=RunStatus.RUNNING, + deployment_num=1, + ) + await create_job( + session=session, + run=run, + status=JobStatus.RUNNING, + deployment_num=1, + registered=True, + ready=True, + replica_num=0, + replica_group_name="router", + ) + # Out-of-date worker replica that never became ready. + old_worker_job = await create_job( + session=session, + run=run, + status=JobStatus.RUNNING, + deployment_num=0, + registered=False, + ready=False, + replica_num=1, + replica_group_name="worker", + ) + old_spec = get_job_spec(old_worker_job) + old_spec.commands = ["echo old worker!"] + old_worker_job.job_spec_data = old_spec.model_dump_json() + await session.commit() + + lock_run(run) + await session.commit() + + await worker.process(run_to_pipeline_item(run)) + + await session.refresh(run) + assert run.status == RunStatus.RUNNING + + await session.refresh(old_worker_job) + assert old_worker_job.status == JobStatus.TERMINATING + assert old_worker_job.termination_reason == JobTerminationReason.SCALED_DOWN + + async def test_service_router_rolling_deployment_terminates_out_of_date_worker_once_replacement_ready( + self, test_db, session: AsyncSession, worker: RunWorker + ) -> None: + project = await create_project(session=session) + user = await create_user(session=session) + repo = await create_repo(session=session, project_id=project.id) + run_spec = get_run_spec( + repo_id=repo.name, + run_name="service-run", + configuration=_router_worker_service_configuration(), + ) + run = await create_run( + session=session, + project=project, + repo=repo, + user=user, + run_name="service-run", + run_spec=run_spec, + status=RunStatus.RUNNING, + deployment_num=1, + ) + await create_job( + session=session, + run=run, + status=JobStatus.RUNNING, + deployment_num=1, + registered=True, + ready=True, + replica_num=0, + replica_group_name="router", + ) + # Out-of-date worker replica, still ready and serving traffic (surge in progress). + old_worker_job = await create_job( + session=session, + run=run, + status=JobStatus.RUNNING, + deployment_num=0, + registered=False, + ready=True, + replica_num=1, + replica_group_name="worker", + ) + old_spec = get_job_spec(old_worker_job) + old_spec.commands = ["echo old worker!"] + old_worker_job.job_spec_data = old_spec.model_dump_json() + # Up-to-date surge worker replica has become ready → the old replica is now excess. + await create_job( + session=session, + run=run, + status=JobStatus.RUNNING, + deployment_num=1, + registered=False, + ready=True, + replica_num=2, + replica_group_name="worker", + ) + await session.commit() + + lock_run(run) + await session.commit() + + await worker.process(run_to_pipeline_item(run)) + + await session.refresh(run) + assert run.status == RunStatus.RUNNING + + await session.refresh(old_worker_job) + assert old_worker_job.status == JobStatus.TERMINATING + assert old_worker_job.termination_reason == JobTerminationReason.SCALED_DOWN + async def test_service_removed_group_cleanup( self, test_db, session: AsyncSession, worker: RunWorker ) -> None: @@ -1126,3 +1330,22 @@ async def test_service_removed_group_cleanup( await session.refresh(old_group_job) assert old_group_job.status == JobStatus.TERMINATING assert old_group_job.termination_reason == JobTerminationReason.SCALED_DOWN + + +def _router_worker_service_configuration() -> ServiceConfiguration: + return ServiceConfiguration.model_validate( + { + "type": "service", + "port": 8000, + "image": "ubuntu", + "replicas": [ + {"name": "worker", "commands": ["echo worker"], "count": 1}, + { + "name": "router", + "router": {"type": "sglang"}, + "commands": ["echo router"], + "count": 1, + }, + ], + } + )