fixing load balancer (#14)

astronomer · Jul 21, 2024 · d3241ff · d3241ff
1 parent b4d1a7b
commit d3241ff
Showing 1 changed file with 13 additions and 2 deletions.
diff --git a/ray_provider/hooks/ray.py b/ray_provider/hooks/ray.py
@@ -286,10 +286,21 @@ def wait_for_load_balancer(
 
                 if lb_details:
                     hostname = lb_details["ip_or_hostname"]
-                    if all(self._is_port_open(hostname, port["port"]) for port in lb_details["ports"]):
+                    all_ports_open = True
+                    for port in lb_details["ports"]:
+                        if not self._is_port_open(hostname, port["port"]):
+                            self.log.info(f"Port {port['port']} is not open yet.")
+                            all_ports_open = False
+                            break
+
+                    if all_ports_open:
+                        self.log.info("All ports are open. LoadBalancer is ready.")
                         return lb_details
+                    else:
+                        self.log.info("Not all ports are open. Waiting...")
+                else:
+                    self.log.info("LoadBalancer details not available yet.")
 
-                    self.log.info("Not all ports are open. Waiting...")
             except AirflowException:
                 self.log.info("LoadBalancer service is not available yet...")