I'm currently trying to create a model serving end point around a model I've recently created. I'm trying to wrap my head around an error. The model is defined as below
class MY-MODEL(mlflow.pyfunc.PythonModel):
def load_context(self, context):
load()
def predict(self, model_input: pd.DataFrame) -> pd.DataFrame:
return logic()
mlflow.models.set_model(MY-MODEL())
In a notebook I'm able to create an instance, run mlflow.pyfunc.log_model and
mlflow.register_model, aside from anonymizing it looks like below
with mlflow.start_run() as run:
model_info = mlflow.pyfunc.log_model(
name='MY-MODEL',
python_model=MY-MODEL(),
signature=signature,
input_example=model_input,
model_config={
'registered_model_name': REGISTERED_MODEL_NAME
},
artifacts={
'CUSTOM_CONFIG': context.artifacts['CUSTOM_CONFIG'],
'CLUSTER_CONFIG': context.artifacts['CLUSTER_CONFIG'],
},
code_paths=["/Workspace/code/"] <- MY-MODEL.py is here
)
registered_model_info = mlflow.register_model(
model_uri=f'runs:/{run.info.run_id}/MY-MODEL',
name=REGISTERED_MODEL_NAME
)
The error comes when I try to serve the model. In the startup logs I see
[vfk87] [2025-12-30 21:21:56 +0000] [32] [INFO] Starting gunicorn 23.0.0
[vfk87] [2025-12-30 21:21:56 +0000] [32] [INFO] Listening at: http://0.0.0.0:8080 (32)
[vfk87] [2025-12-30 21:21:56 +0000] [32] [INFO] Using worker: gevent
[vfk87] [2025-12-30 21:21:56 +0000] [33] [INFO] Booting worker with pid: 33
[vfk87] [2025-12-30 21:21:59 +0000] An error occurred while loading the model: No module named 'MY-MODEL'
[vfk87] [2025-12-30 21:21:59 +0000] Traceback (most recent call last):
[vfk87] [2025-12-30 21:21:59 +0000] File "/opt/conda/envs/mlflow-env/lib/python3.12/site-packages/mlflowserving/scoring_server/__init__.py", line 233, in get_model_option_or_exit
[vfk87] [2025-12-30 21:21:59 +0000] self.model = self.model_future.result()
[vfk87] [2025-12-30 21:21:59 +0000] ^^^^^^^^^^^^^^^^^^^^^^^^^^
[vfk87] [2025-12-30 21:21:59 +0000] File "/opt/conda/envs/mlflow-env/lib/python3.12/concurrent/futures/_base.py", line 449, in result
[vfk87] [2025-12-30 21:21:59 +0000] return self.__get_result()
[vfk87] [2025-12-30 21:21:59 +0000] ^^^^^^^^^^^^^^^^^^^
[vfk87] [2025-12-30 21:21:59 +0000] File "/opt/conda/envs/mlflow-env/lib/python3.12/concurrent/futures/_base.py", line 401, in __get_result
[vfk87] [2025-12-30 21:21:59 +0000] raise self._exception
[vfk87] [2025-12-30 21:21:59 +0000] File "/opt/conda/envs/mlflow-env/lib/python3.12/concurrent/futures/thread.py", line 58, in run
[vfk87] [2025-12-30 21:21:59 +0000] result = self.fn(*self.args, **self.kwargs)
[vfk87] [2025-12-30 21:21:59 +0000] ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[vfk87] [2025-12-30 21:21:59 +0000] File "/opt/conda/envs/mlflow-env/lib/python3.12/site-packages/mlflowserving/scoring_server/__init__.py", line 143, in _load_model_closure
[vfk87] [2025-12-30 21:21:59 +0000] model = load_model_fn(path)
[vfk87] [2025-12-30 21:21:59 +0000] ^^^^^^^^^^^^^^^^^^^
[vfk87] [2025-12-30 21:21:59 +0000] File "/opt/conda/envs/mlflow-env/lib/python3.12/site-packages/mlflow/tracing/provider.py", line 435, in wrapper
[vfk87] [2025-12-30 21:21:59 +0000] is_func_called, result = True, f(*args, **kwargs)
[vfk87] [2025-12-30 21:21:59 +0000] ^^^^^^^^^^^^^^^^^^
[vfk87] [2025-12-30 21:21:59 +0000] File "/opt/conda/envs/mlflow-env/lib/python3.12/site-packages/mlflow/pyfunc/__init__.py", line 1183, in load_model
[vfk87] [2025-12-30 21:21:59 +0000] raise e
[vfk87] [2025-12-30 21:21:59 +0000] File "/opt/conda/envs/mlflow-env/lib/python3.12/site-packages/mlflow/pyfunc/__init__.py", line 1166, in load_model
[vfk87] [2025-12-30 21:21:59 +0000] model_impl = importlib.import_module(conf[MAIN])._load_pyfunc(data_path, model_config)
[vfk87] [2025-12-30 21:21:59 +0000] ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[vfk87] [2025-12-30 21:21:59 +0000] File "/opt/conda/envs/mlflow-env/lib/python3.12/site-packages/mlflow/pyfunc/model.py", line 1243, in _load_pyfunc
[vfk87] [2025-12-30 21:21:59 +0000] context, python_model, signature = _load_context_model_and_signature(model_path, model_config)
[vfk87] [2025-12-30 21:21:59 +0000] ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[vfk87] [2025-12-30 21:21:59 +0000] File "/opt/conda/envs/mlflow-env/lib/python3.12/site-packages/mlflow/pyfunc/model.py", line 1224, in _load_context_model_and_signature
[vfk87] [2025-12-30 21:21:59 +0000] python_model = _maybe_decompress_cloudpickle_load(
[vfk87] [2025-12-30 21:21:59 +0000] ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[vfk87] [2025-12-30 21:21:59 +0000] File "/opt/conda/envs/mlflow-env/lib/python3.12/site-packages/mlflow/pyfunc/model.py", line 824, in _maybe_decompress_cloudpickle_load
[vfk87] [2025-12-30 21:21:59 +0000] return cloudpickle.load(f)
[vfk87] [2025-12-30 21:21:59 +0000] ^^^^^^^^^^^^^^^^^^^
[vfk87] [2025-12-30 21:21:59 +0000] ModuleNotFoundError: No module named 'MY-MODEL'
I suspect the issue is around mlflow and how I'm defining/registering the python class it but I'm not seeing anything in documentation that would indicate that. Any guidance or a point in the right direction would be more than welcome as I'm running out of ideas.
