upload

#!/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.config import c as config
from src.python.debug import debug_break  # noqa
from src.python.utils import (
    colorize_info,
    colorize_result,
    deployments,
    format_app_name,
    read_tf_output,
    shell_command,
)

APP_NAMES = [
    "isaac",
    "ovami",
]


@click.command()
@click.option(
    "--debug/--no-debug",
    default=False,
    show_default=True,
)
@click.argument(
    "deployment_name",
    required=True,
    type=click.Choice(deployments()),
)
@click.option(
    "--remote-dir",
    default=config["default_remote_uploads_dir"],
    prompt=False,
    show_default=True,
    help="Remote directory to upload results to.",
)
@click.option(
    "--delete/--no-delete",
    default=True,
    show_default=True,
    help="Alow deleting when syncing results?",
)
@click.option(
    "--app",
    default="all",
    type=click.Choice(["all"] + APP_NAMES),
    help="Application to upload to.",
)
@click.argument("app", required=False)
def main(
    app: str,
    debug: bool,
    delete: bool,
    remote_dir: str,
    deployment_name: str,
):
    apps = APP_NAMES if app == "all" else [app]

    for app_name in apps:
        ip = read_tf_output(deployment_name, f"{app_name}_ip", verbose=debug)

        if ip in ["NA", ""]:
            # show "not found" message if:
            # - app is explicitly specified
            # - debug is enabled
            if app != "all" or debug:
                click.echo(
                    colorize_info(f"* No {format_app_name(app_name)} instance found")
                )
            continue

        ssh_key_file = f"{config['state_dir']}/{deployment_name}/key.pem"
        uploads_dir = config["uploads_dir"]
        os.makedirs(uploads_dir, exist_ok=True)

        # strip trailing slashes
        remote_dir = remote_dir.rstrip("/")
        uploads_dir = uploads_dir.rstrip("/")

        click.echo(
            colorize_info(
                f'* Uploading from "{uploads_dir}" to "{remote_dir}" @ '
                + f"{format_app_name(app_name)} instance ({ip})..."
            )
        )

        # make sure remote directory exists
        shell_command(
            f"ssh -i {ssh_key_file} -o StrictHostKeyChecking=no ubuntu@{ip}"
            + f" \"[ ! -d '{remote_dir}' ] && mkdir -p '{remote_dir}'\"",
            verbose=debug,
            exit_on_error=False,
            capture_output=False,
        )

        rsync_command = (
            "rsync -avz --progress --rsync-path 'rsync'"
            + (" --delete" if delete else "")
            + f" -e 'ssh -i {ssh_key_file} -o StrictHostKeyChecking=no'"
            + f' "{uploads_dir}/" ubuntu@{ip}:"{remote_dir}"'
        )

        shell_command(
            rsync_command,
            verbose=debug,
            exit_on_error=True,
            capture_output=False,
        )

        # get size of output_dir
        output_dir_size = shell_command(
            f"ssh -i {ssh_key_file} -o StrictHostKeyChecking=no ubuntu@{ip}"
            + f" \"du -hs '{remote_dir}'\"",
            verbose=debug,
            exit_on_error=True,
            capture_output=True,
        ).stdout.decode()
        output_dir_size = output_dir_size.split("\t")[0]

        click.echo(
            colorize_result(
                f'* Uploaded {output_dir_size} to "{remote_dir}" @ {format_app_name(app_name)} instance ({ip})'
            )
        )


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)