python-misc-utils 0.2__py3-none-any.whl

py_misc_utils/pipeline.py

@@ -0,0 +1,98 @@
+from . import assert_checks as tas
+from . import core_utils as cu
+
+
+class Pipeline:
+
+  def __init__(self, *elems):
+    self._elems = list(elems)
+
+  def __len__(self):
+    return len(self._elems)
+
+  def __getitem__(self, i):
+    return self._elems[i]
+
+  def __iter__(self):
+    return iter(self._elems)
+
+  def append(self, elem):
+    tas.check(callable(elem), msg=f'Pipeline elements must be callable: {type(elem)}')
+    self._elems.append(elem)
+
+    return len(self._elems) - 1
+
+  def extend(self, elems):
+    for elem in elems:
+      self.append(elem)
+
+  def pop(self, i=None):
+    return self._elems.pop(i) if i is not None else self._elems.pop()
+
+  def elems(self):
+    return tuple(self._elems)
+
+  def _apply(self, elem, data):
+    if cu.is_iterator(data):
+      if isinstance(elem, IterElement):
+        return elem(data)
+      else:
+        return _iter_process(elem, data)
+    elif isinstance(elem, IterElement):
+      return elem([data])
+    else:
+      return elem(data)
+
+  def __call__(self, x):
+    y = x
+    for elem in self._elems:
+      y = self._apply(elem, y)
+
+    return y
+
+  def __repr__(self):
+    return '\n'.join(f'[{i}] {repr(elem)}' for i, elem in enumerate(self._elems))
+
+  def clone(self):
+    # If a Pipeline elements has a state, it must implement the clone() API.
+    elems = [cu.maybe_call_dv(elem, 'clone', elem) for elem in self._elems]
+
+    return Pipeline(*elems)
+
+  def flush(self):
+    y = None
+    for elem in self._elems:
+      flush_fn = getattr(elem, 'flush', None)
+      if callable(flush_fn):
+        y = flush_fn(y or ())
+      elif y is not None:
+        y = self._apply(elem, y)
+
+    return y
+
+
+# Exception thrown by members of iterative pipelines when they want to stop the
+# stream of data, signaling that nothing more will be allowed through it.
+# Returning an empty iterator/generator will not cut it, as this is a legal
+# return value in case of batching elements.
+# Think about a pipeline element which is a filter which should pass through the
+# first N samples, for example.
+class HaltedPipeline(Exception):
+  pass
+
+
+# The Pipeline can also be used with data which is returned as iterators, where
+# there is not a 1:1 mapping between input and output.
+# Think about a pipeline element which absorbs data, and return batches of it (or
+# absorbs text and emits token indices).
+# The non-iterator approach would not work, as for many inputs there are no ouputs
+# at all (till the batch size is reached).
+# When used in such fashion, pipeline elements whould inherit from IterElement.
+class IterElement:
+  pass
+
+
+def _iter_process(elem, data):
+  for x in data:
+    yield elem(x)
+

py_misc_utils/remap_pickle.py

@@ -0,0 +1,50 @@
+import functools
+import io
+import pickle
+
+from . import alog
+from . import module_utils as mu
+
+
+class Unpickler(pickle.Unpickler):
+
+  def __init__(self, *args,
+               remaps=None,
+               safe_refs=None,
+               **kwargs):
+    super().__init__(*args, **kwargs)
+    self._remaps = remaps or dict()
+    self._safe_refs = set(safe_refs) if safe_refs is not None else None
+
+  def find_class(self, module, name):
+    fqname = f'{module}.{name}'
+    remap = self._remaps.get(fqname, fqname)
+    if remap != fqname:
+      alog.debug(f'Unpickle remapping: {fqname} -> {remap}')
+    elif self._safe_refs is not None and fqname not in self._safe_refs:
+      alog.xraise(RuntimeError, f'Unsafe reference: {fqname}')
+
+    return mu.import_module_names(remap)[0]
+
+
+def load(*args, **kwargs):
+  unpickler = Unpickler(*args, **kwargs)
+
+  return unpickler.load()
+
+
+def loads(data, *args, **kwargs):
+  memfd = io.BytesIO(data)
+  unpickler = Unpickler(memfd, *args, **kwargs)
+
+  return unpickler.load()
+
+
+def make_module(**kwargs):
+  module = mu.clone_module('pickle')
+
+  module.load = functools.partial(load, **kwargs)
+  module.loads = functools.partial(loads, **kwargs)
+
+  return module
+

py_misc_utils/resource_manager.py

@@ -0,0 +1,155 @@
+import collections
+import functools
+import logging
+import multiprocessing
+import multiprocessing.managers as mpmgr
+import os
+import queue
+import tempfile
+import threading
+import time
+import weakref
+
+from . import alog
+from . import daemon_process as dp
+
+
+class ResourceManager:
+
+  def __init__(self):
+    self._lock = threading.Lock()
+    self._resources = collections.defaultdict(dict)
+
+  def get(self, cls, ctor, name, *args, **kwargs):
+    alog.debug(f'Get resource {cls}.{name}')
+    with self._lock:
+      cdict = self._resources[cls]
+      res = cdict.get(name)
+      if res is not None:
+        res = res()
+      if res is None:
+        alog.debug(f'Creating resource {cls}.{name}')
+        res = ctor(*args, **kwargs)
+        cdict[name] = weakref.ref(res)
+
+      return res
+
+  def delete(self, cls, name):
+    alog.debug(f'Remove resource {cls}.{name}')
+    with self._lock:
+      cdict = self._resources[cls]
+      cdict.pop(name, None)
+
+
+_RESMGR = ResourceManager()
+
+def get_resource_manager():
+  return _RESMGR
+
+
+def _create_manager(*args, register_fn=None, **kwargs):
+  # https://github.com/python/cpython/blob/2f56c68dec97002fdd8563a0e4977b75eb191ab9/Lib/multiprocessing/managers.py#L1043
+  # https://github.com/python/cpython/blob/4cba0e66c29b46afbb1eee1d0428f5a2f5b891bb/Lib/multiprocessing/managers.py#L189
+  manager = mpmgr.SyncManager(*args, **kwargs)
+
+  resmgr = get_resource_manager()
+
+  manager.register('get_lock',
+                   functools.partial(resmgr.get, 'LOCKS', threading.Lock),
+                   mpmgr.AcquirerProxy)
+  manager.register('rm_lock',
+                   functools.partial(resmgr.delete, 'LOCKS'))
+
+  manager.register('get_event',
+                   functools.partial(resmgr.get, 'EVENTS', threading.Event),
+                   mpmgr.EventProxy)
+  manager.register('rm_event',
+                   functools.partial(resmgr.delete, 'EVENTS'))
+
+  manager.register('get_condition',
+                   functools.partial(resmgr.get, 'CONDITIONS', threading.Condition),
+                   mpmgr.ConditionProxy)
+  manager.register('rm_condition',
+                   functools.partial(resmgr.delete, 'CONDITIONS'))
+
+  manager.register('get_barrier',
+                   functools.partial(resmgr.get, 'BARRIERS', threading.Barrier),
+                   mpmgr.BarrierProxy)
+  manager.register('rm_barrier',
+                   functools.partial(resmgr.delete, 'BARRIERS'))
+
+  manager.register('get_queue',
+                   functools.partial(resmgr.get, 'QUEUES', queue.Queue))
+  manager.register('rm_queue',
+                   functools.partial(resmgr.delete, 'QUEUES'))
+
+  manager.register('get_lifo',
+                   functools.partial(resmgr.get, 'LIFOS', queue.LifoQueue))
+  manager.register('rm_lifo',
+                   functools.partial(resmgr.delete, 'LIFOS'))
+
+  manager.register('get_namespace',
+                   functools.partial(resmgr.get, 'NAMESPACES', mpmgr.Namespace),
+                   mpmgr.NamespaceProxy)
+  manager.register('rm_namespace',
+                   functools.partial(resmgr.delete, 'NAMESPACES'))
+
+  if register_fn is not None:
+    register_fn(manager, resmgr)
+
+  return manager
+
+
+def _get_logdir():
+  logdir = os.path.join(tempfile.gettempdir(), 'log')
+  os.makedirs(logdir, exist_ok=True)
+
+  return logdir
+
+
+def _server_runner(name, *args, **kwargs):
+  alog.basic_setup(log_level=os.getenv('RESMGR_LOG_LEVEL', 'INFO'),
+                   log_file=os.path.join(_get_logdir(), f'{name}.log'))
+
+  alog.info(f'[{name}] server starting')
+  manager = _create_manager(*args, **kwargs)
+
+  try:
+    server = manager.get_server()
+    server.serve_forever()
+  except Exception as ex:
+    alog.error(f'[{name}] server start failed: {ex}')
+  finally:
+    alog.info(f'[{name}] server gone!')
+
+
+def get_manager(name, *args, **kwargs):
+  daemon = dp.Daemon(name)
+  while daemon.getpid() is None:
+    alog.info(f'[{name}] Starting server daemon')
+    try:
+      daemon.start(functools.partial(_server_runner, name, *args, **kwargs))
+    except FileExistsError as ex:
+      pass
+    time.sleep(0.5)
+
+  alog.info(f'[{name}] Connecting to server')
+
+  manager = _create_manager(*args, **kwargs)
+  while True:
+    try:
+      manager.connect()
+      alog.info(f'[{name}] Connected to the manager')
+      break
+    except Exception as ex:
+      alog.debug(f'[{name}] Connection failed, retrying ...: {ex}')
+      time.sleep(0.5)
+
+  return manager
+
+
+def stop_manager(name):
+  daemon = dp.Daemon(name)
+
+  return daemon.stop()
+

py_misc_utils/rnd_utils.py

@@ -0,0 +1,56 @@
+import binascii
+import random
+import string
+import struct
+
+import numpy as np
+
+from . import assert_checks as tas
+
+
+def compute_seed(seed):
+  if isinstance(seed, int):
+    seed = binascii.crc32(struct.pack('=q', seed))
+  elif isinstance(seed, float):
+    seed = binascii.crc32(struct.pack('=d', seed))
+  elif isinstance(seed, bytes):
+    seed = binascii.crc32(seed)
+  elif isinstance(seed, str):
+    seed = binascii.crc32(seed.encode())
+  else:
+    seed = binascii.crc32(struct.pack('=Q', hash(seed)))
+
+  return seed
+
+
+def manual_seed(seed):
+  cseed = compute_seed(seed)
+
+  np.random.seed(cseed)
+  random.seed(cseed)
+
+  return cseed
+
+
+def choices(weights, n):
+  return random.choices(range(len(weights)), weights=weights, k=n)
+
+
+def shuffle(args):
+  return random.sample(args, k=len(args))
+
+
+def uniform(center, delta=None, pct=None):
+  if pct is not None:
+    delta = abs(center * pct)
+
+  tas.check_is_not_none(delta, msg=f'Either delta or pct must be provided')
+
+  return random.uniform(center - delta, center + delta)
+
+
+def rand_string(n):
+  rng = random.SystemRandom()
+
+  return ''.join(rng.choices(string.ascii_lowercase + string.digits, k=n))
+

py_misc_utils/run_once.py

@@ -0,0 +1,19 @@
+import functools
+import threading
+
+
+def run_once(fn):
+
+  @functools.wraps(fn)
+  def wrapper(*args, **kwargs):
+    with wrapper._lock:
+      if not wrapper._has_run:
+        wrapper._has_run = True
+
+        return fn(*args, **kwargs)
+
+  wrapper._lock = threading.Lock()
+  wrapper._has_run = False
+
+  return wrapper
+

py_misc_utils/scheduler.py

@@ -0,0 +1,135 @@
+import collections
+import heapq
+import os
+import threading
+import time
+import uuid
+
+from . import alog
+from . import executor as xe
+from . import fin_wrap as fw
+from . import timegen as tg
+from . import utils as ut
+
+
+Event = collections.namedtuple(
+    'Event',
+    'time, sequence, ref, action, argument, kwargs')
+
+
+class Scheduler:
+
+  def __init__(self, timegen=None, executor=None, max_workers=None, name='Scheduler'):
+    self._queue = []
+    self._sequence = 0
+    self._lock = threading.Lock()
+    self._cond = threading.Condition(lock=self._lock)
+    self.timegen = tg.TimeGen() if timegen is None else timegen
+
+    if executor is not None:
+      self.executor = executor
+    else:
+      executor = xe.Executor(max_threads=max_workers, name_prefix=name)
+      fw.fin_wrap(self, 'executor', executor, finfn=executor.shutdown)
+
+    self._runner = threading.Thread(target=self._run, daemon=True)
+    self._runner.start()
+
+  def _run_event(self, event):
+    try:
+      event.action(*event.argument, **event.kwargs)
+    except Exception as ex:
+      alog.exception(ex, exmsg=f'Exception while running scheduled action')
+
+  def _run(self):
+    while True:
+      now, event = self.timegen.now(), None
+      with self._lock:
+        timeout = (self._queue[0].time - now) if self._queue else None
+        if timeout is None or timeout > 0:
+          self.timegen.wait(self._cond, timeout=timeout)
+        else:
+          event = heapq.heappop(self._queue)
+
+      if event is not None:
+        self.executor.submit(self._run_event, event)
+
+  def gen_unique_ref(self):
+    return str(uuid.uuid4())
+
+  def enterabs(self, ts, action, ref=None, argument=(), kwargs={}):
+    with self._lock:
+      event = Event(time=ts,
+                    sequence=self._sequence,
+                    ref=ref,
+                    action=action,
+                    argument=argument,
+                    kwargs=kwargs)
+      self._sequence += 1
+
+      heapq.heappush(self._queue, event)
+      if id(event) == id(self._queue[0]):
+        self._cond.notify()
+
+    return event
+
+  def enter(self, delay, action, ref=None, argument=(), kwargs={}):
+    return self.enterabs(self.timegen.now() + delay, action,
+                         ref=ref,
+                         argument=argument,
+                         kwargs=kwargs)
+
+  def _cancel_fn(self, fn):
+    events = []
+    with self._lock:
+      pos = []
+      for i, qe in enumerate(self._queue):
+        if fn(qe):
+          events.append(qe)
+          pos.append(i)
+      if pos:
+        # Positions are added in asceending order above, here we pop them in
+        # descending order to avoid invalidating positions.
+        for i in range(len(pos) - 1, -1, -1):
+          self._queue.pop(pos[i])
+
+        heapq.heapify(self._queue)
+
+    return events
+
+  def cancel(self, event):
+    if isinstance(event, (list, tuple)):
+      ids = set([id(e) for e in event])
+      return self._cancel_fn(lambda qe: id(qe) in ids)
+
+    return self._cancel_fn(lambda qe: id(qe) == id(event))
+
+  def ref_cancel(self, ref):
+    if isinstance(ref, (list, tuple)):
+      refs = set(ref)
+      return self._cancel_fn(lambda qe: qe.ref in refs)
+
+    return self._cancel_fn(lambda qe: qe.ref == ref)
+
+  def get_events(self, fn):
+    events = []
+    with self._lock:
+      for qe in self._queue:
+        if fn(qe):
+          events.append(qe)
+
+    return events
+
+
+_LOCK = threading.Lock()
+_SCHEDULER = None
+
+def common_scheduler():
+  global _SCHEDULER
+
+  with _LOCK:
+    if _SCHEDULER is None:
+      _SCHEDULER = Scheduler(executor=xe.common_executor())
+
+    return _SCHEDULER
+