





'''simple docstring'''




import inspect
import os
import unittest
from pathlib import Path

import torch

import accelerate
from accelerate.test_utils import execute_subprocess_async
from accelerate.test_utils.testing import run_command






class UpperCAmelCase							(		unittest.TestCase							):
    _A :			Optional[Any]						   =       inspect.getfile(accelerate.test_utils							)
    _A :			int						   =       os.path.sep.join(mod_file.split(os.path.sep							)[:-1] + ["""scripts""", """test_cli.py"""]							)

    _A :			Tuple						   =       ["""accelerate""", """launch"""]
    _A :			Optional[int]						   =       Path.home() / """.cache/huggingface/accelerate"""
    _A :			Union[str, Any]						   =       """default_config.yaml"""
    _A :			Tuple						   =       config_folder / config_file
    _A :			str						   =       config_folder / """_default_config.yaml"""

    _A :			int						   =       Path("""tests/test_configs"""							)



    @classmethod
    def 					__lowerCamelCase						(  cls					):
     if cls.config_path.is_file():
      cls.config_path.rename(cls.changed_path					)



    @classmethod
    def 					__lowerCamelCase						(  cls					):
     if cls.changed_path.is_file():
      cls.changed_path.rename(cls.config_path					)



    def 					__lowerCamelCase						(  self					):
     __UpperCAmelCase					=   self.base_cmd
     if torch.cuda.is_available() and (torch.cuda.device_count() > 1):
      cmd += ["--multi_gpu"]
     execute_subprocess_async(cmd + [self.test_file_path]     ,			env=os.environ.copy()					)



    def 					__lowerCamelCase						(  self					):
     for config in sorted(self.test_config_path.glob('**/*.yaml'					)					):
      with self.subTest(config_file=__A					):
       execute_subprocess_async(
           self.base_cmd + ['--config_file', str(__A					), self.test_file_path]     ,			env=os.environ.copy()					)




    def 					__lowerCamelCase						(  self					):
     execute_subprocess_async(['accelerate', 'test']     ,			env=os.environ.copy()					)






class UpperCAmelCase							(		unittest.TestCase							):
    _A :			Tuple						   =       """test-tpu"""
    _A :			int						   =       """us-central1-a"""
    _A :			Dict						   =       """ls"""
    _A :			Optional[int]						   =       ["""accelerate""", """tpu-config"""]
    _A :			Dict						   =       """cd /usr/share"""
    _A :			Dict						   =       """tests/test_samples/test_command_file.sh"""
    _A :			Optional[Any]						   =       """Running gcloud compute tpus tpu-vm ssh"""



    def 					__lowerCamelCase						(  self					):
     __UpperCAmelCase					=   run_command(
         self.cmd
         + ['--command', self.command, '--tpu_zone', self.tpu_zone, '--tpu_name', self.tpu_name, '--debug']     ,			return_stdout=__A     ,			)
     self.assertIn(
         f'{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; ls --worker all'     ,			__A     ,			)



    def 					__lowerCamelCase						(  self					):
     __UpperCAmelCase					=   run_command(
         self.cmd
         + [
             '--config_file',
             'tests/test_configs/0_12_0.yaml',
             '--command',
             self.command,
             '--tpu_zone',
             self.tpu_zone,
             '--tpu_name',
             self.tpu_name,
             '--debug',
         ]     ,			return_stdout=__A     ,			)
     self.assertIn(
         f'{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; ls --worker all'     ,			__A     ,			)



    def 					__lowerCamelCase						(  self					):
     __UpperCAmelCase					=   run_command(
         self.cmd + ['--config_file', 'tests/test_configs/latest.yaml', '--debug']     ,			return_stdout=__A					)
     self.assertIn(
         f'{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; echo "hello world"; echo "this is a second command" --worker all'     ,			__A     ,			)



    def 					__lowerCamelCase						(  self					):
     __UpperCAmelCase					=   run_command(
         self.cmd + ['--config_file', 'tests/test_configs/latest.yaml', '--command', self.command, '--debug']     ,			return_stdout=__A     ,			)
     self.assertIn(
         f'{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; ls --worker all'     ,			__A     ,			)



    def 					__lowerCamelCase						(  self					):
     __UpperCAmelCase					=   run_command(
         self.cmd
         + [
             '--config_file',
             'tests/test_configs/latest.yaml',
             '--command',
             self.command,
             '--command',
             'echo "Hello World"',
             '--debug',
         ]     ,			return_stdout=__A     ,			)
     self.assertIn(
         f'{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; ls; echo "Hello World" --worker all'     ,			__A     ,			)



    def 					__lowerCamelCase						(  self					):
     __UpperCAmelCase					=   run_command(
         self.cmd
         + ['--config_file', 'tests/test_configs/latest.yaml', '--command_file', self.command_file, '--debug']     ,			return_stdout=__A     ,			)
     self.assertIn(
         f'{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; echo "hello world"; echo "this is a second command" --worker all'     ,			__A     ,			)



    def 					__lowerCamelCase						(  self					):
     __UpperCAmelCase					=   run_command(
         self.cmd
         + [
             '--config_file',
             'tests/test_configs/0_12_0.yaml',
             '--command_file',
             self.command_file,
             '--tpu_zone',
             self.tpu_zone,
             '--tpu_name',
             self.tpu_name,
             '--debug',
         ]     ,			return_stdout=__A     ,			)
     self.assertIn(
         f'{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; echo "hello world"; echo "this is a second command" --worker all'     ,			__A     ,			)



    def 					__lowerCamelCase						(  self					):
     __UpperCAmelCase					=   run_command(
         self.cmd + ['--config_file', 'tests/test_configs/latest.yaml', '--install_accelerate', '--debug']     ,			return_stdout=__A     ,			)
     self.assertIn(
         f'{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; pip install accelerate -U; echo "hello world"; echo "this is a second command" --worker all'     ,			__A     ,			)




    def 					__lowerCamelCase						(  self					):
     __UpperCAmelCase					=   run_command(
         self.cmd
         + [
             '--config_file',
             'tests/test_configs/latest.yaml',
             '--install_accelerate',
             '--accelerate_version',
             '12.0.0',
             '--debug',
         ]     ,			return_stdout=__A     ,			)
     self.assertIn(
         f'{self.gcloud} test-tpu --zone us-central1-a --command {self.base_output}; pip install accelerate==12.0.0; echo "hello world"; echo "this is a second command" --worker all'     ,			__A     ,			)












'''simple docstring'''




import multiprocessing
from typing import TYPE_CHECKING, Optional, Union

from .. import Dataset, Features, config
from ..formatting import query_table
from ..packaged_modules.sql.sql import Sql
from ..utils import logging
from .abc import AbstractDatasetInputStream


if TYPE_CHECKING:
    import sqlitea

    import sqlalchemy






class UpperCAmelCase							(		UpperCAmelCase_							):



    def __init__(  self     ,			__A     ,			__A     ,			__A = None     ,			__A = None     ,			__A = False     ,			**__A     ,			):
     super().__init__(features=__A     ,			cache_dir=__A     ,			keep_in_memory=__A     ,			**__A					)
     __UpperCAmelCase					=   Sql(
         cache_dir=__A     ,			features=__A     ,			sql=__A     ,			con=__A     ,			**__A     ,			)




    def 					__lowerCamelCase						(  self					):
     __UpperCAmelCase					=   None
     __UpperCAmelCase					=   None
     __UpperCAmelCase					=   None
     __UpperCAmelCase					=   None

     self.builder.download_and_prepare(
         download_config=__A     ,			download_mode=__A     ,			verification_mode=__A     ,			base_path=__A     ,			)

     # Build dataset for splits
     __UpperCAmelCase					=   self.builder.as_dataset(
         split='train'     ,			verification_mode=__A     ,			in_memory=self.keep_in_memory					)
     return dataset






class UpperCAmelCase							:



    def __init__(  self     ,			__A     ,			__A     ,			__A     ,			__A = None     ,			__A = None     ,			**__A     ,			):
     if num_proc is not None and num_proc <= 0:
      raise ValueError(f'num_proc {num_proc} must be an integer > 0.'					)

     __UpperCAmelCase					=   dataset
     __UpperCAmelCase					=   name
     __UpperCAmelCase					=   con
     __UpperCAmelCase					=   batch_size if batch_size else config.DEFAULT_MAX_BATCH_SIZE
     __UpperCAmelCase					=   num_proc
     __UpperCAmelCase					=   to_sql_kwargs



    def 					__lowerCamelCase						(  self					):
     __UpperCAmelCase					=   self.to_sql_kwargs.pop('sql'     ,			__A					)
     __UpperCAmelCase					=   self.to_sql_kwargs.pop('con'     ,			__A					)
     __UpperCAmelCase					=   self.to_sql_kwargs.pop('index'     ,			__A					)

     __UpperCAmelCase					=   self._write(index=__A     ,			**self.to_sql_kwargs					)
     return written



    def 					__lowerCamelCase						(  self     ,			__A					):
     __UpperCAmelCase					,    __UpperCAmelCase					,    __UpperCAmelCase					=   args
     __UpperCAmelCase					=   {**to_sql_kwargs, 'if_exists': 'append'} if offset > 0 else to_sql_kwargs
     __UpperCAmelCase					=   query_table(
         table=self.dataset.data     ,			key=slice(__A     ,			offset + self.batch_size					)     ,			indices=self.dataset._indices     ,			)
     __UpperCAmelCase					=   batch.to_pandas()
     __UpperCAmelCase					=   df.to_sql(self.name     ,			self.con     ,			index=__A     ,			**__A					)
     return num_rows or len(__A					)




    def 					__lowerCamelCase						(  self     ,			__A     ,			**__A					):
     __UpperCAmelCase					=   0

     if self.num_proc is None or self.num_proc == 1:
      for offset in logging.tqdm(
          range(0     ,			len(self.dataset					)     ,			self.batch_size					)     ,			unit='ba'     ,			disable=not logging.is_progress_bar_enabled()     ,			desc='Creating SQL from Arrow format'     ,			):
       written += self._batch_sql((offset, index, to_sql_kwargs)					)
     else:
      __UpperCAmelCase					,    __UpperCAmelCase					=   len(self.dataset					), self.batch_size
      with multiprocessing.Pool(self.num_proc					) as pool:
       for num_rows in logging.tqdm(
           pool.imap(
               self._batch_sql     ,			[(offset, index, to_sql_kwargs) for offset in range(0     ,			__A     ,			__A					)]     ,			)     ,			total=(num_rows // batch_size) + 1 if num_rows % batch_size else num_rows // batch_size     ,			unit='ba'     ,			disable=not logging.is_progress_bar_enabled()     ,			desc='Creating SQL from Arrow format'     ,			):
        written += num_rows

     return written





