CODEWITHSUNDEEP
pythonamazon-web-servicesamazon-sagemakermlflow
Tiago Bachiega de Almeida
Tiago Bachiega de Almeida

Reputation: 121

Tracking SageMaker Estimator with MLFlow

I'm working on a version tracking system for a ML project and want to use MLflow to do so. My project uses AWS Sagemaker's DeepAR for forecast.

What I want to do is very simple. I'm trying do log the Sagemaker DeepAR model (Sagemaker Estimator) with MLFlow. As it doesn't have a "log_model" funcion in it's "mlflow.sagemaker" module, I tried to use the "mlflow.pyfunc" module to do the log. Unfortunatelly it didn't worked. How can I log the Sagemaker model and get the cloudpickle and yaml files generated by MLFlow?

My code for now:

mlflow.pyfunc.log_model(model)

Where model is a sagemaker.estimator.Estimator object and the error I get from the code is

mlflow.exceptions.MlflowException: Either `loader_module` or `python_model` must be specified. A `loader_module` should be a python module. A `python_model` should be a subclass of PythonModel

I know AWS Sagemaker logs my models, but it is really important to my project to do the log with MLFlow too.

Upvotes: 0

Views: 945

Answers (1)

Vlad Tokarev
Vlad Tokarev

Reputation: 78

You cannot use pyfunc to store Any type object.

You should either specify one of loader_module as shown in the example below or you must write the wrapper that implements PythonModel interface and provides logic to deserialize your model from previously-stored artifacts as described here https://www.mlflow.org/docs/latest/models.html#example-saving-an-xgboost-model-in-mlflow-format

example with loader:

    model_uri = 'model.pkl'

    with open(model_uri, 'wb') as f:
        pickle.dump(model, f)

    mlflow.log_artifact(model_uri, 'model')

    mlflow.pyfunc.log_model(
        'model', loader_module='mlflow.sklearn', data_path='model.pkl', code_path=['src'], conda_env='environment.yml'
    )

I think PythonModel is the better way for you because of mlflow doesn't have a built-in loader for SageMaker DeepAR model.

Nonetheless, You must have the knowledge how to restore SageMaker model from artifacts, because I am not sure that is possible at all, cuz of some built-in SageMaker algorithms are blackboxes.

You can also may be interested in container that allow you to run any MLFlow projects inside Sagemaker: https://github.com/odahu/sagemaker-mlflow-container

Upvotes: 1

Related Questions