



# 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 subprocess

from packaging.version import Version, parse

from accelerate.commands.config.config_args import default_config_file, load_config_from_file


SCREAMING_SNAKE_CASE_:int					=						"Run commands across TPU VMs for initial setup before running `accelerate launch`."




def      __UpperCamelCase     (  _lowerCAmelCase=None		)	->  Union[str, Any]:



    """simple docstring"""






    if subparsers is not None:
        A			:	Dict										=       subparsers.add_parser("""tpu-config"""				,					description=_description		)
    else:
        A			:	str										=       argparse.ArgumentParser("""Accelerate tpu-config command"""				,					description=_description		)
    # Core arguments
    A			:	Tuple										=       parser.add_argument_group(
        """Config Arguments"""				,					"""Arguments that can be configured through `accelerate config`."""		)
    config_args.add_argument(
        """--config_file"""				,					type=_lowerCAmelCase				,					default=_lowerCAmelCase				,					help="""Path to the config file to use for accelerate."""				,					)
    config_args.add_argument(
        """--tpu_name"""				,					default=_lowerCAmelCase				,					help="""The name of the TPU to use. If not specified, will use the TPU specified in the config file."""				,					)
    config_args.add_argument(
        """--tpu_zone"""				,					default=_lowerCAmelCase				,					help="""The zone of the TPU to use. If not specified, will use the zone specified in the config file."""				,					)
    A			:	List[str]										=       parser.add_argument_group("""TPU Arguments"""				,					"""Arguments for options ran inside the TPU."""		)
    pod_args.add_argument(
        """--use_alpha"""				,					action="""store_true"""				,					help="""Whether to use `gcloud alpha` when running the TPU training script instead of `gcloud`."""				,					)
    pod_args.add_argument(
        """--command_file"""				,					default=_lowerCAmelCase				,					help="""The path to the file containing the commands to run on the pod on startup."""				,					)
    pod_args.add_argument(
        """--command"""				,					action="""append"""				,					nargs="""+"""				,					help="""A command to run on the pod. Can be passed multiple times."""				,					)
    pod_args.add_argument(
        """--install_accelerate"""				,					action="""store_true"""				,					help="""Whether to install accelerate on the pod. Defaults to False."""				,					)
    pod_args.add_argument(
        """--accelerate_version"""				,					default="""latest"""				,					help="""The version of accelerate to install on the pod. If not specified, will use the latest pypi version. Specify \'dev\' to install from GitHub."""				,					)
    pod_args.add_argument(
        """--debug"""				,					action="""store_true"""				,					help="""If set, will print the command that would be run instead of running it."""		)

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




def      __UpperCamelCase     (  _lowerCAmelCase		)	->  Dict:



    """simple docstring"""






    A			:	str										=       None

    # Get the default from the config file if it exists.
    if args.config_file is not None or os.path.isfile(_lowerCAmelCase		):
        A			:	Dict										=       load_config_from_file(args.config_file		)
        if not args.command_file and defaults.command_file is not None and not args.command:
            A			:	Any										=       defaults.command_file
        if not args.command and defaults.commands is not None:
            A			:	List[Any]										=       defaults.commands
        if not args.tpu_name:
            A			:	str										=       defaults.tpu_name
        if not args.tpu_zone:
            A			:	str										=       defaults.tpu_zone
    if args.accelerate_version == "dev":
        A			:	Union[str, Any]										=       '''git+https://github.com/huggingface/accelerate.git'''
    elif args.accelerate_version == "latest":
        A			:	List[Any]										=       '''accelerate -U'''
    elif isinstance(parse(args.accelerate_version		)				,					_lowerCAmelCase		):
        A			:	List[Any]										=       f'''accelerate=={args.accelerate_version}'''

    if not args.command_file and not args.command:
        raise ValueError("""You must specify either a command file or a command to run on the pod."""		)

    if args.command_file:
        with open(args.command_file				,					"""r"""		) as f:
            A			:	int										=       [f.read().splitlines()]

             # To turn list of lists into list of strings
    if isinstance(args.command[0]				,					_lowerCAmelCase		):
        A			:	Optional[Any]										=       [line for cmd in args.command for line in cmd]
    # Default to the shared folder and install accelerate
    A			:	Union[str, Any]										=       ['''cd /usr/share''']
    if args.install_accelerate:
        new_cmd += [f'''pip install {args.accelerate_version}''']
    new_cmd += args.command
    A			:	Optional[int]										=       '''; '''.join(_lowerCAmelCase		)

    # Then send it to gcloud
    # Eventually try to use google-api-core to do this instead of subprocess
    A			:	Any										=       ['''gcloud''']
    if args.use_alpha:
        cmd += ["alpha"]
    cmd += [
        "compute",
        "tpus",
        "tpu-vm",
        "ssh",
        args.tpu_name,
        "--zone",
        args.tpu_zone,
        "--command",
        args.command,
        "--worker",
        "all",
    ]
    if args.debug:
        print(f'''Running {" ".join(_lowerCAmelCase		)}'''		)
        return
    subprocess.run(_lowerCAmelCase		)
    print("""Successfully setup pod."""		)




def      __UpperCamelCase     (  )	->  Optional[Any]:



    """simple docstring"""






    A			:	List[str]										=       tpu_command_parser()
    A			:	Tuple										=       parser.parse_args()

    tpu_command_launcher(_lowerCAmelCase		)







import os
from pathlib import Path




def      __UpperCamelCase     (  )	->  Any:



    """simple docstring"""






    from torch.utils.cpp_extension import load

    A			:	Any										=       Path(_lowerCAmelCase		).resolve().parent.parent.parent / """kernels""" / """deformable_detr"""
    A			:	int										=       [
        root / filename
        for filename in [
            """vision.cpp""",
            os.path.join("""cpu"""				,					"""ms_deform_attn_cpu.cpp"""		),
            os.path.join("""cuda"""				,					"""ms_deform_attn_cuda.cu"""		),
        ]
    ]

    load(
        """MultiScaleDeformableAttention"""				,					_lowerCAmelCase				,					with_cuda=_lowerCAmelCase				,					extra_include_paths=[str(_lowerCAmelCase		)]				,					extra_cflags=["""-DWITH_CUDA=1"""]				,					extra_cuda_cflags=[
            """-DCUDA_HAS_FP16=1""",
            """-D__CUDA_NO_HALF_OPERATORS__""",
            """-D__CUDA_NO_HALF_CONVERSIONS__""",
            """-D__CUDA_NO_HALF2_OPERATORS__""",
        ]				,					)

    import MultiScaleDeformableAttention as MSDA

    return MSDA


