# Path: compiler_opt/rl/compilation_runner.py
# class RewardStat:
# class DataClassJSONEncoder(json.JSONEncoder):
# class ProcessKilledError(Exception):
# class ProcessCancellationToken:
# class WorkerCancellationManager:
# class CompilationRunner:
#   def default(self, o):
# def _postprocessing_sequence_example(
#     sequence_example: tf.train.SequenceExample, moving_average_reward: float,
#     policy_reward: float) -> tf.train.SequenceExample:
# def get_command_line_for_bundle(cmd_file: str,
#                                 ir_file: str,
#                                 thinlto: Optional[str] = None) -> List[str]:
#   def __init__(self):
#   def __init__(self):
#   def signal(self):
#   def wait(self):
# def kill_process_ignore_exceptions(p: subprocess.Popen):
#   def __init__(self):
#   def register_process(self, p: subprocess.Popen):
#   def signal(self):
#   def unregister_process(self, p: subprocess.Popen):
# def start_cancellable_process(
#     cmdline: List[str],
#     timeout: float,
#     cancellation_manager: Optional[WorkerCancellationManager],
#     want_output: bool = False) -> Optional[bytes]:
#   def init_pool():
#   def _get_pool():
#   def __init__(self,
#                clang_path: str,
#                llvm_size_path: str,
#                launcher_path: Optional[str] = None,
#                moving_average_decay_rate: float = 1):
#   def _get_cancellation_manager(
#       self, cancellation_token: Optional[ProcessCancellationToken]
#   ) -> Optional[WorkerCancellationManager]:
#     def signaler():
#   def collect_data(
#       self,
#       file_paths: Tuple[str, ...],
#       tf_policy_path: str,
#       reward_stat: Optional[Dict[str, RewardStat]],
#       cancellation_token: Optional[ProcessCancellationToken] = None
#   ) -> Tuple[List[str], Dict[str, RewardStat], List[float]]:
#   def _compile_fn(
#       self, file_paths: Tuple[str, ...], tf_policy_path: str, reward_only: bool,
#       cancellation_manager: Optional[WorkerCancellationManager]
#   ) -> Dict[str, Tuple[tf.train.SequenceExample, float]]:
#   _POOL: concurrent.futures.ThreadPoolExecutor = None
# 
# Path: compiler_opt/rl/data_collector.py
# DEADLINE_IN_SECONDS = 30
# WAIT_TERMINATION = ((0.9, 0), (0.8, 0.5), (0, 1))
# REWARD_QUANTILE_MONITOR = (0.1, 0.5, 1, 2, 3, 4, 5, 6, 8, 10, 20, 30, 40, 50,
#                            60, 70, 80, 90, 95, 99, 99.5, 99.9)
# def build_distribution_monitor(data: Sequence[float]) -> Dict[str, float]:
#   def collect_data(
#       self, policy_path: str
#   ) -> Tuple[Iterator[trajectory.Trajectory], Dict[str, float]]:
#   def on_dataset_consumed(self,
#                           dataset_iterator: Iterator[trajectory.Trajectory]):
#   def __init__(self,  # pylint: disable=dangerous-default-value
#                num_modules: int,
#                deadline: float = DEADLINE_IN_SECONDS,
#                thresholds: Tuple[Tuple[float, float], ...] = WAIT_TERMINATION):
#   def _should_exit(self, collected: int):
#   def wait(self, get_num_finished_work):
#   def waited_time(self):
# class DataCollector(metaclass=abc.ABCMeta):
# class EarlyExitChecker:
# 
# Path: compiler_opt/rl/local_data_collector.py
# class LocalDataCollector(data_collector.DataCollector):
#   def __init__(
#       self,
#       file_paths: Tuple[Tuple[str, ...], ...],
#       num_workers: int,
#       num_modules: int,
#       runner: compilation_runner.CompilationRunner,
#       parser: Callable[[List[str]], Iterator[trajectory.Trajectory]],
#       reward_stat_map: Dict[str, Optional[Dict[str,
#                                                compilation_runner.RewardStat]]],
#       exit_checker_ctor=data_collector.EarlyExitChecker):
#   def close_pool(self):
#   def _join_pending_jobs(self):
#   def _schedule_jobs(self, policy_path, sampled_file_paths):
#     def error_callback(e):
#   def collect_data(
#       self, policy_path: str
#   ) -> Tuple[Iterator[trajectory.Trajectory], Dict[str, Dict[str, float]]]:
#     def wait_for_termination():
#       def get_num_finished_work():
#   def on_dataset_consumed(self,
#                           dataset_iterator: Iterator[trajectory.Trajectory]):


import collections
import multiprocessing as mp
import subprocess
import tensorflow as tf
from unittest import mock
from tf_agents.system import system_multiprocessing as multiprocessing
from compiler_opt.rl import compilation_runner
from compiler_opt.rl import data_collector
from compiler_opt.rl import local_data_collector
        cancellation_manager = self._get_cancellation_manager(
            cancellation_token)
        try:
          compilation_runner.start_cancellable_process(['sleep', '3600s'], 3600,
                                                       cancellation_manager)
        except compilation_runner.ProcessKilledError as e:
          with self._lock:
            self._killed.value += 1
          raise e
        except subprocess.TimeoutExpired as e:
          with self._lock:
            self._timedout.value += 1
          raise e
        with self._lock:
          self._living.value += 1
        return [1], {}, [2.3]

    mock_compilation_runner = Sleeper(killed, timedout, living)

    def parser(_):
      pass

    class QuickExiter(data_collector.EarlyExitChecker):

      def __init__(self, num_modules):
        data_collector.EarlyExitChecker.__init__(self, num_modules=num_modules)

      def wait(self, _):
        return False


# Path: compiler_opt/rl/compilation_runner.py
# class RewardStat:
# class DataClassJSONEncoder(json.JSONEncoder):
# class ProcessKilledError(Exception):
# class ProcessCancellationToken:
# class WorkerCancellationManager:
# class CompilationRunner:
#   def default(self, o):
# def _postprocessing_sequence_example(
#     sequence_example: tf.train.SequenceExample, moving_average_reward: float,
#     policy_reward: float) -> tf.train.SequenceExample:
# def get_command_line_for_bundle(cmd_file: str,
#                                 ir_file: str,
#                                 thinlto: Optional[str] = None) -> List[str]:
#   def __init__(self):
#   def __init__(self):
#   def signal(self):
#   def wait(self):
# def kill_process_ignore_exceptions(p: subprocess.Popen):
#   def __init__(self):
#   def register_process(self, p: subprocess.Popen):
#   def signal(self):
#   def unregister_process(self, p: subprocess.Popen):
# def start_cancellable_process(
#     cmdline: List[str],
#     timeout: float,
#     cancellation_manager: Optional[WorkerCancellationManager],
#     want_output: bool = False) -> Optional[bytes]:
#   def init_pool():
#   def _get_pool():
#   def __init__(self,
#                clang_path: str,
#                llvm_size_path: str,
#                launcher_path: Optional[str] = None,
#                moving_average_decay_rate: float = 1):
#   def _get_cancellation_manager(
#       self, cancellation_token: Optional[ProcessCancellationToken]
#   ) -> Optional[WorkerCancellationManager]:
#     def signaler():
#   def collect_data(
#       self,
#       file_paths: Tuple[str, ...],
#       tf_policy_path: str,
#       reward_stat: Optional[Dict[str, RewardStat]],
#       cancellation_token: Optional[ProcessCancellationToken] = None
#   ) -> Tuple[List[str], Dict[str, RewardStat], List[float]]:
#   def _compile_fn(
#       self, file_paths: Tuple[str, ...], tf_policy_path: str, reward_only: bool,
#       cancellation_manager: Optional[WorkerCancellationManager]
#   ) -> Dict[str, Tuple[tf.train.SequenceExample, float]]:
#   _POOL: concurrent.futures.ThreadPoolExecutor = None
# 
# Path: compiler_opt/rl/data_collector.py
# DEADLINE_IN_SECONDS = 30
# WAIT_TERMINATION = ((0.9, 0), (0.8, 0.5), (0, 1))
# REWARD_QUANTILE_MONITOR = (0.1, 0.5, 1, 2, 3, 4, 5, 6, 8, 10, 20, 30, 40, 50,
#                            60, 70, 80, 90, 95, 99, 99.5, 99.9)
# def build_distribution_monitor(data: Sequence[float]) -> Dict[str, float]:
#   def collect_data(
#       self, policy_path: str
#   ) -> Tuple[Iterator[trajectory.Trajectory], Dict[str, float]]:
#   def on_dataset_consumed(self,
#                           dataset_iterator: Iterator[trajectory.Trajectory]):
#   def __init__(self,  # pylint: disable=dangerous-default-value
#                num_modules: int,
#                deadline: float = DEADLINE_IN_SECONDS,
#                thresholds: Tuple[Tuple[float, float], ...] = WAIT_TERMINATION):
#   def _should_exit(self, collected: int):
#   def wait(self, get_num_finished_work):
#   def waited_time(self):
# class DataCollector(metaclass=abc.ABCMeta):
# class EarlyExitChecker:
# 
# Path: compiler_opt/rl/local_data_collector.py
# class LocalDataCollector(data_collector.DataCollector):
#   def __init__(
#       self,
#       file_paths: Tuple[Tuple[str, ...], ...],
#       num_workers: int,
#       num_modules: int,
#       runner: compilation_runner.CompilationRunner,
#       parser: Callable[[List[str]], Iterator[trajectory.Trajectory]],
#       reward_stat_map: Dict[str, Optional[Dict[str,
#                                                compilation_runner.RewardStat]]],
#       exit_checker_ctor=data_collector.EarlyExitChecker):
#   def close_pool(self):
#   def _join_pending_jobs(self):
#   def _schedule_jobs(self, policy_path, sampled_file_paths):
#     def error_callback(e):
#   def collect_data(
#       self, policy_path: str
#   ) -> Tuple[Iterator[trajectory.Trajectory], Dict[str, Dict[str, float]]]:
#     def wait_for_termination():
#       def get_num_finished_work():
#   def on_dataset_consumed(self,
#                           dataset_iterator: Iterator[trajectory.Trajectory]):


# Path: compiler_opt/rl/local_data_collector_test.py
import collections
import multiprocessing as mp
import subprocess
import tensorflow as tf
from unittest import mock
from tf_agents.system import system_multiprocessing as multiprocessing
from compiler_opt.rl import compilation_runner
from compiler_opt.rl import data_collector
from compiler_opt.rl import local_data_collector

        cancellation_manager = self._get_cancellation_manager(
            cancellation_token)
        try:
          compilation_runner.start_cancellable_process(['sleep', '3600s'], 3600,
                                                       cancellation_manager)
        except compilation_runner.ProcessKilledError as e:
          with self._lock:
            self._killed.value += 1
          raise e
        except subprocess.TimeoutExpired as e:
          with self._lock:
            self._timedout.value += 1
          raise e
        with self._lock:
          self._living.value += 1
        return [1], {}, [2.3]

    mock_compilation_runner = Sleeper(killed, timedout, living)

    def parser(_):
      pass

    class QuickExiter(data_collector.EarlyExitChecker):

      def __init__(self, num_modules):
        data_collector.EarlyExitChecker.__init__(self, num_modules=num_modules)

      def wait(self, _):
        return False


    collector = local_data_collector.LocalDataCollector(