factorize and add helpers to wait on scenarios/jobs

Author: FChataigner

dataikuapi/dss/job.py

@@ -1,3 +1,6 @@
+import time
+import sys
+from dataikuapi.utils import DataikuException
 
 class DSSJob(object):
     """
@@ -40,3 +43,24 @@ def get_log(self, activity=None):
             params={
                 "activity" : activity
             })
+
+class DSSJobWaiter(object):
+    """
+    Helper to wait for a job's completion
+    """    
+    def __init__(self, job):
+        self.job = job
+
+    def wait(self, no_fail=False):
+        job_state = self.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 = self.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

dataikuapi/dss/project.py

@@ -3,7 +3,7 @@
 from .recipe import DSSRecipe
 from .managedfolder import DSSManagedFolder
 from .savedmodel import DSSSavedModel
-from .job import DSSJob
+from .job import DSSJob, DSSJobWaiter
 from .scenario import DSSScenario
 from .apiservice import DSSAPIService
 import sys
@@ -291,18 +291,8 @@ def start_job_and_wait(self, definition, no_fail=False):
         """
         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
+        waiter = DSSJobWaiter(job)
+        return waiter.wait(no_fail)
 
     def new_job_definition_builder(self, job_type='NON_RECURSIVE_FORCED_BUILD'):
         return JobDefinitionBuilder(self.project_key, job_type)

dataikuapi/dss/scenario.py

@@ -57,24 +57,9 @@ def run_and_wait(self, params={}, no_fail=False):
             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)
-        print("Scenario (ID: '%s') run started " % scenario_run.run.get('runId', 'not found'))
-        while not scenario_run.run.get('result', False):
-            scenario_run = trigger_fire.get_scenario_run()
-            time.sleep(5)
-        return scenario_run
+        scenario_run = trigger_fire.wait_for_scenario_run(no_fail)
+        waiter = DSSScenarioRunWaiter(scenario_run, trigger_fire)
+        return waiter.wait(no_fail)
 
     def get_last_runs(self, limit=10, only_finished_runs=False):
         """
@@ -215,6 +200,24 @@ def get_duration(self):
         duration = (end_time - self.get_start_time()).total_seconds()
         return duration
 
+class DSSScenarioRunWaiter(object):
+    """
+    Helper to wait for a job's completion
+    """    
+    def __init__(self, scenario_run, trigger_fire):
+        self.trigger_fire = trigger_fire
+        self.scenario_run = scenario_run
+
+    def wait(self, no_fail=False):
+        while not self.scenario_run.run.get('result', False):
+            self.scenario_run = self.trigger_fire.get_scenario_run()
+            time.sleep(5)
+        outcome = self.scenario_run.run.get('result', None).get('outcome', 'UNKNOWN')
+        if outcome == 'SUCCESS' or no_fail:
+            return self.scenario_run
+        else:
+            raise DataikuException("Scenario run returned status %s" % outcome)
+
 class DSSTriggerFire(object):
     """
     The activation of a trigger on the DSS instance
@@ -254,3 +257,19 @@ def is_cancelled(self, refresh=False):
                 'triggerRunId' : self.run_id
             })
         return self.trigger_fire["cancelled"]
+
+    def wait_for_scenario_run(self, no_fail=False):
+        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 self.is_cancelled(refresh=refresh_trigger_counter == 0):
+                if no_fail:
+                    return None
+                else:
+                    raise DataikuException("Scenario run has been cancelled")
+            scenario_run = self.get_scenario_run()
+            time.sleep(5)
+        return scenario_run