add workflow/{workflow_id}/tasks endpoint

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, WorkflowRunHandler
+from jupyter_scheduler.workflows import WorkflowsHandler, WorkflowsRunHandler, WorkflowsTasksHandler
 
 from .handlers import (
     BatchJobHandler,
@@ -37,11 +37,15 @@ class SchedulerApp(ExtensionApp):
         (r"scheduler/job_definitions/%s/jobs" % JOB_DEFINITION_ID_REGEX, JobFromDefinitionHandler),
         (r"scheduler/runtime_environments", RuntimeEnvironmentsHandler),
         (r"scheduler/config", ConfigHandler),
-        (r"scheduler/worklows", WorkflowHandler),
-        (fr"scheduler/worklows/{WORKFLOW_ID_REGEX}", WorkflowHandler),
+        (r"scheduler/worklows", WorkflowsHandler),
+        (rf"scheduler/worklows/{WORKFLOW_ID_REGEX}", WorkflowsHandler),
         (
             rf"scheduler/worklows/{WORKFLOW_ID_REGEX}/run",
-            WorkflowRunHandler,
+            WorkflowsRunHandler,
+        ),
+        (
+            rf"scheduler/worklows/{WORKFLOW_ID_REGEX}/tasks",
+            WorkflowsTasksHandler,
         ),
     ]
 

jupyter_scheduler/scheduler.py

@@ -473,7 +473,7 @@ def copy_input_folder(self, input_uri: str, nb_copy_to_path: str) -> List[str]:
             destination_dir=staging_dir,
         )
 
-    def create_job(self, model: CreateJob) -> str:
+    def create_job(self, model: CreateJob, run: bool = True) -> str:
         if not model.job_definition_id and not self.file_exists(model.input_uri):
             raise InputUriError(model.input_uri)
 
@@ -514,6 +514,9 @@ def create_job(self, model: CreateJob) -> str:
             else:
                 self.copy_input_file(model.input_uri, staging_paths["input"])
 
+            if not run:
+                return job.job_id
+
             # The MP context forces new processes to not be forked on Linux.
             # This is necessary because `asyncio.get_event_loop()` is bugged in
             # forked processes in Python versions below 3.12. This method is

jupyter_scheduler/workflows.py

@@ -14,11 +14,11 @@
     ExtensionHandlerMixin,
     JobHandlersMixin,
 )
-from jupyter_scheduler.models import Status
+from jupyter_scheduler.models import CreateJob, Status, UpdateJob
 from jupyter_scheduler.pydantic_v1 import BaseModel, ValidationError
 
 
-class WorkflowHandler(ExtensionHandlerMixin, JobHandlersMixin, APIHandler):
+class WorkflowsHandler(ExtensionHandlerMixin, JobHandlersMixin, APIHandler):
     @authenticated
     async def post(self):
         payload = self.get_json_body()
@@ -60,7 +60,76 @@ async def get(self, workflow_id: str = None):
             self.finish(workflow.json())
 
 
-class WorkflowRunHandler(ExtensionHandlerMixin, JobHandlersMixin, APIHandler):
+class WorkflowsTasksHandler(ExtensionHandlerMixin, JobHandlersMixin, APIHandler):
+    @authenticated
+    async def post(self, workflow_id: str):
+        print("WorkflowsTasksHandler post")
+        payload = self.get_json_body()
+        if workflow_id != payload.get("workflow_id"):
+            raise HTTPError(
+                400,
+                "Error during workflow job creation. workflow_id in the URL and payload don't match.",
+            )
+        try:
+            job_id = await ensure_async(self.scheduler.create_job(CreateJob(**payload), run=False))
+        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 workflow job."
+            ) from e
+        else:
+            self.finish(json.dumps(dict(job_id=job_id)))
+
+    @authenticated
+    async def patch(self, workflow_id: str, job_id: str):
+        payload = self.get_json_body()
+        if workflow_id != payload.get("workflow_id", None):
+            raise HTTPError(
+                400,
+                "Error during workflow job creation. workflow_id in the URL and payload don't match.",
+            )
+        status = payload.get("status")
+        status = Status(status) if status else None
+
+        if status and status != Status.STOPPED:
+            raise HTTPError(
+                500,
+                "Invalid value for field 'status'. Workflow job status can only be updated to status 'STOPPED' after creation.",
+            )
+        try:
+            if status:
+                await ensure_async(self.scheduler.stop_job(job_id))
+            else:
+                await ensure_async(self.scheduler.update_job(job_id, UpdateJob(**payload)))
+        except ValidationError as e:
+            self.log.exception(e)
+            raise HTTPError(500, 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 while updating the workflow job."
+            ) from e
+        else:
+            self.set_status(204)
+            self.finish()
+
+
+class WorkflowsRunHandler(ExtensionHandlerMixin, JobHandlersMixin, APIHandler):
     @authenticated
     async def post(self, workflow_id: str):
         try:
@@ -79,7 +148,9 @@ async def post(self, workflow_id: str):
             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
+            raise HTTPError(
+                500, "Unexpected error occurred during attempt to run a workflow."
+            ) from e
         else:
             self.finish(json.dumps(dict(workflow_id=workflow_id)))