wehub-resource-sync
260 lines
10 KiB
Python
260 lines
10 KiB
Python
"""
|
|
This script should be executed in a fresh python interpreter process using `subprocess`.
|
|
"""
|
|
|
|
import argparse
|
|
import builtins
|
|
import functools
|
|
import importlib
|
|
import json
|
|
import os
|
|
import sys
|
|
|
|
import mlflow
|
|
from mlflow.models.model import MLMODEL_FILE_NAME, Model
|
|
from mlflow.pyfunc import MAIN
|
|
from mlflow.utils._spark_utils import _prepare_subprocess_environ_for_creating_local_spark_session
|
|
from mlflow.utils.exception_utils import get_stacktrace
|
|
from mlflow.utils.file_utils import write_to
|
|
from mlflow.utils.requirements_utils import (
|
|
DATABRICKS_MODULES_TO_PACKAGES,
|
|
MLFLOW_MODULES_TO_PACKAGES,
|
|
)
|
|
|
|
|
|
def _get_top_level_module(full_module_name):
|
|
return full_module_name.split(".")[0]
|
|
|
|
|
|
def _get_second_level_module(full_module_name):
|
|
return ".".join(full_module_name.split(".")[:2])
|
|
|
|
|
|
class _CaptureImportedModules:
|
|
"""
|
|
A context manager to capture imported modules by temporarily applying a patch to
|
|
`builtins.__import__` and `importlib.import_module`.
|
|
|
|
If `record_full_module` is set to `False`, it only captures top level modules
|
|
for inferring python package purpose.
|
|
If `record_full_module` is set to `True`, it captures full module name for all
|
|
imported modules and sub-modules. This is used in automatic model code path inference.
|
|
"""
|
|
|
|
def __init__(self, record_full_module=False):
|
|
self.imported_modules = set()
|
|
self.original_import = None
|
|
self.original_import_module = None
|
|
self.record_full_module = record_full_module
|
|
|
|
def _wrap_import(self, original):
|
|
@functools.wraps(original)
|
|
def wrapper(name, globals=None, locals=None, fromlist=(), level=0):
|
|
is_absolute_import = level == 0
|
|
if not self.record_full_module and is_absolute_import:
|
|
self._record_imported_module(name)
|
|
|
|
result = original(name, globals, locals, fromlist, level)
|
|
|
|
if self.record_full_module:
|
|
if is_absolute_import:
|
|
parent_modules = name.split(".")
|
|
else:
|
|
parent_modules = globals["__name__"].split(".")
|
|
if level > 1:
|
|
parent_modules = parent_modules[: -(level - 1)]
|
|
|
|
if fromlist:
|
|
for from_name in fromlist:
|
|
full_modules = parent_modules + [from_name]
|
|
full_module_name = ".".join(full_modules)
|
|
if full_module_name in sys.modules:
|
|
self._record_imported_module(full_module_name)
|
|
else:
|
|
# in the case that `from_name` is a function or a class
|
|
# then record the parent_module
|
|
self._record_imported_module(".".join(parent_modules))
|
|
else:
|
|
full_module_name = ".".join(parent_modules)
|
|
self._record_imported_module(full_module_name)
|
|
|
|
return result
|
|
|
|
return wrapper
|
|
|
|
def _wrap_import_module(self, original):
|
|
@functools.wraps(original)
|
|
def wrapper(name, *args, **kwargs):
|
|
self._record_imported_module(name)
|
|
return original(name, *args, **kwargs)
|
|
|
|
return wrapper
|
|
|
|
def _record_imported_module(self, full_module_name):
|
|
if self.record_full_module:
|
|
self.imported_modules.add(full_module_name)
|
|
return
|
|
|
|
# If the module is an internal module (prefixed by "_") or is the "databricks"
|
|
# module, which is populated by many different packages, don't record it (specific
|
|
# module imports within the databricks namespace are still recorded and mapped to
|
|
# their corresponding packages)
|
|
if full_module_name.startswith("_") or full_module_name == "databricks":
|
|
return
|
|
|
|
top_level_module = _get_top_level_module(full_module_name)
|
|
second_level_module = _get_second_level_module(full_module_name)
|
|
|
|
if top_level_module == "databricks":
|
|
# Multiple packages populate the `databricks` module namespace on Databricks;
|
|
# to avoid bundling extraneous Databricks packages into model dependencies, we
|
|
# scope each module to its relevant package
|
|
if second_level_module in DATABRICKS_MODULES_TO_PACKAGES:
|
|
self.imported_modules.add(second_level_module)
|
|
return
|
|
|
|
for databricks_module in DATABRICKS_MODULES_TO_PACKAGES:
|
|
if full_module_name.startswith(databricks_module):
|
|
self.imported_modules.add(databricks_module)
|
|
return
|
|
|
|
# special casing for mlflow extras since they may not be required by default
|
|
if top_level_module == "mlflow":
|
|
if second_level_module in MLFLOW_MODULES_TO_PACKAGES:
|
|
self.imported_modules.add(second_level_module)
|
|
return
|
|
|
|
self.imported_modules.add(top_level_module)
|
|
|
|
def __enter__(self):
|
|
# Patch `builtins.__import__` and `importlib.import_module`
|
|
self.original_import = builtins.__import__
|
|
self.original_import_module = importlib.import_module
|
|
builtins.__import__ = self._wrap_import(self.original_import)
|
|
importlib.import_module = self._wrap_import_module(self.original_import_module)
|
|
return self
|
|
|
|
def __exit__(self, *_, **__):
|
|
# Revert the patches
|
|
builtins.__import__ = self.original_import
|
|
importlib.import_module = self.original_import_module
|
|
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--model-path", required=True)
|
|
parser.add_argument("--flavor", required=True)
|
|
parser.add_argument("--output-file", required=True)
|
|
parser.add_argument("--sys-path", required=True)
|
|
parser.add_argument("--module-to-throw", required=False)
|
|
parser.add_argument("--error-file", required=False)
|
|
parser.add_argument("--record-full-module", default=False, action="store_true")
|
|
return parser.parse_args()
|
|
|
|
|
|
def store_imported_modules(
|
|
cap_cm, model_path, flavor, output_file, error_file=None, record_full_module=False
|
|
):
|
|
# 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):
|
|
mlflow_model = Model.load(model_path)
|
|
pyfunc_conf = mlflow_model.flavors.get(mlflow.pyfunc.FLAVOR_NAME)
|
|
input_example = mlflow_model.load_input_example(model_path)
|
|
params = mlflow_model.load_input_example_params(model_path)
|
|
|
|
def load_model_and_predict(original_load_fn, *args, **kwargs):
|
|
model = original_load_fn(*args, **kwargs)
|
|
if input_example is not None:
|
|
try:
|
|
model.predict(input_example, params=params)
|
|
except Exception as e:
|
|
if error_file:
|
|
stack_trace = get_stacktrace(e)
|
|
write_to(
|
|
error_file,
|
|
"Failed to run predict on input_example, dependencies "
|
|
"introduced in predict are not captured.\n" + stack_trace,
|
|
)
|
|
else:
|
|
raise e
|
|
return model
|
|
|
|
if record_full_module:
|
|
# Note: if we want to record all imported modules
|
|
# (for inferring code_paths purpose),
|
|
# The `importlib.import_module(pyfunc_conf[MAIN])` invocation
|
|
# must be wrapped with `cap_cm` context manager,
|
|
# because `pyfunc_conf[MAIN]` might also be a module loaded from
|
|
# code_paths.
|
|
with cap_cm:
|
|
# `mlflow.pyfunc.load_model` internally invokes
|
|
# `importlib.import_module(pyfunc_conf[MAIN])`
|
|
mlflow.pyfunc.load_model(model_path)
|
|
else:
|
|
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 load_model_and_predict(original, *args, **kwargs)
|
|
|
|
loader_module._load_pyfunc = _load_pyfunc_patch
|
|
try:
|
|
mlflow.pyfunc.load_model(model_path)
|
|
finally:
|
|
loader_module._load_pyfunc = original
|
|
# 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(output_file, "\n".join(cap_cm.imported_modules))
|
|
|
|
|
|
def main():
|
|
args = parse_args()
|
|
model_path = args.model_path
|
|
flavor = args.flavor
|
|
output_file = args.output_file
|
|
error_file = args.error_file
|
|
# Mirror `sys.path` of the parent process
|
|
sys.path = json.loads(args.sys_path)
|
|
|
|
if flavor == mlflow.spark.FLAVOR_NAME:
|
|
# Create a local spark environment within the subprocess
|
|
from mlflow.utils._spark_utils import _create_local_spark_session_for_loading_spark_model
|
|
|
|
_prepare_subprocess_environ_for_creating_local_spark_session()
|
|
_create_local_spark_session_for_loading_spark_model()
|
|
|
|
cap_cm = _CaptureImportedModules(record_full_module=args.record_full_module)
|
|
store_imported_modules(
|
|
cap_cm,
|
|
model_path,
|
|
flavor,
|
|
output_file,
|
|
error_file,
|
|
record_full_module=args.record_full_module,
|
|
)
|
|
|
|
# Clean up a spark session created by `mlflow.spark._load_pyfunc`
|
|
if flavor == mlflow.spark.FLAVOR_NAME:
|
|
from mlflow.utils._spark_utils import _get_active_spark_session
|
|
|
|
if spark := _get_active_spark_session():
|
|
try:
|
|
spark.stop()
|
|
except Exception:
|
|
# Swallow unexpected exceptions
|
|
pass
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|