Loading model from S3 with mlflow throws AttributeError: 'PassThroughProfiler' object has no attribute '_stage'

Hello, I have used mlflow for logging my experiments, aswell as my models. The model that I am using is pytorch-forecastingTemporalFusionTransformer model trained on kaggle/colab gpus. After my experimentation, I needed to load the model from S3 bucket where I was storing my runs and use my selected model for predictions. However, when I try to load the model when writing the production script I get the error: AttributeError: 'PassThroughProfiler' object has no attribute '_stage'.

My producion script (so far) is this:

import os
import mlflow

MLFLOW_TRACKING_URI='model-uri'

AWS_ACCESS_KEY_ID = 'aws-access-key-id'
AWS_SECRET_ACCESS_KEY = 'aws-secret-acces-key'

MODEL_URI = 's3:path'

os.environ["AWS_ACCESS_KEY_ID"] = AWS_ACCESS_KEY_ID
os.environ["AWS_SECRET_ACCESS_KEY"] = AWS_SECRET_ACCESS_KEY

mlflow.set_tracking_uri(MLFLOW_TRACKING_URI)

print('tmp')
if __name__ == '__main__':

    model = mlflow.pytorch.load_model(MODEL_URI)
    

The full stack trace is this:

Traceback (most recent call last):
  File "/home/petar/Praksa/Production script/production_script.py", line 24, in <module>
    model = mlflow.pytorch.load_model(MODEL_URI)
  File "/home/petar/Praksa/Production script/production_venv/lib/python3.8/site-packages/mlflow/pytorch/__init__.py", line 714, in load_model
    return _load_model(path=torch_model_artifacts_path, **kwargs)
  File "/home/petar/Praksa/Production script/production_venv/lib/python3.8/site-packages/mlflow/pytorch/__init__.py", line 626, in _load_model
    return torch.load(model_path, **kwargs)
  File "/home/petar/Praksa/Production script/production_venv/lib/python3.8/site-packages/torch/serialization.py", line 607, in load
    return _load(opened_zipfile, map_location, pickle_module, **pickle_load_args)
  File "/home/petar/Praksa/Production script/production_venv/lib/python3.8/site-packages/torch/serialization.py", line 882, in _load
    result = unpickler.load()
TypeError: an integer is required (got type bytes)
Exception ignored in: <function Profiler.__del__ at 0x7f98d8c59040>
Traceback (most recent call last):
  File "/home/petar/Praksa/Production script/production_venv/lib/python3.8/site-packages/pytorch_lightning/profilers/profiler.py", line 170, in __del__
    self.teardown(stage=self._stage)
AttributeError: 'PassThroughProfiler' object has no attribute '_stage'

In the stack trace above, it says that the error is thrown in the profiler of pytorch-lightning.
However, if I try to load the same model in a kaggle notebook it runs fine and returns the model.

The dependencies used in the notebook and in my script are the same and are as follows:

-f https://download.pytorch.org/whl/torch_stable.html
pytorch-forecasting==0.9.0
protobuf==3.20.0
torch==1.9.0+cu111
torchvision==0.10.0+cu111 
torchtext==0.10.0
pytorch-lightning==1.7.7
pandas==1.2.5 
numpy==1.23.4
mlflow==1.21.0
scikit-learn==0.24.2
psycopg2-binary==2.9.2
boto3

Any help towards resolving this issue is welcomed.

Solution
The reason why the model could not be loaded was because of the python version incompatibility between the dependencies. The model was trained on pytohn3.7.12, and it’s depencencies in their name have cp37, suggesting that these dependencies are compatible with python3.7. I on the other hand, was using python3.8.10 when trying to load the model, for which I had dependencies with cp38 in their name.

1 Like