diff --git a/mkdocs/docs/concepts/gateways.md b/mkdocs/docs/concepts/gateways.md index c9c35ac45..c75240067 100644 --- a/mkdocs/docs/concepts/gateways.md +++ b/mkdocs/docs/concepts/gateways.md @@ -55,7 +55,7 @@ Provisioning... A gateway requires a `domain` to be specified in the configuration before creation. The domain is used to generate service endpoints (e.g. `.`). -Once the gateway is created and assigned a hostname, configure your DNS by adding a wildcard record for `*.` (e.g. `*.example.com`). The record should point to the gateway's hostname and should be of type `A` if the hostname is an IP address (most cases), or of type `CNAME` if the hostname is another domain (some private gateways and Kubernetes). +Once the gateway is created and assigned a hostname, configure your DNS by adding a wildcard record for `*.` (e.g. `*.example.com`). The record should point to the gateway's hostname and should be of type `A` if the hostname is an IP address (most cases), or of type `CNAME` if the hostname is another domain (load balancers, Kubernetes). ??? info "Project name interpolation" You can use the `${{ run.project_name }}` variable to include the service’s project name in the domain name. This is especially useful when [exporting](exports.md) the gateway to multiple projects, as it ensures each importer receives a unique domain name. @@ -76,6 +76,51 @@ You can create gateways with the `aws`, `azure`, `gcp`, or `kubernetes` backends Gateways in `kubernetes` backend require an external load balancer. Managed Kubernetes solutions usually include a load balancer. For self-hosted Kubernetes, you must provide a load balancer by yourself. +### Load balancer + +The optional `load_balancer` property allows you to provision a load balancer in front of the gateway, which is useful for balancing requests between multiple gateway [replicas](#replicas), or for using certain [certificate](#certificate) types, such as AWS ACM. + +Currently, only AWS Application Load Balancer (ALB) is supported: + +
+ +```yaml +type: gateway +name: example-gateway +backend: aws +region: eu-west-1 +domain: example.com +replicas: 2 +load_balancer: + type: alb +certificate: + type: acm + arn: arn:aws:acm:eu-west-1:164099421079:certificate/3670388f-f43b-4872-aaf8-907b107a170d +``` + +
+ +??? info "Requirements" + An ALB gateway requires: + + - The `aws` backend. + - Either `certificate: { type: acm, ... }` or `certificate: null`. + - A VPC with at least two subnets in different availability zones. If `public_ip: False`, subnets must be private and have a route to a NAT gateway. + +The provisioned load balancer provides a hostname you can add to your DNS records. Replica hostnames do not need to be added to DNS. + +
+ +``` +$ dstack gateway list + NAME BACKEND HOSTNAME DOMAIN DEFAULT STATUS + example-gateway dstack-6t7i1b03-lb-338524206.eu-west-1.elb.amazonaws.com example.com ✓ running + replica=0 aws (eu-west-1) 34.246.162.72 running + replica=1 aws (eu-west-1) 52.18.222.190 running +``` + +
+ ### Certificate By default, when you run a service with a gateway, `dstack` provisions an SSL certificate via Let's Encrypt for the configured domain. This automatically enables HTTPS for the service endpoint. @@ -89,7 +134,7 @@ If you disable [public IP](#public-ip) (e.g. to make the gateway private) or if * `lets-encrypt` (default) — Automatic certificates via [Let's Encrypt](https://letsencrypt.org/). Requires a [public IP](#public-ip). * `acm` — Certificates managed by [AWS Certificate Manager](https://aws.amazon.com/certificate-manager/). AWS-only. TLS is terminated at the load balancer, not at the gateway, and HTTP requests are redirected to HTTPS by the ALB. - Requires a VPC with at least two subnets in different availability zones to provision a load balancer. If `public_ip: False`, subnets must be private and have a route to NAT gateway. + Implies `load_balancer: { type: alb }`. * `null` — No certificate. Services will use HTTP. ### Public IP @@ -155,7 +200,7 @@ replicas: 2 -To balance requests between gateway replicas, add DNS records for each replica or set up a load balancer outside of `dstack`. Replica hostnames are displayed in `dstack` CLI and UI. +To balance requests between gateway replicas, add DNS records for each replica, use a natively-supported [load balancer](#load-balancer), or set up a load balancer outside of `dstack`. Replica hostnames are displayed in `dstack` CLI and UI.
diff --git a/mkdocs/docs/reference/dstack.yml/gateway.md b/mkdocs/docs/reference/dstack.yml/gateway.md index 06d0433f4..65a3409c5 100644 --- a/mkdocs/docs/reference/dstack.yml/gateway.md +++ b/mkdocs/docs/reference/dstack.yml/gateway.md @@ -29,3 +29,13 @@ Set to `null` to disable certificates (e.g. for [private gateways](../../concept show_root_heading: false type: required: true + +### `load_balancer` + +=== "ALB" + + #SCHEMA# dstack._internal.core.models.gateways.ALBGatewayLoadBalancer + overrides: + show_root_heading: false + type: + required: true diff --git a/src/dstack/_internal/core/backends/aws/compute.py b/src/dstack/_internal/core/backends/aws/compute.py index 78b956f5b..1e0570e50 100644 --- a/src/dstack/_internal/core/backends/aws/compute.py +++ b/src/dstack/_internal/core/backends/aws/compute.py @@ -94,7 +94,11 @@ class AWSGatewayBackendData(CoreModel): lb_arn: str tg_arn: str listener_arn: str - http_listener_arn: Optional[str] = None # None for old gateways + """Primary listener""" + http_listener_arn: Optional[str] = None + """Listener for the HTTP->HTTPS redirection. + `None` for `certificate: null` gateways and for pre-0.20.17 gateways that have no redirection + """ class AWSVolumeBackendData(CoreModel): @@ -598,9 +602,8 @@ def create_gateway_load_balancer( self, configuration: GatewayLoadBalancerConfiguration, ) -> GatewayLoadBalancerData: - """Creates an ALB, target group, and listeners for a gateway with an ACM certificate.""" - assert configuration.certificate is not None - assert configuration.certificate.type == "acm" + """Creates an ALB, target group, and listeners for a gateway.""" + assert configuration.certificate is None or configuration.certificate.type == "acm" ec2_client = self.session.client("ec2", region_name=configuration.region) elb_client = self.session.client("elbv2", region_name=configuration.region) @@ -636,7 +639,8 @@ def create_gateway_load_balancer( ) if len(lb_subnets_ids) < 2: raise ComputeError( - "Deploying gateway with ACM certificate requires at least two subnets in different AZs" + "Deploying a gateway with a load balancer requires at least two subnets" + " in different AZs" ) # Using short names as LB and target groups have length limit of 32. @@ -668,43 +672,66 @@ def create_gateway_load_balancer( tg_arn = response["TargetGroups"][0]["TargetGroupArn"] logger.debug("Created Target Group for gateway %s", configuration.gateway_name) - logger.debug("Creating HTTPS ALB listener for gateway %s...", configuration.gateway_name) - response = elb_client.create_listener( - LoadBalancerArn=lb_arn, - Protocol="HTTPS", - Port=443, - SslPolicy="ELBSecurityPolicy-2016-08", - Certificates=[ - {"CertificateArn": configuration.certificate.arn}, - ], - DefaultActions=[ - { - "Type": "forward", - "TargetGroupArn": tg_arn, - } - ], - ) - listener_arn = response["Listeners"][0]["ListenerArn"] - logger.debug("Created HTTPS ALB listener for gateway %s", configuration.gateway_name) + if configuration.certificate is not None: + logger.debug( + "Creating HTTPS ALB listener for gateway %s...", configuration.gateway_name + ) + response = elb_client.create_listener( + LoadBalancerArn=lb_arn, + Protocol="HTTPS", + Port=443, + SslPolicy="ELBSecurityPolicy-2016-08", + Certificates=[ + {"CertificateArn": configuration.certificate.arn}, + ], + DefaultActions=[ + { + "Type": "forward", + "TargetGroupArn": tg_arn, + } + ], + ) + listener_arn = response["Listeners"][0]["ListenerArn"] + logger.debug("Created HTTPS ALB listener for gateway %s", configuration.gateway_name) - logger.debug("Creating HTTP ALB listener for gateway %s...", configuration.gateway_name) - response = elb_client.create_listener( - LoadBalancerArn=lb_arn, - Protocol="HTTP", - Port=80, - DefaultActions=[ - { - "Type": "redirect", - "RedirectConfig": { - "Protocol": "HTTPS", - "Port": "443", - "StatusCode": "HTTP_301", - }, - } - ], - ) - http_listener_arn = response["Listeners"][0]["ListenerArn"] - logger.debug("Created HTTP ALB listener for gateway %s", configuration.gateway_name) + logger.debug( + "Creating HTTP ALB listener for gateway %s...", configuration.gateway_name + ) + response = elb_client.create_listener( + LoadBalancerArn=lb_arn, + Protocol="HTTP", + Port=80, + DefaultActions=[ + { + "Type": "redirect", + "RedirectConfig": { + "Protocol": "HTTPS", + "Port": "443", + "StatusCode": "HTTP_301", + }, + } + ], + ) + http_listener_arn = response["Listeners"][0]["ListenerArn"] + logger.debug("Created HTTP ALB listener for gateway %s", configuration.gateway_name) + else: + logger.debug( + "Creating HTTP ALB listener for gateway %s...", configuration.gateway_name + ) + response = elb_client.create_listener( + LoadBalancerArn=lb_arn, + Protocol="HTTP", + Port=80, + DefaultActions=[ + { + "Type": "forward", + "TargetGroupArn": tg_arn, + } + ], + ) + listener_arn = response["Listeners"][0]["ListenerArn"] + http_listener_arn = None + logger.debug("Created HTTP ALB listener for gateway %s", configuration.gateway_name) return GatewayLoadBalancerData( hostname=lb_dns_name, diff --git a/src/dstack/_internal/core/compatibility/gateways.py b/src/dstack/_internal/core/compatibility/gateways.py index d76c6cece..6df9d1633 100644 --- a/src/dstack/_internal/core/compatibility/gateways.py +++ b/src/dstack/_internal/core/compatibility/gateways.py @@ -67,5 +67,7 @@ def _get_gateway_configuration_excludes( if configuration.default is None: configuration_excludes["default"] = True + if configuration.load_balancer is None: + configuration_excludes["load_balancer"] = True return configuration_excludes diff --git a/src/dstack/_internal/core/models/gateways.py b/src/dstack/_internal/core/models/gateways.py index d9decdaea..2fa85a655 100644 --- a/src/dstack/_internal/core/models/gateways.py +++ b/src/dstack/_internal/core/models/gateways.py @@ -36,7 +36,13 @@ class LetsEncryptGatewayCertificate(CoreModel): class ACMGatewayCertificate(CoreModel): type: Annotated[ - Literal["acm"], Field(description="Certificates by AWS Certificate Manager (ACM)") + Literal["acm"], + Field( + description=( + "Certificates by AWS Certificate Manager (ACM)." + " Implies `load_balancer: { type: alb }`" + ) + ), ] = "acm" arn: Annotated[ str, Field(description="The ARN of the wildcard ACM certificate for the domain") @@ -52,6 +58,13 @@ class GatewayCertificate(RootModel[Annotated[AnyGatewayCertificate, Field(discri pass +class ALBGatewayLoadBalancer(CoreModel): + type: Annotated[Literal["alb"], Field(description="AWS Application Load Balancer")] = "alb" + + +AnyGatewayLoadBalancer = Union[ALBGatewayLoadBalancer] + + class GatewayConfiguration(CoreModel): type: Literal["gateway"] = "gateway" name: Annotated[Optional[str], Field(description="The gateway name")] = None @@ -92,6 +105,16 @@ class GatewayConfiguration(CoreModel): ), ] = None public_ip: Annotated[bool, Field(description="Allocate public IP for the gateway")] = True + load_balancer: Annotated[ + Optional[AnyGatewayLoadBalancer], + Field( + discriminator="type", + description=( + "The load balancer configuration." + " Set to `type: alb` to front the gateway with an AWS Application Load Balancer" + ), + ), + ] = None certificate: Annotated[ Optional[AnyGatewayCertificate], Field( diff --git a/src/dstack/_internal/core/services/gateways.py b/src/dstack/_internal/core/services/gateways.py index 446c7b7fb..3fc244ffa 100644 --- a/src/dstack/_internal/core/services/gateways.py +++ b/src/dstack/_internal/core/services/gateways.py @@ -1,4 +1,8 @@ -from dstack._internal.core.models.gateways import GatewayConfiguration +from dstack._internal.core.models.gateways import ( + ALBGatewayLoadBalancer, + AnyGatewayLoadBalancer, + GatewayConfiguration, +) from dstack._internal.core.services.diff import ModelDiff, diff_models @@ -9,3 +13,13 @@ def diff_gateway_configurations(old: GatewayConfiguration, new: GatewayConfigura # default=None => default should stay unchanged => shouldn't be in the diff reset={"default"} if new.default is None else {}, ) + + +def get_effective_load_balancer( + configuration: GatewayConfiguration, +) -> AnyGatewayLoadBalancer | None: + if configuration.load_balancer is not None: + return configuration.load_balancer + if configuration.certificate is not None and configuration.certificate.type == "acm": + return ALBGatewayLoadBalancer() + return None diff --git a/src/dstack/_internal/server/background/pipeline_tasks/gateways.py b/src/dstack/_internal/server/background/pipeline_tasks/gateways.py index 60e36e452..377f82dd7 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/gateways.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/gateways.py @@ -16,6 +16,7 @@ GatewayReplicaStatus, GatewayStatus, ) +from dstack._internal.core.services.gateways import get_effective_load_balancer from dstack._internal.server.background.pipeline_tasks.base import ( NOW_PLACEHOLDER, Fetcher, @@ -349,7 +350,7 @@ class _SubmittedResult: async def _process_submitted_gateway(gateway_model: GatewayModel) -> _SubmittedResult: configuration = gateways_services.get_gateway_configuration(gateway_model) update_map: _GatewayUpdateMap = {} - if configuration.certificate is not None and configuration.certificate.type == "acm": + if get_effective_load_balancer(configuration) is not None: try: ( _, diff --git a/src/dstack/_internal/server/services/gateways/__init__.py b/src/dstack/_internal/server/services/gateways/__init__.py index 753572d3e..e76f530f9 100644 --- a/src/dstack/_internal/server/services/gateways/__init__.py +++ b/src/dstack/_internal/server/services/gateways/__init__.py @@ -1204,6 +1204,18 @@ def _validate_gateway_configuration(configuration: GatewayConfiguration): f"Cannot provision {replicas} gateway replicas. This server allows at most {GATEWAY_MAX_REPLICAS}" ) + if configuration.load_balancer is not None: + if configuration.load_balancer.type == "alb": + if configuration.backend != BackendType.AWS: + raise ServerClientError( + "`load_balancer: { type: alb }` is supported for `aws` backend only" + ) + if configuration.certificate is not None and configuration.certificate.type != "acm": + raise ServerClientError( + "`load_balancer: { type: alb }` can only be used with `certificate: null` or" + " `certificate: { type: acm }`" + ) + if configuration.certificate is not None: if configuration.certificate.type == "lets-encrypt" and not configuration.public_ip: raise ServerClientError( diff --git a/src/dstack/_internal/server/testing/common.py b/src/dstack/_internal/server/testing/common.py index bf4e4f027..18aef0b0a 100644 --- a/src/dstack/_internal/server/testing/common.py +++ b/src/dstack/_internal/server/testing/common.py @@ -49,6 +49,7 @@ from dstack._internal.core.models.gateways import ( GATEWAY_REPLICAS_DEFAULT, AnyGatewayCertificate, + AnyGatewayLoadBalancer, GatewayConfiguration, GatewayReplicaConfiguration, GatewayReplicaStatus, @@ -669,6 +670,7 @@ async def create_gateway( forbid_new_services: bool = False, populate_configuration: bool = True, certificate: Optional[AnyGatewayCertificate] = LetsEncryptGatewayCertificate(), + load_balancer: Optional[AnyGatewayLoadBalancer] = None, hostname: Optional[str] = None, backend_data: Optional[str] = None, ) -> GatewayModel: @@ -689,6 +691,7 @@ async def create_gateway( domain=wildcard_domain, replicas=replicas, certificate=certificate, + load_balancer=load_balancer, ).model_dump_json() gateway = GatewayModel( project_id=project_id, diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_gateways.py b/src/tests/_internal/server/background/pipeline_tasks/test_gateways.py index a6674fefb..2953744be 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_gateways.py +++ b/src/tests/_internal/server/background/pipeline_tasks/test_gateways.py @@ -1,6 +1,7 @@ import asyncio import uuid from datetime import datetime, timedelta, timezone +from typing import Optional from unittest.mock import Mock, patch import pytest @@ -13,6 +14,7 @@ from dstack._internal.core.models.backends.base import BackendType from dstack._internal.core.models.gateways import ( ACMGatewayCertificate, + ALBGatewayLoadBalancer, GatewayLoadBalancerData, GatewayReplicaStatus, GatewayStatus, @@ -494,8 +496,29 @@ async def test_submitted_to_provisioning( assert replicas[1].replica_num == 1 assert all(r.ip_address is None for r in replicas) - async def test_submitted_to_provisioning_creates_load_balancer_for_acm_gateway( - self, test_db, session: AsyncSession, worker: GatewayWorker + @pytest.mark.parametrize( + "certificate,load_balancer", + [ + pytest.param(None, ALBGatewayLoadBalancer(), id="alb-without-certificate"), + pytest.param( + ACMGatewayCertificate(arn="arn:aws:acm:us:1:certificate/x"), + ALBGatewayLoadBalancer(), + id="alb-with-acm", + ), + pytest.param( + ACMGatewayCertificate(arn="arn:aws:acm:us:1:certificate/x"), + None, + id="acm-with-implicit-alb", + ), + ], + ) + async def test_submitted_to_provisioning_creates_load_balancer( + self, + test_db, + session: AsyncSession, + worker: GatewayWorker, + certificate: Optional[ACMGatewayCertificate], + load_balancer: Optional[ALBGatewayLoadBalancer], ): project = await create_project(session=session) backend = await create_backend(session=session, project_id=project.id) @@ -504,7 +527,8 @@ async def test_submitted_to_provisioning_creates_load_balancer_for_acm_gateway( project_id=project.id, backend_id=backend.id, status=GatewayStatus.SUBMITTED, - certificate=ACMGatewayCertificate(arn="arn:aws:acm:us:1:certificate/x"), + certificate=certificate, + load_balancer=load_balancer, ) gateway.lock_token = uuid.uuid4() gateway.lock_expires_at = datetime(2025, 1, 2, 3, 4, tzinfo=timezone.utc) @@ -528,6 +552,7 @@ async def test_submitted_to_provisioning_creates_load_balancer_for_acm_gateway( create_lb_mock = backend_mock.compute.return_value.create_gateway_load_balancer create_lb_mock.assert_called_once() assert create_lb_mock.call_args.args[0].gateway_name == gateway.name + assert create_lb_mock.call_args.args[0].certificate == certificate await session.refresh(gateway) assert gateway.status == GatewayStatus.PROVISIONING diff --git a/src/tests/_internal/server/routers/test_gateways.py b/src/tests/_internal/server/routers/test_gateways.py index 667c01714..ca7d27cf9 100644 --- a/src/tests/_internal/server/routers/test_gateways.py +++ b/src/tests/_internal/server/routers/test_gateways.py @@ -109,6 +109,7 @@ async def test_list( "domain": gateway.wildcard_domain, "default": False, "public_ip": True, + "load_balancer": None, "certificate": {"type": "lets-encrypt"}, "tags": None, "replicas": None, @@ -195,6 +196,7 @@ async def test_get( "domain": gateway.wildcard_domain, "default": False, "public_ip": True, + "load_balancer": None, "certificate": {"type": "lets-encrypt"}, "tags": None, "replicas": None, @@ -536,6 +538,7 @@ async def test_create_gateway(self, test_db, session: AsyncSession, client: Asyn "domain": None, "default": True, "public_ip": True, + "load_balancer": None, "certificate": {"type": "lets-encrypt"}, "tags": None, "replicas": None, @@ -626,6 +629,7 @@ async def test_create_gateway_without_name( "domain": None, "default": True, "public_ip": True, + "load_balancer": None, "certificate": {"type": "lets-encrypt"}, "tags": None, "replicas": None, @@ -754,6 +758,30 @@ async def test_create_gateway_with_invalid_domain_interpolation( "Cannot provision 4 gateway replicas. This server allows at most 3", id="replicas-exceed-max", ), + pytest.param( + { + "type": "gateway", + "name": "test", + "backend": "gcp", + "region": "us", + "certificate": None, + "load_balancer": {"type": "alb"}, + }, + "`load_balancer: { type: alb }` is supported for `aws` backend only", + id="load-balancer-non-aws-backend", + ), + pytest.param( + { + "type": "gateway", + "name": "test", + "backend": "aws", + "region": "us", + "load_balancer": {"type": "alb"}, + }, + "`load_balancer: { type: alb }` can only be used with `certificate: null` or" + " `certificate: { type: acm }`", + id="load-balancer-with-lets-encrypt-cert", + ), ], ) async def test_invalid_configuration_rejected( @@ -871,6 +899,7 @@ async def test_set_default_gateway( "domain": gateway.wildcard_domain, "default": True, "public_ip": True, + "load_balancer": None, "certificate": {"type": "lets-encrypt"}, "tags": None, "replicas": None, @@ -1267,6 +1296,7 @@ async def test_set_wildcard_domain( "domain": "new.example", "default": False, "public_ip": True, + "load_balancer": None, "certificate": {"type": "lets-encrypt"}, "tags": None, "replicas": None,