genesis-flow 1.0.0__py3-none-any.whl

mlflow/utils/import_hooks/__init__.py

@@ -0,0 +1,363 @@
+"""
+NOTE: The contents of this file have been inlined from the wrapt package's source code
+https://github.com/GrahamDumpleton/wrapt/blob/1.12.1/src/wrapt/importer.py.
+Some modifications, have been made in order to:
+    - avoid duplicate registration of import hooks
+    - inline functions from dependent wrapt submodules rather than importing them.
+
+This module implements a post import hook mechanism styled after what is described in PEP-369.
+Note that it doesn't cope with modules being reloaded.
+It also extends the functionality to support custom hooks for import errors
+(as opposed to only successful imports).
+"""
+
+import importlib  # noqa: F401
+import sys
+import threading
+
+string_types = (str,)
+
+
+# from .decorators import synchronized
+# NOTE: Instead of using this import (from wrapt's decorator module, see
+# https://github.com/GrahamDumpleton/wrapt/blob/68316bea668fd905a4acb21f37f12596d8c30d80/src/wrapt/decorators.py#L430-L456),
+# we define a decorator with similar behavior that acquires a lock while calling the decorated
+# function
+def synchronized(lock):
+    def decorator(f):
+        # See e.g. https://www.python.org/dev/peps/pep-0318/#examples
+        def new_fn(*args, **kwargs):
+            with lock:
+                return f(*args, **kwargs)
+
+        return new_fn
+
+    return decorator
+
+
+# The dictionary registering any post import hooks to be triggered once
+# the target module has been imported. Once a module has been imported
+# and the hooks fired, the list of hooks recorded against the target
+# module will be truncated but the list left in the dictionary. This
+# acts as a flag to indicate that the module had already been imported.
+
+_post_import_hooks = {}
+_post_import_hooks_lock = threading.RLock()
+
+# A dictionary for any import hook error handlers to be triggered when the
+# target module import fails.
+
+_import_error_hooks = {}
+_import_error_hooks_lock = threading.RLock()
+
+_import_hook_finder_init = False
+
+# Register a new post import hook for the target module name. This
+# differs from the PEP-369 implementation in that it also allows the
+# hook function to be specified as a string consisting of the name of
+# the callback in the form 'module:function'. This will result in a
+# proxy callback being registered which will defer loading of the
+# specified module containing the callback function until required.
+
+
+def _create_import_hook_from_string(name):
+    def import_hook(module):
+        module_name, function = name.split(":")
+        attrs = function.split(".")
+        __import__(module_name)
+        callback = sys.modules[module_name]
+        for attr in attrs:
+            callback = getattr(callback, attr)
+        return callback(module)
+
+    return import_hook
+
+
+def register_generic_import_hook(hook, name, hook_dict, overwrite):
+    # Create a deferred import hook if hook is a string name rather than
+    # a callable function.
+
+    if isinstance(hook, string_types):
+        hook = _create_import_hook_from_string(hook)
+
+    # Automatically install the import hook finder if it has not already
+    # been installed.
+
+    global _import_hook_finder_init
+    if not _import_hook_finder_init:
+        _import_hook_finder_init = True
+        sys.meta_path.insert(0, ImportHookFinder())
+
+    # Determine if any prior registration of an import hook for
+    # the target modules has occurred and act appropriately.
+
+    hooks = hook_dict.get(name, None)
+
+    if hooks is None:
+        # No prior registration of import hooks for the target
+        # module. We need to check whether the module has already been
+        # imported. If it has we fire the hook immediately and add an
+        # empty list to the registry to indicate that the module has
+        # already been imported and hooks have fired. Otherwise add
+        # the post import hook to the registry.
+
+        module = sys.modules.get(name, None)
+
+        if module is not None:
+            hook_dict[name] = []
+            hook(module)
+
+        else:
+            hook_dict[name] = [hook]
+
+    elif hooks == []:
+        # A prior registration of import hooks for the target
+        # module was done and the hooks already fired. Fire the hook
+        # immediately.
+
+        module = sys.modules[name]
+        hook(module)
+
+    else:
+        # A prior registration of import hooks for the target
+        # module was done but the module has not yet been imported.
+
+        def hooks_equal(existing_hook, hook):
+            if hasattr(existing_hook, "__name__") and hasattr(hook, "__name__"):
+                return existing_hook.__name__ == hook.__name__
+            else:
+                return False
+
+        if overwrite:
+            hook_dict[name] = [
+                existing_hook
+                for existing_hook in hook_dict[name]
+                if not hooks_equal(existing_hook, hook)
+            ]
+
+        hook_dict[name].append(hook)
+
+
+@synchronized(_import_error_hooks_lock)
+def register_import_error_hook(hook, name, overwrite=True):
+    """
+    Args:
+        hook: A function or string entrypoint to invoke when the specified module is imported
+            and an error occurs.
+        name: The name of the module for which to fire the hook at import error detection time.
+        overwrite: Specifies the desired behavior when a preexisting hook for the same
+            function / entrypoint already exists for the specified module. If `True`,
+            all preexisting hooks matching the specified function / entrypoint will be
+            removed and replaced with a single instance of the specified `hook`.
+    """
+    register_generic_import_hook(hook, name, _import_error_hooks, overwrite)
+
+
+@synchronized(_post_import_hooks_lock)
+def register_post_import_hook(hook, name, overwrite=True):
+    """
+    Args:
+        hook: A function or string entrypoint to invoke when the specified module is imported.
+        name: The name of the module for which to fire the hook at import time.
+        overwrite: Specifies the desired behavior when a preexisting hook for the same
+            function / entrypoint already exists for the specified module. If `True`,
+            all preexisting hooks matching the specified function / entrypoint will be
+            removed and replaced with a single instance of the specified `hook`.
+    """
+    register_generic_import_hook(hook, name, _post_import_hooks, overwrite)
+
+
+@synchronized(_post_import_hooks_lock)
+def get_post_import_hooks(name):
+    return _post_import_hooks.get(name)
+
+
+# Register post import hooks defined as package entry points.
+
+
+def _create_import_hook_from_entrypoint(entrypoint):
+    def import_hook(module):
+        __import__(entrypoint.module_name)
+        callback = sys.modules[entrypoint.module_name]
+        for attr in entrypoint.attrs:
+            callback = getattr(callback, attr)
+        return callback(module)
+
+    return import_hook
+
+
+def discover_post_import_hooks(group):
+    # New in 3.9: https://docs.python.org/3/library/importlib.resources.html#importlib.resources.files
+    if sys.version_info.major > 2 and sys.version_info.minor > 8:
+        from importlib.resources import files  # clint: disable=lazy-builtin-import
+
+        for entrypoint in (
+            resource.name for resource in files(group).iterdir() if resource.is_file()
+        ):
+            callback = _create_import_hook_from_entrypoint(entrypoint)
+            register_post_import_hook(callback, entrypoint.name)
+    else:
+        from importlib.resources import contents  # clint: disable=lazy-builtin-import
+
+        for entrypoint in contents(group):
+            callback = _create_import_hook_from_entrypoint(entrypoint)
+            register_post_import_hook(callback, entrypoint.name)
+
+
+# Indicate that a module has been loaded. Any post import hooks which
+# were registered against the target module will be invoked. If an
+# exception is raised in any of the post import hooks, that will cause
+# the import of the target module to fail.
+
+
+@synchronized(_post_import_hooks_lock)
+def notify_module_loaded(module):
+    name = getattr(module, "__name__", None)
+    hooks = _post_import_hooks.get(name)
+
+    if hooks:
+        _post_import_hooks[name] = []
+
+        for hook in hooks:
+            hook(module)
+
+
+@synchronized(_import_error_hooks_lock)
+def notify_module_import_error(module_name):
+    hooks = _import_error_hooks.get(module_name)
+
+    if hooks:
+        # Error hooks differ from post import hooks, in that we don't clear the
+        # hook as soon as it fires.
+        for hook in hooks:
+            hook(module_name)
+
+
+# A custom module import finder. This intercepts attempts to import
+# modules and watches out for attempts to import target modules of
+# interest. When a module of interest is imported, then any post import
+# hooks which are registered will be invoked.
+
+
+class _ImportHookChainedLoader:
+    def __init__(self, loader):
+        self.loader = loader
+
+    def load_module(self, fullname):
+        try:
+            module = self.loader.load_module(fullname)
+            notify_module_loaded(module)
+        except (ImportError, AttributeError):
+            notify_module_import_error(fullname)
+            raise
+
+        return module
+
+
+class ImportHookFinder:
+    def __init__(self):
+        self.in_progress = {}
+
+    @synchronized(_post_import_hooks_lock)
+    @synchronized(_import_error_hooks_lock)
+    def find_module(self, fullname, path=None):
+        # If the module being imported is not one we have registered
+        # import hooks for, we can return immediately. We will
+        # take no further part in the importing of this module.
+
+        if fullname not in _post_import_hooks and fullname not in _import_error_hooks:
+            return None
+
+        # When we are interested in a specific module, we will call back
+        # into the import system a second time to defer to the import
+        # finder that is supposed to handle the importing of the module.
+        # We set an in progress flag for the target module so that on
+        # the second time through we don't trigger another call back
+        # into the import system and cause a infinite loop.
+
+        if fullname in self.in_progress:
+            return None
+
+        self.in_progress[fullname] = True
+
+        # Now call back into the import system again.
+
+        try:
+            # For Python 3 we need to use find_spec().loader
+            # from the importlib.util module. It doesn't actually
+            # import the target module and only finds the
+            # loader. If a loader is found, we need to return
+            # our own loader which will then in turn call the
+            # real loader to import the module and invoke the
+            # post import hooks.
+            try:
+                import importlib.util  # clint: disable=lazy-builtin-import
+
+                loader = importlib.util.find_spec(fullname).loader
+            # If an ImportError (or AttributeError) is encountered while finding the module,
+            # notify the hooks for import errors
+            except (ImportError, AttributeError):
+                notify_module_import_error(fullname)
+                loader = importlib.find_loader(fullname, path)
+            if loader:
+                return _ImportHookChainedLoader(loader)
+        finally:
+            del self.in_progress[fullname]
+
+    @synchronized(_post_import_hooks_lock)
+    @synchronized(_import_error_hooks_lock)
+    def find_spec(self, fullname, path, target=None):
+        # If the module being imported is not one we have registered
+        # import hooks for, we can return immediately. We will
+        # take no further part in the importing of this module.
+
+        if fullname not in _post_import_hooks and fullname not in _import_error_hooks:
+            return None
+
+        # When we are interested in a specific module, we will call back
+        # into the import system a second time to defer to the import
+        # finder that is supposed to handle the importing of the module.
+        # We set an in progress flag for the target module so that on
+        # the second time through we don't trigger another call back
+        # into the import system and cause a infinite loop.
+
+        if fullname in self.in_progress:
+            return None
+
+        self.in_progress[fullname] = True
+
+        # Now call back into the import system again.
+
+        try:
+            import importlib.util  # clint: disable=lazy-builtin-import
+
+            spec = importlib.util.find_spec(fullname)
+            # Replace the module spec's loader with a wrapped version that executes import
+            # hooks when the module is loaded
+            spec.loader = _ImportHookChainedLoader(spec.loader)
+            return spec
+        except (ImportError, AttributeError):
+            notify_module_import_error(fullname)
+        finally:
+            del self.in_progress[fullname]
+
+
+# Decorator for marking that a function should be called as a post
+# import hook when the target module is imported.
+# If error_handler is True, then apply the marked function as an import hook
+# for import errors (instead of successful imports).
+# It is assumed that all error hooks are added during driver start-up,
+# and thus added prior to any import calls. If an error hook is added
+# after a module has already failed the import, there's no guarantee
+# that the hook will fire.
+
+
+def when_imported(name, error_handler=False):
+    def register(hook):
+        if error_handler:
+            register_import_error_hook(hook, name)
+        else:
+            register_post_import_hook(hook, name)
+        return hook
+
+    return register

mlflow/utils/lazy_load.py

@@ -0,0 +1,51 @@
+"""Utility to lazy load modules."""
+
+import importlib
+import sys
+import types
+
+
+class LazyLoader(types.ModuleType):
+    """Class for module lazy loading.
+
+    This class helps lazily load modules at package level, which avoids pulling in large
+    dependencies like `tensorflow` or `torch`. This class is mirrored from wandb's LazyLoader:
+    https://github.com/wandb/wandb/blob/79b2d4b73e3a9e4488e503c3131ff74d151df689/wandb/sdk/lib/lazyloader.py#L9
+    """
+
+    def __init__(self, local_name, parent_module_globals, name):
+        self._local_name = local_name
+        self._parent_module_globals = parent_module_globals
+
+        self._module = None
+        super().__init__(str(name))
+
+    def _load(self):
+        """Load the module and insert it into the parent's globals."""
+        if self._module:
+            # If already loaded, return the loaded module.
+            return self._module
+
+        # Import the target module and insert it into the parent's namespace
+        module = importlib.import_module(self.__name__)
+        self._parent_module_globals[self._local_name] = module
+        sys.modules[self._local_name] = module
+
+        # Update this object's dict so that if someone keeps a reference to the `LazyLoader`,
+        # lookups are efficient (`__getattr__` is only called on lookups that fail).
+        self.__dict__.update(module.__dict__)
+
+        return module
+
+    def __getattr__(self, item):
+        module = self._load()
+        return getattr(module, item)
+
+    def __dir__(self):
+        module = self._load()
+        return dir(module)
+
+    def __repr__(self):
+        if not self._module:
+            return f"<module '{self.__name__} (Not loaded yet)'>"
+        return repr(self._module)

mlflow/utils/logging_utils.py

@@ -0,0 +1,168 @@
+import contextlib
+import logging
+import logging.config
+import re
+import sys
+
+from mlflow.environment_variables import MLFLOW_LOGGING_LEVEL
+
+# Logging format example:
+# 2018/11/20 12:36:37 INFO mlflow.sagemaker: Creating new SageMaker endpoint
+LOGGING_LINE_FORMAT = "%(asctime)s %(levelname)s %(name)s: %(message)s"
+LOGGING_DATETIME_FORMAT = "%Y/%m/%d %H:%M:%S"
+
+
+class MlflowLoggingStream:
+    """
+    A Python stream for use with event logging APIs throughout MLflow (`eprint()`,
+    `logger.info()`, etc.). This stream wraps `sys.stderr`, forwarding `write()` and
+    `flush()` calls to the stream referred to by `sys.stderr` at the time of the call.
+    It also provides capabilities for disabling the stream to silence event logs.
+    """
+
+    def __init__(self):
+        self._enabled = True
+
+    def write(self, text):
+        if self._enabled:
+            sys.stderr.write(text)
+
+    def flush(self):
+        if self._enabled:
+            sys.stderr.flush()
+
+    @property
+    def enabled(self):
+        return self._enabled
+
+    @enabled.setter
+    def enabled(self, value):
+        self._enabled = value
+
+
+MLFLOW_LOGGING_STREAM = MlflowLoggingStream()
+
+
+def disable_logging():
+    """
+    Disables the `MlflowLoggingStream` used by event logging APIs throughout MLflow
+    (`eprint()`, `logger.info()`, etc), silencing all subsequent event logs.
+    """
+    MLFLOW_LOGGING_STREAM.enabled = False
+
+
+def enable_logging():
+    """
+    Enables the `MlflowLoggingStream` used by event logging APIs throughout MLflow
+    (`eprint()`, `logger.info()`, etc), emitting all subsequent event logs. This
+    reverses the effects of `disable_logging()`.
+    """
+    MLFLOW_LOGGING_STREAM.enabled = True
+
+
+class MlflowFormatter(logging.Formatter):
+    """
+    Custom Formatter Class to support colored log
+    ANSI characters might not work natively on older Windows, so disabling the feature for win32.
+    See https://github.com/borntyping/python-colorlog/blob/dfa10f59186d3d716aec4165ee79e58f2265c0eb/colorlog/escape_codes.py#L16C8-L16C31
+    """
+
+    # Copied from color log package https://github.com/borntyping/python-colorlog/blob/dfa10f59186d3d716aec4165ee79e58f2265c0eb/colorlog/escape_codes.py#L33-L50
+    COLORS = {
+        "black": 30,
+        "red": 31,
+        "green": 32,
+        "yellow": 33,
+        "blue": 34,
+        "purple": 35,
+        "cyan": 36,
+        "white": 37,
+        "light_black": 90,
+        "light_red": 91,
+        "light_green": 92,
+        "light_yellow": 93,
+        "light_blue": 94,
+        "light_purple": 95,
+        "light_cyan": 96,
+        "light_white": 97,
+    }
+    RESET = "\033[0m"
+
+    def format(self, record):
+        if color := getattr(record, "color", None):
+            if color in self.COLORS and sys.platform != "win32":
+                color_code = self._escape(self.COLORS[color])
+                return f"{color_code}{super().format(record)}{self.RESET}"
+        return super().format(record)
+
+    def _escape(self, code: int) -> str:
+        return f"\033[{code}m"
+
+
+def _configure_mlflow_loggers(root_module_name):
+    logging.config.dictConfig(
+        {
+            "version": 1,
+            "disable_existing_loggers": False,
+            "formatters": {
+                "mlflow_formatter": {
+                    "()": MlflowFormatter,
+                    "format": LOGGING_LINE_FORMAT,
+                    "datefmt": LOGGING_DATETIME_FORMAT,
+                },
+            },
+            "handlers": {
+                "mlflow_handler": {
+                    "formatter": "mlflow_formatter",
+                    "class": "logging.StreamHandler",
+                    "stream": MLFLOW_LOGGING_STREAM,
+                },
+            },
+            "loggers": {
+                root_module_name: {
+                    "handlers": ["mlflow_handler"],
+                    "level": (MLFLOW_LOGGING_LEVEL.get() or "INFO").upper(),
+                    "propagate": False,
+                },
+            },
+        }
+    )
+
+
+def eprint(*args, **kwargs):
+    print(*args, file=MLFLOW_LOGGING_STREAM, **kwargs)
+
+
+class LoggerMessageFilter(logging.Filter):
+    def __init__(self, module: str, filter_regex: re.Pattern):
+        super().__init__()
+        self._pattern = filter_regex
+        self._module = module
+
+    def filter(self, record):
+        if record.name == self._module and self._pattern.search(record.msg):
+            return False
+        return True
+
+
+@contextlib.contextmanager
+def suppress_logs(module: str, filter_regex: re.Pattern):
+    """
+    Context manager that suppresses log messages from the specified module that match the specified
+    regular expression. This is useful for suppressing expected log messages from third-party
+    libraries that are not relevant to the current test.
+    """
+    logger = logging.getLogger(module)
+    filter = LoggerMessageFilter(module=module, filter_regex=filter_regex)
+    logger.addFilter(filter)
+    try:
+        yield
+    finally:
+        logger.removeFilter(filter)
+
+
+def _debug(s: str) -> None:
+    """
+    Debug function to test logging level.
+    """
+    logging.getLogger(__name__).debug(s)

mlflow/utils/mime_type_utils.py

@@ -0,0 +1,58 @@
+import os
+import pathlib
+from mimetypes import guess_type
+
+from mlflow.version import IS_TRACING_SDK_ONLY
+
+
+# TODO: Create a module to define constants to avoid circular imports
+#  and move MLMODEL_FILE_NAME and MLPROJECT_FILE_NAME in the module.
+def get_text_extensions():
+    exts = [
+        "txt",
+        "log",
+        "err",
+        "cfg",
+        "conf",
+        "cnf",
+        "cf",
+        "ini",
+        "properties",
+        "prop",
+        "hocon",
+        "toml",
+        "yaml",
+        "yml",
+        "xml",
+        "json",
+        "js",
+        "py",
+        "py3",
+        "csv",
+        "tsv",
+        "md",
+        "rst",
+    ]
+
+    if not IS_TRACING_SDK_ONLY:
+        from mlflow.models.model import MLMODEL_FILE_NAME
+        from mlflow.projects._project_spec import MLPROJECT_FILE_NAME
+
+        exts.extend([MLMODEL_FILE_NAME, MLPROJECT_FILE_NAME])
+
+    return exts
+
+
+def _guess_mime_type(file_path):
+    filename = pathlib.Path(file_path).name
+    extension = os.path.splitext(filename)[-1].replace(".", "")
+    # for MLmodel/mlproject with no extensions
+    if extension == "":
+        extension = filename
+    if extension in get_text_extensions():
+        return "text/plain"
+    mime_type, _ = guess_type(filename)
+    if not mime_type:
+        # As a fallback, if mime type is not detected, treat it as a binary file
+        return "application/octet-stream"
+    return mime_type