Databricks ML/AI 기능 전체 가이드

from databricks.feature_engineering import FeatureEngineeringClient, FeatureLookup

fe = FeatureEngineeringClient()

# 피처 룩업 정의
feature_lookups = [
    FeatureLookup(
        table_name="catalog.schema.sensor_features",
        feature_names=["avg_temperature", "vibration_rms"],
        lookup_key="device_id"
    )
]

# 학습 데이터셋 생성 (라벨 + 피처 자동 조인)
training_set = fe.create_training_set(
    df=labels_df,
    feature_lookups=feature_lookups,
    label="failure"
)

from databricks import automl

summary = automl.classify(
    dataset=train_df,
    target_col="failure",
    primary_metric="f1",
    timeout_minutes=30
)

# 최고 성능 모델의 노트북 경로
print(summary.best_trial.notebook_path)

from pyspark.ml.torch.distributor import TorchDistributor

def train_fn():
    import torch
    import torch.distributed as dist
    # torch.distributed가 자동으로 초기화됨
    # 일반 PyTorch DDP 학습 코드 작성
    ...

distributor = TorchDistributor(
    num_processes=4,       # 총 GPU 프로세스 수
    local_mode=False,      # True면 드라이버 노드만 사용
    use_gpu=True           # GPU 사용 여부
)
distributor.run(train_fn)

import mlflow

with mlflow.start_run(run_name="xgboost-v2"):
    mlflow.log_param("max_depth", 6)
    mlflow.log_param("learning_rate", 0.1)

    model = train_model(...)

    mlflow.log_metric("f1_score", 0.92)
    mlflow.log_metric("auc", 0.95)

    # 모델 + 피처 정보 함께 저장
    mlflow.sklearn.log_model(model, "model")

import mlflow

# Unity Catalog에 모델 등록
mlflow.set_registry_uri("databricks-uc")
model_uri = "runs:/abc123/model"
mlflow.register_model(model_uri, "catalog.schema.fraud_detection")

# Alias 지정
from mlflow import MlflowClient
client = MlflowClient()
client.set_registered_model_alias(
    name="catalog.schema.fraud_detection",
    alias="Champion",
    version=3
)

import requests

# 서빙 엔드포인트 호출 예시
endpoint_url = "https://<workspace-url>/serving-endpoints/<endpoint-name>/invocations"
headers = {"Authorization": f"Bearer {token}"}

payload = {
    "dataframe_records": [
        {"device_id": "sensor-001", "temperature": 85.2, "pressure": 1.2}
    ]
}

response = requests.post(endpoint_url, json=payload, headers=headers)
print(response.json())

import mlflow

# Spark UDF로 대규모 배치 추론
model_uri = "models:/catalog.schema.fraud_detection@Champion"
predict_udf = mlflow.pyfunc.spark_udf(spark, model_uri)

predictions = (
    spark.table("catalog.schema.transactions")
    .withColumn("prediction", predict_udf())
)
predictions.write.mode("overwrite").saveAsTable("catalog.schema.predictions")

from openai import OpenAI

client = OpenAI(
    api_key=databricks_token,
    base_url="https://<workspace-url>/serving-endpoints"
)

response = client.chat.completions.create(
    model="databricks-meta-llama-3-1-70b-instruct",
    messages=[{"role": "user", "content": "매출 하락 원인을 분석해주세요."}]
)

import mlflow
from databricks.agents import ChatAgent

class MyRAGAgent(ChatAgent):
    def predict(self, messages, context=None):
        # 1. 사용자 질문에서 검색 쿼리 추출
        query = messages[-1]["content"]

        # 2. Vector Search로 관련 문서 검색
        docs = vector_search.similarity_search(query, num_results=5)

        # 3. LLM에 컨텍스트와 함께 질문
        response = llm.chat(
            system="검색된 문서를 기반으로 답변하세요.",
            context=docs,
            question=query
        )
        return response

# Agent를 MLflow에 기록
with mlflow.start_run():
    mlflow.pyfunc.log_model("agent", python_model=MyRAGAgent())

import mlflow

results = mlflow.evaluate(
    model="runs:/abc123/agent",
    data=eval_dataset,            # 질문 + (선택) 정답 + (선택) 검색 문서
    model_type="databricks-agent"  # Agent 전용 평가 모드 활성화
)

# 전체 메트릭 확인
print(results.metrics)

# 개별 행 단위 상세 결과 (점수 + 이유)
display(results.tables["eval_results"])

# databricks.yml 예시
bundle:
  name: fraud-detection-pipeline

resources:
  jobs:
    training_job:
      name: "fraud-detection-training"
      tasks:
        - task_key: feature_engineering
          notebook_task:
            notebook_path: ./notebooks/01_feature_engineering.py
        - task_key: model_training
          depends_on:
            - task_key: feature_engineering
          notebook_task:
            notebook_path: ./notebooks/02_model_training.py

targets:
  dev:
    workspace:
      host: https://dev-workspace.cloud.databricks.com
  prod:
    workspace:
      host: https://prod-workspace.cloud.databricks.com
    run_as:
      service_principal_name: "ml-pipeline-sp"

# CLI로 배포
databricks bundle deploy --target prod
databricks bundle run training_job --target prod