From 3a8b84967b6d4cdd0a5cf55ea30af42e1c06b706 Mon Sep 17 00:00:00 2001
From: Yash Gorana <yash.gorana@hotmail.com>
Date: Sat, 3 Feb 2024 02:40:59 +0530
Subject: [PATCH] [k8s] update kr8s with latest bugfix

---
 packages/syft/setup.cfg                       |  2 +-
 .../syft/src/syft/custom_worker/runner_k8s.py | 27 ++-----------------
 2 files changed, 3 insertions(+), 26 deletions(-)

diff --git a/packages/syft/setup.cfg b/packages/syft/setup.cfg
index e098c8c9f4e..a740dc3c79b 100644
--- a/packages/syft/setup.cfg
+++ b/packages/syft/setup.cfg
@@ -62,7 +62,7 @@ syft =
     numpy>=1.23.5,<=1.24.4
     pandas==1.5.3
     docker==6.1.3
-    kr8s==0.13.0
+    kr8s==0.13.1
     PyYAML==6.0.1
     azure-storage-blob==12.19
 
diff --git a/packages/syft/src/syft/custom_worker/runner_k8s.py b/packages/syft/src/syft/custom_worker/runner_k8s.py
index dd1266b98ed..9e7d01e9e26 100644
--- a/packages/syft/src/syft/custom_worker/runner_k8s.py
+++ b/packages/syft/src/syft/custom_worker/runner_k8s.py
@@ -1,5 +1,4 @@
 # stdlib
-from time import sleep
 from typing import Dict
 from typing import List
 from typing import Optional
@@ -56,7 +55,8 @@ def create_pool(
             )
 
             # wait for replicas to be available and ready
-            self.wait(deployment, available_replicas=replicas)
+            status_path = "{.status.availableReplicas}"
+            deployment.wait(f"jsonpath='{status_path}'={replicas}")
         except Exception:
             raise
         finally:
@@ -116,29 +116,6 @@ def get_pod_env_vars(self, pod: Union[str, Pod]) -> Optional[List[Dict]]:
         pod = KubeUtils.resolve_pod(self.client, pod)
         return KubeUtils.get_pod_env(pod)
 
-    def wait(
-        self,
-        deployment: StatefulSet,
-        available_replicas: int,
-        timeout: int = 300,
-    ) -> None:
-        # TODO: Report wait('jsonpath=') bug to kr8s
-        # Until then this is the substitute implementation
-
-        if available_replicas <= 0:
-            return
-
-        while True:
-            if timeout == 0:
-                raise TimeoutError("Timeout waiting for replicas")
-
-            deployment.refresh()
-            if deployment.status.availableReplicas == available_replicas:
-                break
-
-            timeout -= 1
-            sleep(1)
-
     def _create_image_pull_secret(
         self,
         pool_name: str,