thds.mops 3.6.20250219172032__py3-none-any.whl

thds/mops/pure/tools/summarize/cli.py

@@ -0,0 +1,293 @@
+import argparse
+import json
+import statistics
+import sys
+import typing as ty
+from functools import reduce
+from pathlib import Path
+from typing import Dict, List, Literal, Optional, Set, TypedDict
+
+from thds.mops.pure.core.memo.function_memospace import parse_memo_uri
+from thds.mops.pure.tools.summarize import run_summary
+
+SortOrder = Literal["name", "time"]
+
+
+class FunctionSummary(TypedDict):
+    total_calls: int
+    cache_hits: int
+    executed: int
+    error_count: int
+    timestamps: List[str]
+    runner_prefixes: Set[str]
+    pipeline_ids: Set[str]
+    function_logic_keys: Set[str]
+    invoked_by: List[str]
+    invoker_code_version: List[str]
+    remote_code_version: List[str]
+    total_runtime_minutes: List[float]  # minutes
+    remote_runtime_minutes: List[float]  # minutes
+    uris_in_rvalue: List[str]
+
+
+def _empty_summary() -> FunctionSummary:
+    return {
+        "total_calls": 0,
+        "cache_hits": 0,
+        "executed": 0,
+        "timestamps": [],
+        "runner_prefixes": set(),
+        "pipeline_ids": set(),
+        "function_logic_keys": set(),
+        "error_count": 0,
+        "invoked_by": list(),
+        "invoker_code_version": list(),
+        "remote_code_version": list(),
+        "total_runtime_minutes": list(),
+        "remote_runtime_minutes": list(),
+        "uris_in_rvalue": list(),
+    }
+
+
+def _process_log_file(log_file: Path) -> Dict[str, FunctionSummary]:
+    """
+    Process a single JSON log file and return a partial summary.
+    :param log_file: Path to the log file
+    :return: A dictionary with the function names as keys and their execution summaries as values
+    """
+    partial_summary: Dict[str, FunctionSummary] = {}
+    with log_file.open("r") as f:
+        try:
+            log_entry: run_summary.LogEntry = json.load(f)
+        except json.JSONDecodeError:
+            print(f"Error reading log file '{log_file}'")
+            return dict()
+
+        function_name = log_entry["function_name"]
+        if function_name not in partial_summary:
+            partial_summary[function_name] = _empty_summary()
+
+        summary = partial_summary[function_name]
+
+        summary["total_calls"] += 1
+        if log_entry["status"] in ("memoized", "awaited"):
+            summary["cache_hits"] += 1
+        else:
+            summary["executed"] += 1
+        summary["error_count"] += int(log_entry.get("was_error") or 0)
+        summary["timestamps"].append(log_entry["timestamp"])
+        summary["uris_in_rvalue"].extend(log_entry.get("uris_in_rvalue") or tuple())
+
+        mu_parts = parse_memo_uri(
+            log_entry["memo_uri"], runner_prefix=log_entry.get("runner_prefix", "")
+        )
+
+        summary["runner_prefixes"].add(mu_parts.runner_prefix)
+        summary["pipeline_ids"].add(mu_parts.pipeline_id)
+        summary["function_logic_keys"].add(mu_parts.function_logic_key)
+
+        # new metadata stuff below:
+        def append_if_exists(key: str) -> None:
+            if key in log_entry:
+                summary[key].append(log_entry[key])  # type: ignore
+
+        for key in (
+            "invoked_by",
+            "invoker_code_version",
+            "remote_code_version",
+            "total_runtime_minutes",
+            "remote_runtime_minutes",
+        ):
+            append_if_exists(key)
+
+    return partial_summary
+
+
+def _combine_summaries(
+    acc: Dict[str, FunctionSummary], partial: Dict[str, FunctionSummary]
+) -> Dict[str, FunctionSummary]:
+    """
+    Combine two summaries into one
+    :param acc: the accumulator summary
+    :param partial: A partial summary to be combined with the accumulator
+    :return: the combined summary
+    """
+    for function_name, data in partial.items():
+        if function_name not in acc:
+            acc[function_name] = _empty_summary()
+        acc[function_name]["total_calls"] += data["total_calls"]
+        acc[function_name]["cache_hits"] += data["cache_hits"]
+        acc[function_name]["executed"] += data["executed"]
+        acc[function_name]["error_count"] += data["error_count"]
+        acc[function_name]["timestamps"].extend(data["timestamps"])
+        acc[function_name]["runner_prefixes"].update(data["runner_prefixes"])
+        acc[function_name]["pipeline_ids"].update(data["pipeline_ids"])
+        acc[function_name]["function_logic_keys"].update(data["function_logic_keys"])
+        acc[function_name]["uris_in_rvalue"].extend(data["uris_in_rvalue"])
+
+        for key in (
+            "invoked_by",
+            "invoker_code_version",
+            "remote_code_version",
+            "total_runtime_minutes",
+            "remote_runtime_minutes",
+        ):
+            acc[function_name][key].extend(data[key])  # type: ignore
+
+    return acc
+
+
+def _format_summary(summary: Dict[str, FunctionSummary], sort_by: SortOrder, uri_limit: int = 10) -> str:
+    """
+    Format a summary into a readable report
+    """
+    template = (
+        "Function '{function_name}':\n"
+        "  Total calls: {total_calls}\n"
+        "  Cache hits: {cache_hits}\n"
+        "  Executed: {executed}\n"
+        "  Error count: {error_count}\n"
+        "  Timestamps: {timestamps}\n"
+        "  Runner Prefixes: {runner_prefixes}\n"
+        "  Pipeline IDs: {pipeline_ids}\n"
+        "  Function Logic Keys: {function_logic_keys}\n"
+        "  Function Runtime minutes: {function_runtimes}\n"
+        "  Wall clock minutes: {wall_clock_runtimes}\n"
+        "  Invoked by: {invokers}\n"
+        "  Invoker code versions: {invoker_code_version}\n"
+        "  Remote code versions: {remote_code_version}\n"
+    )
+    report_lines = []
+
+    sorted_items = (
+        sorted(summary.items(), key=lambda item: item[0])
+        if sort_by == "name"
+        else sorted(summary.items(), key=lambda item: min(item[1]["timestamps"]))
+    )
+
+    for function_name, data in sorted_items:
+
+        def first_and_last_n(
+            obj_set: ty.Collection[str], n: int
+        ) -> ty.Tuple[ty.List[str], ty.List[str], int]:
+            """take the first n and the last n, unless they would overlap, in which case take the whole list"""
+            if len(obj_set) <= n * 2:
+                return list(obj_set), list(), 0
+            obj_list = list(obj_set)
+            return obj_list[:n], obj_list[-n:], len(obj_set) - n * 2
+
+        def and_more(obj_set: ty.Collection[str], max_count: int = 4) -> str:
+            if max_count < 1:
+                return ""
+            if max_count == 1:
+                max_count = 2  # stupid, but keeps the code simpler.
+            the_first, the_last, remaining_count = first_and_last_n(obj_set, max_count // 2)
+            return ", ".join(
+                [
+                    *the_first,
+                    *([f"...skipping {remaining_count} more..."] if remaining_count else list()),
+                    *the_last,
+                ]
+            )
+
+        def describe(fs: FunctionSummary, key: str) -> str:
+            numlist: ty.List[float] = fs[key]  # type: ignore
+            if not numlist:
+                return ""
+
+            avg = sum(numlist) / len(numlist)
+            maxi = max(numlist)
+            mini = min(numlist)
+            pstddev = statistics.pstdev(numlist)
+            return f"avg: {avg:.2f}, min: {mini:.2f}, max: {maxi:.2f}, pstdev: {pstddev:.2f}"
+
+        report_lines.append(
+            template.format(
+                function_name=function_name,
+                total_calls=data["total_calls"],
+                cache_hits=data["cache_hits"],
+                executed=data["executed"],
+                error_count=data["error_count"],
+                timestamps=and_more(sorted(data["timestamps"])),
+                runner_prefixes=and_more(data["runner_prefixes"]),
+                pipeline_ids=", ".join(data["pipeline_ids"]),
+                function_logic_keys=", ".join(data["function_logic_keys"]),
+                function_runtimes=describe(data, "remote_runtime_minutes"),
+                wall_clock_runtimes=describe(data, "total_runtime_minutes"),
+                invokers=", ".join(sorted(set(data["invoked_by"]))),
+                invoker_code_version=", ".join(sorted(set(data["invoker_code_version"]))),
+                remote_code_version=", ".join(sorted(set(data["remote_code_version"]))),
+            )
+        )
+        n_uris = and_more(
+            sorted(data["uris_in_rvalue"]),
+            max_count=uri_limit if uri_limit >= 0 else sys.maxsize,
+        ).replace(", ", "\n     ")
+        if n_uris:
+            report_lines.append(f"  URIs in return value(s):\n     {n_uris}\n")
+    return "\n".join(report_lines)
+
+
+def _auto_find_run_directory() -> ty.Optional[Path]:
+    mops_root = run_summary.MOPS_SUMMARY_DIR()
+    if not mops_root.exists():
+        raise ValueError(f"No mops summary root directory found at {mops_root}.")
+    if not mops_root.is_dir():
+        raise RuntimeError(
+            "Mops summary root is not a directory! "
+            f"Delete {mops_root} to allow mops to recreate it on the next run."
+        )
+    for directory in sorted(mops_root.iterdir(), key=lambda x: x.name, reverse=True):
+        if directory.is_dir() and list(directory.glob("*.json")):
+            # needs to have some files for it to count for anything
+            return directory
+
+    print("No pipeline run directories found.")
+    return None
+
+
+def summarize(
+    run_directory: Optional[str] = None, sort_by: SortOrder = "name", uri_limit: int = 10
+) -> None:
+    run_directory_path = Path(run_directory) if run_directory else _auto_find_run_directory()
+    if not run_directory_path:
+        return
+
+    print(f"Summarizing pipeline run '{run_directory_path}'\n")
+    log_files = list(run_directory_path.glob("*.json"))
+
+    partial_summaries = map(_process_log_file, log_files)
+
+    summary: Dict[str, FunctionSummary] = reduce(_combine_summaries, partial_summaries, {})
+
+    report = _format_summary(summary, sort_by, uri_limit)
+    print(report)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Summarize mops pipeline run logs.")
+    parser.add_argument(
+        "run_directory",
+        nargs="?",
+        type=str,
+        default=None,
+        help="Path to the pipeline run directory. If not provided, the latest run directory will be used.",
+    )
+    parser.add_argument(
+        "--sort-by",
+        choices=["name", "time"],
+        default="time",
+        help="Sort the summary by function name or by the first call time",
+    )
+    parser.add_argument(
+        "--uri-limit",
+        type=int,
+        default=10,
+        help=(
+            "Limit the number of Source URIs printed in the summary for each function."
+            " Grep for lines beginning with 5 spaces to get only the URIs."
+            " Negative numbers (e.g. -1) mean no limit."
+        ),
+    )
+    args = parser.parse_args()
+    summarize(args.run_directory, args.sort_by, args.uri_limit)

thds/mops/pure/tools/summarize/run_summary.py

@@ -0,0 +1,143 @@
+import datetime as dt
+import json
+import os
+import pickle
+import typing as ty
+import uuid
+from pathlib import Path
+
+from thds.core import config, log, pickle_visit, source
+from thds.mops.pure.core.memo import function_memospace
+from thds.mops.pure.core.metadata import get_invoked_by
+from thds.mops.pure.core.types import T
+
+from ...core import metadata
+
+MOPS_SUMMARY_DIR = config.item("thds.mops.summary.dir", default=Path(".mops/summary"), parse=Path)
+RUN_NAME = config.item(
+    "thds.mops.summary.run_name",
+    default=f"{dt.datetime.utcnow().isoformat()}-pid{os.getpid()}-{get_invoked_by()}",
+)
+
+InvocationType = ty.Literal["memoized", "invoked", "awaited"]
+
+logger = log.getLogger(__name__)
+
+
+class LogEntryV1(ty.TypedDict):
+    function_name: str
+    memo_uri: str
+    timestamp: str  # more or less "when did this complete?"
+    status: InvocationType  # old name that we're retaining for compatibility
+
+
+class LogEntry(LogEntryV1, total=False):
+    runner_prefix: str  # includes env and any prefixes like mops2-mpf
+    pipeline_id: str
+    function_logic_key: str
+    was_error: bool
+
+    total_runtime_minutes: float
+    remote_runtime_minutes: float
+    invoked_by: str
+    invoker_code_version: str
+    remote_code_version: str
+
+    uris_in_rvalue: ty.List[str]
+
+
+def create_mops_run_directory() -> Path:
+    # Define the root directory for mops logs
+    mops_root = MOPS_SUMMARY_DIR()
+    # Use run name if set, otherwise fallback to orchestrator datetime
+    run_name = RUN_NAME()
+    # Create a subdirectory named with the orchestrator datetime and run identifier
+    run_directory = mops_root / run_name
+    try:
+        run_directory.mkdir(parents=True, exist_ok=True)
+    except Exception:
+        if mops_root.exists() and not mops_root.is_dir():
+            # this is going to cause errors later on!
+            logger.error(
+                f"mops summary directory must be a directory: '{mops_root}'"
+                " Please delete this file and allow mops to recreate it!"
+            )
+        else:
+            raise
+
+    return run_directory
+
+
+def _generate_log_filename(run_directory: Path) -> Path:
+    """Generate a log filename using the current timestamp and a short UUID, ensuring uniqueness"""
+    timestamp = dt.datetime.utcnow().strftime("%Y%m%d%H%M%S")
+    short_uuid = str(uuid.uuid4())[:8]
+    filename = f"{timestamp}-{short_uuid}.json"
+    return run_directory / filename
+
+
+def _extract_source_uris(result: ty.Any) -> ty.Set[str]:
+    sources: ty.List[source.Source] = list()
+
+    def extract_source(unknown: ty.Any) -> None:
+        if isinstance(unknown, source.Source):
+            sources.append(unknown)
+
+    try:
+        pickle_visit.recursive_visit(extract_source, result)
+    except pickle.PicklingError:
+        pass
+    except Exception as exc:
+        logger.warning(f'Unexpected error trying to extract source URIs from "%s"; {exc}', result)
+
+    return {source.uri for source in sources}
+
+
+def log_function_execution(
+    run_directory: ty.Optional[Path],
+    func: ty.Callable[..., T],
+    memo_uri: str,
+    itype: InvocationType,
+    metadata: ty.Optional[metadata.ResultMetadata] = None,
+    runner_prefix: str = "",
+    was_error: bool = False,
+    return_value: ty.Any = None,
+) -> None:
+    if not run_directory:
+        logger.debug("Not writing function summary for %s", memo_uri)
+        return
+
+    log_file = _generate_log_filename(run_directory)
+    func_module = func.__module__
+    func_name = func.__name__
+    full_function_name = f"{func_module}:{func_name}"
+
+    parts = function_memospace.parse_memo_uri(memo_uri, runner_prefix)
+
+    log_entry: LogEntry = {
+        "function_name": full_function_name,
+        "memo_uri": memo_uri,
+        "runner_prefix": parts.runner_prefix,
+        "pipeline_id": parts.pipeline_id,
+        "function_logic_key": parts.function_logic_key,
+        "timestamp": dt.datetime.utcnow().isoformat(),
+        "status": itype,
+        "was_error": was_error,
+    }
+    if metadata:
+        log_entry["total_runtime_minutes"] = metadata.result_wall_minutes
+        log_entry["remote_runtime_minutes"] = metadata.remote_wall_minutes
+        log_entry["invoked_by"] = metadata.invoked_by
+        log_entry["invoker_code_version"] = metadata.invoker_code_version
+        log_entry["remote_code_version"] = metadata.remote_code_version
+        # we don't bother with invoked_at or remote_started_at because they can be
+        # inferred from the timestamp and the wall times
+    if source_uris := _extract_source_uris(return_value):
+        log_entry["uris_in_rvalue"] = sorted(source_uris)
+
+    try:
+        assert not log_file.exists(), f"Log file '{log_file}' should not already exist"
+        with log_file.open("w") as f:
+            json.dump(log_entry, f, indent=2)
+    except Exception:
+        logger.exception(f"Unable to write mops function invocation log file at '{log_file}'")

thds/mops/testing/deferred_imports.py

@@ -0,0 +1,81 @@
+import ast
+import itertools
+import re
+import sys
+import typing as ty
+from contextlib import contextmanager
+
+from thds.core.log import getLogger
+
+
+def module_name_re(modules: ty.Collection[str]) -> ty.Pattern[str]:
+    name = "|".join(modules)
+    return re.compile(rf"^({name})(?:\.|$)")
+
+
+def module_names_from_import_statement(import_stmt: str) -> ty.Set[str]:
+    statements = ast.parse(import_stmt).body
+
+    def _extract_imports(imp: ty.Any) -> ty.Iterable[str]:
+        names: ty.Iterable[ty.Optional[str]]
+        if isinstance(imp, ast.Import):
+            names = (n.name for n in imp.names)
+        elif isinstance(imp, ast.ImportFrom):
+            names = (imp.module,)
+        else:
+            names = ()
+        return filter(None, names)
+
+    def _extract_ancestors(module: str) -> ty.Iterable[str]:
+        parts = module.split(".")
+        return (".".join(parts[:i]) for i in range(1, len(parts) + 1))
+
+    imported_modules = itertools.chain.from_iterable(map(_extract_imports, statements))
+    all_imported_modules = itertools.chain.from_iterable(map(_extract_ancestors, imported_modules))
+    return set(all_imported_modules)
+
+
+@contextmanager
+def clear_and_restore_import_cache(module_name_filter: ty.Callable[[str], ty.Any]) -> ty.Iterator[None]:
+    already_imported = [name for name in sys.modules if module_name_filter(name)]
+    if already_imported:
+        getLogger(__name__).debug(
+            "Clearing the following from sys.modules matching %s:\n  %s",
+            module_name_filter,
+            "\n  ".join(already_imported),
+        )
+    to_restore = {name: sys.modules.pop(name) for name in already_imported}
+    try:
+        yield
+    finally:
+        sys.modules.update(to_restore)
+
+
+def assert_dev_deps_not_imported(import_statement: str, forbidden_modules: ty.Collection[str]) -> None:
+    """One of the primary features of `mops` is to provide global memoization of pure function calls
+    using remote storage mechanisms. Sometimes, as a library author, you'd like to pre-compute the
+    result of such a function call, memoizing it and making it available to downstream users without
+    requiring them to perform the computation themselves. As such, it is useful to export a public
+    interface where such functions can be imported and called to achieve a cache hit and download the
+    result locally, _without_ requiring that all the dependencies needed to _compute_ the result be
+    present; only `mops` itself need be present to fetch the memoized result. This function can be used
+    in your test suite to assert that this condition is met for any import statements that a downstream
+    user might use to access your memoized functions.
+
+    :param import_statement: The import statement to test, as a string
+    :param forbidden_modules: Module names that should _not_ be imported in the course of executing
+      `import_statement`.
+    :raises AssertionError: When any of the `forbidden_modules` or their submodules were imported in the
+      course of executing `import_statement`
+    """
+    is_forbidden = module_name_re(forbidden_modules).match
+    # ensure that we clear the cache of the actually imported modules, lest we get a spurious pass
+    # due to the interpreter not evaluating them again!
+    imported_modules = module_names_from_import_statement(import_statement)
+    will_be_imported = module_name_re(imported_modules).match
+    with clear_and_restore_import_cache(lambda name: is_forbidden(name) or will_be_imported(name)):
+        exec(import_statement, {}, {})
+        mistakenly_imported = [name for name in sys.modules if is_forbidden(name)]
+        assert (
+            not mistakenly_imported
+        ), f"Modules {', '.join(mistakenly_imported)} were imported on execution of {import_statement!r}"

thds.mops-3.6.20250219172032.dist-info/METADATA

@@ -0,0 +1,42 @@
+Metadata-Version: 2.2
+Name: thds.mops
+Version: 3.6.20250219172032
+Summary: ML Ops tools for Trilliant Health
+Author: Trilliant Health
+Description-Content-Type: text/markdown
+Requires-Dist: ansicolors
+Requires-Dist: azure-core
+Requires-Dist: azure-identity
+Requires-Dist: azure-storage-file-datalake
+Requires-Dist: cachetools
+Requires-Dist: tblib<3.0.0,>=2.0.0
+Requires-Dist: thds.adls>=3.1
+Requires-Dist: thds.core>=1.32
+Requires-Dist: thds.humenc>=1.0
+Requires-Dist: tomli; python_version < "3.11"
+Provides-Extra: k8s
+Requires-Dist: kubernetes>=18.20; extra == "k8s" and extra == "k8s"
+
+# `mops`
+
+`mops` is a Python library for ML Operations.
+
+`mops` solves for three core issues:
+
+- Transfer of
+  [pure](https://github.com/TrilliantHealth/trilliant-data-science/libs/mops/docs/pure_functions.adoc)
+  function execution to
+  [remote](https://github.com/TrilliantHealth/trilliant-data-science/libs/mops/docs/remote.adoc)
+  execution environments with more &| different compute resources
+- [Efficient](https://github.com/TrilliantHealth/trilliant-data-science/libs/mops/docs/optimizations.adoc)
+  transfer of large blob data to/from other environments.
+- [Memoization](https://github.com/TrilliantHealth/trilliant-data-science/libs/mops/docs/memoization.adoc)
+  — i.e. _reproducibility and fault tolerance_ — for individual functions.
+
+It is used by
+[decorating or wrapping your pure function and then calling it](https://github.com/TrilliantHealth/trilliant-data-science/libs/mops/docs/basic_usage.adoc)
+like a normal function.
+
+### read the docs
+
+[Browse our full documentation here.](https://github.com/TrilliantHealth/trilliant-data-science/libs/mops/README.adoc)