






'''simple docstring'''



from __future__ import annotations

from random import random
from typing import Generic, TypeVar

__UpperCamelCase      :       List[Any]      			=      TypeVar("""KT""")
__UpperCamelCase      :       List[str]      			=      TypeVar("""VT""")
class 							__SCREAMING_SNAKE_CASE       (							Generic[KT, VT]     ):




      def __init__(  self ,   lowerCamelCase = "root" ,   lowerCamelCase = None      )							->Tuple:



         '''simple docstring'''
         __a				     =    key
         __a				     =    value
         __a				     =    []




      def __repr__(  self      )							->str:



         '''simple docstring'''
         return F"""Node({self.key}: {self.value})"""






      @property
      def  __UpperCamelCase							(  self      )							->int:



         '''simple docstring'''
         return len(self.forward      )







class 							__SCREAMING_SNAKE_CASE       (							Generic[KT, VT]     ):




      def __init__(  self ,   lowerCamelCase = 0.5 ,   lowerCamelCase = 16      )							->Dict:



         '''simple docstring'''
         __a				     =    Node[KT, VT]()
         __a				     =    0
         __a				     =    p
         __a				     =    max_level




      def __str__(  self      )							->str:



         '''simple docstring'''
         __a				     =    list(self      )

         if len(lowerCamelCase      ) == 0:
            return F"""SkipList(level={self.level})"""

         __a				     =    max((len(str(lowerCamelCase      )      ) for item in items) ,   default=4      )
         __a				     =    max(lowerCamelCase ,   4      ) + 4

         __a				     =    self.head
         __a				     =    []

         __a				     =    node.forward.copy()
         lines.append(F"""[{node.key}]""".ljust(lowerCamelCase ,   '-'      ) + '* ' * len(lowerCamelCase      )      )
         lines.append(' ' * label_size + '| ' * len(lowerCamelCase      )      )

         while len(node.forward      ) != 0:
            __a				     =    node.forward[0]

            lines.append(
                F"""[{node.key}]""".ljust(lowerCamelCase ,   '-'      )
                + ' '.join(str(n.key      ) if n.key == node.key else '|' for n in forwards      )      )
            lines.append(' ' * label_size + '| ' * len(lowerCamelCase      )      )
            __a				     =    node.forward

         lines.append('None'.ljust(lowerCamelCase      ) + '* ' * len(lowerCamelCase      )      )
         return F"""SkipList(level={self.level})\n""" + "\n".join(lowerCamelCase      )




      def __iter__(  self      )							->Tuple:



         '''simple docstring'''
         __a				     =    self.head

         while len(node.forward      ) != 0:
            yield node.forward[0].key
            __a				     =    node.forward[0]




      def  __UpperCamelCase							(  self      )							->int:



         '''simple docstring'''
         __a				     =    1
         while random() < self.p and level < self.max_level:
            level += 1

         return level




      def  __UpperCamelCase							(  self ,   lowerCamelCase      )							->tuple[Node[KT, VT] | None, list[Node[KT, VT]]]:



         '''simple docstring'''
         __a				     =    []

         __a				     =    self.head

         for i in reversed(range(self.level      )      ):
            # i < node.level - When node level is lesser than `i` decrement `i`.
            # node.forward[i].key < key - Jumping to node with key value higher
            #                             or equal to searched key would result
            #                             in skipping searched key.
            while i < node.level and node.forward[i].key < key:
               __a				     =    node.forward[i]
            # Each leftmost node (relative to searched node) will potentially have to
            # be updated.
            update_vector.append(lowerCamelCase      )

         update_vector.reverse()  # Note that we were inserting values in reverse order.

         # len(node.forward) != 0 - If current node doesn't contain any further
         #                          references then searched key is not present.
         # node.forward[0].key == key - Next node key should be equal to search key
         #                              if key is present.
         if len(node.forward      ) != 0 and node.forward[0].key == key:
            return node.forward[0], update_vector
         else:
            return None, update_vector




      def  __UpperCamelCase							(  self ,   lowerCamelCase      )							->Dict:



         '''simple docstring'''
         __a     ,			__a				     =    self._locate_node(lowerCamelCase      )

         if node is not None:
            for i, update_node in enumerate(lowerCamelCase      ):
               # Remove or replace all references to removed node.
               if update_node.level > i and update_node.forward[i].key == key:
                  if node.level > i:
                     __a				     =    node.forward[i]
                  else:
                     __a				     =    update_node.forward[:i]




      def  __UpperCamelCase							(  self ,   lowerCamelCase ,   lowerCamelCase      )							->Tuple:



         '''simple docstring'''
         __a     ,			__a				     =    self._locate_node(lowerCamelCase      )
         if node is not None:
            __a				     =    value
         else:
            __a				     =    self.random_level()

            if level > self.level:
               # After level increase we have to add additional nodes to head.
               for _ in range(self.level - 1 ,   lowerCamelCase      ):
                  update_vector.append(self.head      )
               __a				     =    level

            __a				     =    Node(lowerCamelCase ,   lowerCamelCase      )

            for i, update_node in enumerate(update_vector[:level]      ):
               # Change references to pass through new node.
               if update_node.level > i:
                  new_node.forward.append(update_node.forward[i]      )

               if update_node.level < i + 1:
                  update_node.forward.append(lowerCamelCase      )
               else:
                  __a				     =    new_node






      def  __UpperCamelCase							(  self ,   lowerCamelCase      )							->VT | None:



         '''simple docstring'''
         __a     ,			__a				     =    self._locate_node(lowerCamelCase      )

         if node is not None:
            return node.value

         return None







def    __UpperCAmelCase	(				)  ->			Tuple:

   """simple docstring"""
   __a				     =    SkipList()
   skip_list.insert('Key1',					3    )
   skip_list.insert('Key2',					12    )
   skip_list.insert('Key3',					41    )
   skip_list.insert('Key4',					-19    )

   __a				     =    skip_list.head
   __a				     =    {}
   while node.level != 0:
      __a				     =    node.forward[0]
      __a				     =    node.value

   assert len(SCREAMING_SNAKE_CASE__    ) == 4
   assert all_values["Key1"] == 3
   assert all_values["Key2"] == 12
   assert all_values["Key3"] == 41
   assert all_values["Key4"] == -19



def    __UpperCAmelCase	(				)  ->			Any:

   """simple docstring"""
   __a				     =    SkipList()
   skip_list.insert('Key1',					10    )
   skip_list.insert('Key1',					12    )

   skip_list.insert('Key5',					7    )
   skip_list.insert('Key7',					10    )
   skip_list.insert('Key10',					5    )

   skip_list.insert('Key7',					7    )
   skip_list.insert('Key5',					5    )
   skip_list.insert('Key10',					10    )

   __a				     =    skip_list.head
   __a				     =    {}
   while node.level != 0:
      __a				     =    node.forward[0]
      __a				     =    node.value

   if len(SCREAMING_SNAKE_CASE__    ) != 4:
      print()
   assert len(SCREAMING_SNAKE_CASE__    ) == 4
   assert all_values["Key1"] == 12
   assert all_values["Key7"] == 7
   assert all_values["Key5"] == 5
   assert all_values["Key10"] == 10



def    __UpperCAmelCase	(				)  ->			Union[str, Any]:

   """simple docstring"""
   __a				     =    SkipList()
   assert skip_list.find('Some key'    ) is None



def    __UpperCAmelCase	(				)  ->			int:

   """simple docstring"""
   __a				     =    SkipList()

   skip_list.insert('Key2',					20    )
   assert skip_list.find('Key2'    ) == 20

   skip_list.insert('Some Key',					10    )
   skip_list.insert('Key2',					8    )
   skip_list.insert('V',					13    )

   assert skip_list.find('Y'    ) is None
   assert skip_list.find('Key2'    ) == 8
   assert skip_list.find('Some Key'    ) == 10
   assert skip_list.find('V'    ) == 13



def    __UpperCAmelCase	(				)  ->			str:

   """simple docstring"""
   __a				     =    SkipList()
   skip_list.delete('Some key'    )

   assert len(skip_list.head.forward    ) == 0



def    __UpperCAmelCase	(				)  ->			Union[str, Any]:

   """simple docstring"""
   __a				     =    SkipList()

   skip_list.insert('Key1',					12    )
   skip_list.insert('V',					13    )
   skip_list.insert('X',					14    )
   skip_list.insert('Key2',					15    )

   skip_list.delete('V'    )
   skip_list.delete('Key2'    )

   assert skip_list.find('V'    ) is None
   assert skip_list.find('Key2'    ) is None



def    __UpperCAmelCase	(				)  ->			Union[str, Any]:

   """simple docstring"""
   __a				     =    SkipList()

   skip_list.insert('Key1',					12    )
   skip_list.insert('V',					13    )
   skip_list.insert('X',					14    )
   skip_list.insert('Key2',					15    )

   skip_list.delete('V'    )
   assert skip_list.find('V'    ) is None
   assert skip_list.find('X'    ) == 14
   assert skip_list.find('Key1'    ) == 12
   assert skip_list.find('Key2'    ) == 15

   skip_list.delete('X'    )
   assert skip_list.find('V'    ) is None
   assert skip_list.find('X'    ) is None
   assert skip_list.find('Key1'    ) == 12
   assert skip_list.find('Key2'    ) == 15

   skip_list.delete('Key1'    )
   assert skip_list.find('V'    ) is None
   assert skip_list.find('X'    ) is None
   assert skip_list.find('Key1'    ) is None
   assert skip_list.find('Key2'    ) == 15

   skip_list.delete('Key2'    )
   assert skip_list.find('V'    ) is None
   assert skip_list.find('X'    ) is None
   assert skip_list.find('Key1'    ) is None
   assert skip_list.find('Key2'    ) is None



def    __UpperCAmelCase	(				)  ->			Optional[Any]:

   """simple docstring"""
   __a				     =    SkipList()

   skip_list.insert('Key1',					12    )
   skip_list.insert('V',					13    )
   skip_list.insert('X',					142    )
   skip_list.insert('Key2',					15    )

   skip_list.delete('X'    )

   def traverse_keys(SCREAMING_SNAKE_CASE__:	Optional[Any]    ):
      yield node.key
      for forward_node in node.forward:
         yield from traverse_keys(SCREAMING_SNAKE_CASE__    )

   assert len(set(traverse_keys(skip_list.head    )    )    ) == 4



def    __UpperCAmelCase	(				)  ->			str:

   """simple docstring"""
   def is_sorted(SCREAMING_SNAKE_CASE__:	Tuple    ):
      return all(next_item >= item for item, next_item in zip(SCREAMING_SNAKE_CASE__,					lst[1:]    )    )

   __a				     =    SkipList()
   for i in range(10    ):
      skip_list.insert(SCREAMING_SNAKE_CASE__,					SCREAMING_SNAKE_CASE__    )
   assert is_sorted(list(SCREAMING_SNAKE_CASE__    )    )
   skip_list.delete(5    )
   skip_list.delete(8    )
   skip_list.delete(2    )
   assert is_sorted(list(SCREAMING_SNAKE_CASE__    )    )
   skip_list.insert(-12,					-12    )
   skip_list.insert(77,					77    )
   assert is_sorted(list(SCREAMING_SNAKE_CASE__    )    )



def    __UpperCAmelCase	(				)  ->			Tuple:

   """simple docstring"""
   for _ in range(100    ):
      # Repeat test 100 times due to the probabilistic nature of skip list
      # random values == random bugs
      test_insert()
      test_insert_overrides_existing_value()

      test_searching_empty_list_returns_none()
      test_search()

      test_deleting_item_from_empty_list_do_nothing()
      test_deleted_items_are_not_founded_by_find_method()
      test_delete_removes_only_given_key()
      test_delete_doesnt_leave_dead_nodes()

      test_iter_always_yields_sorted_values()



def    __UpperCAmelCase	(				)  ->			Optional[Any]:

   """simple docstring"""
   __a				     =    SkipList()
   skip_list.insert(2,					'2'    )
   skip_list.insert(4,					'4'    )
   skip_list.insert(6,					'4'    )
   skip_list.insert(4,					'5'    )
   skip_list.insert(8,					'4'    )
   skip_list.insert(9,					'4'    )

   skip_list.delete(4    )

   print(SCREAMING_SNAKE_CASE__    )


if __name__ == "__main__":
 import doctest

 doctest.testmod()
 main()







'''simple docstring'''



# Copyright 2021 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 json
import os

from ...utils.constants import SAGEMAKER_PARALLEL_EC2_INSTANCES, TORCH_DYNAMO_MODES
from ...utils.dataclasses import ComputeEnvironment, SageMakerDistributedType
from ...utils.imports import is_botoa_available
from .config_args import SageMakerConfig
from .config_utils import (
    DYNAMO_BACKENDS,
    _ask_field,
    _ask_options,
    _convert_dynamo_backend,
    _convert_mixed_precision,
    _convert_sagemaker_distributed_mode,
    _convert_yes_no_to_bool,
)


if is_botoa_available():
 import botoa  # noqa: F401



def    __UpperCAmelCase	(				SCREAMING_SNAKE_CASE__:	Union[str, Any]    )  ->			List[str]:

   """simple docstring"""
   __a				     =    botoa.client('iam'    )

   __a				     =    {
       'Version': '2012-10-17',
       'Statement': [
           {'Effect': 'Allow', 'Principal': {'Service': 'sagemaker.amazonaws.com'}, 'Action': 'sts:AssumeRole'}
       ],
   }
   try:
      # create the role, associated with the chosen trust policy
      iam_client.create_role(
          RoleName=SCREAMING_SNAKE_CASE__,					AssumeRolePolicyDocument=json.dumps(SCREAMING_SNAKE_CASE__,					indent=2    )    )
      __a				     =    {
          'Version': '2012-10-17',
          'Statement': [
              {
                  'Effect': 'Allow',
                  'Action': [
                      'sagemaker:*',
                      'ecr:GetDownloadUrlForLayer',
                      'ecr:BatchGetImage',
                      'ecr:BatchCheckLayerAvailability',
                      'ecr:GetAuthorizationToken',
                      'cloudwatch:PutMetricData',
                      'cloudwatch:GetMetricData',
                      'cloudwatch:GetMetricStatistics',
                      'cloudwatch:ListMetrics',
                      'logs:CreateLogGroup',
                      'logs:CreateLogStream',
                      'logs:DescribeLogStreams',
                      'logs:PutLogEvents',
                      'logs:GetLogEvents',
                      's3:CreateBucket',
                      's3:ListBucket',
                      's3:GetBucketLocation',
                      's3:GetObject',
                      's3:PutObject',
                  ],
                  'Resource': '*',
              }
          ],
      }
      # attach policy to role
      iam_client.put_role_policy(
          RoleName=SCREAMING_SNAKE_CASE__,					PolicyName=f"""{role_name}_policy_permission""",					PolicyDocument=json.dumps(SCREAMING_SNAKE_CASE__,					indent=2    ),					)
   except iam_client.exceptions.EntityAlreadyExistsException:
      print(f"""role {role_name} already exists. Using existing one"""    )



def    __UpperCAmelCase	(				SCREAMING_SNAKE_CASE__:	List[str]    )  ->			Optional[Any]:

   """simple docstring"""
   __a				     =    botoa.client('iam'    )
   return iam_client.get_role(RoleName=SCREAMING_SNAKE_CASE__    )["Role"]["Arn"]



def    __UpperCAmelCase	(				)  ->			Any:

   """simple docstring"""
   __a				     =    _ask_options(
       'How do you want to authorize?',					['AWS Profile', 'Credentials (AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY) '],					SCREAMING_SNAKE_CASE__,					)
   __a				     =    None
   if credentials_configuration == 0:
      __a				     =    _ask_field('Enter your AWS Profile name: [default] ',					default='default'    )
      __a				     =    aws_profile
   else:
      print(
          'Note you will need to provide AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY when you launch you training script with,'
          '`accelerate launch --aws_access_key_id XXX --aws_secret_access_key YYY`'    )
      __a				     =    _ask_field('AWS Access Key ID: '    )
      __a				     =    aws_access_key_id

      __a				     =    _ask_field('AWS Secret Access Key: '    )
      __a				     =    aws_secret_access_key

   __a				     =    _ask_field('Enter your AWS Region: [us-east-1]',					default='us-east-1'    )
   __a				     =    aws_region

   __a				     =    _ask_options(
       'Do you already have an IAM Role for executing Amazon SageMaker Training Jobs?',					['Provide IAM Role name', 'Create new IAM role using credentials'],					SCREAMING_SNAKE_CASE__,					)
   if role_management == 0:
      __a				     =    _ask_field('Enter your IAM role name: '    )
   else:
      __a				     =    'accelerate_sagemaker_execution_role'
      print(f"""Accelerate will create an iam role \"{iam_role_name}\" using the provided credentials"""    )
      _create_iam_role_for_sagemaker(SCREAMING_SNAKE_CASE__    )

   __a				     =    _ask_field(
       'Do you want to use custom Docker image? [yes/NO]: ',					_convert_yes_no_to_bool,					default=SCREAMING_SNAKE_CASE__,					error_message='Please enter yes or no.',					)
   __a				     =    None
   if is_custom_docker_image:
      __a				     =    _ask_field('Enter your Docker image: ',					lambda SCREAMING_SNAKE_CASE__    : str(SCREAMING_SNAKE_CASE__    ).lower()    )

   __a				     =    _ask_field(
       'Do you want to provide SageMaker input channels with data locations? [yes/NO]: ',					_convert_yes_no_to_bool,					default=SCREAMING_SNAKE_CASE__,					error_message='Please enter yes or no.',					)
   __a				     =    None
   if is_sagemaker_inputs_enabled:
      __a				     =    _ask_field(
          'Enter the path to the SageMaker inputs TSV file with columns (channel_name, data_location): ',					lambda SCREAMING_SNAKE_CASE__    : str(SCREAMING_SNAKE_CASE__    ).lower(),					)

   __a				     =    _ask_field(
       'Do you want to enable SageMaker metrics? [yes/NO]: ',					_convert_yes_no_to_bool,					default=SCREAMING_SNAKE_CASE__,					error_message='Please enter yes or no.',					)
   __a				     =    None
   if is_sagemaker_metrics_enabled:
      __a				     =    _ask_field(
          'Enter the path to the SageMaker metrics TSV file with columns (metric_name, metric_regex): ',					lambda SCREAMING_SNAKE_CASE__    : str(SCREAMING_SNAKE_CASE__    ).lower(),					)

   __a				     =    _ask_options(
       'What is the distributed mode?',					['No distributed training', 'Data parallelism'],					_convert_sagemaker_distributed_mode,					)
   __a				     =    {}
   __a				     =    _ask_field(
       'Do you wish to optimize your script with torch dynamo?[yes/NO]:',					_convert_yes_no_to_bool,					default=SCREAMING_SNAKE_CASE__,					error_message='Please enter yes or no.',					)
   if use_dynamo:
      __a				     =    'dynamo_'
      __a				     =    _ask_options(
          'Which dynamo backend would you like to use?',					[x.lower() for x in DYNAMO_BACKENDS],					_convert_dynamo_backend,					default=2,					)
      __a				     =    _ask_field(
          'Do you want to customize the defaults sent to torch.compile? [yes/NO]: ',					_convert_yes_no_to_bool,					default=SCREAMING_SNAKE_CASE__,					error_message='Please enter yes or no.',					)

      if use_custom_options:
         __a				     =    _ask_options(
             'Which mode do you want to use?',					SCREAMING_SNAKE_CASE__,					lambda SCREAMING_SNAKE_CASE__    : TORCH_DYNAMO_MODES[int(SCREAMING_SNAKE_CASE__    )],					default='default',					)
         __a				     =    _ask_field(
             'Do you want the fullgraph mode or it is ok to break model into several subgraphs? [yes/NO]: ',					_convert_yes_no_to_bool,					default=SCREAMING_SNAKE_CASE__,					error_message='Please enter yes or no.',					)
         __a				     =    _ask_field(
             'Do you want to enable dynamic shape tracing? [yes/NO]: ',					_convert_yes_no_to_bool,					default=SCREAMING_SNAKE_CASE__,					error_message='Please enter yes or no.',					)
   __a				     =    'Which EC2 instance type you want to use for your training?'
   if distributed_type != SageMakerDistributedType.NO:
      __a				     =    _ask_options(
          SCREAMING_SNAKE_CASE__,					SCREAMING_SNAKE_CASE__,					lambda SCREAMING_SNAKE_CASE__    : SAGEMAKER_PARALLEL_EC2_INSTANCES[int(SCREAMING_SNAKE_CASE__    )]    )
   else:
      eca_instance_query += "? [ml.p3.2xlarge]:"
      __a				     =    _ask_field(SCREAMING_SNAKE_CASE__,					lambda SCREAMING_SNAKE_CASE__    : str(SCREAMING_SNAKE_CASE__    ).lower(),					default='ml.p3.2xlarge'    )

   __a				     =    1
   if distributed_type in (SageMakerDistributedType.DATA_PARALLEL, SageMakerDistributedType.MODEL_PARALLEL):
      __a				     =    _ask_field(
          'How many machines do you want use? [1]: ',					SCREAMING_SNAKE_CASE__,					default=1,					)

   __a				     =    _ask_options(
       'Do you wish to use FP16 or BF16 (mixed precision)?',					['no', 'fp16', 'bf16', 'fp8'],					_convert_mixed_precision,					)

   if use_dynamo and mixed_precision == "no":
      print(
          'Torch dynamo used without mixed precision requires TF32 to be efficient. Accelerate will enable it by default when launching your scripts.'    )

   return SageMakerConfig(
       image_uri=SCREAMING_SNAKE_CASE__,					compute_environment=ComputeEnvironment.AMAZON_SAGEMAKER,					distributed_type=SCREAMING_SNAKE_CASE__,					use_cpu=SCREAMING_SNAKE_CASE__,					dynamo_config=SCREAMING_SNAKE_CASE__,					eca_instance_type=SCREAMING_SNAKE_CASE__,					profile=SCREAMING_SNAKE_CASE__,					region=SCREAMING_SNAKE_CASE__,					iam_role_name=SCREAMING_SNAKE_CASE__,					mixed_precision=SCREAMING_SNAKE_CASE__,					num_machines=SCREAMING_SNAKE_CASE__,					sagemaker_inputs_file=SCREAMING_SNAKE_CASE__,					sagemaker_metrics_file=SCREAMING_SNAKE_CASE__,					)