Merge pull request #9249 from OpenMined/shubham/pool-different-image

Re-creating pool with different image spins up with last attached image
OpenMined · Sep 3, 2024 · 8f09272 · 8f09272
2 parents 90b00b7 + 2f5a804
commit 8f09272
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/packages/syft/src/syft/custom_worker/runner_k8s.py b/packages/syft/src/syft/custom_worker/runner_k8s.py
@@ -13,7 +13,7 @@
 from .k8s import get_kr8s_client
 
 JSONPATH_AVAILABLE_REPLICAS = "{.status.availableReplicas}"
-CREATE_POOL_TIMEOUT_SEC = 180
+CREATE_POOL_TIMEOUT_SEC = 380
 SCALE_POOL_TIMEOUT_SEC = 60
 
 
@@ -96,9 +96,11 @@ def delete_pool(self, pool_name: str) -> bool:
         selector = {"app.kubernetes.io/component": pool_name}
         for _set in self.client.get("statefulsets", label_selector=selector):
             _set.delete(propagation_policy="Foreground")
+            _set.wait(conditions="delete")
 
         for _secret in self.client.get("secrets", label_selector=selector):
             _secret.delete(propagation_policy="Foreground")
+            _secret.wait(conditions="delete")
 
         return True
 

diff --git a/packages/syft/src/syft/service/worker/worker_pool_service.py b/packages/syft/src/syft/service/worker/worker_pool_service.py
@@ -623,7 +623,10 @@ def delete(
         )
 
         if IN_KUBERNETES:
+            # Scale the workers to zero
             self.scale(context=context, number=0, pool_id=uid)
+            runner = KubernetesRunner()
+            runner.delete_pool(pool_name=worker_pool.name)
         else:
             workers = (
                 worker.resolve_with_context(context=context).unwrap()