thds.mops 3.9.20250722163657__py3-none-any.whl → 3.9.20250722164625__py3-none-any.whl

thds/mops/pure/core/lock/write.py

@@ -1,5 +1,6 @@
 import os
 import typing as ty
+from dataclasses import dataclass
 from datetime import datetime, timedelta
 
 from thds.core import hostname, log
@@ -10,38 +11,37 @@ from .types import LockContents
 logger = log.getLogger(__name__)
 
 
-def make_lock_contents(
-    writer_id: str, expire: timedelta
-) -> ty.Callable[[ty.Optional[datetime]], LockContents]:
-    """Impure - Resets written_at to 'right now' to keep the lock 'live'."""
-    write_count = 0
-    first_written_at = ""
+@dataclass
+class LockEmitter:
+    writer_id: str
+    expire: timedelta
 
-    assert (
-        "/" not in writer_id
-    ), f"{writer_id} should not contain a slash - maybe you passed a URI instead?"
+    write_count: int = 0
+    first_written_at: str = ""
 
-    def lock_contents(first_acquired_at: ty.Optional[datetime]) -> LockContents:
-        nonlocal write_count, first_written_at
-        write_count += 1
+    def __post_init__(self) -> None:
+        assert (
+            "/" not in self.writer_id
+        ), f"{self.writer_id} should not contain a slash - maybe you passed a URI instead?"
+
+    def __call__(self, first_acquired_at: ty.Optional[datetime]) -> LockContents:
+        self.write_count += 1
         now = _funcs.utc_now().isoformat()
-        first_written_at = first_written_at or now
+        self.first_written_at = self.first_written_at or now
 
         return {
-            "writer_id": writer_id,
+            "writer_id": self.writer_id,
             "written_at": now,
-            "expire_s": expire.total_seconds(),
+            "expire_s": self.expire.total_seconds(),
             # debug stuff:
-            "write_count": write_count,
+            "write_count": self.write_count,
             "hostname": hostname.friendly(),
             "pid": str(os.getpid()),
-            "first_written_at": first_written_at,
+            "first_written_at": self.first_written_at,
             "first_acquired_at": first_acquired_at.isoformat() if first_acquired_at else "",
             "released_at": "",
         }
 
-    return lock_contents
-
 
 class LockfileWriter:
     """The core purpose of this class is to allow setting of first_acquired_at immediately

thds/mops/pure/core/memo/__init__.py

@@ -1,4 +1,4 @@
-from . import calls, unique_name_for_function  # noqa: F401
+from . import calls, results, unique_name_for_function  # noqa: F401
 from .function_memospace import (  # noqa
     args_kwargs_content_address,
     make_function_memospace,

thds/mops/pure/core/memo/results.py

@@ -76,8 +76,8 @@ class RequiredResultNotFound(Exception):
 
 def check_if_result_exists(
     memo_uri: str,
-    rerun_excs: bool = False,
-    before_raise: ty.Callable[[], ty.Any] = lambda: None,
+    check_for_exception: bool = False,
+    before_raise: ty.Optional[ty.Callable[[], ty.Any]] = None,
 ) -> ty.Union[None, Success, Error]:
     fs = lookup_blob_store(memo_uri)
     value_uri = fs.join(memo_uri, RESULT)
@@ -86,14 +86,15 @@ def check_if_result_exists(
 
     required_msg = _should_require_result(memo_uri)
     if required_msg:  # might be custom or the default. either way it indicates a required result.
-        before_raise()
+        if before_raise:
+            before_raise()
         error_msg = f"Required a result for {ORANGE(memo_uri)} but that result was not found"
         # i'm tired of visually scanning for these memo_uris in logs.
         if required_msg != _NO_MSG:
             error_msg += f": {required_msg}"
         raise RequiredResultNotFound(error_msg, memo_uri)
 
-    if rerun_excs:
+    if not check_for_exception:
         return None
 
     error_uri = fs.join(memo_uri, EXCEPTION)

thds/mops/pure/core/use_runner.py

@@ -4,21 +4,29 @@ You can transfer control to a Runner without this, but decorators are a Pythonic
 """
 
 import typing as ty
+from contextlib import contextmanager
 from functools import wraps
 
 from thds.core import log, stack_context
+from thds.mops._utils.names import full_name_and_callable
 
-from .entry.runner_registry import entry_count
 from .types import Runner
 
+_USE_RUNNER_BYPASS = stack_context.StackContext[set[str]]("use_runner_bypass", set())
+# use this in a Runner remote entry point to allow the remote function call
+# to bypass any use_runner decorator. Also necessary in case somebody is doing advanced
+# things like using a remote runner to run a manifest of _other_ remote functions...
+
 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()
+@contextmanager
+def unwrap_use_runner(f: F) -> ty.Iterator[None]:
+    full_name, _ = full_name_and_callable(f)
+    with _USE_RUNNER_BYPASS.set({full_name}):
+        # this is a no-op if the function is not wrapped
+        yield
 
 
 def use_runner(runner: Runner, skip: ty.Callable[[], bool] = lambda: False) -> ty.Callable[[F], F]:
@@ -34,9 +42,15 @@ def use_runner(runner: Runner, skip: ty.Callable[[], bool] = lambda: False) -> t
     def deco(f: F) -> F:
         @wraps(f)
         def __use_runner_wrapper(*args, **kwargs):  # type: ignore
-            if _is_runner_entry() or skip():
+            def should_bypass() -> bool:
+                if skip():
+                    return True
+                full_name, _ = full_name_and_callable(f)
+                return full_name in _USE_RUNNER_BYPASS()
+
+            if should_bypass():
                 logger.debug("Calling function %s directly...", f)
-                with FUNCTION_UNWRAP_COUNT.set(FUNCTION_UNWRAP_COUNT() + 1):
+                with unwrap_use_runner(f):
                     return f(*args, **kwargs)
 
             logger.debug("Forwarding local function %s call to runner...", f)

thds/mops/pure/pickling/mprunner.py

@@ -11,7 +11,7 @@ from functools import partial
 
 from typing_extensions import Self
 
-from thds.core import cache, log
+from thds.core import cache, futures, log
 from thds.core.stack_context import StackContext
 
 from ..._utils.once import Once
@@ -20,7 +20,7 @@ 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 ..runner.types import FutureShim, Shim, ShimBuilder
 from ..tools.summarize import run_summary
 from . import _pickle, pickles, sha256_b64
 
@@ -32,7 +32,7 @@ _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:
+def mp_shim(base_shim: ty.Union[Shim, FutureShim], shim_args: ty.Sequence[str]) -> ty.Any:
     return base_shim((RUNNER_NAME, *shim_args))
 
 
@@ -48,7 +48,7 @@ class MemoizingPicklingRunner:
 
     def __init__(
         self,
-        shim: ty.Union[ShimBuilder, Shim],
+        shim: ty.Union[ShimBuilder, Shim, FutureShim],
         blob_storage_root: uris.UriResolvable,
         *,
         rerun_exceptions: bool = True,
@@ -159,20 +159,14 @@ class MemoizingPicklingRunner:
             ),
         )
 
-    def _wrap_shim_builder(self, func: F, args: Args, kwargs: Kwargs) -> Shim:
+    def _wrap_shim_builder(self, func: F, args: Args, kwargs: Kwargs) -> ty.Union[Shim, FutureShim]:
         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.
+    def submit(self, func: ty.Callable[..., T], *args: ty.Any, **kwargs: ty.Any) -> futures.PFuture[T]:
+        """Now that mops supports Futures, we can have an 'inner' API that returns a PFuture.
 
-        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.
+        We are trying to mimic the interface that concurrent.futures.Executors provide.
         """
         logger.debug("Preparing to run function via remote shim")
         with _ARGS_CONTEXT.set(args), _KWARGS_CONTEXT.set(kwargs):
@@ -192,3 +186,16 @@ class MemoizingPicklingRunner:
                 args,
                 kwargs,
             )
+
+    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.
+        """
+        return self.submit(func, *args, **kwargs).result()

thds/mops/pure/pickling/pickles.py

@@ -60,19 +60,30 @@ class PicklableFunction:
     def __repr__(self) -> str:
         return str(self)
 
+    def _resolve(self) -> ty.Callable:
+        """Resolve the function if it hasn't been resolved yet."""
+        if self.f is None:
+            logger.debug(f"Dynamically importing function {str(self)}")
+            if self.fmod == "__main__":
+                self.f = get_main_module_function(self.fname)  # type: ignore
+            else:
+                mod = importlib.import_module(self.fmod)
+            self.f = getattr(mod, self.fname)
+            assert self.f is not None
+            return self.f
+        return self.f
+
+    @property
+    def func(self) -> ty.Callable:
+        """This is a property so we aren't ruining backward pickle compatibility."""
+        return self._resolve()
+
     @property
     def __name__(self) -> str:
         return self.fname
 
     def __call__(self, *args: ty.Any, **kwargs: ty.Any) -> ty.Any:
-        logger.debug(f"Dynamically importing function {str(self)}")
-        if self.fmod == "__main__":
-            self.f = get_main_module_function(self.fname)  # type: ignore
-        else:
-            mod = importlib.import_module(self.fmod)
-            self.f = getattr(mod, self.fname)
-        assert self.f
-        return self.f(*args, **kwargs)
+        return self._resolve()(*args, **kwargs)
 
 
 class UnpickleSimplePickleFromUri:

thds/mops/pure/pickling/remote.py

@@ -12,6 +12,7 @@ from ..core.memo import results
 from ..core.serialize_big_objs import ByIdRegistry, ByIdSerializer
 from ..core.serialize_paths import CoordinatingPathSerializer
 from ..core.types import Args, BlobStore, Kwargs, T
+from ..core.use_runner import unwrap_use_runner
 from ..runner import strings
 from . import _pickle, mprunner, pickles, sha256_b64
 
@@ -123,7 +124,8 @@ def run_pickled_invocation(memo_uri: str, *metadata_args: str) -> None:
     def do_work_return_result() -> object:
         # ONLY failures in this code should transmit an EXCEPTION
         # back to the orchestrator side.
-        return func(*args, **kwargs)
+        with unwrap_use_runner(func):
+            return func(*args, **kwargs)
 
     route_return_value_or_exception(
         _ResultExcWithMetadataChannel(

thds/mops/pure/runner/get_results.py

@@ -0,0 +1,106 @@
+import concurrent.futures
+import threading
+import typing as ty
+from dataclasses import dataclass
+from pathlib import Path
+
+from thds.core import futures, log
+
+from ...config import max_concurrent_network_ops
+from ..core import lock, memo
+from ..core.types import NoResultAfterShimSuccess
+from ..tools.summarize import run_summary
+from . import types
+
+
+class ResultAndInvocationType(ty.NamedTuple):
+    value_or_error: ty.Union[memo.results.Success, memo.results.Error]
+    invoc_type: run_summary.InvocationType
+
+
+def unwrap_value_or_error(
+    get_meta_and_result: types.GetMetaAndResult,
+    run_directory: ty.Optional[Path],
+    runner_prefix: str,
+    args_kwargs_uris: ty.Collection[str],
+    memo_uri: str,
+    result_and_itype: ResultAndInvocationType,
+) -> ty.Any:  # the result value
+    result = result_and_itype.value_or_error
+    metadata = None
+    value_t = None
+    try:
+        if isinstance(result, memo.results.Success):
+            metadata, value_t = get_meta_and_result("value", result.value_uri)
+            return value_t
+        else:
+            assert isinstance(result, memo.results.Error), "Must be Error or Success"
+            metadata, exc = get_meta_and_result("EXCEPTION", result.exception_uri)
+            raise exc
+    finally:
+        run_summary.log_function_execution(
+            *(run_directory, memo_uri, result_and_itype.invoc_type),
+            metadata=metadata,
+            runner_prefix=runner_prefix,
+            was_error=not isinstance(result, memo.results.Success),
+            return_value=value_t,
+            args_kwargs_uris=args_kwargs_uris,
+        )
+
+
+_AFTER_INVOCATION_SEMAPHORE = threading.BoundedSemaphore(int(max_concurrent_network_ops()) * 3)
+# _IN prioritizes retrieving the result of a Shim that has completed.
+logger = log.getLogger(__name__)
+T = ty.TypeVar("T")
+
+
+@dataclass
+class PostShimResultGetter(ty.Generic[T]):
+    """Must be serializable on its own, so we can pass it across process boundaries
+    to serve as a foundation for a cross-process Future.
+
+    Happily, this should not be terribly difficult, as the 'state' of a mops function
+    is predicted entirely on the memo URI, which is a string.
+    """
+
+    memo_uri: str
+    partially_applied_unwrap_value_or_error: ty.Callable[[str, ResultAndInvocationType], T]
+    release_lock: ty.Optional[ty.Callable[[], None]] = None
+
+    def __call__(self, _shim_result: ty.Any) -> T:
+        """Check if the result exists, and return it if it does.
+
+        This is the future 'translator' that allows us to chain a shim future to be a result future.
+        """
+        memo_uri = self.memo_uri
+
+        try:
+            with _AFTER_INVOCATION_SEMAPHORE:
+                value_or_error = memo.results.check_if_result_exists(memo_uri, check_for_exception=True)
+                if not value_or_error:
+                    raise NoResultAfterShimSuccess(
+                        f"The shim for {memo_uri} exited cleanly, but no result or exception was found."
+                    )
+                return self.partially_applied_unwrap_value_or_error(
+                    memo_uri, ResultAndInvocationType(value_or_error, "invoked")
+                )
+        finally:
+            if self.release_lock is not None:
+                try:
+                    self.release_lock()
+                except Exception:
+                    logger.exception("Failed to release lock after shim result retrieval.")
+
+
+def lock_maintaining_future(
+    lock_acquired: lock.LockAcquired,
+    post_shim_result_getter: PostShimResultGetter[futures.R1],
+    inner_future: futures.PFuture[futures.R],
+) -> concurrent.futures.Future[futures.R1]:
+    """Create a Future that will be used to retrieve the result of a shim invocation.
+
+    This Future will be used to retrieve the result of a shim invocation, and will
+    maintain the lock while it is being retrieved.
+    """
+    post_shim_result_getter.release_lock = lock.maintain_to_release(lock_acquired)
+    return futures.chain_futures(inner_future, concurrent.futures.Future(), post_shim_result_getter)

thds/mops/pure/runner/local.py

@@ -5,28 +5,32 @@ import threading
 import time
 import typing as ty
 from datetime import datetime, timedelta, timezone
+from functools import partial
 from pathlib import Path
 
-from thds.core import config, log, scope
+from thds.core import futures, log, scope
 from thds.termtool.colorize import colorized, make_colorized_out
 
 from ...config import max_concurrent_network_ops
 from ..core import deferred_work, lock, memo, metadata, pipeline_id_mask, uris
+from ..core.lock.maintain import MAINTAIN_LOCKS  # noqa: F401
 from ..core.partial import unwrap_partial
-from ..core.types import Args, Kwargs, NoResultAfterShimSuccess, T
+from ..core.types import Args, Kwargs, T
 from ..tools.summarize import run_summary
 from . import strings, types
-
-MAINTAIN_LOCKS = config.item("thds.mops.pure.local.maintain_locks", default=True, parse=config.tobool)
-
-# these two semaphores allow us to prioritize getting meaningful units
+from .get_results import (
+    PostShimResultGetter,
+    ResultAndInvocationType,
+    lock_maintaining_future,
+    unwrap_value_or_error,
+)
+
+# this semaphore (and a similar one in get_results) allow us to prioritize getting a single unit
 # of progress _complete_, rather than issuing many instructions to the
 # underlying client and allowing it to randomly order the operations
 # such that it takes longer to get a full unit of work complete.
 _BEFORE_INVOCATION_SEMAPHORE = threading.BoundedSemaphore(int(max_concurrent_network_ops()))
-# _OUT prioritizes uploading a single invocation and its dependencies so the Shim can start running.
-_AFTER_INVOCATION_SEMAPHORE = threading.BoundedSemaphore(int(max_concurrent_network_ops()))
-# _IN prioritizes retrieving the result of a Shim that has completed.
+# _BEFORE prioritizes uploading a single invocation and its dependencies so the Shim can start running.
 
 _DarkBlue = colorized(fg="white", bg="#00008b")
 _GreenYellow = colorized(fg="black", bg="#adff2f")
@@ -44,9 +48,9 @@ def invoke_via_shim_or_return_memoized(  # noqa: C901
     get_meta_and_result: types.GetMetaAndResult,
     run_directory: ty.Optional[Path] = None,
     calls_registry: ty.Mapping[ty.Callable, ty.Collection[ty.Callable]] = dict(),  # noqa: B006
-) -> ty.Callable[[bool, str, ty.Callable[..., T], Args, Kwargs], T]:
+) -> ty.Callable[[bool, str, ty.Callable[..., T], Args, Kwargs], futures.PFuture[T]]:
     @scope.bound
-    def create_invocation__check_result__wait_shim(
+    def create_invocation_and_result_future(
         rerun_exceptions: bool,
         function_memospace: str,
         # by allowing the caller to set the function memospace, we allow 'redirects' to look up an old result by name.
@@ -54,7 +58,7 @@ def invoke_via_shim_or_return_memoized(  # noqa: C901
         func: ty.Callable[..., T],
         args_: Args,
         kwargs_: Kwargs,
-    ) -> T:
+    ) -> futures.PFuture[T]:
         """This is the generic local runner. Its core abstractions are:
 
         - serializers of some sort (for the function and its arguments)
@@ -89,16 +93,13 @@ def invoke_via_shim_or_return_memoized(  # noqa: C901
         )
 
         # Define some important and reusable 'chunks of work'
-
-        class ResultAndInvocationType(ty.NamedTuple):
-            value_or_error: ty.Union[memo.results.Success, memo.results.Error]
-            invoc_type: run_summary.InvocationType
-
-        def check_result(
+        def check_result_exists(
             invoc_type: run_summary.InvocationType,
         ) -> ty.Union[ResultAndInvocationType, None]:
             result = memo.results.check_if_result_exists(
-                memo_uri, rerun_excs=rerun_exceptions, before_raise=debug_required_result_failure
+                memo_uri,
+                check_for_exception=not rerun_exceptions,
+                before_raise=debug_required_result_failure,
             )
             if not result:
                 return None
@@ -108,28 +109,6 @@ def invoke_via_shim_or_return_memoized(  # noqa: C901
             )
             return ResultAndInvocationType(result, invoc_type)
 
-        def unwrap_value_or_error(result_and_itype: ResultAndInvocationType) -> T:
-            result = result_and_itype.value_or_error
-            metadata = None
-            value_t = None
-            try:
-                if isinstance(result, memo.results.Success):
-                    metadata, value_t = get_meta_and_result("value", result.value_uri)
-                    return ty.cast(T, value_t)
-                else:
-                    assert isinstance(result, memo.results.Error), "Must be Error or Success"
-                    metadata, exc = get_meta_and_result("EXCEPTION", result.exception_uri)
-                    raise exc
-            finally:
-                run_summary.log_function_execution(
-                    *(run_directory, memo_uri, result_and_itype.invoc_type),
-                    metadata=metadata,
-                    runner_prefix=function_memospace.split(pipeline_id)[0],
-                    was_error=not isinstance(result, memo.results.Success),
-                    return_value=value_t,
-                    args_kwargs=(args, kwargs),
-                )
-
         def acquire_lock() -> ty.Optional[lock.LockAcquired]:
             return lock.acquire(fs.join(memo_uri, "lock"), expire=timedelta(seconds=88))
 
@@ -154,6 +133,14 @@ def invoke_via_shim_or_return_memoized(  # noqa: C901
 
             inspect_and_log(memo_uri)
 
+        p_unwrap_value_or_error = partial(
+            unwrap_value_or_error,
+            get_meta_and_result,
+            run_directory,
+            function_memospace.split(pipeline_id)[0],  # runner_prefix
+            run_summary.extract_source_uris((args, kwargs)),
+        )
+
         # the network ops being grouped by _BEFORE_INVOCATION include one or more
         # download attempts (consider possible Paths) plus
         # one or more uploads (embedded Paths & Sources/refs, and then invocation).
@@ -162,9 +149,9 @@ def invoke_via_shim_or_return_memoized(  # noqa: C901
 
             # it's possible that our result may already exist from a previous run of this pipeline id.
             # we can short-circuit the entire process by looking for that result and returning it immediately.
-            result = check_result("memoized")
+            result = check_result_exists("memoized")
             if result:
-                return unwrap_value_or_error(result)
+                return futures.resolved(p_unwrap_value_or_error(memo_uri, result))
 
             lock_owned = acquire_lock()
             # if no result exists, the vastly most common outcome here will be acquiring
@@ -175,10 +162,6 @@ def invoke_via_shim_or_return_memoized(  # noqa: C901
         # LOCK LOOP: entering this loop (where we attempt to acquire the lock) is the common non-memoized case
         while not result:
             if lock_owned:
-                if MAINTAIN_LOCKS():
-                    release_lock = lock.launch_daemon_lock_maintainer(lock_owned)
-                else:
-                    release_lock = lock_owned.release
                 break  # we own the invocation - invoke the shim ourselves (below)
 
             # getting to this point ONLY happens if we failed to acquire the lock, which
@@ -190,57 +173,45 @@ def invoke_via_shim_or_return_memoized(  # noqa: C901
             time.sleep(22)
 
             with _BEFORE_INVOCATION_SEMAPHORE:
-                result = check_result("awaited")
+                result = check_result_exists("awaited")
                 if result:
                     _LogAwaitedResult(
                         f"{val_or_res} for {memo_uri} was found after waiting for the lock."
                     )
-                    return unwrap_value_or_error(result)
+                    return futures.resolved(p_unwrap_value_or_error(memo_uri, result))
 
                 lock_owned = acquire_lock()  # still inside the semaphore, as it's a network op
 
-        assert release_lock is not None
         assert lock_owned is not None
         # if/when we acquire the lock, we move forever into 'run this ourselves mode'.
         # If something about our invocation fails,
         # we fail just as we would have previously, without any attempt to go
         # 'back' to waiting for someone else to compute the result.
 
-        try:
-            with _BEFORE_INVOCATION_SEMAPHORE:
-                _LogNewInvocation(f"Invoking {memo_uri}")
-                upload_invocation_and_deps()
-
-            # can't hold the semaphore while we block on the shim, though.
-            shim_ex = None
-            shim = shim_builder(func, args_, kwargs_)
-            shim(  # ACTUAL INVOCATION (handoff to remote shim) HAPPENS HERE
-                (
-                    memo_uri,
-                    *metadata.format_invocation_cli_args(
-                        metadata.InvocationMetadata.new(pipeline_id, invoked_at, lock_owned.writer_id)
-                    ),
-                )
+        future_result_getter = PostShimResultGetter[T](memo_uri, p_unwrap_value_or_error)
+
+        with _BEFORE_INVOCATION_SEMAPHORE:
+            _LogNewInvocation(f"Invoking {memo_uri}")
+            upload_invocation_and_deps()
+
+        # can't hold the semaphore while we block on the shim, though.
+        shim = shim_builder(func, args_, kwargs_)
+        future_or_shim_result = shim(  # ACTUAL INVOCATION (handoff to remote shim) HAPPENS HERE
+            (
+                memo_uri,
+                *metadata.format_invocation_cli_args(
+                    metadata.InvocationMetadata.new(pipeline_id, invoked_at, lock_owned.writer_id)
+                ),
             )
-        except Exception as ex:
-            # network or similar errors are very common and hard to completely eliminate.
-            # We know that if a result (or error) exists, then the network failure is
-            # not important, because results in blob storage are atomically populated (either fully there or not)
-            logger.exception("Error awaiting shim. Optimistically checking for result.")
-            shim_ex = ex
-
-        finally:
-            release_lock()
-
-        # the network ops being grouped by _AFTER_INVOCATION include one or more downloads.
-        with _AFTER_INVOCATION_SEMAPHORE:
-            value_or_error = memo.results.check_if_result_exists(memo_uri)
-            if not value_or_error:
-                if shim_ex:
-                    raise shim_ex  # re-raise the underlying exception rather than making up our own.
-                raise NoResultAfterShimSuccess(
-                    f"The shim for {memo_uri} exited cleanly, but no result or exception was found."
-                )
-            return unwrap_value_or_error(ResultAndInvocationType(value_or_error, "invoked"))
-
-    return create_invocation__check_result__wait_shim
+        )
+        if hasattr(future_or_shim_result, "add_done_callback"):
+            # if the shim returns a Future, we wrap it.
+            logger.debug("Shim returned a Future; wrapping it for post-shim result retrieval.")
+            return futures.make_lazy(lock_maintaining_future)(
+                lock_owned, future_result_getter, future_or_shim_result
+            )
+        else:  # it's a synchronous shim - just process the result directly.
+            future_result_getter.release_lock = lock.maintain_to_release(lock_owned)
+            return futures.resolved(future_result_getter(future_or_shim_result))
+
+    return create_invocation_and_result_future

thds/mops/pure/runner/shim_builder.py

@@ -2,24 +2,24 @@ import inspect
 import typing as ty
 
 from ..core.types import Args, F, Kwargs
-from .types import Shim, ShimBuilder
+from .types import FutureShim, Shim, ShimBuilder
 
 
 class _static_shim_builder:
-    def __init__(self, shim: Shim) -> None:
+    def __init__(self, shim: ty.Union[Shim, FutureShim]) -> None:
         self.shim = shim
 
-    def __call__(self, _f: F, _args: Args, _kwargs: Kwargs) -> Shim:
+    def __call__(self, _f: F, _args: Args, _kwargs: Kwargs) -> ty.Union[Shim, FutureShim]:
         return self.shim
 
     def __repr__(self) -> str:
         return f"<static_shim_builder for {self.shim}>"
 
 
-def make_builder(shim: ty.Union[Shim, ShimBuilder]) -> ShimBuilder:
+def make_builder(shim_or_builder: ty.Union[Shim, ShimBuilder, FutureShim]) -> ShimBuilder:
     """If you have a Shim and you want to make it into the simplest possible ShimBuilder."""
 
-    if len(inspect.signature(shim).parameters) == 3:
-        return ty.cast(ShimBuilder, shim)
+    if len(inspect.signature(shim_or_builder).parameters) == 3:
+        return ty.cast(ShimBuilder, shim_or_builder)
 
-    return _static_shim_builder(ty.cast(Shim, shim))
+    return _static_shim_builder(ty.cast(Shim, shim_or_builder))