genesis-flow 1.0.0__py3-none-any.whl

mlflow/tracking/metric_value_conversion_utils.py

@@ -0,0 +1,93 @@
+import sys
+
+from mlflow.exceptions import INVALID_PARAMETER_VALUE, MlflowException
+
+
+def _is_module_imported(module_name: str) -> bool:
+    return module_name in sys.modules
+
+
+def _try_get_item(x):
+    try:
+        return x.item()
+    except Exception as e:
+        raise MlflowException(
+            f"Failed to convert metric value to float: {e}",
+            error_code=INVALID_PARAMETER_VALUE,
+        )
+
+
+def _converter_requires(module_name: str):
+    """Wrapper function that checks if specified `module_name` is already imported before
+    invoking wrapped function.
+    """
+
+    def decorator(func):
+        def wrapper(x):
+            if not _is_module_imported(module_name):
+                return x
+
+            return func(x)
+
+        return wrapper
+
+    return decorator
+
+
+def convert_metric_value_to_float_if_possible(x) -> float:
+    if x is None or type(x) == float:
+        return x
+
+    converter_fns_to_try = [
+        convert_metric_value_to_float_if_ndarray,
+        convert_metric_value_to_float_if_tensorflow_tensor,
+        convert_metric_value_to_float_if_torch_tensor,
+    ]
+
+    for converter_fn in converter_fns_to_try:
+        possible_float = converter_fn(x)
+        if type(possible_float) == float:
+            return possible_float
+
+    try:
+        return float(x)
+    except ValueError:
+        return x  # let backend handle conversion if possible
+
+
+@_converter_requires("numpy")
+def convert_metric_value_to_float_if_ndarray(x):
+    import numpy as np
+
+    if isinstance(x, np.ndarray):
+        return float(_try_get_item(x))
+
+    return x
+
+
+@_converter_requires("torch")
+def convert_metric_value_to_float_if_torch_tensor(x):
+    import torch
+
+    if isinstance(x, torch.Tensor):
+        extracted_tensor_val = x.detach().cpu()
+
+        return float(_try_get_item(extracted_tensor_val))
+
+    return x
+
+
+@_converter_requires("tensorflow")
+def convert_metric_value_to_float_if_tensorflow_tensor(x):
+    import tensorflow as tf
+
+    if isinstance(x, tf.Tensor):
+        try:
+            return float(x)
+        except Exception as e:
+            raise MlflowException(
+                f"Failed to convert metric value to float: {e!r}",
+                error_code=INVALID_PARAMETER_VALUE,
+            )
+
+    return x

mlflow/tracking/multimedia.py

@@ -0,0 +1,206 @@
+"""
+Internal module implementing multi-media objects and utilities in MLflow. Multi-media objects are
+exposed to users at the top-level :py:mod:`mlflow` module.
+"""
+
+import warnings
+from typing import TYPE_CHECKING, Any, Optional, Union
+
+if TYPE_CHECKING:
+    import numpy
+    import PIL
+
+
+COMPRESSED_IMAGE_SIZE = 256
+
+
+def compress_image_size(
+    image: "PIL.Image.Image", max_size: Optional[int] = COMPRESSED_IMAGE_SIZE
+) -> "PIL.Image.Image":
+    """
+    Scale the image to fit within a square with length `max_size` while maintaining
+    the aspect ratio.
+    """
+    # scale the image to max(width, height) <= compressed_file_max_size
+    width, height = image.size
+    if width > height:
+        new_width = max_size
+        new_height = int(height * (new_width / width))
+    else:
+        new_height = max_size
+        new_width = int(width * (new_height / height))
+    return image.resize((new_width, new_height))
+
+
+def convert_to_pil_image(image: Union["numpy.ndarray", list[Any]]) -> "PIL.Image.Image":
+    """
+    Convert a numpy array to a PIL image.
+    """
+    import numpy as np
+
+    try:
+        from PIL import Image
+    except ImportError as exc:
+        raise ImportError(
+            "Pillow is required to serialize a numpy array as an image. "
+            "Please install it via: `pip install Pillow`"
+        ) from exc
+
+    def _normalize_to_uint8(x):
+        is_int = np.issubdtype(x.dtype, np.integer)
+        low = 0
+        high = 255 if is_int else 1
+        if x.min() < low or x.max() > high:
+            if is_int:
+                raise ValueError(
+                    "Integer pixel values out of acceptable range [0, 255]. "
+                    f"Found minimum value {x.min()} and maximum value {x.max()}. "
+                    "Ensure all pixel values are within the specified range."
+                )
+            else:
+                warnings.warn(
+                    "Float pixel values out of acceptable range [0.0, 1.0]. "
+                    f"Found minimum value {x.min()} and maximum value {x.max()}. "
+                    "Rescaling values to [0.0, 1.0] with min/max scaler.",
+                    stacklevel=2,
+                )
+                # Min-max scaling
+                x = (x - x.min()) / (x.max() - x.min())
+
+        # float or bool
+        if not is_int:
+            x = x * 255
+
+        return x.astype(np.uint8)
+
+    # Ref.: https://numpy.org/doc/stable/reference/generated/numpy.dtype.kind.html#numpy-dtype-kind
+    valid_data_types = {
+        "b": "bool",
+        "i": "signed integer",
+        "u": "unsigned integer",
+        "f": "floating",
+    }
+
+    if image.dtype.kind not in valid_data_types:
+        raise TypeError(
+            f"Invalid array data type: '{image.dtype}'. "
+            f"Must be one of {list(valid_data_types.values())}"
+        )
+
+    if image.ndim not in [2, 3]:
+        raise ValueError(f"`image` must be a 2D or 3D array but got image shape: {image.shape}")
+
+    if (image.ndim == 3) and (image.shape[2] not in [1, 3, 4]):
+        raise ValueError(f"Invalid channel length: {image.shape[2]}. Must be one of [1, 3, 4]")
+
+    # squeeze a 3D grayscale image since `Image.fromarray` doesn't accept it.
+    if image.ndim == 3 and image.shape[2] == 1:
+        image = image[:, :, 0]
+
+    image = _normalize_to_uint8(image)
+    return Image.fromarray(image)
+
+
+# MLflow media object: Image
+class Image:
+    """
+    `mlflow.Image` is an image media object that provides a lightweight option
+    for handling images in MLflow.
+    The image can be a numpy array, a PIL image, or a file path to an image. The image is
+    stored as a PIL image and can be logged to MLflow using `mlflow.log_image` or
+    `mlflow.log_table`.
+
+    Args:
+        image: Image can be a numpy array, a PIL image, or a file path to an image.
+
+    .. code-block:: python
+        :caption: Example
+
+        import mlflow
+        import numpy as np
+        from PIL import Image
+
+        # Create an image as a numpy array
+        image = np.zeros((100, 100, 3), dtype=np.uint8)
+        image[:, :50] = [255, 128, 0]
+        # Create an Image object
+        image_obj = mlflow.Image(image)
+        # Convert the Image object to a list of pixel values
+        pixel_values = image_obj.to_list()
+    """
+
+    def __init__(self, image: Union["numpy.ndarray", "PIL.Image.Image", str, list[Any]]):
+        import numpy as np
+
+        try:
+            from PIL import Image
+        except ImportError as exc:
+            raise ImportError(
+                "`mlflow.Image` requires Pillow to serialize a numpy array as an image. "
+                "Please install it via: `pip install Pillow`."
+            ) from exc
+
+        if isinstance(image, str):
+            self.image = Image.open(image)
+        elif isinstance(image, (list, np.ndarray)):
+            self.image = convert_to_pil_image(np.array(image))
+        elif isinstance(image, Image.Image):
+            self.image = image
+        else:
+            raise TypeError(
+                f"Unsupported image object type: {type(image)}. "
+                "`image` must be one of numpy.ndarray, "
+                "PIL.Image.Image, or a filepath to an image."
+            )
+        self.size = self.image.size
+
+    def to_list(self):
+        """
+        Convert the image to a list of pixel values.
+
+        Returns:
+            List of pixel values.
+        """
+        return list(self.image.getdata())
+
+    def to_array(self):
+        """
+        Convert the image to a numpy array.
+
+        Returns:
+            Numpy array of pixel values.
+        """
+        import numpy as np
+
+        return np.array(self.image)
+
+    def to_pil(self):
+        """
+        Convert the image to a PIL image.
+
+        Returns:
+            PIL image.
+        """
+        return self.image
+
+    def save(self, path: str):
+        """
+        Save the image to a file.
+
+        Args:
+            path: File path to save the image.
+        """
+        self.image.save(path)
+
+    def resize(self, size: tuple[int, int]):
+        """
+        Resize the image to the specified size.
+
+        Args:
+            size: Size to resize the image to.
+
+        Returns:
+            A copy of the resized image object.
+        """
+        image = self.image.resize(size)
+        return Image(image)

mlflow/tracking/registry.py

@@ -0,0 +1,86 @@
+import warnings
+from abc import ABCMeta
+
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+from mlflow.utils.plugins import get_entry_points
+from mlflow.utils.uri import get_uri_scheme
+
+
+class UnsupportedModelRegistryStoreURIException(MlflowException):
+    """Exception thrown when building a model registry store with an unsupported URI"""
+
+    def __init__(self, unsupported_uri, supported_uri_schemes):
+        message = (
+            " Model registry functionality is unavailable; got unsupported URI"
+            f" '{unsupported_uri}' for model registry data storage. Supported URI schemes are:"
+            f" {supported_uri_schemes}."
+            " See https://www.mlflow.org/docs/latest/tracking.html#storage for how to run"
+            " an MLflow server against one of the supported backend storage locations."
+        )
+        super().__init__(message, error_code=INVALID_PARAMETER_VALUE)
+        self.supported_uri_schemes = supported_uri_schemes
+
+
+class StoreRegistry:
+    """
+    Abstract class defining a scheme-based registry for store implementations.
+
+    This class allows the registration of a function or class to provide an
+    implementation for a given scheme of `store_uri` through the `register`
+    methods. Implementations declared though the entrypoints can be automatically
+    registered through the `register_entrypoints` method.
+
+    When instantiating a store through the `get_store` method, the scheme of
+    the store URI provided (or inferred from environment) will be used to
+    select which implementation to instantiate, which will be called with same
+    arguments passed to the `get_store` method.
+    """
+
+    __metaclass__ = ABCMeta
+
+    def __init__(self, group_name):
+        self._registry = {}
+        self.group_name = group_name
+
+    def register(self, scheme, store_builder):
+        self._registry[scheme] = store_builder
+
+    def register_entrypoints(self):
+        """Register tracking stores provided by other packages"""
+        for entrypoint in get_entry_points(self.group_name):
+            try:
+                self.register(entrypoint.name, entrypoint.load())
+            except (AttributeError, ImportError) as exc:
+                warnings.warn(
+                    'Failure attempting to register store for scheme "{}": {}'.format(
+                        entrypoint.name, str(exc)
+                    ),
+                    stacklevel=2,
+                )
+
+    def get_store_builder(self, store_uri):
+        """Get a store from the registry based on the scheme of store_uri
+
+        Args:
+            store_uri: The store URI. If None, it will be inferred from the environment. This
+                URI is used to select which tracking store implementation to instantiate
+                and is passed to the constructor of the implementation.
+
+        Returns:
+            A function that returns an instance of
+            ``mlflow.store.{tracking|model_registry}.AbstractStore`` that fulfills the store
+            URI requirements.
+        """
+        scheme = (
+            store_uri
+            if store_uri in {"databricks", "databricks-uc", "uc"}
+            else get_uri_scheme(store_uri)
+        )
+        try:
+            store_builder = self._registry[scheme]
+        except KeyError:
+            raise UnsupportedModelRegistryStoreURIException(
+                unsupported_uri=store_uri, supported_uri_schemes=list(self._registry.keys())
+            )
+        return store_builder

mlflow/tracking/request_auth/abstract_request_auth_provider.py

@@ -0,0 +1,34 @@
+from abc import ABC, abstractmethod
+
+from mlflow.utils.annotations import developer_stable
+
+
+@developer_stable
+class RequestAuthProvider(ABC):
+    """
+    Abstract base class for specifying custom request auth to add to outgoing requests
+
+    When a request is sent, MLflow will iterate through all registered RequestAuthProviders.
+    For each provider where ``get_name`` matches auth provider name, MLflow calls the ``get_auth``
+    method on the provider to compute request auth.
+
+    The resulting request auth will then be added and sent with the request.
+    """
+
+    @abstractmethod
+    def get_name(self):
+        """Get the name of the request auth provider.
+
+        Returns:
+            str of request auth provider name.
+        """
+
+    @abstractmethod
+    def get_auth(self):
+        """
+        Generate request auth object (e.g., `requests.auth import HTTPBasicAuth`). See
+        https://requests.readthedocs.io/en/latest/user/authentication/ for more details.
+
+        Returns:
+            request auth object.
+        """

mlflow/tracking/request_auth/registry.py

@@ -0,0 +1,60 @@
+import warnings
+
+from mlflow.utils.plugins import get_entry_points
+
+REQUEST_AUTH_PROVIDER_ENTRYPOINT = "mlflow.request_auth_provider"
+
+
+class RequestAuthProviderRegistry:
+    def __init__(self):
+        self._registry = []
+
+    def register(self, request_auth_provider):
+        self._registry.append(request_auth_provider())
+
+    def register_entrypoints(self):
+        for entrypoint in get_entry_points(REQUEST_AUTH_PROVIDER_ENTRYPOINT):
+            try:
+                self.register(entrypoint.load())
+            except (AttributeError, ImportError) as exc:
+                warnings.warn(
+                    'Failure attempting to register request auth provider "{}": {}'.format(
+                        entrypoint.name, str(exc)
+                    ),
+                    stacklevel=2,
+                )
+
+    def __iter__(self):
+        return iter(self._registry)
+
+
+_request_auth_provider_registry = RequestAuthProviderRegistry()
+_request_auth_provider_registry.register_entrypoints()
+
+
+def fetch_auth(request_auth):
+    """
+    Find the request auth from registered providers based on the auth provider's name.
+    The auth provider's name can be provided through environment variable `MLFLOW_TRACKING_AUTH`.
+
+    This function iterates through all request auth providers in the registry. Additional context
+    providers can be registered as described in
+    :py:class:`mlflow.tracking.request_auth.RequestAuthProvider`.
+
+    Args:
+        request_auth: The name of request auth provider.
+
+    Returns:
+        The auth object.
+    """
+
+    for auth_provider in _request_auth_provider_registry:
+        if auth_provider.get_name() == request_auth:
+            return auth_provider.get_auth()
+
+    warnings.warn(
+        f"Could not find any registered plugin for {request_auth}. "
+        "No authentication header will be added. Please check your "
+        "provider documentation for installing the right plugin or "
+        "correct provider name."
+    )

mlflow/tracking/request_header/abstract_request_header_provider.py

@@ -0,0 +1,36 @@
+from abc import ABCMeta, abstractmethod
+
+from mlflow.utils.annotations import developer_stable
+
+
+@developer_stable
+class RequestHeaderProvider:
+    """
+    Abstract base class for specifying custom request headers to add to outgoing requests
+    (e.g. request headers specifying the environment from which mlflow is running).
+
+    When a request is sent, MLflow will iterate through all registered RequestHeaderProviders.
+    For each provider where ``in_context`` returns ``True``, MLflow calls the ``request_headers``
+    method on the provider to compute request headers.
+
+    All resulting request headers will then be merged together and sent with the request.
+    """
+
+    __metaclass__ = ABCMeta
+
+    @abstractmethod
+    def in_context(self):
+        """Determine if MLflow is running in this context.
+
+        Returns:
+            bool indicating if in this context.
+
+        """
+
+    @abstractmethod
+    def request_headers(self):
+        """Generate context-specific request headers.
+
+        Returns:
+            dict of request headers.
+        """

mlflow/tracking/request_header/databricks_request_header_provider.py

@@ -0,0 +1,38 @@
+from mlflow.tracking.request_header.abstract_request_header_provider import RequestHeaderProvider
+from mlflow.utils import databricks_utils
+
+
+class DatabricksRequestHeaderProvider(RequestHeaderProvider):
+    """
+    Provides request headers indicating the type of Databricks environment from which a request
+    was made.
+    """
+
+    def in_context(self):
+        return (
+            databricks_utils.is_in_cluster()
+            or databricks_utils.is_in_databricks_notebook()
+            or databricks_utils.is_in_databricks_job()
+        )
+
+    def request_headers(self):
+        request_headers = {}
+        if databricks_utils.is_in_databricks_notebook():
+            request_headers["notebook_id"] = databricks_utils.get_notebook_id()
+        if databricks_utils.is_in_databricks_job():
+            request_headers["job_id"] = databricks_utils.get_job_id()
+            request_headers["job_run_id"] = databricks_utils.get_job_run_id()
+            request_headers["job_type"] = databricks_utils.get_job_type()
+        if databricks_utils.is_in_cluster():
+            request_headers["cluster_id"] = databricks_utils.get_cluster_id()
+        command_run_id = databricks_utils.get_command_run_id()
+        if command_run_id is not None:
+            request_headers["command_run_id"] = command_run_id
+        workload_id = databricks_utils.get_workload_id()
+        workload_class = databricks_utils.get_workload_class()
+        if workload_id is not None:
+            request_headers["workload_id"] = workload_id
+        if workload_class is not None:
+            request_headers["workload_class"] = workload_class
+
+        return request_headers

mlflow/tracking/request_header/default_request_header_provider.py

@@ -0,0 +1,17 @@
+from mlflow import __version__
+from mlflow.tracking.request_header.abstract_request_header_provider import RequestHeaderProvider
+
+_USER_AGENT = "User-Agent"
+_DEFAULT_HEADERS = {_USER_AGENT: f"mlflow-python-client/{__version__}"}
+
+
+class DefaultRequestHeaderProvider(RequestHeaderProvider):
+    """
+    Provides default request headers for outgoing request.
+    """
+
+    def in_context(self):
+        return True
+
+    def request_headers(self):
+        return dict(**_DEFAULT_HEADERS)

mlflow/tracking/request_header/registry.py

@@ -0,0 +1,79 @@
+import logging
+import warnings
+
+from mlflow.tracking.request_header.databricks_request_header_provider import (
+    DatabricksRequestHeaderProvider,
+)
+from mlflow.tracking.request_header.default_request_header_provider import (
+    DefaultRequestHeaderProvider,
+)
+from mlflow.utils.plugins import get_entry_points
+
+_logger = logging.getLogger(__name__)
+
+
+class RequestHeaderProviderRegistry:
+    def __init__(self):
+        self._registry = []
+
+    def register(self, request_header_provider):
+        self._registry.append(request_header_provider())
+
+    def register_entrypoints(self):
+        """Register tracking stores provided by other packages"""
+        for entrypoint in get_entry_points("mlflow.request_header_provider"):
+            try:
+                self.register(entrypoint.load())
+            except (AttributeError, ImportError) as exc:
+                warnings.warn(
+                    'Failure attempting to register request header provider "{}": {}'.format(
+                        entrypoint.name, str(exc)
+                    ),
+                    stacklevel=2,
+                )
+
+    def __iter__(self):
+        return iter(self._registry)
+
+
+_request_header_provider_registry = RequestHeaderProviderRegistry()
+_request_header_provider_registry.register(DatabricksRequestHeaderProvider)
+_request_header_provider_registry.register(DefaultRequestHeaderProvider)
+
+_request_header_provider_registry.register_entrypoints()
+
+
+def resolve_request_headers(request_headers=None):
+    """Generate a set of request headers from registered providers.
+
+    Request headers are resolved in the order that providers are registered. Argument headers are
+    applied last. This function iterates through all request header providers in the registry.
+    Additional context providers can be registered as described in
+    :py:class:`mlflow.tracking.request_header.RequestHeaderProvider`.
+
+    Args:
+        request_headers: A dictionary of request headers to override. If specified, headers passed
+            in this argument will override those inferred from the context.
+
+    Returns:
+        A dictionary of resolved headers.
+    """
+
+    all_request_headers = {}
+    for provider in _request_header_provider_registry:
+        try:
+            if provider.in_context():
+                # all_request_headers.update(provider.request_headers())
+                for header, value in provider.request_headers().items():
+                    all_request_headers[header] = (
+                        f"{all_request_headers[header]} {value}"
+                        if header in all_request_headers
+                        else value
+                    )
+        except Exception as e:
+            _logger.warning("Encountered unexpected error during resolving request headers: %s", e)
+
+    if request_headers is not None:
+        all_request_headers.update(request_headers)
+
+    return all_request_headers