wehub-resource-sync
146 lines
4.0 KiB
Python
146 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
|
|
|
|
_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_java=False,
|
|
install_mlflow=False,
|
|
mlflow_home=None,
|
|
base_image=None,
|
|
):
|
|
pass
|
|
|
|
def generate_dockerfile(
|
|
self,
|
|
model_uri,
|
|
output_dir,
|
|
install_java=False,
|
|
install_mlflow=False,
|
|
mlflow_home=None,
|
|
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(
|
|
_r_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,
|
|
synchronous=True,
|
|
stdout=None,
|
|
stderr=None,
|
|
):
|
|
"""
|
|
Generate R model locally.
|
|
"""
|
|
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(
|
|
_r_quote(model_path), port, _r_quote(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 _r_quote(str(s))
|
|
|
|
|
|
def _r_quote(s: str) -> str:
|
|
# The `command` is passed directly as an argv element to `Rscript -e`, so
|
|
# the string is parsed by R, not the shell. Escape backslashes and single
|
|
# quotes so the result is a safe R single-quoted string literal.
|
|
escaped = s.replace("\\", "\\\\").replace("'", "\\'")
|
|
return f"'{escaped}'"
|