deploy-aws

#!/usr/bin/env python3

# region copyright
# Copyright 2023-2026 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.aws import aws_validate_credentials
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",
        # AWS L40S - https://aws.amazon.com/ec2/instance-types/g6e/
        "g6e.xlarge",
        "g6e.2xlarge",
        "g6e.4xlarge",
        "g6e.8xlarge",
        "g6e.12xlarge",
        "g6e.16xlarge",
        "g6e.24xlarge",
        "g6e.48xlarge",
        # AWS RTX PRO 6000 Blackwell Server Edition - https://aws.amazon.com/ec2/instance-types/g7e/
        "g7e.2xlarge",
        "g7e.4xlarge",
        "g7e.8xlarge",
        "g7e.12xlarge",
        "g7e.24xlarge",
        "g7e.48xlarge",
    ]

    @staticmethod
    def ovkit_instance_type_callback(ctx, param, value):
        """
        Called after parsing --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 availability_zone_callback(ctx, param, value):
        """
        Called after parsing --availability-zone option.
        Accepts a full AZ name ("us-west-2b"), a bare zone letter ("b"),
        or an empty value (auto-select the first zone offering the instance type).
        """

        value = (value or "").strip().lower()

        if value == "":
            return ""

        region = ctx.params.get("region", "")

        # accept a bare zone suffix (e.g. "b") and prefix it with the region
        if region and re.fullmatch(r"[a-z]", value):
            value = f"{region}{value}"

        # sanity check: the zone must belong to the selected region
        if region and not value.startswith(region):
            raise click.BadParameter(
                colorize_error(
                    f'Availability zone "{value}" is not in region "{region}".'
                )
            )

        return value

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

        # --instance-type
        self.params.insert(
            # insert before --ingress-cidrs option
            self.param_index("ingress_cidrs"),
            click.core.Option(
                (
                    "--instance-type",
                    "--isaac-workstation-instance-type",
                    "isaac_workstation_instance_type",
                ),
                prompt=colorize_prompt(
                    "* Instance Type "
                    + "(G5, G4dn, G6 and G6e are supported, see https://docs.aws.amazon.com/dlami/latest/devguide/gpu.html)"
                ),
                show_default=True,
                default=config["aws_default_isaac_workstation_instance_type"],
                callback=DeployAWSCommand.ovkit_instance_type_callback,
                help="Instance type. Currently supported: "
                + ", ".join(self.AWS_OVKIT_INSTANCE_TYPES)
                + ".",
            ),
        )

        # --region
        self.params.insert(
            # insert before --ingres-cidrs option
            self.param_index("ingress_cidrs"),
            click.core.Option(
                ("--region",),
                show_default=True,
                default="US East 1",
                help="AWS Region. Can be entered as 'us-east-1' or 'US East 1'."
                " See https://aws.amazon.com/about-aws/global-infrastructure/regions_az/",
                prompt=colorize_prompt(
                    "* AWS Region (e.g. us-east-1 or US East 1,"
                    " see https://aws.amazon.com/about-aws/global-infrastructure/regions_az/)"
                ),
                callback=lambda ctx, param, value: value.lower().replace(" ", "-"),
            ),
        )

        # --availability-zone
        self.params.insert(
            # insert before --ingress-cidrs option (and after --region, so the
            # region is parsed before the AZ callback runs)
            self.param_index("ingress_cidrs"),
            click.core.Option(
                ("--availability-zone", "--az"),
                default="",
                show_default=True,
                callback=DeployAWSCommand.availability_zone_callback,
                prompt=colorize_prompt(
                    "* Availability Zone (e.g. us-west-2b, or just \"b\"; "
                    "leave empty to auto-select. Use this to route around "
                    '"InsufficientInstanceCapacity" errors in a given zone)'
                ),
                help="AWS availability zone for the instance, e.g. 'us-west-2b'."
                " Leave empty to auto-select the first zone that offers the"
                " instance type. Useful to avoid per-zone GPU capacity shortages.",
            ),
        )

        # 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()

        # update region in AWS config if specified
        if self.params.get("region"):
            from src.python.aws import aws_cli_set

            aws_cli_set("region", self.params["region"], verbose=self.params["debug"])

        if self.existing_behavior != "run_ansible":
            # create tfvars file, deal with existing deployment
            self.create_tfvars(
                {
                    "region": self.params["region"],
                    "availability_zone": self.params.get("availability_zone", ""),
                }
            )

            # 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")

            # display selected AMI
            ami_id = self.tf_output("isaac_workstation_ami_id", default="")
            ami_name = self.tf_output("isaac_workstation_ami_name", default="")
            if ami_id:
                click.echo(colorize_info(f"* Base image: {ami_id} ({ami_name})"))

        # 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"):
        # set DEBUG env var if --debug option is used, so that pre-click code can also see debug mode
        debug = "--debug" in sys.argv or os.environ.get("DEBUG", "0") == "1"
        # validate AWS credentials and prompt for new ones if needed
        aws_validate_credentials(verbose=debug)
        # run command
        main()
    else:
        # we're outside, start docker container first
        shell_command(f"./run '{' '.join(sys.argv)}'", verbose=True)