From 87b3c2798629f07d622d9ffa707a6feac3eae369 Mon Sep 17 00:00:00 2001
From: Michael Duong <michaelduong22@gmail.com>
Date: Mon, 3 Nov 2025 23:11:26 +0000
Subject: [PATCH] clean up prediction app and add logging

---
 deployment/handlers/prediction_app.py | 46 +++++++++++++++++----------
 1 file changed, 29 insertions(+), 17 deletions(-)

diff --git a/deployment/handlers/prediction_app.py b/deployment/handlers/prediction_app.py
index 9d98882..5b9a807 100644
--- a/deployment/handlers/prediction_app.py
+++ b/deployment/handlers/prediction_app.py
@@ -47,6 +47,23 @@ def upload_dataframe_to_s3(df, bucket, s3_file_name):
         return False
 
 
+def warming_up_invocation(model_filepath: str):
+    """
+    Function to handle warm up invocations
+    """
+    import pandas as pd
+
+    model = model_factory(settings.build_model["model_type"])
+    model_filepath = settings.build_model["model_save_filepath"]
+    model.load_model(model_filepath)
+
+    warmup_df = pd.DataFrame(columns=model.model.original_features)
+    warmup_df = pd.concat([warmup_df.T, pd.DataFrame([0] * len(warmup_df.T))], axis=1).T
+    warmup_df.fillna(0, inplace=True)
+
+    model.predict(data=warmup_df)
+
+
 def handler(event, context):
     """
     Take in event and trigger the prediction pipeline
@@ -77,32 +94,27 @@ def handler(event, context):
 
         model_filepath = build_model_params["model_save_filepath"]
 
-        if "testing" in body:
-            storage_filepath = body["file_location"].replace(
-                ".parquet", "_output.parquet"
-            )
-        elif "warm" in body:
-            logger.info("Warm up invocation - skipping prediction")
+        if "warm" in body:
+            logger.info("Warm up invocation - synthetic prediction")
 
-            import pandas as pd
+            warming_up_invocation(model_filepath=model_filepath)
 
-            model.load_model(model_filepath)
-
-            warmup_df = pd.DataFrame(columns=model.model.original_features)
-            warmup_df = pd.concat(
-                [warmup_df.T, pd.DataFrame([0] * len(warmup_df.T))], axis=1
-            ).T
-            warmup_df.fillna(0, inplace=True)
-
-            model.predict(data=warmup_df)
             return {
                 "statusCode": 200,
                 "body": json.dumps(
                     {
-                        "message": "Successfully processed warm up invocation",
+                        "message": "Successfully warmed up invocation",
                     }
                 ),
             }
+
+        if "testing" in body:
+            logger.info(
+                "Testing invocation for CI/CD - save file to same location in S3"
+            )
+            storage_filepath = body["file_location"].replace(
+                ".parquet", "_output.parquet"
+            )
         else:
             storage_filepath = f"s3://{PREDICTIONS_BUCKET}/{portfolio_id}/{property_id}/{created_at}.parquet"