def main():
args = parse_args()
model_path = args.model_path
flavor = args.flavor
# Mirror `sys.path` of the parent process
sys.path = json.loads(args.sys_path)
cap_cm = _CaptureImportedModules()
# If `model_path` refers to an MLflow model directory, load the model using
# `mlflow.pyfunc.load_model`
if os.path.isdir(model_path) and MLMODEL_FILE_NAME in os.listdir(
model_path):
pyfunc_conf = Model.load(model_path).flavors.get(
mlflow.pyfunc.FLAVOR_NAME)
loader_module = importlib.import_module(pyfunc_conf[MAIN])
original = loader_module._load_pyfunc
@functools.wraps(original)
def _load_pyfunc_patch(*args, **kwargs):
with cap_cm:
return original(*args, **kwargs)
loader_module._load_pyfunc = _load_pyfunc_patch
mlflow.pyfunc.load_model(model_path)
# Otherwise, load the model using `mlflow.<flavor>._load_pyfunc`. For models that don't contain
# pyfunc flavor (e.g. scikit-learn estimator that doesn't implement a `predict` method),
# we need to directly pass a model data path to this script.
else:
with cap_cm:
importlib.import_module(f"mlflow.{flavor}")._load_pyfunc(
model_path)
# Store the imported modules in `output_file`
write_to(args.output_file, "\n".join(cap_cm.imported_modules))
def main():
args = parse_args()
model_path = args.model_path
flavor = args.flavor
# Mirror `sys.path` of the parent process
sys.path = json.loads(args.sys_path)
if flavor == mlflow.spark.FLAVOR_NAME and is_in_databricks_runtime():
# Clear 'PYSPARK_GATEWAY_PORT' and 'PYSPARK_GATEWAY_SECRET' to enforce launching a new JVM
# gateway before calling `mlflow.spark._load_pyfunc` that creates a new spark session
# if it doesn't exist.
os.environ.pop("PYSPARK_GATEWAY_PORT", None)
os.environ.pop("PYSPARK_GATEWAY_SECRET", None)
os.environ["SPARK_DIST_CLASSPATH"] = "/databricks/jars/*"
cap_cm = _CaptureImportedModules()
# If `model_path` refers to an MLflow model directory, load the model using
# `mlflow.pyfunc.load_model`
if os.path.isdir(model_path) and MLMODEL_FILE_NAME in os.listdir(
model_path):
pyfunc_conf = Model.load(model_path).flavors.get(
mlflow.pyfunc.FLAVOR_NAME)
loader_module = importlib.import_module(pyfunc_conf[MAIN])
original = loader_module._load_pyfunc
@functools.wraps(original)
def _load_pyfunc_patch(*args, **kwargs):
with cap_cm:
return original(*args, **kwargs)
loader_module._load_pyfunc = _load_pyfunc_patch
mlflow.pyfunc.load_model(model_path)
# Otherwise, load the model using `mlflow.<flavor>._load_pyfunc`. For models that don't contain
# pyfunc flavor (e.g. scikit-learn estimator that doesn't implement a `predict` method),
# we need to directly pass a model data path to this script.
else:
with cap_cm:
importlib.import_module(f"mlflow.{flavor}")._load_pyfunc(
model_path)
# Store the imported modules in `output_file`
write_to(args.output_file, "\n".join(cap_cm.imported_modules))
# Clean up a spark session created by `mlflow.spark._load_pyfunc`
if flavor == mlflow.spark.FLAVOR_NAME:
from pyspark.sql import SparkSession
spark = SparkSession._instantiatedSession
if spark:
try:
spark.stop()
except Exception:
# Swallow unexpected exceptions
pass
def main():
args = parse_args()
model_path = args.model_path
flavor = args.flavor
# Mirror `sys.path` of the parent process
sys.path = json.loads(args.sys_path)
if flavor == mlflow.spark.FLAVOR_NAME and is_in_databricks_runtime():
try:
# pylint: disable=import-error
from dbruntime.spark_connection import initialize_spark_connection
initialize_spark_connection()
except Exception as e:
raise Exception(
"Attempted to initialize a spark session to load the spark model, but failed"
) from e
cap_cm = _CaptureImportedModules()
# If `model_path` refers to an MLflow model directory, load the model using
# `mlflow.pyfunc.load_model`
if os.path.isdir(model_path) and MLMODEL_FILE_NAME in os.listdir(
model_path):
pyfunc_conf = Model.load(model_path).flavors.get(
mlflow.pyfunc.FLAVOR_NAME)
loader_module = importlib.import_module(pyfunc_conf[MAIN])
original = loader_module._load_pyfunc
@functools.wraps(original)
def _load_pyfunc_patch(*args, **kwargs):
with cap_cm:
return original(*args, **kwargs)
loader_module._load_pyfunc = _load_pyfunc_patch
mlflow.pyfunc.load_model(model_path)
# Otherwise, load the model using `mlflow.<flavor>._load_pyfunc`. For models that don't contain
# pyfunc flavor (e.g. scikit-learn estimator that doesn't implement a `predict` method),
# we need to directly pass a model data path to this script.
else:
with cap_cm:
importlib.import_module(f"mlflow.{flavor}")._load_pyfunc(
model_path)
# Store the imported modules in `output_file`
write_to(args.output_file, "\n".join(cap_cm.imported_modules))
def build_image_local_from_model_uri(self,
model_uri,
base_image,
mlflow_home=None,
**kwargs):
"""build PythonModel Backed service image from model_uri
:param base_image: image base from which build model image
:param mlflow_home: mllfow local copy used to startup the model service in container
if None install from pip.
:param model_uri: directory contains pyfunc model filesystem.
<"pyfunc-filename-system"
https://mlflow.org/docs/latest/python_api/mlflow.pyfunc.html#pyfunc-filename-system>_
"""
with tempfile.TemporaryDirectory() as tmp_dir:
if ModelsArtifactRepository.is_models_uri(model_uri):
underlying_model_uri = ModelsArtifactRepository.get_underlying_uri(
model_uri)
else:
underlying_model_uri = model_uri
local_path = _download_artifact_from_uri(append_to_uri_path(
underlying_model_uri, MLMODEL_FILE_NAME),
output_path=tmp_dir)
model_meta = Model.load(local_path)
flavor_name, flavor_backend = get_flavor_backend(
model_meta, **kwargs)
if flavor_name is None:
raise TypeError("no suitable backend was found for the model")
if not flavor_backend.can_build_image():
raise AttributeError(
'flavor {} not support build image'.format(flavor_name))
# always intall mlflow for override office mlflow package in container
return_code = flavor_backend.build_image(model_uri,
self.image_name,
install_mlflow=True,
mlflow_home=mlflow_home,
base_image=base_image)
return True if not return_code else False
def __generate_mar_file(
self, model_name, version, model_file, handler, requirements_file, extra_files, model_uri
):
"""
Generates mar file using the torch archiver in the specified model store path
"""
valid_file_suffixes = [".pt", ".pth"]
extra_files_list = []
req_file_path = None
model_path = None
if not os.path.isfile(model_uri):
path = Path(_download_artifact_from_uri(model_uri))
# Derive model.pth or state_dict.pth file path
for root, dirs, files in os.walk(path, topdown=False):
for name in files:
if Path(name).suffix in valid_file_suffixes:
model_path = os.path.join(root, name)
# For eager mode models, derive extra files, requirement files
if "state_dict.pth" not in model_path:
model_config = path / "MLmodel"
model = Model.load(model_config)
model_json = json.loads(Model.to_json(model))
try:
if model_json["flavors"]["pytorch"]["extra_files"]:
for extra_file in model_json["flavors"]["pytorch"]["extra_files"]:
extra_files_list.append(os.path.join(path, extra_file["path"]))
except KeyError:
pass
try:
if model_json["flavors"]["pytorch"]["requirements_file"]:
req_file_path = os.path.join(
path, model_json["flavors"]["pytorch"]["requirements_file"]["path"]
)
except KeyError:
pass
if model_path is None:
raise RuntimeError(
"Model file does not have a valid suffix. Expected to be one of "
+ ", ".join(valid_file_suffixes)
)
model_uri = model_path
export_path = self.server_config["export_path"]
if export_path:
model_store = export_path
else:
model_store = "model_store"
if not os.path.isdir(model_store):
os.makedirs(model_store)
cmd = (
"torch-model-archiver --force --model-name {} "
"--version {} --serialized-file {} "
"--handler {} --export-path {}".format(
model_name, version, model_uri, handler, model_store
)
)
if model_file:
cmd += " --model-file {file_path}".format(file_path=model_file)
extra_files_str = ""
if extra_files_list:
extra_files_str += ",".join(extra_files_list)
if extra_files:
if extra_files_list:
extra_files_str = "{base_string},{user_defined_string}".format(
base_string=extra_files_str,
user_defined_string=str(extra_files).replace("'", ""),
)
else:
extra_files_str = str(extra_files).replace("'", "")
if extra_files_str:
cmd = "{cmd} --extra-files '{extra_files}'".format(cmd=cmd, extra_files=extra_files_str)
if requirements_file:
cmd = "{cmd} -r {path}".format(cmd=cmd, path=requirements_file)
elif req_file_path:
cmd = "{cmd} -r {path}".format(cmd=cmd, path=req_file_path)
return_code = subprocess.Popen(cmd, shell=True).wait()
if return_code != 0:
_logger.error(
"Error when attempting to load and parse JSON cluster spec from file %s",
cmd,
)
raise Exception("Unable to create mar file")
if export_path:
mar_file = "{}/{}.mar".format(export_path, model_name)
else:
mar_file = "{}/{}/{}.mar".format(os.getcwd(), model_store, model_name)
if os.path.isfile(mar_file):
print("{} file generated successfully".format(mar_file))
return mar_file
def extract_signature(self, mlmodel_file):
self.mlmodel = Model.load(mlmodel_file)
model_json = json.loads(Model.to_json(self.mlmodel))
if "signature" not in model_json.keys():
raise Exception("Model Signature not found")