cosmos-framework/cosmos_framework/utils/launch.py

# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: OpenMDW-1.1

import argparse
import os
import sys
import time

import torch
from omegaconf import OmegaConf

from cosmos_framework.utils.config import Config
from cosmos_framework.utils import distributed, log
from cosmos_framework.utils.cluster_env import get_cluster_env
from cosmos_framework.utils.easy_io import easy_io
from cosmos_framework.utils.env_parsers.cred_env_parser import CRED_ENVS
from cosmos_framework.utils.wandb_util import set_wandb_job_info

# Global variable to track S3 readiness
S3_READY = False


def log_reproducible_setup(config: Config, args: argparse.Namespace) -> None:
    """
    Configures the environment for reproducibility of experiments by setting up
    S3 backends for storage, logging important job details, and saving configuration and
    environment details both locally and on S3.
    This function is crucial for ensuring that all aspects of the computational environment are captured and can be
    replicated for future runs or analysis.

    Parameters:
        config (Config): A configuration object containing all the settings necessary
                         for the job, including paths and credentials.
        args (argparse.Namespace): An argparse namespace containing the command line
                                   arguments passed to the script. This includes configurations
                                   and any overrides specified at runtime.

    Actions:
        - Sets up S3 backend for storing user data and other outputs.
        - Logs job paths and critical information regarding job execution.
        - Saves the job configuration locally only for the main node in a distributed setting.
        - Captures and logs command-line execution details.
        - Optionally reads git commit and branch information if available and logs them.
        - Saves both job environment information and launch details locally and syncs these to S3.
        - Supports conditional integration with Weights & Biases (wandb) for experiment tracking.

    Notes:
        - The function is designed to run within a distributed environment where certain actions
          (like saving configurations) are restricted to the main node (rank 0).
        - It uses the 'easy_io' module for interacting with S3, ensuring files are written and
          read correctly from the object store.
        - It leverages OmegaConf for saving YAML configurations
        - git information is read from 'git_commit.txt' and 'git_branch.txt' files if they exist.
        - snapshot codebase is saved as 'codebase.zip' if it exists in the current directory.

    Raises:
        FileNotFoundError: If specific files like 'git_commit.txt' or 'codebase.zip' are expected
                           but not found.
        IOError: If there are issues in file handling operations, particularly with file
                 reading/writing.
    """

    run_timestamp = f"{time.strftime('%Y-%m-%d_%H-%M-%S')}"
    time_tensor = torch.ByteTensor(bytearray(run_timestamp, "utf-8")).cuda()
    distributed.broadcast(time_tensor, 0)
    run_timestamp = time_tensor.cpu().numpy().tobytes().decode("utf-8")

    global S3_READY
    if os.path.exists(config.checkpoint.save_to_object_store.credentials) or CRED_ENVS.APP_ENV in [
        "prod",
        "dev",
        "stg",
    ]:
        easy_io.set_s3_backend(
            backend_args={
                "backend": "s3",
                "path_mapping": {
                    "s3://timestamps_rundir/": f"s3://{config.checkpoint.save_to_object_store.bucket}/{config.job.path}/job_runs/{run_timestamp}/",
                    "s3://rundir/": f"s3://{config.checkpoint.save_to_object_store.bucket}/{config.job.path}/",
                },
                "s3_credential_path": config.checkpoint.save_to_object_store.credentials,
            }
        )
        S3_READY = True
    else:
        log.warning("S3 credentials not found. Skipping easy_io S3 setup.")

    log.warning(f"Job path: {config.job.path}")
    job_info = get_cluster_env()
    # save cfg to local
    if distributed.get_rank() == 0:
        job_local_path = config.job.path_local
        log.critical(f"Job local path: {job_local_path}")
        os.makedirs(config.job.path_local, exist_ok=True)
        launch_info = {
            "cmd": " ".join(sys.argv),
            "args_cfg_path": args.config,
            "args_override": args.opts,
        }

        job_info["job_local_path"] = str(job_local_path)
        job_info["s3"] = f"s3://{config.checkpoint.save_to_object_store.bucket}/{config.job.path}/"
        # optional read git_commit.txt and save git commit id
        if os.path.exists("git_commit.txt"):
            with open("git_commit.txt", "r") as f:
                job_info["commit_id"] = f.read().strip()
                log.critical(f"Commit id: {job_info['commit_id']}")
        if os.path.exists("git_branch.txt"):
            with open("git_branch.txt", "r") as f:
                job_info["git_branch"] = f.read().strip()
                log.critical(f"git branch: {job_info['git_branch']}")
        if os.path.exists("git_diff.txt"):
            with open("git_diff.txt", "r") as f:
                job_info["git_diff"] = f.read().strip()
                log.critical(f"git diff: {job_info['git_diff']}")

        with open(f"{job_local_path}/job_env.yaml", "w") as f:
            OmegaConf.save(job_info, f)
        with open(f"{job_local_path}/launch_info.yaml", "w") as f:
            OmegaConf.save(launch_info, f)
        set_wandb_job_info(job_info)

        # by default, we upload run in ngc and slurm
        if config.upload_reproducible_setup:
            # sync to s3
            if S3_READY:
                log.critical(
                    f"Uploading reproducible setup to s3://{config.checkpoint.save_to_object_store.bucket}/{config.job.path}/job_runs/{run_timestamp}/"
                )

                config_pkl_save_fp = f"{config.job.path_local}/config.pkl"
                easy_io.copyfile_from_local(
                    config_pkl_save_fp, f"s3://timestamps_rundir/{config_pkl_save_fp.split('/')[-1]}"
                )
                config_yaml_save_fp = config_pkl_save_fp.replace(".pkl", ".yaml")
                easy_io.copyfile_from_local(
                    config_yaml_save_fp, f"s3://timestamps_rundir/{config_yaml_save_fp.split('/')[-1]}"
                )
                easy_io.copyfile_from_local(f"{job_local_path}/job_env.yaml", "s3://timestamps_rundir/job_env.yaml")
                easy_io.copyfile_from_local(
                    f"{job_local_path}/launch_info.yaml",
                    "s3://timestamps_rundir/launch_info.yaml",
                )
                if os.path.exists("codebase.zip"):
                    easy_io.copyfile_from_local("codebase.zip", "s3://timestamps_rundir/codebase.zip")
                if os.path.exists("code.tar.gz"):
                    easy_io.copyfile_from_local("code.tar.gz", "s3://timestamps_rundir/code.tar.gz")
                if os.path.exists("git_diff.txt"):
                    easy_io.copyfile_from_local("git_diff.txt", "s3://timestamps_rundir/git_diff.txt")
                if easy_io.exists("s3://rundir/job_history.yaml"):
                    job_history = easy_io.load("s3://rundir/job_history.yaml")
                else:
                    job_history = {}
                job_history[len(job_history)] = {
                    "timestamp": run_timestamp,
                    "reproduce_dir": f"s3://{config.checkpoint.save_to_object_store.bucket}/{config.job.path}/job_runs/{run_timestamp}/",
                    **launch_info,
                }
                print(job_history)
                easy_io.dump(job_history, "s3://rundir/job_history.yaml")
            else:
                log.warning("S3 credentials not found. Skipping upload of reproducible setup.")

    # save per rank cluster information to s3
    if config.upload_reproducible_setup:
        if S3_READY:
            easy_io.dump(job_info, f"s3://timestamps_rundir/cluster_env/RANK_{distributed.get_rank():06d}.yaml")