Merge pull request #18 from dataiku/feature/scenarios-and-jobs-run-and-wait

Scenarios and jobs run and wait

Author: cstenac

dataikuapi/dss/project.py

@@ -1,3 +1,4 @@
+import time
 from dataset import DSSDataset
 from recipe import DSSRecipe
 from managedfolder import DSSManagedFolder
@@ -9,6 +10,8 @@
 import os.path as osp
 from .future import DSSFuture
 from .notebook import DSSNotebook
+from dataikuapi.utils import DataikuException
+
 
 class DSSProject(object):
     """
@@ -276,7 +279,33 @@ def start_job(self, definition):
         job_def = self.client._perform_json("POST", "/projects/%s/jobs/" % self.project_key, body = definition)
         return DSSJob(self.client, self.project_key, job_def['id'])
 
-
+    def start_job_and_wait(self, definition, no_fail=False):
+        """
+        Create a new job. Wait the end of the job to complete.
+        
+        Args:
+            definition: the definition for the job to create. The definition must contain the type of job (RECURSIVE_BUILD, 
+            NON_RECURSIVE_FORCED_BUILD, RECURSIVE_FORCED_BUILD, RECURSIVE_MISSING_ONLY_BUILD) and a list of outputs to build.
+            Optionally, a refreshHiveMetastore field can specify whether to re-synchronize the Hive metastore for recomputed
+            HDFS datasets.
+        """
+        job_def = self.client._perform_json("POST", "/projects/%s/jobs/" % self.project_key, body = definition)
+        job = DSSJob(self.client, self.project_key, job_def['id'])
+        job_state = job.get_status().get("baseStatus", {}).get("state", "")
+        sleep_time = 2
+        while job_state not in ["DONE", "ABORTED", "FAILED"]:
+            sleep_time = 300 if sleep_time >= 300 else sleep_time * 2
+            time.sleep(sleep_time)
+            job_state = job.get_status().get("baseStatus", {}).get("state", "")
+            if job_state in ["ABORTED", "FAILED"]:
+                if no_fail:
+                    break
+                else:
+                    raise DataikuException("Job run did not finish. Status: %s" % (job_state))
+        return job_state
+
+    def new_job_definition_builder(self, job_type='NON_RECURSIVE_FORCED_BUILD'):
+        return JobDefinitionBuilder(self.project_key, job_type)
 
     ########################################################
     # Variables
@@ -380,7 +409,7 @@ def list_imported_bundles(self):
                 "/projects/%s/bundles/imported" % self.project_key)
 
     def import_bundle_from_archive(self, archive_path):
-        return self.client._perform_empty("POST",
+        return self.client._perform_json("POST",
                 "/projects/%s/bundles/imported/actions/importFromArchive" % (self.project_key),
                  params = { "archivePath" : osp.abspath(archive_path) })
 
@@ -545,3 +574,45 @@ def set_tags(self, tags={}):
         @param obj: must be a modified version of the object returned by list_tags
         """
         return self.client._perform_empty("PUT", "/projects/%s/tags" % self.project_key, body = tags)
+
+
+class JobDefinitionBuilder(object):
+    def __init__(self, project_key, job_type="NON_RECURSIVE_FORCED_BUILD"):
+        """
+        Create a helper to build a job definition
+
+        :param project_key: the project in which the job is launched
+        :param job_type: the build type for the job  RECURSIVE_BUILD, NON_RECURSIVE_FORCED_BUILD, 
+                            RECURSIVE_FORCED_BUILD, RECURSIVE_MISSING_ONLY_BUILD
+
+        """
+        self.project_key = project_key
+        self.definition = {'type':job_type, 'refreshHiveMetastore':False, 'outputs':[]}
+
+    def with_type(self, job_type):
+        """
+        Sets the build type
+
+        :param job_type: the build type for the job  RECURSIVE_BUILD, NON_RECURSIVE_FORCED_BUILD, 
+                            RECURSIVE_FORCED_BUILD, RECURSIVE_MISSING_ONLY_BUILD
+        """
+        self.definition['type'] = job_type
+        return self
+
+    def with_refresh_metastore(self, refresh_metastore):
+        """
+        Sets whether the hive tables built by the job should have their definitions
+        refreshed after the corresponding dataset is built
+        """
+        self.definition['refreshHiveMetastore'] = refresh_metastore
+        return self
+
+    def with_output(self, name, object_type=None, object_project_key=None, partition=None):
+        """
+        Adds an item to build in the definition
+        """
+        self.definition['outputs'].append({'type':object_type, 'id':name, 'projectKey':object_project_key, 'partition':partition})
+        return self
+
+    def get_definition(self):
+        return self.definition

dataikuapi/dss/scenario.py

@@ -1,4 +1,7 @@
 from datetime import datetime
+import time
+from dataikuapi.utils import DataikuException
+
 
 class DSSScenario(object):
     """
@@ -26,6 +29,52 @@ def run(self, params={}):
             "POST", "/projects/%s/scenarios/%s/run" % (self.project_key, self.id), body=params)
         return DSSTriggerFire(self, trigger_fire)
 
+    def get_trigger_fire(self, trigger_id, trigger_run_id):
+        """
+        Requests a trigger of the run of a scenario
+
+        Args:
+            trigger_id:  Id of trigger
+            trigger_run_id: Id of the run of the trigger
+
+        Returns:
+            A :class:`dataikuapi.dss.admin.DSSTriggerFire` trigger handle
+        """
+        trigger_fire = self.client._perform_json(
+            "GET", "/projects/%s/scenarios/trigger/%s/%s" % (self.project_key, self.id, trigger_id), params={
+                'triggerRunId' : trigger_run_id
+            })
+        return DSSTriggerFire(self, trigger_fire)
+
+    def run_and_wait(self, params={}, no_fail=False):
+        """
+        Requests a run of the scenario, which will start after a few seconds. Wait the end of the run to complete.
+
+        Args:
+            params: additional parameters that will be passed to the scenario through trigger params
+
+        Returns:
+            A :class:`dataikuapi.dss.admin.DSSScenarioRun` run handle
+        """
+        trigger_fire = self.run(params)
+        scenario_run = None
+        refresh_trigger_counter = 0
+        while scenario_run is None:
+            refresh_trigger_counter += 1
+            if refresh_trigger_counter == 10:
+                refresh_trigger_counter = 0
+            if trigger_fire.is_cancelled(refresh=refresh_trigger_counter == 0):
+                if no_fail:
+                    return None
+                else:
+                    raise DataikuException("Scenario run has been cancelled")
+            scenario_run = trigger_fire.get_scenario_run()
+            time.sleep(5)
+        while not scenario_run.run.get('result', False):
+            scenario_run = trigger_fire.get_scenario_run()
+            time.sleep(5)
+        return scenario_run
+
     def get_last_runs(self, limit=10, only_finished_runs=False):
         """
         Get the list of the last runs of the scenario.
@@ -191,3 +240,16 @@ def get_scenario_run(self):
             return None
         else:
             return DSSScenarioRun(self.client, run['scenarioRun'])
+
+    def is_cancelled(self, refresh=False):
+        """
+        Whether the trigger has been cancelled
+
+        :param refresh: get the state of the trigger from the backend
+        """
+        if refresh == True:
+            self.trigger_fire = self.client._perform_json(
+            "GET", "/projects/%s/scenarios/trigger/%s/%s" % (self.project_key, self.scenario_id, self.trigger_id), params={
+                'triggerRunId' : self.run_id
+            })
+        return self.trigger_fire["cancelled"]