


import os
import sys
import tempfile

import torch

from .state import AcceleratorState
from .utils import PrecisionType, PrepareForLaunch, is_mps_available, patch_environment




def       a_ (       SCREAMING_SNAKE_CASE__   : List[Any]						,     SCREAMING_SNAKE_CASE__   : Optional[Any]=()						,     SCREAMING_SNAKE_CASE__   : int=None						,     SCREAMING_SNAKE_CASE__   : Union[str, Any]="no"						,     SCREAMING_SNAKE_CASE__   : int="29500"	):
       '''simple docstring'''





       _lowerCamelCase			:						List[str]            =False
       _lowerCamelCase			:						str            =False
       if any(key.startswith('KAGGLE'	) for key in os.environ.keys()	):
              _lowerCamelCase			:						List[Any]            =True
       elif "IPython" in sys.modules:
              _lowerCamelCase			:						List[str]            ='google.colab' in str(sys.modules['IPython'].get_ipython()	)

       try:
              _lowerCamelCase			:						List[Any]            =PrecisionType(mixed_precision.lower()	)
       except ValueError:
              raise ValueError(
                  F'''Unknown mixed_precision mode: {args.mixed_precision.lower()}. Choose between {PrecisionType.list()}.'''	)

       if (in_colab or in_kaggle) and (os.environ.get('TPU_NAME'						,     SCREAMING_SNAKE_CASE__	) is not None):
              # TPU launch
              import torch_xla.distributed.xla_multiprocessing as xmp

              if len(AcceleratorState._shared_state	) > 0:
                     raise ValueError(
                         'To train on TPU in Colab or Kaggle Kernel, the `Accelerator` should only be initialized inside '
                         'your training function. Restart your notebook and make sure no cells initializes an '
                         '`Accelerator`.'	)
              if num_processes is None:
                     _lowerCamelCase			:						Any            =8

              _lowerCamelCase			:						List[str]            =PrepareForLaunch(SCREAMING_SNAKE_CASE__						,     distributed_type='TPU'	)
              print(F'''Launching a training on {num_processes} TPU cores.'''	)
              xmp.spawn(SCREAMING_SNAKE_CASE__						,     args=SCREAMING_SNAKE_CASE__						,     nprocs=SCREAMING_SNAKE_CASE__						,     start_method='fork'	)
       elif in_colab:
              # No need for a distributed launch otherwise as it's either CPU or one GPU.
              if torch.cuda.is_available():
                     print('Launching training on one GPU.'	)
              else:
                     print('Launching training on one CPU.'	)
              function(*SCREAMING_SNAKE_CASE__	)
       else:
              if num_processes is None:
                     raise ValueError(
                         'You have to specify the number of GPUs you would like to use, add `num_processes=...` to your call.'	)

              if num_processes > 1:
                     # Multi-GPU launch
                     from torch.multiprocessing import start_processes
                     from torch.multiprocessing.spawn import ProcessRaisedException

                     if len(AcceleratorState._shared_state	) > 0:
                            raise ValueError(
                                'To launch a multi-GPU training from your notebook, the `Accelerator` should only be initialized '
                                'inside your training function. Restart your notebook and make sure no cells initializes an '
                                '`Accelerator`.'	)

                     if torch.cuda.is_initialized():
                            raise ValueError(
                                'To launch a multi-GPU training from your notebook, you need to avoid running any instruction '
                                'using `torch.cuda` in any cell. Restart your notebook and make sure no cells use any CUDA '
                                'function.'	)

                     # torch.distributed will expect a few environment variable to be here. We set the ones common to each
                     # process here (the other ones will be set be the launcher).
                     with patch_environment(
                         world_size=SCREAMING_SNAKE_CASE__						,     master_addr='127.0.01'						,     master_port=SCREAMING_SNAKE_CASE__						,     mixed_precision=SCREAMING_SNAKE_CASE__	):
                            _lowerCamelCase			:						Optional[int]            =PrepareForLaunch(SCREAMING_SNAKE_CASE__						,     distributed_type='MULTI_GPU'	)
                            print(F'''Launching training on {num_processes} GPUs.'''	)
                            try:
                                   start_processes(SCREAMING_SNAKE_CASE__						,     args=SCREAMING_SNAKE_CASE__						,     nprocs=SCREAMING_SNAKE_CASE__						,     start_method='fork'	)
                            except ProcessRaisedException as e:
                                   if "Cannot re-initialize CUDA in forked subprocess" in e.args[0]:
                                          raise RuntimeError(
                                              'CUDA has been initialized before the `notebook_launcher` could create a forked subprocess. '
                                              'This likely stems from an outside import causing issues once the `notebook_launcher()` is called. '
                                              'Please review your imports and test them when running the `notebook_launcher()` to identify '
                                              'which one is problematic.'	) from e

              else:
                     # No need for a distributed launch otherwise as it's either CPU, GPU or MPS.
                     if is_mps_available():
                            _lowerCamelCase			:						Dict            ='1'
                            print('Launching training on MPS.'	)
                     elif torch.cuda.is_available():
                            print('Launching training on one GPU.'	)
                     else:
                            print('Launching training on CPU.'	)
                     function(*SCREAMING_SNAKE_CASE__	)




def       a_ (       SCREAMING_SNAKE_CASE__   : Union[str, Any]						,     SCREAMING_SNAKE_CASE__   : List[Any]=()						,     SCREAMING_SNAKE_CASE__   : Tuple=2	):
       '''simple docstring'''





       from torch.multiprocessing import start_processes

       with tempfile.NamedTemporaryFile() as tmp_file:
              # torch.distributed will expect a few environment variable to be here. We set the ones common to each
              # process here (the other ones will be set be the launcher).
              with patch_environment(
                  world_size=SCREAMING_SNAKE_CASE__						,     master_addr='127.0.01'						,     master_port='29500'						,     accelerate_mixed_precision='no'						,     accelerate_debug_rdv_file=tmp_file.name						,     accelerate_use_cpu='yes'						,     ):
                     _lowerCamelCase			:						List[Any]            =PrepareForLaunch(SCREAMING_SNAKE_CASE__						,     debug=SCREAMING_SNAKE_CASE__	)
                     start_processes(SCREAMING_SNAKE_CASE__						,     args=SCREAMING_SNAKE_CASE__						,     nprocs=SCREAMING_SNAKE_CASE__						,     start_method='fork'	)






# Copyright 2022 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import os
import platform

import numpy as np
import psutil
import torch

from accelerate import __version__ as version
from accelerate.commands.config import default_config_file, load_config_from_file

from ..utils import is_npu_available, is_xpu_available




def       a_ (       SCREAMING_SNAKE_CASE__   : str=None	):
       '''simple docstring'''





       if subparsers is not None:
              _lowerCamelCase			:						str            =subparsers.add_parser('env'	)
       else:
              _lowerCamelCase			:						List[Any]            =argparse.ArgumentParser('Accelerate env command'	)

       parser.add_argument(
           '--config_file'						,     default=SCREAMING_SNAKE_CASE__						,     help='The config file to use for the default values in the launching script.'	)

       if subparsers is not None:
              parser.set_defaults(func=SCREAMING_SNAKE_CASE__	)
       return parser




def       a_ (       SCREAMING_SNAKE_CASE__   : int	):
       '''simple docstring'''





       _lowerCamelCase			:						int            =torch.__version__
       _lowerCamelCase			:						Optional[Any]            =torch.cuda.is_available()
       _lowerCamelCase			:						int            =is_xpu_available()
       _lowerCamelCase			:						Union[str, Any]            =is_npu_available()

       _lowerCamelCase			:						Optional[int]            ='Not found'
       # Get the default from the config file.
       if args.config_file is not None or os.path.isfile(SCREAMING_SNAKE_CASE__	):
              _lowerCamelCase			:						Any            =load_config_from_file(args.config_file	).to_dict()

       _lowerCamelCase			:						List[str]            ={
           '`Accelerate` version': version,
           'Platform': platform.platform(),
           'Python version': platform.python_version(),
           'Numpy version': np.__version__,
           'PyTorch version (GPU?)': F'''{pt_version} ({pt_cuda_available})''',
           'PyTorch XPU available': str(SCREAMING_SNAKE_CASE__	),
           'PyTorch NPU available': str(SCREAMING_SNAKE_CASE__	),
           'System RAM': F'''{psutil.virtual_memory().total / 1_024 ** 3:.2f} GB''',
       }
       if pt_cuda_available:
              _lowerCamelCase			:						List[Any]            =torch.cuda.get_device_name()

       print('\nCopy-and-paste the text below in your GitHub issue\n'	)
       print('\n'.join([F'''- {prop}: {val}''' for prop, val in info.items()]	)	)

       print('- `Accelerate` default config:' if args.config_file is None else '- `Accelerate` config passed:'	)
       _lowerCamelCase			:						Dict            =(
           '\n'.join([F'''\t- {prop}: {val}''' for prop, val in accelerate_config.items()]	)
           if isinstance(SCREAMING_SNAKE_CASE__						,     SCREAMING_SNAKE_CASE__	)
           else F'''\t{accelerate_config}'''
       )
       print(SCREAMING_SNAKE_CASE__	)

       _lowerCamelCase			:						str            =accelerate_config

       return info




def       a_ (       ):
       '''simple docstring'''





       _lowerCamelCase			:						Optional[int]            =env_command_parser()
       _lowerCamelCase			:						Any            =parser.parse_args()
       env_command(SCREAMING_SNAKE_CASE__	)
       return 0


if __name__ == "__main__":
 raise SystemExit(main())


