ML/AI 성능 최적화와 GPU 사이징

# Model Serving 엔드포인트 설정 (최적화 버전)
import mlflow

# 모델 등록 시 최적화 힌트
mlflow.pyfunc.log_model(
    artifact_path="model",
    python_model=my_model,
    pip_requirements=["scikit-learn==1.3.0"],  # 최소 의존성
    metadata={"serving_optimization": "latency"}
)

import mlflow
from pyspark.sql.functions import struct

# Spark UDF로 배치 추론 병렬화
model_uri = "models:/fraud_detection/production"
predict_udf = mlflow.pyfunc.spark_udf(spark, model_uri)

# 대용량 데이터셋에 병렬 추론 적용
predictions = (
    spark.table("catalog.schema.transactions")
    .withColumn("prediction", predict_udf(struct("feature1", "feature2", "feature3")))
)

# 최적화 팁: 파티션 수를 코어 수의 2~3배로 설정
predictions.repartition(spark.sparkContext.defaultParallelism * 2)

# Vector Search 쿼리 최적화
results = index.similarity_search(
    query_text="고객 이탈 예측 방법",
    columns=["content", "source", "page"],  # 필요한 컬럼만 반환
    filters={"category": "ml", "year": 2025},  # 메타데이터 필터로 범위 축소
    num_results=5  # 필요한 만큼만 요청
)

# GPU 사용률 확인 (노트북에서)
import subprocess
result = subprocess.run(['nvidia-smi', '--query-gpu=utilization.gpu,memory.used,memory.total',
                        '--format=csv,noheader,nounits'], capture_output=True, text=True)
print(result.stdout)

# GPU 사용률 지속 모니터링 (학습 중)
# nvidia-smi dmon -s u -d 5  # 5초 간격으로 사용률 출력

from pyspark.ml.torch.distributor import TorchDistributor

def train_fn():
    import torch
    import torch.distributed as dist
    from torch.nn.parallel import DistributedDataParallel as DDP
    
    # 분산 환경 초기화 (TorchDistributor가 자동 설정)
    local_rank = int(os.environ.get("LOCAL_RANK", 0))
    device = torch.device(f"cuda:{local_rank}")
    
    model = MyModel().to(device)
    model = DDP(model, device_ids=[local_rank])
    
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
    
    for epoch in range(10):
        for batch in train_loader:
            loss = model(batch)
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()
    
    return model

# Databricks에서 분산 학습 실행
distributor = TorchDistributor(
    num_processes=4,     # 총 GPU 수
    local_mode=False,    # 멀티 노드
    use_gpu=True
)
model = distributor.run(train_fn)

# Mixed Precision Training (메모리 절감 + 속도 향상)
from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

for batch in train_loader:
    optimizer.zero_grad()
    
    with autocast():  # FP16으로 연산
        output = model(batch)
        loss = criterion(output, labels)
    
    scaler.scale(loss).backward()  # 그래디언트 스케일링
    scaler.step(optimizer)
    scaler.update()

# Gradient Accumulation (큰 배치 효과를 작은 GPU로)
accumulation_steps = 4
for i, batch in enumerate(train_loader):
    loss = model(batch) / accumulation_steps
    loss.backward()
    
    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

# 단일 노드: Pandas + scikit-learn (데이터 < 100GB)
import pandas as pd
from sklearn.ensemble import GradientBoostingClassifier
import mlflow

# Spark 테이블을 Pandas로 변환 (데이터가 메모리에 들어갈 때)
df = spark.table("catalog.schema.features").toPandas()

with mlflow.start_run():
    model = GradientBoostingClassifier(n_estimators=100)
    model.fit(df[features], df[target])
    mlflow.sklearn.log_model(model, "model")

# 분산: Spark ML (데이터 > 100GB)
from pyspark.ml.classification import GBTClassifier
from pyspark.ml.feature import VectorAssembler

assembler = VectorAssembler(inputCols=features, outputCol="features")
gbt = GBTClassifier(featuresCol="features", labelCol="target", maxIter=100)

pipeline = Pipeline(stages=[assembler, gbt])
model = pipeline.fit(spark.table("catalog.schema.features"))

# 대용량 데이터에서 Feature Engineering은 Spark로,
# 학습은 단일 노드로 하는 하이브리드 패턴

# Step 1: Spark로 대규모 Feature Engineering (분산)
features_df = (
    spark.table("catalog.schema.raw_data")  # 1TB
    .withColumn("feature_1", ...)
    .withColumn("feature_2", ...)
    .groupBy("entity_id")
    .agg(...)
)
features_df.write.saveAsTable("catalog.schema.features")  # 50GB

# Step 2: 학습용 샘플링 (50GB → 5GB)
sample_df = features_df.sample(fraction=0.1).toPandas()

# Step 3: 단일 노드에서 학습 (최적화된 라이브러리 활용)
import xgboost as xgb
model = xgb.XGBClassifier(tree_method="gpu_hist", gpu_id=0)
model.fit(sample_df[feature_cols], sample_df[target_col])

# Step 4: 전체 데이터에 Spark UDF로 추론 (분산)
predict_udf = mlflow.pyfunc.spark_udf(spark, "models:/my_model/production")
predictions = features_df.withColumn("prediction", predict_udf(struct(*feature_cols)))