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

thds/mops/k8s/batching.py

@@ -1,198 +0,0 @@
-"""The basic idea of this module is that different threads can submit _parts_ of a job to a batcher,
-and immediately get the job name back, while the batcher itself defers creating the job until the
-batch is full, or when the process exits.
-
-The theory is that will get used in processes whose only responsibility is to create jobs,
-so waiting on atexit to create the final batch is not an issue.
-
-If you want a batcher that has a more context-manager-like behavior, you can write one of
-those, but it wouldn't work well with a concurrent.futures Executor-style approach, since
-those don't have an explicit shutdown procedure that we can hook to call __exit__.
-"""
-
-import atexit
-import concurrent.futures
-import itertools
-import multiprocessing
-import threading
-import typing as ty
-
-from thds.core import cpus, futures, log
-
-from . import _launch, counts
-
-T = ty.TypeVar("T")
-logger = log.getLogger(__name__)
-
-
-class _AtExitBatcher(ty.Generic[T]):
-    def __init__(self, batch_processor: ty.Callable[[ty.Collection[T]], None]) -> None:
-        self.batch: list[T] = []
-        self._registered = False
-        self._lock = threading.RLock()
-        self._batch_processor = batch_processor
-
-    def add(self, item: T) -> None:
-        with self._lock:
-            if not self._registered:
-                atexit.register(self.process)
-                # ensure we flush on process exit, since we don't know how many items are coming
-                self._registered = True
-            self.batch.append(item)
-
-    def process(self) -> None:
-        if self.batch:
-            with self._lock:
-                if self.batch:
-                    self._batch_processor(self.batch)
-                    self.batch = []
-
-
-class K8sJobBatchingShim(_AtExitBatcher[str]):
-    """Thread-safe for use within a single process by multiple threads."""
-
-    def __init__(
-        self,
-        submit_func: ty.Callable[[ty.Collection[str]], ty.Any],
-        max_batch_size: int,
-        job_counter: counts.MpValue[int],
-        name_prefix: str = "",
-    ) -> None:
-        """submit_func in particular should be a closure around whatever setup you need to
-        do to call back into a function that is locally wrapped with a k8s shim that will
-        ultimately call k8s.launch. Notably, you
-        """
-        super().__init__(self._process_batch)
-        self._max_batch_size = max_batch_size
-        self._job_counter = job_counter
-        self._job_name = ""
-        self._name_prefix = name_prefix
-        self._submit_func = submit_func
-
-    def _get_new_name(self) -> str:
-        # counts.inc takes a multiprocess lock. do not forget this!
-        job_num = counts.inc(self._job_counter)
-        return _launch.construct_job_name(self._name_prefix, counts.to_name(job_num))
-
-    def add_to_named_job(self, mops_invocation: ty.Sequence[str]) -> str:
-        """Returns job name for the invocation."""
-        with self._lock:
-            if not self._job_name:
-                self._job_name = self._get_new_name()
-            if len(self.batch) >= self._max_batch_size:
-                self.process()
-                self._job_name = self._get_new_name()
-            super().add(" ".join(mops_invocation))
-            return self._job_name
-
-    def _process_batch(self, batch: ty.Collection[str]) -> None:
-        with _launch.JOB_NAME.set(self._job_name):
-            log_lvl = logger.warning if len(batch) < self._max_batch_size else logger.info
-            log_lvl(f"Processing batch of len {len(batch)} with job name {self._job_name}")
-            self._submit_func(batch)
-
-
-F = ty.TypeVar("F", bound=ty.Callable)
-FunctionDecorator = ty.Callable[[F], F]
-
-
-_BATCHER: ty.Optional[K8sJobBatchingShim] = None
-
-
-def init_batcher(
-    submit_func: ty.Callable[[ty.Collection[str]], ty.Any],
-    func_max_batch_size: int,
-    job_counter: counts.MpValue[int],
-    name_prefix: str = "",
-) -> None:
-    # for use with multiprocessing pool initializer
-    global _BATCHER
-    if _BATCHER is not None:
-        logger.warning("Batcher is already initialized; reinitializing will reset the job name.")
-        return
-
-    _BATCHER = K8sJobBatchingShim(submit_func, func_max_batch_size, job_counter, name_prefix)
-
-
-def init_batcher_with_unpicklable_submit_func(
-    make_submit_func: ty.Callable[[T], ty.Callable[[ty.Collection[str]], ty.Any]],
-    submit_func_arg: T,
-    func_max_batch_size: int,
-    job_counter: counts.MpValue[int],
-    name_prefix: str = "",
-) -> None:
-    """Use this if you want to have an unpicklable submit function - because applying make_submit_func(submit_func_arg)
-    will happen inside the pool worker process after all the pickling/unpickling has happened.
-    """
-    return init_batcher(
-        make_submit_func(submit_func_arg), func_max_batch_size, job_counter, name_prefix=name_prefix
-    )
-
-
-def make_counting_process_pool_executor(
-    make_submit_func: ty.Callable[[T], ty.Callable[[ty.Collection[str]], ty.Any]],
-    submit_func_arg: T,
-    max_batch_size: int,
-    name_prefix: str = "",
-    max_workers: int = 0,
-) -> concurrent.futures.ProcessPoolExecutor:
-    """Creates a ProcessPoolExecutor that uses the batching shim for job submission.
-
-    We are introducing this because we see segfaults prior to Python 3.12 related to this issue:
-    https://github.com/python/cpython/issues/77377
-
-    And it would seem that this had to do with creating mp.Values using a 'fork' start
-    method, and then passing those to a ProcessPoolExecutor with
-    mp_context=multiprolcessing.get_context('spawn'). So we can help you avoid that by creating
-    the mp.Value for you, alongside its ProcessPoolExecutor.
-
-    NOTE!!
-
-    You should only have one of these per process at a time, because we're doing spooky
-    things with the Job Counter.  In fact, you should probably only create one of these
-    _ever_ within a single logical 'application'.
-
-    If you fail to heed this advice, you will get weird launched/finished counts at a
-    minimum. Although these job counts are not mission-critical, you _will_ be confused.
-    """
-    start_method: str = "spawn"
-    # 'spawn' prevents weird batch processing deadlocks that seem to only happen on Linux with 'fork'.
-    # it is strongly recommended to use 'spawn' for this reason.
-
-    mp_context = multiprocessing.get_context(start_method)
-    launch_count = mp_context.Value("i", 0)
-    # even though i want to assign this to a global, I also want to prevent
-    # any possible race condition where i somehow use a different thread's LAUNCH_COUNT
-    # when i create the ProcessPoolExecutor a few lines below.
-    counts.LAUNCH_COUNT = launch_count
-    counts.FINISH_COUNT = mp_context.Value("i", 0)  # we don't use this here; we just reset it to zero.
-    # SPOOKY - reset the global finish counter and make it be the same 'type'
-    return concurrent.futures.ProcessPoolExecutor(
-        max_workers=max_workers or cpus.available_cpu_count(),
-        initializer=init_batcher_with_unpicklable_submit_func,
-        initargs=(make_submit_func, submit_func_arg, max_batch_size, launch_count, name_prefix),
-        mp_context=mp_context,
-    )
-
-
-def shim(args: ty.Sequence[str]) -> futures.PFuture[bool]:
-    # This thing needs to return a lazy Uncertain Future that contains a job name, so that Job can be polled on
-    # ... but the job does not exist yet! So the batcher is in charge of creating the job name
-    # upfront, and then ensuring that it gets used when the job is created.
-    assert _BATCHER is not None, "Batcher must be initialized before using the batching shim."
-    job_name = _BATCHER.add_to_named_job(args)
-    return _launch.create_lazy_job_logging_future(job_name)
-
-
-def batched(iterable: ty.Iterable[T], n: int, *, strict: bool = False) -> ty.Iterator[tuple[T, ...]]:
-    """Just a utility for pre-batching if you're using multiprocessing to create batches."""
-    # TODO get rid of this when we go to Python 3.12+ which has itertools.batched
-    #
-    # batched('ABCDEFG', 3) → ABC DEF G
-    if n < 1:
-        raise ValueError("n must be at least one")
-    iterator = iter(iterable)
-    while batch := tuple(itertools.islice(iterator, n)):
-        if strict and len(batch) != n:
-            raise ValueError("batched(): incomplete batch")
-        yield batch

thds/mops/k8s/counts.py

@@ -1,28 +0,0 @@
-import multiprocessing as mp
-import typing as ty
-
-T = ty.TypeVar("T")
-
-
-class MpValue(ty.Protocol[T]):
-    def get_lock(self) -> ty.Any:
-        ...
-
-    value: T
-
-
-def inc(mp_val: MpValue[int]) -> int:
-    with mp_val.get_lock():
-        mp_val.value += 1
-        return mp_val.value
-
-
-LAUNCH_COUNT = mp.Value("i", 0)
-FINISH_COUNT = mp.Value("i", 0)
-# these are spooky - they're global and mutable, and may in fact get overwritten by code
-# using specific multiprocessing contexts.
-
-
-def to_name(count: int) -> str:
-    """Convert a count to a name."""
-    return f"{count:0>4}"

thds/mops/k8s/job_future.py

@@ -1,109 +0,0 @@
-import threading
-import typing as ty
-
-from kubernetes import client
-
-from thds.core import futures, log
-from thds.termtool.colorize import colorized
-
-from . import config, counts, uncertain_future
-from .jobs import is_job_failed, is_job_succeeded, job_source
-
-logger = log.getLogger(__name__)
-
-UNUSUAL = colorized(fg="white", bg="yellow")
-SUCCEEDED = colorized(fg="white", bg="blue")
-FAILED = colorized(fg="white", bg="red")
-
-
-_FINISHED_JOBS = set[str]()
-_FINISHED_JOBS_LOCK = threading.Lock()
-
-
-def _check_newly_finished(job_name: str, namespace: str = "") -> str:
-    # I don't believe it's possible to ever have a Job that both succeeds and fails.
-    namespace = namespace or config.k8s_namespace()
-    job_full = f"{namespace}/{job_name}"
-    if job_full in _FINISHED_JOBS:
-        return ""
-
-    with _FINISHED_JOBS_LOCK:
-        if job_full in _FINISHED_JOBS:
-            return ""
-
-        _FINISHED_JOBS.add(job_full)
-
-    launched = counts.LAUNCH_COUNT.value
-    return f"- ({launched - counts.inc(counts.FINISH_COUNT)} unfinished of {launched})"
-
-
-class K8sJobFailedError(Exception):
-    """Raised by `launch` when a Job is seen to terminate in a Failed state."""
-
-
-def make_job_completion_future(job_name: str, *, namespace: str = "") -> futures.PFuture[bool]:
-    """This is a natural boundary for a serializable lazy future - something that represents
-    work being done across process boundaries (since Kubernetes jobs will be listed via an API.
-
-    If True is returned, the Job has definitely succeeded.
-
-    If False is returned, the Job may have succeeded but we saw no evidence of it.
-
-    If the Job definitely failed, an Exception will be raised.
-    """
-
-    JOB_SEEN = False
-
-    def job_completion_interpreter(
-        job: ty.Optional[client.models.V1Job], last_seen_at: float
-    ) -> ty.Union[uncertain_future.NotYetDone, bool]:
-        nonlocal JOB_SEEN
-        if not job:
-            if JOB_SEEN:
-                logger.warning(
-                    UNUSUAL(f"Previously-seen job {job_name} no longer exists - assuming success!")
-                )
-                # we hereby indicate an unusual success to the Future waiter.
-                return False
-
-            time_since_last_seen = uncertain_future.official_timer() - last_seen_at
-            if time_since_last_seen > config.k8s_watch_object_stale_seconds():
-                # this is 5 minutes by default as of 2025-07-15.
-                raise TimeoutError(
-                    f"Job {job_name} has not been seen for {time_since_last_seen:.1f} seconds - assuming failure!"
-                )
-
-            # we don't know what's going on but things aren't truly stale yet.
-            return uncertain_future.NotYetDone()
-
-        JOB_SEEN = True
-
-        if is_job_succeeded(job):
-            newly_succeeded = _check_newly_finished(job_name, namespace)
-            if newly_succeeded:
-                logger.info(SUCCEEDED(f"Job {job_name} Succeeded! {newly_succeeded}"))
-            return True
-
-        if is_job_failed(job):
-            newly_failed = _check_newly_finished(job_name, namespace)
-            if newly_failed:
-                logger.error(FAILED(f"Job {job_name} Failed! {newly_failed}"))
-            raise K8sJobFailedError(f"Job {job_name} has failed with status: {job.status}")
-
-        return uncertain_future.NotYetDone()  # job is still in progress
-
-    return job_source().create_future(
-        job_completion_interpreter,
-        job_name,
-        namespace=namespace or config.k8s_namespace(),
-    )
-
-
-def make_lazy_completion_future(job_name: str, *, namespace: str = "") -> futures.LazyFuture[bool]:
-    """This is a convenience function that will create a job completion future and then
-    immediately process it, returning the result. See docs on function above.
-    """
-    return futures.make_lazy(make_job_completion_future)(
-        job_name,
-        namespace=namespace or config.k8s_namespace(),
-    )

thds/mops/k8s/uncertain_future.py

@@ -1,160 +0,0 @@
-import collections
-import threading
-import time
-import typing as ty
-
-# we use concurrent.futures.Future as an implementation detail, but it's communicated
-# as core.futures.PFuture to give us the flexibility to change the implementation later if needed.
-from concurrent.futures import Future
-from dataclasses import dataclass
-from uuid import uuid4
-
-from typing_extensions import Self
-
-from thds import core
-
-R_0 = ty.TypeVar("R_0", contravariant=True)  # R-naught - the thing that might resolve a Future.
-# a value for this type may never be None.
-
-R = ty.TypeVar("R")
-# the Result type of the Future. These are allowed to be None, since some Futures may
-# resolve but not return a value.
-
-
-class NotYetDone:
-    pass
-
-
-_LastSeenAt = float  # type alias for the last seen time of the Future, in seconds since epoch
-
-
-FutureInterpreter = ty.Callable[[ty.Optional[R_0], _LastSeenAt], ty.Union[R, NotYetDone]]
-# a FutureInterpreter is a function that takes an object R_0 and the time.monotonic() at
-# which it was last seen, and returns either NotYetDone (if the status is still in progress) or
-# the actual Future result of type R, or, if the status is failure,
-# _raises_ an appropriate Exception.
-
-
-class _FutureInterpretationShim(ty.Generic[R_0, R]):
-    def __init__(self, interpreter: FutureInterpreter[R_0, ty.Union[NotYetDone, R]]) -> None:
-        self.future = Future[R]()
-        self._interpreter = interpreter
-        self._id = uuid4().hex  # has an id so it can be hashed and therefore easily found in a set
-
-    def __hash__(self) -> int:
-        return hash(self._id)
-
-    def __call__(self, r_0: ty.Optional[R_0], last_seen_at: float) -> ty.Optional[Self]:
-        """First and foremost - this _must_ be treated as an object that the creator
-        is ultimately responsible for calling on a semi-regular basis. It represents a
-        likely deadlock for the holder of the Future if it is never called.
-
-        Return False if the Future is still in progress and should not be unregistered.
-        Return True if the Future is done and should be unregistered.
-        """
-        try:
-            interpretation = self._interpreter(r_0, last_seen_at)
-            if isinstance(interpretation, NotYetDone):
-                return None  # do nothing and do not unregister - the status is still in progress.
-
-            self.future.set_result(interpretation)
-        except Exception as e:
-            self.future.set_exception(e)
-
-        return self
-
-
-K = ty.TypeVar("K")  # Key type for the UncertainFuturesTracker
-
-
-@dataclass
-class _FuturesState(ty.Generic[R_0]):
-    """Represents a single 'observable' that may have multiple Futures (and therefore interpretations) associated with it."""
-
-    futshims: list[_FutureInterpretationShim[R_0, ty.Any]]
-    last_seen_at: float
-
-
-def official_timer() -> float:
-    # we don't need any particular meaning to the time.
-    return time.monotonic()
-
-
-class UncertainFuturesTracker(ty.Generic[K, R_0]):
-    """This class represents a kind of Future where we cannot be guaranteed that we will ever see
-    any further information about it, because we do not control the source of the data.
-
-    A good example would be a Kubernetes object that we are watching - we may _think_ that a Job will be created,
-    but there are race conditions galore in terms of actually looking for that object.
-
-    However, if we _do_ see it at a some point, then we can interpret future 'missingness'
-    as a tentative success.
-
-    The danger with this uncertainty is that Futures represent implicit deadlocks - if we
-    never resolve the Future, then a caller may be waiting for it forever. Therefore, we
-    ask the original requestor of the Future to specify how long they are willing to wait
-    to get a result, after which point we will resolve the Future as an exception.
-    """
-
-    def __init__(self, allowed_stale_seconds: float) -> None:
-        self._keyed_futures_state = collections.OrderedDict[K, _FuturesState[R_0]]()
-        self._lock = threading.Lock()  # i don't trust ordered dict operations to be thread-safe.
-        self._check_stale_seconds = allowed_stale_seconds
-
-    def create(self, key: K, interpreter: FutureInterpreter[R_0, R]) -> core.futures.PFuture[R]:
-        futshim = _FutureInterpretationShim(interpreter)
-        with self._lock:
-            if key not in self._keyed_futures_state:
-                self._keyed_futures_state[key] = _FuturesState(
-                    [futshim],
-                    last_seen_at=official_timer() + self._check_stale_seconds,
-                    # we provide a double margin for objects that we have never seen before.
-                )
-                self._keyed_futures_state.move_to_end(key, last=False)
-                # never seen and therefore should be at the beginning (most stale)
-            else:
-                # maintain our ordered dict so we can handle garbage collection of stale Futures.
-                self._keyed_futures_state[key].futshims.append(futshim)
-
-        return futshim.future
-
-    def update(self, key: ty.Optional[K], r_0: ty.Optional[R_0]) -> None:
-        """Update the keyed Futures based on their interpreters.
-
-        Also check any stale Futures - Futures that have not seen an update (via their key) in a while.
-
-        If `key` is None, we will update all Futures that have been created so far.
-        """
-
-        def check_resolution(fut_state: _FuturesState[R_0], inner_r_0: ty.Optional[R_0]) -> None:
-            for future_shim_that_is_done in core.parallel.yield_results(
-                [
-                    core.thunks.thunking(futshim)(inner_r_0, fut_state.last_seen_at)
-                    for futshim in fut_state.futshims
-                ],
-                progress_logger=core.log.getLogger(__name__).debug,
-                named="UncertainFuturesTracker.update",
-            ):
-                if future_shim_that_is_done is not None:
-                    # the Future is done, so we can remove it from the list of Futures.
-                    fut_state.futshims.remove(future_shim_that_is_done)
-
-        if key is not None:
-            with self._lock:
-                if key not in self._keyed_futures_state:
-                    self._keyed_futures_state[key] = _FuturesState(list(), last_seen_at=official_timer())
-                else:
-                    # maintain our ordered dict so we can handle garbage collection of stale Futures.
-                    self._keyed_futures_state.move_to_end(key)
-                    self._keyed_futures_state[key].last_seen_at = official_timer()
-
-            fut_state = self._keyed_futures_state[key]
-            check_resolution(fut_state, r_0)
-
-        # 'garbage collect' any Futures that haven't been updated in a while.
-        for futs_state in self._keyed_futures_state.values():
-            if futs_state.last_seen_at + self._check_stale_seconds < official_timer():
-                check_resolution(futs_state, None)
-            else:  # these are ordered, so once we see one that's not stale, we can stop checking.
-                # this prevents us from having to do O(N) checks for every update.
-                break

thds/mops/pure/runner/get_results.py

@@ -1,106 +0,0 @@
-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)