From 319bc04331d16b019a22a9f2b1fe3044123408d9 Mon Sep 17 00:00:00 2001
From: Indrajit Bhosale <iamindrajitb@gmail.com>
Date: Mon, 14 Oct 2024 23:41:25 -0700
Subject: [PATCH] Add CalculateAndReportLoadTime to reduce size of OnLoadFinal

---
 .../model_lifecycle.cc                        | 26 ++++++++++++-------
 .../model_lifecycle.h                         |  2 ++
 2 files changed, 18 insertions(+), 10 deletions(-)

diff --git a/src/model_repository_manager/model_lifecycle.cc b/src/model_repository_manager/model_lifecycle.cc
index d43e2c986..450e7365a 100644
--- a/src/model_repository_manager/model_lifecycle.cc
+++ b/src/model_repository_manager/model_lifecycle.cc
@@ -804,16 +804,7 @@ ModelLifeCycle::OnLoadFinal(
       loaded.second->state_ = ModelReadyState::READY;
       loaded.second->state_reason_.clear();
 #ifdef TRITON_ENABLE_METRICS
-      auto reporter = loaded.second->model_->MetricReporter();
-      const uint64_t now_ns =
-          std::chrono::duration_cast<std::chrono::nanoseconds>(
-              std::chrono::steady_clock::now().time_since_epoch())
-              .count();
-      uint64_t time_to_load_ns = now_ns - loaded.second->load_start_ns_;
-      std::chrono::duration<double> time_to_load =
-          std::chrono::duration_cast<std::chrono::duration<double>>(
-              std::chrono::nanoseconds(time_to_load_ns));
-      ReportModelLoadTime(reporter, time_to_load);
+      CalculateAndReportLoadTime(loaded.second);
 #endif  // TRITON_ENABLE_METRICS
       auto bit = background_models_.find((uintptr_t)loaded.second);
       // Check if the version model is loaded in background, if so,
@@ -859,6 +850,21 @@ ModelLifeCycle::OnLoadFinal(
   }
 }
 
+void
+ModelLifeCycle::CalculateAndReportLoadTime(ModelInfo* loaded_model_info)
+{
+  auto reporter = loaded_model_info->model_->MetricReporter();
+  const uint64_t now_ns =
+      std::chrono::duration_cast<std::chrono::nanoseconds>(
+          std::chrono::steady_clock::now().time_since_epoch())
+          .count();
+  uint64_t time_to_load_ns = now_ns - loaded_model_info->load_start_ns_;
+  std::chrono::duration<double> time_to_load =
+      std::chrono::duration_cast<std::chrono::duration<double>>(
+          std::chrono::nanoseconds(time_to_load_ns));
+  ReportModelLoadTime(reporter, time_to_load);
+}
+
 void
 ModelLifeCycle::ReportModelLoadTime(
     std::shared_ptr<MetricModelReporter> reporter,
diff --git a/src/model_repository_manager/model_lifecycle.h b/src/model_repository_manager/model_lifecycle.h
index de4271d95..46962b19e 100644
--- a/src/model_repository_manager/model_lifecycle.h
+++ b/src/model_repository_manager/model_lifecycle.h
@@ -314,6 +314,8 @@ class ModelLifeCycle {
       ModelInfo* model_info, const bool is_update,
       const std::function<void(Status)>& OnComplete,
       std::shared_ptr<LoadTracker> load_tracker);
+  // Calculate time to load model
+  void ModelLifeCycle::CalculateAndReportLoadTime(ModelInfo* loaded_model_info);
   // Report Load time per model metrics
   void ReportModelLoadTime(
       std::shared_ptr<MetricModelReporter> reporter,