genesis-flow 1.0.0__py3-none-any.whl

mlflow/pyfunc/_mlflow_pyfunc_backend_predict.py

@@ -0,0 +1,61 @@
+"""
+This script should be executed in a fresh python interpreter process using `subprocess`.
+"""
+
+import argparse
+
+from mlflow.pyfunc.scoring_server import _predict
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model-uri", required=True)
+    parser.add_argument("--input-path", required=False)
+    parser.add_argument("--output-path", required=False)
+    parser.add_argument("--content-type", required=True)
+    return parser.parse_args()
+
+
+# Guidance for fixing missing module error
+_MISSING_MODULE_HELP_MSG = (
+    "Exception occurred while running inference: {e}"
+    "\n\n"
+    "\033[93m[Hint] It appears that your MLflow Model doesn't contain the required "
+    "dependency '{missing_module}' to run model inference. When logging a model, MLflow "
+    "detects dependencies based on the model flavor, but it is possible that some "
+    "dependencies are not captured. In this case, you can manually add dependencies "
+    "using the `extra_pip_requirements` parameter of `mlflow.pyfunc.log_model`.\033[0m"
+    """
+
+\033[1mSample code:\033[0m
+    ----
+    mlflow.pyfunc.log_model(
+        artifact_path="model",
+        python_model=your_model,
+        extra_pip_requirements=["{missing_module}==x.y.z"]
+    )
+    ----
+
+    For mode guidance on fixing missing dependencies, please refer to the MLflow docs:
+    https://www.mlflow.org/docs/latest/deployment/index.html#how-to-fix-dependency-errors-when-serving-my-model
+"""
+)
+
+
+def main():
+    args = parse_args()
+
+    try:
+        _predict(
+            model_uri=args.model_uri,
+            input_path=args.input_path if args.input_path else None,
+            output_path=args.output_path if args.output_path else None,
+            content_type=args.content_type,
+        )
+    except ModuleNotFoundError as e:
+        message = _MISSING_MODULE_HELP_MSG.format(e=str(e), missing_module=e.name)
+        raise RuntimeError(message) from e
+
+
+if __name__ == "__main__":
+    main()

mlflow/pyfunc/backend.py

@@ -0,0 +1,523 @@
+import ctypes
+import json
+import logging
+import os
+import pathlib
+import shlex
+import signal
+import subprocess
+import sys
+import warnings
+from pathlib import Path
+
+from mlflow import pyfunc
+from mlflow.exceptions import MlflowException
+from mlflow.models import FlavorBackend, Model, docker_utils
+from mlflow.models.docker_utils import PYTHON_SLIM_BASE_IMAGE, UBUNTU_BASE_IMAGE
+from mlflow.pyfunc import (
+    ENV,
+    _extract_conda_env,
+    _mlflow_pyfunc_backend_predict,
+    mlserver,
+    scoring_server,
+)
+from mlflow.tracking.artifact_utils import _download_artifact_from_uri
+from mlflow.utils import env_manager as em
+from mlflow.utils.conda import get_conda_bin_executable, get_or_create_conda_env
+from mlflow.utils.environment import Environment, _get_pip_install_mlflow, _PythonEnv
+from mlflow.utils.file_utils import (
+    TempDir,
+    get_or_create_nfs_tmp_dir,
+    get_or_create_tmp_dir,
+    path_to_local_file_uri,
+)
+from mlflow.utils.model_utils import _get_all_flavor_configurations
+from mlflow.utils.nfs_on_spark import get_nfs_cache_root_dir
+from mlflow.utils.os import is_windows
+from mlflow.utils.process import ShellCommandException, cache_return_value_per_process
+from mlflow.utils.virtualenv import _get_or_create_virtualenv
+from mlflow.version import VERSION
+
+_logger = logging.getLogger(__name__)
+
+_STDIN_SERVER_SCRIPT = Path(__file__).parent.joinpath("stdin_server.py")
+
+# Flavors that require Java to be installed in the environment
+JAVA_FLAVORS = {"johnsnowlabs", "h2o", "spark"}
+
+# Some flavor requires additional packages to be installed in the environment
+FLAVOR_SPECIFIC_APT_PACKAGES = {
+    "lightgbm": ["libgomp1"],
+    "paddle": ["libgomp1"],
+}
+
+# Directory to store loaded model inside the Docker context directory
+_MODEL_DIR_NAME = "model_dir"
+LOCAL_ENV_MANAGER_ERROR_MESSAGE = "We cannot use 'LOCAL' environment manager "
+"for your model configuration. Please specify a virtualenv or conda environment "
+"manager instead with `--env-manager` argument."
+
+
+def _set_mlflow_config_env(command_env, model_config):
+    if model_config:
+        command_env[scoring_server.SERVING_MODEL_CONFIG] = json.dumps(model_config)
+    return command_env
+
+
+class PyFuncBackend(FlavorBackend):
+    """
+    Flavor backend implementation for the generic python models.
+    """
+
+    def __init__(  # noqa: D417
+        self,
+        config,
+        env_manager,
+        workers=1,
+        install_mlflow=False,
+        create_env_root_dir=False,
+        env_root_dir=None,
+        **kwargs,
+    ):
+        """
+        Args:
+            env_manager: Environment manager to use for preparing the environment. If None,
+                MLflow will automatically pick the env manager based on the model's flavor
+                configuration for generate_dockerfile. It can't be None for other methods.
+            env_root_dir: Root path for conda env. If None, use Conda's default environments
+                directory. Note if this is set, conda package cache path becomes
+                "{env_root_dir}/conda_cache_pkgs" instead of the global package cache
+                path, and pip package cache path becomes
+                "{env_root_dir}/pip_cache_pkgs" instead of the global package cache
+                path.
+        """
+        super().__init__(config=config, **kwargs)
+        self._nworkers = workers or 1
+        if env_manager == em.CONDA and ENV not in config:
+            warnings.warn(
+                "Conda environment is not specified in config `env`. Using local environment."
+            )
+            env_manager = em.LOCAL
+        self._env_manager = env_manager
+        self._install_mlflow = install_mlflow
+        self._env_id = os.environ.get("MLFLOW_HOME", VERSION) if install_mlflow else None
+        self._create_env_root_dir = create_env_root_dir
+        self._env_root_dir = env_root_dir
+        self._environment = None
+
+    def prepare_env(
+        self, model_uri, capture_output=False, pip_requirements_override=None, extra_envs=None
+    ):
+        if self._environment is not None:
+            return self._environment
+
+        @cache_return_value_per_process
+        def _get_or_create_env_root_dir(should_use_nfs):
+            if should_use_nfs:
+                root_tmp_dir = get_or_create_nfs_tmp_dir()
+            else:
+                root_tmp_dir = get_or_create_tmp_dir()
+
+            envs_root_dir = os.path.join(root_tmp_dir, "envs")
+            os.makedirs(envs_root_dir, exist_ok=True)
+            return envs_root_dir
+
+        local_path = _download_artifact_from_uri(model_uri)
+        if self._create_env_root_dir:
+            if self._env_root_dir is not None:
+                raise Exception("env_root_dir can not be set when create_env_root_dir=True")
+            nfs_root_dir = get_nfs_cache_root_dir()
+            env_root_dir = _get_or_create_env_root_dir(nfs_root_dir is not None)
+        else:
+            env_root_dir = self._env_root_dir
+
+        if self._env_manager in {em.VIRTUALENV, em.UV}:
+            activate_cmd = _get_or_create_virtualenv(
+                local_path,
+                self._env_id,
+                env_root_dir=env_root_dir,
+                capture_output=capture_output,
+                pip_requirements_override=pip_requirements_override,
+                env_manager=self._env_manager,
+            )
+            self._environment = Environment(activate_cmd, extra_env=extra_envs)
+        elif self._env_manager == em.CONDA:
+            conda_env_path = os.path.join(local_path, _extract_conda_env(self._config[ENV]))
+            self._environment = get_or_create_conda_env(
+                conda_env_path,
+                env_id=self._env_id,
+                capture_output=capture_output,
+                env_root_dir=env_root_dir,
+                pip_requirements_override=pip_requirements_override,
+                extra_envs=extra_envs,
+            )
+
+        elif self._env_manager == em.LOCAL:
+            raise Exception("Prepare env should not be called with local env manager!")
+        else:
+            raise Exception(f"Unexpected env manager value '{self._env_manager}'")
+
+        if self._install_mlflow:
+            self._environment.execute(_get_pip_install_mlflow())
+        else:
+            self._environment.execute('python -c ""')
+
+        return self._environment
+
+    def predict(
+        self,
+        model_uri,
+        input_path,
+        output_path,
+        content_type,
+        pip_requirements_override=None,
+        extra_envs=None,
+    ):
+        """
+        Generate predictions using generic python model saved with MLflow. The expected format of
+        the input JSON is the MLflow scoring format.
+        Return the prediction results as a JSON.
+        """
+        local_path = _download_artifact_from_uri(model_uri)
+        # NB: Absolute windows paths do not work with mlflow apis, use file uri to ensure
+        # platform compatibility.
+        local_uri = path_to_local_file_uri(local_path)
+
+        if self._env_manager != em.LOCAL:
+            predict_cmd = [
+                "python",
+                _mlflow_pyfunc_backend_predict.__file__,
+                "--model-uri",
+                str(local_uri),
+                "--content-type",
+                shlex.quote(str(content_type)),
+            ]
+            if input_path:
+                predict_cmd += ["--input-path", shlex.quote(str(input_path))]
+            if output_path:
+                predict_cmd += ["--output-path", shlex.quote(str(output_path))]
+
+            if pip_requirements_override and self._env_manager == em.CONDA:
+                # Conda use = instead of == for version pinning
+                pip_requirements_override = [
+                    pip_req.replace("==", "=") for pip_req in pip_requirements_override
+                ]
+
+            environment = self.prepare_env(
+                local_path,
+                pip_requirements_override=pip_requirements_override,
+                extra_envs=extra_envs,
+            )
+
+            try:
+                environment.execute(" ".join(predict_cmd))
+            except ShellCommandException as e:
+                raise MlflowException(
+                    f"{e}\n\nAn exception occurred while running model prediction within a "
+                    f"{self._env_manager} environment. You can find the error message "
+                    f"from the prediction subprocess by scrolling above."
+                ) from None
+        else:
+            if pip_requirements_override:
+                raise MlflowException(
+                    "`pip_requirements_override` is not supported for local env manager."
+                    "Please use conda or virtualenv instead."
+                )
+            scoring_server._predict(local_uri, input_path, output_path, content_type)
+
+    def serve(
+        self,
+        model_uri,
+        port,
+        host,
+        timeout,
+        enable_mlserver,
+        synchronous=True,
+        stdout=None,
+        stderr=None,
+        model_config=None,
+    ):
+        """
+        Serve pyfunc model locally.
+        """
+        local_path = _download_artifact_from_uri(model_uri)
+
+        server_implementation = mlserver if enable_mlserver else scoring_server
+        command, command_env = server_implementation.get_cmd(
+            local_path, port, host, timeout, self._nworkers
+        )
+        _set_mlflow_config_env(command_env, model_config)
+
+        if sys.platform.startswith("linux"):
+
+            def setup_sigterm_on_parent_death():
+                """
+                Uses prctl to automatically send SIGTERM to the command process when its parent is
+                dead.
+
+                This handles the case when the parent is a PySpark worker process.
+                If a user cancels the PySpark job, the worker process gets killed, regardless of
+                PySpark daemon and worker reuse settings.
+                We use prctl to ensure the command process receives SIGTERM after spark job
+                cancellation.
+                The command process itself should handle SIGTERM properly.
+                This is a no-op on macOS because prctl is not supported.
+
+                Note:
+                When a pyspark job canceled, the UDF python process are killed by signal "SIGKILL",
+                This case neither "atexit" nor signal handler can capture SIGKILL signal.
+                prctl is the only way to capture SIGKILL signal.
+                """
+                try:
+                    libc = ctypes.CDLL("libc.so.6")
+                    # Set the parent process death signal of the command process to SIGTERM.
+                    libc.prctl(1, signal.SIGTERM)  # PR_SET_PDEATHSIG, see prctl.h
+                except OSError as e:
+                    # TODO: find approach for supporting MacOS/Windows system which does
+                    #  not support prctl.
+                    warnings.warn(f"Setup libc.prctl PR_SET_PDEATHSIG failed, error {e!r}.")
+
+        else:
+            setup_sigterm_on_parent_death = None
+
+        if not is_windows():
+            # Add "exec" before the starting scoring server command, so that the scoring server
+            # process replaces the bash process, otherwise the scoring server process is created
+            # as a child process of the bash process.
+            # Note we in `mlflow.pyfunc.spark_udf`, use prctl PR_SET_PDEATHSIG to ensure scoring
+            # server process being killed when UDF process exit. The PR_SET_PDEATHSIG can only
+            # send signal to the bash process, if the scoring server process is created as a
+            # child process of the bash process, then it cannot receive the signal sent by prctl.
+            # TODO: For Windows, there's no equivalent things of Unix shell's exec. Windows also
+            #  does not support prctl. We need to find an approach to address it.
+            command = "exec " + command
+
+        if self._env_manager != em.LOCAL:
+            return self.prepare_env(local_path).execute(
+                command,
+                command_env,
+                stdout=stdout,
+                stderr=stderr,
+                preexec_fn=setup_sigterm_on_parent_death,
+                synchronous=synchronous,
+            )
+        else:
+            _logger.info("=== Running command '%s'", command)
+
+            if not is_windows():
+                command = ["bash", "-c", command]
+
+            child_proc = subprocess.Popen(
+                command,
+                env=command_env,
+                preexec_fn=setup_sigterm_on_parent_death,
+                stdout=stdout,
+                stderr=stderr,
+            )
+
+            if synchronous:
+                rc = child_proc.wait()
+                if rc != 0:
+                    raise Exception(
+                        f"Command '{command}' returned non zero return code. Return code = {rc}"
+                    )
+                return 0
+            else:
+                return child_proc
+
+    def serve_stdin(
+        self,
+        model_uri,
+        stdout=None,
+        stderr=None,
+        model_config=None,
+    ):
+        local_path = _download_artifact_from_uri(model_uri)
+
+        command_env = os.environ.copy()
+        _set_mlflow_config_env(command_env, model_config)
+
+        return self.prepare_env(local_path).execute(
+            command=f"python {_STDIN_SERVER_SCRIPT} --model-uri {local_path}",
+            command_env=command_env,
+            stdin=subprocess.PIPE,
+            stdout=stdout,
+            stderr=stderr,
+            synchronous=False,
+        )
+
+    def can_score_model(self):
+        if self._env_manager == em.LOCAL:
+            # noconda => already in python and dependencies are assumed to be installed.
+            return True
+        conda_path = get_conda_bin_executable("conda")
+        try:
+            p = subprocess.Popen(
+                [conda_path, "--version"], stdout=subprocess.PIPE, stderr=subprocess.PIPE
+            )
+            _, _ = p.communicate()
+            return p.wait() == 0
+        except FileNotFoundError:
+            # Can not find conda
+            return False
+
+    def build_image(
+        self,
+        model_uri,
+        image_name,
+        install_java=False,
+        install_mlflow=False,
+        mlflow_home=None,
+        enable_mlserver=False,
+        base_image=None,
+    ):
+        with TempDir() as tmp:
+            cwd = tmp.path()
+            self.generate_dockerfile(
+                model_uri=model_uri,
+                output_dir=cwd,
+                install_java=install_java,
+                install_mlflow=install_mlflow,
+                mlflow_home=mlflow_home,
+                enable_mlserver=enable_mlserver,
+                base_image=base_image,
+            )
+
+            _logger.info("Building docker image with name %s", image_name)
+            docker_utils.build_image_from_context(context_dir=cwd, image_name=image_name)
+
+    def generate_dockerfile(
+        self,
+        model_uri,
+        output_dir,
+        install_java=False,
+        install_mlflow=False,
+        mlflow_home=None,
+        enable_mlserver=False,
+        base_image=None,
+    ):
+        os.makedirs(output_dir, exist_ok=True)
+        _logger.debug("Created all folders in path", extra={"output_directory": output_dir})
+
+        if model_uri:
+            model_cwd = os.path.join(output_dir, _MODEL_DIR_NAME)
+            pathlib.Path(model_cwd).mkdir(parents=True, exist_ok=True)
+            model_path = _download_artifact_from_uri(model_uri, output_path=model_cwd)
+            base_image = base_image or self._get_base_image(model_path, install_java)
+            env_manager = self._env_manager or em.LOCAL
+
+            if base_image.startswith("python"):
+                # we can directly use local env for python image
+                if env_manager in [em.CONDA, em.VIRTUALENV]:
+                    # we can directly use ubuntu image for conda and virtualenv
+                    base_image = UBUNTU_BASE_IMAGE
+            elif base_image == UBUNTU_BASE_IMAGE:
+                env_manager = self._env_manager or em.VIRTUALENV
+                # installing python on ubuntu image is problematic and not recommended officially
+                # , so we recommend using conda or virtualenv instead on ubuntu image
+                if env_manager == em.LOCAL:
+                    raise MlflowException.invalid_parameter_value(LOCAL_ENV_MANAGER_ERROR_MESSAGE)
+
+            copy_src = os.path.relpath(model_path, start=output_dir)
+            model_install_steps = self._model_installation_steps(
+                copy_src,
+                model_path,
+                env_manager,
+                install_mlflow,
+                enable_mlserver,
+            )
+            entrypoint = f"from mlflow.models import container as C; C._serve('{env_manager}')"
+
+        # if no model_uri specified, user must use virtualenv or conda env based on ubuntu image
+        else:
+            base_image = base_image or UBUNTU_BASE_IMAGE
+            env_manager = self._env_manager or em.VIRTUALENV
+            if env_manager == em.LOCAL:
+                raise MlflowException.invalid_parameter_value(LOCAL_ENV_MANAGER_ERROR_MESSAGE)
+
+            model_install_steps = ""
+            # If model_uri is not specified, dependencies are installed at runtime
+            entrypoint = (
+                self._get_install_pyfunc_deps_cmd(env_manager, install_mlflow, enable_mlserver)
+                + f" C._serve('{env_manager}')"
+            )
+
+        dockerfile_text = docker_utils.generate_dockerfile(
+            output_dir=output_dir,
+            base_image=base_image,
+            model_install_steps=model_install_steps,
+            entrypoint=entrypoint,
+            env_manager=env_manager,
+            mlflow_home=mlflow_home,
+            enable_mlserver=enable_mlserver,
+            # always disable env creation at runtime for pyfunc
+            disable_env_creation_at_runtime=True,
+        )
+        _logger.debug("generated dockerfile at {output_dir}", extra={"dockerfile": dockerfile_text})
+
+    def _get_base_image(self, model_path: str, install_java: bool) -> str:
+        """
+        Determine the base image to use for the Dockerfile.
+
+        We use Python slim base image when all the following conditions are met:
+          1. Model URI is specified by the user
+          2. Model flavor does not require Java
+          3. Python version is specified in the model
+
+        Returns:
+            Either the Ubuntu base image or the Python slim base image.
+        """
+        # Check if the model requires Java
+        if not install_java:
+            flavors = _get_all_flavor_configurations(model_path).keys()
+            if java_flavors := JAVA_FLAVORS & flavors:
+                _logger.info(f"Detected java flavors {java_flavors}, installing Java in the image")
+                install_java = True
+
+        # Use ubuntu base image if Java is required
+        if install_java:
+            return UBUNTU_BASE_IMAGE
+
+        # Get Python version from MLmodel
+        try:
+            env_conf = Model.load(model_path).flavors[pyfunc.FLAVOR_NAME][pyfunc.ENV][em.VIRTUALENV]
+            python_env_config_path = os.path.join(model_path, env_conf)
+
+            python_env = _PythonEnv.from_yaml(python_env_config_path)
+            return PYTHON_SLIM_BASE_IMAGE.format(version=python_env.python)
+        except Exception as e:
+            _logger.warning(
+                f"Failed to determine Python version from {model_path}. "
+                f"Defaulting to {UBUNTU_BASE_IMAGE}. Error: {e}"
+            )
+            return UBUNTU_BASE_IMAGE
+
+    def _model_installation_steps(
+        self, copy_src, model_path, env_manager, install_mlflow, enable_mlserver
+    ):
+        # Copy model to image if model_uri is specified
+        steps = (
+            "# Copy model to image and install dependencies\n"
+            f"COPY {copy_src} /opt/ml/model\nRUN python -c "
+        )
+        steps += (
+            f'"{self._get_install_pyfunc_deps_cmd(env_manager, install_mlflow, enable_mlserver)}"'
+        )
+
+        # Install flavor-specific dependencies if needed
+        flavors = _get_all_flavor_configurations(model_path).keys()
+        for flavor in flavors:
+            if flavor in FLAVOR_SPECIFIC_APT_PACKAGES:
+                packages = " ".join(FLAVOR_SPECIFIC_APT_PACKAGES[flavor])
+                steps += f"\nRUN apt-get install -y --no-install-recommends {packages}"
+
+        return steps
+
+    def _get_install_pyfunc_deps_cmd(
+        self, env_manager: str, install_mlflow: bool, enable_mlserver: bool
+    ):
+        return (
+            "from mlflow.models import container as C; "
+            f"C._install_pyfunc_deps('/opt/ml/model', install_mlflow={install_mlflow}, "
+            f"enable_mlserver={enable_mlserver}, env_manager='{env_manager}');"
+        )

mlflow/pyfunc/context.py

@@ -0,0 +1,78 @@
+import contextlib
+from contextvars import ContextVar
+from dataclasses import dataclass
+from typing import Any, Optional
+
+# A thread local variable to store the context of the current prediction request.
+# This is particularly used to associate logs/traces with a specific prediction request in the
+# caller side. The context variable is intended to be set by the called before invoking the
+# predict method, using the set_prediction_context context manager.
+_PREDICTION_REQUEST_CTX = ContextVar("mlflow_prediction_request_context", default=None)
+
+
+@dataclass
+class Context:
+    # A unique identifier for the current prediction request.
+    request_id: Optional[str] = None
+    # Whether the current prediction request is as a part of MLflow model evaluation.
+    is_evaluate: bool = False
+    # The schema of the dependencies to be added into the tag of trace info.
+    dependencies_schemas: Optional[dict[str, Any]] = None
+    # The logged model ID associated with the current prediction request
+    model_id: Optional[str] = None
+    # The model serving endpoint name where the prediction request is made
+    endpoint_name: Optional[str] = None
+
+    def __init__(
+        self,
+        request_id: Optional[str] = None,
+        is_evaluate: bool = False,
+        dependencies_schemas: Optional[dict[str, Any]] = None,
+        model_id: Optional[str] = None,
+        endpoint_name: Optional[str] = None,
+        # Accept extra kwargs so we don't need to worry backward compatibility
+        # when adding new attributes to the Context class
+        **kwargs,
+    ):
+        self.request_id = request_id
+        self.is_evaluate = is_evaluate
+        self.dependencies_schemas = dependencies_schemas
+        self.model_id = model_id
+        self.endpoint_name = endpoint_name
+
+    def update(self, **kwargs):
+        for key, value in kwargs.items():
+            if hasattr(self, key):
+                setattr(self, key, value)
+            else:
+                raise AttributeError(f"Context has no attribute named '{key}'")
+
+
+@contextlib.contextmanager
+def set_prediction_context(context: Optional[Context]):
+    """
+    Set the context for the current prediction request. The context will be set as a thread-local
+    variable and will be accessible globally within the same thread.
+
+    Args:
+        context: The context for the current prediction request.
+    """
+    if context and not isinstance(context, Context):
+        raise TypeError(f"Expected context to be an instance of Context, but got: {context}")
+
+    token = _PREDICTION_REQUEST_CTX.set(context)
+    try:
+        yield
+    finally:
+        _PREDICTION_REQUEST_CTX.reset(token)
+
+
+def get_prediction_context() -> Optional[Context]:
+    """
+    Get the context for the current prediction request. The context is thread-local and is set
+    using the set_prediction_context context manager.
+
+    Returns:
+        The context for the current prediction request, or None if no context is set.
+    """
+    return _PREDICTION_REQUEST_CTX.get()