"""
    ########################################################################################################################################################
    # This algorithm is part of a code library to accelerate Sapiens Technology® Artificial Intelligence models, and its disclosure, distribution,         #
    # or reverse engineering without the company's prior consent will be subject to legal proceedings and actions pursued by our legal department.         #
    ########################################################################################################################################################
"""
import argparse
import os
import subprocess
import sys
from ast import literal_eval
from shutil import which
from typing import Any, Dict, List, Tuple
import torch
from ..commands.config.config_args import SageMakerConfig
from ..utils import (DynamoBackend, PrecisionType, is_fp8_available, is_ipex_available, is_mlu_available, is_musa_available, is_npu_available, is_torch_xla_available, is_xpu_available)
from ..utils.constants import DEEPSPEED_MULTINODE_LAUNCHERS
from ..utils.other import is_port_in_use, merge_dicts
from .dataclasses import DistributedType, SageMakerDistributedType
def _filter_args(args, parser, default_args=[]):
    new_args, _ = parser.parse_known_args(default_args)
    for key, value in vars(args).items():
        if key in vars(new_args).keys(): setattr(new_args, key, value)
    return new_args
def _get_mpirun_args():
    mpi_apps = [x for x in ["mpirun", "mpiexec"] if which(x)]
    if len(mpi_apps) == 0: raise OSError("mpirun or mpiexec were not found. Ensure that Intel MPI, Open MPI, or MVAPICH are installed.")
    mpi_app = mpi_apps[0]
    mpirun_version = subprocess.check_output([mpi_app, "--version"])
    if b"Open MPI" in mpirun_version: return mpi_app, "--hostfile", "-n", "--npernode", "--bind-to"
    else: return mpi_app, "-f", "-n", "-ppn", ""
def setup_fp8_env(args: argparse.Namespace, current_env: Dict[str, str]):
    prefix = "SAPIENS_ACCELERATOR_"
    for arg in vars(args):
        if arg.startswith("fp8_"):
            value = getattr(args, arg)
            if value is not None: current_env[f"{prefix}{arg.upper()}"] = str(getattr(args, arg))
    return current_env
def prepare_simple_launcher_cmd_env(args: argparse.Namespace) -> Tuple[List[str], Dict[str, str]]:
    cmd = []
    if args.no_python and args.module: raise ValueError("--module and --no_python cannot be used together")
    if args.mpirun_hostfile is not None:
        mpi_app_name, hostfile_arg, num_proc_arg, proc_per_node_arg, bind_to_arg = _get_mpirun_args()
        mpirun_ccl = getattr(args, "mpirun_ccl", None)
        bind_to = getattr(args, "bind-to", "socket")
        num_machines = args.num_machines
        num_processes = getattr(args, "num_processes", None)
        nproc_per_node = str(num_processes // num_machines) if num_processes and num_machines else "1"
        cmd += [mpi_app_name, hostfile_arg, args.mpirun_hostfile, proc_per_node_arg, nproc_per_node]
        if num_processes: cmd += [num_proc_arg, str(num_processes)]
        if bind_to_arg: cmd += [bind_to_arg, bind_to]
    if not args.no_python:
        cmd.append(sys.executable)
        if args.module: cmd.append("-m")
    cmd.append(args.training_script)
    cmd.extend(args.training_script_args)
    current_env = os.environ.copy()
    current_env["SAPIENS_ACCELERATOR_USE_CPU"] = str(args.cpu or args.use_cpu)
    if args.debug: current_env["SAPIENS_ACCELERATOR_DEBUG_MODE"] = "true"
    if args.gpu_ids != "all" and args.gpu_ids is not None:
        if is_xpu_available(): current_env["ZE_AFFINITY_MASK"] = args.gpu_ids
        elif is_mlu_available(): current_env["MLU_VISIBLE_DEVICES"] = args.gpu_ids
        elif is_musa_available(): current_env["MUSA_VISIBLE_DEVICES"] = args.gpu_ids
        elif is_npu_available(): current_env["ASCEND_RT_VISIBLE_DEVICES"] = args.gpu_ids
        else: current_env["CUDA_VISIBLE_DEVICES"] = args.gpu_ids
    if args.num_machines > 1:
        current_env["MASTER_ADDR"] = args.main_process_ip
        current_env["MASTER_PORT"] = str(args.main_process_port)
        if args.mpirun_hostfile is not None: current_env["CCL_WORKER_COUNT"] = str(mpirun_ccl)
    elif args.num_processes > 1:
        current_env["MASTER_ADDR"] = args.main_process_ip if args.main_process_ip is not None else "127.0.0.1"
        current_env["MASTER_PORT"] = str(args.main_process_port) if args.main_process_port is not None else "29500"
    try: mixed_precision = PrecisionType(args.mixed_precision.lower())
    except ValueError: raise ValueError(f"Unknown mixed_precision mode: {args.mixed_precision.lower()}. Choose between {PrecisionType.list()}.")
    current_env["SAPIENS_ACCELERATOR_MIXED_PRECISION"] = str(mixed_precision)
    if args.mixed_precision.lower() == "fp8":
        if not is_fp8_available(): raise RuntimeError("FP8 is not available on this machine. Please ensure that either Transformer Engine or MSAMP is installed.")
        current_env = setup_fp8_env(args, current_env)
    try: dynamo_backend = DynamoBackend(args.dynamo_backend.upper())
    except ValueError: raise ValueError(f"Unknown dynamo backend: {args.dynamo_backend.upper()}. Choose between {DynamoBackend.list()}.")
    current_env["SAPIENS_ACCELERATOR_DYNAMO_BACKEND"] = dynamo_backend.value
    current_env["SAPIENS_ACCELERATOR_DYNAMO_MODE"] = args.dynamo_mode
    current_env["SAPIENS_ACCELERATOR_DYNAMO_USE_FULLGRAPH"] = str(args.dynamo_use_fullgraph)
    current_env["SAPIENS_ACCELERATOR_DYNAMO_USE_DYNAMIC"] = str(args.dynamo_use_dynamic)
    current_env["OMP_NUM_THREADS"] = str(args.num_cpu_threads_per_process)
    if is_ipex_available():
        current_env["SAPIENS_ACCELERATOR_USE_IPEX"] = str(args.ipex).lower()
        current_env["SAPIENS_ACCELERATOR_USE_XPU"] = str(args.use_xpu).lower()
    if args.enable_cpu_affinity: current_env["SAPIENS_ACCELERATOR_CPU_AFFINITY"] = "1"
    return cmd, current_env
def prepare_multi_gpu_env(args: argparse.Namespace) -> Dict[str, str]:
    num_processes = args.num_processes
    num_machines = args.num_machines
    main_process_ip = args.main_process_ip
    main_process_port = args.main_process_port
    if num_machines > 1:
        args.nproc_per_node = str(num_processes // num_machines)
        args.nnodes = str(num_machines)
        args.node_rank = int(args.machine_rank)
        if getattr(args, "same_network", False):
            args.master_addr = str(main_process_ip)
            args.master_port = str(main_process_port)
        else: args.rdzv_endpoint = f"{main_process_ip}:{main_process_port}"
    else:
        args.nproc_per_node = str(num_processes)
        if main_process_port is not None: args.master_port = str(main_process_port)
    if main_process_port is None: main_process_port = 29500
    need_port_check = num_machines <= 1 or int(args.machine_rank) == 0
    if need_port_check and is_port_in_use(main_process_port): raise ConnectionError(f"Tried to launch distributed communication on port `{main_process_port}`, but another process is utilizing it. Please specify a different port (such as using the `--main_process_port` flag or specifying a different `main_process_port` in your config file) and rerun your script. To automatically use the next open port (on a single node), you can set this to `0`.")
    if args.module and args.no_python: raise ValueError("--module and --no_python cannot be used together")
    elif args.module: args.module = True
    elif args.no_python: args.no_python = True
    current_env = os.environ.copy()
    if args.debug: current_env["SAPIENS_ACCELERATOR_DEBUG_MODE"] = "true"
    gpu_ids = getattr(args, "gpu_ids", "all")
    if gpu_ids != "all" and args.gpu_ids is not None:
        if is_xpu_available(): current_env["ZE_AFFINITY_MASK"] = gpu_ids
        elif is_mlu_available(): current_env["MLU_VISIBLE_DEVICES"] = gpu_ids
        elif is_musa_available(): current_env["MUSA_VISIBLE_DEVICES"] = gpu_ids
        elif is_npu_available(): current_env["ASCEND_RT_VISIBLE_DEVICES"] = gpu_ids
        else: current_env["CUDA_VISIBLE_DEVICES"] = gpu_ids
    mixed_precision = args.mixed_precision.lower()
    try: mixed_precision = PrecisionType(mixed_precision)
    except ValueError: raise ValueError(f"Unknown mixed_precision mode: {mixed_precision}. Choose between {PrecisionType.list()}.")
    current_env["SAPIENS_ACCELERATOR_MIXED_PRECISION"] = str(mixed_precision)
    if args.mixed_precision.lower() == "fp8":
        if not is_fp8_available(): raise RuntimeError("FP8 is not available on this machine. Please ensure that either Transformer Engine or MSAMP is installed.")
        current_env = setup_fp8_env(args, current_env)
    try: dynamo_backend = DynamoBackend(args.dynamo_backend.upper())
    except ValueError: raise ValueError(f"Unknown dynamo backend: {args.dynamo_backend.upper()}. Choose between {DynamoBackend.list()}.")
    current_env["SAPIENS_ACCELERATOR_DYNAMO_BACKEND"] = dynamo_backend.value
    current_env["SAPIENS_ACCELERATOR_DYNAMO_MODE"] = args.dynamo_mode
    current_env["SAPIENS_ACCELERATOR_DYNAMO_USE_FULLGRAPH"] = str(args.dynamo_use_fullgraph)
    current_env["SAPIENS_ACCELERATOR_DYNAMO_USE_DYNAMIC"] = str(args.dynamo_use_dynamic)
    if args.use_fsdp:
        current_env["SAPIENS_ACCELERATOR_USE_FSDP"] = "true"
        if args.fsdp_cpu_ram_efficient_loading and not args.fsdp_sync_module_states: raise ValueError("When using `--fsdp_cpu_ram_efficient_loading` set `--fsdp_sync_module_states` to `True`")
        current_env["FSDP_SHARDING_STRATEGY"] = str(args.fsdp_sharding_strategy)
        current_env["FSDP_OFFLOAD_PARAMS"] = str(args.fsdp_offload_params).lower()
        current_env["FSDP_MIN_NUM_PARAMS"] = str(args.fsdp_min_num_params)
        if args.fsdp_auto_wrap_policy is not None: current_env["FSDP_AUTO_WRAP_POLICY"] = str(args.fsdp_auto_wrap_policy)
        if args.fsdp_transformer_layer_cls_to_wrap is not None: current_env["FSDP_TRANSFORMER_CLS_TO_WRAP"] = str(args.fsdp_transformer_layer_cls_to_wrap)
        if args.fsdp_backward_prefetch is not None: current_env["FSDP_BACKWARD_PREFETCH"] = str(args.fsdp_backward_prefetch)
        if args.fsdp_state_dict_type is not None: current_env["FSDP_STATE_DICT_TYPE"] = str(args.fsdp_state_dict_type)
        current_env["FSDP_FORWARD_PREFETCH"] = str(args.fsdp_forward_prefetch).lower()
        current_env["FSDP_USE_ORIG_PARAMS"] = str(args.fsdp_use_orig_params).lower()
        current_env["FSDP_CPU_RAM_EFFICIENT_LOADING"] = str(args.fsdp_cpu_ram_efficient_loading).lower()
        current_env["FSDP_SYNC_MODULE_STATES"] = str(args.fsdp_sync_module_states).lower()
        current_env["FSDP_ACTIVATION_CHECKPOINTING"] = str(args.fsdp_activation_checkpointing).lower()
    if args.use_megatron_lm:
        prefix = "MEGATRON_LM_"
        current_env["SAPIENS_ACCELERATOR_USE_MEGATRON_LM"] = "true"
        current_env[prefix + "TP_DEGREE"] = str(args.megatron_lm_tp_degree)
        current_env[prefix + "PP_DEGREE"] = str(args.megatron_lm_pp_degree)
        current_env[prefix + "GRADIENT_CLIPPING"] = str(args.megatron_lm_gradient_clipping)
        if args.megatron_lm_num_micro_batches is not None: current_env[prefix + "NUM_MICRO_BATCHES"] = str(args.megatron_lm_num_micro_batches)
        if args.megatron_lm_sequence_parallelism is not None: current_env[prefix + "SEQUENCE_PARALLELISM"] = str(args.megatron_lm_sequence_parallelism)
        if args.megatron_lm_recompute_activations is not None: current_env[prefix + "RECOMPUTE_ACTIVATIONS"] = str(args.megatron_lm_recompute_activations)
        if args.megatron_lm_use_distributed_optimizer is not None: current_env[prefix + "USE_DISTRIBUTED_OPTIMIZER"] = str(args.megatron_lm_use_distributed_optimizer)
    current_env["OMP_NUM_THREADS"] = str(args.num_cpu_threads_per_process)
    if args.enable_cpu_affinity: current_env["SAPIENS_ACCELERATOR_CPU_AFFINITY"] = "1"
    return current_env
def prepare_deepspeed_cmd_env(args: argparse.Namespace) -> Tuple[List[str], Dict[str, str]]:
    num_processes = args.num_processes
    num_machines = args.num_machines
    main_process_ip = args.main_process_ip
    main_process_port = args.main_process_port
    cmd = None
    if args.deepspeed_multinode_launcher is None: args.deepspeed_multinode_launcher = DEEPSPEED_MULTINODE_LAUNCHERS[0]
    if num_machines > 1 and args.deepspeed_multinode_launcher != DEEPSPEED_MULTINODE_LAUNCHERS[1]:
        cmd = ["deepspeed", "--no_local_rank"]
        cmd.extend(["--hostfile", str(args.deepspeed_hostfile), "--launcher", str(args.deepspeed_multinode_launcher)])
        if args.deepspeed_exclusion_filter is not None: cmd.extend(["--exclude", str(args.deepspeed_exclusion_filter)])
        elif args.deepspeed_inclusion_filter is not None: cmd.extend(["--include", str(args.deepspeed_inclusion_filter)])
        else: cmd.extend(["--num_gpus", str(args.num_processes // args.num_machines)])
        if main_process_ip: cmd.extend(["--master_addr", str(main_process_ip)])
        cmd.extend(["--master_port", str(main_process_port)])
        if args.module and args.no_python: raise ValueError("--module and --no_python cannot be used together")
        elif args.module: cmd.append("--module")
        elif args.no_python: cmd.append("--no_python")
        cmd.append(args.training_script)
        cmd.extend(args.training_script_args)
    elif num_machines > 1 and args.deepspeed_multinode_launcher == DEEPSPEED_MULTINODE_LAUNCHERS[1]:
        args.nproc_per_node = str(num_processes // num_machines)
        args.nnodes = str(num_machines)
        args.node_rank = int(args.machine_rank)
        if getattr(args, "same_network", False):
            args.master_addr = str(main_process_ip)
            args.master_port = str(main_process_port)
        else: args.rdzv_endpoint = f"{main_process_ip}:{main_process_port}"
    else:
        args.nproc_per_node = str(num_processes)
        if main_process_port is not None: args.master_port = str(main_process_port)
    if main_process_port is None: main_process_port = 29500
    need_port_check = num_machines <= 1 or int(args.machine_rank) == 0
    if need_port_check and is_port_in_use(main_process_port): raise ConnectionError(f"Tried to launch distributed communication on port `{main_process_port}`, but another process is utilizing it. Please specify a different port (such as using the `--main_process_port` flag or specifying a different `main_process_port` in your config file) and rerun your script. To automatically use the next open port (on a single node), you can set this to `0`.")
    if args.module and args.no_python: raise ValueError("--module and --no_python cannot be used together")
    elif args.module: args.module = True
    elif args.no_python: args.no_python = True
    current_env = os.environ.copy()
    if args.debug: current_env["SAPIENS_ACCELERATOR_DEBUG_MODE"] = "true"
    gpu_ids = getattr(args, "gpu_ids", "all")
    if gpu_ids != "all" and args.gpu_ids is not None:
        if is_xpu_available(): current_env["ZE_AFFINITY_MASK"] = gpu_ids
        elif is_mlu_available(): current_env["MLU_VISIBLE_DEVICES"] = gpu_ids
        elif is_musa_available(): current_env["MUSA_VISIBLE_DEVICES"] = gpu_ids
        elif is_npu_available(): current_env["ASCEND_RT_VISIBLE_DEVICES"] = gpu_ids
        else: current_env["CUDA_VISIBLE_DEVICES"] = gpu_ids
    try: mixed_precision = PrecisionType(args.mixed_precision.lower())
    except ValueError: raise ValueError(f"Unknown mixed_precision mode: {args.mixed_precision.lower()}. Choose between {PrecisionType.list()}.")
    current_env["PYTHONPATH"] = env_var_path_add("PYTHONPATH", os.path.abspath("."))
    current_env["SAPIENS_ACCELERATOR_MIXED_PRECISION"] = str(mixed_precision)
    if args.mixed_precision.lower() == "fp8":
        if not is_fp8_available(): raise RuntimeError("FP8 is not available on this machine. Please ensure that either Transformer Engine or MSAMP is installed.")
        current_env = setup_fp8_env(args, current_env)
    current_env["SAPIENS_ACCELERATOR_CONFIG_DS_FIELDS"] = str(args.deepspeed_fields_from_sapiens_accelerator_config).lower()
    current_env["SAPIENS_ACCELERATOR_USE_DEEPSPEED"] = "true"
    if args.zero_stage is not None: current_env["SAPIENS_ACCELERATOR_DEEPSPEED_ZERO_STAGE"] = str(args.zero_stage)
    if args.gradient_accumulation_steps is not None: current_env["SAPIENS_ACCELERATOR_GRADIENT_ACCUMULATION_STEPS"] = str(args.gradient_accumulation_steps)
    if args.gradient_clipping is not None: current_env["SAPIENS_ACCELERATOR_GRADIENT_CLIPPING"] = str(args.gradient_clipping).lower()
    if args.offload_optimizer_device is not None: current_env["SAPIENS_ACCELERATOR_DEEPSPEED_OFFLOAD_OPTIMIZER_DEVICE"] = str(args.offload_optimizer_device).lower()
    if args.offload_param_device is not None: current_env["SAPIENS_ACCELERATOR_DEEPSPEED_OFFLOAD_PARAM_DEVICE"] = str(args.offload_param_device).lower()
    if args.zero3_init_flag is not None: current_env["SAPIENS_ACCELERATOR_DEEPSPEED_ZERO3_INIT"] = str(args.zero3_init_flag).lower()
    if args.zero3_save_16bit_model is not None: current_env["SAPIENS_ACCELERATOR_DEEPSPEED_ZERO3_SAVE_16BIT_MODEL"] = str(args.zero3_save_16bit_model).lower()
    if args.deepspeed_config_file is not None: current_env["SAPIENS_ACCELERATOR_DEEPSPEED_CONFIG_FILE"] = str(args.deepspeed_config_file)
    if args.enable_cpu_affinity: current_env["SAPIENS_ACCELERATOR_CPU_AFFINITY"] = "1"
    if args.deepspeed_moe_layer_cls_names is not None: current_env["SAPIENS_ACCELERATOR_DEEPSPEED_MOE_LAYER_CLS_NAMES"] = str(args.deepspeed_moe_layer_cls_names)
    return cmd, current_env
def prepare_tpu(args: argparse.Namespace, current_env: Dict[str, str], pod: bool = False) -> Tuple[argparse.Namespace, Dict[str, str]]:
    if args.mixed_precision == "bf16" and is_torch_xla_available(check_is_tpu=True):
        if args.downcast_bf16: current_env["XLA_DOWNCAST_BF16"] = "1"
        else: current_env["XLA_USE_BF16"] = "1"
    if args.debug: current_env["SAPIENS_ACCELERATOR_DEBUG_MODE"] = "true"
    if pod:
        args.vm = args.tpu_vm
        args.tpu = args.tpu_name
    return args, current_env
def _convert_nargs_to_dict(nargs: List[str]) -> Dict[str, str]:
    if len(nargs) < 0: return {}
    def _infer_type(s):
        try:
            s = float(s)
            if s // 1 == s: return int(s)
            return s
        except ValueError: return s
    parser = argparse.ArgumentParser()
    _, unknown = parser.parse_known_args(nargs)
    for index, argument in enumerate(unknown):
        if argument.startswith(("-", "--")):
            action = None
            if index + 1 < len(unknown):
                if unknown[index + 1].startswith(("-", "--")): raise ValueError("SageMaker doesn’t support argparse actions for `store_true` or `store_false`. Please define explicit types")
            else: raise ValueError("SageMaker doesn’t support argparse actions for `store_true` or `store_false`. Please define explicit types")
            if action is None: parser.add_argument(argument, type=_infer_type)
            else: parser.add_argument(argument, action=action)
    return {key: (literal_eval(value) if value in ("True", "False") else value) for key, value in parser.parse_args(nargs).__dict__.items()}
def prepare_sagemager_args_inputs(sagemaker_config: SageMakerConfig, args: argparse.Namespace) -> Tuple[argparse.Namespace, Dict[str, Any]]:
    os.environ["AWS_DEFAULT_REGION"] = sagemaker_config.region
    if sagemaker_config.profile is not None: os.environ["AWS_PROFILE"] = sagemaker_config.profile
    elif args.aws_access_key_id is not None and args.aws_secret_access_key is not None:
        os.environ["AWS_ACCESS_KEY_ID"] = args.aws_access_key_id
        os.environ["AWS_SECRET_ACCESS_KEY"] = args.aws_secret_access_key
    else: raise OSError("You need to provide an aws_access_key_id and aws_secret_access_key when not using aws_profile")
    source_dir = os.path.dirname(args.training_script)
    if not source_dir: source_dir = "."
    entry_point = os.path.basename(args.training_script)
    if not entry_point.endswith(".py"): raise ValueError(f'Your training script should be a python script and not "{entry_point}"')
    hyperparameters = _convert_nargs_to_dict(args.training_script_args)
    try: mixed_precision = PrecisionType(args.mixed_precision.lower())
    except ValueError: raise ValueError(f"Unknown mixed_precision mode: {args.mixed_precision.lower()}. Choose between {PrecisionType.list()}.")
    try: dynamo_backend = DynamoBackend(args.dynamo_backend.upper())
    except ValueError: raise ValueError(f"Unknown dynamo backend: {args.dynamo_backend.upper()}. Choose between {DynamoBackend.list()}.")
    environment = {"SAPIENS_ACCELERATOR_USE_SAGEMAKER": "true", "SAPIENS_ACCELERATOR_MIXED_PRECISION": str(mixed_precision), "SAPIENS_ACCELERATOR_DYNAMO_BACKEND": dynamo_backend.value,
    "SAPIENS_ACCELERATOR_DYNAMO_MODE": args.dynamo_mode, "SAPIENS_ACCELERATOR_DYNAMO_USE_FULLGRAPH": str(args.dynamo_use_fullgraph), "SAPIENS_ACCELERATOR_DYNAMO_USE_DYNAMIC": str(args.dynamo_use_dynamic),
    "SAPIENS_ACCELERATOR_SAGEMAKER_DISTRIBUTED_TYPE": sagemaker_config.distributed_type.value}
    if args.mixed_precision.lower() == "fp8":
        if not is_fp8_available(): raise RuntimeError("FP8 is not available on this machine. Please ensure that either Transformer Engine or MSAMP is installed.")
        environment = setup_fp8_env(args, environment)
    distribution = None
    if sagemaker_config.distributed_type == SageMakerDistributedType.DATA_PARALLEL: distribution = {"smdistributed": {"dataparallel": {"enabled": True}}}
    sagemaker_inputs = None
    if sagemaker_config.sagemaker_inputs_file is not None:
        print(f"Loading SageMaker Inputs from {sagemaker_config.sagemaker_inputs_file} file")
        sagemaker_inputs = {}
        with open(sagemaker_config.sagemaker_inputs_file) as file:
            for i, line in enumerate(file):
                if i == 0: continue
                l = line.split("\t")
                sagemaker_inputs[l[0]] = l[1].strip()
    sagemaker_metrics = None
    if sagemaker_config.sagemaker_metrics_file is not None:
        print(f"Loading SageMaker Metrics from {sagemaker_config.sagemaker_metrics_file} file")
        sagemaker_metrics = []
        with open(sagemaker_config.sagemaker_metrics_file) as file:
            for i, line in enumerate(file):
                if i == 0: continue
                l = line.split("\t")
                metric_dict = {"Name": l[0], "Regex": l[1].strip()}
                sagemaker_metrics.append(metric_dict)
    args = {"image_uri": sagemaker_config.image_uri, "entry_point": entry_point, "source_dir": source_dir, "role": sagemaker_config.iam_role_name,
    "transformers_version": sagemaker_config.transformers_version, "pytorch_version": sagemaker_config.pytorch_version, "py_version": sagemaker_config.py_version,
    "base_job_name": sagemaker_config.base_job_name, "instance_count": sagemaker_config.num_machines, "instance_type": sagemaker_config.ec2_instance_type,
    "debugger_hook_config": False, "distribution": distribution, "hyperparameters": hyperparameters, "environment": environment, "metric_definitions": sagemaker_metrics}
    if sagemaker_config.additional_args is not None: args = merge_dicts(sagemaker_config.additional_args, args)
    return args, sagemaker_inputs
def env_var_path_add(env_var_name, path_to_add):
    paths = [p for p in os.environ.get(env_var_name, "").split(":") if len(p) > 0]
    paths.append(str(path_to_add))
    return ":".join(paths)
class PrepareForLaunch:
    def __init__(self, launcher, distributed_type="NO", debug=False):
        self.launcher = launcher
        self.distributed_type = DistributedType(distributed_type)
        self.debug = debug
    def __call__(self, index, *args):
        if self.debug:
            world_size = int(os.environ.get("WORLD_SIZE"))
            rdv_file = os.environ.get("SAPIENS_ACCELERATOR_DEBUG_RDV_FILE")
            torch.distributed.init_process_group("gloo", rank=index, store=torch.distributed.FileStore(rdv_file, world_size), world_size=world_size)
        elif self.distributed_type in (DistributedType.MULTI_GPU, DistributedType.MULTI_MLU, DistributedType.MULTI_MUSA, DistributedType.MULTI_NPU,
        DistributedType.MULTI_XPU, DistributedType.MULTI_CPU):
            os.environ["LOCAL_RANK"] = str(index)
            nproc = int(os.environ.get("NPROC", 1))
            node_rank = int(os.environ.get("NODE_RANK", 0))
            os.environ["RANK"] = str(nproc * node_rank + index)
        os.environ["FORK_LAUNCHED"] = str(1)
        self.launcher(*args)
"""
    ########################################################################################################################################################
    # This algorithm is part of a code library to accelerate Sapiens Technology® Artificial Intelligence models, and its disclosure, distribution,         #
    # or reverse engineering without the company's prior consent will be subject to legal proceedings and actions pursued by our legal department.         #
    ########################################################################################################################################################
"""
