Description: This template creates the resources necessary for an anyscale cloud.
Transform: AWS::LanguageExtensions
Parameters:
  EnvironmentName:
    Description: Anyscale deploy environment. Used in resource names and tags.
    Type: String

  CloudID:
    Description: ID of the anyscale cloud.
    Type: String

  VpcCIDR:
    Description: CIDR for the anyscale cloud VPC.
    Type: String
    Default: 10.0.0.0/16

  AnyscaleAWSAccountID:
    Description: Anyscale control plane AWS account.
    Type: String
    Default: 525325868955

  AnyscaleCrossAccountIAMRoleName:
    Description: Name of the cross account IAM role.
    Type: String

  AnyscaleCrossAccountIAMPolicySteadyState:
    Description: Stead state IAM policy document
    Type: String

  AnyscaleCrossAccountIAMPolicyServiceSteadyState:
    Description: Stead state IAM policy document for services
    Type: String

  AnyscaleCrossAccountIAMPolicyInitialRun:
    Description: Initial run IAM policy document
    Type: String

  AnyscaleSecurityGroupName:
    Description: Name of the anyscale security group
    Type: String
    Default: "anyscale-security-group"

  ClusterNodeIAMRoleName:
    Description: Name of the data plane cluster IAM role
    Type: String

Resources:
  VPC:
    Type: AWS::EC2::VPC
    Properties:
      CidrBlock: !Ref VpcCIDR
      EnableDnsSupport: true
      EnableDnsHostnames: true
      Tags:
        - Key: Name
          Value: !Ref AWS::StackName
        - Key: anyscale-cloud-id
          Value: !Ref CloudID
        - Key: anyscale-deploy-environment
          Value: !Ref EnvironmentName

  InternetGateway:
    Type: AWS::EC2::InternetGateway
    Properties:
      Tags:
        - Key: Name
          Value: !Ref AWS::StackName
        - Key: anyscale-cloud-id
          Value: !Ref CloudID
        - Key: anyscale-deploy-environment
          Value: !Ref EnvironmentName

  InternetGatewayAttachment:
    Type: AWS::EC2::VPCGatewayAttachment
    Properties:
      InternetGatewayId: !Ref InternetGateway
      VpcId: !Ref VPC


  Subnet0:
    Type: AWS::EC2::Subnet
    Properties:
        VpcId: !Ref VPC
        AvailabilityZone: us-east-1a
        CidrBlock: 10.0.0.0/19
        MapPublicIpOnLaunch: true
        Tags:
        - Key: Name
          Value: mock_cfn_stack_name-subnet-us-east-1a
        - Key: anyscale-cloud-id
          Value: mock_cloud_id

  Subnet1:
    Type: AWS::EC2::Subnet
    Properties:
        VpcId: !Ref VPC
        AvailabilityZone: us-east-1b
        CidrBlock: 10.0.32.0/19
        MapPublicIpOnLaunch: true
        Tags:
        - Key: Name
          Value: mock_cfn_stack_name-subnet-us-east-1b
        - Key: anyscale-cloud-id
          Value: mock_cloud_id

  Subnet2:
    Type: AWS::EC2::Subnet
    Properties:
        VpcId: !Ref VPC
        AvailabilityZone: us-east-1c
        CidrBlock: 10.0.64.0/19
        MapPublicIpOnLaunch: true
        Tags:
        - Key: Name
          Value: mock_cfn_stack_name-subnet-us-east-1c
        - Key: anyscale-cloud-id
          Value: mock_cloud_id

  Subnet3:
    Type: AWS::EC2::Subnet
    Properties:
        VpcId: !Ref VPC
        AvailabilityZone: us-east-1d
        CidrBlock: 10.0.96.0/19
        MapPublicIpOnLaunch: true
        Tags:
        - Key: Name
          Value: mock_cfn_stack_name-subnet-us-east-1d
        - Key: anyscale-cloud-id
          Value: mock_cloud_id

  Subnet4:
    Type: AWS::EC2::Subnet
    Properties:
        VpcId: !Ref VPC
        AvailabilityZone: us-east-1e
        CidrBlock: 10.0.128.0/19
        MapPublicIpOnLaunch: true
        Tags:
        - Key: Name
          Value: mock_cfn_stack_name-subnet-us-east-1e
        - Key: anyscale-cloud-id
          Value: mock_cloud_id

  Subnet5:
    Type: AWS::EC2::Subnet
    Properties:
        VpcId: !Ref VPC
        AvailabilityZone: us-east-1f
        CidrBlock: 10.0.160.0/19
        MapPublicIpOnLaunch: true
        Tags:
        - Key: Name
          Value: mock_cfn_stack_name-subnet-us-east-1f
        - Key: anyscale-cloud-id
          Value: mock_cloud_id

  PublicRouteTable:
    Type: AWS::EC2::RouteTable
    Properties:
      VpcId: !Ref VPC
      Tags:
        - Key: Name
          Value: !Sub ${AWS::StackName} Public Routes
        - Key: anyscale-cloud-id
          Value: !Ref CloudID
        - Key: anyscale-deploy-environment
          Value: !Ref EnvironmentName

  DefaultPublicRoute:
    Type: AWS::EC2::Route
    DependsOn: InternetGatewayAttachment
    Properties:
      RouteTableId: !Ref PublicRouteTable
      DestinationCidrBlock: 0.0.0.0/0
      GatewayId: !Ref InternetGateway


  Subnet0RouteTableAssociation:
    Type: AWS::EC2::SubnetRouteTableAssociation
    Properties:
        RouteTableId: !Ref PublicRouteTable
        SubnetId: !Ref Subnet0

  Subnet1RouteTableAssociation:
    Type: AWS::EC2::SubnetRouteTableAssociation
    Properties:
        RouteTableId: !Ref PublicRouteTable
        SubnetId: !Ref Subnet1

  Subnet2RouteTableAssociation:
    Type: AWS::EC2::SubnetRouteTableAssociation
    Properties:
        RouteTableId: !Ref PublicRouteTable
        SubnetId: !Ref Subnet2

  Subnet3RouteTableAssociation:
    Type: AWS::EC2::SubnetRouteTableAssociation
    Properties:
        RouteTableId: !Ref PublicRouteTable
        SubnetId: !Ref Subnet3

  Subnet4RouteTableAssociation:
    Type: AWS::EC2::SubnetRouteTableAssociation
    Properties:
        RouteTableId: !Ref PublicRouteTable
        SubnetId: !Ref Subnet4

  Subnet5RouteTableAssociation:
    Type: AWS::EC2::SubnetRouteTableAssociation
    Properties:
        RouteTableId: !Ref PublicRouteTable
        SubnetId: !Ref Subnet5

  AnyscaleSecurityGroup:
    Type: AWS::EC2::SecurityGroup
    Properties:
      GroupName: !Ref AnyscaleSecurityGroupName
      GroupDescription: "Anyscale security group"
      VpcId: !Ref VPC
      SecurityGroupIngress:
        # For https
        - IpProtocol: "tcp"
          FromPort: 443
          ToPort: 443
          CidrIp: "0.0.0.0/0"
        # For ssh
        - IpProtocol: "tcp"
          FromPort: 22
          ToPort: 22
          CidrIp: "0.0.0.0/0"
      Tags:
        - Key: anyscale-cloud-id
          Value: !Ref CloudID
        - Key: anyscale-deploy-environment
          Value: !Ref EnvironmentName

  # This allows ray cluster nodes to talk to each other since all nodes have the same security group attached.
  # This also allows ray cluster nodes to talk to EFS since they have the same security group attached.
  AnyscaleSecurityGroupIntraClusterIngressRule:
    Type: AWS::EC2::SecurityGroupIngress
    Properties:
      IpProtocol: "-1"
      GroupId: !Ref AnyscaleSecurityGroup
      SourceSecurityGroupId: !Ref AnyscaleSecurityGroup

  S3GatewayEndpoint:
    Type: 'AWS::EC2::VPCEndpoint'
    Properties:
      RouteTableIds:
        - !Ref PublicRouteTable
      ServiceName: !Sub 'com.amazonaws.${AWS::Region}.s3'
      VpcId: !Ref VPC

  S3Bucket:
    Type: AWS::S3::Bucket
    Description: Anyscale managed s3 bucket.
    DeletionPolicy: Retain
    Properties:
      BucketName: !Sub anyscale-${EnvironmentName}-data-${AWS::StackName}
      CorsConfiguration:
        CorsRules:
          - AllowedHeaders:
              - '*'
            AllowedMethods:
              - GET
            AllowedOrigins:
              - https://console.anyscale.com
            Id: AnyscaleCORSRule
      PublicAccessBlockConfiguration:
        BlockPublicAcls: True
        BlockPublicPolicy: True
        IgnorePublicAcls: True
        RestrictPublicBuckets: True
      Tags:
        - Key: anyscale-cloud-id
          Value: !Ref CloudID
        - Key: anyscale-deploy-environment
          Value: !Ref EnvironmentName

  S3BucketPolicy:
    Type: AWS::S3::BucketPolicy
    Description: Bucket policy that allow ray autoscaler role to access the bucket.
    Properties:
      Bucket: !Ref S3Bucket
      PolicyDocument:
        Version: 2012-10-17
        Statement:
          - Action:
              - "s3:PutObject"
              - "s3:DeleteObject"
              - "s3:GetObject"
              - "s3:ListBucket"
            Effect: Allow
            Resource:
              - !Sub arn:aws:s3:::${S3Bucket}
              - !Sub arn:aws:s3:::${S3Bucket}/*
            Principal:
              AWS:
                - !GetAtt
                  - nodeRole
                  - Arn
                - !GetAtt
                  - customerRole
                  - Arn
  EFS:
    Type: AWS::EFS::FileSystem
    Properties:
        BackupPolicy:
          Status: ENABLED
        Encrypted: true
        LifecyclePolicies:
          - TransitionToIA: AFTER_60_DAYS
        PerformanceMode: generalPurpose
        Encrypted: true
        ThroughputMode: bursting
        FileSystemTags:
          - Key: anyscale-cloud-id
            Value: !Ref CloudID
          - Key: anyscale-deploy-environment
            Value: !Ref EnvironmentName


  EFSMountTarget0:
    Type: AWS::EFS::MountTarget
    Properties:
        FileSystemId: !Ref EFS
        SecurityGroups:
          - !Ref AnyscaleSecurityGroup
        SubnetId: !Ref Subnet0

  EFSMountTarget1:
    Type: AWS::EFS::MountTarget
    Properties:
        FileSystemId: !Ref EFS
        SecurityGroups:
          - !Ref AnyscaleSecurityGroup
        SubnetId: !Ref Subnet1

  EFSMountTarget2:
    Type: AWS::EFS::MountTarget
    Properties:
        FileSystemId: !Ref EFS
        SecurityGroups:
          - !Ref AnyscaleSecurityGroup
        SubnetId: !Ref Subnet2

  EFSMountTarget3:
    Type: AWS::EFS::MountTarget
    Properties:
        FileSystemId: !Ref EFS
        SecurityGroups:
          - !Ref AnyscaleSecurityGroup
        SubnetId: !Ref Subnet3

  EFSMountTarget4:
    Type: AWS::EFS::MountTarget
    Properties:
        FileSystemId: !Ref EFS
        SecurityGroups:
          - !Ref AnyscaleSecurityGroup
        SubnetId: !Ref Subnet4

  EFSMountTarget5:
    Type: AWS::EFS::MountTarget
    Properties:
        FileSystemId: !Ref EFS
        SecurityGroups:
          - !Ref AnyscaleSecurityGroup
        SubnetId: !Ref Subnet5

  customerRole:
    Type: 'AWS::IAM::Role'
    Properties:
      RoleName: !Ref AnyscaleCrossAccountIAMRoleName
      AssumeRolePolicyDocument:
        Statement:
          - Action: 'sts:AssumeRole'
            Effect: Allow
            Principal:
              AWS: !Ref AnyscaleAWSAccountID
            Sid: 'AnyscaleControlPlaneAssumeRole'
            Condition:
              StringEquals:
                  sts:ExternalId: !Ref CloudID
        Version: 2012-10-17
      Path: /
      Tags:
        - Key: Name
          Value: !Ref AWS::StackName
        - Key: anyscale-cloud-id
          Value: !Ref CloudID
        - Key: anyscale-deploy-environment
          Value: !Ref EnvironmentName

  nodeRole:
    Type: 'AWS::IAM::Role'
    Properties:
      RoleName: !Ref ClusterNodeIAMRoleName
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/AmazonEC2ContainerRegistryReadOnly
      AssumeRolePolicyDocument:
        Statement:
          - Action: 'sts:AssumeRole'
            Effect: Allow
            Principal:
              Service: ec2.amazonaws.com
            Sid: 'Allow'
        Version: 2012-10-17
      Path: /
      Tags:
        - Key: Name
          Value: !Ref AWS::StackName
        - Key: anyscale-cloud-id
          Value: !Ref CloudID
        - Key: anyscale-deploy-environment
          Value: !Ref EnvironmentName

  IamInstanceProfile:
    Type: AWS::IAM::InstanceProfile
    Properties:
      InstanceProfileName: !Ref nodeRole
      Path: /
      Roles:
        - !Ref nodeRole

  IAMPermissionEC2SteadyState:
    Type: AWS::IAM::Policy
    Properties:
      PolicyDocument: !Ref AnyscaleCrossAccountIAMPolicySteadyState
      PolicyName: Anyscale_IAM_Policy_Steady_State
      Roles:
        - !Ref customerRole

  IAMPermissionServiceSteadyState:
    Type: AWS::IAM::Policy
    Properties:
      PolicyDocument: !Ref AnyscaleCrossAccountIAMPolicyServiceSteadyState
      PolicyName: Anyscale_IAM_Policy_Service_Steady_State
      Roles:
        - !Ref customerRole

  IAMPermissionEC2InitialRun:
    Type: AWS::IAM::Policy
    Properties:
      PolicyDocument: !Ref AnyscaleCrossAccountIAMPolicyInitialRun
      PolicyName: Anyscale_IAM_Policy_Initial_Setup
      Roles:
        - !Ref customerRole

Outputs:
  VPC:
    Description: A reference to the created VPC
    Value: !Ref VPC

  SubnetsWithAvailabilityZones:
    Description: A list of the subnets with their availability zones
    Value:
      Fn::ToJsonString:
        [{"subnet_id": !Ref Subnet0, "availability_zone": !GetAtt Subnet0.AvailabilityZone},{"subnet_id": !Ref Subnet1, "availability_zone": !GetAtt Subnet1.AvailabilityZone},{"subnet_id": !Ref Subnet2, "availability_zone": !GetAtt Subnet2.AvailabilityZone},{"subnet_id": !Ref Subnet3, "availability_zone": !GetAtt Subnet3.AvailabilityZone},{"subnet_id": !Ref Subnet4, "availability_zone": !GetAtt Subnet4.AvailabilityZone},{"subnet_id": !Ref Subnet5, "availability_zone": !GetAtt Subnet5.AvailabilityZone}]

  AnyscaleSecurityGroup:
    Description: Anyscale Security group
    Value: !Ref AnyscaleSecurityGroup

  S3Bucket:
    Description: Anyscale managed S3 bucket
    Value: !Ref S3Bucket

  EFS:
    Description: Anyscale managed EFS
    Value: !Ref EFS

  EFSMountTargetIP:
    Description: Anyscale managed EFS mount target
    Value: !GetAtt
      - EFSMountTarget0
      - IpAddress

  AnyscaleIAMRole:
    Description: ARN of the cross-account IAM role
    Value: !GetAtt
      - customerRole
      - Arn

  NodeIamRole:
    Description: ARN of the node IAM role
    Value: !GetAtt
      - nodeRole
      - Arn
