From f863da40e84ace880fa495cfce9a93377c5c5659 Mon Sep 17 00:00:00 2001
From: Daniel Gafni <danielgafni16@gmail.com>
Date: Mon, 16 Sep 2024 21:01:46 +0200
Subject: [PATCH] wait until RayCluster is ready

---
 README.md                      | 4 +++-
 dagster_ray/kuberay/configs.py | 4 ++--
 dagster_ray/kuberay/pipes.py   | 9 +++++++++
 3 files changed, 14 insertions(+), 3 deletions(-)

diff --git a/README.md b/README.md
index c5f1177..33bddf7 100644
--- a/README.md
+++ b/README.md
@@ -142,7 +142,9 @@ def my_asset(context: AssetExecutionContext, pipes_rayjob_client: PipesRayJobCli
     pipes_rayjob_client.run(
         context=context,
         ray_job={
-            # RayJob manifest goes here, only .metadata.name is not required and will be generated if not provided
+            # RayJob manifest goes here
+            # .metadata.name is not required and will be generated if not provided
+            # *.container.image is not required and will be set to the current `dagster/image` tag if not provided
             # full reference: https://ray-project.github.io/kuberay/reference/api/#rayjob
             ...
         },
diff --git a/dagster_ray/kuberay/configs.py b/dagster_ray/kuberay/configs.py
index 94277b8..95da7ea 100644
--- a/dagster_ray/kuberay/configs.py
+++ b/dagster_ray/kuberay/configs.py
@@ -13,8 +13,8 @@
     "env": [],
     "envFrom": [],
     "resources": {
-        "limits": {"cpu": "1000m", "memory": "1Gi"},
-        "requests": {"cpu": "1000m", "memory": "1Gi"},
+        "limits": {"cpu": "50m", "memory": "0.1Gi"},
+        "requests": {"cpu": "50m", "memory": "0.1Gi"},
     },
 }
 DEFAULT_HEAD_GROUP_SPEC = {
diff --git a/dagster_ray/kuberay/pipes.py b/dagster_ray/kuberay/pipes.py
index 096790f..96e48cc 100644
--- a/dagster_ray/kuberay/pipes.py
+++ b/dagster_ray/kuberay/pipes.py
@@ -172,6 +172,15 @@ def _start(self, context: OpExecutionContext, ray_job: Dict[str, Any]) -> Dict[s
             namespace=ray_job["metadata"]["namespace"],
         )
 
+        self.client.ray_cluster_client.wait_until_ready(
+            name=self.client.get_ray_cluster_name(
+                name=name,
+                namespace=namespace,
+            ),
+            namespace=namespace,
+            timeout=self.timeout,
+        )
+
         self.client.wait_until_running(
             name=name,
             namespace=namespace,