deploy-aws

#!/usr/bin/env python3

# region copyright
# Copyright 2023 NVIDIA 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.
# endregion

import os
import re
import sys

import click

from src.python.config import c as config
from src.python.debug import debug_break, debug_start  # noqa
from src.python.deploy_command import DeployCommand
from src.python.deployer import Deployer
from src.python.utils import (
    colorize_error,
    colorize_info,
    colorize_prompt,
    shell_command,
)


class DeployAWSCommand(DeployCommand):
    """
    Defines options specific for "deploy-aws" command.
    """

    # supported AWS instance types for DRIVE Sim/Isaac Sim
    # @see https://aws.amazon.com/ec2/instance-types/g5/
    AWS_OVKIT_INSTANCE_TYPES = [
        # T4 - https://aws.amazon.com/ec2/instance-types/g4/
        "g4dn.2xlarge",
        "g4dn.4xlarge",
        "g4dn.8xlarge",
        "g4dn.16xlarge",
        "g4dn.12xlarge",
        "g4dn.metal",
        # A10G - https://aws.amazon.com/ec2/instance-types/g5/
        "g5.2xlarge",
        "g5.4xlarge",
        "g5.8xlarge",
        "g5.16xlarge",
        "g5.12xlarge",
        "g5.24xlarge",
        "g5.48xlarge",
        # L4 - https://aws.amazon.com/ec2/instance-types/g6
        "g6.2xlarge",
        "g6.4xlarge",
        "g6.8xlarge",
        "g6.16xlarge",
        "gr6.4xlarge",
        "gr6.8xlarge",
        "g6.12xlarge",
        "g6.24xlarge",
        "g6.48xlarge",
    ]

    @staticmethod
    def ovkit_instance_type_callback(ctx, param, value):
        """
        Called after parsing --isaac-instance-type option
        """

        # warn of "g5.xlarge" being too small
        if "g5.xlarge" == value:
            raise click.BadParameter(colorize_error("g5.xlarge is not supported. "))

        # warn of other unsupported types
        if value not in DeployAWSCommand.AWS_OVKIT_INSTANCE_TYPES:
            raise click.BadParameter(
                colorize_error(
                    f"Invalid instance type: {value}. Choose one of: "
                    + ", ".join(DeployAWSCommand.AWS_OVKIT_INSTANCE_TYPES)
                    + "."
                )
            )

        return value

    @staticmethod
    def aws_access_key_id_callback(ctx, param, value):
        """
        Called after parsing --aws-access-key-id option
        """

        if not re.match("^[A-Z0-9]{20}$", value):
            raise click.BadParameter(
                colorize_error(
                    "Invalid AWS Access Key ID. "
                    + "It should be 20 characters long and contain only uppercase letters and digits."
                )
            )

        return value

    @staticmethod
    def aws_secret_access_key_callback(ctx, param, value):
        """
        Called after parsing --aws-secret-access-key option
        """

        if not re.match("^[A-Za-z0-9/+]{40}$", value):
            raise click.BadParameter(
                colorize_error(
                    "Invalid AWS Secret Access Key. "
                    + "It should be 40 characters long and contain only letters, digits, '/' and '+' signs."
                )
            )

        return value

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        # --isaac-instance-type
        self.params.insert(
            # insert after --isaac option
            self.param_index("isaac") + 1,
            click.core.Option(
                ("--isaac-instance-type",),
                prompt=colorize_prompt(
                    "* Isaac Sim Instance Type "
                    + "(https://docs.aws.amazon.com/dlami/latest/devguide/gpu.html, "
                    + "G5, G4dn and G6 instances are supported)"
                ),
                show_default=True,
                default=config["aws_default_isaac_instance_type"],
                callback=DeployAWSCommand.ovkit_instance_type_callback,
                help="Isaac Sim instance type. Currently supported: "
                + ", ".join(self.AWS_OVKIT_INSTANCE_TYPES)
                + ".",
            ),
        )

        # --region
        self.params.insert(
            # insert before --isaac option
            self.param_index("isaac"),
            click.core.Option(
                ("--region",),
                show_default=True,
                default="US East 1",
                prompt=colorize_prompt("* AWS Region"),
                callback=lambda ctx, param, value: value.lower().replace(" ", "-"),
            ),
        )

        # --aws-access-key-id
        self.params.insert(
            len(self.params),
            click.core.Option(
                ("--aws-access-key-id",),
                prompt=colorize_prompt("* AWS Access Key ID"),
                default=os.environ.get("AWS_ACCESS_KEY_ID", ""),
                callback=DeployAWSCommand.aws_access_key_id_callback,
                show_default="AWS_ACCESS_KEY_ID environment variable",
            ),
        )

        # --aws-secret-access-key
        self.params.insert(
            len(self.params),
            click.core.Option(
                ("--aws-secret-access-key",),
                prompt=colorize_prompt("* AWS Secret Access Key"),
                default=os.environ.get("AWS_SECRET_ACCESS_KEY", ""),
                callback=DeployAWSCommand.aws_secret_access_key_callback,
                show_default="AWS_SECRET_ACCESS_KEY environment variable",
            ),
        )

        # --aws-session-token
        self.params.insert(
            len(self.params),
            click.core.Option(
                ("--aws-session-token",),
                prompt=colorize_prompt(
                    "* AWS Session Token (if using temporary credentials, otherwise leave empty)"
                ),
                default=os.environ.get("AWS_SESSION_TOKEN", ""),
                show_default="AWS_SESSION_TOKEN environment variable",
            ),
        )

        # --ovami
        # deploy OV AMI instance?
        self.params.insert(
            len(self.params),
            click.core.Option(
                ("--ovami/--no-ovami",),
                default=False,
                prompt=False,
                hidden=True,
            ),
        )

        # defaults

        self.params[self.param_index("from_image")].default = config[
            "aws_default_from_image"
        ]


class AWSDeployer(Deployer):
    """
    Deploys stuff to Azure
    """

    def __init__(self, params, config):
        super().__init__(params, config)

    def _output_deployment_info(self, print_text=True):
        extra_info = ""

        self.output_deployment_info(extra_text=extra_info, print_text=print_text)

    def main(self):
        # ask what to do if deployment already exists
        self.ask_existing_behavior()

        # check if ngc api key is valid and has access to Isaac Sim
        if self.params["isaac"]:
            self.validate_ngc_api_key(self.params["isaac_image"])

        if self.existing_behavior != "run_ansible":
            # create tfvars file, deal with existing deployment
            self.create_tfvars(
                {
                    "region": self.params["region"],
                    "ovami_enabled": self.params["ovami"],
                    "aws_access_key_id": self.params["aws_access_key_id"],
                    "aws_secret_access_key": self.params["aws_secret_access_key"],
                    "aws_session_token": self.params["aws_session_token"],
                }
            )

            # run terraform
            click.echo(colorize_info("* Running Terraform..."))
            self.initialize_terraform(cwd=f"{config['terraform_dir']}/aws")
            self.run_terraform(cwd=f"{config['terraform_dir']}/aws")

        # export ssh key from terraform
        self.export_ssh_key()

        # create ansible inventory file
        self.create_ansible_inventory()

        # save instructions
        self._output_deployment_info(print_text=False)

        # run ansible
        self.run_all_ansible()

        # upload user data
        if self.params["upload"]:
            self.upload_user_data()

        # print info for the user
        self._output_deployment_info()


@click.command(cls=DeployAWSCommand)
def main(**params):
    AWSDeployer(params, config).main()


if __name__ == "__main__":
    if os.path.exists("/.dockerenv"):
        # we're in docker, run command
        main()
    else:
        # we're outside, start docker container first
        shell_command(f"./run '{' '.join(sys.argv)}'", verbose=True)