#!/usr/bin/env python3
# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
import argparse
import importlib
import os
import sys

import megengine as mge
import megengine.distributed as dist
from basecore.config import ConfigDict
from loguru import logger

from basecls.models import build_model, load_model, sync_model
from basecls.utils import registers, set_nccl_env, set_num_threads, setup_logger


def default_parser() -> argparse.ArgumentParser:
    """Build args parser for training script.

    Returns:
        The args parser.
    """
    parser = argparse.ArgumentParser()
    parser.add_argument("-f", "--file", type=str, help="training process description file")
    parser.add_argument(
        "--resume", action="store_true", help="resume training from saved checkpoint or not"
    )
    parser.add_argument(
        "opts",
        default=None,
        help="Modify config options using the command-line",
        nargs=argparse.REMAINDER,
    )
    return parser


@logger.catch
def worker(args: argparse.Namespace):
    """Worker function for training script.

    Args:
        args: args for training script.
    """
    logger.info(f"Init process group for gpu{dist.get_rank()} done")

    sys.path.append(os.path.dirname(args.file))
    module_name = os.path.splitext(os.path.basename(args.file))[0]
    current_network = importlib.import_module(module_name)
    cfg = current_network.Cfg()
    cfg.merge(args.opts)
    cfg.resume = args.resume
    if cfg.output_dir is None:
        cfg.output_dir = f"./logs_{module_name}"
    cfg.output_dir = os.path.abspath(cfg.output_dir)

    cfg.set_mode("freeze")

    if 
dist.get_rank()