thds.core 0.0.1__py3-none-any.whl → 1.31.20250116223856__py3-none-any.whl

thds/core/parallel.py

@@ -0,0 +1,200 @@
+"""Some utilities for running things in parallel - potentially large numbers of things.
+"""
+
+import concurrent.futures
+import itertools
+import typing as ty
+from collections import defaultdict
+from dataclasses import dataclass
+from uuid import uuid4
+
+from thds.core import concurrency, config, files, log
+
+PARALLEL_OFF = config.item("off", default=False, parse=config.tobool)
+# if you want to simplify a stack trace, this may be your friend
+
+R = ty.TypeVar("R")
+T_co = ty.TypeVar("T_co", covariant=True)
+
+
+logger = log.getLogger(__name__)
+
+
+class IterableWithLen(ty.Protocol[T_co]):
+    def __iter__(self) -> ty.Iterator[T_co]:
+        ...  # pragma: no cover
+
+    def __len__(self) -> int:
+        ...  # pragma: no cover
+
+
+class IteratorWithLen(ty.Generic[R]):
+    """Suitable for a case where you know how many elements you have
+    and you want to be able to represent that somewhere else but you
+    don't want to 'realize' all the elements upfront.
+    """
+
+    def __init__(self, length: int, iterable: ty.Iterable[R]):
+        self._length = length
+        self._iterator = iter(iterable)
+
+    def __len__(self) -> int:
+        return self._length
+
+    def __iter__(self) -> ty.Iterator[R]:
+        return self
+
+    def __next__(self) -> R:
+        return next(self._iterator)
+
+    @staticmethod
+    def chain(a: IterableWithLen[R], b: IterableWithLen[R]) -> "IteratorWithLen[R]":
+        return IteratorWithLen(len(a) + len(b), itertools.chain(a, b))
+
+    @staticmethod
+    def from_iwl(iwl: IterableWithLen[R]) -> "IteratorWithLen[R]":
+        return IteratorWithLen(len(iwl), iwl)
+
+
+def try_len(iterable: ty.Iterable[R]) -> ty.Optional[int]:
+    try:
+        return len(iterable)  # type: ignore
+    except TypeError:
+        return None
+
+
+@dataclass
+class Error:
+    error: Exception
+
+
+H = ty.TypeVar("H", bound=ty.Hashable)
+
+
+def yield_all(
+    thunks: ty.Iterable[ty.Tuple[H, ty.Callable[[], R]]],
+    *,
+    executor_cm: ty.Optional[ty.ContextManager[concurrent.futures.Executor]] = None,
+) -> ty.Iterator[ty.Tuple[H, ty.Union[R, Error]]]:
+    """Stream your results so that you don't have to load them all into memory at the same
+    time (necessarily). Also, yield (rather than raise) Exceptions, wrapped as Errors.
+
+    Additionally, if your iterable has a length and you do not provide
+    a pre-sized Executor, we will create a ThreadPoolExecutor with the
+    same size as your iterable. If you want to throttle the number of
+    parallel tasks, you should provide your own Executor - and for
+    most mops purposes it should be a ThreadPoolExecutor.
+    """
+    files.bump_limits()
+    len_or_none = try_len(thunks)
+
+    if PARALLEL_OFF() or len_or_none == 1:
+        # don't actually transfer this to an executor we only have one task.
+        for key, thunk in thunks:
+            try:
+                yield key, thunk()
+            except Exception as e:
+                yield key, Error(e)
+        return  # we're done here
+
+    executor_cm = executor_cm or concurrent.futures.ThreadPoolExecutor(
+        max_workers=len_or_none or None, **concurrency.initcontext()
+    )  # if len_or_none turns out to be zero, swap in a None which won't kill the executor
+    with executor_cm as executor:
+        keys_onto_futures = {key: executor.submit(thunk) for key, thunk in thunks}
+        future_ids_onto_keys = {id(future): key for key, future in keys_onto_futures.items()}
+        for future in concurrent.futures.as_completed(keys_onto_futures.values()):
+            thunk_key = future_ids_onto_keys[id(future)]
+            try:
+                yield thunk_key, ty.cast(R, future.result())
+            except Exception as e:
+                yield thunk_key, Error(e)
+
+
+def failfast(results: ty.Iterable[ty.Tuple[H, ty.Union[R, Error]]]) -> ty.Iterator[ty.Tuple[H, R]]:
+    for key, res in results:
+        if isinstance(res, Error):
+            raise res.error
+        yield key, res
+
+
+def xf_mapping(thunks: ty.Mapping[H, ty.Callable[[], R]]) -> ty.Iterator[ty.Tuple[H, R]]:
+    return failfast(yield_all(IteratorWithLen(len(thunks), thunks.items())))
+
+
+def create_keys(iterable: ty.Iterable[R]) -> ty.Iterator[ty.Tuple[str, R]]:
+    """Use this if you wanted to call yield_all with a list or other sequence that
+    has no keys, and you don't need to 'track' the correspondence of input thunks to
+    output results.
+    """
+    with_keys: ty.Iterable[ty.Tuple[str, R]] = ((uuid4().hex, item) for item in iterable)
+    try:
+        return IteratorWithLen(len(iterable), with_keys)  # type: ignore[arg-type]
+    except TypeError:
+        return iter(with_keys)
+
+
+def yield_results(
+    thunks: ty.Iterable[ty.Callable[[], R]],
+    *,
+    executor_cm: ty.Optional[ty.ContextManager[concurrent.futures.Executor]] = None,
+    error_fmt: ty.Callable[[str], str] = lambda x: x,
+    success_fmt: ty.Callable[[str], str] = lambda x: x,
+    named: str = "",
+    progress_logger: ty.Callable[[str], ty.Any] = logger.info,
+) -> ty.Iterator[R]:
+    """Yield only the successful results of your Callables/Thunks.
+
+    If your iterable has a length, we will be able to log progress
+    information. In most cases, this will be advantageous for you.
+
+    Each task will fail or succeed separately without impacting other tasks.
+
+    However, if any Exceptions are raised in any task, an Exception
+    will be raised at the end of execution to indicate that not all
+    tasks were successful. If you wish to capture Exceptions alongside
+    results, use `yield_all` instead.
+    """
+
+    exceptions: ty.List[Exception] = list()
+
+    num_tasks = try_len(thunks)
+    num_tasks_log = "" if not num_tasks else f" of {num_tasks}"
+    named = f" {named} " if named else " result "
+
+    for i, (_key, res) in enumerate(
+        yield_all(create_keys(thunks), executor_cm=executor_cm),
+        start=1,
+    ):
+        if not isinstance(res, Error):
+            errors = error_fmt(f"; {len(exceptions)} tasks have raised exceptions") if exceptions else ""
+            progress_logger(success_fmt(f"Yielding{named}{i}{num_tasks_log} {errors}"))
+            yield res
+        else:
+            exceptions.append(res.error)
+            logger.exception(error_fmt(f"Task {i}{num_tasks_log} errored with {str(res.error)}"))
+
+    summarize_exceptions(error_fmt, exceptions)
+
+
+def summarize_exceptions(
+    error_fmt: ty.Callable[[str], str],
+    exceptions: ty.List[Exception],
+) -> None:
+    if exceptions:
+        # summarize them and raise a final exception
+        by_type = defaultdict(list)
+        for exc in exceptions:
+            by_type[type(exc)].append(exc)
+            logger.error(error_fmt("EXCEPTION"), exc_info=(type(exc), exc, exc.__traceback__))
+
+        most_common_type = None
+        max_count = 0
+        for _type, excs in by_type.items():
+            logger.error(error_fmt(f"{len(excs)} tasks failed with exception: " + str(_type)))
+            if len(excs) > max_count:
+                max_count = len(excs)
+                most_common_type = _type
+
+        logger.info("Raising one of the most common exception type.")
+        raise by_type[most_common_type][0]  # type: ignore

thds/core/pickle_visit.py

@@ -0,0 +1,24 @@
+import typing as ty
+from io import BytesIO
+from pickle import Pickler
+
+
+def recursive_visit(visitor: ty.Callable[[ty.Any], ty.Any], obj: ty.Any) -> None:
+    """A hilarious abuse of pickle to do nearly effortless recursive object 'visiting' in
+    Python for us.  In other words, if you want to 'see' everything inside an object but
+    don't actually care about serializing.
+
+    This can only work for objects that are fully recursively picklable. If yours isn't,
+    a pickling error will be raised and only some of the object will be visited.
+    """
+
+    class PickleVisit(Pickler):
+        def __init__(self, file):
+            super().__init__(file)
+            self.file = file
+
+        def reducer_override(self, obj: ty.Any):
+            visitor(obj)
+            return NotImplemented
+
+    PickleVisit(BytesIO()).dump(obj)

thds/core/prof.py

@@ -0,0 +1,276 @@
+"""Intentionally unsophisticated memory and CPU profiler that is
+triggered by logging.  Used for sanity-checking against the results
+provided by other, black-box memory profilers.
+
+Essentially, this can be added to any logger to force a sample of
+memory and CPU usage on every logging statement, with the context
+provided by the logger.
+
+Since logging often happens at opportune times anyway, this is a
+fairly easy and low-overhead way of getting the 'history' of your
+process to be output within and alongside the logs.
+
+This will do nothing if you don't have psutil installed.
+If you do have psutil installed, it will be enabled automatically,
+but without any patched loggers, you'll get no profiling output.
+
+You can patch your loggers or wrap them.
+
+Patching is implicit, and will try to intercept all calls you make to
+`core.getLogger`. You must import this module before any others and
+call the monkey_patch_core_getLogger function. Alternatively, you can
+set the TH_PROF_ALL_LOGGERS environment variable and this
+monkey-patching will be done automatically.
+
+To wrap a logger, simply use the output of `wrap_logger(YOUR_LOGGER)`
+as your logger. It will automatically output profiling information on
+every usage.
+"""
+import contextlib
+import csv
+import logging
+import os
+import pathlib
+import random
+import string
+import sys
+import typing as ty
+from datetime import datetime
+from timeit import default_timer
+
+from thds.core.stack_context import StackContext
+
+F = ty.TypeVar("F", bound=ty.Callable)
+Decorator = ty.Callable[[F], F]
+
+try:
+    import psutil  # type: ignore
+except ImportError:
+    psutil = None
+
+
+def corelog():
+    """We need core.log's basicConfig call to not happen on startup"""
+    import thds.core.log as log
+
+    return log
+
+
+_IS_ENABLED = bool(psutil)
+
+_PROF = "PROF>"
+_MSG = "MSG>"
+_TAG_LEN = 6
+
+
+_TRUE_START = default_timer()
+
+
+def _get_time():
+    return default_timer() - _TRUE_START
+
+
+class NoPsutilError(Exception):
+    pass
+
+
+def _get_proc(pid=-1, cache=dict()):  # noqa: B006
+    """If we want CPU statistics, we need to cache these things"""
+    if not psutil:
+        raise NoPsutilError("")
+    if pid < 0:
+        pid = os.getpid()
+    if pid not in cache:
+        cache[pid] = psutil.Process(pid)
+    return cache[pid]
+
+
+def _get_mem_mb() -> float:
+    try:
+        return _get_proc().memory_info().rss / 10**6
+    except NoPsutilError:
+        return 0.0
+
+
+def _get_cpu_percent(pid: int = -1) -> float:
+    try:
+        proc = _get_proc(pid)
+        return proc.cpu_percent() + sum(_get_cpu_percent(c.pid) for c in proc.children())
+    except NoPsutilError:
+        return 0.0
+
+
+_PID: int = 0
+_CSV_OUT = None
+_CSV_WRITER = None
+_PROFS_DIR = pathlib.Path("th-profiles")
+
+
+def _open_csv_writer():
+    global _PID, _CSV_OUT, _CSV_WRITER
+    cur_pid = os.getpid()
+    if _PID != cur_pid:  # we have forked a new process - open a new file
+        parent_pid = "" if _PID == 0 else f"{_PID}_PARENT_"
+        _PROFS_DIR.mkdir(exist_ok=True)
+        csvf = _PROFS_DIR / (
+            f"th-prof-{datetime.utcnow().isoformat()}-"
+            f"{'-'.join(sys.argv).replace('/', '_')}___{parent_pid}{cur_pid}.csv"
+        )
+        _CSV_OUT = csvf.open("w")
+        _CSV_WRITER = csv.writer(_CSV_OUT)
+        _CSV_WRITER.writerow(("tag", "time", "mem", "cpu", "time_d", "mem_d", "descriptors", "msg"))
+        _PID = cur_pid
+
+
+def _write_record(*row: str) -> None:
+    if _IS_ENABLED:
+        _open_csv_writer()
+        _CSV_WRITER.writerow(row)  # type: ignore
+        _CSV_OUT.flush()  # type: ignore
+
+
+def _dt(dt: float) -> str:
+    return f"{dt:+9.1f}s"
+
+
+def _t(t: float) -> str:
+    return f"{t:9.1f}s"
+
+
+def _cpu(cpu: float) -> str:
+    return f"{cpu:7.1f}%"
+
+
+def _dm(dm: float) -> str:
+    return f"{dm:+9.1f}m"
+
+
+def _m(m: float) -> str:
+    return f"{m:9.1f}m"
+
+
+class ThProfiler:
+    """A ContextManager that outputs time and memory on enter, exit, and
+    whenever it is called in between.
+
+    This is not designed for tight loops - it makes kernel calls, and
+    therefore is most suitable for attaching to relatively beefy parts
+    of your program that you want general purpose, basic profiling
+    info logged for.
+
+    It _is_ designed to be used in concert with `scope.enter` - using
+    with statements for profiling code that might later be removed is
+    very ugly and a Bad Idea™.
+    """
+
+    def __init__(self, tag: str):
+        self.tag = tag
+        self.start_mem = _get_mem_mb()
+        self.start_t = _get_time()
+
+    def _desc(self, **more_desc) -> str:
+        return "; ".join([f"{k}: {v}" for k, v in dict(corelog()._LOG_CONTEXT(), **more_desc).items()])
+
+    def _parts(self) -> ty.Tuple[str, ...]:
+        now_m = _get_mem_mb()
+        now_t = _get_time()
+        now_cpu = _get_cpu_percent()
+        delta_t = now_t - self.start_t
+        delta_m = now_m - self.start_mem
+        return _t(now_t), _m(now_m), _cpu(now_cpu), _dt(delta_t), _dm(delta_m)
+
+    def _make_msg(self, msg: str, **more_desc) -> ty.Tuple[str, ...]:
+        descriptors = self._desc(**more_desc)
+        parts = self._parts()
+        _write_record(self.tag, *parts, descriptors, msg)
+        if msg:
+            msg = _MSG + " " + msg
+        return tuple(filter(None, (_PROF, self.tag, *parts, descriptors, msg)))
+
+    def __call__(self, msg: str = "", **descriptors) -> str:
+        """Outputs the delta between now and when this ThProfiler was created"""
+        if not _IS_ENABLED:
+            return msg
+        return " ".join(self._make_msg(msg, **descriptors))
+
+
+_TH_PROFILER = StackContext(
+    "__th_profiler",
+    ThProfiler("." * _TAG_LEN)
+    # This root profiler exists simply to make sure that there's
+    # always a profiling context even if no function has been
+    # annotated with `context`.
+)
+
+
+def _make_tag(parent: str = "", _cache: dict = {None: 0}) -> str:  # noqa: B006
+    # try to cycle through string prefixes for reasonability
+    try:
+        src = string.ascii_uppercase
+        ascending_prefix = src[_cache[None] % len(src)]
+        if parent:
+            return ascending_prefix + parent[:-1]
+        end = "".join(random.choice(src) for _ in range(_TAG_LEN - 1))
+        return ascending_prefix + end
+    finally:
+        _cache[None] += 1  # increment prefix
+
+
+@contextlib.contextmanager
+def profile(profname: str, tag: str = ""):
+    """Establish a new profiling context around a given function, with its
+    own 'start values' for time and memory, and its own 'tag' for
+    identifiability.
+    """
+
+    if not _IS_ENABLED:
+        yield
+        return
+
+    with corelog().logger_context(profname=profname):
+        with _TH_PROFILER.set(ThProfiler(tag or _make_tag(_TH_PROFILER().tag))):
+            yield
+
+
+class ProfilingLoggerAdapter(logging.LoggerAdapter):
+    """Use this to replace your current logger - it remains a logger, but now prepends profiling information!"""
+
+    def process(self, msg, kwargs):
+        kw_logger = corelog().KwLogger(self.logger, dict())
+        msg, kwargs = kw_logger.process(msg, kwargs)
+        if not _IS_ENABLED:
+            return msg, kwargs
+        extra = kwargs.get("extra") or dict()
+        th_kw = extra.get(corelog()._TH_REC_CTXT) or dict()
+        return _TH_PROFILER()(msg, **th_kw), kwargs  # type: ignore
+
+
+def wrap_logger(logger: logging.Logger) -> ProfilingLoggerAdapter:
+    return ProfilingLoggerAdapter(logger, dict())
+
+
+def monkey_patch_core_getLogger():
+    """Only call this if you want all your loggers replaced by profiling loggers.
+
+    For this to take effect, you must call it *before* imports of modules that define loggers.
+    """
+    if not _IS_ENABLED:
+        return
+    import thds.core.log
+
+    thds.core.log.__dict__["getLogger"] = lambda name: wrap_logger(logging.getLogger(name))
+
+
+if "TH_PROF_ALL_LOGGERS" in os.environ:
+    if not _IS_ENABLED:
+        print(
+            "Warning: psutil is not installed but you have set TH_PROF_ALL_LOGGERS. Please pip install psutil"
+        )
+    else:
+        print("Monkey patching all loggers because of TH_PROF_ALL_LOGGERS")
+        logging.basicConfig(
+            level=os.environ.get("LOGLEVEL", logging.INFO),
+            style="{",
+            format="{name:<45} - {levelname:^8} - {message}",
+        )
+        monkey_patch_core_getLogger()

thds/core/progress.py

@@ -0,0 +1,112 @@
+import math
+import typing as ty
+from datetime import timedelta
+from functools import partial, wraps
+from timeit import default_timer
+
+from typing_extensions import ParamSpec
+
+from thds.core import log, scope
+
+
+def _smooth_number(num):
+    if num == 0:
+        return 0
+
+    # Determine the order of magnitude of the number
+    magnitude = int(math.log10(abs(num)))
+
+    # Find the nearest power of 10
+    power_of_10 = 10**magnitude
+
+    # Scale down the number to the range 1-10
+    scaled_num = num / power_of_10
+
+    # Round scaled number to nearest 1, 2, or 5 times 10^n
+    if scaled_num < 1.5:
+        return 1 * power_of_10
+    elif scaled_num < 3:
+        return 2 * power_of_10
+    elif scaled_num < 7:
+        return 5 * power_of_10
+    return 10 * power_of_10
+
+
+logger = log.getLogger(__name__.replace("ud.shared.", ""))
+T = ty.TypeVar("T")
+_progress_scope = scope.Scope("progress")
+
+
+def calc_report_every(target_interval: float, total: int, sec_elapsed: float) -> int:
+    seconds_per_item = sec_elapsed / total
+    rate = 1 / seconds_per_item
+    target_rate = 1 / target_interval
+    return _smooth_number(int(rate / target_rate) or 1)
+
+
+def _name(obj: ty.Any) -> str:
+    if hasattr(obj, "__class__"):
+        return obj.__class__.__name__
+    if hasattr(obj, "__name__"):
+        return obj.__name__
+    if obj is not None:
+        return str(obj)[:20]
+    return "item"
+
+
+@_progress_scope.bound
+def report(
+    it: ty.Iterable[T],
+    *,
+    name: str = "",
+    roughly_every_s: timedelta = timedelta(seconds=20),
+) -> ty.Iterator[T]:
+    """Report round-number progress roughly every so often..."""
+    iterator = iter(it)
+    total = 0
+    start = default_timer()
+    report_every = 0
+    last_report = 0
+    frequency = roughly_every_s.total_seconds()
+
+    try:
+        first_item = next(iterator)
+        name = name or _name(first_item)
+        _progress_scope.enter(log.logger_context(P=name))
+        yield first_item
+        total = 1
+    except StopIteration:
+        pass
+
+    for total, item in enumerate(iterator, start=2):
+        yield item
+
+        if not report_every:
+            elapsed = default_timer() - start
+            if elapsed > frequency * 0.5:
+                report_every = calc_report_every(frequency, total, elapsed)
+        elif report_every and (total % report_every == 0) and (total - last_report >= report_every):
+            elapsed = default_timer() - start
+            # once we have our first report_every value, don't get the time on every iteration
+            if total >= elapsed:
+                rate_str = f"{total / elapsed:10,.0f}/s"
+            else:
+                rate_str = f"{elapsed / total:10,.0f}s/{name}"
+            logger.info(f"Processed {total:12,d} in {elapsed:6,.1f}s at {rate_str}")
+            last_report = total
+            report_every = calc_report_every(frequency, total, elapsed)
+
+    _log = logger.info if total > 0 else logger.warning
+    elapsed = default_timer() - start
+    _log(f"FINISHED {total:12,d} {name} in {elapsed:6,.1f}s at {total / elapsed:10,.0f}/s")
+
+
+P = ParamSpec("P")
+
+
+def _report_gen(f: ty.Callable[P, ty.Iterator[T]], *args: P.args, **kwargs: P.kwargs) -> ty.Iterator[T]:
+    yield from report(f(*args, **kwargs))
+
+
+def report_gen(f: ty.Callable[P, ty.Iterator[T]]) -> ty.Callable[P, ty.Iterator[T]]:
+    return wraps(f)(partial(_report_gen, f))

thds/core/protocols.py

@@ -0,0 +1,17 @@
+import typing as ty
+from types import TracebackType
+
+_T_co = ty.TypeVar("_T_co", covariant=True)
+
+
+class ContextManager(ty.Protocol[_T_co]):
+    def __enter__(self) -> _T_co:
+        ...
+
+    def __exit__(
+        self,
+        exc_type: ty.Optional[ty.Type[BaseException]],
+        exc_value: ty.Optional[BaseException],
+        traceback: ty.Optional[TracebackType],
+    ) -> ty.Optional[bool]:
+        ...

thds/core/scaling.py

@@ -0,0 +1,39 @@
+"""A 'scale group' is a string that represents a configuration 'grouping' various systems
+by a shared input scaling factor.
+
+This module only provides a way for other modules to add to the active set of scale groups.
+Other modules/projects must provide their own meaning for these names.
+"""
+
+import contextlib
+import typing as ty
+
+from thds.core.stack_context import StackContext
+
+_SCALE_GROUP_PRIORITY: StackContext[ty.Tuple[str, ...]] = StackContext("_SCALE_GROUP_PRIORITY", ("",))
+
+
+@contextlib.contextmanager
+def push_scale_group(*scale_group_names: str) -> ty.Iterator[ty.Tuple[str, ...]]:
+    """Puts some scale group strings at the head of the list, such that they will be
+    'found first', before all other scale groups. First argument is the highest overall priority.
+
+    The scale groups are popped off the stack when the context manager exits.
+
+    You may want to use this in conjunction with `core.scope.enter` (and
+    `core.scope.bound`) in order to avoid introducing extra layers of nesting, while still
+    making sure that you only set the scale group for your context.
+    """
+    for scale_group_name in scale_group_names:
+        if not scale_group_name:
+            raise ValueError(f"Scale group name must be a non-empty string; got {scale_group_names}")
+
+    with _SCALE_GROUP_PRIORITY.set(
+        (*scale_group_names, *(sz for sz in _SCALE_GROUP_PRIORITY() if sz not in scale_group_names))
+    ):
+        yield _SCALE_GROUP_PRIORITY()
+
+
+def active_scale_groups() -> ty.Tuple[str, ...]:
+    """Returns the active scale groups, in order of priority."""
+    return _SCALE_GROUP_PRIORITY()