이 글에서 다루는 것

MLflow Registry를 단일 소스로 삼아 Airflow DAG에서 Triton에 모델을 자동 배포하고, 검증 체인(load/ready/infer) 통과 후에만 운영 확정하며, 실패 시 자동 롤백하는 파이프라인을 구축한 과정

선수지식

이 단계에서 해결하려는 문제

실무 환경에서 모델 배포는 “새 모델을 올리는 작업"이 아니라 “현재 운영 상태를 안전하게 갱신하는 상태 전이(State Transition)“에 가깝다. 이번 단계에서는 MLflow Registry를 단일 소스로 삼아 Triton Inference Server에 모델을 자동 배포하고, 로딩/헬스 체크/실제 추론 검증을 모두 통과한 경우에만 운영 모델을 확정(commit)하며, 중간 단계에서 하나라도 실패하면 이전 운영 상태로 자동 복구되는 최소 롤백 구조를 구현했다.

🎯 핵심 목표 요약

  • MLflow Registry에 등록된 모델을 alias 기준으로 선택
  • Airflow DAG로 배포 → 검증 → 확정(commit) 흐름 강제
  • 검증 실패 시 이전 운영 상태로 자동 복구
  • Triton inference latency를 Prometheus/Grafana로 관측

1️⃣ 전체 아키텍처 요약 (파일/책임 단위)

1) Airflow DAG 레포 구조

airflow-dags-dev/
└─ dags/
   ├─ dag_mlflow_to_triton_deploy_dev.py   # 배포 흐름 오케스트레이션/롤백 트리거
   └─ ml_code/
      ├─ train_model.py                    # 학습 + MLflow logging + ONNX artifact 생성
      └─ triton_deploy.py                  # 배포/검증/commit/rollback 로직 집약

2) Triton Model Repository 구조 (NFS, RWX)

/model-repo/dev/
└─ best_model/
   ├─ config.pbtxt
   ├─ current.json
   ├─ 12/
   │  └─ model.onnx
   ├─ 13/
   │  └─ model.onnx
   └─ 13.failed_20251231T...

설계 의도

  • current.json: 현재 운영 중인 모델 버전을 정의하는 단일 진실(Single Source of Truth). FastAPI, 운영 판단, 롤백 기준이 모두 이 파일을 기준으로 동작
  • 실패한 버전은 삭제하지 않고 격리 → 원인 분석 및 재현 가능성 유지

2️⃣ 전체 실행 흐름 (상태 전이 기준)

Triton을 선택한 이유:

  • 모델 로딩을 명시적으로 제어 가능
  • 로드 실패 시에도 기존 모델 유지
  • 배포 실패가 곧바로 서비스 장애로 이어지지 않음

실행 흐름

[Trigger] Airflow UI에서 DAG 수동 실행 (params.alias = "A")

DAG: dag_mlflow_to_triton_deploy_dev.py
  1) snapshot_current()     - 기존 current.json 스냅샷
  2) materialize()          - MLflow → ONNX → model-repo 배치
  3) triton_load()          - Triton explicit load
  4) triton_ready()         - ready endpoint 검증
  5) triton_infer_smoke()   - 샘플 추론 smoke test
  6) commit_current()       - 성공 시 운영 기준 확정

  X) rollback_minimal()     - 중간 실패 시 자동 복구

3️⃣ DAG 오케스트레이션 핵심

with DAG(
    dag_id="mlflow_to_triton_min_dev",
    schedule=None,
    params={"alias": "A"},
) as dag:

    t0 = PythonOperator(task_id="snapshot_current",
        python_callable=snapshot_current)
    t1 = PythonOperator(task_id="materialize_repo",
        python_callable=materialize,
        op_kwargs={"alias": "{{ params.alias }}"})
    t2 = PythonOperator(task_id="triton_load",
        python_callable=triton_load)
    t_ready = PythonOperator(task_id="triton_ready",
        python_callable=triton_ready)
    t_smoke = PythonOperator(task_id="triton_infer_smoke",
        python_callable=triton_infer_smoke)
    t3 = PythonOperator(task_id="commit_current",
        python_callable=commit_current)
    t_rb = PythonOperator(task_id="rollback_minimal",
        python_callable=rollback_minimal,
        trigger_rule=TriggerRule.ONE_FAILED)

    t0 >> t1 >> t2 >> t_ready >> t_smoke >> t3
    [t1, t2, t_ready, t_smoke] >> t_rb

설계 포인트

  • rollback task는 검증 체인 어디서든 실패 시 실행
  • commit은 모든 검증 성공 이후에만 허용
  • rollback의 목적은 원인 분석이 아니라 운영 상태 보호

4️⃣ MLflow → Triton Materialize 단계

def materialize(ti, alias: str = "A", **_):
    v, run_id = select_by_alias(model, alias)
    ver_dir = os.path.join(model_dir, str(v))
    os.makedirs(ver_dir, exist_ok=True)

    local = mlflow.artifacts.download_artifacts(
        artifact_uri=f"runs:/{run_id}/onnx/model.onnx")
    shutil.copyfile(local, os.path.join(ver_dir, "model.onnx"))

    with open(os.path.join(model_dir, "config.pbtxt"), "w") as f:
        f.write(CONFIG_TEMPLATE.format(model=model))

config.pbtxt를 버전별로 두면 버전 전환 시 설정 불일치 위험이 증가한다. model root 단일 config로 운영 안정성과 예측 가능성을 확보했다.

5️⃣ 배포 검증 단계

1) Ready Check

GET /v2/models/{model}/ready

Triton 내부 로딩 상태 검증

2) Smoke Inference

POST /v2/models/{model}/infer

실제 입력으로 추론 가능 여부까지 검증. “로드 성공 but 추론 실패” 케이스 차단.

3) Commit (운영 확정)

{
  "active_version": 13,
  "run_id": "...",
  "alias": "A",
  "updated_at_utc": "20251231T025902"
}

이 파일이 현재 운영 모델의 단일 기준이다. 사람이 아니라 시스템이 해석 가능한 운영 상태.

6️⃣ 최소 롤백(minimal rollback) 전략

def rollback_minimal(ti, **_):
    # 1) current.json 복구
    # 2) 실패한 버전 디렉터리 격리
    # 3) Triton unload → load (best-effort)

왜 최소 롤백인가: Triton은 load 실패 시 기존 모델을 자동 유지한다. 핵심은 운영 기준 복구와 실패 결과를 다음 시도에 영향 주지 않게 정리하는 것이다. 불필요한 복잡도 없이 안정성을 확보하는 구조다.

예를 들어 ONNX 출력 shape 변경으로 smoke inference가 실패할 경우, 해당 버전은 .failed_*로 격리되고 운영 기준은 이전 상태로 자동 복구된다.

7️⃣ Observability: Inference Latency

Triton Metrics 예시

nv_inference_request_duration_us
nv_inference_compute_infer_duration_us
nv_inference_queue_duration_us
  • Prometheus가 /metrics 수집
  • Grafana에서 p95 latency, request count, 실패율 추이 관측

이 단계에서는 정확도보다 운영 이상을 가장 먼저 감지할 수 있는 지표로 latency를 핵심 관측 대상으로 삼았다.

🧠 이번 단계에서 확보된 개념

항목내용
TriggerAirflow DAG 수동 실행 (alias 기반)
Source of TruthMLflow Registry → 배포 확정은 current.json
Sync 방식MLflow artifact → 파일 배치 → Triton explicit load
Rollback 조건load / ready / infer 중 하나라도 실패
ObservabilityTriton inference latency (p95) + 실패율

설계 판단 (Why This Way?)

current.json 파일 하나로 운영 모델을 정의하여 MLflow 장애와 독립적인 롤백을 가능하게 하고, 실패 버전은 삭제 대신 .failed_* 디렉터리로 격리하여 원인 분석용 아티팩트를 보존합니다. 검증 체인을 load/ready/infer 3단계로 나누어 각 단계의 실패 원인을 빠르게 격리합니다.

다음에 읽을 글

→ Triton 서빙 플랫폼 - Alerting 운영 표준 매뉴얼 — Triton 전용 알림 체계