From 92dc02bae86f17ea1f92913b96db79b3f62c6727 Mon Sep 17 00:00:00 2001
From: Kihyo Moon <36657450+hyoxt121@users.noreply.github.com>
Date: Fri, 6 Dec 2024 08:44:04 +0900
Subject: [PATCH] [SCP] open port support

---
 sky/provision/scp/instance.py | 48 +++++++++++++++++++++++++----------
 1 file changed, 35 insertions(+), 13 deletions(-)

diff --git a/sky/provision/scp/instance.py b/sky/provision/scp/instance.py
index f8a939b1d70..f4b4658fbf6 100644
--- 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