thds.mops 3.6.20250219172032__py3-none-any.whl

thds/mops/__about__.py

@@ -0,0 +1,8 @@
+from thds.core import meta
+
+__version__ = meta.get_version("thds.mops")
+__commit__ = meta.read_metadata(__name__).git_commit
+
+
+def backward_compatible_with() -> int:
+    return 2  # v2 is the current major version.

thds/mops/__init__.py

@@ -0,0 +1,3 @@
+from . import parallel  # noqa
+from .__about__ import __commit__, __version__  # noqa
+from ._utils.temp import tempdir  # noqa

thds/mops/_compat.py

@@ -0,0 +1,6 @@
+# compatibility shims
+
+try:
+    import tomllib  # type: ignore [import-not-found] # noqa: F401
+except ImportError:
+    import tomli as tomllib  # noqa: F401

thds/mops/_utils/colorize.py

@@ -0,0 +1,110 @@
+import itertools
+import random
+import typing as ty
+from functools import partial
+
+from colors import color, csscolors
+
+pref = "\033["
+reset = f"{pref}0m"
+
+_RESERVED_COLORS = [
+    "black",
+    # Various whitish-looking colors
+    "aliceblue",
+    "antiquewhite",
+    "floralwhite",
+    "ghostwhite",
+    "ivory",
+    "white",
+    "whitesmoke",
+    "snow",
+    "seashell",
+    "mintcream",
+    "honeydew",
+    "azure",
+    "beige",
+    "cornsilk",
+    "floralwhite",
+    # These are pretty illegible on a black background
+    "darkblue",
+    "indigo",
+    "mediumblue",
+    "navy",
+    "purple",
+]
+
+_PREFERRED_COLORS = [
+    "mediumseagreen",
+    "cornflowerblue",
+    "gold",
+    "salmon",
+    "violet",
+    "limegreen",
+    "dodgerblue",
+    "goldenrod",
+    "indianred",
+    "fuchsia",
+    "forestgreen",
+    "royalblue",
+    "yellow",
+    "chocolate",
+    "palevioletred",
+    "mediumspringgreen",
+    "deepskyblue",
+    "khaki",
+    "red",
+    "deeppink",
+    "seagreen",
+    "cyan",
+    "greenyellow",
+    "sandybrown",
+    "orchid",
+    "lightgreen",
+    "steelblue",
+    "darkgoldenrod",
+    "coral",
+    "darkorchid",
+]
+
+
+def _start_from(color_list: ty.List[str], index: int) -> ty.List[str]:
+    return color_list[index:] + color_list[:index]
+
+
+def _preferred_randgreen_start() -> ty.List[str]:
+    return _start_from(_PREFERRED_COLORS, random.randint(0, 6) * 5)
+
+
+def _all_colors() -> ty.List[str]:
+    forbidden_colors = {csscolors.css_colors[name] for name in _RESERVED_COLORS}
+    used_colors = {csscolors.css_colors[name] for name in _PREFERRED_COLORS}
+    assert len(used_colors) == len(_PREFERRED_COLORS)  # assert no RGB dupes in the preferred list
+    all_colors = list(csscolors.css_colors.items())
+    random.shuffle(all_colors)
+    return _preferred_randgreen_start() + [
+        name
+        for name, rgb in all_colors
+        if rgb not in used_colors
+        and not used_colors.add(rgb)  # type: ignore
+        and rgb not in forbidden_colors
+    ]
+
+
+next_color = ty.cast(ty.Callable[[], str], partial(next, itertools.cycle(_all_colors())))
+
+
+def colorized(fg: str, bg: str = "", style: str = "") -> ty.Callable[[str], str]:
+    def colorize(s: str) -> str:
+        return color(s, fg=fg, bg=bg, style=style)
+
+    return colorize
+
+
+def make_colorized_out(
+    colorized: ty.Callable[[str], str], *, fmt_str: str = "{}", out: ty.Callable[[str], ty.Any] = print
+) -> ty.Callable[[str], None]:
+    def _out(s: str) -> None:
+        out(colorized(fmt_str.format(s)))
+
+    return _out

thds/mops/_utils/config_tree.py

@@ -0,0 +1,167 @@
+import inspect
+import types
+import typing as ty
+from functools import partial
+
+from thds import core
+
+from .names import full_name_and_callable
+
+IGNORED_PACKAGES = ["thds.mops"]
+# if a library needs to build on top of mops, it can put itself into this
+# list and we'll ignore it when looking for the 'true calling frame'
+
+
+def _get_first_external_module(ignore_packages: ty.Collection[str] = IGNORED_PACKAGES) -> str:
+    frame = inspect.currentframe()
+    if not frame:
+        return ""
+
+    while frame := frame.f_back:  # type: ignore
+        module_name = frame.f_globals["__name__"]
+        is_ignored = False
+        for ignore_package in ignore_packages:
+            if module_name.startswith(ignore_package):
+                is_ignored = True
+                break
+
+        if not is_ignored:
+            return module_name
+    return ""  # fallback if no external caller found
+
+
+Pathable = ty.Union[str, types.ModuleType, ty.Callable, None]
+_NONE = object()
+V = ty.TypeVar("V")
+logger = core.log.getLogger(__name__)
+
+
+def to_dotted_path(pathable: Pathable) -> str:
+    if isinstance(pathable, str):
+        return pathable
+
+    if pathable is None:
+        if not (module_path := _get_first_external_module()):
+            raise ValueError(f"Found no module outside mops within {pathable}")
+        return module_path
+
+    if isinstance(pathable, types.ModuleType):
+        return pathable.__name__
+
+    return full_name_and_callable(pathable)[0].replace("--", ".")
+
+
+class ConfigTree(ty.Generic[V]):
+    """This is a cute little utility class for applying homogeneously-typed configuration
+    following hierarchical (tree-like) paths.
+
+    Generally, the config closest to the 'leaf' path will be used, but there is also a
+    'mask' option to override subtrees.
+
+    Imagine you have some modules:
+    - foo.bar.baz.materialize
+    - foo.bar.quux.materialize
+    - foo.george.materialize
+    - foo.steve.materialize
+
+    in each of which you have several materialization functions using mops.
+
+    Some API might construct one of these objects to afford you a way to 'set' the config
+    at each level of your hierarchy.
+
+    Inside foo.bar.__init__.py, you could call
+
+    - `the_api.setv(a_config_object)`
+
+    and this would set the config for anything where the module path to it included foo.bar.
+
+    But if foo.bar.baz.materialize wanted to set config for everything inside itself,
+    at the top of that module you'd call
+
+    - `the_api.setv(diff_config_object)`
+
+    and this would set the config for that module only.
+
+    If you need to _override_ the config for an entire subtree, we call this masking.
+    You can call the_api.setv(value, mask=True) to mask the subtree.
+
+    - `the_api.setv(value, 'foo', mask=True)`
+
+    will mask everything under foo, including bar, george, and steve.
+
+    This isn't truly limited to modules, either - you can pass any module _or_ callable
+    in to setv as the object from which you want us to derive a dot-separated path.
+    Or you can pass in an arbitrary dot-separated string and we'll use it verbatim.
+    """
+
+    def __init__(self, name: str, parse: ty.Optional[ty.Callable[[ty.Any], V]] = None):
+        """If provided, parse must be an idempotent parser. In other words, parse(parse(x)) == parse(x)"""
+        self.registry = core.config.ConfigRegistry(name)
+        self.parse = parse or (lambda v: v)
+        self._make_config = partial(
+            core.config.ConfigItem[V], registry=self.registry, name_transform=lambda s: s, parse=parse
+        )
+
+    def getv(self, path: str, default: V = ty.cast(V, _NONE)) -> V:
+        parts = [*path.split(".")]
+        mask = "__mask"
+        for i in range(0, len(parts) + 1):
+            prefix = ".".join([mask, *parts[:i]])
+            # we do an 'in' check b/c the value might not be truthy, or even non-None
+            if prefix in self.registry:
+                return self.registry[prefix]()
+
+        # If not masked, fall back to normal hierarchical lookup
+        return self._get_most_specific_v(path, parts, default)
+
+    def _get_most_specific_v(
+        self, path: str, parts: ty.Sequence[str], default: V = ty.cast(V, _NONE)
+    ) -> V:
+        for i in range(len(parts), -1, -1):
+            prefix = ".".join(parts[:i])
+            # we do an 'in' check b/c the value might not be truthy, or even non-None
+            if prefix in self.registry:
+                return self.registry[prefix]()
+        assert prefix == ""
+
+        if default is not _NONE:
+            return default
+
+        name = self.registry.name
+        raise RuntimeError(f"No {name} configuration matches {path} and no global config was set")
+
+    def setv(
+        self, value: V, pathable: Pathable = None, *, mask: bool = False
+    ) -> core.config.ConfigItem[V]:
+        """Set the value for the given Pathable, or the current module if no Pathable is given.
+        By default, greater overlap in paths will supersede less overlap.
+
+        mask=True will override any 'more specific' config below it in the hierarchy.
+        """
+        config_path = to_dotted_path(pathable)
+        if mask:
+            config_path = ".".join(filter(None, ["__mask", config_path]))
+            log_msg = "Masking all [%s] config under '%s' with %s"
+        else:
+            log_msg = "Setting [%s] '%s' to %s"
+        logger.debug(log_msg, self.registry.name, config_path, value)
+        if config_item := self.registry.get(config_path):
+            config_item.set_global(self.parse(value))
+        else:
+            config_item = self._make_config(config_path, default=value)  # also registers the ConfigItem
+        return config_item
+
+    def __setitem__(self, key: str, value: V) -> None:
+        self.setv(value, pathable=key)
+
+    def load_config(self, config: ty.Mapping[str, ty.Any]) -> None:
+        """Loads things with an inner key matching this name into the config."""
+        mask_name = f".__mask.{self.registry.name}"
+        conf_name = f".{self.registry.name}"
+        logger.debug("Loading config for %s", self.registry.name)
+        for key, value in core.config.flatten_config(config).items():
+            if key.endswith(conf_name):
+                self.setv(value, key[: -len(conf_name)], mask=key.endswith(mask_name))
+
+    def __repr__(self) -> str:
+        return f"ConfigTree('{self.registry.name}', {list(self.registry.items())})"

thds/mops/_utils/exception.py

@@ -0,0 +1,16 @@
+import contextlib
+import typing as ty
+
+
+@contextlib.contextmanager
+def catch(allow: ty.Callable[[Exception], bool]) -> ty.Iterator:
+    """try-except but flexible. Catch only Exceptions matching the filter.
+
+    Useful for libraries like azure where all the Exceptions have the
+    same type.
+    """
+    try:
+        yield
+    except Exception as e:
+        if not allow(e):
+            raise

thds/mops/_utils/locked_cache.py

@@ -0,0 +1,78 @@
+import functools
+import typing as ty
+from threading import Lock, RLock
+from typing import Optional, Union
+
+from cachetools import keys
+
+try:
+    from cachetools.func import _CacheInfo  # type: ignore
+except ImportError:
+    # this moved between 5.2.1 and 5.3.
+    from cachetools import _CacheInfo  # type: ignore
+
+
+F = ty.TypeVar("F", bound=ty.Callable)
+
+
+def locked_cached(
+    cache: ty.Any, typed: bool = False, lock: Optional[Union[RLock, Lock]] = None
+) -> ty.Callable[[F], F]:
+    """Like cachetools.func._cache, except it locks the actual
+    function call but does _not_ lock reading from the cache the first
+    time, so most of the time, cache hits are nearly free, but you
+    don't call the function more than once for the same arguments.
+    """
+    maxsize = cache.maxsize
+
+    def decorator(func: F) -> F:
+        key = keys.typedkey if typed else keys.hashkey
+        hits = misses = 0
+        _lock = lock or RLock()
+
+        def wrapper(*args, **kwargs):  # type: ignore
+            nonlocal hits, misses
+            k = key(*args, **kwargs)
+
+            # optimistic lookup on a cache that is threadsafe for reads
+            try:
+                v = cache[k]
+                hits += 1
+                return v
+            except KeyError:
+                with _lock:
+                    try:
+                        v = cache[k]
+                        hits += 1
+                        return v
+                    except KeyError:
+                        misses += 1
+
+                    v = func(*args, **kwargs)
+                    # in case of a race, prefer the item already in the cache
+                    try:
+                        return cache.setdefault(k, v)
+                    except ValueError:
+                        return v  # value too large
+
+        def cache_info() -> _CacheInfo:
+            with _lock:
+                maxsize = cache.maxsize
+                currsize = cache.currsize
+            return _CacheInfo(hits, misses, maxsize, currsize)
+
+        def cache_clear() -> None:
+            nonlocal hits, misses
+            with _lock:
+                try:
+                    cache.clear()
+                finally:
+                    hits = misses = 0
+
+        wrapper.cache_info = cache_info  # type: ignore
+        wrapper.cache_clear = cache_clear  # type: ignore
+        wrapper.cache_parameters = lambda: {"maxsize": maxsize, "typed": typed}  # type: ignore
+        functools.update_wrapper(wrapper, func)
+        return ty.cast(F, wrapper)
+
+    return decorator

thds/mops/_utils/names.py

@@ -0,0 +1,23 @@
+import typing as ty
+
+
+def full_name_and_callable(func: ty.Any) -> ty.Tuple[str, ty.Callable]:
+    """return {module}--{name} for an actual (non-wrapped) function or class,
+    plus the unwrapped callable itself.
+    """
+    if hasattr(func, "func"):  # support functools.partial
+        return full_name_and_callable(func.func)
+
+    module = func.__module__
+    try:
+        name = func.__name__
+    except AttributeError:
+        try:
+            # for some reason, __name__ does not exist on instances of objects,
+            # nor does it exist as a 'member' of the __class__ attribute, but
+            # we can just pull it out directly like this for callable classes.
+            name = func.__class__.__name__
+        except AttributeError:
+            name = "MOPS_UNKNOWN_NAME"
+
+    return f"{module}--{name}", func

thds/mops/_utils/on_slow.py

@@ -0,0 +1,28 @@
+import typing as ty
+from functools import wraps
+from timeit import default_timer
+
+from thds.core import log
+
+from .colorize import colorized, make_colorized_out
+
+F = ty.TypeVar("F", bound=ty.Callable)
+logger = log.getLogger(__name__)
+_SLOW = colorized(fg="yellow", bg="black")
+LogSlow = make_colorized_out(_SLOW, out=logger.warning)
+
+
+def on_slow(callback: ty.Callable[[float], None], slow_seconds: float = 3.0) -> ty.Callable[[F], F]:
+    def deco(f: F) -> F:
+        @wraps(f)
+        def wrapper(*args, **kwargs):  # type: ignore
+            start_time = default_timer()
+            r = f(*args, **kwargs)
+            elapsed_s = default_timer() - start_time
+            if elapsed_s > slow_seconds:
+                callback(elapsed_s)
+            return r
+
+        return ty.cast(F, wrapper)
+
+    return deco

thds/mops/_utils/once.py

@@ -0,0 +1,30 @@
+import threading
+import typing as ty
+
+FNone = ty.TypeVar("FNone", bound=ty.Callable[[], None])
+
+
+class Once:
+    """Uses unique IDs to guarantee that an operation has only run
+    once in the lifetime of this object, and waits for it to be complete.
+
+    Is a potential source of memory leaks, since each event will be
+    stored until the entire Once object is disposed.
+    """
+
+    def __init__(self) -> None:
+        self.lock = threading.Lock()
+        self.events: ty.Dict[ty.Hashable, threading.Event] = dict()
+
+    def run_once(self, run_id: ty.Hashable, f: FNone) -> None:
+        needs_run = False
+        if run_id not in self.events:
+            with self.lock:
+                if run_id not in self.events:
+                    needs_run = True
+                    self.events[run_id] = threading.Event()
+        if needs_run:
+            f()
+            self.events[run_id].set()
+        else:
+            self.events[run_id].wait()

thds/mops/_utils/temp.py

@@ -0,0 +1,32 @@
+"""Utility to make returning files via their Paths less confusing in the application code."""
+import typing as ty
+from contextlib import contextmanager
+from pathlib import Path
+
+from thds.core import lazy, scope, tmp
+
+
+@contextmanager
+def _temp_dir() -> ty.Iterator[Path]:
+    with tmp.temppath_same_fs() as p:
+        p.mkdir()
+        yield p
+
+
+_FOREVER_SCOPE = scope.Scope("until_mops_exit")
+_SINGLE_REMOTE_TMP_DIR = lazy.Lazy(lambda: _FOREVER_SCOPE.enter(_temp_dir()))
+# there's really no obvious reason why you'd ever need more than one of these as long as
+# you're giving your actual output files names, so we create one as a global for general
+# use.
+
+
+def new_tempdir() -> Path:
+    return _FOREVER_SCOPE.enter(_temp_dir())
+
+
+def tempdir() -> Path:
+    """Lazily creates a global/shared temporary directory and returns it as a Path.
+
+    The files will get cleaned up when the interpreter exits.
+    """
+    return _SINGLE_REMOTE_TMP_DIR()

thds/mops/config.py

@@ -0,0 +1,60 @@
+import os
+import typing as ty
+from pathlib import Path
+
+from thds.core import config, log
+from thds.mops._compat import tomllib
+
+logger = log.getLogger(__name__)
+
+
+def find_first_upward_mops_toml() -> ty.Optional[Path]:
+    current = Path.cwd()
+    while True:
+        try:
+            mops_path = current / ".mops.toml"
+            if mops_path.is_file() and os.access(mops_path, os.R_OK):
+                return mops_path.resolve()
+            if current == current.parent:  # At root
+                return None
+            current = current.parent
+        except PermissionError:
+            return None
+
+
+def first_found_config_file() -> ty.Optional[Path]:
+    paths = [
+        Path(os.environ.get("MOPS_CONFIG", "")),
+        find_first_upward_mops_toml(),
+        Path(f"{Path.home()}/.mops.toml"),
+    ]
+    for path in paths:
+        if path and path.is_file():
+            return path
+    return None
+
+
+def load(config_file: ty.Optional[Path], name: str = "mops") -> ty.Dict[str, ty.Any]:
+    if config_file:
+        logger.debug("Loading %s config from %s", name, config_file)
+        return tomllib.load(open(config_file, "rb"))
+    return dict()
+
+
+max_concurrent_network_ops = config.item("mops.max_concurrent_network_ops", 8, parse=int)
+# 8 clients has been obtained experimentally via the `stress_test`
+# application running on a Mac M1 laptop running 200 parallel 5 second
+# tasks, though no significant difference was obtained between 5 and
+# 20 clients. Running a similar stress test from your orchestrator may
+# be a good idea if you are dealing with hundreds of micro (<20
+# second) remote tasks.
+
+open_files_limit = config.item("mops.resources.max_open_files", 10000)
+
+
+def _filter_to_known_mops_config(config: ty.Dict[str, ty.Any]) -> ty.Dict[str, ty.Any]:
+    return {k: v for k, v in config.items() if k.startswith("mops.") or k.startswith("thds.mops")}
+
+
+# load this after creating the config items
+config.set_global_defaults(_filter_to_known_mops_config(load(first_found_config_file())))

thds/mops/impure/__init__.py

@@ -0,0 +1,2 @@
+from .keyfunc import nil_args  # noqa
+from .runner import KeyedLocalRunner  # noqa

thds/mops/impure/keyfunc.py

@@ -0,0 +1,14 @@
+"""Impure keyfunctions that are useful for common cases."""
+import typing as ty
+
+from ..pure.core.memo.keyfunc import Args, Keyfunc, Kwargs
+from ..pure.core.memo.overwrite_params import parameter_overwriter
+
+
+def nil_args(*named_parameters: str) -> Keyfunc:
+    def nil_args_impure_keyfunc(
+        c: ty.Callable, args: Args, kwargs: Kwargs
+    ) -> ty.Tuple[ty.Callable, Args, Kwargs]:
+        return c, *parameter_overwriter(c, {name: None for name in named_parameters})(args, kwargs)
+
+    return nil_args_impure_keyfunc

thds/mops/impure/runner.py

@@ -0,0 +1,73 @@
+"""Builds on top of the pure.MemoizingPicklingRunner to provide
+impure, customizable memoization.
+"""
+
+import typing as ty
+
+from typing_extensions import ParamSpec
+
+from thds.core import log
+from thds.core.stack_context import StackContext
+
+from ..pure.core.memo.keyfunc import ArgsOnlyKeyfunc, Keyfunc, autowrap_args_only_keyfunc
+from ..pure.core.types import Args, Kwargs
+from ..pure.core.uris import UriResolvable
+from ..pure.pickling.mprunner import NO_REDIRECT, MemoizingPicklingRunner, Redirect
+from ..pure.runner.simple_shims import samethread_shim
+
+logger = log.getLogger(__name__)
+
+
+R = ty.TypeVar("R")
+P = ParamSpec("P")
+F_Args_Kwargs = ty.Tuple[ty.Callable, Args, Kwargs]
+_ORIGINAL_F_ARGS_KWARGS: StackContext[ty.Optional[F_Args_Kwargs]] = StackContext("f_args_kwargs", None)
+
+
+def _perform_original_invocation(*_args: ty.Any, **_kwargs: ty.Any) -> ty.Any:
+    f_args_kwargs = _ORIGINAL_F_ARGS_KWARGS()
+    assert (
+        f_args_kwargs is not None
+    ), "_perform_original_invocation() must be called from within a runner"
+    f, args, kwargs = f_args_kwargs
+    return f(*args, **kwargs)
+
+
+class KeyedLocalRunner(MemoizingPicklingRunner):
+    """The only purpose for using this is to reify/memoize your results.
+
+    Allows changing the memoization key, at the expense of
+    (theoretical) purity, since now we're memoizing on something you
+    made up, rather than something directly derived from the full set
+    of arguments passed to your function.
+
+    When the 'remote' side is reached, the original (args, kwargs)
+    will be passed to the result of change_function, or the original
+    function if change_function is the default (identity).
+
+    This runs the 'remote' function in the same process - your
+    function, if no memoized result is found, will execute in the same
+    thread where it was originally called. This runner will use the
+    return values of change_key_elements _only_ for the purposes of
+    keying the cache.
+    """
+
+    def __init__(
+        self,
+        blob_storage_root: UriResolvable,
+        *,
+        keyfunc: ty.Union[ArgsOnlyKeyfunc, Keyfunc],
+        redirect: Redirect = NO_REDIRECT,
+    ):
+        self._impure_keyfunc = autowrap_args_only_keyfunc(keyfunc)
+        self._pre_pickle_redirect = redirect
+        super().__init__(
+            samethread_shim,
+            blob_storage_root,
+            redirect=lambda _f, _args, _kwargs: _perform_original_invocation,
+        )
+
+    def __call__(self, raw_func: ty.Callable[P, R], raw_args: P.args, raw_kwargs: P.kwargs) -> R:
+        actual_function_to_call = self._pre_pickle_redirect(raw_func, raw_args, raw_kwargs)
+        with _ORIGINAL_F_ARGS_KWARGS.set((actual_function_to_call, raw_args, raw_kwargs)):
+            return super().__call__(*self._impure_keyfunc(raw_func, raw_args, raw_kwargs))