__main__.py

# Copyright 2016-2024, Pulumi Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import pulumi
import json
import pulumi_aws as aws
import pulumi_docker as docker

config = pulumi.Config()
vpc_cidr = config.get("vpc-cidr")
if vpc_cidr is None:
    vpc_cidr = "10.0.0.0/16"
subnet1_cidr = config.get("subnet-1-cidr")
if subnet1_cidr is None:
    subnet1_cidr = "10.0.0.0/24"
subnet2_cidr = config.get("subnet-2-cidr")
if subnet2_cidr is None:
    subnet2_cidr = "10.0.1.0/24"
container_context = config.get("container-context")
if container_context is None:
    container_context = "."
container_file = config.get("container-file")
if container_file is None:
    container_file = "./Dockerfile"
open_api_key = config.get("open-api-key")
if open_api_key is None:
    open_api_key = "CHANGEME"

region = aws.get_region()

availability_zones = [
    f"{region.name}a",
    f"{region.name}b",
]
current = aws.get_caller_identity_output()
pulumi_project = pulumi.get_project()
pulumi_stack = pulumi.get_stack()
langserve_ecr_repository = aws.ecr.Repository(
    "langserve-ecr-repository",
    name=f"{pulumi_project}-{pulumi_stack}",
    force_delete=True,
)
token = aws.ecr.get_authorization_token_output(registry_id=langserve_ecr_repository.registry_id)
account_id = current.account_id
langserve_ecr_life_cycle_policy = aws.ecr.LifecyclePolicy(
    "langserve-ecr-life-cycle-policy",
    repository=langserve_ecr_repository.name,
    policy=json.dumps(
        {
            "rules": [
                {
                    "rulePriority": 1,
                    "description": "Expire images when they are more than 10 available",
                    "selection": {
                        "tagStatus": "any",
                        "countType": "imageCountMoreThan",
                        "countNumber": 10,
                    },
                    "action": {
                        "type": "expire",
                    },
                }
            ],
        }
    ),
)
langserve_ecr_image = docker.Image(
    "langserve-ecr-image",
    build=docker.DockerBuildArgs(
        platform="linux/amd64",
        context=container_context,
        dockerfile=container_file,
    ),
    image_name=langserve_ecr_repository.repository_url,
    registry=docker.RegistryArgs(
        server=langserve_ecr_repository.repository_url,
        username=token.user_name,
        password=pulumi.Output.secret(token.password),
    ),
)
langserve_vpc = aws.ec2.Vpc(
    "langserve-vpc",
    cidr_block=vpc_cidr,
    enable_dns_hostnames=True,
    enable_dns_support=True,
    instance_tenancy="default",
    tags={
        "Name": f"{pulumi_project}-{pulumi_stack}",
    },
)
langserve_rt = aws.ec2.RouteTable(
    "langserve-rt",
    vpc_id=langserve_vpc.id,
    tags={
        "Name": f"{pulumi_project}-{pulumi_stack}",
    },
)
langserve_igw = aws.ec2.InternetGateway(
    "langserve-igw",
    vpc_id=langserve_vpc.id,
    tags={
        "Name": f"{pulumi_project}-{pulumi_stack}",
    },
)
langserve_route = aws.ec2.Route(
    "langserve-route",
    route_table_id=langserve_rt.id,
    destination_cidr_block="0.0.0.0/0",
    gateway_id=langserve_igw.id,
)
langserve_subnet1 = aws.ec2.Subnet(
    "langserve-subnet1",
    vpc_id=langserve_vpc.id,
    cidr_block=subnet1_cidr,
    availability_zone=availability_zones[0],
    map_public_ip_on_launch=True,
    tags={
        "Name": f"{pulumi_project}-{pulumi_stack}-1",
    },
)
langserve_subnet2 = aws.ec2.Subnet(
    "langserve-subnet2",
    vpc_id=langserve_vpc.id,
    cidr_block=subnet2_cidr,
    availability_zone=availability_zones[1],
    map_public_ip_on_launch=True,
    tags={
        "Name": f"{pulumi_project}-{pulumi_stack}-2",
    },
)
langserve_subnet1_rt_assoc = aws.ec2.RouteTableAssociation(
    "langserve-subnet1-rt-assoc",
    subnet_id=langserve_subnet1.id,
    route_table_id=langserve_rt.id,
)
langserve_subnet2_rt_assoc = aws.ec2.RouteTableAssociation(
    "langserve-subnet2-rt-assoc",
    subnet_id=langserve_subnet2.id,
    route_table_id=langserve_rt.id,
)
langserve_ecs_cluster = aws.ecs.Cluster(
    "langserve-ecs-cluster",
    configuration=aws.ecs.ClusterConfigurationArgs(
        execute_command_configuration=aws.ecs.ClusterConfigurationExecuteCommandConfigurationArgs(
            logging="DEFAULT",
        ),
    ),
    settings=[
        aws.ecs.ClusterSettingArgs(
            name="containerInsights",
            value="disabled",
        )
    ],
    tags={
        "Name": f"{pulumi_project}-{pulumi_stack}",
    },
)
langserve_cluster_capacity_providers = aws.ecs.ClusterCapacityProviders(
    "langserve-cluster-capacity-providers",
    cluster_name=langserve_ecs_cluster.name,
    capacity_providers=[
        "FARGATE",
        "FARGATE_SPOT",
    ],
)
langserve_security_group = aws.ec2.SecurityGroup(
    "langserve-security-group",
    vpc_id=langserve_vpc.id,
    ingress=[
        aws.ec2.SecurityGroupIngressArgs(
            protocol="tcp",
            from_port=80,
            to_port=80,
            cidr_blocks=["0.0.0.0/0"],
        )
    ],
    egress=[
        aws.ec2.SecurityGroupEgressArgs(
            protocol="-1",
            from_port=0,
            to_port=0,
            cidr_blocks=["0.0.0.0/0"],
        )
    ],
)
langserve_load_balancer = aws.lb.LoadBalancer(
    "langserve-load-balancer",
    load_balancer_type="application",
    security_groups=[langserve_security_group.id],
    subnets=[
        langserve_subnet1.id,
        langserve_subnet2.id,
    ],
)
langserve_target_group = aws.lb.TargetGroup(
    "langserve-target-group",
    port=80,
    protocol="HTTP",
    target_type="ip",
    vpc_id=langserve_vpc.id,
)
langserve_listener = aws.lb.Listener(
    "langserve-listener",
    load_balancer_arn=langserve_load_balancer.arn,
    port=80,
    protocol="HTTP",
    default_actions=[
        aws.lb.ListenerDefaultActionArgs(
            type="forward",
            target_group_arn=langserve_target_group.arn,
        )
    ],
)
langserve_log_group = aws.cloudwatch.LogGroup("langserve-log-group", retention_in_days=7)
langserve_key = aws.kms.Key(
    "langserve-key",
    description="Key for encrypting secrets",
    enable_key_rotation=True,
    policy=account_id.apply(
        lambda account_id: json.dumps(
            {
                "Version": "2012-10-17",
                "Statement": [
                    {
                        "Effect": "Allow",
                        "Principal": {
                            "AWS": f"arn:aws:iam::{account_id}:root",
                        },
                        "Action": [
                            "kms:Create*",
                            "kms:Describe*",
                            "kms:Enable*",
                            "kms:List*",
                            "kms:Put*",
                            "kms:Update*",
                            "kms:Revoke*",
                            "kms:Disable*",
                            "kms:Get*",
                            "kms:Delete*",
                            "kms:ScheduleKeyDeletion",
                            "kms:CancelKeyDeletion",
                            "kms:Tag*",
                            "kms:UntagResource",
                        ],
                        "Resource": "*",
                    },
                    {
                        "Effect": "Allow",
                        "Principal": {
                            "AWS": f"arn:aws:iam::{account_id}:root",
                        },
                        "Action": [
                            "kms:Encrypt",
                            "kms:Decrypt",
                            "kms:ReEncrypt*",
                            "kms:GenerateDataKey*",
                            "kms:DescribeKey",
                        ],
                        "Resource": "*",
                    },
                ],
            }
        )
    ),
    tags={
        "pulumi-application": pulumi_project,
        "pulumi-environment": pulumi_stack,
    },
)
langserve_ssm_parameter = aws.ssm.Parameter(
    "langserve-ssm-parameter",
    type="SecureString",
    value=open_api_key,
    key_id=langserve_key.key_id,
    name=f"/pulumi/{pulumi_project}/{pulumi_stack}/OPENAI_API_KEY",
    tags={
        "pulumi-application": pulumi_project,
        "pulumi-environment": pulumi_stack,
    },
)

langserve_execution_role = aws.iam.Role(
    "langserve-execution-role",
    assume_role_policy=json.dumps(
        {
            "Statement": [
                {
                    "Action": "sts:AssumeRole",
                    "Effect": "Allow",
                    "Principal": {
                        "Service": "ecs-tasks.amazonaws.com",
                    },
                }
            ],
            "Version": "2012-10-17",
        }
    ),
    inline_policies=[
        aws.iam.RoleInlinePolicyArgs(
            name=f"{pulumi_project}-{pulumi_stack}-service-secrets-policy",
            policy=pulumi.Output.all(langserve_ssm_parameter.arn, langserve_key.arn).apply(
                lambda args: json.dumps(
                    {
                        "Version": "2012-10-17",
                        "Statement": [
                            {
                                "Action": ["ssm:GetParameters"],
                                "Condition": {
                                    "StringEquals": {
                                        "ssm:ResourceTag/pulumi-application": pulumi_project,
                                        "ssm:ResourceTag/pulumi-environment": pulumi_stack,
                                    },
                                },
                                "Effect": "Allow",
                                "Resource": [args[0]],
                            },
                            {
                                "Action": ["kms:Decrypt"],
                                "Condition": {
                                    "StringEquals": {
                                        "aws:ResourceTag/pulumi-application": pulumi_project,
                                        "aws:ResourceTag/pulumi-environment": pulumi_stack,
                                    },
                                },
                                "Effect": "Allow",
                                "Resource": [args[1]],
                                "Sid": "DecryptTaggedKMSKey",
                            },
                        ],
                    }
                )
            ),
        )
    ],
    managed_policy_arns=["arn:aws:iam::aws:policy/service-role/AmazonECSTaskExecutionRolePolicy"],
)

langserve_task_role = aws.iam.Role(
    "langserve-task-role",
    assume_role_policy=json.dumps(
        {
            "Statement": [
                {
                    "Action": "sts:AssumeRole",
                    "Effect": "Allow",
                    "Principal": {
                        "Service": "ecs-tasks.amazonaws.com",
                    },
                }
            ],
            "Version": "2012-10-17",
        }
    ),
    inline_policies=[
        aws.iam.RoleInlinePolicyArgs(
            name="ExecuteCommand",
            policy=json.dumps(
                {
                    "Version": "2012-10-17",
                    "Statement": [
                        {
                            "Action": [
                                "ssmmessages:CreateControlChannel",
                                "ssmmessages:OpenControlChannel",
                                "ssmmessages:CreateDataChannel",
                                "ssmmessages:OpenDataChannel",
                            ],
                            "Effect": "Allow",
                            "Resource": "*",
                        },
                        {
                            "Action": [
                                "logs:CreateLogStream",
                                "logs:DescribeLogGroups",
                                "logs:DescribeLogStreams",
                                "logs:PutLogEvents",
                            ],
                            "Effect": "Allow",
                            "Resource": "*",
                        },
                    ],
                }
            ),
        ),
        aws.iam.RoleInlinePolicyArgs(
            name="DenyIAM",
            policy=json.dumps(
                {
                    "Version": "2012-10-17",
                    "Statement": [
                        {
                            "Action": "iam:*",
                            "Effect": "Deny",
                            "Resource": "*",
                        }
                    ],
                }
            ),
        ),
    ],
)
langserve_task_definition = aws.ecs.TaskDefinition(
    "langserve-task-definition",
    family=f"{pulumi_project}-{pulumi_stack}",
    cpu="256",
    memory="512",
    network_mode="awsvpc",
    execution_role_arn=langserve_execution_role.arn,
    task_role_arn=langserve_task_role.arn,
    requires_compatibilities=["FARGATE"],
    container_definitions=pulumi.Output.all(
        langserve_ecr_image.repo_digest,
        langserve_ssm_parameter.name,
        langserve_log_group.name,
    ).apply(
        lambda args: json.dumps(
            [
                {
                    "name": f"{pulumi_project}-{pulumi_stack}-service",
                    "image": args[0],
                    "cpu": 0,
                    "portMappings": [
                        {
                            "name": "target",
                            "containerPort": 8080,
                            "hostPort": 8080,
                            "protocol": "tcp",
                        }
                    ],
                    "essential": True,
                    "secrets": [
                        {
                            "name": "OPENAI_API_KEY",
                            "valueFrom": args[1],
                        }
                    ],
                    "logConfiguration": {
                        "logDriver": "awslogs",
                        "options": {
                            "awslogs-group": args[2],
                            "awslogs-region": region.name,
                            "awslogs-stream-prefix": "pulumi-langserve",
                        },
                    },
                }
            ]
        )
    ),
)
langserve_ecs_security_group = aws.ec2.SecurityGroup(
    "langserve-ecs-security-group",
    vpc_id=langserve_vpc.id,
    ingress=[
        aws.ec2.SecurityGroupIngressArgs(
            protocol="-1",
            from_port=0,
            to_port=0,
            cidr_blocks=["0.0.0.0/0"],
        )
    ],
    egress=[
        aws.ec2.SecurityGroupEgressArgs(
            protocol="-1",
            from_port=0,
            to_port=0,
            cidr_blocks=["0.0.0.0/0"],
        )
    ],
)
langserve_service_discovery_namespace = aws.servicediscovery.PrivateDnsNamespace(
    "langserve-service-discovery-namespace",
    name=f"{pulumi_stack}.{pulumi_project}.local",
    vpc=langserve_vpc.id,
)
langserve_service = aws.ecs.Service(
    "langserve-service",
    cluster=langserve_ecs_cluster.arn,
    task_definition=langserve_task_definition.arn,
    desired_count=1,
    launch_type="FARGATE",
    network_configuration=aws.ecs.ServiceNetworkConfigurationArgs(
        assign_public_ip=True,
        security_groups=[langserve_ecs_security_group.id],
        subnets=[
            langserve_subnet1.id,
            langserve_subnet2.id,
        ],
    ),
    load_balancers=[
        aws.ecs.ServiceLoadBalancerArgs(
            target_group_arn=langserve_target_group.arn,
            container_name=f"{pulumi_project}-{pulumi_stack}-service",
            container_port=8080,
        )
    ],
    scheduling_strategy="REPLICA",
    service_connect_configuration=aws.ecs.ServiceServiceConnectConfigurationArgs(
        enabled=True,
        namespace=langserve_service_discovery_namespace.arn,
    ),
    tags={
        "Name": f"{pulumi_project}-{pulumi_stack}",
    },
)
pulumi.export("url", langserve_load_balancer.dns_name.apply(lambda dns_name: f"http://{dns_name}"))