genesis-flow 1.0.0__py3-none-any.whl

mlflow/pyfunc/dbconnect_artifact_cache.py

@@ -0,0 +1,144 @@
+import json
+import os
+import subprocess
+import tarfile
+
+from mlflow.utils.databricks_utils import is_in_databricks_runtime
+from mlflow.utils.file_utils import get_or_create_tmp_dir
+
+_CACHE_MAP_FILE_NAME = "db_connect_artifact_cache.json"
+
+
+class DBConnectArtifactCache:
+    """
+    Manages Databricks Connect artifacts cache.
+    Note it doesn't support OSS Spark Connect.
+
+    This class can be used in the following environment:
+     - Databricks shared cluster python notebook REPL
+     - Databricks Serverless python notebook REPL
+     - Databricks connect client python REPL that connects to remote Databricks Serverless
+     - Databricks connect client python REPL that connects to remote Databricks shared cluster
+
+    .. code-block:: python
+        :caption: Example
+
+        # client side code
+        db_artifact_cache = DBConnectArtifactCache.get_or_create()
+        db_artifact_cache.add_artifact_archive("archive1", "/tmp/archive1.tar.gz")
+
+
+        @pandas_udf(...)
+        def my_udf(x):
+            # we can get the unpacked archive files in `archive1_unpacked_dir`
+            archive1_unpacked_dir = db_artifact_cache.get("archive1")
+    """
+
+    _global_cache = None
+
+    @staticmethod
+    def get_or_create(spark):
+        if (
+            DBConnectArtifactCache._global_cache is None
+            or spark is not DBConnectArtifactCache._global_cache._spark
+        ):
+            DBConnectArtifactCache._global_cache = DBConnectArtifactCache(spark)
+            cache_file = os.path.join(get_or_create_tmp_dir(), _CACHE_MAP_FILE_NAME)
+            if is_in_databricks_runtime() and os.path.exists(cache_file):
+                # In databricks runtime (shared cluster or Serverless), when you restart the
+                # notebook REPL by %restart_python or dbutils.library.restartPython(), the
+                # DBConnect session is still preserved. So in this case, we can reuse the cached
+                # artifact files.
+                # So that when adding artifact, the cache map is serialized to local disk file
+                # `db_connect_artifact_cache.json` and after REPL restarts,
+                # `DBConnectArtifactCache` restores the cache map by loading data from the file.
+                with open(cache_file) as f:
+                    DBConnectArtifactCache._global_cache._cache = json.load(f)
+        return DBConnectArtifactCache._global_cache
+
+    def __init__(self, spark):
+        self._spark = spark
+        self._cache = {}
+
+    def __getstate__(self):
+        """
+        The `DBConnectArtifactCache` instance is created in Databricks Connect client side,
+        and it will be pickled to Databricks Connect UDF sandbox
+        (see `get_unpacked_artifact_dir` method), but Spark Connect client object is
+        not pickle-able, we need to skip this field.
+        """
+        state = self.__dict__.copy()
+        # Don't pickle `_spark`
+        del state["_spark"]
+        return state
+
+    def __setstate__(self, state):
+        self.__dict__.update(state)
+        self._spark = None
+
+    def has_cache_key(self, cache_key):
+        return cache_key in self._cache
+
+    def add_artifact_archive(self, cache_key, artifact_archive_path):
+        """
+        Add an artifact archive file to Databricks connect cache.
+        The archive file must be 'tar.gz' format.
+        You can only call this method in Databricks Connect client side.
+        """
+        if not artifact_archive_path.endswith(".tar.gz"):
+            raise RuntimeError(
+                "'add_artifact_archive' only supports archive file in 'tar.gz' format."
+            )
+
+        archive_file_name = os.path.basename(artifact_archive_path)
+        if cache_key not in self._cache:
+            self._spark.addArtifact(artifact_archive_path, archive=True)
+            self._cache[cache_key] = archive_file_name
+
+        if is_in_databricks_runtime():
+            with open(os.path.join(get_or_create_tmp_dir(), _CACHE_MAP_FILE_NAME), "w") as f:
+                json.dump(self._cache, f)
+
+    def get_unpacked_artifact_dir(self, cache_key):
+        """
+        Get unpacked artifact directory path, you can only call this method
+        inside Databricks Connect spark UDF sandbox.
+        """
+        if cache_key not in self._cache:
+            raise RuntimeError(f"The artifact '{cache_key}' does not exist.")
+        archive_file_name = self._cache[cache_key]
+
+        if session_id := os.environ.get("DB_SESSION_UUID"):
+            return (
+                f"/local_disk0/.ephemeral_nfs/artifacts/{session_id}/archives/{archive_file_name}"
+            )
+
+        # If 'DB_SESSION_UUID' environment variable does not exist, it means it is running
+        # in a dedicated mode Spark cluster.
+        return os.path.join(os.getcwd(), archive_file_name)
+
+
+def archive_directory(input_dir, archive_file_path):
+    """
+    Archive the `input_dir` directory, save the archive file to `archive_file_path`,
+    the generated archive file is 'tar.gz' format.
+    Note: all symlink files in the input directory are kept as it is in the archive file.
+    """
+
+    archive_file_path = os.path.abspath(archive_file_path)
+    # Note: `shutil.make_archive` doesn't work because it replaces symlink files with
+    #  the file symlink pointing to, which is not the expected behavior in our usage.
+    #  We need to pack the python and virtualenv environment, which contains a bunch of
+    #  symlink files.
+    subprocess.check_call(
+        ["tar", "-czf", archive_file_path, *os.listdir(input_dir)],
+        cwd=input_dir,
+    )
+    return archive_file_path
+
+
+def extract_archive_to_dir(archive_path, dest_dir):
+    os.makedirs(dest_dir, exist_ok=True)
+    with tarfile.open(archive_path, "r") as tar:
+        tar.extractall(path=dest_dir)
+    return dest_dir

mlflow/pyfunc/loaders/__init__.py

@@ -0,0 +1,7 @@
+import mlflow.pyfunc.loaders.chat_agent
+import mlflow.pyfunc.loaders.chat_model
+import mlflow.pyfunc.loaders.code_model
+from mlflow.utils.pydantic_utils import IS_PYDANTIC_V2_OR_NEWER
+
+if IS_PYDANTIC_V2_OR_NEWER:
+    import mlflow.pyfunc.loaders.responses_agent  # noqa: F401

mlflow/pyfunc/loaders/chat_agent.py

@@ -0,0 +1,117 @@
+from typing import Any, Generator, Optional
+
+import pydantic
+
+from mlflow.exceptions import MlflowException
+from mlflow.models.utils import _convert_llm_ndarray_to_list
+from mlflow.protos.databricks_pb2 import INTERNAL_ERROR
+from mlflow.pyfunc.model import (
+    _load_context_model_and_signature,
+)
+from mlflow.types.agent import (
+    ChatAgentChunk,
+    ChatAgentMessage,
+    ChatAgentResponse,
+    ChatContext,
+)
+from mlflow.types.type_hints import model_validate
+from mlflow.utils.annotations import experimental
+
+
+def _load_pyfunc(model_path: str, model_config: Optional[dict[str, Any]] = None):
+    _, chat_agent, _ = _load_context_model_and_signature(model_path, model_config)
+    return _ChatAgentPyfuncWrapper(chat_agent)
+
+
+@experimental(version="2.20.0")
+class _ChatAgentPyfuncWrapper:
+    """
+    Wrapper class that converts dict inputs to pydantic objects accepted by :class:`~ChatAgent`.
+    """
+
+    def __init__(self, chat_agent):
+        """
+        Args:
+            chat_agent: An instance of a subclass of :class:`~ChatAgent`.
+        """
+        self.chat_agent = chat_agent
+
+    def get_raw_model(self):
+        """
+        Returns the underlying model.
+        """
+        return self.chat_agent
+
+    def _convert_input(
+        self, model_input
+    ) -> tuple[list[ChatAgentMessage], Optional[ChatContext], Optional[dict[str, Any]]]:
+        import pandas
+
+        if isinstance(model_input, dict):
+            dict_input = model_input
+        elif isinstance(model_input, pandas.DataFrame):
+            dict_input = {
+                k: _convert_llm_ndarray_to_list(v[0])
+                for k, v in model_input.to_dict(orient="list").items()
+            }
+        else:
+            raise MlflowException(
+                "Unsupported model input type. Expected a dict or pandas.DataFrame, but got "
+                f"{type(model_input)} instead.",
+                error_code=INTERNAL_ERROR,
+            )
+
+        messages = [ChatAgentMessage(**message) for message in dict_input.get("messages", [])]
+        context = ChatContext(**dict_input["context"]) if "context" in dict_input else None
+        custom_inputs = dict_input.get("custom_inputs", None)
+
+        return messages, context, custom_inputs
+
+    def _response_to_dict(self, response, pydantic_class) -> dict[str, Any]:
+        if isinstance(response, pydantic_class):
+            return response.model_dump_compat(exclude_none=True)
+        try:
+            model_validate(pydantic_class, response)
+        except pydantic.ValidationError as e:
+            raise MlflowException(
+                message=(
+                    f"Model returned an invalid response. Expected a {pydantic_class.__name__} "
+                    f"object or dictionary with the same schema. Pydantic validation error: {e}"
+                ),
+                error_code=INTERNAL_ERROR,
+            ) from e
+        return response
+
+    def predict(self, model_input: dict[str, Any], params=None) -> dict[str, Any]:
+        """
+        Args:
+            model_input: A dict with the
+                :py:class:`ChatAgentRequest <mlflow.types.agent.ChatAgentRequest>` schema.
+            params: Unused in this function, but required in the signature because
+                `load_model_and_predict` in `utils/_capture_modules.py` expects a params field
+
+        Returns:
+            A dict with the (:py:class:`ChatAgentResponse <mlflow.types.agent.ChatAgentResponse>`)
+                schema.
+        """
+        messages, context, custom_inputs = self._convert_input(model_input)
+        response = self.chat_agent.predict(messages, context, custom_inputs)
+        return self._response_to_dict(response, ChatAgentResponse)
+
+    def predict_stream(
+        self, model_input: dict[str, Any], params=None
+    ) -> Generator[dict[str, Any], None, None]:
+        """
+        Args:
+            model_input: A dict with the
+                :py:class:`ChatAgentRequest <mlflow.types.agent.ChatAgentRequest>` schema.
+            params: Unused in this function, but required in the signature because
+                `load_model_and_predict` in `utils/_capture_modules.py` expects a params field
+
+        Returns:
+            A generator over dicts with the
+                (:py:class:`ChatAgentChunk <mlflow.types.agent.ChatAgentChunk>`) schema.
+        """
+        messages, context, custom_inputs = self._convert_input(model_input)
+        for response in self.chat_agent.predict_stream(messages, context, custom_inputs):
+            yield self._response_to_dict(response, ChatAgentChunk)

mlflow/pyfunc/loaders/chat_model.py

@@ -0,0 +1,125 @@
+import inspect
+import logging
+from typing import Any, Generator, Optional
+
+from mlflow.exceptions import MlflowException
+from mlflow.models.utils import _convert_llm_ndarray_to_list
+from mlflow.protos.databricks_pb2 import INTERNAL_ERROR
+from mlflow.pyfunc.model import (
+    _load_context_model_and_signature,
+)
+from mlflow.types.llm import ChatCompletionChunk, ChatCompletionResponse, ChatMessage, ChatParams
+
+_logger = logging.getLogger(__name__)
+
+
+def _load_pyfunc(model_path: str, model_config: Optional[dict[str, Any]] = None):
+    context, chat_model, signature = _load_context_model_and_signature(model_path, model_config)
+    return _ChatModelPyfuncWrapper(chat_model=chat_model, context=context, signature=signature)
+
+
+class _ChatModelPyfuncWrapper:
+    """
+    Wrapper class that converts dict inputs to pydantic objects accepted by :class:`~ChatModel`.
+    """
+
+    def __init__(self, chat_model, context, signature):
+        """
+        Args:
+            chat_model: An instance of a subclass of :class:`~ChatModel`.
+            context: A :class:`~PythonModelContext` instance containing artifacts that
+                        ``chat_model`` may use when performing inference.
+            signature: :class:`~ModelSignature` instance describing model input and output.
+        """
+        self.chat_model = chat_model
+        self.context = context
+        self.signature = signature
+
+    def get_raw_model(self):
+        """
+        Returns the underlying model.
+        """
+        return self.chat_model
+
+    def _convert_input(self, model_input):
+        import pandas
+
+        if isinstance(model_input, dict):
+            dict_input = model_input
+        elif isinstance(model_input, pandas.DataFrame):
+            dict_input = {
+                k: _convert_llm_ndarray_to_list(v[0])
+                for k, v in model_input.to_dict(orient="list").items()
+            }
+        else:
+            raise MlflowException(
+                "Unsupported model input type. Expected a dict or pandas.DataFrame, "
+                f"but got {type(model_input)} instead.",
+                error_code=INTERNAL_ERROR,
+            )
+
+        messages = [ChatMessage.from_dict(message) for message in dict_input.pop("messages", [])]
+        params = ChatParams.from_dict(dict_input)
+        return messages, params
+
+    def predict(
+        self, model_input: dict[str, Any], params: Optional[dict[str, Any]] = None
+    ) -> dict[str, Any]:
+        """
+        Args:
+            model_input: Model input data in the form of a chat request.
+            params: Additional parameters to pass to the model for inference.
+                       Unused in this implementation, as the params are handled
+                       via ``self._convert_input()``.
+
+        Returns:
+            Model predictions in :py:class:`~ChatCompletionResponse` format.
+        """
+        messages, params = self._convert_input(model_input)
+        parameters = inspect.signature(self.chat_model.predict).parameters
+        if "context" in parameters or len(parameters) == 3:
+            response = self.chat_model.predict(self.context, messages, params)
+        else:
+            response = self.chat_model.predict(messages, params)
+        return self._response_to_dict(response)
+
+    def _response_to_dict(self, response: ChatCompletionResponse) -> dict[str, Any]:
+        if not isinstance(response, ChatCompletionResponse):
+            raise MlflowException(
+                "Model returned an invalid response. Expected a ChatCompletionResponse, but "
+                f"got {type(response)} instead.",
+                error_code=INTERNAL_ERROR,
+            )
+        return response.to_dict()
+
+    def _streaming_response_to_dict(self, response: ChatCompletionChunk) -> dict[str, Any]:
+        if not isinstance(response, ChatCompletionChunk):
+            raise MlflowException(
+                "Model returned an invalid response. Expected a ChatCompletionChunk, but "
+                f"got {type(response)} instead.",
+                error_code=INTERNAL_ERROR,
+            )
+        return response.to_dict()
+
+    def predict_stream(
+        self, model_input: dict[str, Any], params: Optional[dict[str, Any]] = None
+    ) -> Generator[dict[str, Any], None, None]:
+        """
+        Args:
+            model_input: Model input data in the form of a chat request.
+            params: Additional parameters to pass to the model for inference.
+                       Unused in this implementation, as the params are handled
+                       via ``self._convert_input()``.
+
+        Returns:
+            Generator over model predictions in :py:class:`~ChatCompletionChunk` format.
+        """
+        messages, params = self._convert_input(model_input)
+        parameters = inspect.signature(self.chat_model.predict_stream).parameters
+        if "context" in parameters or len(parameters) == 3:
+            stream = self.chat_model.predict_stream(self.context, messages, params)
+        else:
+            stream = self.chat_model.predict_stream(messages, params)
+
+        for response in stream:
+            yield self._streaming_response_to_dict(response)

mlflow/pyfunc/loaders/code_model.py

@@ -0,0 +1,31 @@
+from typing import Any, Optional
+
+from mlflow.pyfunc.loaders.chat_agent import _ChatAgentPyfuncWrapper
+from mlflow.pyfunc.loaders.chat_model import _ChatModelPyfuncWrapper
+from mlflow.pyfunc.model import (
+    ChatAgent,
+    ChatModel,
+    _load_context_model_and_signature,
+    _PythonModelPyfuncWrapper,
+)
+
+try:
+    from mlflow.pyfunc.model import ResponsesAgent
+
+    IS_RESPONSES_AGENT_AVAILABLE = True
+except ImportError:
+    IS_RESPONSES_AGENT_AVAILABLE = False
+
+
+def _load_pyfunc(local_path: str, model_config: Optional[dict[str, Any]] = None):
+    context, model, signature = _load_context_model_and_signature(local_path, model_config)
+    if isinstance(model, ChatModel):
+        return _ChatModelPyfuncWrapper(model, context, signature)
+    elif isinstance(model, ChatAgent):
+        return _ChatAgentPyfuncWrapper(model)
+    elif IS_RESPONSES_AGENT_AVAILABLE and isinstance(model, ResponsesAgent):
+        from mlflow.pyfunc.loaders.responses_agent import _ResponsesAgentPyfuncWrapper
+
+        return _ResponsesAgentPyfuncWrapper(model)
+    else:
+        return _PythonModelPyfuncWrapper(model, context, signature)

mlflow/pyfunc/loaders/responses_agent.py

@@ -0,0 +1,112 @@
+from typing import Any, Generator, Optional
+
+import pydantic
+
+from mlflow.exceptions import MlflowException
+from mlflow.models.utils import _convert_llm_ndarray_to_list
+from mlflow.protos.databricks_pb2 import INTERNAL_ERROR
+from mlflow.pyfunc.model import _load_context_model_and_signature
+from mlflow.types.type_hints import model_validate
+from mlflow.utils.annotations import experimental
+from mlflow.utils.pydantic_utils import IS_PYDANTIC_V2_OR_NEWER
+
+if not IS_PYDANTIC_V2_OR_NEWER:
+    raise ImportError(
+        "ResponsesAgent and its pydantic classes are not supported in pydantic v1. "
+        "Please upgrade to pydantic v2 or newer to use ResponsesAgent.",
+    )
+
+from mlflow.types.responses import (
+    ResponsesAgentRequest,
+    ResponsesAgentResponse,
+    ResponsesAgentStreamEvent,
+)
+
+
+def _load_pyfunc(model_path: str, model_config: Optional[dict[str, Any]] = None):
+    _, responses_agent, _ = _load_context_model_and_signature(model_path, model_config)
+    return _ResponsesAgentPyfuncWrapper(responses_agent)
+
+
+@experimental(version="3.0.0")
+class _ResponsesAgentPyfuncWrapper:
+    """
+    Wrapper class that converts dict inputs to pydantic objects accepted by
+    :class:`~ResponsesAgent`.
+    """
+
+    def __init__(self, responses_agent):
+        self.responses_agent = responses_agent
+
+    def get_raw_model(self):
+        """
+        Returns the underlying model.
+        """
+        return self.responses_agent
+
+    def _convert_input(self, model_input) -> ResponsesAgentRequest:
+        import pandas
+
+        if isinstance(model_input, pandas.DataFrame):
+            model_input = {
+                k: _convert_llm_ndarray_to_list(v[0])
+                for k, v in model_input.to_dict(orient="list").items()
+            }
+        elif not isinstance(model_input, dict):
+            raise MlflowException(
+                "Unsupported model input type. Expected a dict or pandas.DataFrame, but got "
+                f"{type(model_input)} instead.",
+                error_code=INTERNAL_ERROR,
+            )
+        return ResponsesAgentRequest(**model_input)
+
+    def _response_to_dict(self, response, pydantic_class) -> dict[str, Any]:
+        if isinstance(response, pydantic_class):
+            return response.model_dump_compat(exclude_none=True)
+        try:
+            model_validate(pydantic_class, response)
+        except pydantic.ValidationError as e:
+            raise MlflowException(
+                message=(
+                    f"Model returned an invalid response. Expected a {pydantic_class.__name__} "
+                    f"object or dictionary with the same schema. Pydantic validation error: {e}"
+                ),
+                error_code=INTERNAL_ERROR,
+            ) from e
+        return response
+
+    def predict(self, model_input: dict[str, Any], params=None) -> dict[str, Any]:
+        """
+        Args:
+            model_input: A dict with the
+                :py:class:`ResponsesRequest <mlflow.types.responses.ResponsesRequest>` schema.
+            params: Unused in this function, but required in the signature because
+                `load_model_and_predict` in `utils/_capture_modules.py` expects a params field
+
+        Returns:
+            A dict with the
+            (:py:class:`ResponsesResponse <mlflow.types.responses.ResponsesResponse>`)
+            schema.
+        """
+        request = self._convert_input(model_input)
+        response = self.responses_agent.predict(request)
+        return self._response_to_dict(response, ResponsesAgentResponse)
+
+    def predict_stream(
+        self, model_input: dict[str, Any], params=None
+    ) -> Generator[dict[str, Any], None, None]:
+        """
+        Args:
+            model_input: A dict with the
+                :py:class:`ResponsesRequest <mlflow.types.responses.ResponsesRequest>` schema.
+            params: Unused in this function, but required in the signature because
+                `load_model_and_predict` in `utils/_capture_modules.py` expects a params field
+
+        Returns:
+            A generator over dicts with the
+                (:py:class:`ResponsesStreamEvent <mlflow.types.responses.ResponsesStreamEvent>`)
+                schema.
+        """
+        request = self._convert_input(model_input)
+        for response in self.responses_agent.predict_stream(request):
+            yield self._response_to_dict(response, ResponsesAgentStreamEvent)

mlflow/pyfunc/mlserver.py

@@ -0,0 +1,46 @@
+import logging
+import os
+from typing import Optional
+
+_logger = logging.getLogger(__name__)
+
+MLServerMLflowRuntime = "mlserver_mlflow.MLflowRuntime"
+MLServerDefaultModelName = "mlflow-model"
+
+
+def get_cmd(
+    model_uri: str,
+    port: Optional[int] = None,
+    host: Optional[str] = None,
+    timeout: Optional[int] = None,
+    nworkers: Optional[int] = None,
+    model_name: Optional[str] = None,
+    model_version: Optional[str] = None,
+) -> tuple[str, dict[str, str]]:
+    cmd = f"mlserver start {model_uri}"
+
+    cmd_env = os.environ.copy()
+
+    if port:
+        cmd_env["MLSERVER_HTTP_PORT"] = str(port)
+
+    if host:
+        cmd_env["MLSERVER_HOST"] = host
+
+    if timeout:
+        _logger.warning("Timeout is not yet supported in MLServer.")
+
+    if nworkers:
+        cmd_env["MLSERVER_PARALLEL_WORKERS"] = str(nworkers)
+
+    # give precedence to user env var input
+    cmd_env["MLSERVER_MODEL_NAME"] = (
+        cmd_env.get("MLSERVER_MODEL_NAME") or model_name or MLServerDefaultModelName
+    )
+    if model_version and not cmd_env.get("MLSERVER_MODEL_VERSION"):
+        cmd_env["MLSERVER_MODEL_VERSION"] = model_version
+
+    cmd_env["MLSERVER_MODEL_IMPLEMENTATION"] = MLServerMLflowRuntime
+    cmd_env["MLSERVER_MODEL_URI"] = model_uri
+
+    return cmd, cmd_env