MLflow 및 작업 실행 페이지를 사용하여 실행 추적

Important

이 기능은 공개 미리보기 단계에 있습니다.

air run로 제출하는 각 워크로드는 Databricks 작업 실행인 동시에 MLflow 실행이기도 합니다:

  • 작업 실행(작업 영역 작업 및 파이프라인 페이지에 표시됨)은 상태, 컴퓨팅, 재시도 및 드라이버 출력의 실행을 추적합니다.
  • MLflow 실행은 매개 변수, 메트릭, 시스템 메트릭 및 아티팩트 등 실험을 추적합니다.

하나의 제출은 하나의 작업 실행과 하나의 MLflow 실행을 만듭니다. 재시도는 새 MLflow 실행을 만듭니다.

실험 및 실행

두 워크로드 YAML 필드는 MLflow에서 실행이 표시되는 방식을 제어합니다.

experiment_name: my-training # Creates or appends to this MLflow experiment
mlflow_run_name: baseline-lr3e5 # Names the MLflow run for this submission
compute:
  num_accelerators: 8
  accelerator_type: GPU_8xH100
command: torchrun --nproc_per_node=8 train.py
max_retries: 2
  • experiment_name (필수): 이름이 없는 경우 이 이름으로 MLflow 실험을 만들거나 기존 실험에 새 실행을 추가합니다. 하나의 실험에는 여러 실행이 포함됩니다.
  • mlflow_run_name (선택 사항): 실행 이름을 설정합니다. 생략하면 실행 이름은 기본적으로 실험 이름(experiment_name)으로 설정됩니다.
  • max_retries (선택 사항): 각 재시도는 동일한 실험에서 실행되는 새 MLflow이므로 시도를 비교할 수 있습니다. 원본 제출과 그 재시도는 동일한 작업 실행을 공유합니다.

메트릭을 보여 주는 MLflow 실행 페이지

다음 세 곳에서 실행할 수 있습니다.

  • 작업: 작업 실행 페이지에는 실행 목록이 표시되며, 각 실행에는 해당 MLflow 실행과 실험으로 연결되는 링크가 있습니다.
  • MLflow: 실험 페이지에는 MLflow 실험이 나열됩니다.
  • 이전 워크로드: air get run <job-run-id> 실행의 작업, 실험 및 MLflow 실행에 대한 클릭 가능한 링크를 출력합니다. air list runs 에서는 이전 실행을 나열하고 필터링하여 특정 실행을 찾을 수 있습니다.
air get run <job-run-id> # Links to the job, experiment, and MLflow run
air list runs # List previous runs; filter to find a specific run

시스템 메트릭

GPU, CPU 및 메모리 시스템 메트릭은 실행될 때마다 자동으로 캡처됩니다. 구성이 필요하지 않습니다. MLflow 실행의 시스템 메트릭 탭에서 해당 메트릭을 확인할 수 있습니다.

MLflow 실행의 시스템 메트릭 탭(GPU/CPU/메모리)

사용자 지정 메트릭 기록

플랫폼은 MLflow 실행을 만들고 환경 변수를 통해 학습 프로세스에 해당 ID를 MLFLOW_RUN_ID 노출합니다. MLflow 추적 API를 사용하여 해당 실행에 고유한 매개 변수, 메트릭 및 아티팩트를 기록합니다.

분산(다중 노드) 워크로드에서 모든 노드는 동일한 MLflow 실행을 공유합니다. rank-0 프로세스에서만 로그하므로 각 메트릭은 한 번 기록됩니다.

import os

import mlflow

# Log from rank 0 only; all nodes share the same MLFLOW_RUN_ID.
if os.environ.get("RANK", "0") == "0":
    with mlflow.start_run(run_id=os.environ["MLFLOW_RUN_ID"]):
        mlflow.log_param("learning_rate", 3e-4)
        for step, loss in enumerate(training_losses):
            mlflow.log_metric("train_loss", loss, step=step)

로그 및 아티팩트

다음을 사용하여 실행 로그를 스트리밍하거나 다운로드합니다.air logs

air logs <job-run-id> # Stream logs from node 0
air logs <job-run-id> --node 2 # Logs from a specific node
air logs <job-run-id> --download-to ./logs/ # Download instead of streaming

로그는 MLflow 실행 시 아티팩트로도 사용할 수 있습니다. 모델 검사점을 유지하려면 Unity 카탈로그 볼륨에 씁니다. 검사점 패턴 및 볼륨 관리는 실험 추적 및 관찰 가능성을 참조하세요.

추가 리소스

  • Last updated on