split create_workflow and run_workflow endpoints and methods

Author: andrii-i

jupyter_scheduler/extension.py

@@ -6,7 +6,7 @@
 from traitlets import Bool, Type, Unicode, default
 
 from jupyter_scheduler.orm import create_tables
-from jupyter_scheduler.workflows import WorkflowHandler
+from jupyter_scheduler.workflows import WorkflowHandler, WorkflowRunHandler
 
 from .handlers import (
     BatchJobHandler,
@@ -38,9 +38,10 @@ class SchedulerApp(ExtensionApp):
         (r"scheduler/runtime_environments", RuntimeEnvironmentsHandler),
         (r"scheduler/config", ConfigHandler),
         (r"scheduler/worklows", WorkflowHandler),
-        (r"scheduler/worklows/{}/run".format(WORKFLOW_ID_REGEX, JOB_ID_REGEX), WorkflowHandler),
-        (r"scheduler/worklows/%s/jobs" % WORKFLOW_ID_REGEX, WorkflowHandler),
-        (r"scheduler/worklows/{}/jobs/{}".format(WORKFLOW_ID_REGEX, JOB_ID_REGEX), WorkflowHandler),
+        (
+            r"scheduler/worklows/{}/run".format(WORKFLOW_ID_REGEX),
+            WorkflowRunHandler,
+        ),
     ]
 
     drop_tables = Bool(False, config=True, help="Drop the database tables before starting.")

jupyter_scheduler/scheduler.py

@@ -532,23 +532,18 @@ def create_job(self, model: CreateJob) -> str:
         return job_id
 
     def create_workflow(self, model: CreateWorkflow) -> str:
-
         with self.db_session() as session:
-
             workflow = Workflow(**model.dict(exclude_none=True))
-
             session.add(workflow)
             session.commit()
+            return workflow.workflow_id
 
-            execution_manager = self.execution_manager_class(
-                workflow_id=workflow.workflow_id,
-                db_url=self.db_url,
-            )
-            execution_manager.process_workflow()
-            session.commit()
-
-            workflow_id = workflow.workflow_id
-
+    def run_workflow(self, workflow_id: str) -> str:
+        execution_manager = self.execution_manager_class(
+            workflow_id=workflow_id,
+            db_url=self.db_url,
+        )
+        execution_manager.process_workflow()
         return workflow_id
 
     def update_job(self, job_id: str, model: UpdateJob):

jupyter_scheduler/workflows.py

@@ -45,6 +45,30 @@ async def post(self):
             self.finish(json.dumps(dict(workflow_id=workflow_id)))
 
 
+class WorkflowRunHandler(ExtensionHandlerMixin, JobHandlersMixin, APIHandler):
+    @authenticated
+    async def post(self, workflow_id: str):
+        try:
+            workflow_id = await ensure_async(self.scheduler.run_workflow(workflow_id))
+        except ValidationError as e:
+            self.log.exception(e)
+            raise HTTPError(500, str(e)) from e
+        except InputUriError as e:
+            self.log.exception(e)
+            raise HTTPError(500, str(e)) from e
+        except IdempotencyTokenError as e:
+            self.log.exception(e)
+            raise HTTPError(409, str(e)) from e
+        except SchedulerError as e:
+            self.log.exception(e)
+            raise HTTPError(500, str(e)) from e
+        except Exception as e:
+            self.log.exception(e)
+            raise HTTPError(500, "Unexpected error occurred during creation of a workflow.") from e
+        else:
+            self.finish(json.dumps(dict(workflow_id=workflow_id)))
+
+
 class CreateWorkflow(BaseModel):
     tasks: List[str]