From c11b2f0a9468d8e535604111b50d8f91c055baaf Mon Sep 17 00:00:00 2001
From: Brandon Thomas <echelon@gmail.com>
Date: Sun, 23 Jul 2023 21:46:26 -0400
Subject: [PATCH] better logging

---
 .../src/job/job_loop/process_single_job.rs    | 42 ++++++++++++++-----
 1 file changed, 31 insertions(+), 11 deletions(-)

diff --git a/crates/service/job/inference_job/src/job/job_loop/process_single_job.rs b/crates/service/job/inference_job/src/job/job_loop/process_single_job.rs
index e44f1e78ef..c598c327ce 100644
--- a/crates/service/job/inference_job/src/job/job_loop/process_single_job.rs
+++ b/crates/service/job/inference_job/src/job/job_loop/process_single_job.rs
@@ -19,17 +19,6 @@ pub async fn process_single_job(
   job_dependencies: &JobDependencies,
   job: &AvailableInferenceJob,
 ) -> Result<ProcessSingleJobSuccessCase, ProcessSingleJobError> {
-
-  let mut maybe_keepalive_redis = job_dependencies
-      .maybe_keepalive_redis_pool
-      .as_ref()
-      .map(|redis| redis.get())
-      .transpose()
-      .map_err(|err| ProcessSingleJobError::Other(anyhow!("redis pool error: {:?}", err)))?;
-
-  // TODO(bt, 2023-01-11): Restore an optional status logger
-  //let mut redis_logger = RedisJobStatusLogger::new_generic_download(&mut redis, job.download_job_token.as_str());
-
   let mut force_execution = false;
 
   // Some jobs have "routing tags". These ensure that jobs only execute on certain hosts.
@@ -84,6 +73,14 @@ pub async fn process_single_job(
     return Ok(ProcessSingleJobSuccessCase::LockNotObtained)
   }
 
+  process_single_job_wrap_with_logs(job_dependencies, job).await
+}
+
+async fn process_single_job_wrap_with_logs(
+  job_dependencies: &JobDependencies,
+  job: &AvailableInferenceJob,
+) -> Result<ProcessSingleJobSuccessCase, ProcessSingleJobError> {
+
   println!("\n  ----------------------------------------- JOB START -----------------------------------------  \n");
 
   info!("Beginning work on job ({}): {:?}", job.id.0, job.inference_job_token);
@@ -91,6 +88,29 @@ pub async fn process_single_job(
   info!("Job model type: {:?}", job.maybe_model_type);
   info!("Job model token: {:?}", job.maybe_model_token);
 
+  let result = do_process_single_job(job_dependencies, job).await;
+
+  println!("\n  ----------------------------------------- JOB END -----------------------------------------  \n");
+
+  result
+}
+
+async fn do_process_single_job(
+  job_dependencies: &JobDependencies,
+  job: &AvailableInferenceJob,
+) -> Result<ProcessSingleJobSuccessCase, ProcessSingleJobError> {
+
+  // TODO(bt, 2023-07-23): Redis pool management probably belongs at near the outermost loop.
+  let mut maybe_keepalive_redis = job_dependencies
+      .maybe_keepalive_redis_pool
+      .as_ref()
+      .map(|redis| redis.get())
+      .transpose()
+      .map_err(|err| ProcessSingleJobError::Other(anyhow!("redis pool error: {:?}", err)))?;
+
+  // TODO(bt, 2023-01-11): Restore an optional status logger
+  //let mut redis_logger = RedisJobStatusLogger::new_generic_download(&mut redis, job.download_job_token.as_str());
+
   let job_start_time = Instant::now();
 
   // ==================== HANDLE KEEPALIVE (OPTIONAL) ==================== //