add workflow handler and endpoint

jupyter_scheduler/executors.py

@@ -4,17 +4,22 @@
 import tarfile
 import traceback
 from abc import ABC, abstractmethod
-from typing import Dict
+from functools import lru_cache
+from typing import Dict, List
 
 import fsspec
 import nbconvert
 import nbformat
 from nbconvert.preprocessors import CellExecutionError, ExecutePreprocessor
+from prefect import flow, task
+from prefect.futures import as_completed
+from prefect_dask.task_runners import DaskTaskRunner
 
 from jupyter_scheduler.models import DescribeJob, JobFeature, Status
-from jupyter_scheduler.orm import Job, create_session
+from jupyter_scheduler.orm import Job, Workflow, create_session
 from jupyter_scheduler.parameterize import add_parameters
 from jupyter_scheduler.utils import get_utc_timestamp
+from jupyter_scheduler.workflows import DescribeWorkflow
 
 
 class ExecutionManager(ABC):
@@ -29,14 +34,29 @@ class ExecutionManager(ABC):
  _model = None
  _db_session = None
 
- def __init__(self, job_id: str, root_dir: str, db_url: str, staging_paths: Dict[str, str]):
+ def __init__(
+ self,
+ job_id: str,
+ workflow_id: str,
+ root_dir: str,
+ db_url: str,
+ staging_paths: Dict[str, str],
+ ):
  self.job_id = job_id
+ self.workflow_id = workflow_id
  self.staging_paths = staging_paths
  self.root_dir = root_dir
  self.db_url = db_url
 
  @property
  def model(self):
+ if self.workflow_id:
+ with self.db_session() as session:
+ workflow = (
+ session.query(Workflow).filter(Workflow.workflow_id == self.workflow_id).first()
+ )
+ self._model = DescribeWorkflow.from_orm(workflow)
+ return self._model
  if self._model is None:
  with self.db_session() as session:
  job = session.query(Job).filter(Job.job_id == self.job_id).first()
@@ -65,6 +85,18 @@ def process(self):
  else:
  self.on_complete()
 
+ def process_workflow(self):
+
+ self.before_start_workflow()
+ try:
+ self.execute_workflow()
+ except CellExecutionError as e:
+ self.on_failure_workflow(e)
+ except Exception as e:
+ self.on_failure_workflow(e)
+ else:
+ self.on_complete_workflow()
+
  @abstractmethod
  def execute(self):
  """Performs notebook execution,
@@ -74,6 +106,11 @@ def execute(self):
  """
  pass
 
+ @abstractmethod
+ def execute_workflow(self):
+ """Performs workflow execution"""
+ pass
+
  @classmethod
  @abstractmethod
  def supported_features(cls) -> Dict[JobFeature, bool]:
@@ -98,6 +135,15 @@ def before_start(self):
  )
  session.commit()
 
+ def before_start_workflow(self):
+ """Called before start of execute"""
+ workflow = self.model
+ with self.db_session() as session:
+ session.query(Workflow).filter(Workflow.workflow_id == workflow.workflow_id).update(
+ {"status": Status.IN_PROGRESS}
+ )
+ session.commit()
+
  def on_failure(self, e: Exception):
  """Called after failure of execute"""
  job = self.model
@@ -109,6 +155,17 @@ def on_failure(self, e: Exception):
 
  traceback.print_exc()
 
+ def on_failure_workflow(self, e: Exception):
+ """Called after failure of execute"""
+ workflow = self.model
+ with self.db_session() as session:
+ session.query(Workflow).filter(Workflow.workflow_id == workflow.workflow_id).update(
+ {"status": Status.FAILED, "status_message": str(e)}
+ )
+ session.commit()
+
+ traceback.print_exc()
+
  def on_complete(self):
  """Called after job is completed"""
  job = self.model
@@ -118,10 +175,40 @@ def on_complete(self):
  )
  session.commit()
 
+ def on_complete_workflow(self):
+ workflow = self.model
+ with self.db_session() as session:
+ session.query(Workflow).filter(Workflow.workflow_id == workflow.workflow_id).update(
+ {"status": Status.COMPLETED}
+ )
+ session.commit()
+
 
 class DefaultExecutionManager(ExecutionManager):
  """Default execution manager that executes notebooks"""
 
+ @task(task_run_name="{task_id}")
+ def execute_task(task_id: str):
+ print(f"Task {task_id} executed")
+ return task_id
+
+ @flow(task_runner=DaskTaskRunner())
+ def execute_workflow(self):
+ workflow: DescribeWorkflow = self.model
+ tasks = {task["id"]: task for task in workflow.tasks}
+
+ # create Prefect tasks, use caching to ensure Prefect tasks are created before wait_for is called on them
+ @lru_cache(maxsize=None)
+ def make_task(task_id, execute_task):
+ deps = tasks[task_id]["dependsOn"]
+ return execute_task.submit(
+ task_id, wait_for=[make_task(dep_id, execute_task) for dep_id in deps]
+ )
+
+ final_tasks = [make_task(task_id, self.execute_task) for task_id in tasks]
+ for future in as_completed(final_tasks):
+ print(future.result())
+
  def execute(self):
  job = self.model
 

jupyter_scheduler/extension.py

@@ -6,6 +6,7 @@
 from traitlets import Bool, Type, Unicode, default
 
 from jupyter_scheduler.orm import create_tables
+from jupyter_scheduler.workflows import WorkflowHandler
 
 from .handlers import (
  BatchJobHandler,
@@ -35,6 +36,7 @@ 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),
  ]
 
  drop_tables = Bool(False, config=True, help="Drop the database tables before starting.")

jupyter_scheduler/orm.py

@@ -107,6 +107,16 @@ class Job(CommonColumns, Base):
  # Any default values specified for new columns will be ignored during the migration process.
 
 
+class Workflow(Base):
+ __tablename__ = "workflows"
+ __table_args__ = {"extend_existing": True}
+ workflow_id = Column(String(36), primary_key=True, default=generate_uuid)
+ tasks = Column(JsonType(1024))
+ status = Column(String(64), default=Status.STOPPED)
+ # All new columns added to this table must be nullable to ensure compatibility during database migrations.
+ # Any default values specified for new columns will be ignored during the migration process.
+
+
 class JobDefinition(CommonColumns, Base):
  __tablename__ = "job_definitions"
  __table_args__ = {"extend_existing": True}

jupyter_scheduler/scheduler.py

@@ -40,12 +40,13 @@
  UpdateJob,
  UpdateJobDefinition,
 )
-from jupyter_scheduler.orm import Job, JobDefinition, create_session
+from jupyter_scheduler.orm import Job, JobDefinition, Workflow, create_session
 from jupyter_scheduler.utils import (
  copy_directory,
  create_output_directory,
  create_output_filename,
 )
+from jupyter_scheduler.workflows import CreateWorkflow
 
 
 class BaseScheduler(LoggingConfigurable):
@@ -111,6 +112,10 @@ def create_job(self, model: CreateJob) -> str:
  """
  raise NotImplementedError("must be implemented by subclass")
 
+ def create_workflow(self, model: CreateWorkflow) -> str:
+ """Creates a new workflow record, may trigger execution of the workflow."""
+ raise NotImplementedError("must be implemented by subclass")
+
  def update_job(self, job_id: str, model: UpdateJob):
  """Updates job metadata in the persistence store,
  for example name, status etc. In case of status
@@ -526,6 +531,27 @@ 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()
+
+ execution_manager = self.execution_manager_class(
+ workflow_id=workflow.workflow_id,
+ root_dir=self.root_dir,
+ db_url=self.db_url,
+ )
+ execution_manager.process_workflow()
+ session.commit()
+
+ workflow_id = workflow.workflow_id
+
+ return workflow_id
+
  def update_job(self, job_id: str, model: UpdateJob):
  with self.db_session() as session:
  session.query(Job).filter(Job.job_id == job_id).update(model.dict(exclude_none=True))

jupyter_scheduler/workflows.py

@@ -0,0 +1,57 @@
+import json
+from typing import List
+
+from jupyter_server.utils import ensure_async
+from tornado.web import HTTPError, authenticated
+
+from jupyter_scheduler.exceptions import (
+ IdempotencyTokenError,
+ InputUriError,
+ SchedulerError,
+)
+from jupyter_scheduler.handlers import (
+ APIHandler,
+ ExtensionHandlerMixin,
+ JobHandlersMixin,
+)
+from jupyter_scheduler.models import Status
+from jupyter_scheduler.pydantic_v1 import BaseModel, ValidationError
+
+
+class WorkflowHandler(ExtensionHandlerMixin, JobHandlersMixin, APIHandler):
+ @authenticated
+ async def post(self):
+ payload = self.get_json_body()
+ try:
+ workflow_id = await ensure_async(
+ self.scheduler.create_workflow(CreateWorkflow(**payload))
+ )
+ self.log.info(payload)
+ print(payload)
+ 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]
+
+
+class DescribeWorkflow(BaseModel):
+ workflow_id: str
+ tasks: List[str] = None
+ status: Status = Status.CREATED