Getting origin ml task/trained model: improve error handling

Author: Louis Pouillot

dataikuapi/dss/ml.py

@@ -528,19 +528,19 @@ def save_user_meta(self):
 
     def get_origin_analysis_trained_model(self):
         """
-        Fetch details about the model in an analysis, this model has been exported from
+        Fetch details about the model in an analysis, this model has been exported from. Returns None if the
+        deployed trained model does not have a origin analysis trained model.
 
-        :rtype: DSSTrainedModelDetails
+        :rtype: DSSTrainedModelDetails | None
         """
         if self.saved_model is None:
             return self
         else:
-            if "smOrigin" not in self.get_raw():
-                raise DataikuException("Unknow ")
-            fmi = self.get_raw()["smOrigin"]["fullModelId"]
-            origin_ml_task = DSSMLTask.from_full_model_id(self.saved_model.client, fmi,
-                                                          project_key=self.saved_model.project_key)
-            return origin_ml_task.get_trained_model_details(fmi)
+            fmi = self.get_raw().get("smOrigin", {}).get("fullModelId")
+            if fmi is not None:
+                origin_ml_task = DSSMLTask.from_full_model_id(self.saved_model.client, fmi,
+                                                              project_key=self.saved_model.project_key)
+                return origin_ml_task.get_trained_model_details(fmi)
 
 class DSSTreeNode(object):
     def __init__(self, tree, i):
@@ -1446,10 +1446,13 @@ class DSSMLTask(object):
 
     @staticmethod
     def from_full_model_id(client, fmi, project_key=None):
-        match = re.match(u'A-([\w]+)-([\w]+)-([\w]+)-s\w+-pp\w+-m\w+', fmi)
-        if project_key is None:
-            project_key = match.group(1)
-        return DSSMLTask(client, project_key, match.group(2), match.group(3))
+        match = re.match("^A-(\w+)-(\w+)-(\w+)-(s[0-9]+)-(pp[0-9]+(-part-(\w+)|-base)?)-(m[0-9]+)$", fmi)
+        if match is None:
+            return DataikuException("Invalid model id: {}".format(fmi))
+        else:
+            if project_key is None:
+                project_key = match.group(1)
+            return DSSMLTask(client, project_key, match.group(2), match.group(3))
 
     """A handle to interact with a MLTask for prediction or clustering in a DSS visual analysis"""
     def __init__(self, client, project_key, analysis_id, mltask_id):

dataikuapi/dss/savedmodel.py

@@ -93,15 +93,16 @@ def delete_versions(self, versions, remove_intermediate=True):
 
     def get_origin_ml_task(self):
         """
-        Fetch the last ML task that has been exported to this saved model
+        Fetch the last ML task that has been exported to this saved model. Returns None if the saved model
+        does not have an origin ml task.
 
-        :rtype: DSSMLTask
+        :rtype: DSSMLTask | None
         """
-        definition = self.get_definition()
-        fmi = definition["lastExportedFrom"]
+        fmi = self.get_definition().get("lastExportedFrom")
+        if fmi is not None:
+            origin_ml_task = DSSMLTask.from_full_model_id(self.client, fmi, project_key=self.project_key)
+            return origin_ml_task.get_trained_model_details(fmi)
 
-        origin_ml_task = DSSMLTask.from_full_model_id(self.client, fmi, project_key=self.project_key)
-        return origin_ml_task.get_trained_model_details(fmi)
 
     ########################################################
     # Metrics