mlflow: [BUG] TypeError: cannot pickle 'weakref' object

Willingness to contribute

The MLflow Community encourages bug fix contributions. Would you or another member of your organization be willing to contribute a fix for this bug to the MLflow code base?

• Yes. I can contribute a fix for this bug independently.
• Yes. I would be willing to contribute a fix for this bug with guidance from the MLflow community.
• No. I cannot contribute a bug fix at this time.

System information

• Have I written custom code (as opposed to using a stock example script provided in MLflow): Yes
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): MacOS 11.1
• MLflow installed from (source or binary):
• MLflow version (run mlflow --version): 1.13.1
• Python version: 3.8
• npm version, if running the dev UI:
• Exact command to reproduce:

Describe the problem

I’ve defined a basic Keras model and used the custom pyfunc flavor to create a model class and save this model to local file. I’ve tried several configurations but all lead to a “TypeError: cannot pickle ‘weakref’ object”.

Code to reproduce issue

        model = keras.Sequential()
        model.add(Bidirectional(LSTM(128, input_shape=[10,100])))
        model.add(Dropout(rate=0.5))
        model.add(Dense(6, activation='softmax'))
        model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['acc'])
        history = model.fit(x_train, y_train, epochs=64, validation_split=0.2, shuffle=True)

        class NewModel(mlflow.pyfunc.PythonModel):
            def __init__(self, model):
                self.model = model
            def predict(self, context, x_input):
                return self.model.predict(x_input)
        
        mlflow.pyfunc.save_model(
            path = './_model/test_model',
            python_model = NewModel(model)
        )

Other info / logs

Traceback (most recent call last):
  File "/Users/roderickmacintosh/gDrive/MAIN/model_training/src/venv/lib/python3.8/site-packages/mlflow/pyfunc/model.py", line 138, in _save_model_with_class_artifacts_params
    cloudpickle.dump(python_model, out)
  File "/Users/roderickmacintosh/gDrive/MAIN/model_training/src/venv/lib/python3.8/site-packages/cloudpickle/cloudpickle_fast.py", line 55, in dump
    CloudPickler(
  File "/Users/roderickmacintosh/gDrive/MAIN/model_training/src/venv/lib/python3.8/site-packages/cloudpickle/cloudpickle_fast.py", line 563, in dump
    return Pickler.dump(self, obj)
TypeError: cannot pickle 'weakref' object

What component(s), interfaces, languages, and integrations does this bug affect?

Components

• area/artifacts: Artifact stores and artifact logging
• area/build: Build and test infrastructure for MLflow
• area/docs: MLflow documentation pages
• area/examples: Example code
• area/model-registry: Model Registry service, APIs, and the fluent client calls for Model Registry
• area/models: MLmodel format, model serialization/deserialization, flavors
• area/projects: MLproject format, project running backends
• area/scoring: Local serving, model deployment tools, spark UDFs
• area/server-infra: MLflow server, JavaScript dev server
• area/tracking: Tracking Service, tracking client APIs, autologging

Interface

• area/uiux: Front-end, user experience, JavaScript, plotting
• area/docker: Docker use across MLflow’s components, such as MLflow Projects and MLflow Models
• area/sqlalchemy: Use of SQLAlchemy in the Tracking Service or Model Registry
• area/windows: Windows support

Language

• language/r: R APIs and clients
• language/java: Java APIs and clients
• language/new: Proposals for new client languages

Integrations

• integrations/azure: Azure and Azure ML integrations
• integrations/sagemaker: SageMaker integrations
• integrations/databricks: Databricks integrations