scripts/train.py

# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
import builtins
import os
import sys

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# Expose third_party/soma/tools/shape_convert.py as `tools.shape_convert`
# (used by gem_private.gem.callbacks.converter.convert_soma_to_smpl).
_repo_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
_soma_tools = os.path.join(_repo_root, "third_party", "soma")
if os.path.isdir(_soma_tools) and _soma_tools not in sys.path:
    sys.path.insert(0, _soma_tools)


from datetime import datetime

import hydra
import pytorch_lightning as pl
import torch
import torch.distributed as dist
import wandb
import yaml
from hydra.core.config_search_path import ConfigSearchPath
from hydra.core.plugins import Plugins
from hydra.plugins.search_path_plugin import SearchPathPlugin
from omegaconf import DictConfig, ListConfig, OmegaConf
from pytorch_lightning.callbacks.checkpoint import Checkpoint
from pytorch_lightning.loggers.tensorboard import TensorBoardLogger

from gem.callbacks.autoresume_callback import AutoResume, AutoResumeCallback
from gem.utils.net_utils import get_resume_ckpt_path, load_pretrained_model
from gem.utils.pylogger import Log
from gem.utils.tools import (
    find_last_version,
    get_checkpoint_path,
    rsync_file_from_remote,
)
from gem.utils.vis.rich_logger import print_cfg

OmegaConf.register_new_resolver("eval", builtins.eval)


def _get_rank():
    # SLURM_PROCID can be set even if SLURM is not managing the multiprocessing,
    # therefore LOCAL_RANK needs to be checked first
    rank_keys = ("RANK", "LOCAL_RANK", "SLURM_PROCID", "JSM_NAMESPACE_RANK")
    for key in rank_keys:
        rank = os.environ.get(key)
        if rank is not None:
            return int(rank)
    # None to differentiate whether an environment variable was set at all
    return 0


global_rank = _get_rank()


def wandb_run_exists():
    return isinstance(wandb.run, wandb.sdk.wandb_run.Run)


def get_callbacks(cfg: DictConfig) -> list:
    """Parse and instantiate all the callbacks in the config.

    Supports both flat and nested callback configs. Only nodes containing
    a `_target_` are instantiated.
    """
    if not hasattr(cfg, "callbacks") or cfg.callbacks is None:
        return None

    def _collect_callback_nodes(node):
        collected = []
        if node is None:
            return collected
        # Dict-like node
        if isinstance(node, DictConfig | dict):
            # direct instantiable config
            if "_target_" in node:
                collected.append(node)
            else:
                for child in node.values():
                    collected.extend(_collect_callback_nodes(child))
        # List-like node
        elif isinstance(node, ListConfig | list | tuple):
            for child in node:
                collected.extend(_collect_callback_nodes(child))
        # primitives are ignored
        return collected

    enable_checkpointing = cfg.pl_trainer.get("enable_checkpointing", True)
    callbacks = []
    for cb_conf in _collect_callback_nodes(cfg.callbacks):
        cb = hydra.utils.instantiate(cb_conf, _recursive_=False)
        if not enable_checkpointing and isinstance(cb, Checkpoint):
            continue
        callbacks.append(cb)
    return callbacks


def train(cfg: DictConfig) -> None:
    """Train/Test"""
    Log.info(f"[Exp Name]: {cfg.exp_name}")
    # use total batch size
    if cfg.task == "fit":
        Log.info(
            f"[GPU x Batch] = {cfg.pl_trainer.devices} x {cfg.data.loader_opts.train.batch_size}"
        )
    num_nodes = cfg.pl_trainer.get("num_nodes", 1)
    cfg.num_test_data *= cfg.pl_trainer.devices * num_nodes
    if (
        "imgfeat_motionx" in cfg.test_datasets
        and "max_num_motions" in cfg.test_datasets.imgfeat_motionx
    ):
        cfg.test_datasets.imgfeat_motionx.max_num_motions *= cfg.pl_trainer.devices * num_nodes
    pl.seed_everything(cfg.seed)
    torch.cuda.set_device(global_rank % 8)  # for tinycudann default memory
    wandb_run = None
    version = None

    if cfg.get("timing", False):
        os.environ["DEBUG_TIMING"] = "TRUE"

    if AutoResume is not None:
        details = AutoResume.get_resume_details()
        if details:
            cfg.resume_mode = "last"
            if "wandb_id" in details:
                wandb_run = details["wandb_id"]
                version = int(details["version"])
            print(
                f"[Auto Resume] Loading. checkpoint: {details['checkpoint']} wandb_id: {details.get('wandb_id', None)}"
            )

    if cfg.task == "test" and not cfg.get("no_checkpoint", False):
        test_cp = cfg.get("test_checkpoint", "last")
        remote_run_dir = cfg.output_dir.replace("outputs", cfg.remote_results_path)
        version = find_last_version(remote_run_dir, cp=test_cp)
        checkpoint_dir = f"{remote_run_dir}/version_{version}/checkpoints"
        remote_ckpt_path = get_checkpoint_path(checkpoint_dir, test_cp)
        if cfg.get("rsync_ckpt", False):
            cfg.ckpt_path = remote_ckpt_path.replace(cfg.remote_results_path, "outputs")
            if not os.path.exists(cfg.ckpt_path):
                print(f"rsyncing from remote: {remote_ckpt_path}")
                print(f"output_dir: {cfg.output_dir}")
                rsync_file_from_remote(
                    cfg.ckpt_path,
                    remote_run_dir,
                    cfg.output_dir,
                    hostname=cfg.get("remote_hostname", None),
                )
        else:
            cfg.ckpt_path = remote_ckpt_path
        print("ckpt path:", cfg.ckpt_path)
        cfg.output_dir = f"{cfg.output_dir}/version_{version}"
        cfg.logger.name = f"{cfg.exp_name}_v{version}_{datetime.now().strftime('%Y%m%d%H%M%S')}"
    else:
        run_root_dir = cfg.output_dir
        if version is None and cfg.resume_mode == "last":
            version = find_last_version(run_root_dir, cp="last")

    # preparation
    datamodule: pl.LightningDataModule = hydra.utils.instantiate(cfg.data, _recursive_=False)
    model: pl.LightningModule = hydra.utils.instantiate(cfg.model, _recursive_=False)

    if (
        cfg.get("pretrain_ckpt", None) is not None
        and cfg.ckpt_path is None
        and cfg.resume_mode is None
    ):
        cfg.ckpt_path = cfg.pretrain_ckpt

    wandb_cfg = OmegaConf.to_container(cfg, resolve=True)
    if cfg.ckpt_path is not None:
        if cfg.get("rsync_ckpt", False) and not os.path.exists(cfg.ckpt_path):
            print(f"rsyncing from remote: {cfg.ckpt_path}")
            cfg.ckpt_path = cfg.ckpt_path.replace(cfg.remote_results_path, "outputs")
            local_dir = cfg.ckpt_path.split("/version_")[0]
            os.makedirs(local_dir, exist_ok=True)
            rsync_file_from_remote(
                cfg.ckpt_path,
                cfg.remote_results_path,
                "outputs",
                hostname=cfg.get("remote_hostname", None),
            )

        ckpt = load_pretrained_model(model, cfg.ckpt_path)
        print(f"Loaded pretrained model from {cfg.ckpt_path}")
        if ckpt is not None:
            wandb_cfg["pretrained_ckpt_info"] = {
                "global_step": ckpt["global_step"],
                "epoch": ckpt["epoch"],
            }
            print("pretrained ckpt info:", wandb_cfg["pretrained_ckpt_info"])

    # PL callbacks and logger
    if cfg.task == "fit":
        if global_rank == 0:
            tb_logger = TensorBoardLogger(run_root_dir, version=version, name="")
            version = tb_logger.version
            os.makedirs(tb_logger.log_dir, exist_ok=True)
            cfg.output_dir = tb_logger.log_dir

            slurm_job_id = int(os.environ.get("SLURM_JOB_ID", "-1"))
            run_name = (
                f"{cfg.exp_name}_v{version}_{slurm_job_id}"
                if slurm_job_id > 0
                else f"{cfg.exp_name}_v{version}"
            )
            cfg.logger.name = run_name
            # cfg.logger.version = version  # shouldn't set version for Wandb

            if cfg.resume_mode == "last" and os.path.exists(f"{tb_logger.log_dir}/meta.yaml"):
                meta = yaml.safe_load(open(f"{tb_logger.log_dir}/meta.yaml"))
                if wandb_run is None:
                    wandb_run = meta["wandb_run"]
            if wandb_run is None:
                wandb_run = (
                    f"{cfg.exp_name.replace('/', '_')}_{datetime.now().strftime('%Y%m%d%H%M%S')}"
                )
            cfg.logger.id = wandb_run

        if cfg.pl_trainer.devices > 1 and "RANK" in os.environ:
            dist.init_process_group("nccl")
            dist.barrier()

        if global_rank != 0:
            if version is None:
                version = find_last_version(run_root_dir, cp=None)
            cfg.output_dir = f"{run_root_dir}/version_{version}"

    callbacks = get_callbacks(cfg)
    has_ckpt_cb = any([isinstance(cb, Checkpoint) for cb in callbacks])
    if not has_ckpt_cb and cfg.pl_trainer.get("enable_checkpointing", True):
        Log.warning("No checkpoint-callback found. Disabling PL auto checkpointing.")
        cfg.pl_trainer = {**cfg.pl_trainer, "enable_checkpointing": False}
    if AutoResume is not None:
        callbacks.append(AutoResumeCallback(version))

    cfg.logger.config = wandb_cfg
    if cfg.use_wandb:
        logger = hydra.utils.instantiate(cfg.logger, _recursive_=False, _convert_="partial")
    else:
        logger = tb_logger

    if cfg.task == "fit" and global_rank == 0:
        # wandb.config.update({"cfg": OmegaConf.to_container(cfg)}, allow_val_change=True)
        assert cfg.logger.id is not None
        meta = {"wandb_run": cfg.logger.id}
        yaml.safe_dump(meta, open(f"{tb_logger.log_dir}/meta.yaml", "w"))
        print("saved meta:", meta)

    # PL-Trainer
    if cfg.task == "test":
        Log.info("Test mode forces full-precision.")
        cfg.pl_trainer = {**cfg.pl_trainer, "precision": 32}
    trainer = pl.Trainer(
        accelerator="gpu",
        logger=logger if logger is not None else False,
        callbacks=callbacks,
        **cfg.pl_trainer,
    )

    print("=" * 20)
    print("version:", version)

    if cfg.task == "fit":
        resume_path = None
        if cfg.resume_mode is not None:
            save_dir = cfg.output_dir + "/checkpoints"
            resume_path = get_resume_ckpt_path(cfg.resume_mode, ckpt_dir=save_dir)
        Log.info("Start Fitting...")
        trainer.fit(
            model,
            datamodule.train_dataloader(),
            datamodule.val_dataloader(),
            ckpt_path=resume_path,
        )
    elif cfg.task == "test":
        Log.info("Start Testing...")
        trainer.test(model, datamodule.test_dataloader())
    else:
        raise ValueError(f"Unknown task: {cfg.task}")

    Log.info("End of script.")


class _GemPrivateSearchPath(SearchPathPlugin):
    def manipulate_search_path(self, search_path: ConfigSearchPath) -> None:
        search_path.append("file", "gem_private/configs")


Plugins.instance().register(_GemPrivateSearchPath)


@hydra.main(version_base="1.3", config_path="../configs", config_name="train")
def main(cfg) -> None:
    print_cfg(cfg, use_rich=True)
    train(cfg)


if __name__ == "__main__":
    main()