Update job model to be used with workflows

Author: andrii-i

jupyter_scheduler/executors.py

@@ -19,7 +19,7 @@
 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 DescribeTask, DescribeWorkflow
+from jupyter_scheduler.workflows import DescribeWorkflow
 
 
 class ExecutionManager(ABC):

jupyter_scheduler/models.py

@@ -42,6 +42,7 @@ def __str__(self) -> str:
 
 
 class Status(str, Enum):
+    DRAFT = "DRAFT"
     CREATED = "CREATED"
     QUEUED = "QUEUED"
     IN_PROGRESS = "IN_PROGRESS"
@@ -86,6 +87,8 @@ class CreateJob(BaseModel):
     output_filename_template: Optional[str] = OUTPUT_FILENAME_TEMPLATE
     compute_type: Optional[str] = None
     package_input_folder: Optional[bool] = None
+    depends_on: Optional[str] = None
+    workflow_id: str = None
 
     @root_validator
     def compute_input_filename(cls, values) -> Dict:
@@ -148,6 +151,8 @@ class DescribeJob(BaseModel):
     downloaded: bool = False
     package_input_folder: Optional[bool] = None
     packaged_files: Optional[List[str]] = []
+    depends_on: Optional[str] = None
+    workflow_id: str = None
 
     class Config:
         orm_mode = True
@@ -193,6 +198,7 @@ class UpdateJob(BaseModel):
     status: Optional[Status] = None
     name: Optional[str] = None
     compute_type: Optional[str] = None
+    depends_on: Optional[str] = None
 
 
 class DeleteJob(BaseModel):

jupyter_scheduler/orm.py

@@ -103,6 +103,8 @@ class Job(CommonColumns, Base):
     url = Column(String(256), default=generate_jobs_url)
     pid = Column(Integer)
     idempotency_token = Column(String(256))
+    depends_on = Column(JsonType(1024))
+    workflow_id = Column(String(36))
     # 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.
 
@@ -112,7 +114,7 @@ class Workflow(Base):
     __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)
+    status = Column(String(64), default=Status.DRAFT)
     # 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.
 

jupyter_scheduler/workflows.py

@@ -57,9 +57,10 @@ class Config:
         orm_mode = True
 
 
-class DescribeTask(BaseModel):
-    dependsOn: List[str] = []
-    status: Status = Status.CREATED
+class UpdateWorkflow(BaseModel):
+    workflow_id: str
+    tasks: List[str] = None
+    status: Status = None
 
     class Config:
         orm_mode = True

src/handler.ts

@@ -372,6 +372,7 @@ export namespace Scheduler {
     timezone?: string;
     active?: boolean;
     input_uri?: string;
+    depends_on?: string[];
   }
 
   export interface IDescribeJobDefinition {
@@ -418,6 +419,8 @@ export namespace Scheduler {
     output_formats?: string[];
     compute_type?: string;
     package_input_folder?: boolean;
+    depends_on?: string[];
+    workflow_id?: string;
   }
 
   export interface ICreateJobFromDefinition {
@@ -467,6 +470,8 @@ export namespace Scheduler {
     end_time?: number;
     downloaded: boolean;
     package_input_folder?: boolean;
+    depends_on?: string[];
+    workflow_id?: string;
   }
 
   export interface ICreateJobResponse {

src/model.ts

@@ -100,6 +100,8 @@ export interface ICreateJobModel
   // Is the create button disabled due to a submission in progress?
   createInProgress?: boolean;
   packageInputFolder?: boolean;
+  dependsOn?: string[];
+  workflowId?: string;
 }
 
 export const defaultScheduleFields: ModelWithScheduleFields = {
@@ -312,6 +314,8 @@ export interface IJobDetailModel {
   job_files: Scheduler.IJobFile[];
   downloaded: boolean;
   packageInputFolder?: boolean;
+  dependsOn?: string[];
+  workflowId?: string;
 }
 
 export interface IJobDefinitionModel {
@@ -388,7 +392,9 @@ export function convertDescribeJobtoJobDetail(
     startTime: describeJob.start_time,
     endTime: describeJob.end_time,
     downloaded: describeJob.downloaded,
-    packageInputFolder: describeJob.package_input_folder
+    packageInputFolder: describeJob.package_input_folder,
+    dependsOn: describeJob.depends_on,
+    workflowId: describeJob.workflow_id
   };
 }