26Y02m4

"""
[2단계-Polars 완치판] 멀티 코어 기반 일별 자원 마감 정산 및 AIStor 전송 파이프라인
"""
import os
import argparse
import boto3
from botocore.exceptions import ClientError
import polars as pl
from datetime import datetime, timedelta, timezone
from pathlib import Path

import config
from config import (
    RAW_DIR, MERGED_DIR, CLUSTER_NODE_PATTERNS, 
    get_workload_type, classify_node_infrastructure
)

KST = timezone(timedelta(hours=9))
W_CPU, W_MEM, W_PV = 1.0, 0.11, 0.02

def parse_arguments():
    parser = argparse.ArgumentParser(description="FinOps Polars-Powered Reprocessing Engine")
    parser.add_argument("--days", type=int, default=None)
    parser.add_argument("--start-date", type=str, default=None)
    parser.add_argument("--end-date", type=str, default=None)
    parser.add_argument("--cluster", type=str, default="ALL", choices=["COMPUTE", "STORAGE", "ALL"])
    parser.add_argument("--workload-domain", type=str, default=None)
    return parser.parse_args()

def get_minio_client():
    return boto3.client(
        "s3", endpoint_url=os.getenv("MINIO_ENDPOINT"), aws_access_key_id=os.getenv("MINIO_ACCESS_KEY"),
        aws_secret_access_key=os.getenv("MINIO_SECRET_KEY"), region_name="us-east-1",
        config=boto3.session.Config(signature_version="s3v4")
    )

def sync_down_raw_from_minio(s3_client, bucket_name, start_dt, end_dt):
    chunk_delta = timedelta(hours=1)
    current_start = start_dt
    print(f"\n📡 [AIStor 스마트 캐시 엔진] 원천 데이터 가용성 체크 시작...")

    while current_start < end_dt:
        chunk_str = current_start.strftime("%Y%m%d_%H")
        f_name = f"prom_raw_{chunk_str}.parquet"
        object_key = f"raw/{f_name}"
        local_path = RAW_DIR / f_name
        
        if local_path.exists() and local_path.stat().st_size > 0:
            print(f"   ⚡ [Cache Hit] 로컬 캐시 자산 활용 (다운로드 스킵): {f_name}")
        else:
            try:
                s3_client.head_object(Bucket=bucket_name, Key=object_key)
                print(f"   📥 [Cache Miss] AIStor 백업으로부터 다운로드: {object_key}")
                s3_client.download_file(bucket_name, object_key, str(local_path))
            except ClientError as e:
                if e.response['Error']['Code'] != "404":
                    print(f"   ❌ AIStor 통신 실패 에러 ({f_name}): {str(e)}")
            except Exception as e:
                print(f"   ⚠️ 일반 다운로드 예외 발생 ({f_name}): {str(e)}")
                
        current_start += chunk_delta

def main():
    args = parse_arguments()
    s3_client = get_minio_client()
    bucket_name = os.getenv("MINIO_RAW_BUCKET", "devops-test")
    now_kst = datetime.now(KST)
    
    if args.start_date and args.end_date:
        start_dt = datetime.strptime(args.start_date, "%Y-%m-%d").replace(tzinfo=KST)
        end_dt = (datetime.strptime(args.end_date, "%Y-%m-%d") + timedelta(days=1)).replace(tzinfo=KST)
    elif args.days:
        start_dt = (now_kst - timedelta(days=args.days)).replace(hour=0, minute=0, second=0, microsecond=0)
        end_dt = now_kst
    else:
        yesterday = (now_kst - timedelta(days=1)).date()
        start_dt = datetime.combine(yesterday, datetime.min.time()).replace(tzinfo=KST)
        end_dt = datetime.combine(now_kst.date(), datetime.min.time()).replace(tzinfo=KST)

    sync_down_raw_from_minio(s3_client, bucket_name, start_dt, end_dt)

    raw_files = list(RAW_DIR.glob("prom_raw_*.parquet"))
    if not raw_files:
        print("\n🚨 [중단] 정산 가동할 물리 원천 파일이 존재하지 않습니다.")
        return

    print(f"\n📊 [Polars Lazy Engine 기동] {len(raw_files)}개 소스 파티션 스캔 링킹...")
    lf_raw = pl.scan_parquet(str(RAW_DIR / "prom_raw_*.parquet"))

    start_naive = start_dt.replace(tzinfo=None)
    end_naive = end_dt.replace(tzinfo=None)
    
    lf_filtered = lf_raw.filter(
        (pl.col("timestamp") >= start_naive) & (pl.col("timestamp") < end_naive)
    ).with_columns(
        pl.col("timestamp").dt.strftime("%Y%m%d").alias("date")
    )

    print("⏳ 대용량 인프라 정산 메트릭 다차원 분산 집계 중 (Rust Core 스레드 가동)...")
    df_raw = lf_filtered.collect()
    
    if df_raw.is_empty():
        print("⚠️  해당 기간 내 정산할 데이터 로우가 없습니다.")
        return

    df_raw = df_raw.with_columns([
        pl.col("node").map_elements(
            lambda x: {"cluster_type": classify_node_infrastructure(x)[0], "workload_domain": classify_node_infrastructure(x)[1]},
            return_dtype=pl.Struct({"cluster_type": pl.String, "workload_domain": pl.String})
        ).alias("infra_struct"),
        pl.col("pod").map_elements(get_workload_type, return_dtype=pl.String).alias("workload_type")
    ]).with_columns([
        pl.col("infra_struct").struct.field("cluster_type").alias("cluster_type"),
        pl.col("infra_struct").struct.field("workload_domain").alias("workload_domain")
    ]).drop("infra_struct")

    target_dates = df_raw["date"].unique().to_list()
    target_clusters = ["COMPUTE", "STORAGE"] if args.cluster.upper() == "ALL" else [args.cluster.upper()]

    print(f"🎯 분할 연산 타겟 스케줄러 가동 -> 일자 풀: {target_dates} | 클러스터 풀: {target_clusters}")

    GB_DIV = 1024 ** 3

    for date_chunk in target_dates:
        for cluster_chunk in target_clusters:
            
            df_slice = df_raw.filter((pl.col("date") == date_chunk) & (pl.col("cluster_type") == cluster_chunk))
            if df_slice.is_empty():
                continue

            if args.workload_domain:
                df_slice = df_slice.filter(pl.col("workload_domain") == args.workload_domain)
                if df_slice.is_empty(): continue

            print(f"   ⏳ [Polars 집계 연산] {date_chunk} ➡️ {cluster_chunk} 고성능 그룹바이 컴파일...")

            target_fields = ["cpu_request", "cpu_limit", "cpu_usage", "cpu_throttled", "mem_request", "mem_limit", "mem_usage", "mem_rss", "oom_event", "pv_capacity", "pv_used"]
            for col in target_fields:
                if col not in df_slice.columns:
                    df_slice = df_slice.with_columns(pl.lit(0.0).alias(col))

            # [테이블 1] 상세 메트릭 원부 집계
            df_pod = df_slice.group_by([
                "date", "cluster_type", "workload_domain", "namespace", "workload_type", "node", "pod", "container"
            ]).agg([
                pl.len().alias("minutes_running"),
                pl.col("cpu_request").max().alias("cpu_request_max"),
                pl.col("cpu_limit").max().alias("cpu_limit_max"),
                pl.col("cpu_usage").quantile(0.95).alias("cpu_usage_p95"),
                pl.col("cpu_throttled").max().alias("cpu_throttled_max"),
                (pl.col("mem_request").max() / GB_DIV).alias("mem_request_max"),
                (pl.col("mem_limit").max() / GB_DIV).alias("mem_limit_max"),
                (pl.col("mem_usage").quantile(0.95) / GB_DIV).alias("mem_usage_p95"),
                (pl.col("mem_rss").quantile(0.95) / GB_DIV).alias("mem_rss_p95"),
                pl.col("oom_event").sum().alias("oom_strike_sum"),
                (pl.col("pv_capacity").max() / GB_DIV).alias("pv_capacity_max"),
                (pl.col("pv_used").quantile(0.95) / GB_DIV).alias("pv_used_p95")
            ]).fill_null(0.0)

            df_pod = df_pod.with_columns([
                (pl.col("cpu_request_max") * (pl.col("minutes_running") / 60.0)).alias("cpu_allocated_core_hours"),
                (pl.col("cpu_usage_p95") * (pl.col("minutes_running") / 60.0)).alias("cpu_usage_core_hours"),
                (pl.col("mem_request_max") * (pl.col("minutes_running") / 60.0)).alias("mem_allocated_gb_hours"),
                (pl.col("mem_usage_p95") * (pl.col("minutes_running") / 60.0)).alias("mem_usage_gb_hours"),
                (pl.col("pv_capacity_max") * (pl.col("minutes_running") / 60.0)).alias("pv_allocated_gb_hours"),
                (pl.col("pv_usage_p95") * (pl.col("minutes_running") / 60.0)).alias("pv_usage_gb_hours"),
                (pl.col("oom_strike_sum") > 0).alias("is_oom_killed"),
                (pl.col("cpu_request_max") == 0).alias("has_no_request"),
                (pl.col("cpu_limit_max") == 0).alias("has_no_limit"),
                ((pl.col("cpu_usage_p95") - pl.col("cpu_request_max")).clip(lower_bound=0)).alias("cpu_shortage_cores")
            ]).with_columns([
                ((pl.col("cpu_allocated_core_hours") - pl.col("cpu_usage_core_hours")).clip(lower_bound=0)).alias("cpu_waste_core_hours"),
                ((pl.col("mem_allocated_gb_hours") - pl.col("mem_usage_gb_hours")).clip(lower_bound=0)).alias("mem_waste_gb_hours"),
                ((pl.col("pv_allocated_gb_hours") - pl.col("pv_usage_gb_hours")).clip(lower_bound=0)).alias("pv_waste_gb_hours"),
            ])

            df_pod = df_pod.with_columns(
                pl.when(pl.col("is_oom_killed")).then(pl.lit("💥 OOM장애발생"))
                .when((pl.col("cpu_shortage_cores") > 0.5) | (pl.col("cpu_throttled_max") > 0.2)).then(pl.lit("⚠️ Request부족"))
                .when((pl.col("cpu_waste_core_hours") > 10) | (pl.col("pv_waste_gb_hours") > 50)).then(pl.lit("📉 과다할당"))
                .otherwise(pl.lit("✅ 최적화완료")).alias("status")
            )

            # 📅 [테이블 2] Namespace 사용량 요약 빌드
            df_daily_ns = df_pod.group_by(["date", "namespace"]).agg([
                pl.col("cpu_usage_core_hours").sum().alias("cpu_used_ch"),
                pl.col("cpu_allocated_core_hours").sum().alias("cpu_alloc_ch"),
                pl.col("cpu_waste_core_hours").sum().alias("cpu_waste_ch"),
                pl.col("mem_usage_gb_hours").sum().alias("mem_used_gh"),
                pl.col("mem_allocated_gb_hours").sum().alias("mem_alloc_gh"),
                pl.col("pv_usage_gb_hours").sum().alias("pv_used_gh"),
                pl.col("pv_allocated_gb_hours").sum().alias("pv_alloc_gh")
            ]).with_columns(
                ((pl.col("cpu_alloc_ch") * W_CPU) + (pl.col("mem_alloc_gh") * W_MEM) + (pl.col("pv_alloc_gh") * W_PV)).round(1).alias("final_usage_score")
            )

            # 📉 [테이블 3] 파레토 원부 롤업 빌드 
            df_ns = df_pod.group_by("namespace").agg([
                pl.col("cpu_waste_core_hours").sum().alias("total_waste_core_hours"),
                pl.col("minutes_running").sum().alias("minutes_running_sum"),
                pl.col("container").count().alias("container_cnt"),
                pl.col("cpu_allocated_core_hours").sum().alias("total_allocated_core_hours")
            ]).sort("total_waste_core_hours", descending=True)
            
            global_total_waste = df_ns["total_waste_core_hours"].sum() or 0.1
            df_ns = df_ns.with_columns([
                pl.lit(date_chunk).alias("date"),
                ((pl.col("total_waste_core_hours") / global_total_waste) * 100).round(2).alias("waste_share_pct")
            ]).with_columns(
                pl.col("waste_share_pct").cum_sum().round(2).alias("waste_cumsum_pct")
            )

            # ─── 💾 [로컬 저장 명세 고정 바인딩] ───
            cluster_partition_dir = MERGED_DIR / cluster_chunk
            cluster_partition_dir.mkdir(parents=True, exist_ok=True)
            
            target_output_file = cluster_partition_dir / f"daily_enriched_{cluster_chunk}_{date_chunk}.parquet"
            ns_output_file     = cluster_partition_dir / f"daily_ns_usage_{cluster_chunk}_{date_chunk}.parquet"
            pareto_output_file = cluster_partition_dir / f"pareto_ns_{cluster_chunk}_{date_chunk}.parquet"
            
            df_pod.write_parquet(target_output_file)
            df_daily_ns.write_parquet(ns_output_file)
            df_ns.write_parquet(pareto_output_file) # ◀ 파레토 로컬 디스크 정착
            
            print(f"     💾 [로컬 저장] 파티션 3대 원부 컴파일 성료 (클러스터: {cluster_chunk})")

            # ─── 📡 [AIStor 영구 적재 가드레일 주입] ───
            print(f"     Submitting artifacts to AIStor Tables storage system...")
            try:
                s3_client.upload_file(str(target_output_file), bucket_name, f"merged/{cluster_chunk}/{target_output_file.name}")
                s3_client.upload_file(str(ns_output_file), bucket_name, f"merged/{cluster_chunk}/{ns_output_file.name}")
                s3_client.upload_file(str(pareto_output_file), bucket_name, f"merged/{cluster_chunk}/{pareto_output_file.name}")
                print(f"     ✅ [AIStor 업로드 완수] 테이블 자산 적재 완료 -> {pareto_output_file.name}")
            except Exception as e:
                print(f"     ❌ [AIStor 통신장애] 오브젝트 전송 실패: {str(e)}")

    print("\n🏁 === [Step2 Polars 배치 완료] 로컬 및 AIStor 테이블 연동이 완벽히 마감되었습니다. ===")

if __name__ == "__main__":
    main()