Renaming set_run_classes to set_run_inference_info. Adding prediction_type, code_env-name and target.

Author: lpenet

dataikuapi/dss/mlflow.py

@@ -134,33 +134,67 @@ def clean_experiment_tracking_db(self):
         """
         self.client._perform_raw("DELETE", "/api/2.0/mlflow/extension/clean-db/%s" % self.project_key)
 
-    def set_run_classes(self, run_id, classes):
+    def set_run_inference_info(self, run_id, prediction_type, classes, code_env_name, target):
         """
-        Stores the classes of the target of classification models trained in the specified run. This information is leveraged
-        to prefill the classes when deploying using the GUI an MLflow model as a version of a DSS Saved Model.
+        Sets the prediction_type of the model, and the classes, if it is a classification model.
 
+        prediction_type must be one of:
+        - NON_TABULAR if the model is not tabular
+        - REGRESSION
+        - BINARY_CLASSIFICATION
+        - MULTICLASS
+
+        Classes must be specified if and only if the model is a BINARY_CLASSIFICATION or MULTICLASS model.
+
+        This information is leveraged to filter saved models on their prediction type and prefill the classes
+        when deploying using the GUI an MLflow model as a version of a DSS Saved Model.
+
+        :param prediction_type: prediction type (see doc)
+        :type prediction_type: str
         :param run_id: run_id for which to set the classes
         :type run_id: str
-        :param classes: ordered list of classes
+        :param classes: ordered list of classes (not for all prediction types, see doc)
         :type classes: list(str)
-        """
+        :param code_env_name: name of an adequate DSS python code environment
+        :type code_env_name: str
+        :param target: name of the target
+        :type target: str
+        """
+        if prediction_type not in {"NON_TABULAR", "REGRESSION", "BINARY_CLASSIFICATION", "MULTICLASS"}:
+            raise ValueError('Invalid prediction type: {}'.format(prediction_type))
+
+        if classes and prediction_type not in {"BINARY_CLASSIFICATION", "MULTICLASS"}:
+            raise ValueError('Classes can be specified only for BINARY_CLASSIFICATION or MULTICLASS prediction types')
+        if prediction_type in {"BINARY_CLASSIFICATION", "MULTICLASS"}:
+            if not classes:
+                raise ValueError('Classes must be specified for {} prediction type'.format(prediction_type))
+            if not isinstance(classes, list):
+                raise ValueError('Wrong type for classes: {}'.format(type(classes)))
+            for cur_class in classes:
+                if not cur_class:
+                    raise ValueError('class can not be None')
+                if not isinstance(cur_class, str):
+                    raise ValueError('Wrong type for class {}: {}'.format(cur_class, type(cur_class)))
+
+        if code_env_name and not isinstance(code_env_name, str):
+            raise ValueError('code_env_name must be a string')
+        if target and not isinstance(target, str):
+            raise ValueError('target must be a string')
+
+        params = {
+            "run_id": run_id,
+            "prediction_type": prediction_type
+        }
+
+        if classes:
+            params["classes"] = json.dumps(classes)
+        if code_env_name:
+            params["code_env_name"] = code_env_name
+        if target:
+            params["target"] = target
 
-        if not classes:
-            raise ValueError('Parameter classes must be defined')
-        if not isinstance(classes, list):
-            raise ValueError('Wrong type for classes: {}'.format(type(classes)))
-        for cur_class in classes:
-            if not cur_class:
-                raise ValueError('class can not be None')
-            if not isinstance(cur_class, str):
-                raise ValueError('Wrong type for class {}: {}'.format(cur_class, type(cur_class)))
         self.client._perform_http(
-            "POST", "/api/2.0/mlflow/runs/set-tag",
+            "POST", "/api/2.0/mlflow/extension/set-run-inference-info",
             headers={"x-dku-mlflow-project-key": self.project_key},
-            body={
-                "run_id": run_id,
-                "run_uuid": run_id,
-                "key": "dku-ext.targetClasses",
-                "value": json.dumps(classes)
-            }
+            body=params
         )