Merge PR #148 train & list ML Task Queues

from feature/dss100-mltask-queues

Author: instanceofme

dataikuapi/dss/ml.py

@@ -3348,7 +3348,7 @@ def get_settings(self):
         else:
             return DSSClusteringMLTaskSettings(self.client, self.project_key, self.analysis_id, self.mltask_id, settings)
 
-    def train(self, session_name=None, session_description=None):
+    def train(self, session_name=None, session_description=None, run_queue=False):
         """
         Trains models for this ML Task
         
@@ -3366,7 +3366,7 @@ def train(self, session_name=None, session_description=None):
         :return: A list of model identifiers
         :rtype: list of strings
         """
-        train_ret = self.start_train(session_name, session_description)
+        train_ret = self.start_train(session_name, session_description, run_queue)
         self.wait_train_complete()
         return self.get_trained_models_ids(session_id = train_ret["sessionId"])
 
@@ -3395,7 +3395,7 @@ def ensemble(self, model_ids=None, method=None):
         return train_ret
 
 
-    def start_train(self, session_name=None, session_description=None):
+    def start_train(self, session_name=None, session_description=None, run_queue=False):
         """
         Starts asynchronously a new train session for this ML Task.
 
@@ -3406,7 +3406,8 @@ def start_train(self, session_name=None, session_description=None):
         """
         session_info = {
                             "sessionName" : session_name,
-                            "sessionDescription" : session_description
+                            "sessionDescription" : session_description,
+                            "runQueue": run_queue
                         }
 
         return self.client._perform_json(
@@ -3521,6 +3522,16 @@ def delete_trained_model(self, model_id):
         self.client._perform_empty(
             "DELETE", "/projects/%s/models/lab/%s/%s/models/%s" % (self.project_key, self.analysis_id, self.mltask_id, model_id))
 
+    def train_queue(self):
+        """
+        Trains this MLTask's queue
+
+        :return: A dict including the next sessionID to be trained in the queue
+        :rtype dict
+        """
+        return self.client._perform_json(
+            "POST", "/projects/%s/models/lab/%s/%s/actions/train-queue" % (self.project_key, self.analysis_id, self.mltask_id))
+
     def deploy_to_flow(self, model_id, model_name, train_dataset, test_dataset=None, redo_optimization=True):
         """
         Deploys a trained model from this ML Task to a saved model + train recipe in the Flow.
@@ -3606,3 +3617,17 @@ def guess(self, prediction_type=None, reguess_level=None):
             "PUT",
             "/projects/%s/models/lab/%s/%s/guess" % (self.project_key, self.analysis_id, self.mltask_id),
             params = obj)
+
+
+class DSSMLTaskQueues(object):
+    """
+    Iterable listing of MLTask queues
+    """
+    def __init__(self, data):
+        self.data = data
+
+    def __iter__(self):
+        return self.data["queues"].__iter__()
+
+    def get_raw(self):
+        return self.data

dataikuapi/dss/project.py

@@ -16,7 +16,7 @@
 from .macro import DSSMacro
 from .wiki import DSSWiki
 from .discussion import DSSObjectDiscussions
-from .ml import DSSMLTask
+from .ml import DSSMLTask, DSSMLTaskQueues
 from .analysis import DSSAnalysis
 from .flow import DSSProjectFlow
 from .app import DSSAppManifest
@@ -638,6 +638,15 @@ def get_ml_task(self, analysis_id, mltask_id):
         """
         return DSSMLTask(self.client, self.project_key, analysis_id, mltask_id)
 
+    def list_mltask_queues(self):
+        """
+        List non-empty ML task queues in this project 
+        
+        :returns: an iterable :class:`DSSMLTaskQueues` listing of MLTask queues (each a dict)
+        :rtype: :class:`DSSMLTaskQueues`
+        """ 
+        data = self.client._perform_json("GET", "/projects/%s/models/labs/mltask-queues" % self.project_key)
+        return DSSMLTaskQueues(data)
 
     def create_analysis(self, input_dataset):
         """