genesis-flow 1.0.0__py3-none-any.whl

mlflow/utils/__init__.py

@@ -0,0 +1,283 @@
+import inspect
+import logging
+import socket
+import subprocess
+import uuid
+from contextlib import closing
+from itertools import islice
+from sys import version_info
+
+from mlflow.utils.pydantic_utils import IS_PYDANTIC_V2_OR_NEWER  # noqa: F401
+
+PYTHON_VERSION = f"{version_info.major}.{version_info.minor}.{version_info.micro}"
+
+
+_logger = logging.getLogger(__name__)
+
+
+def get_major_minor_py_version(py_version):
+    return ".".join(py_version.split(".")[:2])
+
+
+def reraise(tp, value, tb=None):
+    # Taken from: https://github.com/benjaminp/six/blob/1.15.0/six.py#L694-L700
+    try:
+        if value is None:
+            value = tp()
+        if value.__traceback__ is not tb:
+            raise value.with_traceback(tb)
+        raise value
+    finally:
+        value = None
+        tb = None
+
+
+def chunk_list(l, chunk_size):
+    for i in range(0, len(l), chunk_size):
+        yield l[i : i + chunk_size]
+
+
+def _chunk_dict(d, chunk_size):
+    """
+    Splits a dictionary into chunks of the specified size.
+    Taken from: https://stackoverflow.com/a/22878842
+    """
+    it = iter(d)
+    for _ in range(0, len(d), chunk_size):
+        yield {k: d[k] for k in islice(it, chunk_size)}
+
+
+def _truncate_and_ellipsize(value, max_length):
+    """
+    Truncates the string representation of the specified value to the specified
+    maximum length, if necessary. The end of the string is ellipsized if truncation occurs
+    """
+    value = str(value)
+    if len(value) > max_length:
+        return value[: (max_length - 3)] + "..."
+    else:
+        return value
+
+
+def _truncate_dict(d, max_key_length=None, max_value_length=None):
+    """
+    Truncates keys and/or values in a dictionary to the specified maximum length.
+    Truncated items will be converted to strings and ellipsized.
+    """
+    key_is_none = max_key_length is None
+    val_is_none = max_value_length is None
+
+    if key_is_none and val_is_none:
+        raise ValueError("Must specify at least either `max_key_length` or `max_value_length`")
+
+    truncated = {}
+    for k, v in d.items():
+        should_truncate_key = (not key_is_none) and (len(str(k)) > max_key_length)
+        should_truncate_val = (not val_is_none) and (len(str(v)) > max_value_length)
+
+        new_k = _truncate_and_ellipsize(k, max_key_length) if should_truncate_key else k
+        if should_truncate_key:
+            # Use the truncated key for warning logs to avoid noisy printing to stdout
+            msg = f"Truncated the key `{new_k}`"
+            _logger.warning(msg)
+
+        new_v = _truncate_and_ellipsize(v, max_value_length) if should_truncate_val else v
+        if should_truncate_val:
+            # Use the truncated key and value for warning logs to avoid noisy printing to stdout
+            msg = f"Truncated the value of the key `{new_k}`. Truncated value: `{new_v}`"
+            _logger.warning(msg)
+
+        truncated[new_k] = new_v
+
+    return truncated
+
+
+def merge_dicts(dict_a, dict_b, raise_on_duplicates=True):
+    """This function takes two dictionaries and returns one singular merged dictionary.
+
+    Args:
+        dict_a: The first dictionary.
+        dict_b: The second dictionary.
+        raise_on_duplicates: If True, the function raises ValueError if there are duplicate keys.
+            Otherwise, duplicate keys in `dict_b` will override the ones in `dict_a`.
+
+    Returns:
+        A merged dictionary.
+
+    """
+    duplicate_keys = dict_a.keys() & dict_b.keys()
+    if raise_on_duplicates and len(duplicate_keys) > 0:
+        raise ValueError(f"The two merging dictionaries contains duplicate keys: {duplicate_keys}.")
+    return {**dict_a, **dict_b}
+
+
+def _get_fully_qualified_class_name(obj):
+    """
+    Obtains the fully qualified class name of the given object.
+    """
+    return obj.__class__.__module__ + "." + obj.__class__.__name__
+
+
+def _inspect_original_var_name(var, fallback_name):
+    """
+    Inspect variable name, will search above frames and fetch the same instance variable name
+    in the most outer frame.
+    If inspect failed, return fallback_name
+    """
+    if var is None:
+        return fallback_name
+    try:
+        original_var_name = fallback_name
+
+        frame = inspect.currentframe().f_back
+        while frame is not None:
+            arg_info = inspect.getargvalues(frame)
+
+            fixed_args = [arg_info.locals[arg_name] for arg_name in arg_info.args]
+            varlen_args = list(arg_info.locals[arg_info.varargs]) if arg_info.varargs else []
+            keyword_args = (
+                list(arg_info.locals[arg_info.keywords].values()) if arg_info.keywords else []
+            )
+
+            all_args = fixed_args + varlen_args + keyword_args
+
+            # check whether `var` is in arg list first. If yes, go to check parent frame.
+            if any(var is arg for arg in all_args):
+                # the var is passed in from caller, check parent frame.
+                frame = frame.f_back
+                continue
+
+            for var_name, var_val in frame.f_locals.items():
+                if var_val is var:
+                    original_var_name = var_name
+                    break
+
+            break
+
+        return original_var_name
+
+    except Exception:
+        return fallback_name
+
+
+def find_free_port():
+    """
+    Find free socket port on local machine.
+    """
+    with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s:
+        s.bind(("", 0))
+        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        return s.getsockname()[1]
+
+
+def check_port_connectivity():
+    port = find_free_port()
+    try:
+        with subprocess.Popen(
+            ["nc", "-l", "-p", str(port)],
+            stdout=subprocess.DEVNULL,
+            stderr=subprocess.DEVNULL,
+        ) as server:
+            with subprocess.Popen(
+                ["nc", "-zv", "localhost", str(port)],
+                stdout=subprocess.DEVNULL,
+                stderr=subprocess.DEVNULL,
+            ) as client:
+                client.wait()
+                server.terminate()
+                return client.returncode == 0
+    except Exception as e:
+        _logger.warning("Failed to check port connectivity: %s", e)
+        return False
+
+
+def is_iterator(obj):
+    """
+    Args:
+        obj: Any object.
+
+    Returns:
+        Boolean representing whether or not 'obj' is an iterator.
+    """
+    return (hasattr(obj, "__next__") or hasattr(obj, "next")) and hasattr(obj, "__iter__")
+
+
+def _is_in_ipython_notebook():
+    try:
+        from IPython import get_ipython
+
+        return get_ipython() is not None
+    except Exception:
+        return False
+
+
+def get_results_from_paginated_fn(paginated_fn, max_results_per_page, max_results=None):
+    """Gets results by calling the ``paginated_fn`` until either no more results remain or
+    the specified ``max_results`` threshold has been reached.
+
+    Args:
+        paginated_fn: This function is expected to take in the number of results to retrieve
+            per page and a pagination token, and return a PagedList object.
+        max_results_per_page: The maximum number of results to retrieve per page.
+        max_results: The maximum number of results to retrieve overall. If unspecified,
+            all results will be retrieved.
+
+    Returns:
+        Returns a list of entities, as determined by the paginated_fn parameter, with no more
+        entities than specified by max_results.
+
+    """
+    all_results = []
+    next_page_token = None
+    returns_all = max_results is None
+    while returns_all or len(all_results) < max_results:
+        num_to_get = max_results_per_page if returns_all else max_results - len(all_results)
+        if num_to_get < max_results_per_page:
+            page_results = paginated_fn(num_to_get, next_page_token)
+        else:
+            page_results = paginated_fn(max_results_per_page, next_page_token)
+        all_results.extend(page_results)
+        if hasattr(page_results, "token") and page_results.token:
+            next_page_token = page_results.token
+        else:
+            break
+    return all_results
+
+
+class AttrDict(dict):
+    """
+    Dict-like object that exposes its keys as attributes.
+
+    Examples
+    --------
+    >>> d = AttrDict({"a": 1, "b": 2})
+    >>> d.a
+    1
+    >>> d = AttrDict({"a": 1, "b": {"c": 3, "d": 4}})
+    >>> d.b.c
+    3
+    """
+
+    def __getattr__(self, attr):
+        try:
+            value = self[attr]
+        except KeyError:
+            raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{attr}'")
+        if isinstance(value, dict):
+            return AttrDict(value)
+        return value
+
+
+def get_parent_module(module):
+    return module[0 : module.rindex(".")]
+
+
+def is_uuid(s: str) -> bool:
+    """
+    Returns True if the specified string is a UUID, False otherwise.
+    """
+    try:
+        uuid.UUID(s)
+        return True
+    except ValueError:
+        return False

mlflow/utils/_capture_modules.py

@@ -0,0 +1,256 @@
+"""
+This script should be executed in a fresh python interpreter process using `subprocess`.
+"""
+
+import argparse
+import builtins
+import functools
+import importlib
+import json
+import os
+import sys
+
+import mlflow
+from mlflow.models.model import MLMODEL_FILE_NAME, Model
+from mlflow.pyfunc import MAIN
+from mlflow.utils._spark_utils import _prepare_subprocess_environ_for_creating_local_spark_session
+from mlflow.utils.exception_utils import get_stacktrace
+from mlflow.utils.file_utils import write_to
+from mlflow.utils.requirements_utils import (
+    DATABRICKS_MODULES_TO_PACKAGES,
+    MLFLOW_MODULES_TO_PACKAGES,
+)
+
+
+def _get_top_level_module(full_module_name):
+    return full_module_name.split(".")[0]
+
+
+def _get_second_level_module(full_module_name):
+    return ".".join(full_module_name.split(".")[:2])
+
+
+class _CaptureImportedModules:
+    """
+    A context manager to capture imported modules by temporarily applying a patch to
+    `builtins.__import__` and `importlib.import_module`.
+
+    If `record_full_module` is set to `False`, it only captures top level modules
+    for inferring python package purpose.
+    If `record_full_module` is set to `True`, it captures full module name for all
+    imported modules and sub-modules. This is used in automatic model code path inference.
+    """
+
+    def __init__(self, record_full_module=False):
+        self.imported_modules = set()
+        self.original_import = None
+        self.original_import_module = None
+        self.record_full_module = record_full_module
+
+    def _wrap_import(self, original):
+        @functools.wraps(original)
+        def wrapper(name, globals=None, locals=None, fromlist=(), level=0):
+            is_absolute_import = level == 0
+            if not self.record_full_module and is_absolute_import:
+                self._record_imported_module(name)
+
+            result = original(name, globals, locals, fromlist, level)
+
+            if self.record_full_module:
+                if is_absolute_import:
+                    parent_modules = name.split(".")
+                else:
+                    parent_modules = globals["__name__"].split(".")
+                    if level > 1:
+                        parent_modules = parent_modules[: -(level - 1)]
+
+                if fromlist:
+                    for from_name in fromlist:
+                        full_modules = parent_modules + [from_name]
+                        full_module_name = ".".join(full_modules)
+                        if full_module_name in sys.modules:
+                            self._record_imported_module(full_module_name)
+                else:
+                    full_module_name = ".".join(parent_modules)
+                    self._record_imported_module(full_module_name)
+
+            return result
+
+        return wrapper
+
+    def _wrap_import_module(self, original):
+        @functools.wraps(original)
+        def wrapper(name, *args, **kwargs):
+            self._record_imported_module(name)
+            return original(name, *args, **kwargs)
+
+        return wrapper
+
+    def _record_imported_module(self, full_module_name):
+        if self.record_full_module:
+            self.imported_modules.add(full_module_name)
+            return
+
+        # If the module is an internal module (prefixed by "_") or is the "databricks"
+        # module, which is populated by many different packages, don't record it (specific
+        # module imports within the databricks namespace are still recorded and mapped to
+        # their corresponding packages)
+        if full_module_name.startswith("_") or full_module_name == "databricks":
+            return
+
+        top_level_module = _get_top_level_module(full_module_name)
+        second_level_module = _get_second_level_module(full_module_name)
+
+        if top_level_module == "databricks":
+            # Multiple packages populate the `databricks` module namespace on Databricks;
+            # to avoid bundling extraneous Databricks packages into model dependencies, we
+            # scope each module to its relevant package
+            if second_level_module in DATABRICKS_MODULES_TO_PACKAGES:
+                self.imported_modules.add(second_level_module)
+                return
+
+            for databricks_module in DATABRICKS_MODULES_TO_PACKAGES:
+                if full_module_name.startswith(databricks_module):
+                    self.imported_modules.add(databricks_module)
+                    return
+
+        # special casing for mlflow extras since they may not be required by default
+        if top_level_module == "mlflow":
+            if second_level_module in MLFLOW_MODULES_TO_PACKAGES:
+                self.imported_modules.add(second_level_module)
+                return
+
+        self.imported_modules.add(top_level_module)
+
+    def __enter__(self):
+        # Patch `builtins.__import__` and `importlib.import_module`
+        self.original_import = builtins.__import__
+        self.original_import_module = importlib.import_module
+        builtins.__import__ = self._wrap_import(self.original_import)
+        importlib.import_module = self._wrap_import_module(self.original_import_module)
+        return self
+
+    def __exit__(self, *_, **__):
+        # Revert the patches
+        builtins.__import__ = self.original_import
+        importlib.import_module = self.original_import_module
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model-path", required=True)
+    parser.add_argument("--flavor", required=True)
+    parser.add_argument("--output-file", required=True)
+    parser.add_argument("--sys-path", required=True)
+    parser.add_argument("--module-to-throw", required=False)
+    parser.add_argument("--error-file", required=False)
+    parser.add_argument("--record-full-module", default=False, action="store_true")
+    return parser.parse_args()
+
+
+def store_imported_modules(
+    cap_cm, model_path, flavor, output_file, error_file=None, record_full_module=False
+):
+    # If `model_path` refers to an MLflow model directory, load the model using
+    # `mlflow.pyfunc.load_model`
+    if os.path.isdir(model_path) and MLMODEL_FILE_NAME in os.listdir(model_path):
+        mlflow_model = Model.load(model_path)
+        pyfunc_conf = mlflow_model.flavors.get(mlflow.pyfunc.FLAVOR_NAME)
+        input_example = mlflow_model.load_input_example(model_path)
+        params = mlflow_model.load_input_example_params(model_path)
+
+        def load_model_and_predict(original_load_fn, *args, **kwargs):
+            model = original_load_fn(*args, **kwargs)
+            if input_example is not None:
+                try:
+                    model.predict(input_example, params=params)
+                except Exception as e:
+                    if error_file:
+                        stack_trace = get_stacktrace(e)
+                        write_to(
+                            error_file,
+                            "Failed to run predict on input_example, dependencies "
+                            "introduced in predict are not captured.\n" + stack_trace,
+                        )
+                    else:
+                        raise e
+            return model
+
+        if record_full_module:
+            # Note: if we want to record all imported modules
+            # (for inferring code_paths purpose),
+            # The `importlib.import_module(pyfunc_conf[MAIN])` invocation
+            # must be wrapped with `cap_cm` context manager,
+            # because `pyfunc_conf[MAIN]` might also be a module loaded from
+            # code_paths.
+            with cap_cm:
+                # `mlflow.pyfunc.load_model` internally invokes
+                # `importlib.import_module(pyfunc_conf[MAIN])`
+                mlflow.pyfunc.load_model(model_path)
+        else:
+            loader_module = importlib.import_module(pyfunc_conf[MAIN])
+            original = loader_module._load_pyfunc
+
+            @functools.wraps(original)
+            def _load_pyfunc_patch(*args, **kwargs):
+                with cap_cm:
+                    return load_model_and_predict(original, *args, **kwargs)
+
+            loader_module._load_pyfunc = _load_pyfunc_patch
+            try:
+                mlflow.pyfunc.load_model(model_path)
+            finally:
+                loader_module._load_pyfunc = original
+    # Otherwise, load the model using `mlflow.<flavor>._load_pyfunc`.
+    # For models that don't contain pyfunc flavor (e.g. scikit-learn estimator
+    # that doesn't implement a `predict` method),
+    # we need to directly pass a model data path to this script.
+    else:
+        with cap_cm:
+            importlib.import_module(f"mlflow.{flavor}")._load_pyfunc(model_path)
+
+    # Store the imported modules in `output_file`
+    write_to(output_file, "\n".join(cap_cm.imported_modules))
+
+
+def main():
+    args = parse_args()
+    model_path = args.model_path
+    flavor = args.flavor
+    output_file = args.output_file
+    error_file = args.error_file
+    # Mirror `sys.path` of the parent process
+    sys.path = json.loads(args.sys_path)
+
+    if flavor == mlflow.spark.FLAVOR_NAME:
+        # Create a local spark environment within the subprocess
+        from mlflow.utils._spark_utils import _create_local_spark_session_for_loading_spark_model
+
+        _prepare_subprocess_environ_for_creating_local_spark_session()
+        _create_local_spark_session_for_loading_spark_model()
+
+    cap_cm = _CaptureImportedModules(record_full_module=args.record_full_module)
+    store_imported_modules(
+        cap_cm,
+        model_path,
+        flavor,
+        output_file,
+        error_file,
+        record_full_module=args.record_full_module,
+    )
+
+    # Clean up a spark session created by `mlflow.spark._load_pyfunc`
+    if flavor == mlflow.spark.FLAVOR_NAME:
+        from mlflow.utils._spark_utils import _get_active_spark_session
+
+        spark = _get_active_spark_session()
+        if spark:
+            try:
+                spark.stop()
+            except Exception:
+                # Swallow unexpected exceptions
+                pass
+
+
+if __name__ == "__main__":
+    main()

mlflow/utils/_capture_transformers_modules.py

@@ -0,0 +1,75 @@
+"""
+This script should be executed in a fresh python interpreter process using `subprocess`.
+"""
+
+import json
+import os
+import sys
+
+import mlflow
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+from mlflow.utils._capture_modules import (
+    _CaptureImportedModules,
+    parse_args,
+    store_imported_modules,
+)
+
+
+class _CaptureImportedModulesForHF(_CaptureImportedModules):
+    """
+    A context manager to capture imported modules by temporarily applying a patch to
+    `builtins.__import__` and `importlib.import_module`.
+    Used for 'transformers' flavor only.
+    """
+
+    def __init__(self, module_to_throw, record_full_module=False):
+        super().__init__(record_full_module=record_full_module)
+        self.module_to_throw = module_to_throw
+
+    def _record_imported_module(self, full_module_name):
+        if full_module_name == self.module_to_throw or full_module_name.startswith(
+            f"{self.module_to_throw}."
+        ):
+            raise ImportError(f"Disabled package {full_module_name}")
+        return super()._record_imported_module(full_module_name)
+
+
+def main():
+    args = parse_args()
+    model_path = args.model_path
+    flavor = args.flavor
+    output_file = args.output_file
+    module_to_throw = args.module_to_throw
+    # Mirror `sys.path` of the parent process
+    sys.path = json.loads(args.sys_path)
+
+    if flavor != mlflow.transformers.FLAVOR_NAME:
+        raise MlflowException(
+            f"This script is only applicable to '{mlflow.transformers.FLAVOR_NAME}' flavor, "
+            "if you're applying other flavors, please use _capture_modules script.",
+        )
+
+    if module_to_throw == "":
+        raise MlflowException("Please specify the module to throw.")
+    elif module_to_throw == "tensorflow":
+        if os.environ.get("USE_TORCH", None) != "TRUE":
+            raise MlflowException(
+                "The environment variable USE_TORCH has to be set to TRUE to disable Tensorflow.",
+                error_code=INVALID_PARAMETER_VALUE,
+            )
+    elif module_to_throw == "torch":
+        if os.environ.get("USE_TF", None) != "TRUE":
+            raise MlflowException(
+                "The environment variable USE_TF has to be set to TRUE to disable Pytorch.",
+                error_code=INVALID_PARAMETER_VALUE,
+            )
+
+    cap_cm = _CaptureImportedModulesForHF(
+        module_to_throw, record_full_module=args.record_full_module
+    )
+    store_imported_modules(cap_cm, model_path, flavor, output_file)
+
+
+if __name__ == "__main__":
+    main()