diff --git a/src/dstack/_internal/server/background/pipeline_tasks/instances/termination.py b/src/dstack/_internal/server/background/pipeline_tasks/instances/termination.py index a4bf6d329..c4cce13b2 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/instances/termination.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/instances/termination.py @@ -12,7 +12,6 @@ from dstack._internal.server.services import backends as backends_services from dstack._internal.server.services.instances import get_instance_provisioning_data from dstack._internal.server.services.runner.pool import ( - InstanceConnectionKey, instance_connection_pool, ) from dstack._internal.utils.common import get_current_datetime, run_async @@ -82,7 +81,7 @@ async def terminate_instance(instance_model: InstanceModel) -> ProcessResult: ) if job_provisioning_data is not None: - instance_connection_pool.drop(InstanceConnectionKey.from_jpd(job_provisioning_data)) + instance_connection_pool.drop_by_jpd(job_provisioning_data) result.instance_update_map["deleted"] = True result.instance_update_map["deleted_at"] = NOW_PLACEHOLDER diff --git a/src/dstack/_internal/server/services/runner/pool.py b/src/dstack/_internal/server/services/runner/pool.py index b91d1d312..e3a012967 100644 --- a/src/dstack/_internal/server/services/runner/pool.py +++ b/src/dstack/_internal/server/services/runner/pool.py @@ -130,6 +130,12 @@ def drop(self, key: InstanceConnectionKey) -> None: except Exception: logger.exception("Failed to close instance connection %s", key) + def drop_by_jpd(self, jpd: JobProvisioningData, jrd: Optional[JobRuntimeData] = None): + if jpd.hostname is None or jpd.ssh_port is None: + return + key = InstanceConnectionKey.from_jpd(jpd, jrd) + self.drop(key) + def startup_cleanup(self) -> None: """ Removes connection dirs left by a previous server process (e.g. after SIGKILL).