feat(azure): implement compaction request handling and logging in merge process

Author: pwh9882

collector/spot-dataset/azure/batch/merge/merge_data.py

@@ -2,6 +2,7 @@
 import os
 import argparse
 import boto3
+import json
 import pickle
 import gzip
 import gc
@@ -52,6 +53,29 @@ def _stage_log(stage: str, *, extra: str = "") -> None:
     suffix = f" {extra}" if extra else ""
     Logger.info(f"[TITANS/{PROVIDER}] {stage} rss_mb={_rss_mb():.1f}{suffix}")
 
+
+def _queue_compaction_request(
+    request_path: str,
+    *,
+    hot_key: str,
+    timestamp: datetime,
+    timeout_seconds: float,
+) -> None:
+    """Persist a compaction request for the shell to run in a fresh process."""
+    request_file = Path(request_path)
+    request_file.parent.mkdir(parents=True, exist_ok=True)
+    request_file.write_text(
+        json.dumps(
+            {
+                "provider": PROVIDER,
+                "hot_key": hot_key,
+                "timestamp": timestamp.isoformat(),
+                "timeout_seconds": timeout_seconds,
+            }
+        ),
+        encoding="utf-8",
+    )
+
 def merge_if_saving_price_sps_df(price_saving_if_df, sps_df, az=True):
     # Ensure join keys are present and types match
     join_df = pd.merge(price_saving_if_df, sps_df, on=['InstanceTier', 'InstanceType', 'Region'], how='outer')
@@ -469,13 +493,30 @@ def main():
                             extra=f"elapsed_s={time.time() - hot_started:.2f} hot_key={hot_key}",
                         )
                         if hot_key:
-                            compact_started = time.time()
-                            _stage_log("run_compaction start", extra=f"hot_key={hot_key}")
-                            run_compaction(hot_key, ts_utc, provider=PROVIDER, timeout_seconds=30.0, s3_client=titans_s3)
-                            _stage_log(
-                                "run_compaction end",
-                                extra=f"elapsed_s={time.time() - compact_started:.2f} hot_key={hot_key}",
-                            )
+                            request_path = os.environ.get("TITANS_COMPACTION_REQUEST_PATH", "").strip()
+                            if request_path:
+                                _stage_log(
+                                    "run_compaction handoff start",
+                                    extra=f"hot_key={hot_key} request_path={request_path}",
+                                )
+                                _queue_compaction_request(
+                                    request_path,
+                                    hot_key=hot_key,
+                                    timestamp=ts_utc,
+                                    timeout_seconds=30.0,
+                                )
+                                _stage_log(
+                                    "run_compaction handoff end",
+                                    extra=f"hot_key={hot_key} request_path={request_path}",
+                                )
+                            else:
+                                compact_started = time.time()
+                                _stage_log("run_compaction start", extra=f"hot_key={hot_key}")
+                                run_compaction(hot_key, ts_utc, provider=PROVIDER, timeout_seconds=30.0, s3_client=titans_s3)
+                                _stage_log(
+                                    "run_compaction end",
+                                    extra=f"elapsed_s={time.time() - compact_started:.2f} hot_key={hot_key}",
+                                )
                         _stage_log("success")
                         Logger.info(f"[TITANS/{PROVIDER}] Successfully uploaded")
                     else:

collector/spot-dataset/azure/batch/scripts/run_collection.sh

@@ -2,6 +2,8 @@
 set -e
 
 TIMESTAMP=$1
+SAFE_TIMESTAMP=$(printf '%s' "$TIMESTAMP" | tr -c 'A-Za-z0-9_.-' '_')
+COMPACTION_REQUEST_FILE="/tmp/titans_compaction_${SAFE_TIMESTAMP}.json"
 
 # Capture start time
 START_TIME_READABLE=$(date "+%Y-%m-%d %H:%M:%S")
@@ -86,7 +88,15 @@ if [ $STATUS_SPS -eq 0 ] && [ $STATUS_IF -eq 0 ] && [ $STATUS_PRICE -eq 0 ]; the
         echo "Found SPS Key: $SPS_KEY"
 
         echo "Starting Merge Job..."
-        python3 collector/spot-dataset/azure/batch/merge/merge_data.py --sps_key "$SPS_KEY"
+        rm -f "$COMPACTION_REQUEST_FILE"
+        TITANS_COMPACTION_REQUEST_PATH="$COMPACTION_REQUEST_FILE" \
+            python3 collector/spot-dataset/azure/batch/merge/merge_data.py --sps_key "$SPS_KEY"
+
+        if [ -f "$COMPACTION_REQUEST_FILE" ]; then
+            echo "Starting TITANS Warm Compaction Job..."
+            python3 collector/titans_common/run_compaction_request.py --request "$COMPACTION_REQUEST_FILE"
+            rm -f "$COMPACTION_REQUEST_FILE"
+        fi
 
         MERGE_END_TIME_EPOCH=$(date +%s)
         MERGE_DURATION=$((MERGE_END_TIME_EPOCH - COLLECTION_END_TIME_EPOCH))