skypilot-nightly 1.0.0.dev20250502__py3-none-any.whl → 1.0.0.dev20251203__py3-none-any.whl

sky/clouds/utils/scp_utils.py

@@ -3,7 +3,6 @@
 This module contains a set of rest api functions accessing SCP Open-API.
 """
 import base64
-import datetime
 from functools import wraps
 import hashlib
 import hmac
@@ -171,29 +170,30 @@ class SCPClient:
         self.secret_key = self._credentials['secret_key']
         self.project_id = self._credentials['project_id']
         self.client_type = 'OpenApi'
-        self.timestamp = ''
-        self.signature = ''
-
-        self.headers = {
-            'X-Cmp-AccessKey': f'{self.access_key}',
-            'X-Cmp-ClientType': f'{self.client_type}',
-            'X-Cmp-ProjectId': f'{self.project_id}',
-            'X-Cmp-Timestamp': f'{self.timestamp}',
-            'X-Cmp-Signature': f'{self.signature}'
+
+    def _signed_headers(self, method: str, url: str) -> dict:
+        timestamp = str(int(time.time() * 1000))
+        signature = self.get_signature(method=method,
+                                       url=url,
+                                       timestamp=timestamp)
+        return {
+            'X-Cmp-AccessKey': self.access_key,
+            'X-Cmp-ClientType': self.client_type,
+            'X-Cmp-ProjectId': self.project_id,
+            'X-Cmp-Timestamp': timestamp,
+            'X-Cmp-Signature': signature,
         }
 
     def create_instance(self, instance_config):
         """Launch new instances."""
-        url = f'{API_ENDPOINT}/virtual-server/v3/virtual-servers'
+        url = f'{API_ENDPOINT}/virtual-server/v4/virtual-servers'
         return self._post(url, instance_config)
 
     @_retry
     def _get(self, url, contents_key='contents'):
         method = 'GET'
-        self.set_timestamp()
-        self.set_signature(url=url, method=method)
-
-        response = requests.get(url, headers=self.headers)
+        headers = self._signed_headers(method, url)
+        response = requests.get(url, headers=headers)
         raise_scp_error(response)
         if contents_key is not None:
             return response.json().get(contents_key, [])
@@ -203,135 +203,167 @@ class SCPClient:
     @_retry
     def _post(self, url, request_body):
         method = 'POST'
-        self.set_timestamp()
-        self.set_signature(url=url, method=method)
-
-        response = requests.post(url, json=request_body, headers=self.headers)
-
+        headers = self._signed_headers(method, url)
+        response = requests.post(url, json=request_body, headers=headers)
         raise_scp_error(response)
         return response.json()
 
     @_retry
     def _delete(self, url, request_body=None):
         method = 'DELETE'
-        self.set_timestamp()
-        self.set_signature(url=url, method=method)
+        headers = self._signed_headers(method, url)
         if request_body:
-            response = requests.delete(url,
-                                       json=request_body,
-                                       headers=self.headers)
-
+            response = requests.delete(url, json=request_body, headers=headers)
         else:
-            response = requests.delete(url, headers=self.headers)
+            response = requests.delete(url, headers=headers)
         raise_scp_error(response)
         return response.json()
 
-    def create_security_group(self, zone_id, vpc, sg_name):
+    def create_security_group(self, zone_id, vpc_id, sg_name):
         url = f'{API_ENDPOINT}/security-group/v3/security-groups'
         request_body = {
             'loggable': False,
             'securityGroupName': sg_name,
             'serviceZoneId': zone_id,
-            'vpcId': vpc,
-            'securityGroupDescription': 'skypilot sg'
-        }
-        return self._post(url, request_body)
-
-    def add_security_group_in_rule(self, sg_id):
-        url = f'{API_ENDPOINT}/security-group/v2/security-groups/{sg_id}/rules'
-        request_body = {
-            'ruleDirection': 'IN',
-            'services': [{
-                'serviceType': 'TCP',
-                'serviceValue': '22'
-            }],
-            'sourceIpAddresses': ['0.0.0.0/0'],
-            'ruleDescription': 'skypilot ssh rule'
-        }
-        return self._post(url, request_body)
-
-    def add_security_group_out_rule(self, sg_id):
-        url = f'{API_ENDPOINT}/security-group/v2/security-groups/{sg_id}/rules'
-        request_body = {
-            'ruleDirection': 'OUT',
-            'services': [{
-                'serviceType': 'TCP',
-                'serviceValue': '21'
-            }, {
-                'serviceType': 'TCP',
-                'serviceValue': '22'
-            }, {
-                'serviceType': 'TCP',
-                'serviceValue': '80'
-            }, {
-                'serviceType': 'TCP',
-                'serviceValue': '443'
-            }],
-            'destinationIpAddresses': ['0.0.0.0/0'],
-            'ruleDescription': 'skypilot out rule'
-        }
-        return self._post(url, request_body)
-
-    def add_firewall_inbound_rule(self, firewall_id, internal_ip):
-        url = f'{API_ENDPOINT}/firewall/v2/firewalls/{firewall_id}/rules'
-        request_body = {
-            'sourceIpAddresses': ['0.0.0.0/0'],
-            'destinationIpAddresses': [internal_ip],
-            'services': [{
-                'serviceType': 'TCP',
-                'serviceValue': '22'
-            }],
-            'ruleDirection': 'IN',
-            'ruleAction': 'ALLOW',
-            'isRuleEnabled': True,
-            'ruleLocationType': 'FIRST',
-            'ruleDescription': 'description'
+            'vpcId': vpc_id,
+            'securityGroupDescription': 'sky security group'
         }
         return self._post(url, request_body)
 
-    def add_firewall_outbound_rule(self, firewall_id, internal_ip):
-        url = f'{API_ENDPOINT}/firewall/v2/firewalls/{firewall_id}/rules'
-        request_body = {
-            'sourceIpAddresses': [internal_ip],
-            'destinationIpAddresses': ['0.0.0.0/0'],
-            'services': [{
-                'serviceType': 'TCP',
-                'serviceValue': '21'
-            }, {
-                'serviceType': 'TCP',
-                'serviceValue': '22'
-            }, {
-                'serviceType': 'TCP',
-                'serviceValue': '80'
-            }, {
-                'serviceType': 'TCP',
-                'serviceValue': '443'
-            }],
-            'ruleDirection': 'OUT',
-            'ruleAction': 'ALLOW',
-            'isRuleEnabled': True,
-            'ruleLocationType': 'FIRST',
-            'ruleDescription': 'description'
-        }
-        return self._post(url, request_body)
+    def _security_group_rule_not_exist(self, sg_id, direction, ports):
+        response = self.get_security_group_rules(sg_id)
+        rules = []
+        for rule in response:
+            rule_direction = rule['ruleDirection']
+            if rule_direction == direction:
+                rules.append(rule)
+        for rule in rules:
+            port_list = ','.join(rule['tcpServices'])
+            port = ','.join(ports)
+            if port == port_list:
+                return False
+        return True
+
+    def add_security_group_rule(self, sg_id, direction,
+                                ports: Optional[List[str]],
+                                cnt: Optional[int]) -> None:
+        if ports is None:
+            if direction == 'IN':
+                if cnt == 1:
+                    ports = ['22']
+                else:
+                    ports = ['22', '6380', '8076', '10001', '11001-11200']
+            else:
+                if cnt == 1:
+                    ports = ['21', '22', '80', '443']
+                else:
+                    ports = [
+                        '21', '22', '80', '443', '6380', '8076', '10001',
+                        '11001-11200'
+                    ]
+        services = []
+        for port in ports:
+            services.append({'serviceType': 'TCP', 'serviceValue': port})
+        if self._security_group_rule_not_exist(sg_id, direction, ports):
+            url = f'{API_ENDPOINT}/security-group/v2/security-groups/{sg_id}/rules'  # pylint: disable=line-too-long
+            if direction == 'IN':
+                target_address = 'sourceIpAddresses'
+            else:
+                target_address = 'destinationIpAddresses'
+            request_body = {
+                'ruleDirection': direction,
+                'services': services,
+                target_address: ['0.0.0.0/0'],
+                'ruleDescription': 'sky security group rule'
+            }
+            self._post(url, request_body)
+        else:
+            return None
+
+    def _firewall_rule_not_exist(self, firewall_id, internal_ip, direction,
+                                 ports):
+        response = self.get_firewall_rules(firewall_id)
+        rules = []
+        for rule in response:
+            if direction == 'IN':
+                if internal_ip == rule['destinationIpAddresses'][0]:
+                    rules.append(rule)
+            else:
+                if internal_ip == rule['sourceIpAddresses'][0]:
+                    rules.append(rule)
+        for rule in rules:
+            port_list = ','.join(rule['tcpServices'])
+            port = ','.join(ports)
+            if port == port_list:
+                return False
+        return True
+
+    def add_firewall_rule(self, firewall_id, internal_ip, direction,
+                          ports: Optional[List[str]], cnt: Optional[int]):
+        if ports is None:
+            if direction == 'IN':
+                if cnt == 1:
+                    ports = ['22']
+                else:
+                    ports = ['22', '6380', '8076', '10001', '11001-11200']
+            else:
+                if cnt == 1:
+                    ports = ['21', '22', '80', '443']
+                else:
+                    ports = [
+                        '21', '22', '80', '443', '6380', '8076', '10001',
+                        '11001-11200'
+                    ]
+        services = []
+        for port in ports:
+            services.append({'serviceType': 'TCP', 'serviceValue': port})
+        if self._firewall_rule_not_exist(firewall_id, internal_ip, direction,
+                                         ports):
+            url = f'{API_ENDPOINT}/firewall/v2/firewalls/{firewall_id}/rules'
+            if direction == 'IN':
+                source_ip = '0.0.0.0/0'
+                destination_ip = internal_ip
+            else:
+                source_ip = internal_ip
+                destination_ip = '0.0.0.0/0'
+            request_body = {
+                'sourceIpAddresses': [source_ip],
+                'destinationIpAddresses': [destination_ip],
+                'services': services,
+                'ruleDirection': direction,
+                'ruleAction': 'ALLOW',
+                'isRuleEnabled': True,
+                'ruleLocationType': 'FIRST',
+                'ruleDescription': 'sky firewall rule'
+            }
+            return self._post(url, request_body)
+        else:
+            return None
 
-    def terminate_instance(self, vm_id):
-        url = f'{API_ENDPOINT}/virtual-server/v2/virtual-servers/{vm_id}'
+    def terminate_instance(self, instance_id):
+        url = f'{API_ENDPOINT}/virtual-server/v2/virtual-servers/{instance_id}'
         return self._delete(url)
 
-    def list_instances(self) -> List[dict]:
+    def get_instances(self) -> List[dict]:
         """List existing instances."""
         url = f'{API_ENDPOINT}/virtual-server/v2/virtual-servers'
         return self._get(url)
 
-    def list_catalog(self) -> Dict[str, Any]:
+    def get_catalog(self) -> Dict[str, Any]:
         """List offered instances and their availability."""
-        response = requests.get(f'{API_ENDPOINT}/instance-types',
-                                headers=self.headers)
+        url = f'{API_ENDPOINT}/instance-types'
+        headers = self._signed_headers('GET', url)
+        response = requests.get(url, headers=headers)
         raise_scp_error(response)
         return response.json().get('data', [])
 
-    def get_signature(self, method: str, url: str) -> str:
+    def get_signature(self,
+                      method: str,
+                      url: str,
+                      timestamp: Optional[str] = None) -> str:
+        if timestamp is None:
+            timestamp = str(int(time.time() * 1000))
+
         url_info = parse.urlsplit(url)
         url = (f'{url_info.scheme}://{url_info.netloc}'
                f'{parse.quote(url_info.path)}')
@@ -341,7 +373,7 @@ class SCPClient:
                     parse.parse_qs(url_info.query).items()))
             url = f'{url}?{parse.urlencode(enc_params)}'
 
-        message = method + url + self.timestamp \
+        message = method + url + timestamp \
                   + self.access_key + self.project_id + self.client_type
         message = bytes(message, 'utf-8')
         secret = bytes(self.secret_key, 'utf-8')
@@ -352,99 +384,131 @@ class SCPClient:
 
         return str(signature)
 
-    def set_timestamp(self) -> None:
-        self.timestamp = str(
-            int(
-                round(
-                    datetime.datetime.timestamp(datetime.datetime.now() -
-                                                datetime.timedelta(minutes=1)) *
-                    1000)))
-        self.headers['X-Cmp-Timestamp'] = self.timestamp
-
-    def set_signature(self, method: str, url: str) -> None:
-
-        self.signature = self.get_signature(url=url, method=method)
-        self.headers['X-Cmp-Signature'] = self.signature
-
-    def list_nic_details(self, virtual_server_id) -> List[dict]:
-        """List existing instances."""
-        url = f'{API_ENDPOINT}/virtual-server/v2/virtual-servers/{virtual_server_id}/nics'  # pylint: disable=line-too-long
+    def get_nic(self, instance_id) -> List[dict]:
+        url = f'{API_ENDPOINT}/virtual-server/v2/virtual-servers/{instance_id}/nics'  # pylint: disable=line-too-long
         return self._get(url)
 
-    def get_external_ip(self, virtual_server_id, ip):
-        nic_details_list = self.list_nic_details(
-            virtual_server_id=virtual_server_id)
-        for nic_details in nic_details_list:
-            if (nic_details['ip'] == ip and
-                    nic_details['subnetType'] == 'PUBLIC'):
-                return nic_details['natIp']
+    def get_external_ip(self, instance_id, ip):
+        nics = self.get_nic(instance_id=instance_id)
+        for nic in nics:
+            if (nic['ip'] == ip and nic['subnetType'] == 'PUBLIC'):
+                return nic['natIp']
         return None
 
-    def list_zones(self) -> List[dict]:
+    def get_zones(self) -> List[dict]:
         url = f'{API_ENDPOINT}/project/v3/projects/{self.project_id}/zones'
         return self._get(url)
 
-    def list_products(self, service_zone_id) -> List[dict]:
-        url = f'{API_ENDPOINT}/product/v2/zones/{service_zone_id}/products'
-        return self._get(url)
-
-    def list_product_groups(self, service_zone_id) -> List[dict]:
-        url = f'{API_ENDPOINT}/product/v2/zones/{service_zone_id}/product-groups'  # pylint: disable=line-too-long
-        return self._get(url)
-
-    def list_vpcs(self, service_zone_id) -> List[dict]:
+    def get_vpcs(self, service_zone_id) -> List[dict]:
         url = f'{API_ENDPOINT}/vpc/v2/vpcs?serviceZoneId={service_zone_id}'
         return self._get(url)
 
-    def list_subnets(self) -> List[dict]:
+    def get_subnets(self) -> List[dict]:
         url = f'{API_ENDPOINT}/subnet/v2/subnets?subnetTypes=PUBLIC'
         return self._get(url)
 
-    def del_security_group(self, sg_id):
+    def delete_security_group(self, sg_id):
         url = f'{API_ENDPOINT}/security-group/v2/security-groups/{sg_id}'
         return self._delete(url)
 
-    def del_firwall_rules(self, firewall_id, rule_id_list):
+    def delete_firewall_rule(self, firewall_id, rule_ids):
         url = f'{API_ENDPOINT}/firewall/v2/firewalls/{firewall_id}/rules'
-        request_body = {'ruleDeletionType': 'PARTIAL', 'ruleIds': rule_id_list}
+        request_body = {'ruleDeletionType': 'PARTIAL', 'ruleIds': rule_ids}
         return self._delete(url, request_body=request_body)
 
-    def list_security_groups(self, vpc_id=None, sg_name=None):
+    def get_security_groups(self, vpc_id=None, sg_name=None):
         url = f'{API_ENDPOINT}/security-group/v2/security-groups'
         parameter = []
         if vpc_id is not None:
             parameter.append('vpcId=' + vpc_id)
         if sg_name is not None:
             parameter.append('securityGroupName=' + sg_name)
-        if parameter:
+        if len(parameter) > 0:
             url = url + '?' + '&'.join(parameter)
         return self._get(url)
 
-    def list_igw(self):
+    def get_internet_gateway(self):
         url = f'{API_ENDPOINT}/internet-gateway/v2/internet-gateways'
         return self._get(url)
 
-    def get_vm_info(self, vm_id):
-        url = f'{API_ENDPOINT}/virtual-server/v3/virtual-servers/{vm_id}'
-        return self._get(url, contents_key=None)
-
-    def get_firewal_rule_info(self, firewall_id, rule_id):
+    def get_firewall_rule_info(self, firewall_id, rule_id):
         url = f'{API_ENDPOINT}/firewall/v2/firewalls/{firewall_id}/rules/{rule_id}'  # pylint: disable=line-too-long
         return self._get(url, contents_key=None)
 
-    def list_firwalls(self):
+    def get_firewalls(self):
         url = f'{API_ENDPOINT}/firewall/v2/firewalls'
         return self._get(url)
 
-    def list_service_zone_names(self):
+    def get_service_zone_names(self):
         url = f'{API_ENDPOINT}/project/v3/projects/{self.project_id}/zones'
         zone_contents = self._get(url)
         return [content['serviceZoneName'] for content in zone_contents]
 
-    def start_instance(self, vm_id):
-        url = f'{API_ENDPOINT}/virtual-server/v2/virtual-servers/{vm_id}/start'
+    def start_instance(self, instance_id):
+        url = f'{API_ENDPOINT}/virtual-server/v2/virtual-servers/{instance_id}/start'  # pylint: disable=line-too-long
         return self._post(url=url, request_body={})
 
-    def stop_instance(self, vm_id):
-        url = f'{API_ENDPOINT}/virtual-server/v2/virtual-servers/{vm_id}/stop'
+    def stop_instance(self, instance_id):
+        url = f'{API_ENDPOINT}/virtual-server/v2/virtual-servers/{instance_id}/stop'  # pylint: disable=line-too-long
         return self._post(url=url, request_body={})
+
+    def get_security_group_rules(self, sg_id):
+        url = f'{API_ENDPOINT}/security-group/v2/security-groups/{sg_id}/rules'
+        return self._get(url)
+
+    def get_firewall_rules(self, firewall_id):
+        url = f'{API_ENDPOINT}/firewall/v2/firewalls/{firewall_id}/rules'
+        return self._get(url)
+
+    def get_instance_info(self, instance_id):
+        url = f'{API_ENDPOINT}/virtual-server/v3/virtual-servers/{instance_id}'
+        return self._get(url=url, contents_key=None)
+
+    def create_vpc(self, zone_id):
+        vpc_name = 'skyvpc' + zone_id[5:10]
+        request_body = {
+            'serviceZoneId': zone_id,
+            'vpcName': vpc_name,
+            'vpcDescription': 'sky vpc'
+        }
+        url = f'{API_ENDPOINT}/vpc/v3/vpcs'
+        return self._post(url, request_body)
+
+    def create_subnet(self, vpc_id, zone_id):
+        subnet_name = 'skysubnet' + zone_id[5:10]
+        request_body = {
+            'subnetCidrBlock': '192.168.0.0/24',
+            'subnetName': subnet_name,
+            'subnetType': 'PUBLIC',
+            'vpcId': vpc_id,
+            'subnetDescription': 'sky subnet'
+        }
+        url = f'{API_ENDPOINT}/subnet/v2/subnets'
+        return self._post(url, request_body)
+
+    def create_internet_gateway(self, vpc_id):
+        request_body = {
+            'firewallEnabled': True,
+            'firewallLoggable': False,
+            'internetGatewayType': 'SHARED',
+            'vpcId': vpc_id,
+            'internetGatewayDescription': 'sky internet gateway'
+        }
+        url = f'{API_ENDPOINT}/internet-gateway/v4/internet-gateways'
+        return self._post(url, request_body)
+
+    def get_vpc_info(self, vpc_id):
+        url = f'{API_ENDPOINT}/vpc/v2/vpcs/{vpc_id}'
+        return self._get(url=url, contents_key=None)
+
+    def get_subnet_info(self, subnet_id):
+        url = f'{API_ENDPOINT}/subnet/v2/subnets/{subnet_id}'
+        return self._get(url=url, contents_key=None)
+
+    def get_internet_gateway_info(self, internet_gateway_id):
+        url = f'{API_ENDPOINT}/internet-gateway/v2/internet-gateways/{internet_gateway_id}'  # pylint: disable=line-too-long
+        return self._get(url=url, contents_key=None)
+
+    def get_key_pairs(self):
+        url = f'{API_ENDPOINT}/key-pair/v1/key-pairs'
+        return self._get(url=url, contents_key=None)