dgxcloud.py

# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# 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.

import base64
import glob
import json
import logging
import os
import subprocess
import tempfile
import time
from dataclasses import dataclass, field
from enum import Enum
from pathlib import Path
from typing import Any, Iterable, Optional

import requests
from invoke.context import Context

from nemo_run.config import RUNDIR_NAME, get_nemorun_home
from nemo_run.core.execution.base import Executor, ExecutorMacros
from nemo_run.core.execution.launcher import FaultTolerance, Launcher, Torchrun
from nemo_run.core.execution.utils import fill_template
from nemo_run.core.frontend.console.api import CONSOLE
from nemo_run.core.packaging.base import Packager
from nemo_run.core.packaging.git import GitArchivePackager

logger = logging.getLogger(__name__)


class DGXCloudState(Enum):
    CREATING = "Creating"
    INITIALIZING = "Initializing"
    RESUMING = "Resuming"
    PENDING = "Pending"
    DELETING = "Deleting"
    RUNNING = "Running"
    UPDATING = "Updating"
    STOPPED = "Stopped"
    STOPPING = "Stopping"
    DEGRADED = "Degraded"
    FAILED = "Failed"
    COMPLETED = "Completed"
    TERMINATING = "Terminating"
    UNKNOWN = "Unknown"


@dataclass(kw_only=True)
class DGXCloudExecutor(Executor):
    """
    Dataclass to configure a DGX Executor.

    This executor integrates with a DGX cloud endpoint for launching jobs
    via a REST API. It acquires an auth token, identifies the project/cluster,
    and launches jobs with a specified command. It can be adapted to meet user
    authentication and job-submission requirements on DGX.
    """

    base_url: str
    kube_apiserver_url: str
    app_id: str
    app_secret: str
    project_name: str
    container_image: str
    pvc_nemo_run_dir: str
    launched_from_cluster: bool = False
    nodes: int = 1
    gpus_per_node: int = 0
    nprocs_per_node: int = 1
    pvc_job_dir: str = field(init=False, default="")
    pvcs: list[dict[str, Any]] = field(default_factory=list)
    distributed_framework: str = "PyTorch"
    custom_spec: dict[str, Any] = field(default_factory=dict)

    def get_auth_token(self) -> Optional[str]:
        url = f"{self.base_url}/token"
        payload = {
            "grantType": "app_token",
            "appId": self.app_id,
            "appSecret": self.app_secret,
        }

        n_attempts = 0
        while n_attempts < 3:
            try:
                response = requests.post(url, json=payload, headers=self._default_headers())
                response_text = response.text.strip()
                auth_token = json.loads(response_text).get("accessToken", None)  # [1]
                if auth_token:
                    return auth_token

                raise ValueError(f"Failed to retrieve auth token; response was: {response_text}")

            except Exception as e:
                logger.error("Failed to retrieve auth token; error was: %s", e)
                time.sleep(10)
                n_attempts += 1

        logger.error("Failed to retrieve auth token after 3 attempts.")
        return None

    def get_project_and_cluster_id(self, token: str) -> tuple[Optional[str], Optional[str]]:
        url = f"{self.base_url}/org-unit/projects"
        headers = self._default_headers(token=token)
        response = requests.get(url, headers=headers)
        projects = json.loads(response.text.strip()).get("projects", [])
        project_id, cluster_id = None, None
        for prj in projects:
            if not self.project_name or prj["name"] == self.project_name:  # [2]
                project_id, cluster_id = prj["id"], prj["clusterId"]
                logger.debug(
                    "Found project '%s' (%s) on cluster '%s'", prj["name"], project_id, cluster_id
                )
                break
        return project_id, cluster_id

    def copy_directory_data_command(self, local_dir_path: str, dest_path: str) -> str:
        with tempfile.TemporaryDirectory() as temp_dir:
            tarball_path = os.path.join(temp_dir, "archive.tar.gz")
            subprocess.run(f"tar -czf {tarball_path} -C {local_dir_path} .", shell=True, check=True)
            with open(tarball_path, "rb") as file:
                file_data = file.read()
            encoded_data = base64.b64encode(file_data).decode("utf-8")

            # Delete and recreate directory if it already exists, command to decode base64 data, save to a file, and extract inside the pod
            cmd = f"rm -rf {dest_path} && mkdir -p {dest_path} && echo {encoded_data} | base64 -d > {dest_path}/archive.tar.gz && tar -xzf {dest_path}/archive.tar.gz -C {dest_path} && rm {dest_path}/archive.tar.gz"
            return cmd

    def create_data_mover_workload(self, token: str, project_id: str, cluster_id: str):
        """
        Creates a CPU only workload to move job directory into PVC using the provided project/cluster IDs.
        """

        cmd = self.copy_directory_data_command(self.job_dir, self.pvc_job_dir)

        url = f"{self.base_url}/workloads/workspaces"
        headers = self._default_headers(token=token)

        payload = {
            "name": "data-mover",
            "useGivenNameAsPrefix": True,
            "projectId": project_id,
            "clusterId": cluster_id,
            "spec": {
                "command": "sh -c",
                "args": f"'{cmd}'",
                "image": "busybox:1.37.0",
                "storage": {"pvc": self.pvcs},
            },
        }

        response = requests.post(url, json=payload, headers=headers)

        logger.debug(
            "Created workload; response code=%s, content=%s",
            response.status_code,
            response.text.strip(),
        )

        return response

    def delete_workload(self, token: str, workload_id: str):
        url = f"{self.base_url}/workloads/workspaces/{workload_id}"
        headers = self._default_headers(token=token)

        response = requests.delete(url, headers=headers)

        logger.debug(
            "Delete interactive workspace; response code=%s, content=%s",
            response.status_code,
            response.text.strip(),
        )
        return response

    def move_data(self, token: str, project_id: str, cluster_id: str, sleep: float = 10) -> None:
        """
        Moves job directory into PVC and deletes the workload after completion
        """

        resp = self.create_data_mover_workload(token, project_id, cluster_id)
        if resp.status_code not in [200, 202]:
            raise RuntimeError(
                f"Failed to create data mover workload, status_code={resp.status_code}, reason={resp.text}"
            )

        resp_json = resp.json()
        workload_id = resp_json["workloadId"]
        status = DGXCloudState(resp_json["actualPhase"])

        logger.info(f"Successfully created data movement workload {workload_id} on DGXCloud")

        while status in [
            DGXCloudState.PENDING,
            DGXCloudState.CREATING,
            DGXCloudState.INITIALIZING,
            DGXCloudState.RUNNING,
        ]:
            time.sleep(sleep)
            status = self.status(workload_id)
            logger.debug(
                f"Polling data movement workload {workload_id}'s status. Current status is: {status}"
            )

        if status is not DGXCloudState.COMPLETED:
            raise RuntimeError(f"Failed to move data to PVC. Workload status is {status}")

        resp = self.delete_workload(token, workload_id)
        if resp.status_code >= 200 and resp.status_code < 300:
            logger.info(
                "Successfully deleted data movement workload %s on DGXCloud with response code %d",
                workload_id,
                resp.status_code,
            )
        else:
            logger.error(
                "Failed to delete data movement workload %s, response code=%d, reason=%s",
                workload_id,
                resp.status_code,
                resp.text,
            )

    def create_training_job(
        self, token: str, project_id: str, cluster_id: str, name: str
    ) -> requests.Response:
        """
        Creates a training job on DGX Cloud using the provided project/cluster IDs.
        For multi-node jobs, creates a distributed workload. Otherwise creates a single-node training.

        Args:
            token: Authentication token for DGX Cloud API
            project_id: ID of the project to create the job in
            cluster_id: ID of the cluster to create the job on
            name: Name for the job

        Returns:
            Response object from the API request
        """
        # Validate inputs
        if not token or not project_id or not cluster_id:
            raise ValueError("Token, project ID, and cluster ID are required")

        if self.nodes < 1:
            raise ValueError("Node count must be at least 1")

        if len(name) >= 35:
            logger.warning(
                "Training name can only be max 35 characters. Shortening name to 35 characters..."
            )
            name = name[:34]

        # Common payload elements
        common_payload = {
            "name": name,
            "useGivenNameAsPrefix": True,
            "projectId": project_id,
            "clusterId": cluster_id,
        }

        # Common spec elements
        common_spec = {
            "command": f"/bin/bash {self.pvc_job_dir}/launch_script.sh",
            "image": self.container_image,
            "compute": {"gpuDevicesRequest": self.gpus_per_node},
            "storage": {"pvc": self.pvcs},
            "environmentVariables": [
                {"name": key, "value": value} for key, value in self.env_vars.items()
            ],
            **self.custom_spec,
        }

        # Determine endpoint and build payload based on node count
        if self.nodes > 1:
            url = f"{self.base_url}/workloads/distributed"

            # Add distributed-specific parameters
            distributed_spec = {
                "distributedFramework": self.distributed_framework,
                "minReplicas": self.nodes,
                "maxReplicas": self.nodes,
                "numWorkers": self.nodes,
            }

            payload = {**common_payload, "spec": {**common_spec, **distributed_spec}}
        else:
            url = f"{self.base_url}/workloads/trainings"
            payload = {**common_payload, "spec": common_spec}

        headers = self._default_headers(token=token)
        response = requests.post(url, json=payload, headers=headers)

        logger.info(json.dumps(payload))
        logger.debug(
            "Created %s job; response code=%s, content=%s",
            "distributed" if self.nodes > 1 else "training",
            response.status_code,
            response.text.strip(),
        )

        return response

    def launch(self, name: str, cmd: list[str]) -> tuple[str, str]:
        name = name.replace("_", "-").replace(".", "-").lower()  # to meet K8s requirements
        logger.info(f"workload name:{name}")
        token = self.get_auth_token()
        if not token:
            raise RuntimeError("Failed to get auth token")

        project_id, cluster_id = self.get_project_and_cluster_id(token)
        if not project_id or not cluster_id:
            raise RuntimeError("Unable to determine project/cluster IDs for job submission")

        # prepare launch script and move data to PVC
        launch_script = f"""
ln -s {self.pvc_job_dir}/ /nemo_run
cd /nemo_run/code
mkdir -p {self.pvc_job_dir}/logs
{" ".join(cmd)} 2>&1 | tee -a {self.pvc_job_dir}/log_$HOSTNAME.out {self.pvc_job_dir}/log-allranks_0.out
"""
        with open(os.path.join(self.job_dir, "launch_script.sh"), "w+") as f:
            f.write(launch_script)

        if not self.launched_from_cluster:
            logger.info("Creating data movement workload")
            self.move_data(token, project_id, cluster_id)

        logger.info("Creating training workload")
        resp = self.create_training_job(token, project_id, cluster_id, name)
        if resp.status_code not in [200, 202]:
            raise RuntimeError(
                f"Failed to create job, status_code={resp.status_code}, reason={resp.text}"
            )

        r_json = resp.json()
        job_id = r_json["workloadId"]
        status = r_json["actualPhase"]
        return job_id, status

    def nnodes(self) -> int:
        return self.nodes

    def nproc_per_node(self) -> int:
        # Default to the number of GPUs specified per node
        # If user doesn't want GPUs, can run multiple processes with CPU only
        if self.gpus_per_node:
            return self.gpus_per_node
        elif self.nprocs_per_node:
            return self.nprocs_per_node
        return 1

    def status(self, job_id: str) -> Optional[DGXCloudState]:
        url = f"{self.base_url}/workloads/{job_id}"
        token = self.get_auth_token()
        if not token:
            logger.error("Failed to retrieve auth token for status request.")
            return None

        headers = self._default_headers(token=token)
        response = requests.get(url, headers=headers)
        if response.status_code != 200:
            return DGXCloudState("Unknown")

        r_json = response.json()
        return DGXCloudState(r_json["phase"])

    def fetch_logs(
        self,
        job_id: str,
        stream: bool,
        stderr: Optional[bool] = None,
        stdout: Optional[bool] = None,
    ) -> Iterable[str]:
        while self.status(job_id) != DGXCloudState.RUNNING:
            logger.info("Waiting for job to start...")
            time.sleep(15)

        cmd = ["tail"]

        if stream:
            cmd.append("-f")

        # setting linked PVC job directory
        nemo_run_home = get_nemorun_home()
        job_subdir = self.job_dir[len(nemo_run_home) + 1 :]  # +1 to remove the initial backslash
        self.pvc_job_dir = os.path.join(self.pvc_nemo_run_dir, job_subdir)

        files = []
        while len(files) < self.nodes:
            files = list(glob.glob(f"{self.pvc_job_dir}/log_*.out"))
            files = [f for f in files if "log-allranks_0" not in f]
            logger.info(
                f"Waiting for {self.nodes + 1 - len(files)} log files to be created in {self.pvc_job_dir}..."
            )
            time.sleep(3)

        cmd.extend(files)

        logger.info(f"Attempting to stream logs with command: {cmd}")

        proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, text=True, bufsize=1)

        if stream:
            while True:
                try:
                    for line in iter(proc.stdout.readline, ""):
                        if (
                            line
                            and not line.rstrip("\n").endswith(".out <==")
                            and line.rstrip("\n") != ""
                        ):
                            yield f"{line}"
                        if proc.poll() is not None:
                            break
                except Exception as e:
                    logger.error(f"Error streaming logs: {e}")
                    time.sleep(3)
                    continue

        else:
            try:
                for line in iter(proc.stdout.readline, ""):
                    if line:
                        yield line.rstrip("\n")
                    if proc.poll() is not None:
                        break
            finally:
                proc.terminate()
                proc.wait(timeout=2)

    def cancel(self, job_id: str):
        # Retrieve the authentication token for the REST calls
        token = self.get_auth_token()
        if not token:
            logger.error("Failed to retrieve auth token for cancellation request.")
            return

        # Build the DELETE request to cancel the job
        url = f"{self.base_url}/workloads/distributed/{job_id}/suspend"
        headers = self._default_headers(token=token)

        response = requests.get(url, headers=headers)
        if response.status_code >= 200 and response.status_code < 300:
            logger.info(
                "Successfully cancelled job %s on DGX with response code %d",
                job_id,
                response.status_code,
            )
        else:
            logger.error(
                "Failed to cancel job %s, response code=%d, reason=%s",
                job_id,
                response.status_code,
                response.text,
            )

    def _setup_launcher(self):
        super()._setup_launcher()
        launcher = self.launcher
        if launcher and isinstance(launcher, (FaultTolerance, Torchrun)):
            self.torchrun_nproc_per_node = self.nprocs_per_node
            self.ntasks_per_node = 1
            CONSOLE.log(
                f"Detected {launcher.__class__.__name__} launcher, setting ntasks_per_node=1 and torchrun_nproc_per_node={self.torchrun_nproc_per_node}"
            )

        if launcher and isinstance(launcher, FaultTolerance):
            base_dir = os.path.join(self.job_dir, Path(self.job_dir).name)
            launcher.cfg_path = os.path.join(base_dir, f"{self.job_name}_ft_cfg.yml")
            launcher.finished_flag_file = os.path.join(
                "/", RUNDIR_NAME, f"{self.job_name}_finished_flag"
            )
            launcher.job_results_file = os.path.join(base_dir, f"{self.job_name}_job_results")

    def cleanup(self, handle: str): ...

    def assign(
        self,
        exp_id: str,
        exp_dir: str,
        task_id: str,
        task_dir: str,
    ):
        self.job_name = task_id
        self.experiment_dir = exp_dir
        self.job_dir = os.path.join(exp_dir, task_dir)
        assert any(
            map(
                lambda x: os.path.commonpath(
                    [os.path.abspath(x["path"]), os.path.abspath(self.pvc_nemo_run_dir)]
                )
                == os.path.abspath(x["path"]),
                self.pvcs,
            )
        ), (
            f"Need to specify at least one PVC containing {self.pvc_nemo_run_dir}. Update your PVC path or pvc_nemo_run_dir."
        )

        # setting linked PVC job directory
        nemo_run_home = get_nemorun_home()
        job_subdir = self.job_dir[len(nemo_run_home) + 1 :]  # +1 to remove the initial backslash
        self.pvc_job_dir = os.path.join(self.pvc_nemo_run_dir, job_subdir)

        logger.info(
            "PVC job directory set as:  %s",
            self.pvc_job_dir,
        )
        self.experiment_id = exp_id

    def get_launcher_prefix(self) -> Optional[list[str]]:
        launcher = self.get_launcher()
        if launcher.nsys_profile:
            return launcher.get_nsys_prefix(profile_dir="/nemo_run")

    def package_configs(self, *cfgs: tuple[str, str]) -> list[str]:
        filenames = []
        basepath = os.path.join(self.job_dir, "configs")
        for name, cfg in cfgs:
            filename = os.path.join(basepath, name)
            os.makedirs(os.path.dirname(filename), exist_ok=True)
            with open(filename, "w") as f:
                f.write(cfg)
            filenames.append(
                os.path.join(
                    "/nemo_run/configs",
                    name,
                )
            )
        return filenames

    def package(self, packager: Packager, job_name: str):
        assert self.experiment_id, "Executor not assigned to an experiment."
        if isinstance(packager, GitArchivePackager):
            output = subprocess.run(
                ["git", "rev-parse", "--show-toplevel"],
                check=True,
                stdout=subprocess.PIPE,
            )
            path = output.stdout.splitlines()[0].decode()
            base_path = Path(path).absolute()
        else:
            base_path = Path(os.getcwd()).absolute()

        local_pkg = packager.package(base_path, self.job_dir, job_name)
        local_code_extraction_path = os.path.join(self.job_dir, "code")
        ctx = Context()
        ctx.run(f"mkdir -p {local_code_extraction_path}")

        if self.get_launcher().nsys_profile:
            remote_nsys_extraction_path = os.path.join(
                self.job_dir, self.get_launcher().nsys_folder
            )
            ctx.run(f"mkdir -p {remote_nsys_extraction_path}")
        if local_pkg:
            ctx.run(
                f"tar -xvzf {local_pkg} -C {local_code_extraction_path} --ignore-zeros", hide=True
            )

    def macro_values(self) -> Optional[ExecutorMacros]:
        return None

    def _default_headers(self, token: Optional[str] = None) -> dict:
        headers = {
            "Accept": "application/json",
            "Content-Type": "application/json",
        }
        if token:
            headers["Authorization"] = f"Bearer {token}"
        return headers


@dataclass(kw_only=True)
class DGXCloudRequest:
    launch_cmd: list[str]
    jobs: list[str]
    executor: DGXCloudExecutor
    max_retries: int
    extra_env: dict[str, str]
    launcher: Optional[Launcher] = None

    def materialize(self) -> str:
        """Creates the content of a DGXC entrypoint script."""

        # 1. Environment Variables
        # Combine executor defaults with extra envs
        env_vars = []
        full_env_vars = self.executor.env_vars | self.extra_env
        for key, value in full_env_vars.items():
            env_vars.append(f"export {key.upper()}={value}")

        # 3. Prepare Template Variables
        vars_to_fill = {
            "max_retries": self.max_retries,
            "env_vars": env_vars,
            "training_command": " ".join(self.launch_cmd),
            "ft_enabled": bool(self.launcher and isinstance(self.launcher, FaultTolerance)),
        }

        # 4. Fault Tolerance Injection
        if self.launcher and isinstance(self.launcher, FaultTolerance):
            assert (
                self.launcher.cfg_path
                and self.launcher.finished_flag_file
                and self.launcher.job_results_file
            ), "Fault Tolerance requires cfg_path, finished_flag_file, and job_results_file"

            vars_to_fill["fault_tol_cfg_path"] = self.launcher.cfg_path
            vars_to_fill["fault_tol_finished_flag_file"] = self.launcher.finished_flag_file
            vars_to_fill["fault_tol_job_results_file"] = self.launcher.job_results_file

        # Render the template
        entrypoint_script = fill_template("dgxc.sh.j2", vars_to_fill)
        return entrypoint_script

    def __repr__(self) -> str:
        return f"""# DGXC Entrypoint Script Request
# Executor: {self.executor.__class__.__name__}
# Jobs: {self.jobs}
# ---------------------------------------------------
{self.materialize()}
"""