thds.mops 3.6.20250219172032__py3-none-any.whl

thds/mops/pure/core/uris.py

@@ -0,0 +1,81 @@
+import importlib.metadata
+import io
+import typing as ty
+from pathlib import Path
+from typing import Callable, Union
+
+from thds.adls import AdlsFqn, AdlsRoot
+from thds.core.stack_context import StackContext
+
+from ..adls.blob_store import get_adls_blob_store
+from .file_blob_store import get_file_blob_store
+from .types import BlobStore
+
+GetBlobStoreForUri = ty.Callable[[str], ty.Optional[BlobStore]]
+
+
+# we add the ADLS blob store and FileBlobStore here because they are the 'blessed'
+# implementations that our internal users rely on.
+# Others can be registered via entry-points.
+_REGISTERED_BLOB_STORES: ty.List[GetBlobStoreForUri] = [
+    get_file_blob_store,
+    get_adls_blob_store,
+]
+
+
+def register_blob_store(get_store: GetBlobStoreForUri) -> None:
+    """Dynamically register a BlobStore implementation."""
+    _REGISTERED_BLOB_STORES.append(get_store)
+
+
+def load_plugin_blobstores() -> None:
+    for entry_point in importlib.metadata.entry_points().get("thds.mops.pure.blob_stores", []):
+        try:
+            register_blob_store(entry_point.load())
+        except Exception as e:
+            print(f"Error loading entry point {entry_point.name}: {e}")
+
+
+def lookup_blob_store(uri: str) -> BlobStore:
+    for get_store in _REGISTERED_BLOB_STORES[::-1]:
+        if store := get_store(uri):
+            return store
+    raise ValueError(f"Unsupported URI: {uri}")
+
+
+def get_root(uri: str) -> str:
+    blob_store = lookup_blob_store(uri)
+    return blob_store.control_root(uri)
+
+
+UriIsh = Union[AdlsRoot, AdlsFqn, str, Path]
+UriResolvable = Union[UriIsh, Callable[[], UriIsh]]
+
+
+def to_lazy_uri(resolvable: UriResolvable) -> Callable[[], str]:
+    if isinstance(resolvable, Path):
+        return lambda: str(resolvable.resolve())
+    if isinstance(resolvable, (str, AdlsRoot, AdlsFqn)):
+        return lambda: str(resolvable)
+    if callable(resolvable):
+        return lambda: str(resolvable())  # type: ignore
+    raise TypeError(type(resolvable))
+
+
+def get_bytes(remote_uri: str, type_hint: str) -> bytes:
+    blob_store = lookup_blob_store(remote_uri)
+    with io.BytesIO() as tb:
+        blob_store.readbytesinto(remote_uri, tb, type_hint=type_hint)
+        tb.seek(0)
+        return tb.read()
+
+
+# ACTIVE_STORAGE_ROOT is meant as a global, non-semantic URI prefix.
+# In other words, it should have nothing to do with your application
+ACTIVE_STORAGE_ROOT: StackContext[str] = StackContext("ACTIVE_STORAGE_ROOT", "")
+# objects referencing this StackContext must be used in the same thread where they were created.
+
+
+def active_storage_root() -> str:
+    assert ACTIVE_STORAGE_ROOT(), "ACTIVE_STORAGE_ROOT must be set before use."
+    return ACTIVE_STORAGE_ROOT()

thds/mops/pure/core/use_runner.py

@@ -0,0 +1,47 @@
+"""use_runner wrapper decorator factory lives here.
+
+You can transfer control to a Runner without this, but decorators are a Pythonic approach.
+"""
+
+import typing as ty
+from functools import wraps
+
+from thds.core import log, stack_context
+
+from .entry.runner_registry import entry_count
+from .types import Runner
+
+logger = log.getLogger(__name__)
+F = ty.TypeVar("F", bound=ty.Callable)
+FUNCTION_UNWRAP_COUNT = stack_context.StackContext("function_unwrap_count", 0)
+
+
+def _is_runner_entry() -> bool:
+    """Function is being called in the context of a Runner."""
+    return entry_count() > FUNCTION_UNWRAP_COUNT()
+
+
+def use_runner(runner: Runner, skip: ty.Callable[[], bool] = lambda: False) -> ty.Callable[[F], F]:
+    """Wrap a function that is pure with respect to its arguments and result.
+
+    Run that function on the provided runner.
+
+    The arguments must be able to be transmitted by the runner to the
+    remote context and not refer to anything that will not be
+    accessible in that context.
+    """
+
+    def deco(f: F) -> F:
+        @wraps(f)
+        def __use_runner_wrapper(*args, **kwargs):  # type: ignore
+            if _is_runner_entry() or skip():
+                logger.debug("Calling function %s directly...", f)
+                with FUNCTION_UNWRAP_COUNT.set(FUNCTION_UNWRAP_COUNT() + 1):
+                    return f(*args, **kwargs)
+
+            logger.debug("Forwarding local function %s call to runner...", f)
+            return runner(f, args, kwargs)
+
+        return ty.cast(F, __use_runner_wrapper)
+
+    return deco

thds/mops/pure/joblib/__init__.py

@@ -0,0 +1 @@
+from .backend import TrillMlParallelJoblibBackend  # noqa

thds/mops/pure/joblib/backend.py

@@ -0,0 +1,81 @@
+"""A mops.pure powered joblib backend.
+
+Additionally, several hacks to work around either bugs or bad behavior in joblib.
+
+mops itself does _not_ force a dependency on joblib.
+
+If you want to use this supplied helper class, you must ensure that
+Joblib is installed, or you will get an ImportError.
+"""
+
+import typing as ty
+
+from joblib._parallel_backends import LokyBackend, SequentialBackend, ThreadingBackend  # type: ignore
+
+from thds.core import files, log
+
+from ..core.types import Runner
+from .batching import patch_joblib_parallel_batching  # noqa
+
+logger = log.getLogger(__name__)
+
+
+class NoMemmapLokyBackend(LokyBackend):
+    """Workaround for joblib/Cython bug exposed via SciKit-Learn.
+
+    https://github.com/scikit-learn/scikit-learn/issues/7981#issuecomment-341879166
+
+    Since this only runs on the remotes, and only then a few
+    CPUs/processes at a time, the extra memcopies being done are
+    no big deal.
+    """
+
+    def configure(self, *args: ty.Any, **kwargs: ty.Any) -> None:
+        kwargs["max_nbytes"] = None
+        return super().configure(*args, **kwargs)
+
+
+class TrillMlParallelJoblibBackend(ThreadingBackend):
+    """A joblib backend that forwards to our MemoizingPicklingRunner.
+
+    Performs simple batching based on the parallelism and oversubscribe factors at construction time.
+
+    Note that you'll likely need to customize pre_dispatch, otherwise
+    there won't be enough created tasks to actually batch anything.
+    """
+
+    supports_sharedmem = False
+    uses_threads = False
+
+    def __init__(
+        self,
+        runner: Runner,
+        parallelism: int,
+        n_cores: int,
+        oversubscribe: int = 10,
+    ):
+        """number of cores should be the number of cores available on the remote system."""
+        files.bump_limits()
+        self.runner = runner
+        self.n_cores = n_cores
+        self._n_jobs = parallelism
+        self.oversubscribe = oversubscribe
+
+    def effective_n_jobs(self, _nj: int) -> int:
+        return self._n_jobs
+
+    def compute_batch_size(self) -> int:
+        return self.n_cores * self.oversubscribe
+
+    def apply_async(self, func: ty.Any, callback: ty.Any = None) -> ty.Any:
+        def call_in_runner() -> ty.Any:
+            return self.runner(func, (), dict())
+
+        return super().apply_async(call_in_runner, callback=callback)
+
+    def get_nested_backend(self) -> ty.Any:
+        nesting_level = getattr(self, "nesting_level", 0) + 1
+        if nesting_level > 1:
+            logger.warning("Using sequential backend")
+            return SequentialBackend(nesting_level=nesting_level), None
+        return NoMemmapLokyBackend(nesting_level=nesting_level), self.n_cores

thds/mops/pure/joblib/batching.py

@@ -0,0 +1,67 @@
+"""Fix joblib batching for use with many parallel tasks running remotely."""
+
+import itertools
+import queue
+import typing as ty
+import unittest.mock
+from collections import defaultdict
+from contextlib import contextmanager
+
+import joblib  # type: ignore
+from joblib.parallel import BatchedCalls  # type: ignore
+
+from thds.core.log import getLogger
+
+logger = getLogger(__name__)
+
+
+def dispatch_one_batch(self: ty.Any, iterator: ty.Iterable[ty.Any]) -> bool:
+    """Joblib batching is truly horrible for running on a remote machine.
+
+    Various things conspire to essentially try to outsmart your
+    backend's batching instructions, and you end up with much smaller
+    batches than desired if you try to launch lots of runtimes in
+    parallel.
+
+    This is an ugly monkey patch, but it _works_.
+    """
+    if not hasattr(self, "__patch_stats"):
+        self.__patch_stats = defaultdict(int)  # type: ignore
+
+    batch_size = self._backend.compute_batch_size()
+    with self._lock:
+        try:
+            tasks = self._ready_batches.get(block=False)
+        except queue.Empty:
+            n_jobs = self._cached_effective_n_jobs
+            islice = list(itertools.islice(iterator, batch_size * n_jobs))
+            if len(islice) == 0:
+                return False
+            self.__patch_stats["tasks"] += len(islice)
+            logger.info(
+                f"Creating new tasks with patched batch size {batch_size}; "
+                f"stats so far: {self.__patch_stats}"
+            )
+            for i in range(0, len(islice), batch_size):
+                self._ready_batches.put(
+                    BatchedCalls(
+                        islice[i : i + batch_size],
+                        self._backend.get_nested_backend(),
+                        self._reducer_callback,
+                        self._pickle_cache,
+                    )
+                )
+            # finally, get one task.
+            tasks = self._ready_batches.get(block=False)
+
+        if len(tasks) == 0:
+            return False
+        self.__patch_stats["batches"] += 1
+        self._dispatch(tasks)
+        return True
+
+
+@contextmanager
+def patch_joblib_parallel_batching() -> ty.Iterator[None]:
+    with unittest.mock.patch.object(joblib.Parallel, "dispatch_one_batch", dispatch_one_batch):
+        yield

thds/mops/pure/pickling/__init__.py

@@ -0,0 +1,3 @@
+from . import mprunner, remote  # noqa: F401
+from .memoize_only import memoize_in  # noqa: F401
+from .mprunner import MemoizingPicklingRunner  # noqa: F401

thds/mops/pure/pickling/_pickle.py

@@ -0,0 +1,193 @@
+"""Utilities built around pickle for the purpose of transferring large amounts of on-disk
+data and also functions."""
+
+import inspect
+import io
+import pickle
+import typing as ty
+from functools import partial
+
+# so we can pickle and re-raise exceptions with remote tracebacks
+from tblib import pickling_support  # type: ignore
+
+from thds.core import hashing, log, source
+
+from ..core import memo, metadata
+from ..core.source import prepare_source_argument, prepare_source_result
+from ..core.types import Args, Deserializer, Kwargs, SerializerHandler
+from ..core.uris import get_bytes
+from .pickles import (
+    PicklableFunction,
+    UnpickleFunctionWithLogicKey,
+    UnpickleSourceHashrefArgument,
+    UnpickleSourceResult,
+    UnpickleSourceUriArgument,
+)
+
+logger = log.getLogger(__name__)
+F = ty.TypeVar("F", bound=ty.Callable)
+
+
+def wrap_f(f: F) -> ty.Union[F, PicklableFunction]:
+    if hasattr(f, "__module__") and hasattr(f, "__name__"):
+        return PicklableFunction(f)
+    return f
+
+
+class _CallbackPickler(pickle.Pickler):
+    def __init__(self, handlers: ty.Sequence[SerializerHandler], *args: ty.Any, **kwargs: ty.Any):
+        super().__init__(*args, **kwargs)
+        self.handlers = handlers
+
+    def persistent_id(self, obj: ty.Any) -> ty.Union[None, ty.Callable]:
+        if isinstance(obj, Exception):
+            pickling_support.install(obj)
+        for handler in self.handlers:
+            pid = handler(obj)
+            if pid is not None:
+                return pid
+        return None
+
+
+class CallableUnpickler(pickle.Unpickler):
+    """Present same interface as pickle.load but support unpickling callable PIDs generated by CallbackPickler."""
+
+    def persistent_load(self, pid: ty.Callable) -> ty.Any:
+        try:
+            return pid()
+        except TypeError as te:
+            # logger.exception("TypeError hit while debugging")
+            # this line should never get hit as long as nobody asks us to unpickle PIDs we don't know about
+            raise pickle.UnpicklingError(f"unsupported persistent object - {te}")  # pragma: no cover
+
+
+class Dumper:
+    """Presents the same interface as pickle.dump but supports
+    arbitrary callback-based unpickling.
+    """
+
+    def __init__(self, *handlers: SerializerHandler):
+        self.handlers = handlers
+
+    def __call__(self, obj: object, file: ty.IO, *args: ty.Any, **kwargs: ty.Any) -> ty.Any:
+        _CallbackPickler(self.handlers, file, *args, **kwargs).dump(obj)
+
+
+def gimme_bytes(pickle_dump: ty.Callable[[object, ty.IO], None], obj: object) -> bytes:
+    with io.BytesIO() as bio:
+        pickle_dump(obj, bio)
+        bio.seek(0)
+        return bio.read()
+
+
+def read_partial_pickle(full_bytes: bytes) -> ty.Tuple[bytes, ty.Any]:
+    # in order to be forward-compatible with v3 of mops, we're introducing a new
+    # wrinkle in the read. Instead of assuming that the data at the URI
+    # _begins_ with a pickle, we are looking for the first possible pickle
+    # and beginning our read there. Mops 3 will be generating some human-readable,
+    # non-pickle metadata and embedding it at the beginning of the file.
+    first_pickle_pos = full_bytes.find(b"\x80")
+    if first_pickle_pos == -1:
+        raise ValueError("Unable to find a pickle in the bytes")
+    return (
+        full_bytes[:first_pickle_pos],
+        CallableUnpickler(io.BytesIO(full_bytes[first_pickle_pos:])).load(),
+    )
+
+
+H = ty.TypeVar("H")
+
+
+def make_read_header_and_object(
+    type_hint: str, xf_header: ty.Optional[ty.Callable[[bytes], H]] = None
+) -> ty.Callable[[str], ty.Tuple[H, ty.Any]]:
+    def read_object(uri: str) -> ty.Tuple[H, ty.Any]:
+        header, unpickled = read_partial_pickle(get_bytes(uri, type_hint=type_hint))
+        return (xf_header or (lambda h: h))(header), unpickled  # type: ignore
+
+    return read_object
+
+
+def read_metadata_and_object(
+    type_hint: str, uri: str
+) -> ty.Tuple[ty.Optional[metadata.ResultMetadata], ty.Any]:
+    def _read_metadata_header(header_bytes: bytes) -> ty.Optional[metadata.ResultMetadata]:
+        if not header_bytes:
+            return None
+        return metadata.parse_result_metadata(header_bytes.decode("utf-8").split("\n"))
+
+    return make_read_header_and_object(type_hint, xf_header=_read_metadata_header)(uri)
+
+
+def freeze_args_kwargs(dumper: Dumper, f: ty.Callable, args: Args, kwargs: Kwargs) -> bytes:
+    """Returns a pickled (args, kwargs) tuple, with pre-bound
+    arguments to normalize different call structures into a
+    canonical/determinstic binding.
+
+    Also binds default arguments, for maximum determinism/explicitness.
+    """
+    bound_arguments = inspect.signature(f).bind(*args, **kwargs)
+    bound_arguments.apply_defaults()
+    return gimme_bytes(dumper, (bound_arguments.args, bound_arguments.kwargs))
+
+
+def unfreeze_args_kwargs(
+    args_kwargs_pickle: bytes, unpickler: ty.Type[pickle.Unpickler] = CallableUnpickler
+) -> ty.Tuple[Args, Kwargs]:
+    """Undoes a freeze_args_kwargs call."""
+    return unpickler(io.BytesIO(args_kwargs_pickle)).load()
+
+
+# SerializerHandlers for Source objects:
+_DeserSource = ty.Callable[[], source.Source]
+
+
+class SourceArgumentPickler:
+    """Only for use on the orchestrator side, when serializing the arguments."""
+
+    def __call__(self, maybe_source: ty.Any) -> ty.Optional[_DeserSource]:
+        if isinstance(maybe_source, source.Source):
+            uri_or_hash = prepare_source_argument(maybe_source)
+            if isinstance(uri_or_hash, hashing.Hash):
+                return ty.cast(_DeserSource, UnpickleSourceHashrefArgument(uri_or_hash))
+            return ty.cast(_DeserSource, UnpickleSourceUriArgument(uri_or_hash))
+            # I do not understand why these casts are necessary to avoid mypy errors.
+            # I think it has something to do with NamedTuples being the underlying
+            # object type that is expected to support __call__(self) -> Foo,
+            # but I haven't recently located a relevant Issue anywhere.
+        return None
+
+
+class SourceResultPickler:
+    """Only for use on the remote side, when serializing the result."""
+
+    def __call__(self, maybe_source: ty.Any) -> ty.Optional[_DeserSource]:
+        if isinstance(maybe_source, source.Source):
+            return ty.cast(_DeserSource, UnpickleSourceResult(*prepare_source_result(maybe_source)))
+
+        return None
+
+
+class NestedFunctionWithLogicKeyPickler:
+    def __call__(self, maybe_function_with_logic_key: ty.Any) -> ty.Optional[Deserializer]:
+        """Returns a pickle 'persistent id' which is a 'kind' of CallableUnpickler.
+
+        ...or None, which means pickle normally.
+        """
+        if not callable(maybe_function_with_logic_key):
+            return None
+
+        if isinstance(maybe_function_with_logic_key, partial):
+            # do not extract from the partial - only a raw function
+            # which will itself be included in the partial when it gets pickled
+            return None
+
+        function_logic_key = memo.extract_function_logic_key_from_docstr(maybe_function_with_logic_key)
+        if not function_logic_key:
+            return None
+
+        return UnpickleFunctionWithLogicKey(  # type: ignore
+            # we must then wrap the function itself so that this does not cause infinite recursion.
+            pickle.dumps(maybe_function_with_logic_key),
+            function_logic_key,
+        )

thds/mops/pure/pickling/memoize_only.py

@@ -0,0 +1,22 @@
+from typing import Callable
+
+from ..core.types import F
+from ..core.uris import UriResolvable
+from ..core.use_runner import use_runner
+from ..runner import simple_shims
+from .mprunner import MemoizingPicklingRunner
+
+
+# this may soon become deprecated in favor of mops.pure.magic(blob_root=...)
+def memoize_in(uri_resolvable: UriResolvable) -> Callable[[F], F]:
+    """A decorator that makes a function globally-memoizable, but running in the current
+    thread.
+
+    This is a good thing to use when the computation derives no
+    advantage from not running locally (i.e. exhibits no parallelism)
+    but you still want memoization.
+
+    This enables nested memoized function calls, which is not (yet)
+    the default for `use_runner`.
+    """
+    return use_runner(MemoizingPicklingRunner(simple_shims.samethread_shim, uri_resolvable))

thds/mops/pure/pickling/mprunner.py

@@ -0,0 +1,173 @@
+"""Provides concrete serialization/deserialization (via pickling) for the basic memoizing runner algorithm.
+
+Contains default config, core 'state' and some rarely-used customization interfaces.
+
+See runner.local.py for the core runner implementation.
+"""
+
+import typing as ty
+from functools import partial
+
+from thds.core import cache, log
+from thds.core.stack_context import StackContext
+
+from ..._utils.once import Once
+from ..core import memo, uris
+from ..core.serialize_big_objs import ByIdRegistry, ByIdSerializer
+from ..core.serialize_paths import CoordinatingPathSerializer
+from ..core.types import Args, F, Kwargs, Serializer, T
+from ..runner import local, shim_builder
+from ..runner.types import Shim, ShimBuilder
+from ..tools.summarize import run_summary
+from . import _pickle, pickles, sha256_b64
+
+RUNNER_NAME = "mops2-mpf"
+Redirect = ty.Callable[[F, Args, Kwargs], F]
+NO_REDIRECT = lambda f, _args, _kwargs: f  # noqa: E731
+_ARGS_CONTEXT = StackContext[ty.Sequence]("args_kwargs", tuple())
+_KWARGS_CONTEXT = StackContext[ty.Mapping]("args_kwargs", dict())
+logger = log.getLogger(__name__)
+
+
+def mp_shim(base_shim: Shim, shim_args: ty.Sequence[str]) -> ty.Any:
+    return base_shim((RUNNER_NAME, *shim_args))
+
+
+def _runner_prefix_for_pickled_functions(storage_root: str) -> str:
+    return uris.lookup_blob_store(storage_root).join(storage_root, RUNNER_NAME)
+
+
+class MemoizingPicklingRunner:
+    """
+    Runs callables in a process as defined by the Shim.
+    This is often a remote process, however a local shim may be provided.
+    """
+
+    def __init__(
+        self,
+        shim: ty.Union[ShimBuilder, Shim],
+        blob_storage_root: uris.UriResolvable,
+        *,
+        rerun_exceptions: bool = True,
+        serialization_registry: ByIdRegistry[ty.Any, Serializer] = ByIdRegistry(),  # noqa: B008
+        redirect: Redirect = NO_REDIRECT,
+    ):
+        """Construct a memoizing shim runner.
+
+        Transmitted Path resources will be content-hash-addressed
+        below the runner_prefix to save storage and increase chances
+        of memoization. Named objects will be treated
+        similarly. Function invocations will be pickled and stored
+        under the current pipeline id since we do not have a way of
+        inferring whether their associated code is safely content-addressable
+        across runs.
+
+        The Shim must forward control in the remote environment to a
+        wrapper that will pull the function and arguments from the URI(s).
+
+        A ShimBuilder will receive the original function and its
+        original arguments, which you can use to determine which
+        concrete Shim implementation to return for the given function
+        call.
+
+        `rerun_exceptions` will cause a pre-existing `exception`
+        result to be ignored, as though Exceptions in your function
+        are the result of transient errors and not an expected return
+        value of a (simulated) pure function. If you do not want this
+        behavior, turn it off.
+
+        `redirect` changes only the function that is actually invoked
+        on the remote side of the runner. It does not change the
+        computed memoization key, which is still based on the original
+        function and the args, kwargs pair passed in. A common use for
+        this would be allowing a contextually-aware function to be
+        invoked in the manner of initializer/initargs, without those
+        additional bits being part of the function invocation and
+        therefore the memoization key, especially where they're not
+        picklable at all.
+        """
+        self._shim_builder = shim_builder.make_builder(shim)
+        self._get_storage_root = uris.to_lazy_uri(blob_storage_root)
+        self._rerun_exceptions = rerun_exceptions
+        self._by_id_registry = serialization_registry
+        self._redirect = redirect
+
+        self._run_directory = run_summary.create_mops_run_directory()
+
+    def shared(self, *objs: ty.Any, **named_objs: ty.Any) -> None:
+        """Set up memoizing pickle serialization for these objects.
+
+        Provided names are used for debugging purposes only.
+        """
+        for obj in objs:
+            self._by_id_registry[obj] = sha256_b64.Sha256B64Pickler()
+        for name, obj in named_objs.items():
+            self._by_id_registry[obj] = sha256_b64.Sha256B64Pickler(name)
+
+    @cache.locking
+    def _get_stateful_dumper(self, _root: str) -> _pickle.Dumper:
+        """We want one of these per blob storage root, because the
+        invocation and result must exist on the same blob store as
+        any other automatically dumped objects, e.g. Paths or named
+        objects, such that the full invocation payload is
+        byte-for-byte identical, since its hash is our memoization
+        key.
+        """
+        return _pickle.Dumper(
+            ByIdSerializer(self._by_id_registry),
+            CoordinatingPathSerializer(sha256_b64.Sha256B64PathStream(), Once()),
+            _pickle.SourceArgumentPickler(),
+            _pickle.NestedFunctionWithLogicKeyPickler(),
+        )
+
+    def _serialize_args_kwargs(
+        self, storage_root: str, func: ty.Callable[..., T], args: Args, kwargs: Kwargs
+    ) -> bytes:
+        # Why do we need func in order to serialize args and kwargs? Because
+        # we use it to bind the arguments to the function first, which makes that part
+        # deterministic and also 'reifies' any default arguments, so we don't have any implicit state.
+        return _pickle.freeze_args_kwargs(self._get_stateful_dumper(storage_root), func, args, kwargs)
+
+    def _serialize_invocation(
+        self, storage_root: str, func: ty.Callable[..., T], args_kwargs: bytes
+    ) -> bytes:
+        return _pickle.gimme_bytes(
+            self._get_stateful_dumper(storage_root),
+            pickles.Invocation(
+                _pickle.wrap_f(self._redirect(func, _ARGS_CONTEXT(), _KWARGS_CONTEXT())),
+                args_kwargs,
+            ),
+        )
+
+    def _wrap_shim_builder(self, func: F, args: Args, kwargs: Kwargs) -> Shim:
+        base_shim = self._shim_builder(func, args, kwargs)
+        return partial(mp_shim, base_shim)
+
+    def __call__(self, func: ty.Callable[..., T], args: Args, kwargs: Kwargs) -> T:
+        """Return result of running this function remotely via the shim.
+
+        Passes data to shim process via pickles in a Blob Store.
+
+        May return cached (previously-computed) results found via the
+        derived function memo URI, which contains the determinstic
+        hashed bytes of all the function arguments, but also
+        additional namespacing including pipeline_id as documented
+        in memo.function_memospace.py.
+        """
+        logger.debug("Preparing to run function via remote shim")
+        with _ARGS_CONTEXT.set(args), _KWARGS_CONTEXT.set(kwargs):
+            return local.invoke_via_shim_or_return_memoized(
+                self._serialize_args_kwargs,
+                self._serialize_invocation,
+                self._wrap_shim_builder,
+                _pickle.read_metadata_and_object,
+                self._run_directory,
+            )(
+                self._rerun_exceptions,
+                memo.make_function_memospace(
+                    _runner_prefix_for_pickled_functions(self._get_storage_root()), func
+                ),
+                func,
+                args,
+                kwargs,
+            )