# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.

"""Megatron global variables."""
import os
import sys
from megatron.training import global_vars
from megatron.training.global_vars import _ensure_var_is_not_initialized

from megatron.addons.function_wrapper import FUNCTION_WRAPPER

# new_feature: enable autoresume
@FUNCTION_WRAPPER.register_wrapper(func_name='megatron.training.global_vars._set_adlr_autoresume')
def _set_adlr_autoresume(args):
    """Initialize ADLR autoresume."""
    #global global_vars._GLOBAL_ADLR_AUTORESUME
    _ensure_var_is_not_initialized(global_vars._GLOBAL_ADLR_AUTORESUME, 'adlr autoresume')

    if args.adlr_autoresume:
        if args.rank == 0:
            print('enabling autoresume ...', flush=True)
        sys.path.append(os.environ.get('SUBMIT_SCRIPTS', '.'))
        try:
            from cntrainkit.auto_resume import AutoResume
        except BaseException:
            print('ADLR autoresume is not available, exiting ...')
            sys.exit()

        global_vars._GLOBAL_ADLR_AUTORESUME = AutoResume
