[SCP] open port support

skypilot-org · Dec 5, 2024 · 92dc02b · 92dc02b
1 parent 06de6de
commit 92dc02b
Showing 1 changed file with 35 additions and 13 deletions.
diff --git a/sky/provision/scp/instance.py b/sky/provision/scp/instance.py
@@ -1,7 +1,7 @@
 """SCP instance provisioning."""
 
+import time
 from typing import Any, Dict, List, Optional
-
 from sky.clouds.utils import scp_utils
 
 
@@ -19,24 +19,46 @@ def open_ports(  # pylint: disable=unused-argument
         sg_id = vm_info['securityGroupIds'][0]['securityGroupId']
         scp_client.add_new_security_group_in_rule(sg_id, ports[0])
         scp_client.add_new_security_group_out_rule(sg_id, ports[0])
+
         vpc_id = vm_info['vpcId']
         firewall_list = scp_client.list_firewalls()
         internal_ip = vm_info['ip']
+
         for firewall in firewall_list:
             if firewall['vpcId'] == vpc_id:
                 firewall_id = firewall['firewallId']
-                rule_info = scp_client.add_new_firewall_inbound_rule(
-                    firewall_id, internal_ip, ports[0])
-                if rule_info is not None:
-                    rule_id = rule_info['resourceId']
-                    scp_client.wait_firewall_inbound_rule_complete(
-                        firewall_id, rule_id)
-                rule_info = scp_client.add_new_firewall_outbound_rule(
-                    firewall_id, internal_ip, ports[0])
-                if rule_info is not None:
-                    rule_id = rule_info['resourceId']
-                    scp_client.wait_firewall_outbound_rule_complete(
-                        firewall_id, rule_id)
+
+                attempts = 0
+                max_attempts = 300
+                while attempts < max_attempts:
+                    try:
+                        rule_info = scp_client.add_new_firewall_inbound_rule(
+                            firewall_id, internal_ip, ports[0])
+                        if rule_info is not None:
+                            rule_id = rule_info['resourceId']
+                            scp_client.wait_firewall_inbound_rule_complete(
+                                firewall_id, rule_id)
+                        break
+                    except Exception as e:
+                        attempts += 1
+                        time.sleep(10)
+                        continue
+
+                attempts = 0
+                max_attempts = 300
+                while attempts < max_attempts:
+                    try:
+                        rule_info = scp_client.add_new_firewall_outbound_rule(
+                            firewall_id, internal_ip, ports[0])
+                        if rule_info is not None:
+                            rule_id = rule_info['resourceId']
+                            scp_client.wait_firewall_outbound_rule_complete(
+                                firewall_id, rule_id)
+                        break
+                    except Exception as e:
+                        attempts += 1
+                        time.sleep(10)
+                        continue
 
 
 def cleanup_ports(  # pylint: disable=unused-argument, pointless-string-statement