Proxy model feature branch API

dataikuapi/dss/project.py

@@ -722,7 +722,7 @@ def get_saved_model(self, sm_id):
         """
         return DSSSavedModel(self.client, self.project_key, sm_id)
 
-    def create_mlflow_pyfunc_model(self, name, prediction_type = None):
+    def create_mlflow_pyfunc_model(self, name, prediction_type=None):
         """
         Creates a new external saved model for storing and managing MLFlow models
 
@@ -732,13 +732,29 @@ def create_mlflow_pyfunc_model(self, name, prediction_type = None):
         :rtype: :class:`dataikuapi.dss.savedmodel.DSSSavedModel`
         """
         model = {
-            "savedModelType" : "MLFLOW_PYFUNC",
-            "predictionType" : prediction_type,
+            "savedModelType": "MLFLOW_PYFUNC",
+            "predictionType": prediction_type,
             "name": name
         }
 
-        id = self.client._perform_json("POST", "/projects/%s/savedmodels/" % self.project_key, body = model)["id"]
-        return self.get_saved_model(id)
+        saved_model_id = self.client._perform_json("POST", "/projects/%s/savedmodels/" % self.project_key, body=model)["id"]
+        return self.get_saved_model(saved_model_id)
+
+    def create_proxy_model(self, name, prediction_type=None):
+        """
+        Creates a new external saved model for storing and managing MLFlow models
+
+        :param string name: Human readable name for the new saved model in the flow
+        :param string prediction_type: Optional (but needed for most operations). One of BINARY_CLASSIFICATION, MULTICLASS or REGRESSION
+        """
+        model = {
+            "savedModelType": "PROXY_MODEL",
+            "predictionType": prediction_type,
+            "name": name
+        }
+
+        saved_model_id = self.client._perform_json("POST", "/projects/%s/savedmodels/" % self.project_key, body=model)["id"]
+        return self.get_saved_model(saved_model_id)
 
     ########################################################
     # Managed folders

dataikuapi/dss/savedmodel.py

@@ -138,7 +138,7 @@ def import_mlflow_version_from_path(self, version_id, path, code_env_name="INHER
                                   this model version.
                                   If value is "INHERIT", the container execution configuration of the project will be used.
                                   If value is "NONE", local execution will be used (no container)
-        :return a :class:MLFlowVersionHandler in order to interact with the new MLFlow model version
+        :return a :class:ExternalModelVersionHandler in order to interact with the new MLFlow model version
         """
         # TODO: Add a check that it's indeed a MLFlow model folder
         import shutil
@@ -151,7 +151,7 @@ def import_mlflow_version_from_path(self, version_id, path, code_env_name="INHER
             with open(archive_filename, "rb") as fp:
                 self.client._perform_empty("POST", "/projects/%s/savedmodels/%s/versions/%s?codeEnvName=%s&containerExecConfigName=%s" % (self.project_key, self.sm_id, version_id, code_env_name, container_exec_config_name),
                                            files={"file": (archive_filename, fp)})
-            return self.get_mlflow_version_handler(version_id)
+            return self.get_external_version_handler(version_id)
         finally:
             shutil.rmtree(archive_temp_dir)
 
@@ -171,7 +171,7 @@ def import_mlflow_version_from_managed_folder(self, version_id, managed_folder,
                                   this model version.
                                   If value is "INHERIT", the container execution configuration of the project will be used.
                                   If value is "NONE", local execution will be used (no container)
-        :return a :class:MLFlowVersionHandler in order to interact with the new MLFlow model version
+        :return a :class:ExternalModelVersionHandler in order to interact with the new MLFlow model version
         """
         # TODO: Add a check that it's indeed a MLFlow model folder
         folder_ref = None
@@ -190,13 +190,40 @@ def import_mlflow_version_from_managed_folder(self, version_id, managed_folder,
             params={"folderRef": folder_ref, "path": path},
             files={"file": (None, None)}  # required for backend-mandated multipart request
         )
-        return self.get_mlflow_version_handler(version_id)
+        return self.get_external_version_handler(version_id)
 
-    def get_mlflow_version_handler(self, version_id):
+    def get_external_version_handler(self, version_id):
         """
-        Returns a :class:MLFlowVersionHandler to interact with a MLFlow model version
+        Returns a :class:ExternalModelVersionHandler to interact with a MLFlow model version
         """
-        return MLFlowVersionHandler(self, version_id)
+        return ExternalModelVersionHandler(self, version_id)
+
+    def create_proxy_model_version(self, version_id, protocol, configuration, credentials, container_exec_config_name="INHERIT"):
+        """
+        Create a new version for this saved model from a path containing a MLFlow model.
+
+        Requires the saved model to have been created using :meth:`dataikuapi.dss.project.DSSProject.create_mlflow_pyfunc_model`.
+
+        :param str version_id: Identifier of the version to create
+        :param str protocol: one of ["KServe", "DSS_API_NODE"]
+        :param str configuration: A dictionnary containing the required params for the selected protocol
+        :param str configuration: A dictionnary containing the required credentials for the selected protocol
+        :param str container_exec_config_name: Name of the containerized execution configuration to use while creating
+                                  this model version.
+                                  If value is "INHERIT", the container execution configuration of the project will be used.
+                                  If value is "NONE", local execution will be used (no container)
+        :return a :class:ExternalModelVersionHandler in order to interact with the new Proxy model version
+        """
+        import json
+        self.client._perform_empty(
+            "POST", "/projects/{project_id}/savedmodels/{saved_model_id}/versions/{version_id}?containerExecConfigName={containerExecConfigName}".format(
+                project_id=self.project_key, saved_model_id=self.sm_id, version_id=version_id, containerExecConfigName=container_exec_config_name
+            ),
+            params={"protocol": protocol, "configuration": json.dumps(configuration), "credentials": json.dumps(credentials)},
+            files={"file": (None, None)}  # required for backend-mandated multipart request
+        )
+        return self.get_external_version_handler(version_id)
+
 
     ########################################################
     # Metrics
@@ -300,10 +327,10 @@ def save(self):
             "/projects/%s/savedmodels/%s/versions/%s/external-ml/metadata" % (self.version_handler.saved_model.project_key, self.version_handler.saved_model.sm_id, self.version_handler.version_id),
             body=self.data)
 
-class MLFlowVersionHandler:
-    """Handler to interact with an imported MLFlow model version"""
+class ExternalModelVersionHandler:
+    """Handler to interact with an imported externel model version (MLflow or Proxy model)"""
     def __init__(self, saved_model, version_id):
-        """Do not call this, use :meth:`DSSSavedModel.get_mlflow_version_handler`"""
+        """Do not call this, use :meth:`DSSSavedModel.get_external_version_handler`"""
         self.saved_model = saved_model
         self.version_id = version_id