SprintDataset.py

"""
Implements the SprintDatasetBase and ExternSprintDataset classes, some Dataset subtypes.
Note that from the main RETURNN process, you probably want ExternSprintDataset instead.
"""

from __future__ import print_function

import atexit
import os
import signal
import sys
try:
  import thread
except ImportError:
  import _thread as thread
from threading import Condition, currentThread, Thread
import time
import numpy
import typing

import TaskSystem
from Dataset import Dataset, DatasetSeq
from CachedDataset2 import CachedDataset2
from Log import log
from TaskSystem import Unpickler, numpy_copy_and_set_unused
from Util import eval_shell_str, interrupt_main, unicode, PY3, BytesIO


class SprintDatasetBase(Dataset):
  """
  In Sprint, we use this object for multiple purposes:
  - Multiple epoch handling via SprintInterface.getSegmentList().
    For this, we get the segment list from Sprint and use the Dataset
    shuffling method.
  - Fill in data which we get via SprintInterface.feedInput*().
    Note that each such input doesn't necessarily correspond to a single
    segment. This depends which type of FeatureExtractor is used in Sprint.
    If we use the BufferedFeatureExtractor in utterance mode, we will get
    one call for every segment and we get also segmentName as parameter.
    Otherwise, we will get batches of fixed size - in that case,
    it doesn't correspond to the segments.
    In any case, we use this data as-is as a new seq.
    Because of that, we cannot really know the number of seqs in advance,
    nor the total number of time frames, etc.

  If you want to use this directly in RETURNN, see ExternSprintDataset.
  """

  SprintCachedSeqsMax = 200
  SprintCachedSeqsMin = 100

  def __init__(self, target_maps=None, str_add_final_zero=False, input_stddev=1.,
               orth_post_process=None, bpe=None, orth_vocab=None,
               suppress_load_seqs_print=False,
               **kwargs):
    """
    :param dict[str,str|dict] target_maps: e.g. {"speaker": "speaker_map.txt"}
    :param bool str_add_final_zero: adds e.g. "orth0" with '\0'-ending
    :param float input_stddev: if != 1, will divide the input "data" by that
    :param str|list[str]|((str)->str)|None orth_post_process: :func:`get_post_processor_function`, applied on orth
    :param None|dict[str] bpe: if given, will be opts for :class:`BytePairEncoding`
    :param None|dict[str] orth_vocab: if given, orth_vocab is applied to orth and orth_classes is an available target`
    :param bool suppress_load_seqs_print: less verbose
    """
    super(SprintDatasetBase, self).__init__(**kwargs)
    self.suppress_load_seqs_print = suppress_load_seqs_print
    if target_maps:
      assert isinstance(target_maps, dict)
      target_maps = target_maps.copy()
      for key, tmap in list(target_maps.items()):
        if isinstance(tmap, (str, unicode)):
          tmap = {l: i for (i, l) in enumerate(open(tmap).read().splitlines())}
        assert isinstance(tmap, dict)  # dict[str,int]
        target_maps[key] = tmap
    self.target_maps = target_maps
    self.str_add_final_zero = str_add_final_zero
    self.input_stddev = input_stddev
    # Note: "orth" is actually the raw bytes of the utf8 string,
    # so it does not make quite sense to associate a single str to each byte.
    # However, some other code might expect that the labels are all strings, not bytes,
    # and the API requires the labels to be strings.
    # The code in Dataset.serialize_data tries to decode this case as utf8 (if possible).
    self.labels["orth"] = [chr(i) for i in range(255)]
    self.orth_post_process = None  # type: typing.Optional[typing.Callable[[str],str]]
    if orth_post_process:
      if callable(orth_post_process):
        self.orth_post_process = orth_post_process
      else:
        from LmDataset import get_post_processor_function
        self.orth_post_process = get_post_processor_function(orth_post_process)
    self.bpe = None
    if bpe:
      from GeneratingDataset import BytePairEncoding
      self.bpe = BytePairEncoding(**bpe)
      self.labels["bpe"] = self.bpe.labels
    self.orth_vocab = None
    if orth_vocab:
      from GeneratingDataset import Vocabulary
      self.orth_vocab = Vocabulary.create_vocab(**orth_vocab)
      self.labels["orth_classes"] = self.orth_vocab.labels
    self.cond = Condition(lock=self.lock)
    self.add_data_thread_id = thread.get_ident()  # This will be created in the Sprint thread.
    self.ready_for_data = False
    self.reached_final_seq = False
    self.reached_final_seq_seen_all = False
    self.multiple_epochs = False
    self._complete_frac = None
    self.sprintEpoch = None  # in SprintInterface.getSegmentList()
    self.crnnEpoch = None  # in CRNN train thread, Engine.train(). set via init_seq_order
    self.predefined_seq_list_order = None  # via init_seq_order
    self.sprintFinalized = False
    self._target_black_list = []  # if we get non numpy arrays and cannot convert them
    self._reset_cache()
    assert self.shuffle_frames_of_nseqs == 0  # Currently broken. But just use Sprint itself to do this.

  def use_multiple_epochs(self):
    """
    Called via SprintInterface.getSegmentList().
    """
    self.multiple_epochs = True

  def set_dimensions(self, input_dim, output_dim):
    """
    :type input_dim: int
    :type output_dim: int

    Called via python_train.
    """
    assert input_dim > 0
    self.num_inputs = input_dim
    self.num_outputs = {"data": (input_dim * self.window, 2)}
    if output_dim > 0:
      self.num_outputs["classes"] = (output_dim, 1)
    if self.bpe:
      self.num_outputs["bpe"] = (self.bpe.num_labels, 1)
    if self.orth_vocab:
      self.num_outputs["orth_classes"] = (self.orth_vocab.num_labels, 1)
    self.num_outputs["orth"] = (256, 1)
    self._base_init()
    # At this point, we are ready for data. In case we don't use the Sprint PythonSegmentOrdering
    # (SprintInterface.getSegmentList()), we must call this at least once.
    if not self.multiple_epochs:
      self.init_sprint_epoch(None)

  def _reset_cache(self):
    self.expected_load_seq_start = 0
    self.requested_load_seq_end = 0
    self.next_seq_to_be_added = 0
    self.reached_final_seq = False
    self.reached_final_seq_seen_all = False
    self._num_timesteps = 0
    self.added_data = []  # type: typing.List[DatasetSeq]
    self.ready_for_data = True

  def init_sprint_epoch(self, epoch):
    """
    :type epoch: int | None
    Called by SprintInterface.getSegmentList() when we start a new epoch.
    We must not call this via self.init_seq_order() because we will already have filled the cache by Sprint
    before the CRNN train thread starts the epoch.
    """
    with self.lock:
      self.sprintEpoch = epoch
      self.sprintFinalized = False
      self._reset_cache()
      self.cond.notify_all()

  def finalize_sprint(self):
    """
    Called when SprintInterface.getSegmentList() ends.
    """
    with self.lock:
      self.sprintFinalized = True
      self.cond.notify_all()

  def init_seq_order(self, epoch=None, seq_list=None):
    """
    Called by CRNN train thread when we enter a new epoch.
    """
    super(SprintDatasetBase, self).init_seq_order(epoch=epoch, seq_list=seq_list)
    with self.lock:
      self.crnnEpoch = epoch
      self.predefined_seq_list_order = seq_list
      self.cond.notify_all()
      # No need to wait/check for Sprint thread here.
      # SprintInterface.getSegmentList() will wait for us.
    return True

  def _cleanup_old_seq_cache(self, seq_end):
    i = 0
    while i < len(self.added_data):
      if self.added_data[i].seq_idx >= seq_end:
        break
      i += 1
    del self.added_data[:i]

  def wait_for_returnn_epoch(self, epoch):
    """
    Called by SprintInterface.
    """
    with self.lock:
      while epoch != self.crnnEpoch:
        assert epoch > self.crnnEpoch
        self.cond.wait()

  def _wait_for_seq_can_pass_check(self, seq_start, seq_end):
    """
    :param int seq_start:
    :param int seq_end:
    :return: True if _waitForSeq can pass/return. False means that we need to wait more (until next signal)
    :rtype: bool
    """
    if self.reached_final_seq:
      return True
    if self._have_seqs_added(seq_start, seq_end):
      return True
    return False

  def _wait_for_seq(self, seq_start, seq_end=None):
    """
    Called by RETURNN train thread.
    Wait until we have seqs [seqStart,..,seqEnd-1] loaded,
    or we now that they will not be loaded anymore because we reached the end.

    :param int seq_start:
    :param int|None seq_end:
    """
    if seq_end is None:
      seq_end = seq_start + 1
    if seq_end > self.requested_load_seq_end:
      self.requested_load_seq_end = seq_end
      self.cond.notify_all()
    if self._wait_for_seq_can_pass_check(seq_start=seq_start, seq_end=seq_end):
      return
    # We need to wait.
    assert thread.get_ident() != self.add_data_thread_id
    if not self.suppress_load_seqs_print:
      print(
        "%s %s: wait for seqs (%i,%i) (last added: %s) (current time: %s)" % (
          self, currentThread().name, seq_start, seq_end, self._latest_added_seq(), time.strftime("%H:%M:%S")),
        file=log.v5)
    while not self._wait_for_seq_can_pass_check(seq_start=seq_start, seq_end=seq_end):
      self.cond.wait()

  def _latest_added_seq(self):
    if not self.added_data:
      return None
    return self.added_data[-1].seq_idx

  def _have_seqs_added(self, start, end=None):
    if end is None:
      end = start + 1
    if start >= end:
      return True
    for data in self.added_data:
      assert start >= data.seq_idx, "%s: We expect that we only ask about the cache of the upcoming seqs." % self
      if data.seq_idx == start:
        start += 1
      if start >= end:
        return True
    return False

  def _get_seq(self, seq_idx):
    """
    :param int seq_idx:
    :rtype: DatasetSeq
    """
    for data in self.added_data:
      if data.seq_idx == seq_idx:
        return data
    return None

  def is_cached(self, start, end):
    """
    :param int start:
    :param int end:
    :rtype: bool
    """
    # Always False, to force that we call self._load_seqs().
    # This is important for our buffer management.
    return False

  def load_seqs(self, start, end):
    """
    Called by RETURNN train thread.

    :param int start:
    :param int end:
    """
    if start == end:
      return
    if not self.suppress_load_seqs_print:
      print("%s load_seqs in %s:" % (self, currentThread().name), start, end, end=' ', file=log.v5)
    with self.lock:
      super(SprintDatasetBase, self).load_seqs(start, end)
      if not self.suppress_load_seqs_print:
        print("first features shape:", self._get_seq(start).features["data"].shape, file=log.v5)

  def _load_seqs(self, start, end):
    """
    Called by RETURNN train thread.
    We expect that start increase monotonic on each call
    for not-yet-loaded data.
    This will already be called with _load_seqs_superset indices.

    :param int start:
    :param int end:
    """
    assert start >= self.expected_load_seq_start
    if start > self.expected_load_seq_start:
      # Cleanup old data.
      self._cleanup_old_seq_cache(start)
      self.expected_load_seq_start = start
      self.cond.notify_all()
    self._wait_for_seq(start, end)

  def add_new_data(self, features, targets=None, segment_name=None):
    """
    Adds a new seq.
    This is called via the Sprint main thread.

    :param numpy.ndarray features: format (input-feature,time) (via Sprint)
    :param dict[str,numpy.ndarray|str] targets: format (time) (idx of output-feature)
    :param str|None segment_name:
    :returns the sorted seq index
    :rtype: int
    """

    # is in format (feature,time)
    assert self.num_inputs == features.shape[0]
    num_frames = features.shape[1]
    # must be in format: (time,feature)
    features = features.transpose()
    assert features.shape == (num_frames, self.num_inputs)
    if self.input_stddev != 1:
      features /= self.input_stddev
    if self.window > 1:
      features = self._sliding_window(features)
      assert features.shape == (num_frames, self.num_inputs * self.window)

    if targets is None:
      targets = {}
    if not isinstance(targets, dict):
      targets = {"classes": targets}
    if "classes" in targets:
      # 'classes' is always the alignment
      assert targets["classes"].shape == (num_frames,), (  # is in format (time,)
        "Number of targets %s does not equal to number of features %s" % (targets["classes"].shape, (num_frames,)))
    if "orth" in targets:
      targets["orth"] = targets["orth"].decode("utf8").strip()
    if "orth" in targets and self.orth_post_process:
      targets["orth"] = self.orth_post_process(targets["orth"])
    if self.bpe:
      assert "orth" in targets
      orth = targets["orth"]
      assert isinstance(orth, (str, unicode))
      assert "bpe" not in targets
      targets["bpe"] = numpy.array(self.bpe.get_seq(orth), dtype="int32")
    if self.orth_vocab:
      assert "orth" in targets
      orth = targets["orth"]
      assert isinstance(orth, (str, unicode))
      assert "orth_classes" not in targets
      targets["orth_classes"] = numpy.array(self.orth_vocab.get_seq(orth), dtype="int32")

    # Maybe convert some targets.
    if self.target_maps:
      for key, target_map in self.target_maps.items():
        assert key in targets
        v = target_map[targets[key]]
        v = numpy.asarray(v)
        if v.ndim == 0:
          v = numpy.zeros((num_frames,), dtype=v.dtype) + v  # add time dimension
        targets[key] = v

    # Maybe remove some targets.
    for key in self._target_black_list:
      if key in targets:
        del targets[key]

    # Check if all targets are valid.
    for key, v in sorted(list(targets.items())):
      if isinstance(v, numpy.ndarray):
        continue  # ok
      if isinstance(v, unicode):
        v = v.encode("utf8")
      if isinstance(v, (str, bytes)):
        if PY3:
          assert isinstance(v, bytes)
          v = list(v)
        else:
          v = list(map(ord, v))
        v = numpy.array(v, dtype="uint8")
        targets[key] = v
        if self.str_add_final_zero:
          v = numpy.append(v, numpy.array([0], dtype=v.dtype))
          assert key + "0" not in targets
          targets[key + "0"] = v
        continue
      print("%s, we will ignore the target %r because it is not a numpy array: %r" % (self, key, v), file=log.v3)
      self._target_black_list += [key]
      del targets[key]

    with self.lock:
      # This gets called in the Sprint main thread.
      # If this is used together with SprintInterface.getSegmentList(), we are always in a state where
      # we just yielded a segment name, thus we are always in a Sprint epoch and thus ready for data.
      assert self.ready_for_data
      assert not self.reached_final_seq
      assert not self.sprintFinalized

      seq_idx = self.next_seq_to_be_added

      if self.predefined_seq_list_order:
        # Note: Only in ExternSprintDataset, we can reliably set the seq order for now.
        assert seq_idx < len(self.predefined_seq_list_order), "seq_idx %i, expected predef num seqs %i" % (
          seq_idx, len(self.predefined_seq_list_order))
        expected_seq_name = self.predefined_seq_list_order[seq_idx]
        if expected_seq_name != segment_name:
          if segment_name in self.predefined_seq_list_order:
            raise Exception("seq_idx %i expected to be tag %r but got tag %r; tag %r is at idx %i" % (
              seq_idx, expected_seq_name, segment_name, segment_name,
              self.predefined_seq_list_order.index(segment_name)))
          raise Exception("seq_idx %i expected to be tag %r but got tag %r; tag %r not found" % (
            seq_idx, expected_seq_name, segment_name, segment_name))

      self.next_seq_to_be_added += 1
      self._num_timesteps += num_frames
      self.cond.notify_all()

      if seq_idx > self.requested_load_seq_end - 1 + self.SprintCachedSeqsMax:
        print("%s add_new_data: seq=%i, len=%i. Cache filled, waiting to get loaded..." % (
          self, seq_idx, num_frames), file=log.v5)
        while seq_idx > self.requested_load_seq_end - 1 + self.SprintCachedSeqsMin:
          assert not self.reached_final_seq
          assert seq_idx + 1 == self.next_seq_to_be_added
          self.cond.wait()

      self.added_data += [DatasetSeq(seq_idx, features, targets, seq_tag=segment_name)]
      self.cond.notify_all()
      return seq_idx

  def finish_sprint_epoch(self, seen_all=True):
    """
    Called by SprintInterface.getSegmentList().
    This is in a state where Sprint asks for the next segment after we just finished an epoch.
    Thus, any upcoming self.add_new_data() call will contain data from a segment in the new epoch.
    Thus, we finish the current epoch in Sprint.
    """
    with self.lock:
      self.reached_final_seq = True
      self.reached_final_seq_seen_all = seen_all
      self.ready_for_data = False
      self.cond.notify_all()

  def _shuffle_frames_in_seqs(self, start, end):
    assert False, "Shuffling in SprintDataset only via Sprint at the moment."

  def get_num_timesteps(self):
    """
    :rtype: int
    """
    with self.lock:
      assert self.reached_final_seq
      return self._num_timesteps

  @property
  def num_seqs(self):
    """
    :rtype: int
    """
    with self.lock:
      if self.predefined_seq_list_order:
        return len(self.predefined_seq_list_order)
      if not self.reached_final_seq:
        raise NotImplementedError
      return self.next_seq_to_be_added

  def have_seqs(self):
    """
    :rtype: bool
    """
    with self.lock:
      if self.next_seq_to_be_added > 0:
        return True
      self._wait_for_seq(0)
      return self.next_seq_to_be_added > 0

  def is_less_than_num_seqs(self, n):
    """
    :param int n:
    :rtype: bool
    """
    with self.lock:
      self._wait_for_seq(n)
      return n < self.next_seq_to_be_added

  def get_data_keys(self):
    """
    :rtype: list[str]
    """
    with self.lock:
      if not self.added_data:
        self._wait_for_seq(0)
      assert self.added_data
      return sorted(self.added_data[0].features.keys())

  def get_target_list(self):
    """
    :rtype: list[str]
    """
    keys = list(self.get_data_keys())
    if "data" in keys:
      keys.remove("data")
    return keys

  def set_complete_frac(self, frac):
    """
    :param float frac:
    """
    self._complete_frac = frac

  def get_complete_frac(self, seq_idx):
    """
    :param int seq_idx:
    :rtype: float
    """
    with self.lock:
      if self.predefined_seq_list_order:
        return float(seq_idx + 1) / len(self.predefined_seq_list_order)
      elif self._complete_frac is not None:
        if not self.next_seq_to_be_added:
          return self._complete_frac
        else:
          # We can do somewhat better. self._complete_frac is for self.next_seq_to_be_added.
          return self._complete_frac * float(seq_idx + 1) / self.next_seq_to_be_added
      else:
        return super(SprintDatasetBase, self).get_complete_frac(seq_idx)

  def get_seq_length(self, sorted_seq_idx):
    """
    :param int sorted_seq_idx:
    :rtype: Util.NumbersDict
    """
    with self.lock:
      self._wait_for_seq(sorted_seq_idx)
      return self._get_seq(sorted_seq_idx).num_frames

  def get_data(self, seq_idx, key):
    """
    :param int seq_idx:
    :param str key:
    :rtype: numpy.ndarray
    """
    with self.lock:
      self._wait_for_seq(seq_idx)
      return self._get_seq(seq_idx).features[key]

  def get_input_data(self, sorted_seq_idx):
    """
    :param int sorted_seq_idx:
    :rtype: numpy.ndarray
    """
    with self.lock:
      self._wait_for_seq(sorted_seq_idx)
      return self._get_seq(sorted_seq_idx).features["data"]

  def get_targets(self, target, sorted_seq_idx):
    """
    :param str target:
    :param int sorted_seq_idx:
    :rtype: numpy.ndarray
    """
    with self.lock:
      self._wait_for_seq(sorted_seq_idx)
      return self._get_seq(sorted_seq_idx).features.get(target, None)

  def get_ctc_targets(self, sorted_seq_idx):
    """
    :param int sorted_seq_idx:
    """
    assert False, "No CTC targets."

  def get_tag(self, sorted_seq_idx):
    """
    :param int sorted_seq_idx:
    :rtype: str
    """
    with self.lock:
      self._wait_for_seq(sorted_seq_idx)
      return self._get_seq(sorted_seq_idx).seq_tag


class ExternSprintDataset(SprintDatasetBase):
  """
  This is a Dataset which you can use directly in RETURNN.
  You can use it to get any type of data from Sprint (RWTH ASR toolkit),
  e.g. you can use Sprint to do feature extraction and preprocessing.

  This class is like SprintDatasetBase, except that we will start an external Sprint instance ourselves
  which will forward the data to us over a pipe.
  The Sprint subprocess will use SprintExternInterface to communicate with us.
  """

  # Do not change the argument names here, to not break existing configs.
  # noinspection PyPep8Naming
  def __init__(self, sprintTrainerExecPath, sprintConfigStr, partitionEpoch=None, **kwargs):
    """
    :param str|list[str] sprintTrainerExecPath:
    :param str | list[str] | ()->str | list[()->str] | ()->list[str] | ()->list[()->str] sprintConfigStr:
      via eval_shell_str
    :param int|None partitionEpoch: deprecated. use partition_epoch instead
    """
    super(ExternSprintDataset, self).__init__(**kwargs)
    self.add_data_thread_id = None
    self.sprint_trainer_exec_path = sprintTrainerExecPath
    self.sprint_config = sprintConfigStr
    if partitionEpoch:
      assert self.partition_epoch == 1, "don't provide partitionEpoch and partition_epoch"
      self.partition_epoch = partitionEpoch
    self._num_seqs = None
    self.child_pid = None  # type: typing.Optional[int]
    self.parent_pid = os.getpid()
    self.reader_thread = None  # type: typing.Optional[Thread]
    self.seq_list_file = None
    self.use_multiple_epochs()
    # There is no generic way to see whether Python is exiting.
    # This is our workaround. We check for it in self.run_inner().
    self.python_exit = False
    atexit.register(self._exit_handler)
    # We don't know about num_outputs yet, but we should.
    # Thus we call Sprint and immediately exit it.
    self._start_child(epoch=None, get_dim_only=True)

  def finish_epoch(self):
    """
    Called at the end of the epoch.
    """
    with self.lock:
      # Reset epoch such that exiting the child will go smoothly.
      super(ExternSprintDataset, self).init_seq_order(epoch=None, seq_list=None)
    # Exit child, before we overwrite anything, such as new epoch or seq_list.
    self._exit_child(wait_thread=True)
    super(ExternSprintDataset, self).finish_epoch()

  def _exit_handler(self):
    """
    Called at exit.
    """
    assert os.getpid() == self.parent_pid
    self.python_exit = True
    self._exit_child(wait_thread=False)

  def _exit_child(self, wait_thread=True):
    """
    :param bool wait_thread:
    """
    if self.child_pid:
      expected_exit_status = 0 if wait_thread and not self.python_exit else None
      if self._join_child(wait=False, expected_exit_status=expected_exit_status) is False:  # Not yet terminated.
        interrupt = not self.reached_final_seq_seen_all or not wait_thread
        if interrupt:
          print("%s: interrupt child proc %s" % (self, self.child_pid), file=log.v5)
          os.kill(self.child_pid, signal.SIGKILL)
          # Also join such that the process is cleaned up, and pipes get closed.
          self._join_child(wait=True, expected_exit_status=None)
          self.child_pid = None
      else:  # child process terminated
        self.child_pid = None
      if wait_thread and self.reader_thread:
        # Load all remaining data so that the reader thread is not waiting in self.add_new_data().
        while self.is_less_than_num_seqs(self.expected_load_seq_start + 1):
          if self.reached_final_seq:  # this is set by the reader thread
            break
          self.load_seqs(self.expected_load_seq_start + 1, self.expected_load_seq_start + 2)
        self.reader_thread.join()
        self.reader_thread = None
      try:
        self.pipe_p2c[1].close()
      except IOError:
        pass
      try:
        self.pipe_c2p[0].close()
      except IOError:
        pass
      if self.child_pid:
        self._join_child(wait=True, expected_exit_status=expected_exit_status)
        self.child_pid = None

  def _start_child(self, epoch, get_dim_only=False):
    """
    :param int|None epoch:
    :param bool get_dim_only:
    """
    assert self.child_pid is None
    assert self.reader_thread is None
    self.pipe_c2p = self._pipe_open()
    self.pipe_p2c = self._pipe_open()
    args = self._build_sprint_args()
    print("%s: epoch" % self, epoch, "exec", args, file=log.v5)

    pid = os.fork()
    if pid == 0:  # child
      # In case we are in some test environment or so, recover the original stdout/stderr.
      sys.stdin = sys.__stdin__
      sys.stdout = sys.__stdout__
      sys.stderr = sys.__stderr__
      import better_exchook
      better_exchook.install()
      # noinspection PyBroadException
      try:
        sys.stdin.close()  # Force no tty stdin.
        self.pipe_c2p[0].close()
        self.pipe_p2c[1].close()
        os.execv(args[0], args)  # Does not return if successful.
        print("%s child exec failed." % self)
      except BaseException:
        print("%s child: Error when starting Sprint %r." % (self, args))
        sys.excepthook(*sys.exc_info())
      finally:
        print("%s child: exit" % self)
        # noinspection PyProtectedMember
        os._exit(1)
        return  # Not reached.

    # parent
    self.pipe_c2p[1].close()
    self.pipe_p2c[0].close()
    self.child_pid = pid

    try:
      init_signal, (input_dim, output_dim, num_segments) = self._read_next_raw()
      assert init_signal == b"init"
      assert isinstance(input_dim, int) and isinstance(output_dim, int)
      # Ignore num_segments. It can be totally different than the real number of sequences.
      self.set_dimensions(input_dim, output_dim)
    except Exception:
      print("%s: Sprint child process (%r) caused an exception." % (self, args), file=log.v1)
      sys.excepthook(*sys.exc_info())
      self._exit_child(wait_thread=False)
      raise Exception("%s Sprint init failed" % self)

    if get_dim_only:
      self._exit_child(wait_thread=False)

    else:
      self.reader_thread = Thread(target=self._reader_thread_proc, args=(pid, epoch),
                                  name="%s reader thread" % self)
      self.reader_thread.daemon = True
      self.reader_thread.start()

  # noinspection PyMethodMayBeStatic
  def _pipe_open(self):
    readend, writeend = os.pipe()
    if hasattr(os, "set_inheritable"):
      # Python 3 by default will close all fds in subprocesses. This will avoid that.
      os.set_inheritable(readend, True)
      os.set_inheritable(writeend, True)
    readend = os.fdopen(readend, "rb", 0)
    writeend = os.fdopen(writeend, "wb", 0)
    return readend, writeend

  @property
  def _my_python_mod_path(self):
    """
    :rtype: str
    """
    return os.path.dirname(os.path.abspath(__file__))

  def _build_sprint_args(self):
    """
    :rtype: list[str]
    """
    config_str = "action:ExternSprintDataset,c2p_fd:%i,p2c_fd:%i" % (
      self.pipe_c2p[1].fileno(), self.pipe_p2c[0].fileno())
    if TaskSystem.SharedMemNumpyConfig["enabled"]:
      config_str += ",EnableAutoNumpySharedMemPickling:True"
    epoch = self.crnnEpoch or 1
    assert epoch >= 1
    if isinstance(self.sprint_trainer_exec_path, (list, tuple)):
      args = list(self.sprint_trainer_exec_path)
    else:
      args = [self.sprint_trainer_exec_path]
    # First the user options. Usually also involves loading some config.
    args += eval_shell_str(self.sprint_config)
    # Now our options. They might overwrite some of the config settings. (That is why we do it after the user opts.)
    args += [
      "--*.seed=%i" % (self._get_random_seed_for_epoch(epoch=epoch) - 1)]
    if self.partition_epoch > 1:
      args += [
        "--*.corpus.partition=%i" % self.partition_epoch,
        "--*.corpus.select-partition=%i" % ((epoch - 1) % self.partition_epoch)]
    args += [
      "--*.python-segment-order=true",
      "--*.python-segment-order-pymod-path=%s" % self._my_python_mod_path,
      "--*.python-segment-order-pymod-name=SprintExternInterface",
      "--*.use-data-source=false",
      "--*.trainer=python-trainer",
      "--*.pymod-path=%s" % self._my_python_mod_path,
      "--*.pymod-name=SprintExternInterface",
      "--*.pymod-config=%s" % config_str]
    if self.predefined_seq_list_order:
      import tempfile
      self.seq_list_file = tempfile.mktemp(prefix="crnn-sprint-predefined-seq-list")
      with open(self.seq_list_file, "w") as f:
        for tag in self.predefined_seq_list_order:
          f.write(tag)
          f.write("\n")
        f.close()
      args += [
        "--*.corpus.segment-order-shuffle=false",
        "--*.corpus.segments.file=%s" % self.seq_list_file,
        "--*.corpus.segment-order=%s" % self.seq_list_file]
    if self.seq_tags_filter is not None:
      assert not self.predefined_seq_list_order
      import tempfile
      self.seq_list_file = tempfile.mktemp(prefix="crnn-sprint-predefined-seq-filter")
      with open(self.seq_list_file, "w") as f:
        for tag in self.seq_tags_filter:
          f.write(tag)
          f.write("\n")
        f.close()
      args += ["--*.corpus.segments.file=%s" % self.seq_list_file]
    return args

  def _read_next_raw(self):
    """
    :return: (data_type, args)
    :rtype: (str, object)
    """
    import struct
    size_raw = self.pipe_c2p[0].read(4)
    if len(size_raw) < 4:
      raise EOFError
    size, = struct.unpack("<i", size_raw)
    assert size > 0, "%s: We expect to get some non-empty package. Invalid Python mod in Sprint?" % (self,)
    stream = BytesIO()
    read_size = 0
    while read_size < size:
      data_raw = self.pipe_c2p[0].read(size - read_size)
      if len(data_raw) == 0:
        raise EOFError("%s: expected to read %i bytes but got EOF after %i bytes" % (self, size, read_size))
      read_size += len(data_raw)
      stream.write(data_raw)
    stream.seek(0)
    try:
      if PY3:
        # encoding is for converting Python2 strings to Python3.
        # Cannot use utf8 because Numpy will also encode the data as strings and there we need it as bytes.
        data_type, args = Unpickler(stream, encoding="bytes").load()
      else:
        data_type, args = Unpickler(stream).load()
    except EOFError:
      raise Exception("%s: parse error of %i bytes (%r)" % (self, size, stream.getvalue()))
    return data_type, args

  def _join_child(self, wait=True, expected_exit_status=None):
    """
    :param bool wait:
    :param int|None expected_exit_status:
    :return: whether the child has exited now
    :rtype: bool
    """
    assert self.child_pid
    options = 0 if wait else os.WNOHANG
    pid, exit_status = os.waitpid(self.child_pid, options)
    if not wait and pid == 0:
      return False
    assert pid == self.child_pid
    if expected_exit_status is not None:
      assert exit_status == expected_exit_status, "%s: Sprint exit code is %i" % (self, exit_status)
    return True

  def _reader_thread_proc(self, child_pid, epoch):
    """
    :param int child_pid:
    :param int epoch:
    """
    try:
      self.add_data_thread_id = thread.get_ident()

      self.init_sprint_epoch(epoch)
      have_seen_the_whole = False

      seq_count = 0
      while not self.python_exit and self.child_pid:
        try:
          data_type, args = self._read_next_raw()
        except (IOError, EOFError):
          with self.lock:
            if epoch != self.crnnEpoch:
              # We have passed on to a new epoch. This is a valid reason that the child has been killed.
              break
            if self.python_exit or not self.child_pid:
              break
          raise

        with self.lock:
          if epoch != self.crnnEpoch:
            break
          if self.python_exit or not self.child_pid:
            break

          if data_type == b"data":
            seq_count += 1
            segment_name, features, targets = args
            if segment_name is not None:
              segment_name = segment_name.decode("utf8")
            assert isinstance(features, numpy.ndarray)
            if isinstance(targets, dict):
              targets = {key.decode("utf8"): value for (key, value) in targets.items()}
            self.add_new_data(
              numpy_copy_and_set_unused(features),
              numpy_copy_and_set_unused(targets),
              segment_name=segment_name)
          elif data_type == b"exit":
            have_seen_the_whole = True
            break
          else:
            assert False, "not handled: (%r, %r)" % (data_type, args)

      if self.seq_list_file:
        try:
          os.remove(self.seq_list_file)
        except Exception as e:
          print("%s: error when removing %r: %r" % (self, self.seq_list_file, e), file=log.v5)
        finally:
          self.seq_list_file = None

      if not self.python_exit:
        with self.lock:
          self.finish_sprint_epoch(seen_all=have_seen_the_whole)
          if have_seen_the_whole:
            self._num_seqs = self.next_seq_to_be_added
      print("%s (proc %i) finished reading epoch %i, seen all %r (finished), num seqs %i" % (
        self, child_pid, epoch, have_seen_the_whole, seq_count), file=log.v5)

    except Exception as exc:
      if not self.python_exit:
        # Catch all standard exceptions.
        # Don't catch KeyboardInterrupt here because that will get send by the main thread
        # when it is exiting. It's never by the user because SIGINT will always
        # trigger KeyboardInterrupt in the main thread only.
        if epoch == self.crnnEpoch:
          with self.lock:
            self.finish_sprint_epoch(seen_all=False)
        try:
          print("%s reader failed (%s)" % (self, exc), file=log.v1)
          sys.excepthook(*sys.exc_info())
          print("")
        finally:
          # Exceptions are fatal. If we can recover, we should handle it in run_inner().
          interrupt_main()

  def init_seq_order(self, epoch=None, seq_list=None):
    """
    :param int epoch:
    :param list[str]|None seq_list:
    :rtype: bool
    """
    if seq_list:
      assert self.partition_epoch == 1, "specifying partition_epoch and using seq_list not supported"
    if epoch is None:
      epoch = 1
    with self.lock:
      if epoch == self.crnnEpoch and self.expected_load_seq_start == 0 and seq_list == self.predefined_seq_list_order:
        return
      # Reset epoch such that exiting the child will go smoothly.
      super(ExternSprintDataset, self).init_seq_order(epoch=None, seq_list=None)
    # Exit child, before we overwrite anything, such as new epoch or seq_list.
    self._exit_child(wait_thread=True)
    with self.lock:  # Lock should not be needed now, but just to make it clean.
      if self._num_seqs:
        self._estimated_num_seqs = self._num_seqs  # last epoch num_seqs is a good estimate
      self._num_seqs = None  # we are not certain whether we have the same num_seqs for this epoch
      super(ExternSprintDataset, self).init_seq_order(epoch=epoch, seq_list=seq_list)
    self._start_child(epoch)
    return True


class SprintCacheDataset(CachedDataset2):
  """
  Can directly read Sprint cache files (and bundle files).
  Supports both cached features and cached alignments.
  For alignments, you need to provide all options for the AllophoneLabeling class, such as allophone file, etc.
  """

  class SprintCacheReader(object):
    """
    Helper class to read a Sprint cache directly.
    """
    def __init__(self, data_key, filename, data_type=None, allophone_labeling=None):
      """
      :param str data_key: e.g. "data" or "classes"
      :param str filename: to Sprint cache archive
      :param str|None data_type: "feat" or "align"
      :param dict[str] allophone_labeling: kwargs for :class:`AllophoneLabeling`
      """
      self.data_key = data_key
      from SprintCache import open_file_archive
      self.sprint_cache = open_file_archive(filename)
      if not data_type:
        if data_key == "data":
          data_type = "feat"
        elif data_key == "classes":
          data_type = "align"
        else:
          # Some sensible defaults.
          if allophone_labeling:
            data_type = "align"
          else:
            data_type = "feat"
      assert data_type in ["feat", "align"]
      self.type = data_type
      self.allophone_labeling = None
      if allophone_labeling:
        from SprintCache import AllophoneLabeling
        self.allophone_labeling = AllophoneLabeling(**allophone_labeling)
        self.sprint_cache.set_allophones(self.allophone_labeling.allophone_file)
      else:
        assert data_type != "align", "need allophone_labeling for 'align' type"
      self.content_keys = [fn for fn in self.sprint_cache.file_list() if not fn.endswith(".attribs")]
      if data_type == "align":
        self.num_labels = self.allophone_labeling.num_labels
        if self.num_labels < 2 ** 7:
          self.dtype = "int8"
        elif self.num_labels < 2 ** 15:
          self.dtype = "int16"
        else:
          assert self.num_labels < 2 ** 31
          self.dtype = "int32"
        self.num_dims = 1
        if self.allophone_labeling.state_tying_by_allo_state_idx:
          self.type = "align_raw"
      elif data_type == "feat":
        self.num_labels = self._get_feature_dim()
        self.num_dims = 2
        self.dtype = "float32"
      else:
        assert False

    def _get_feature_dim(self):
      """
      :rtype: int
      """
      assert self.type == "feat"
      assert self.content_keys
      times, feats = self.sprint_cache.read(self.content_keys[0], "feat")
      assert len(times) == len(feats) > 0
      feat = feats[0]
      assert isinstance(feat, numpy.ndarray)
      assert feat.ndim == 1
      return feat.shape[0]

    def read(self, name):
      """
      :param str name: content-filename for sprint cache
      :return: numpy array of shape (time, [num_labels])
      :rtype: numpy.ndarray
      """
      res = self.sprint_cache.read(name, typ=self.type)
      if self.type == "align":
        label_seq = numpy.array([self.allophone_labeling.get_label_idx(a, s) for (t, a, s) in res], dtype=self.dtype)
        assert label_seq.shape == (len(res),)
        return label_seq
      elif self.type == "align_raw":
        label_seq = numpy.array(
          [self.allophone_labeling.state_tying_by_allo_state_idx[a] for (t, a, s) in res], dtype=self.dtype)
        assert label_seq.shape == (len(res),)
        return label_seq
      elif self.type == "feat":
        times, feats = res
        assert len(times) == len(feats) > 0
        feat_mat = numpy.array(feats, dtype=self.dtype)
        assert feat_mat.shape == (len(times), self.num_labels)
        return feat_mat
      else:
        assert False

  def __init__(self, data, **kwargs):
    """
    :param dict[str,dict[str]] data: data-key -> dict which keys such as filename, see SprintCacheReader constructor
    """
    super(SprintCacheDataset, self).__init__(**kwargs)
    self.data = {key: self.SprintCacheReader(data_key=key, **opts) for (key, opts) in data.items()}
    self.seq_list_original = self.data["data"].content_keys
    self.seq_list_ordered = self.seq_list_original
    self._num_seqs = len(self.seq_list_original)
    self._check_matching_content_list()
    self.num_outputs = {key: (d.num_labels, d.num_dims) for (key, d) in self.data.items()}
    self.num_inputs = self.num_outputs["data"][0]
    self._seq_lens = None

  def _check_matching_content_list(self):
    data0 = self.data["data"]
    assert isinstance(data0, self.SprintCacheReader)
    sorted_list0 = sorted(data0.content_keys)
    for key, data in self.data.items():
      if key == "data":
        continue
      assert isinstance(data, self.SprintCacheReader)
      assert len(data.content_keys) == len(data0.content_keys)
      sorted_list1 = sorted(data.content_keys)
      for i in range(len(data.content_keys)):
        k0 = sorted_list0[i]
        k1 = sorted_list1[i]
        assert k0 == k1

  def init_seq_order(self, epoch=None, seq_list=None):
    """
    :param int epoch:
    :param list[str]|None seq_list:
    :rtype: bool
    """
    assert not seq_list
    need_reinit = self.epoch is None or self.epoch != epoch
    super(SprintCacheDataset, self).init_seq_order(epoch=epoch, seq_list=seq_list)
    if not need_reinit:
      return False
    self._num_seqs = len(self.seq_list_original)
    data0 = self.data["data"]
    assert isinstance(data0, self.SprintCacheReader)

    def get_seq_size(s):
      """
      :param int s:
      :rtype: int
      """
      return data0.sprint_cache.ft[self.seq_list_original[s]].size
    seq_index = self.get_seq_order_for_epoch(epoch, self.num_seqs, get_seq_len=get_seq_size)
    self.seq_list_ordered = [self.seq_list_original[s] for s in seq_index]
    return True

  def get_dataset_seq_for_name(self, name, seq_idx=-1):
    """
    :param str name:
    :param int seq_idx:
    :rtype: DatasetSeq
    """
    data = {key: d.read(name) for (key, d) in self.data.items()}  # type: typing.Dict[str,numpy.ndarray]
    return DatasetSeq(seq_idx=seq_idx, seq_tag=name, features=data["data"], targets=data)

  def _collect_single_seq(self, seq_idx):
    """
    :type seq_idx: int
    :rtype: DatasetSeq | None
    :returns DatasetSeq or None if seq_idx >= num_seqs.
    """
    if seq_idx >= self.num_seqs:
      return None
    seq_tag = self.get_tag(seq_idx)  # type: str
    return self.get_dataset_seq_for_name(seq_idx=seq_idx, name=seq_tag)

  def get_data_keys(self):
    """
    :rtype: list[str]
    """
    return self.data.keys()

  def get_target_list(self):
    """
    :rtype: list[str]
    """
    return [key for (key, d) in self.data if d.type == "align"]

  def get_tag(self, sorted_seq_idx):
    """
    :rtype: str
    """
    return self.seq_list_ordered[sorted_seq_idx]


def demo():
  """
  Demo.
  """
  print("SprintDataset demo.")
  from argparse import ArgumentParser
  from Util import progress_bar_with_time
  from Log import log
  from Config import Config
  from Dataset import init_dataset
  arg_parser = ArgumentParser()
  arg_parser.add_argument("--config", help="config with ExternSprintDataset", required=True)
  arg_parser.add_argument("--sprint_cache_dataset", help="kwargs dict for SprintCacheDataset", required=True)
  arg_parser.add_argument("--max_num_seqs", default=sys.maxsize, type=int)
  arg_parser.add_argument("--action", default="compare", help="compare or benchmark")
  args = arg_parser.parse_args()
  log.initialize(verbosity=[4])
  sprint_cache_dataset_kwargs = eval(args.sprint_cache_dataset)
  assert isinstance(sprint_cache_dataset_kwargs, dict)
  sprint_cache_dataset = SprintCacheDataset(**sprint_cache_dataset_kwargs)
  print("SprintCacheDataset: %r" % sprint_cache_dataset)
  config = Config()
  config.load_file(args.config)
  dataset = init_dataset(config.typed_value("train"))
  print("Dataset via config: %r" % dataset)
  assert sprint_cache_dataset.num_inputs == dataset.num_inputs
  assert tuple(sprint_cache_dataset.num_outputs["classes"]) == tuple(dataset.num_outputs["classes"])
  sprint_cache_dataset.init_seq_order(epoch=1)

  if args.action == "compare":
    print("Iterating through dataset...")
    seq_idx = 0
    dataset.init_seq_order(epoch=1)
    while seq_idx < args.max_num_seqs:
      if not dataset.is_less_than_num_seqs(seq_idx):
        break
      dataset.load_seqs(seq_idx, seq_idx + 1)
      tag = dataset.get_tag(seq_idx)
      assert not tag.startswith("seq-"), "dataset does not provide tag-names for seqs"
      dataset_seq = sprint_cache_dataset.get_dataset_seq_for_name(tag)
      data = dataset.get_data(seq_idx, "data")
      targets = dataset.get_data(seq_idx, "classes")
      assert data.shape == dataset_seq.features["data"].shape
      assert targets.shape == dataset_seq.features["classes"].shape
      assert numpy.allclose(data, dataset_seq.features["data"])
      assert numpy.allclose(targets, dataset_seq.features["classes"])
      seq_idx += 1
      progress_bar_with_time(dataset.get_complete_frac(seq_idx))

    print("Finished through dataset. Num seqs: %i" % seq_idx)
    print("SprintCacheDataset has num seqs: %i." % sprint_cache_dataset.num_seqs)

  elif args.action == "benchmark":
    print("Iterating through dataset...")
    start_time = time.time()
    seq_tags = []
    seq_idx = 0
    dataset.init_seq_order(epoch=1)
    while seq_idx < args.max_num_seqs:
      if not dataset.is_less_than_num_seqs(seq_idx):
        break
      dataset.load_seqs(seq_idx, seq_idx + 1)
      tag = dataset.get_tag(seq_idx)
      assert not tag.startswith("seq-"), "dataset does not provide tag-names for seqs"
      seq_tags.append(tag)
      dataset.get_data(seq_idx, "data")
      dataset.get_data(seq_idx, "classes")
      seq_idx += 1
      progress_bar_with_time(dataset.get_complete_frac(seq_idx))
    print("Finished through dataset. Num seqs: %i, time: %f" % (seq_idx, time.time() - start_time))
    print("SprintCacheDataset has num seqs: %i." % sprint_cache_dataset.num_seqs)
    if hasattr(dataset, "exit_handler"):
      dataset.exit_handler()
    else:
      print("No way to stop any background tasks.")
    del dataset

    start_time = time.time()
    print("Iterating through SprintCacheDataset...")
    for i, tag in enumerate(seq_tags):
      sprint_cache_dataset.get_dataset_seq_for_name(tag)
      progress_bar_with_time(float(i) / len(seq_tags))
    print("Finished through SprintCacheDataset. time: %f" % (time.time() - start_time,))

  else:
    raise Exception("invalid action: %r" % args.action)


if __name__ == "__main__":
  import better_exchook
  better_exchook.install()
  try:
    demo()
  except KeyboardInterrupt:
    print("KeyboardInterrupt.")
    sys.exit(1)