feat: expose submit and cancel job methods as public in scheduler (#1260)

* expose schedule and cancel job methods as public

* address code review comment

Author: milenkovicm

ballista/scheduler/src/scheduler_server/grpc.rs

@@ -457,14 +457,13 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerGrpc
                 plan.display_indent()
             );
 
-            let job_id = self.state.task_manager.generate_job_id();
-
             log::trace!("setting job name: {}", job_name);
-            self.submit_job(&job_id, &job_name, session_ctx, &plan)
+            let job_id = self
+                .submit_job(&job_name, session_ctx, &plan)
                 .await
                 .map_err(|e| {
                     let msg =
-                        format!("Failed to send JobQueued event for {job_id}: {e:?}");
+                        format!("Failed to send JobQueued event for {job_name}: {e:?}");
                     error!("{}", msg);
 
                     Status::internal(msg)
@@ -534,20 +533,12 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerGrpc
         let job_id = request.into_inner().job_id;
         info!("Received cancellation request for job {}", job_id);
 
-        self.query_stage_event_loop
-            .get_sender()
-            .map_err(|e| {
-                let msg = format!("Get query stage event loop error due to {e:?}");
-                error!("{}", msg);
-                Status::internal(msg)
-            })?
-            .post_event(QueryStageSchedulerEvent::JobCancel(job_id))
-            .await
-            .map_err(|e| {
-                let msg = format!("Post to query stage event loop error due to {e:?}");
-                error!("{}", msg);
-                Status::internal(msg)
-            })?;
+        self.cancel_job(job_id).await.map_err(|e| {
+            let msg = format!("Post to query stage event loop error due to {e:?}");
+            error!("{}", msg);
+            Status::internal(msg)
+        })?;
+
         Ok(Response::new(CancelJobResult { cancelled: true }))
     }
 

ballista/scheduler/src/scheduler_server/mod.rs

@@ -163,14 +163,28 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerServer<T
     pub(crate) fn metrics_collector(&self) -> &dyn SchedulerMetricsCollector {
         self.query_stage_scheduler.metrics_collector()
     }
+    /// Cancels job for given job_id
+    pub async fn cancel_job(&self, job_id: String) -> Result<()> {
+        log::debug!("Received cancellation request for job {}", job_id);
 
-    pub(crate) async fn submit_job(
+        self.query_stage_event_loop
+            .get_sender()?
+            .post_event(QueryStageSchedulerEvent::JobCancel(job_id.to_owned()))
+            .await?;
+
+        Ok(())
+    }
+
+    /// Submits a job to executor returning job_id
+    pub async fn submit_job(
         &self,
-        job_id: &str,
         job_name: &str,
         ctx: Arc<SessionContext>,
         plan: &LogicalPlan,
-    ) -> Result<()> {
+    ) -> Result<String> {
+        log::debug!("Received submit request for job {}", job_name);
+        let job_id = self.state.task_manager.generate_job_id();
+
         self.query_stage_event_loop
             .get_sender()?
             .post_event(QueryStageSchedulerEvent::JobQueued {
@@ -180,7 +194,9 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerServer<T
                 plan: Box::new(plan.clone()),
                 queued_at: timestamp_millis(),
             })
-            .await
+            .await?;
+
+        Ok(job_id)
     }
 
     /// It just send task status update event to the channel,
@@ -505,7 +521,7 @@ mod test {
         )
         .await?;
 
-        let status = test.run("job", "", &plan).await.expect("running plan");
+        let (status, job_id) = test.run("", &plan).await.expect("running plan");
 
         match status.status {
             Some(job_status::Status::Successful(SuccessfulJob {
@@ -519,8 +535,8 @@ mod test {
             }
         }
 
-        assert_submitted_event("job", &metrics_collector);
-        assert_completed_event("job", &metrics_collector);
+        assert_submitted_event(&job_id, &metrics_collector);
+        assert_completed_event(&job_id, &metrics_collector);
 
         Ok(())
     }
@@ -580,7 +596,7 @@ mod test {
         )
         .await?;
 
-        let status = test.run("job", "", &plan).await.expect("running plan");
+        let (status, job_id) = test.run("", &plan).await.expect("running plan");
 
         assert!(
             matches!(
@@ -594,8 +610,8 @@ mod test {
             "Expected job status to be failed but it was {status:?}"
         );
 
-        assert_submitted_event("job", &metrics_collector);
-        assert_failed_event("job", &metrics_collector);
+        assert_submitted_event(&job_id, &metrics_collector);
+        assert_failed_event(&job_id, &metrics_collector);
 
         Ok(())
     }
@@ -625,7 +641,7 @@ mod test {
             .into_optimized_plan()?;
 
         // This should fail when we try and create the physical plan
-        let status = test.run("job", "", &plan).await?;
+        let (status, job_id) = test.run("", &plan).await?;
 
         assert!(
             matches!(
@@ -639,8 +655,8 @@ mod test {
             "Expected job status to be failed but it was {status:?}"
         );
 
-        assert_no_submitted_event("job", &metrics_collector);
-        assert_failed_event("job", &metrics_collector);
+        assert_no_submitted_event(&job_id, &metrics_collector);
+        assert_failed_event(&job_id, &metrics_collector);
 
         Ok(())
     }

ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs

@@ -376,9 +376,7 @@ mod tests {
         )
         .await?;
 
-        let job_id = "job-1";
-
-        test.submit(job_id, "", &plan).await?;
+        let job_id = test.submit("", &plan).await?;
 
         test.tick().await?;
 
@@ -398,7 +396,7 @@ mod tests {
             expected, running_jobs
         );
 
-        test.cancel(job_id).await?;
+        test.cancel(&job_id).await?;
 
         let expected = 0usize;
         let success = await_condition(Duration::from_millis(10), 20, || {

ballista/scheduler/src/test_utils.rs

@@ -480,12 +480,7 @@ impl SchedulerTest {
             .await
     }
 
-    pub async fn submit(
-        &mut self,
-        job_id: &str,
-        job_name: &str,
-        plan: &LogicalPlan,
-    ) -> Result<()> {
+    pub async fn submit(&mut self, job_name: &str, plan: &LogicalPlan) -> Result<String> {
         println!("{:?}", self.session_config);
         let ctx = self
             .scheduler
@@ -494,11 +489,9 @@ impl SchedulerTest {
             .create_session(&self.session_config)
             .await?;
 
-        self.scheduler
-            .submit_job(job_id, job_name, ctx, plan)
-            .await?;
+        let job_id = self.scheduler.submit_job(job_name, ctx, plan).await?;
 
-        Ok(())
+        Ok(job_id)
     }
 
     pub async fn post_scheduler_event(
@@ -606,22 +599,20 @@ impl SchedulerTest {
         final_status
     }
 
+    /// Returns job status and job_id
     pub async fn run(
         &mut self,
-        job_id: &str,
         job_name: &str,
         plan: &LogicalPlan,
-    ) -> Result<JobStatus> {
+    ) -> Result<(JobStatus, String)> {
         let ctx = self
             .scheduler
             .state
             .session_manager
             .create_session(&self.session_config)
             .await?;
 
-        self.scheduler
-            .submit_job(job_id, job_name, ctx, plan)
-            .await?;
+        let job_id = self.scheduler.submit_job(job_name, ctx, plan).await?;
 
         let mut receiver = self.status_receiver.take().unwrap();
 
@@ -640,7 +631,7 @@ impl SchedulerTest {
                 .scheduler
                 .state
                 .task_manager
-                .get_job_status(job_id)
+                .get_job_status(&job_id)
                 .await?;
 
             if let Some(JobStatus {
@@ -659,7 +650,7 @@ impl SchedulerTest {
             tokio::time::sleep(Duration::from_millis(100)).await
         };
 
-        final_status
+        final_status.map(|s| (s, job_id))
     }
 }
 
@@ -746,7 +737,7 @@ pub fn assert_submitted_event(job_id: &str, collector: &TestMetricsCollector) {
         .iter()
         .any(|ev| matches!(ev, MetricEvent::Submitted(_, _, _)));
 
-    assert!(found, "{}", "Expected submitted event for job {job_id}");
+    assert!(found, "Expected submitted event for job {job_id}");
 }
 
 pub fn assert_no_submitted_event(job_id: &str, collector: &TestMetricsCollector) {
@@ -755,7 +746,7 @@ pub fn assert_no_submitted_event(job_id: &str, collector: &TestMetricsCollector)
         .iter()
         .any(|ev| matches!(ev, MetricEvent::Submitted(_, _, _)));
 
-    assert!(!found, "{}", "Expected no submitted event for job {job_id}");
+    assert!(!found, "Expected no submitted event for job {job_id}");
 }
 
 pub fn assert_completed_event(job_id: &str, collector: &TestMetricsCollector) {
@@ -782,7 +773,7 @@ pub fn assert_failed_event(job_id: &str, collector: &TestMetricsCollector) {
         .iter()
         .any(|ev| matches!(ev, MetricEvent::Failed(_, _, _)));
 
-    assert!(found, "{}", "Expected failed event for job {job_id}");
+    assert!(found, "Expected failed event for job {job_id}");
 }
 
 pub fn revive_graph_and_complete_next_stage(graph: &mut ExecutionGraph) -> Result<usize> {