




'''simple docstring'''



import os
import sys
from pathlib import Path

import pytest

from datasets import Dataset, IterableDataset
from datasets.distributed import split_dataset_by_node

from .utils import execute_subprocess_async, get_torch_dist_unique_port, require_torch







def 			_UpperCamelCase ()->   Optional[Any]:




     '''simple docstring'''





     __snake_case    					=  Dataset.from_dict({'''i''': range(17	)}	)
     __snake_case    					=  len(_lowerCamelCase	)
     __snake_case    					=  3
     __snake_case    					=  [
         split_dataset_by_node(_lowerCamelCase   ,	rank=_lowerCamelCase   ,	world_size=_lowerCamelCase	) for rank in range(_lowerCamelCase	)
     ]
     assert sum(len(_lowerCamelCase	) for ds in datasets_per_rank	) == full_size
     assert len({tuple(x.values()	) for ds in datasets_per_rank for x in ds}	) == full_size







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




     '''simple docstring'''





     def gen():
          return ({"i": i} for i in range(17	))

     __snake_case    					=  3
     __snake_case    					=  IterableDataset.from_generator(_lowerCamelCase	)
     __snake_case    					=  len(list(_lowerCamelCase	)	)
     __snake_case    					=  [
         split_dataset_by_node(_lowerCamelCase   ,	rank=_lowerCamelCase   ,	world_size=_lowerCamelCase	) for rank in range(_lowerCamelCase	)
     ]
     assert sum(len(list(_lowerCamelCase	)	) for ds in datasets_per_rank	) == full_size
     assert len({tuple(x.values()	) for ds in datasets_per_rank for x in ds}	) == full_size







@pytest.mark.parametrize('''shards_per_node'''   ,	[1, 2, 3]	)
def 			_UpperCamelCase (_lowerCamelCase							:     Dict	)->   Any:




     '''simple docstring'''





     def gen(_lowerCamelCase							:     Tuple	):
          for shard in shards:
               yield from ({"i": i, "shard": shard} for i in range(17	))

     __snake_case    					=  3
     __snake_case    					=  shards_per_node * world_size
     __snake_case    					=  {'''shards''': [f'''shard_{shard_idx}.txt''' for shard_idx in range(_lowerCamelCase	)]}
     __snake_case    					=  IterableDataset.from_generator(_lowerCamelCase   ,	gen_kwargs=_lowerCamelCase	)
     __snake_case    					=  len(list(_lowerCamelCase	)	)
     assert full_ds.n_shards == world_size * shards_per_node
     __snake_case    					=  [
         split_dataset_by_node(_lowerCamelCase   ,	rank=_lowerCamelCase   ,	world_size=_lowerCamelCase	) for rank in range(_lowerCamelCase	)
     ]
     assert [ds.n_shards for ds in datasets_per_rank] == [shards_per_node] * world_size
     assert sum(len(list(_lowerCamelCase	)	) for ds in datasets_per_rank	) == full_size
     assert len({tuple(x.values()	) for ds in datasets_per_rank for x in ds}	) == full_size







def 			_UpperCamelCase ()->   Optional[int]:




     '''simple docstring'''





     def gen():
          return ({"i": i} for i in range(17	))

     __snake_case    					=  2
     __snake_case    					=  IterableDataset.from_generator(_lowerCamelCase	)
     __snake_case    					=  len(list(_lowerCamelCase	)	)

     __snake_case    					=  split_dataset_by_node(_lowerCamelCase   ,	rank=0   ,	world_size=_lowerCamelCase	).shuffle(seed=42	)
     assert len(list(_lowerCamelCase	)	) == 1 + full_size // world_size
     with pytest.raises(_lowerCamelCase	):
          split_dataset_by_node(_lowerCamelCase   ,	rank=0   ,	world_size=_lowerCamelCase	).shuffle()

     __snake_case    					=  split_dataset_by_node(full_ds.shuffle(seed=42	)   ,	rank=0   ,	world_size=_lowerCamelCase	)
     assert len(list(_lowerCamelCase	)	) == 1 + full_size // world_size
     with pytest.raises(_lowerCamelCase	):
          split_dataset_by_node(full_ds.shuffle()   ,	rank=0   ,	world_size=_lowerCamelCase	)







@pytest.mark.parametrize('''streaming'''   ,	[False, True]	)
@require_torch
@pytest.mark.skipif(os.name == '''nt'''   ,	reason='''execute_subprocess_async doesn\'t support windows'''	)
@pytest.mark.integration
def 			_UpperCamelCase (_lowerCamelCase							:     List[Any]	)->   Union[str, Any]:




     '''simple docstring'''





     __snake_case    					=  2
     __snake_case    					=  get_torch_dist_unique_port()
     __snake_case    					=  Path(_lowerCamelCase	).resolve().parent / '''distributed_scripts''' / '''launch_torch_distributed.py'''
     __snake_case    					=  f'''
        -m torch.distributed.launch
        --nproc_per_node={nproc_per_node}
        --master_port={master_port}
        {test_script}
    '''.split()
     __snake_case    					=  f'''
        --streaming={streaming}
    '''.split()
     __snake_case    					=  [sys.executable] + distributed_args + args
     execute_subprocess_async(_lowerCamelCase   ,	env=os.environ.copy()	)







@pytest.mark.parametrize(
    '''nproc_per_node, num_workers'''   ,	[
        (2, 2),  # each node has 2 shards and each worker has 1 shards
        (3, 2),  # each node uses all the shards but skips examples, and each worker has 2 shards
    ]   ,	)
@require_torch
@pytest.mark.skipif(os.name == '''nt'''   ,	reason='''execute_subprocess_async doesn\'t support windows'''	)
@pytest.mark.integration
def 			_UpperCamelCase (_lowerCamelCase							:     List[Any]   ,	_lowerCamelCase							:     Optional[int]	)->   Dict:




     '''simple docstring'''





     __snake_case    					=  True
     __snake_case    					=  get_torch_dist_unique_port()
     __snake_case    					=  Path(_lowerCamelCase	).resolve().parent / '''distributed_scripts''' / '''launch_torch_distributed.py'''
     __snake_case    					=  f'''
        -m torch.distributed.launch
        --nproc_per_node={nproc_per_node}
        --master_port={master_port}
        {test_script}
    '''.split()
     __snake_case    					=  f'''
        --streaming={streaming}
        --num_workers={num_workers}
    '''.split()
     __snake_case    					=  [sys.executable] + distributed_args + args
     execute_subprocess_async(_lowerCamelCase   ,	env=os.environ.copy()	)








'''simple docstring'''



# 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 os
from typing import List







def 			_UpperCamelCase (_lowerCamelCase							:     List[str]   ,	_lowerCamelCase							:     str	)->   str:




     '''simple docstring'''





     if name != "training_function" and name != "main":
          raise ValueError(f'''Incorrect function name passed: {name}, choose either \'main\' or \'training_function\''''	)
     __snake_case					,    __snake_case    					=  [], False
     for line in lines:
          if not found_start and f'''def {name}''' in line:
               __snake_case    					=  True
               good_lines.append(_lowerCamelCase	)
               continue
          if found_start:
               if name == "training_function" and "def main" in line:
                    return good_lines
               if name == "main" and "if __name__" in line:
                    return good_lines
               good_lines.append(_lowerCamelCase	)







def 			_UpperCamelCase (_lowerCamelCase							:     List[str]	)->   int:




     '''simple docstring'''





     return [line for line in lines if not line.lstrip().startswith('''#'''	) and line != "\n"]







def 			_UpperCamelCase (_lowerCamelCase							:     str   ,	_lowerCamelCase							:     str   ,	_lowerCamelCase							:     bool   ,	_lowerCamelCase							:     str = None	)->   Optional[Any]:




     '''simple docstring'''





     with open(_lowerCamelCase   ,	'''r'''	) as f:
          __snake_case    					=  f.readlines()
     with open(os.path.abspath(os.path.join('''examples'''   ,	'''nlp_example.py'''	)	)   ,	'''r'''	) as f:
          __snake_case    					=  f.readlines()
     with open(_lowerCamelCase   ,	'''r'''	) as f:
          __snake_case    					=  f.readlines()
     if secondary_filename is not None:
          with open(_lowerCamelCase   ,	'''r'''	) as f:
               __snake_case    					=  f.readlines()

    # This is our base, we remove all the code from here in our `full_filename` and `feature_filename` to find the new content
     if parser_only:
          __snake_case    					=  clean_lines(get_function_contents_by_name(_lowerCamelCase   ,	'''main'''	)	)
          __snake_case    					=  clean_lines(get_function_contents_by_name(_lowerCamelCase   ,	'''main'''	)	)
          __snake_case    					=  clean_lines(get_function_contents_by_name(_lowerCamelCase   ,	'''main'''	)	)
          if secondary_filename is not None:
               __snake_case    					=  clean_lines(get_function_contents_by_name(_lowerCamelCase   ,	'''main'''	)	)
     else:
          __snake_case    					=  clean_lines(get_function_contents_by_name(_lowerCamelCase   ,	'''training_function'''	)	)
          __snake_case    					=  clean_lines(get_function_contents_by_name(_lowerCamelCase   ,	'''training_function'''	)	)
          __snake_case    					=  clean_lines(get_function_contents_by_name(_lowerCamelCase   ,	'''training_function'''	)	)
          if secondary_filename is not None:
               __snake_case    					=  clean_lines(
                   get_function_contents_by_name(_lowerCamelCase   ,	'''training_function'''	)	)

     __snake_case    					=  '''train_dataloader, eval_dataloader = get_dataloaders(accelerator, batch_size)\n'''

     # Specific code in our script that differs from the full version, aka what is new
     __snake_case    					=  []
     __snake_case    					=  []  # We keep track of the idxs just in case it's a repeated statement
     __snake_case    					=  iter(_lowerCamelCase	)
     for i in range(len(_lowerCamelCase	) - 1	):
          if i not in passed_idxs:
               __snake_case    					=  next(_lowerCamelCase	)
               if (line not in full_file_func) and (line.lstrip() != _dl_line):
                    if "TESTING_MOCKED_DATALOADERS" not in line:
                         new_feature_code.append(_lowerCamelCase	)
                         passed_idxs.append(_lowerCamelCase	)
                    else:
                         # Skip over the `config['num_epochs'] = 2` statement
                         __snake_case    					=  next(_lowerCamelCase	)

    # Extract out just the new parts from the full_file_training_func
     __snake_case    					=  []
     __snake_case    					=  []  # We keep track of the idxs just in case it's a repeated statement
     for i, line in enumerate(_lowerCamelCase	):
          if i not in passed_idxs:
               if (line not in full_file_func) and (line.lstrip() != _dl_line):
                    if "TESTING_MOCKED_DATALOADERS" not in line:
                         new_full_example_parts.append(_lowerCamelCase	)
                         passed_idxs.append(_lowerCamelCase	)

    # Finally, get the overall diff
     __snake_case    					=  [line for line in new_feature_code if line not in new_full_example_parts]
     if secondary_filename is not None:
          __snake_case    					=  [line for line in full_file_contents if line not in secondary_file_func]
          __snake_case    					=  [line for line in diff_from_example if line not in diff_from_two]

     return diff_from_example






