Merge pull request #620 from ddps-lab/titans

Titans 기능 적용

Author: ty-kim7

collector/run_cleanup_orphans.sh

@@ -0,0 +1,21 @@
+#!/bin/bash
+# TITANS Orphan Cleanup Script
+# orphan = manifest에 없는 S3 잔여 파일 (IAM 권한 누락 등으로 삭제 실패 시 누적)
+
+# --- Test 환경 ---
+
+# Dry run (기본값, 삭제하지 않고 목록만 출력)
+TITANS_ENV=test uv run python -m titans_common.cleanup_orphans --year 2026 --month 2 --profile spotrank
+
+# 실제 삭제
+# TITANS_ENV=test uv run python -m titans_common.cleanup_orphans --year 2026 --month 2 --profile spotrank --execute
+
+# --- Production 환경 ---
+
+# TITANS_ENV=production uv run python -m titans_common.cleanup_orphans --year 2026 --month 2 --profile spotrank
+# TITANS_ENV=production uv run python -m titans_common.cleanup_orphans --year 2026 --month 2 --profile spotrank --execute
+
+# --- 특정 provider ---
+
+# TITANS_ENV=test uv run python -m titans_common.cleanup_orphans --year 2026 --month 2 --provider azure --profile spotrank
+# TITANS_ENV=test uv run python -m titans_common.cleanup_orphans --year 2026 --month 2 --provider gcp --profile spotrank

collector/spot-dataset/aws/batch-test/Dockerfile

@@ -52,11 +52,16 @@ RUN pip install --no-cache-dir \
     requests \
     azure-identity \
     azure-core \
-    scikit-learn
+    scikit-learn \
+    polars>=1.37.0 \
+    pyarrow
 
 # Copy utility module (shared by both AWS and Azure)
 COPY utility /app/utility
 
+# Copy titans_common module (TITANS Hot/Warm tier integration)
+COPY collector/titans_common /app/collector/titans_common
+
 # Copy AWS test collector code
 COPY collector/spot-dataset/aws/batch-test /app/collector/spot-dataset/aws/batch-test
 RUN chmod +x /app/collector/spot-dataset/aws/batch-test/scripts/run_collection.sh
@@ -67,6 +72,7 @@ RUN chmod +x /app/collector/spot-dataset/azure/batch-test/scripts/run_collection
 
 # Set PYTHONPATH to include /app so imports work
 ENV PYTHONPATH=/app
+ENV PANDAS_FUTURE_INFER_STRING=0
 
 # Default entrypoint (can be overridden)
 CMD ["python3", "--version"]

collector/spot-dataset/aws/batch-test/infrastructure/main.tf

@@ -102,6 +102,7 @@ resource "aws_iam_policy" "batch_job_policy" {
   policy = jsonencode({
     Version = "2012-10-17"
     Statement = [
+      # Primary bucket (read) and test bucket (write)
       {
         Effect = "Allow"
         Action = [
@@ -112,7 +113,23 @@ resource "aws_iam_policy" "batch_job_policy" {
         ]
         Resource = [
           "arn:aws:s3:::${var.s3_bucket}",
-          "arn:aws:s3:::${var.s3_bucket}/*"
+          "arn:aws:s3:::${var.s3_bucket}/*",
+          "arn:aws:s3:::${var.s3_bucket}-test",
+          "arn:aws:s3:::${var.s3_bucket}-test/*"
+        ]
+      },
+      # TITANS Hot/Warm tier bucket (parquet storage)
+      {
+        Effect = "Allow"
+        Action = [
+          "s3:GetObject",
+          "s3:PutObject",
+          "s3:DeleteObject",
+          "s3:ListBucket"
+        ]
+        Resource = [
+          "arn:aws:s3:::${var.titans_bucket}",
+          "arn:aws:s3:::${var.titans_bucket}/*"
         ]
       },
       {

collector/spot-dataset/aws/batch-test/infrastructure/variables.tf

@@ -25,6 +25,12 @@ variable "s3_bucket" {
   default     = "spotlake"
 }
 
+variable "titans_bucket" {
+  description = "S3 Bucket for TITANS Hot/Warm tier parquet data"
+  type        = string
+  default     = "titans-spotlake-data"
+}
+
 variable "image_uri" {
   description = "Docker Image URI for Batch Jobs"
   type        = string

collector/spot-dataset/aws/batch-test/merge/compare_data.py

@@ -34,8 +34,8 @@ def compare(previous_df, current_df, workload_cols, feature_cols):
                 prev_idx += 1
                 continue
             else:
-                send_slack_message(f"{prev_workload}, {curr_workload} workload error")
-                print(f"{prev_workload}, {curr_workload} workload error")
+                send_slack_message(f"{prev_workload} workload error (current array exhausted)")
+                print(f"{prev_workload} workload error (current array exhausted)")
                 raise Exception("workload error")
         elif prev_idx == len(previous_indices):
             curr_workload = current_values[curr_idx][0]
@@ -44,8 +44,8 @@ def compare(previous_df, current_df, workload_cols, feature_cols):
                 curr_idx += 1
                 continue
             else:
-                send_slack_message(f"{prev_workload}, {curr_workload} workload error")
-                print(f"{prev_workload}, {curr_workload} workload error")
+                send_slack_message(f"{curr_workload} workload error (previous array exhausted)")
+                print(f"{curr_workload} workload error (previous array exhausted)")
                 raise Exception("workload error")
 
         prev_workload = previous_values[prev_idx][0]
@@ -135,7 +135,7 @@ def compare_max_instance(previous_df, new_df, target_capacity):
 
     # Convert to int
     for col in ["SPS", "T2", "T3"]:
-        merged_df[col] = merged_df[col].astype("Int64")
+        merged_df[col] = merged_df[col].astype("int64")
 
     # Drop unnecessary columns
     merged_df.drop(columns=["T3_prev", "T2_prev", "SPS_prev"], inplace=True)

collector/spot-dataset/aws/batch-test/merge/merge_data.py

@@ -6,6 +6,22 @@
 import json
 import pandas as pd
 import argparse
+import os
+import sys
+from pathlib import Path
+
+# TITANS environment setup (Test) - set before imports
+os.environ.setdefault("TITANS_ENV", "test")
+
+# Add titans_common path (merge -> batch-test -> aws -> spot-dataset -> collector)
+COLLECTOR_ROOT = Path(__file__).resolve().parents[4]
+sys.path.insert(0, str(COLLECTOR_ROOT))
+
+from titans_common.upload_titans import upload_hot_tier
+from titans_common.warm_compactor import run_compaction, ConcurrencyConflictError
+from titans_common.utils import prepare_for_upload
+
+PROVIDER = "aws"  # Provider constant
 
 # ------ import user module ------
 from utility.slack_msg_sender import send_slack_message
@@ -235,6 +251,29 @@ def main():
         end_time = datetime.now(timezone.utc)
         print(f"Uploading time to TSDB is {(end_time - start_time).total_seconds() * 1000 / 60000:.2f} min")
 
+        # ------ TITANS Hot tier upload + Warm compaction (Test environment) ------
+        try:
+            # Merge changed_df + removed_df (with Ceased column alignment)
+            # Note: -1 sentinel values are NOT filtered — same behavior as Timestream's
+            # upload_timestream(data.dropna()). The fillna(-1) values pass through to
+            # TITANS, matching Timestream TSDB semantics.
+            combined_df = prepare_for_upload(changed_df, removed_df, pk_columns=workload_cols)
+
+            # Ensure timezone-aware (TIMESTAMP must be timezone-aware)
+            ts_utc = TIMESTAMP if TIMESTAMP.tzinfo else TIMESTAMP.replace(tzinfo=timezone.utc)
+
+            if not combined_df.empty:
+                titans_s3 = boto3.client("s3")
+                hot_key = upload_hot_tier(combined_df, ts_utc, provider=PROVIDER, s3_client=titans_s3)
+                if hot_key:
+                    run_compaction(hot_key, ts_utc, provider=PROVIDER, timeout_seconds=30.0, s3_client=titans_s3)
+                print(f"[TITANS/{PROVIDER}/TEST] Successfully uploaded to test environment")
+
+        except ConcurrencyConflictError as e:
+            print(f"[TITANS/{PROVIDER}/TEST] Concurrency conflict, will retry next cycle: {e}")
+        except Exception as e:
+            print(f"[TITANS/{PROVIDER}/TEST] Failed (non-fatal): {e}")
+
         # ------ Upload Spotlake Query Selector to S3 ------
         start_time = datetime.now(timezone.utc)
         update_query_selector(changed_df)

collector/spot-dataset/aws/batch-test/merge/upload_data.py

@@ -7,6 +7,7 @@
 
 # ------ import user module ------
 from utility.utils import get_region
+from utility.slack_msg_sender import send_slack_message
 
 BUCKET_NAME = "spotlake-test"
 S3_PATH_PREFIX = "rawdata/aws"