Drop connection on instance termination

Author: r4victor

src/dstack/_internal/server/background/pipeline_tasks/instances/termination.py

@@ -11,6 +11,10 @@
 from dstack._internal.server.models import InstanceModel
 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
 from dstack._internal.utils.logging import get_logger
 
@@ -77,6 +81,9 @@ async def terminate_instance(instance_model: InstanceModel) -> ProcessResult:
                     exc_info=not isinstance(exc, BackendError),
                 )
 
+    if job_provisioning_data is not None:
+        instance_connection_pool.drop(InstanceConnectionKey.from_jpd(job_provisioning_data))
+
     result.instance_update_map["deleted"] = True
     result.instance_update_map["deleted_at"] = NOW_PLACEHOLDER
     result.instance_update_map["finished_at"] = NOW_PLACEHOLDER

src/dstack/_internal/server/services/runner/pool.py

@@ -45,7 +45,7 @@ class InstanceConnectionKey:
 
     @staticmethod
     def from_jpd(
-        jpd: JobProvisioningData, jrd: Optional[JobRuntimeData]
+        jpd: JobProvisioningData, jrd: Optional[JobRuntimeData] = None
     ) -> "InstanceConnectionKey":
         assert jpd.hostname is not None and jpd.ssh_port is not None
         container_to_host_port_map = InstanceConnection.get_container_to_host_port_map(jpd, jrd)