[Python] Exclude Private Package Dependency When Logging Model with MLflow in Databricks

I'm deploying a Databricks workflow using Databricks Asset Bundles (DABs). I have a custom package (mypackage) installed on my cluster, which includes multiple CLI commands for different workflows. This package contains all the dependencies required across my workflows. Note that mypackage is a private PyPI repository and isn't accessible externally.

When I use mlflow.pyfunc.log_model() to log my model, MLflow automatically detects and includes mypackage.submodule as a dependency for the model. This causes issues because:

• mypackage is private and cannot be installed from an external environment.
• The model doesn't actually need mypackage or its dependencies for inference—it only needs standard libraries like PyTorch and Pandas.

Here's a simplified version of my code:

import mlflow.pyfunc
from mypackage.submodule import MyModelClass

model = MyModelClass()

mlflow.pyfunc.log_model(
python_model=model,
artifact_path="my_model",
# Other parameters
)


Questions:

• 
  How can I log my model with mlflow.pyfunc.log_model() without including mypackage as a dependency?
  
  
• 
  Is there an MLflow-native way to prevent certain packages from being included in the model's environment?
  
  
• 
  Are there best practices for handling situations where the training environment includes packages that shouldn't be dependencies of the serialized model?
  

Additional Context:

• I'm using Databricks and Databricks Asset Bundles for deployment.
• The model uses standard libraries (e.g., PyTorch, Pandas) needed for inference.
• mypackage is only used during training and workflow orchestration but is not needed for model inference.

What I've tried:

As a workaround, I attempted a hack where I redefine my model classes within the __main__ module and serialize them using cloudpickle, essentially removing references to mypackage. Here's a snippet of what I did:

import inspect
import cloudpickle

# Extract source code of the model classes
model_source = inspect.getsource(MyModelClass)

# Create a main namespace and redefine classes in __main__
main_namespace = {
'__name__': '__main__',
# Include other necessary imports
}
exec(model_source, main_namespace)

# Access the redefined class
MyModelClassMain = main_namespace['MyModelClass']
model = MyModelClassMain()

# Serialize the model
with open("model.pkl", "wb") as f:
cloudpickle.dump(model, f)

# Log the model without including 'mypackage'
mlflow.pyfunc.log_model(
python_model=model,
artifact_path="my_model",
artifacts={"model.pkl": "model.pkl"},
pip_requirements=["torch", "pandas", "cloudpickle"],
# Other parameters
)


While this approach works, it feels like a hacky solution.

Continue reading...