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

thds/mops/pure/core/deferred_work.py

@@ -3,23 +3,15 @@
 # this could be _any_ kind of work, but is only uploads as of initial abstraction.
 # this basic idea was stolen from `pure.core.source` as a form of optimization for
 # uploading Sources and their hashrefs.
-import concurrent.futures
 import typing as ty
 from contextlib import contextmanager
 
 from thds import core
-from thds.core import config, refcount
 from thds.core.stack_context import StackContext
 
 _DEFERRED_INVOCATION_WORK: StackContext[
     ty.Optional[ty.Dict[ty.Hashable, ty.Callable[[], ty.Any]]]
 ] = StackContext("DEFERRED_INVOCATION_WORK", None)
-_MAX_DEFERRED_WORK_THREADS = config.item("max_deferred_work_threads", default=50, parse=int)
-_DEFERRED_WORK_THREADPOOL = refcount.Resource[concurrent.futures.ThreadPoolExecutor](
-    lambda: concurrent.futures.ThreadPoolExecutor(
-        max_workers=_MAX_DEFERRED_WORK_THREADS(), **core.concurrency.initcontext()
-    )
-)
 logger = core.log.getLogger(__name__)
 
 
@@ -32,7 +24,15 @@ def open_context() -> ty.Iterator[None]:
     The idea is that you'd call perform_all() inside your Shim which transfers
     execution to a remote environment, but _not_ call it if you're transferring execution
     to a local environment, as the upload will not be needed.
+
+    This is not re-entrant. If this is called while the dictionary is non-empty, an
+    exception will be raised. This is only because I can think of no reason why anyone
+    would want it to be re-entrant, so it seems better to raise an error. If for some
+    reason re-entrancy were desired, we could just silently pass if the dictionary already
+    has deferred work.
     """
+    existing_work = _DEFERRED_INVOCATION_WORK()
+    assert existing_work is None, f"deferred work context is not re-entrant! {existing_work}"
     with _DEFERRED_INVOCATION_WORK.set(dict()):
         logger.debug("Opening deferred work context")
         yield
@@ -74,13 +74,10 @@ def perform_all() -> None:
     work_items = _DEFERRED_INVOCATION_WORK()
     if work_items:
         logger.info("Performing %s items of deferred work", len(work_items))
-        with _DEFERRED_WORK_THREADPOOL.get() as thread_pool_executor:
-            for key, _ in core.parallel.failfast(
-                core.parallel.yield_all(dict(work_items).items(), executor_cm=thread_pool_executor)
-            ):
-                # consume iterator but don't keep results in memory.
-                logger.debug("Popping deferred work %s from %s", key, id(work_items))
-                work_items.pop(key)
+        for key, _result in core.parallel.yield_all(dict(work_items).items()):
+            # consume iterator but don't keep results in memory.
+            logger.debug("Popping deferred work %s from %s", key, id(work_items))
+            work_items.pop(key)
 
         logger.debug("Done performing deferred work on %s", id(work_items))
         assert not work_items, f"Some deferred work was not performed! {work_items}"

thds/mops/pure/core/entry/runner_registry.py

@@ -5,6 +5,14 @@ In practice we only have a single Runner type registered, the MemoizingPicklingR
 
 import typing as ty
 
+from thds.core import stack_context
+
+RUNNER_ENTRY_COUNT = stack_context.StackContext("runner_entry_count", 0)
+
+
+def entry_count() -> int:
+    return RUNNER_ENTRY_COUNT()
+
 
 class EntryHandler(ty.Protocol):
     def __call__(self, *__args: str) -> ty.Any:
@@ -19,4 +27,5 @@ def register_entry_handler(name: str, mh: EntryHandler) -> None:
 
 
 def run_named_entry_handler(name: str, *args: str) -> None:
-    ENTRY_HANDLERS[name](*args)
+    with RUNNER_ENTRY_COUNT.set(RUNNER_ENTRY_COUNT() + 1):
+        ENTRY_HANDLERS[name](*args)

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

@@ -2,7 +2,6 @@ from ._acquire import acquire  # noqa: F401
 from .maintain import (  # noqa: F401
     CannotMaintainLock,
     launch_daemon_lock_maintainer,
-    maintain_to_release,
     remote_lock_maintain,
 )
 from .types import LockAcquired  # noqa: F401

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

@@ -32,7 +32,7 @@ from thds.core import log
 from . import _funcs
 from .read import get_writer_id, make_read_lockfile
 from .types import LockAcquired, LockContents
-from .write import LockEmitter, LockfileWriter
+from .write import LockfileWriter, make_lock_contents
 
 logger = log.getLogger(__name__)
 
@@ -106,7 +106,7 @@ def acquire(  # noqa: C901
     lockfile_writer = LockfileWriter(
         my_writer_id,
         lock_dir_uri,
-        LockEmitter(my_writer_id, expire),
+        make_lock_contents(my_writer_id, expire),
         expire.total_seconds(),
         debug=debug,
     )

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

@@ -15,14 +15,12 @@ from datetime import datetime, timedelta
 from functools import partial
 from threading import Thread
 
-from thds.core import config, log
+from thds.core import log
 
 from ._funcs import make_lock_uri
 from .read import get_writer_id, make_read_lockfile
 from .types import LockAcquired
-from .write import LockEmitter, LockfileWriter
-
-MAINTAIN_LOCKS = config.item("thds.mops.pure.local.maintain_locks", default=True, parse=config.tobool)
+from .write import LockfileWriter, make_lock_contents
 
 logger = log.getLogger(__name__)
 
@@ -105,7 +103,7 @@ def remote_lock_maintain(lock_dir_uri: str, expected_writer_id: str = "") -> Loc
     lockfile_writer = LockfileWriter(
         current_writer_id,
         lock_dir_uri,
-        LockEmitter(get_writer_id(lock_contents), timedelta(seconds=expire_s)),
+        make_lock_contents(get_writer_id(lock_contents), timedelta(seconds=expire_s)),
         expire_s,
         writer_name="remote",
     )
@@ -150,20 +148,3 @@ def launch_daemon_lock_maintainer(lock_acq: LockAcquired) -> ty.Callable[[], Non
         lock_acq.release()
 
     return stop_maintaining
-
-
-def maintain_to_release(
-    acquired_lock: LockAcquired,
-) -> ty.Callable[[], None]:
-    """Depending on configuration, potentially start maintaining the lock.
-
-    Return a callable that will release the lock when called.
-    """
-    if MAINTAIN_LOCKS():
-        return launch_daemon_lock_maintainer(acquired_lock)
-
-    return acquired_lock.release
-
-
-def no_maintain() -> None:
-    MAINTAIN_LOCKS.set_global(False)

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

@@ -1,6 +1,5 @@
 import os
 import typing as ty
-from dataclasses import dataclass
 from datetime import datetime, timedelta
 
 from thds.core import hostname, log
@@ -11,37 +10,38 @@ from .types import LockContents
 logger = log.getLogger(__name__)
 
 
-@dataclass
-class LockEmitter:
-    writer_id: str
-    expire: timedelta
+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 = ""
 
-    write_count: int = 0
-    first_written_at: str = ""
+    assert (
+        "/" not in writer_id
+    ), f"{writer_id} should not contain a slash - maybe you passed a URI instead?"
 
-    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
+    def lock_contents(first_acquired_at: ty.Optional[datetime]) -> LockContents:
+        nonlocal write_count, first_written_at
+        write_count += 1
         now = _funcs.utc_now().isoformat()
-        self.first_written_at = self.first_written_at or now
+        first_written_at = first_written_at or now
 
         return {
-            "writer_id": self.writer_id,
+            "writer_id": writer_id,
             "written_at": now,
-            "expire_s": self.expire.total_seconds(),
+            "expire_s": expire.total_seconds(),
             # debug stuff:
-            "write_count": self.write_count,
+            "write_count": write_count,
             "hostname": hostname.friendly(),
             "pid": str(os.getpid()),
-            "first_written_at": self.first_written_at,
+            "first_written_at": 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, results, unique_name_for_function  # noqa: F401
+from . import calls, 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,
-    check_for_exception: bool = False,
-    before_raise: ty.Optional[ty.Callable[[], ty.Any]] = None,
+    rerun_excs: bool = False,
+    before_raise: ty.Callable[[], ty.Any] = lambda: None,
 ) -> ty.Union[None, Success, Error]:
     fs = lookup_blob_store(memo_uri)
     value_uri = fs.join(memo_uri, RESULT)
@@ -86,15 +86,14 @@ 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.
-        if before_raise:
-            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 not check_for_exception:
+    if rerun_excs:
         return None
 
     error_uri = fs.join(memo_uri, EXCEPTION)

thds/mops/pure/core/use_runner.py

@@ -4,29 +4,21 @@ 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)
 
 
-@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 _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]:
@@ -42,15 +34,9 @@ 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
-            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():
+            if _is_runner_entry() or skip():
                 logger.debug("Calling function %s directly...", f)
-                with unwrap_use_runner(f):
+                with FUNCTION_UNWRAP_COUNT.set(FUNCTION_UNWRAP_COUNT() + 1):
                     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, futures, log
+from thds.core import cache, 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 FutureShim, Shim, ShimBuilder
+from ..runner.types import 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: ty.Union[Shim, FutureShim], shim_args: ty.Sequence[str]) -> ty.Any:
+def mp_shim(base_shim: Shim, 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, FutureShim],
+        shim: ty.Union[ShimBuilder, Shim],
         blob_storage_root: uris.UriResolvable,
         *,
         rerun_exceptions: bool = True,
@@ -159,14 +159,20 @@ class MemoizingPicklingRunner:
             ),
         )
 
-    def _wrap_shim_builder(self, func: F, args: Args, kwargs: Kwargs) -> ty.Union[Shim, FutureShim]:
+    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 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.
+    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.
 
-        We are trying to mimic the interface that concurrent.futures.Executors provide.
+        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):
@@ -186,16 +192,3 @@ 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,30 +60,19 @@ 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:
-        return self._resolve()(*args, **kwargs)
+        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)
 
 
 class UnpickleSimplePickleFromUri:

thds/mops/pure/pickling/remote.py

@@ -12,7 +12,6 @@ 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
 
@@ -124,8 +123,7 @@ 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.
-        with unwrap_use_runner(func):
-            return func(*args, **kwargs)
+        return func(*args, **kwargs)
 
     route_return_value_or_exception(
         _ResultExcWithMetadataChannel(