mlrun 1.3.3__py3-none-any.whl → 1.4.0__py3-none-any.whl

mlrun/package/utils/_pickler.py

@@ -0,0 +1,234 @@
+# Copyright 2023 Iguazio
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+import importlib
+import os
+import sys
+import tempfile
+import warnings
+from types import ModuleType
+from typing import Any, Dict, Tuple, Union
+
+from mlrun.errors import MLRunInvalidArgumentError
+from mlrun.utils import logger
+
+
+class Pickler:
+    """
+    A static class to pickle objects with multiple modules while capturing the environment of the pickled object. The
+    pickler will raise warnings in case the object is un-pickled in a mismatching environment (different modules
+    and / or python versions)
+    """
+
+    @staticmethod
+    def pickle(
+        obj: Any, pickle_module_name: str, output_path: str = None
+    ) -> Tuple[str, Dict[str, Union[str, None]]]:
+        """
+        Pickle an object using the given module. The pickled object will be saved to file to the given output path.
+
+        :param obj:                The object to pickle.
+        :param pickle_module_name: The pickle module to use. For example: "pickle", "joblib", "cloudpickle".
+        :param output_path:        The output path to save the 'pkl' file to. If not provided, the pickle will be saved
+                                   to a temporary directory. The user is responsible to clean the temporary directory.
+
+        :return: A tuple of the path of the 'pkl' file and the instructions the pickler noted.
+        """
+        # Get the pickle module:
+        pickle_module = importlib.import_module(pickle_module_name)
+        Pickler._validate_pickle_module(pickle_module=pickle_module)
+        pickle_module_version = Pickler._get_module_version(
+            module_name=pickle_module_name
+        )
+
+        # Get the object's module (module name can be extracted usually from the object's class):
+        object_module_name = (
+            obj.__module__.split(".")[0]
+            if hasattr(obj, "__module__")
+            else type(obj).__module__.split(".")[0]
+        )
+        object_module_version = Pickler._get_module_version(
+            module_name=object_module_name
+        )
+
+        # Get the python version:
+        python_version = Pickler._get_python_version()
+
+        # Construct the pickler labels dictionary (versions may not be available):
+        instructions = {
+            "object_module_name": object_module_name,
+            "pickle_module_name": pickle_module_name,
+            "python_version": python_version,
+        }
+        if object_module_version is not None:
+            instructions["object_module_version"] = object_module_version
+        if pickle_module_version is not None:
+            instructions["pickle_module_version"] = pickle_module_version
+
+        # Generate a temporary output path if not provided:
+        if output_path is None:
+            output_path = os.path.join(tempfile.mkdtemp(), "obj.pkl")
+
+        # Pickle the object to file:
+        with open(output_path, "wb") as pkl_file:
+            pickle_module.dump(obj, pkl_file)
+
+        return output_path, instructions
+
+    @staticmethod
+    def unpickle(
+        pickle_path: str,
+        pickle_module_name: str,
+        object_module_name: str = None,
+        python_version: str = None,
+        pickle_module_version: str = None,
+        object_module_version: str = None,
+    ) -> Any:
+        """
+        Unpickle an object using the given instructions. Warnings may be raised in case any of the versions are
+        mismatching (only if provided - not None).
+
+        :param pickle_path:           Path to the 'pkl' file to un-pickle.
+        :param pickle_module_name:    Module to use for unpickling the object.
+        :param object_module_name:    The original object's module. Used to verify the current interpreter object module
+                                      version match the pickled object version before unpickling the object.
+        :param python_version:        The python version in which the original object was pickled. Used to verify the
+                                      current interpreter python version match the pickled object version before
+                                      unpickling the object.
+        :param pickle_module_version: The pickle module version. Used to verify the current interpreter module version
+                                      match the one who pickled the object before unpickling it.
+        :param object_module_version: The original object's module version to match to the interpreter's module version.
+
+        :return: The un-pickled object.
+        """
+        # Check the python version against the pickled object:
+        if python_version is not None:
+            current_python_version = Pickler._get_python_version()
+            if python_version != current_python_version:
+                logger.warn(
+                    f"MLRun is trying to load an object that was pickled on python version "
+                    f"'{python_version}' but the current python version is '{current_python_version}'. "
+                    f"When using pickle, it is recommended to save and load an object on the same python version to "
+                    f"reduce unexpected errors."
+                )
+
+        # Get the pickle module:
+        pickle_module = importlib.import_module(pickle_module_name)
+        Pickler._validate_pickle_module(pickle_module=pickle_module)
+
+        # Check the pickle module against the pickled object (only if the version is given):
+        if pickle_module_version is not None:
+            current_pickle_module_version = Pickler._get_module_version(
+                module_name=pickle_module_name
+            )
+            if pickle_module_version != current_pickle_module_version:
+                logger.warn(
+                    f"MLRun is trying to load an object that was pickled using "
+                    f"{pickle_module_name} version {pickle_module_version} but the current module version is "
+                    f"'{current_pickle_module_version}'. "
+                    f"When using pickle, it is recommended to save and load an "
+                    f"object using the same pickling module version to reduce unexpected errors."
+                )
+
+        # Check the object module against the pickled object (only if the version is given):
+        if object_module_version is not None and object_module_name is not None:
+            current_object_module_version = Pickler._get_module_version(
+                module_name=object_module_name
+            )
+            if object_module_version != current_object_module_version:
+                logger.warn(
+                    f"MLRun is trying to load an object from module {object_module_name} version "
+                    f"{object_module_version} but the current module version is '{current_object_module_version}'. "
+                    f"When using pickle, it is recommended to save and load an object using "
+                    f"the same exact module version to reduce unexpected errors."
+                )
+
+        # Load the object from the pickle file:
+        with open(pickle_path, "rb") as pickle_file:
+            obj = pickle_module.load(pickle_file)
+
+        return obj
+
+    @staticmethod
+    def _validate_pickle_module(pickle_module: ModuleType):
+        """
+        Validate the pickle module to use have a `dump` and `load` functions so the Pickler can use it.
+
+        :param pickle_module: The pickle module tot validate.
+
+        :raise MLRunInvalidArgumentError: If the pickle module is not valid.
+        """
+        for function_name in ["dump", "load"]:
+            if not hasattr(pickle_module, function_name):
+                raise MLRunInvalidArgumentError(
+                    f"A pickle module is expected to have a `{function_name}` function but the provided module "
+                    f"{pickle_module.__name__} does not have it."
+                )
+
+    @staticmethod
+    def _get_module_version(module_name: str) -> Union[str, None]:
+        """
+        Get a module's version. Most updated modules have versions but some don't. In case the version could not be
+        read, None is returned.
+
+        :param module_name: The module's name to get its version.
+
+        :return: The module's version if found and None otherwise.
+        """
+        # First we'll try to get the module version from `importlib`:
+        try:
+            # Since Python 3.8, `version` is part of `importlib.metadata`. Before 3.8, we'll use the module
+            # `importlib_metadata` to get `version`.
+            if (
+                sys.version_info[1] > 7
+            ):  # TODO: Remove once Python 3.7 is not supported.
+                from importlib.metadata import version
+            else:
+                from importlib_metadata import version
+
+            return version(module_name)
+        except (ModuleNotFoundError, importlib.metadata.PackageNotFoundError):
+            # User won't necessarily have the `importlib_metadata` module, so we will ignore it by catching
+            # `ModuleNotFoundError`. `PackageNotFoundError` is ignored as well as this is raised when `version` could
+            # not find the package related to the module.
+            pass
+
+        # Secondly, if importlib could not get the version (most likely 'importlib_metadata' is not installed), we'll
+        # try to use `pkg_resources` to get the version (the version will be found only if the package name is equal to
+        # the module name. For example, if the module name is 'x' then the way we installed the package must be
+        # 'pip install x'):
+        import pkg_resources
+
+        with warnings.catch_warnings():
+            # If a module's package is not found, a `PkgResourcesDeprecationWarning` warning will be raised and then
+            # `DistributionNotFound` exception will be raised, so we ignore them both:
+            warnings.filterwarnings(
+                "ignore", category=pkg_resources.PkgResourcesDeprecationWarning
+            )
+            try:
+                return pkg_resources.get_distribution(module_name).version
+            except pkg_resources.DistributionNotFound:
+                pass
+
+        # The version could not be found.
+        return None
+
+    @staticmethod
+    def _get_python_version() -> str:
+        """
+        Get the current running python's version.
+
+        :return: The python version string.
+        """
+        return f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}"

mlrun/package/utils/_supported_format.py

@@ -0,0 +1,71 @@
+# Copyright 2023 Iguazio
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+from abc import ABC
+from typing import Dict, Generic, List, Type, TypeVar, Union
+
+# A generic type for a supported format handler class type:
+FileHandlerType = TypeVar("FileHandlerType")
+
+
+class SupportedFormat(ABC, Generic[FileHandlerType]):
+    """
+    Library of supported formats by some builtin MLRun packagers.
+    """
+
+    # Add here the all the supported formats in ALL CAPS and their value as a string:
+    ...
+
+    # The map to use in the method `get_format_handler`. A dictionary of string key to a class type to handle that
+    # format. New supported formats and handlers should be added to it:
+    _FORMAT_HANDLERS_MAP: Dict[str, Type[FileHandlerType]] = {}
+
+    @classmethod
+    def get_all_formats(cls) -> List[str]:
+        """
+        Get all supported formats.
+
+        :return: A list of all the supported formats.
+        """
+        return [
+            value
+            for key, value in cls.__dict__.items()
+            if isinstance(value, str) and not key.startswith("_")
+        ]
+
+    @classmethod
+    def get_format_handler(cls, fmt: str) -> Type[FileHandlerType]:
+        """
+        Get the format handler to the provided format (file extension):
+
+        :param fmt: The file extension to get the corresponding handler.
+
+        :return: The handler class.
+        """
+        return cls._FORMAT_HANDLERS_MAP[fmt]
+
+    @classmethod
+    def match_format(cls, path: str) -> Union[str, None]:
+        """
+        Try to match one of the available formats this class holds to a given path.
+
+        :param path: The path to match the format to.
+
+        :return: The matched format if found and None otherwise.
+        """
+        formats = cls.get_all_formats()
+        for fmt in formats:
+            if path.endswith(f".{fmt}"):
+                return fmt
+        return None

mlrun/package/utils/log_hint_utils.py

@@ -0,0 +1,93 @@
+# Copyright 2023 Iguazio
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+import typing
+
+from mlrun.errors import MLRunInvalidArgumentError
+
+
+class LogHintKey:
+    """
+    Known keys for a log hint to have.
+    """
+
+    KEY = "key"
+    ARTIFACT_TYPE = "artifact_type"
+    EXTRA_DATA = "extra_data"
+    METRICS = "metrics"
+
+
+class LogHintUtils:
+    """
+    Static class for utilities functions to process log hints.
+    """
+
+    @staticmethod
+    def parse_log_hint(
+        log_hint: typing.Union[typing.Dict[str, str], str, None]
+    ) -> typing.Union[typing.Dict[str, str], None]:
+        """
+        Parse a given log hint from string to a logging configuration dictionary. The string will be read as the
+        artifact key ('key' in the dictionary) and if the string have a single colon, the following structure is
+        assumed: "<artifact_key> : <artifact_type>".
+
+        If a logging configuration dictionary is received, it will be validated to have a key field.
+
+        None will be returned as None.
+
+        :param log_hint: The log hint to parse.
+
+        :return: The hinted logging configuration.
+
+        :raise MLRunInvalidArgumentError: In case the log hint is not following the string structure or the dictionary
+                                          is missing the key field.
+        """
+        # Check for None value:
+        if log_hint is None:
+            return None
+
+        # If the log hint was provided as a string, construct a dictionary out of it:
+        if isinstance(log_hint, str):
+            # Check if only key is given:
+            if ":" not in log_hint:
+                log_hint = {LogHintKey.KEY: log_hint}
+            # Check for valid "<key> : <artifact type>" pattern:
+            else:
+                if log_hint.count(":") > 1:
+                    raise MLRunInvalidArgumentError(
+                        f"Incorrect log hint pattern. Log hints can have only a single ':' in them to specify the "
+                        f"desired artifact type the returned value will be logged as: "
+                        f"'<artifact_key> : <artifact_type>', but given: {log_hint}"
+                    )
+                # Split into key and type:
+                key, artifact_type = log_hint.replace(" ", "").split(":")
+                if artifact_type == "":
+                    raise MLRunInvalidArgumentError(
+                        f"Incorrect log hint pattern. The ':' in a log hint should specify the desired artifact type "
+                        f"the returned value will be logged as in the following pattern: "
+                        f"'<artifact_key> : <artifact_type>', but no artifact type was given: {log_hint}"
+                    )
+                log_hint = {
+                    LogHintKey.KEY: key,
+                    LogHintKey.ARTIFACT_TYPE: artifact_type,
+                }
+
+        # Validate the log hint dictionary has the mandatory key:
+        if LogHintKey.KEY not in log_hint:
+            raise MLRunInvalidArgumentError(
+                f"A log hint dictionary must include the 'key' - the artifact key (it's name). The following log hint "
+                f"is missing the key: {log_hint}."
+            )
+
+        return log_hint