Allow to create gateways in AWS and Azure

cli/dstack/_internal/backend/aws/compute.py

@@ -2,10 +2,12 @@
 
 from boto3 import Session
 
+import dstack._internal.backend.aws.gateway as gateway
 from dstack._internal.backend.aws import runners
 from dstack._internal.backend.aws import utils as aws_utils
 from dstack._internal.backend.aws.config import AWSConfig
 from dstack._internal.backend.base.compute import Compute
+from dstack._internal.core.gateway import GatewayHead
 from dstack._internal.core.instance import InstanceType, LaunchedInstanceInfo
 from dstack._internal.core.job import Job
 from dstack._internal.core.request import RequestHead
@@ -63,6 +65,33 @@ def cancel_spot_request(self, runner: Runner):
             request_id=runner.request_id,
         )
 
+    def create_gateway(self, instance_name: str, ssh_key_pub: str) -> GatewayHead:
+        instance = gateway.create_gateway_instance(
+            ec2_client=self._get_ec2_client(region=self.backend_config.region_name),
+            subnet_id=self.backend_config.subnet_id,
+            bucket_name=self.backend_config.bucket_name,
+            instance_name=instance_name,
+            ssh_key_pub=ssh_key_pub,
+        )
+        return GatewayHead(
+            instance_name=instance_name,
+            external_ip=instance["PublicIpAddress"],
+            internal_ip=instance["PrivateIpAddress"],
+        )
+
+    def delete_instance(self, instance_name: str):
+        try:
+            instance_id = gateway.get_instance_id(
+                ec2_client=self._get_ec2_client(region=self.backend_config.region_name),
+                instance_name=instance_name,
+            )
+            runners.terminate_instance(
+                ec2_client=self._get_ec2_client(region=self.backend_config.region_name),
+                request_id=instance_id,
+            )
+        except IndexError:
+            return
+
     def _get_ec2_client(self, region: Optional[str] = None):
         if region is None:
             return aws_utils.get_ec2_client(self.session)

cli/dstack/_internal/backend/aws/gateway.py

@@ -0,0 +1,162 @@
+import time
+from typing import Optional
+
+from botocore.client import BaseClient
+
+
+def create_gateway_instance(
+    ec2_client: BaseClient,
+    subnet_id: Optional[str],
+    bucket_name: str,
+    instance_name: str,
+    ssh_key_pub: str,
+    machine_type: str = "t2.micro",
+) -> dict:
+    launch_specification = {}
+    if subnet_id:
+        launch_specification["NetworkInterfaces"] = [
+            {
+                "AssociatePublicIpAddress": True,
+                "DeviceIndex": 0,
+                "SubnetId": subnet_id,
+                "Groups": [gateway_security_group_id(ec2_client, subnet_id, bucket_name)],
+            },
+        ]
+    else:
+        launch_specification["SecurityGroupIds"] = [
+            gateway_security_group_id(ec2_client, subnet_id, bucket_name)
+        ]
+    tags = [
+        {"Key": "Name", "Value": instance_name},
+        {"Key": "owner", "Value": "dstack"},
+        {"Key": "role", "Value": "gateway"},
+        {"Key": "dstack_bucket", "Value": bucket_name},
+    ]
+    response = ec2_client.run_instances(
+        BlockDeviceMappings=[
+            {
+                "DeviceName": "/dev/sda1",
+                "Ebs": {
+                    "VolumeSize": 10,
+                    "VolumeType": "gp2",
+                },
+            }
+        ],
+        ImageId="ami-0cffefff2d52e0a23",  # Ubuntu 22.04 LTS
+        InstanceType=machine_type,
+        MinCount=1,
+        MaxCount=1,
+        UserData=gateway_user_data_script(ssh_key_pub),
+        TagSpecifications=[
+            {
+                "ResourceType": "instance",
+                "Tags": tags,
+            },
+        ],
+        **launch_specification,
+    )
+    return wait_till_running(ec2_client, response["Instances"][0])
+
+
+def gateway_security_group_id(
+    ec2_client: BaseClient, subnet_id: Optional[str], bucket_name: str
+) -> str:
+    name_parts = ["dstack_gateway_sg"]
+    if subnet_id:
+        name_parts.append(subnet_id.replace("-", "_"))
+    name_parts.append(bucket_name.replace("-", "_"))
+    security_group_name = "_".join(name_parts)
+    response = ec2_client.describe_security_groups(
+        Filters=[
+            {
+                "Name": "group-name",
+                "Values": [
+                    security_group_name,
+                ],
+            },
+        ],
+    )
+    if response.get("SecurityGroups"):
+        return response["SecurityGroups"][0]["GroupId"]
+
+    group_specification = {}
+    if subnet_id:
+        subnets_response = ec2_client.describe_subnets(SubnetIds=[subnet_id])
+        group_specification["VpcId"] = subnets_response["Subnets"][0]["VpcId"]
+    security_group = ec2_client.create_security_group(
+        Description="Generated by dstack",
+        GroupName=security_group_name,
+        TagSpecifications=[
+            {
+                "ResourceType": "security-group",
+                "Tags": [
+                    {"Key": "owner", "Value": "dstack"},
+                    {"Key": "role", "Value": "gateway"},
+                ],
+            },
+        ],
+        **group_specification,
+    )
+    security_group_id = security_group["GroupId"]
+    ip_permissions = [
+        {
+            "FromPort": 0,
+            "ToPort": 65535,
+            "IpProtocol": "tcp",
+            "IpRanges": [{"CidrIp": "0.0.0.0/0"}],
+        }
+    ]
+    ec2_client.authorize_security_group_ingress(
+        GroupId=security_group_id, IpPermissions=ip_permissions
+    )
+    ec2_client.authorize_security_group_egress(
+        GroupId=security_group_id,
+        IpPermissions=[
+            {
+                "IpProtocol": "-1",
+            }
+        ],
+    )
+    return security_group_id
+
+
+def wait_till_running(
+    ec2_client: BaseClient, instance: dict, delay: int = 5, attempts: int = 30
+) -> dict:
+    instance_id = instance["InstanceId"]
+    attempt = 0
+    while instance["State"]["Name"] != "running":
+        if attempt >= attempts:
+            raise RuntimeError(f"Instance {instance_id} is not running")
+        time.sleep(delay)
+        attempt += 1
+        desc = ec2_client.describe_instances(InstanceIds=[instance_id])
+        instance = desc["Reservations"][0]["Instances"][0]
+    return instance
+
+
+def get_instance_id(ec2_client: BaseClient, instance_name: str) -> str:
+    desc = ec2_client.describe_instances(
+        Filters=[
+            {
+                "Name": "tag:Name",
+                "Values": [instance_name],
+            }
+        ]
+    )
+    return desc["Reservations"][0]["Instances"][0]["InstanceId"]
+
+
+def gateway_user_data_script(ssh_key_pub: str) -> str:
+    return f"""#!/bin/bash
+sudo apt-get update
+DEBIAN_FRONTEND=noninteractive sudo apt-get install -y -q nginx
+UBUNTU_UID=$(id -u ubuntu)
+UBUNTU_GID=$(id -g ubuntu)
+install -m 700 -o $UBUNTU_UID -g $UBUNTU_GID -d /home/ubuntu/.ssh
+install -m 600 -o $UBUNTU_UID -g $UBUNTU_GID /dev/null /home/ubuntu/.ssh/authorized_keys
+echo "{ssh_key_pub}" > /home/ubuntu/.ssh/authorized_keys
+WWW_UID=$(id -u www-data)
+WWW_GID=$(id -g www-data)
+install -m 700 -o $WWW_UID -g $WWW_GID -d /var/www/.ssh
+install -m 600 -o $WWW_UID -g $WWW_GID /dev/null /var/www/.ssh/authorized_keys"""

cli/dstack/_internal/backend/azure/compute.py

@@ -34,7 +34,9 @@
 from azure.mgmt.keyvault import KeyVaultManagementClient
 from azure.mgmt.network import NetworkManagementClient
 from azure.mgmt.resource import ResourceManagementClient
+from msrestazure.tools import parse_resource_id
 
+import dstack._internal.backend.azure.gateway as gateway
 from dstack import version
 from dstack._internal.backend.azure import utils as azure_utils
 from dstack._internal.backend.azure.config import AzureConfig
@@ -47,6 +49,7 @@
 )
 from dstack._internal.backend.base.config import BACKEND_CONFIG_FILENAME, RUNNER_CONFIG_FILENAME
 from dstack._internal.backend.base.runners import serialize_runner_yaml
+from dstack._internal.core.gateway import GatewayHead
 from dstack._internal.core.instance import InstanceType, LaunchedInstanceInfo
 from dstack._internal.core.job import Job
 from dstack._internal.core.request import RequestHead, RequestStatus
@@ -140,6 +143,43 @@ def terminate_instance(self, runner: Runner):
     def cancel_spot_request(self, runner: Runner):
         self.terminate_instance(runner)
 
+    def create_gateway(self, instance_name: str, ssh_key_pub: str) -> GatewayHead:
+        vm = gateway.create_gateway(
+            compute_client=self._compute_client,
+            network_client=self._network_client,
+            subscription_id=self.azure_config.subscription_id,
+            location=self.azure_config.location,
+            resource_group=self.azure_config.resource_group,
+            network=self.azure_config.network,
+            subnet=self.azure_config.subnet,
+            instance_name=instance_name,
+            ssh_key_pub=ssh_key_pub,
+        )
+        interface = gateway.get_network_interface(
+            network_client=self._network_client,
+            resource_group=self.azure_config.resource_group,
+            interface=parse_resource_id(vm.network_profile.network_interfaces[0].id)[
+                "resource_name"
+            ],
+        )
+        public_ip = gateway.get_public_ip(
+            network_client=self._network_client,
+            resource_group=self.azure_config.resource_group,
+            public_ip=interface.ip_configurations[0].public_ip_address.name,
+        )
+        return GatewayHead(
+            instance_name=instance_name,
+            external_ip=public_ip.ip_address,
+            internal_ip=interface.ip_configurations[0].private_ip_address,
+        )
+
+    def delete_instance(self, instance_name: str):
+        _terminate_instance(
+            compute_client=self._compute_client,
+            resource_group=self.azure_config.resource_group,
+            instance_name=instance_name,
+        )
+
 
 def _get_instance_types(client: ComputeManagementClient, location: str) -> List[InstanceType]:
     instance_types = []