deploy-azure

#!/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 sys

import click

from src.python.azure import azure_login
from src.python.config import c as config
from src.python.debug import debug_break  # 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 DeployAzureCommand(DeployCommand):
    """
    Defines options specific for "deploy-azure" command.
    """

    # azure instace types capable of running isaac
    AZURE_OVKIT_INSTANCE_TYPES = [
        "Standard_NV36ads_A10_v5",
        "Standard_NV36adms_A10_v5",
        "Standard_NV72ads_A10_v5",
    ]

    @staticmethod
    def ovkit_instance_type_callback(ctx, param, value):
        if value not in DeployAzureCommand.AZURE_OVKIT_INSTANCE_TYPES:
            if (
                int(value) > len(DeployAzureCommand.AZURE_OVKIT_INSTANCE_TYPES)
                or int(value) < 1
            ):
                raise click.BadParameter(
                    colorize_error(
                        f"Invalid instance type: {value}. Valid values: "
                        + f"{[i+1 for i in range(len(DeployAzureCommand.AZURE_OVKIT_INSTANCE_TYPES))]}"
                    )
                )

            value = DeployAzureCommand.AZURE_OVKIT_INSTANCE_TYPES[int(value) - 1]

        return value

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

        # --region
        self.params.insert(
            # insert before --isaac option
            self.param_index("isaac"),
            click.core.Option(
                ("--region",),
                prompt=colorize_prompt(
                    "* Region (see "
                    + "https://azure.microsoft.com/en-us/pricing/details/virtual-machines/linux/#pricing"
                    + " for availability of 'NVadsA10_v5' instances)"
                ),
                callback=lambda ctx, param, value: value.lower().replace(" ", ""),
                default="West US 3",
                show_default=True,
                help="Azure region to deploy to. See "
                + "https://azure.microsoft.com/en-us/pricing/details/virtual-machines/linux/#pricing"
                + " for availability of 'NVadsA10_v5' instances.",
            ),
        )

        # --isaac-instance-type
        self.params.insert(
            # insert after --isaac option
            self.param_index("isaac") + 1,
            click.core.Option(
                ("--isaac-instance-type",),
                type=str,
                prompt=colorize_prompt("* Isaac Sim VM Type")
                + f"\n  1) {self.AZURE_OVKIT_INSTANCE_TYPES[0]}\t1x NVIDIA A10\t440 GB RAM"
                + f"\n  2) {self.AZURE_OVKIT_INSTANCE_TYPES[1]}\t1x NVIDIA A10\t880 GB RAM"
                + f"\n  3) {self.AZURE_OVKIT_INSTANCE_TYPES[2]}\t2x NVIDIA A10\t880 GB RAM\n"
                + "",
                default=config["azure_default_isaac_instance_type"],
                show_default=self.AZURE_OVKIT_INSTANCE_TYPES[
                    int(config["azure_default_isaac_instance_type"]) - 1
                ],
                callback=DeployAzureCommand.ovkit_instance_type_callback,
                help="Isaac Sim VM type. Currently supported: "
                + ", ".join(self.AZURE_OVKIT_INSTANCE_TYPES)
                + ".",
            ),
        )

        # --login/--no-login
        self.params.insert(
            len(self.params),
            click.core.Option(
                ("--login/--no-login",),
                default=True,
                show_default=True,
                help="Login or don't login into Azure. "
                + "Login can be skipped if you logged in before, to save time.",
            ),
        )

        # --resource-group
        self.params.insert(
            len(self.params),
            click.core.Option(
                ("--resource-group",),
                type=str,
                default="",
                help="A resource group to be used for the deployment. If empty, it will be created.",
            ),
        )

        # defaults

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


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

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

    def import_resource_group(self, cwd: str):
        """
        Import existing AZ resource group into terraform state
        cwd: directory where terraform scripts are located
        """
        debug = self.params["debug"]
        resource_group = self.params["resource_group"]
        state_dir = self.config["state_dir"]
        deployment_name = self.params["deployment_name"]

        shell_command(
            f"terraform import -var-file={state_dir}/{deployment_name}/.tfvars "
            + f"-state={self.config['state_dir']}/{deployment_name}/.tfstate "
            + "module.common.azurerm_resource_group.isa_rg"
            + f" {resource_group} {' > /dev/null' if not debug else ''}",
            verbose=debug,
            cwd=cwd,
        )

    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"],
                    "resource_group_name": (
                        self.params["resource_group"].split("/")[-1]
                        if len(self.params["resource_group"]) > 0
                        else f"{self.params['prefix']}.{self.params['deployment_name']}.rg"
                    ),
                }
            )

            # login into azure
            if self.params["login"]:
                azure_login(verbose=self.params["debug"])

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

            if len(self.params["resource_group"]) > 0:
                self.import_resource_group(cwd=f"{config['terraform_dir']}/azure")

            self.run_terraform(cwd=f"{config['terraform_dir']}/azure")

        # create ansible inventory file
        self.create_ansible_inventory()

        # export ssh key from terraform
        self.export_ssh_key()

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

        # connection info for the user
        self.output_deployment_info()


@click.command(cls=DeployAzureCommand)
def main(**params):
    AzureDeployer(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)