Files
zenml/venv/lib/python3.9/site-packages/mlflow/rfunc/backend.py
Christian Mantha 2ca0b9ef7c star
2026-03-02 19:10:52 -05:00

135 lines
4.0 KiB
Python

import logging
import os
import re
import subprocess
import sys
from mlflow.exceptions import MlflowException
from mlflow.models import FlavorBackend
from mlflow.tracking.artifact_utils import _download_artifact_from_uri
from mlflow.utils.string_utils import quote
_logger = logging.getLogger(__name__)
class RFuncBackend(FlavorBackend):
"""
Flavor backend implementation for the generic R models.
Predict and serve locally models with 'crate' flavor.
"""
def build_image(
self, model_uri, image_name, install_mlflow, mlflow_home, enable_mlserver, base_image=None
):
pass
def generate_dockerfile(
self, model_uri, output_path, install_mlflow, mlflow_home, enable_mlserver, base_image=None
):
pass
version_pattern = re.compile(r"version ([0-9]+\.[0-9]+\.[0-9]+)")
def predict(
self,
model_uri,
input_path,
output_path,
content_type,
pip_requirements_override=None,
extra_envs=None,
):
"""
Generate predictions using R model saved with MLflow.
Return the prediction results as a JSON.
"""
if pip_requirements_override is not None:
raise MlflowException("pip_requirements_override is not supported in the R backend.")
model_path = _download_artifact_from_uri(model_uri)
str_cmd = (
"mlflow:::mlflow_rfunc_predict(model_path = '{0}', input_path = {1}, "
"output_path = {2}, content_type = {3})"
)
command = str_cmd.format(
quote(model_path),
_str_optional(input_path),
_str_optional(output_path),
_str_optional(content_type),
)
_execute(command, extra_envs=extra_envs)
def serve(
self,
model_uri,
port,
host,
timeout,
enable_mlserver,
synchronous=True,
stdout=None,
stderr=None,
):
"""
Generate R model locally.
NOTE: The `enable_mlserver` parameter is there to comply with the
FlavorBackend interface but is not supported by MLServer yet.
https://github.com/SeldonIO/MLServer/issues/183
"""
if enable_mlserver:
raise Exception("The MLServer inference server is not yet supported in the R backend.")
if timeout:
_logger.warning("Timeout is not yet supported in the R backend.")
if not synchronous:
raise Exception("RBackend does not support call with synchronous=False")
if stdout is not None or stderr is not None:
raise Exception("RBackend does not support redirect stdout/stderr.")
model_path = _download_artifact_from_uri(model_uri)
command = "mlflow::mlflow_rfunc_serve('{}', port = {}, host = '{}')".format(
quote(model_path), port, host
)
_execute(command)
def can_score_model(self):
# `Rscript --version` writes to stderr in R < 4.2.0 but stdout in R >= 4.2.0.
process = subprocess.Popen(
["Rscript", "--version"],
close_fds=True,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
)
stdout, _ = process.communicate()
if process.wait() != 0:
return False
version = self.version_pattern.search(stdout.decode("utf-8"))
if not version:
return False
version = [int(x) for x in version.group(1).split(".")]
return version[0] > 3 or version[0] == 3 and version[1] >= 3
def _execute(command, extra_envs=None):
env = os.environ.copy()
if extra_envs:
env.update(extra_envs)
process = subprocess.Popen(
["Rscript", "-e", command],
env=env,
close_fds=False,
stdin=sys.stdin,
stdout=sys.stdout,
stderr=sys.stderr,
)
if process.wait() != 0:
raise Exception("Command returned non zero exit code.")
def _str_optional(s):
return "NULL" if s is None else f"'{quote(str(s))}'"