thds.mops 3.6.20250219172032__py3-none-any.whl

thds/mops/k8s/namespace.py

@@ -0,0 +1,17 @@
+from getpass import getuser
+
+
+def parse_namespace(input_str: str) -> str:
+    # lowercase and replace all non-alphanumeric characters with dashes
+    return "".join(c if c.isalnum() else "-" for c in input_str.lower())
+
+
+def user_namespace() -> str:
+    try:
+        return getuser()
+    except OSError:
+        return "CICD-Runner"
+
+
+def main() -> None:
+    print(user_namespace())

thds/mops/k8s/node_selection.py

@@ -0,0 +1,58 @@
+import typing as ty
+
+from kubernetes import client
+from typing_extensions import TypedDict
+
+
+class ResourceDefinition(TypedDict, total=False):
+    """
+    This works for both limits and requests.
+
+    https://kubernetes.io/docs/concepts/configuration/manage-resources-containers/
+    """
+
+    memory: str
+    """E.g., 10G"""
+    cpu: str
+    """E.g., 4.5, or 4500m (millicores)"""
+
+
+class NodeNarrowing(TypedDict, total=False):
+    """This is a more transparent interface for selecting nodes that your job can run on.
+
+    You don't have to provide each key, but any key/value you pair you provide must be the proper type.
+    """
+
+    resource_requests: ResourceDefinition
+    resource_limits: ResourceDefinition
+    # https://kubernetes.io/docs/concepts/configuration/manage-resources-containers/
+    node_selector: ty.Mapping[str, str]
+    # https://kubernetes.io/docs/concepts/scheduling-eviction/assign-pod-node/
+    tolerations: ty.Sequence[client.V1Toleration]
+    # https://kubernetes.io/docs/concepts/scheduling-eviction/taint-and-toleration/
+
+
+def tolerates_spot() -> client.V1Toleration:
+    """Return our custom spot instance toleration configuration."""
+    return client.V1Toleration(
+        key="kubernetes.azure.com/scalesetpriority", value="spot", effect="NoSchedule"
+    )
+
+
+def tolerates_gpu() -> client.V1Toleration:
+    """Apply this toleration to enable use of GPUs."""
+    return client.V1Toleration(key="dedicated", value="gpu", effect="NoSchedule")
+
+
+def tolerates_64cpu() -> client.V1Toleration:
+    """These node pools often do not scale up well or quickly, so by
+    default they're disabled. If that changes in the future, or if you
+    are requesting more than 32 CPUs for your Pod, you should apply
+    this toleration.
+    """
+    return client.V1Toleration(key="dedicated", value="64cpu", effect="NoSchedule")
+
+
+def require_gpu() -> NodeNarrowing:
+    """Merge this with any additional NodeNarrowing (e.g. resource_requests) to run on GPUs."""
+    return dict(node_selector={"instance-type": "gpu"}, tolerations=[tolerates_gpu()])

thds/mops/k8s/retry.py

@@ -0,0 +1,75 @@
+import time
+import typing as ty
+from functools import wraps
+
+import urllib3.exceptions
+from kubernetes import client
+
+from . import auth, config
+from ._shared import logger
+
+F = ty.TypeVar("F", bound=ty.Callable)
+
+
+# The first thing you should know about the Kubernetes SDK is that it
+# is riddled with race conditions and timeouts and all kinds of
+# horrible gremlins.  This first function/decorator is an _ongoing_
+# attempt to deal with the fallout from that.  Hopefully from now on
+# we'll be able to consolidate/maintain all of that logic in a single
+# place:
+
+
+_URLLIB_COMMON = (
+    urllib3.exceptions.ProtocolError,
+    urllib3.exceptions.MaxRetryError,
+)
+
+
+def k8s_sdk_retry(
+    get_retry_args_kwargs: ty.Optional[ty.Callable[[int], ty.Tuple[tuple, dict]]] = None,
+    should_retry: ty.Callable[[Exception], bool] = lambda _: False,
+    max_retries: int = 20,
+) -> ty.Callable[[F], F]:
+    """Handle the common cases - lets you decide about uncommon ones."""
+
+    def decorator(f: F) -> F:
+        @wraps(f)
+        def wrapper(*args, **kwargs):  # type: ignore
+            i = 0
+
+            def _raise_if_max(i: int) -> None:
+                if i >= max_retries:
+                    logger.warning(f"Failing after {i} tries")
+                    raise
+
+            while True:
+                try:
+                    return f(*args, **kwargs)
+                except Exception as ex:
+                    # some shared behavior for all exceptions means we want a single except block
+                    _raise_if_max(i)
+                    if isinstance(ex, _URLLIB_COMMON):
+                        # these are extremely common and should always be retried
+                        logger.debug(
+                            "Encountered probable connection timeout - retrying",
+                            exc=str(ex),
+                        )
+                        # necessary b/c https://github.com/kubernetes-client/python/issues/1234
+                    elif isinstance(ex, client.exceptions.ApiException) and ex.reason == "Unauthorized":
+                        # this one is fairly common - who knows why their SDK can't handle this automatically.
+                        #
+                        # https://github.com/kubernetes-client/python/blob/release-18.0/kubernetes/client/exceptions.py?ts=4#L84
+                        logger.info(f"{ex} - retrying after auth failure")
+                        auth.load_config()
+                    elif not should_retry(ex):
+                        raise
+
+                    i += 1
+                    logger.info(f"Will retry after K8S error {str(ex)}; attempt {i}")
+                time.sleep(config.k8s_monitor_delay())
+                if get_retry_args_kwargs:
+                    args, kwargs = get_retry_args_kwargs(i)
+
+        return ty.cast(F, wrapper)
+
+    return decorator

thds/mops/k8s/too_old_resource_version.py

@@ -0,0 +1,42 @@
+import re
+import typing as ty
+
+from kubernetes import client
+
+_TOO_OLD_RESOURCE_VERSION = re.compile(
+    r"Expired: too old resource version: (?P<old>\w+) \((?P<cur>\w+)\)"
+)
+# holy bananas I cannot believe how much K8s' SDK sucks.  this is a
+# standard exception with an known retry semantic that their watchers
+# are apparently unable to handle on their own - I'm staring at their
+# code right now and they don't even attempt to handle this.
+
+
+class TooOldResourceVersion(ty.NamedTuple):
+    old: str
+    cur: str
+    spread: str  # only if the above are actually numbers
+
+
+def parse_too_old_resource_version(
+    exc: Exception,
+) -> ty.Optional[TooOldResourceVersion]:
+    if not isinstance(exc, client.exceptions.ApiException):
+        return None
+    m = _TOO_OLD_RESOURCE_VERSION.match(exc.reason)
+    if m:
+        # this is a completely bonkers thing to have to do
+        # ourselves, but here we are.  I can't find any
+        # documentation on why their SDK doesn't handle this
+        # themselves, and I don't even know why we haven't run
+        # into it before. Regardless, apparently we have to
+        # special-case a retry when there are enough old
+        # events on the server.
+        resource_version = m.group("cur")
+        old = m.group("old")
+        try:
+            spread = str(int(resource_version) - int(m.group("old")))
+        except ValueError:
+            spread = "unknown"
+        return TooOldResourceVersion(old, resource_version, spread)
+    return None

thds/mops/k8s/tools/krsync.py

@@ -0,0 +1,50 @@
+"""Uses rsync to copy files to/from a Kubernetes pod.
+
+The remote pod must have rsync installed.
+
+CLI wrapper with help text for the krsync.sh script, which is usable on its own.
+Thank you, Karl Bunch, who provided the world with this elegant implementation.
+https://serverfault.com/questions/741670/rsync-files-to-a-kubernetes-pod?newreg=22b5f958cdce4e6a9a1a7ce0fc88b546
+
+When addressing the remote, you must specify a pod name, and
+optionally a namespace preceded by '@', and then a colon, followed by
+the path on the remote.
+
+Examples:
+
+krsync ~/my/local.txt pod1234:/root/local_2.txt
+krsync ~/my/local pod1234:~/local_dir -rav  # recursively copies entire directory
+krsync pod1234@my-namespace:/root/my.parquet your.parquet
+krsync prod-udla-0@unified-directory:/var/data/labels.db ./labels.db --container prod-udla-db
+"""
+
+import argparse
+import importlib
+import os
+import subprocess
+import sys
+
+with importlib.resources.path(__package__, "krsync.sh") as p:
+    krsync = str(p.resolve())
+
+
+def main() -> int:
+    remote_path = "pod-name[@namespace]:/remote/path"
+    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawTextHelpFormatter)
+    parser.add_argument("src", help=f"Either a local path or {remote_path}")
+    parser.add_argument("dest", help=f"Either a local path or {remote_path}")
+    parser.add_argument(
+        "--container",
+        "-c",
+        help="Container name - if not provided, will use the default container",
+        default="",
+    )
+    args, rsync_args = parser.parse_known_args()
+    return subprocess.run(
+        ["/bin/bash", krsync, args.src, args.dest, *rsync_args],
+        env=dict(os.environ, KRSYNC_CONTAINER=args.container or ""),
+    ).returncode
+
+
+if __name__ == "__main__":
+    sys.exit(main())

thds/mops/k8s/tools/krsync.sh

@@ -0,0 +1,22 @@
+#!/bin/bash
+# thank you, Karl Bunch:
+# https://serverfault.com/questions/741670/rsync-files-to-a-kubernetes-pod?newreg=22b5f958cdce4e6a9a1a7ce0fc88b546
+if [ -z "$KRSYNC_STARTED" ]; then
+    export KRSYNC_STARTED=true
+    exec rsync --blocking-io --rsh "$0" $@
+fi
+
+# Running as --rsh
+namespace=''
+pod=$1
+shift
+
+# If use uses pod@namespace rsync passes as: {us} -l pod namespace ...
+if [ "X$pod" = "X-l" ]; then
+    pod=$1
+    shift
+    namespace="-n $1"
+    shift
+fi
+
+exec kubectl $namespace exec -i $pod --container "${KRSYNC_CONTAINER}" -- "$@"

thds/mops/k8s/wait_job.py

@@ -0,0 +1,72 @@
+"""Wait for a Job to finish."""
+
+import time
+from datetime import timedelta
+from timeit import default_timer
+
+from thds.core import scope
+from thds.core.log import logger_context
+
+from .._utils.colorize import colorized
+from . import config
+from ._shared import logger
+from .jobs import get_job, is_job_failed, is_job_succeeded
+
+UNUSUAL = colorized(fg="white", bg="yellow")
+
+
+def _max_no_job_wait() -> timedelta:
+    return timedelta(seconds=config.k8s_monitor_max_attempts() * config.k8s_monitor_delay())
+
+
+@scope.bound
+def wait_for_job(job_name: str, short_name: str = "") -> bool:
+    """Return True if Job completed, False if it failed.
+
+    May raise an exception if something truly unusual happened.
+
+    A _lot_ has gone in to trying to make this robust against common
+    failure patterns. My apologies for the resulting shape of the
+    code. :/
+    """
+    scope.enter(logger_context(job=job_name))
+    log_name = f"Job {short_name}" if short_name else "Job"
+    logger.debug(f"Waiting for {log_name} to finish...")
+    start_time = default_timer()
+
+    def _wait_for_job() -> bool:
+        nonlocal start_time
+        found_at_least_once = False
+        while True:
+            time.sleep(0.5 if found_at_least_once else 10.0)
+            job = get_job(job_name)
+            if not job:
+                if found_at_least_once:
+                    logger.warning(UNUSUAL(f"Known job {job_name} no longer exists - assuming success!"))
+                    return True
+                max_wait_seconds = _max_no_job_wait().total_seconds()
+                if default_timer() - start_time > max_wait_seconds:
+                    logger.error(
+                        UNUSUAL(
+                            f"Job {job_name} has not been seen for {max_wait_seconds:.1f} seconds"
+                            " - assuming failure!"
+                        )
+                    )
+                    return False
+
+                logger.debug("%s not found... retrying.", job_name)
+                continue
+
+            if is_job_succeeded(job):
+                return True
+
+            if is_job_failed(job):
+                logger.error(
+                    UNUSUAL(f"A Kubernetes Job is reporting an actual failed status: {job_name}")
+                )
+                return False
+
+            found_at_least_once = True
+            start_time = default_timer()  # restart timer since the job has been found.
+
+    return _wait_for_job()

thds/mops/k8s/warn_image_backoff.py

@@ -0,0 +1,63 @@
+import threading
+import typing as ty
+from datetime import datetime, timezone
+
+from kubernetes import client
+
+from thds.core.log import getLogger
+
+from .._utils.colorize import colorized
+from . import config
+from .watch import K8sList, OneShotLimiter, yield_objects_from_list
+
+logger = getLogger(__name__)
+
+OnCoreEvent = ty.Callable[[client.CoreV1Event], ty.Any]
+
+YIKES = colorized(fg="black", bg="yellow")
+
+
+def _emit_basic(event: client.CoreV1Event) -> None:
+    logger.error(YIKES(event.message))
+
+
+def _warn_image_pull_backoff(namespace: str, on_backoff: OnCoreEvent = _emit_basic) -> None:
+    """Log scary errors when ImagePullBackoff is observed."""
+    start_dt = datetime.now(tz=timezone.utc)
+    for _ns, obj in yield_objects_from_list(
+        namespace,
+        lambda _, __: ty.cast(
+            # do NOT use client.EventsV1Api here - for some reason
+            # it does not return the right 'types' of events.
+            # why? who the heck knows? How much time did I spend
+            # trying to figure this out? Also who knows.
+            K8sList[client.CoreV1Event],
+            client.CoreV1Api().list_namespaced_event,
+        ),
+        object_type_hint="backoff-warnings",
+        field_selector="reason=BackOff",
+    ):
+        if None is obj.last_timestamp or obj.last_timestamp > start_dt:
+            on_backoff(obj)
+
+
+_WARN_IMAGE_PULL_BACKOFF = OneShotLimiter()
+
+
+def start_warn_image_pull_backoff_thread(
+    namespace: str = "", on_backoff: ty.Optional[OnCoreEvent] = None
+) -> None:
+    """Limit 1 thread per namespace per application.
+
+    You can pass an additional message context
+    """
+    namespace = namespace or config.k8s_namespace()
+
+    _WARN_IMAGE_PULL_BACKOFF(
+        namespace,
+        lambda ns: threading.Thread(
+            target=_warn_image_pull_backoff,
+            args=(namespace, on_backoff or _emit_basic),
+            daemon=True,
+        ).start(),
+    )

thds/mops/k8s/watch.py

@@ -0,0 +1,266 @@
+"""K8s SDK watching is very unreliable for lots of reasons.
+
+This is a general-purpose fix for using watchers in a thread reliably.
+"""
+
+import threading
+import time
+import typing as ty
+
+import urllib3
+from kubernetes import client
+from kubernetes import watch as k8s_watch
+
+from thds.core import scope
+from thds.core.log import getLogger, logger_context
+
+from .._utils.colorize import colorized
+from . import config
+from .auth import load_config
+from .too_old_resource_version import parse_too_old_resource_version
+
+logger = getLogger(__name__)
+
+T = ty.TypeVar("T")
+
+
+class V1List(ty.Protocol[T]):
+    api_version: str
+    items: ty.List[T]
+    kind: str
+    metadata: client.models.V1ListMeta
+
+
+class K8sList(ty.Protocol[T]):
+    def __call__(self, *args: ty.Any, namespace: str, **kwargs: ty.Any) -> V1List[T]:
+        ...
+
+
+# If this does not return a K8sList API method, the loop will exit
+GetListMethod = ty.Callable[[str, ty.Optional[Exception]], ty.Optional[K8sList[T]]]
+# if this returns True, the loop will exit.
+OnEvent = ty.Callable[[str, T], ty.Optional[bool]]
+
+
+def yield_objects_from_list(
+    namespace: str,
+    get_list_method: GetListMethod[T],
+    server_timeout: int = 10,
+    object_type_hint: str = "items",
+    init: ty.Optional[ty.Callable[[], None]] = None,
+    **kwargs: ty.Any,
+) -> ty.Iterator[ty.Tuple[str, T]]:
+    ex = None
+    if init:
+        init()
+    while True:
+        try:
+            load_config()
+            list_method = get_list_method(namespace, ex)
+            if not list_method:
+                logger.debug(f"No longer watching {object_type_hint} events in namespace: {namespace}")
+                break
+            initial_list = list_method(namespace=namespace)
+            logger.debug(
+                f"Listed {len(initial_list.items)} {object_type_hint} in namespace: {namespace}"
+            )
+            for object in initial_list.items:
+                yield namespace, object
+
+            if initial_list.metadata._continue:
+                logger.warning(
+                    f"We did not fetch the whole list of {object_type_hint} the first time..."
+                )
+            for evt in k8s_watch.Watch().stream(
+                list_method,
+                namespace=namespace,
+                resource_version=initial_list.metadata.resource_version,
+                **kwargs,
+                _request_timeout=(server_timeout, config.k8s_job_timeout_seconds()),
+            ):
+                object = evt.get("object")
+                if object:
+                    yield namespace, object
+                # once we've received events, let the resource version
+                # be managed automatically if possible.
+        except urllib3.exceptions.ProtocolError:
+            ex = None
+        except urllib3.exceptions.ReadTimeoutError:
+            ex = None
+        except Exception as e:
+            too_old = parse_too_old_resource_version(e)
+            if too_old:
+                logger.debug(f"Immediately retrying {too_old}")
+            else:
+                logger.exception(f"Unexpected exception while listing {object_type_hint}")
+            ex = e
+
+
+def callback_events(on_event: OnEvent[T], event_yielder: ty.Iterable[ty.Tuple[str, T]]) -> None:
+    """Suitable for use with a daemon thread."""
+    for namespace, event in event_yielder:
+        should_exit = on_event(namespace, event)
+        if should_exit:
+            break
+
+
+def _make_name(namespace: str, name: str) -> str:
+    return f"{namespace}/{name}"
+
+
+def _default_get_name(obj: ty.Any) -> str:
+    return obj.metadata.name
+
+
+def _default_get_namespace(obj: ty.Any) -> str:
+    return obj.metadata.namespace
+
+
+STARTING = colorized(fg="white", bg="orange")
+
+
+class OneShotLimiter:
+    """Do an action once per provided name. Does not wait for it to complete."""
+
+    def __init__(self) -> None:
+        self._lock = threading.RLock()
+        self._names: ty.Set[str] = set()
+
+    def __call__(self, name: str, shoot: ty.Callable[[str], ty.Any]) -> None:
+        """Shoot if the name has not already been shot."""
+        if name in self._names:
+            return
+        with self._lock:
+            if name in self._names:
+                return
+            shoot(name)
+            self._names.add(name)
+
+
+def is_stale(api_last_update_time: float, obj_last_seen_time: float) -> bool:
+    now = time.monotonic()
+    allowed_stale_seconds = config.k8s_watch_object_stale_seconds()
+    if (time_since_api_update := now - api_last_update_time) > allowed_stale_seconds:  # noqa: F841
+        # we haven't heard anything from the API in a while; probably
+        # the API is down. Ignore object staleness to avoid false positives.
+        return False
+
+    if not obj_last_seen_time:
+        return False  # false positives aren't worth it
+
+    return (time_since_obj_update := now - obj_last_seen_time) > allowed_stale_seconds  # noqa: F841
+
+
+class WatchingObjectSource(ty.Generic[T]):
+    """Efficiently 'get' objects by reliably watching for changes to all such objects in a given namespace.
+
+    This is network-efficient for observing many different objects,
+    but not memory efficient if you really only need to fetch details
+    for a few objects.
+    """
+
+    def __init__(
+        self,
+        get_list_method: GetListMethod[T],
+        get_name: ty.Callable[[T], str] = ty.cast(  # noqa: B008
+            ty.Callable[[T], str], _default_get_name
+        ),
+        backup_fetch: ty.Optional[ty.Callable[[str, str], T]] = None,
+        typename: str = "object",
+        starting: ty.Callable[[str], str] = STARTING,
+    ) -> None:
+        self.get_list_method = get_list_method
+        self.get_name = get_name
+        self.backup_fetch = backup_fetch
+        self.typename = typename
+        self._objs_by_name: ty.Dict[str, T] = dict()
+        # ^ is a possibly big/expensive local cache of the most recent
+        # state for all of the event type in the namespace.  Don't use
+        # this class if you can't afford the memory overhead of
+        # observing everything in your namespace and keeping the last
+        # known copy of everything forever.
+        self._last_seen_time_by_name: ty.Dict[str, float] = dict()
+        self._last_api_update_time = 0.0
+        self._limiter = OneShotLimiter()
+
+    def _start_thread(self, namespace: str) -> None:
+        threading.Thread(
+            target=callback_events,
+            args=(
+                self._add_object,
+                yield_objects_from_list(
+                    namespace,
+                    self._get_list_method_on_restart,
+                    object_type_hint=self.typename + "s",
+                    init=lambda: logger.info(STARTING(f"Watching {self.typename}s in {namespace}")),
+                ),
+            ),
+            daemon=True,
+        ).start()
+
+    def _add_object(self, namespace: str, obj: T) -> None:
+        """This is where we receive updates from the k8s API."""
+        self._last_api_update_time = time.monotonic()
+
+        if not obj:
+            logger.warning(f"Received null/empty {self.typename}")
+            return
+
+        name = _make_name(namespace, self.get_name(obj))
+        logger.debug(f"{self.typename} {name} updated")
+        self._last_seen_time_by_name[name] = time.monotonic()
+        self._objs_by_name[name] = obj
+
+    def _get_list_method_on_restart(
+        self, namespace: str, exc: ty.Optional[Exception]
+    ) -> ty.Optional[K8sList[T]]:
+        suffix = ""
+        if exc:
+            too_old = parse_too_old_resource_version(exc)
+            if not too_old:
+                logger.exception(f"Not fatal, but sleeping before we retry {self.typename} scraping...")
+                time.sleep(config.k8s_monitor_delay())
+                suffix = f" after {type(exc).__name__}: {exc}"
+                logger.info(f"Watching {self.typename}s in namespace: {namespace}{suffix}")
+        return self.get_list_method(namespace, exc)
+
+    def _is_stale(self, name: str) -> bool:
+        return is_stale(self._last_api_update_time, self._last_seen_time_by_name.get(name) or 0)
+
+    @scope.bound
+    def get(self, obj_name: str, namespace: str = "") -> ty.Optional[T]:
+        namespace = namespace or config.k8s_namespace()
+        name = _make_name(namespace, obj_name)
+        scope.enter(logger_context(name=obj_name, namespace=namespace))
+
+        # first try is looking in our local cache
+        if (obj := self._objs_by_name.get(name)) and not self._is_stale(name):
+            return obj
+
+        # second try is making sure the namespace watcher is running, sleeping, and then looking in the cache again.
+        # This is much more efficient than a manual fetch.
+        self._limiter(namespace, self._start_thread)
+        time.sleep(config.k8s_monitor_delay())
+        if (obj := self._objs_by_name.get(name)) and not self._is_stale(name):
+            return obj
+
+        # if that doesn't work, try a manual fetch.
+        if self.backup_fetch:
+            logger.warning(f"Manually fetching {self.typename}...")
+            # doing a lot of manual fetches may indicate that the k8s API is having trouble keeping up...
+            try:
+                if obj := self.backup_fetch(namespace, obj_name):
+                    self._add_object(namespace, obj)  # updates last seen, too
+                    return obj
+
+            except Exception:
+                logger.exception(f"Unexpected error during manual fetch of {self.typename}.")
+
+        if self._is_stale(name):
+            logger.warning(
+                f"Could not refresh {name}, and our record of it is stale - dropping stale object!"
+            )
+            self._objs_by_name.pop(name, None)
+            self._last_seen_time_by_name.pop(name, None)
+
+        return None

thds/mops/meta.json

@@ -0,0 +1,8 @@
+{
+  "git_commit": "affac0a04726de27f8065525d193df4cd6376b9c",
+  "git_branch": "task/dbxtend/use-pure-magic",
+  "git_is_clean": true,
+  "pyproject_version": "3.6.20250219172032",
+  "thds_user": "peter.gaultney",
+  "misc": {}
+}