Allow writeScoreCode to include nonstandard pickle types

Author: smlindauer

src/sasctl/pzmm/importModel.py

@@ -113,6 +113,7 @@ def pzmmImportModel(
         isH2OModel=False,
         force=False,
         binaryString=None,
+        mlFlowDetails=None
     ):
         """Import model to SAS Model Manager using pzmm submodule.
 
@@ -172,10 +173,16 @@ def pzmmImportModel(
             Sets whether to overwrite models with the same name upon upload. By default False.
         binaryString : string, optional
             Binary string representation of the model object. By default None.
+        mlFlowDetails : dict, optional
+            Model details from an MLFlow model. This dictionary is created by the readMLModelFile function.
+            By default None.       
         """
         # Initialize no score code or binary H2O model flags
         noScoreCode = False
         binaryModel = False
+        
+        if mlFlowDetails is None:
+            mlFlowDetails = {'serialization_format': 'pickle'} 
 
         if pyPath is None:
             pyPath = Path(zPath)
@@ -236,6 +243,7 @@ def getFiles(extensions):
                     isH2OModel=isH2OModel,
                     isBinaryModel=binaryModel,
                     binaryString=binaryString,
+                    pickleType=mlFlowDetails['serialization_format']
                 )
                 print(
                     "Model score code was written successfully to {}.".format(
@@ -299,6 +307,7 @@ def getFiles(extensions):
                     isH2OModel=isH2OModel,
                     isBinaryModel=binaryModel,
                     binaryString=binaryString,
+                    pickleType=mlFlowDetails['serialization_format']
                 )
                 print(
                     "Model score code was written successfully to {} and uploaded to SAS Model Manager".format(

src/sasctl/pzmm/writeScoreCode.py

@@ -28,6 +28,7 @@ def writeScoreCode(
         scoreCAS=True,
         isBinaryModel=False,
         binaryString=None,
+        pickleType=None
     ):
         """
         Writes a Python score code file based on training data used to generate the model
@@ -107,6 +108,8 @@ def writeScoreCode(
             Sets whether the H2O model provided is a binary model or a MOJO model. By default False.
         binaryString : string, optional
             Binary string representation of the model object. By default None.
+        pickleType : string optional
+            Indicator for MLFlow models, which may pickle by non-standard methods. By default None.
         """
         # Check if binary string model
         if binaryString is not None:
@@ -174,10 +177,10 @@ def upload_and_copy_score_resources(model, files):
             cls.pyFile.write(
                 """\n
 import math
-import pickle
+import {pickleType}
 import pandas as pd
 import numpy as np"""
-            )
+            ).format(pickleType=pickleType)
             # In SAS Viya 4.0 and SAS Open Model Manager, a settings.py file is generated that points to the resource
             # location
             if not isViya35:
@@ -219,8 +222,8 @@ def upload_and_copy_score_resources(model, files):
                 cls.pyFile.write(
                     """\n
 with open('/models/resources/viya/{modelID}/{modelFileName}', 'rb') as _pFile:
-    _thisModelFit = pickle.load(_pFile)""".format(
-                        modelID=modelID, modelFileName=modelFileName
+    _thisModelFit = {pickleType}.load(_pFile)""".format(
+                        modelID=modelID, modelFileName=modelFileName, pickleType=pickleType
                     )
                 )
             elif isViya35 and isBinaryModel:
@@ -234,8 +237,8 @@ def upload_and_copy_score_resources(model, files):
                 cls.pyFile.write(
                     """\n
 with open(settings.pickle_path + '{modelFileName}', 'rb') as _pFile:
-    _thisModelFit = pickle.load(_pFile)""".format(
-                        modelFileName=modelFileName
+    _thisModelFit = {pickleType}.load(_pFile)""".format(
+                        modelFileName=modelFileName, pickleType=pickleType
                     )
                 )
             elif not isViya35 and isBinaryModel:
@@ -279,8 +282,8 @@ def score{modelPrefix}({inputVarList}):
                     cls.pyFile.write(
                         """
         with open('/models/resources/viya/{modelID}/{modelFileName}', 'rb') as _pFile:
-            _thisModelFit = pickle.load(_pFile)""".format(
-                            modelID=modelID, modelFileName=modelFileName
+            _thisModelFit = {pickleType}.load(_pFile)""".format(
+                            modelID=modelID, modelFileName=modelFileName, pickleType = pickleType
                         )
                     )
                 elif isViya35 and isH2OModel and not isBinaryModel:
@@ -302,8 +305,8 @@ def score{modelPrefix}({inputVarList}):
                     cls.pyFile.write(
                         """
         with open(settings.pickle_path + '{modelFileName}', 'rb') as _pFile:
-            _thisModelFit = pickle.load(_pFile)""".format(
-                            modelFileName=modelFileName
+            _thisModelFit = {pickleType}.load(_pFile)""".format(
+                            modelFileName=modelFileName, pickleType=pickleType
                         )
                     )
                 elif not isViya35 and isH2OModel: