Spaces:
Running on Zero
Running on Zero
| # 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") | |