Spark + LLM 연동 가이드 - 대규모 데이터에 AI 적용하기

from pyspark.sql import SparkSession
from pyspark.sql.functions import udf, col
from pyspark.sql.types import StringType
import requests
 
spark = SparkSession.builder.appName("SparkLLM").getOrCreate()
 
# LLM 호출 UDF
@udf(returnType=StringType())
def classify_sentiment(text):
    """텍스트 감성 분류 (Ollama 로컬 서버)"""
    if not text:
        return "unknown"
    try:
        response = requests.post(
            "http://ollama-server:11434/api/generate",
            json={
                "model": "llama3.1:8b",
                "prompt": f'다음 리뷰의 감성을 "긍정", "부정", "중립" 중 하나로 분류하세요. 답만 쓰세요.\n\n리뷰: {text}\n\n감성:',
                "stream": False,
                "options": {"temperature": 0.0, "num_predict": 10}
            },
            timeout=30
        )
        return response.json()["response"].strip()
    except Exception as e:
        return f"error: {str(e)}"
 
# 적용
df = spark.read.parquet("s3://data-lake/reviews/")
result = df.withColumn("sentiment", classify_sentiment(col("review_text")))
result.write.parquet("s3://data-lake/reviews_classified/")

from pyspark.sql import Row
import requests
 
def process_partition(partition):
    """파티션 단위 배치 처리 (연결 재사용)"""
    session = requests.Session()  # 커넥션 풀 재사용
    
    batch = []
    results = []
    
    for row in partition:
        batch.append(row)
        
        if len(batch) >= 10:  # 10건씩 배치
            prompts = [r.review_text for r in batch]
            classifications = batch_classify(session, prompts)
            
            for r, cls in zip(batch, classifications):
                results.append(Row(id=r.id, review_text=r.review_text, sentiment=cls))
            batch = []
    
    # 나머지 처리
    if batch:
        prompts = [r.review_text for r in batch]
        classifications = batch_classify(session, prompts)
        for r, cls in zip(batch, classifications):
            results.append(Row(id=r.id, review_text=r.review_text, sentiment=cls))
    
    session.close()
    return iter(results)
 
def batch_classify(session, texts):
    """vLLM 배치 API 호출"""
    response = session.post(
        "http://vllm-server:8000/v1/chat/completions",
        json={
            "model": "meta-llama/Llama-3.1-8B-Instruct",
            "messages": [{"role": "user", "content": f"Classify sentiments: {texts}"}],
            "temperature": 0.0
        }
    )
    return parse_classifications(response.json())
 
# 적용
result_rdd = df.rdd.mapPartitions(process_partition)
result_df = spark.createDataFrame(result_rdd)

from sentence_transformers import SentenceTransformer
from pyspark.sql.types import ArrayType, FloatType
import numpy as np
 
# 브로드캐스트 모델 (각 Executor에서 1번만 로드)
model_broadcast = None
 
def get_model():
    global model_broadcast
    if model_broadcast is None:
        model_broadcast = SentenceTransformer("BAAI/bge-m3")
    return model_broadcast
 
@udf(returnType=ArrayType(FloatType()))
def generate_embedding(text):
    """텍스트를 벡터로 변환"""
    if not text:
        return None
    model = get_model()
    embedding = model.encode(text, normalize_embeddings=True)
    return embedding.tolist()
 
# 대규모 임베딩 생성
df = spark.read.parquet("s3://data-lake/documents/")
embedded = df.withColumn("embedding", generate_embedding(col("content")))
 
# 벡터 DB에 적재할 수 있는 형태로 저장
embedded.write.parquet("s3://data-lake/embeddings/", mode="overwrite")

def embed_partition(partition):
    """파티션 단위 배치 임베딩 (훨씬 효율적)"""
    model = SentenceTransformer("BAAI/bge-m3")
    
    rows = list(partition)
    if not rows:
        return iter([])
    
    texts = [r.content for r in rows]
    
    # 배치로 한 번에 임베딩 (GPU 활용 극대화)
    embeddings = model.encode(texts, batch_size=64, normalize_embeddings=True)
    
    results = []
    for row, emb in zip(rows, embeddings):
        results.append(Row(id=row.id, content=row.content, embedding=emb.tolist()))
    
    return iter(results)
 
# 적용 (파티션 수 = GPU 수에 맞춰 조정)
df = df.repartition(8)  # 8개 파티션 = 8개 GPU
result = spark.createDataFrame(df.rdd.mapPartitions(embed_partition))

# 1. 감성 분류 + 핵심 키워드 추출 + 요약
pipeline_prompt = """다음 고객 리뷰를 분석하세요.
 
리뷰: {review}
 
JSON으로 반환:
{{"sentiment": "긍정/부정/중립", "keywords": ["키워드1", "키워드2"], "summary": "한줄 요약", "category": "배송/품질/가격/서비스/기타"}}
"""

# 서버 로그를 구조화된 데이터로 변환
log_prompt = """다음 로그 라인을 분석하여 JSON으로 변환하세요.
 
로그: {log_line}
 
{{"timestamp": "", "level": "", "service": "", "message": "", "error_type": "", "stack_trace": ""}}
"""