# 6. MLflow 실험 추적

## MLflow란?

**MLflow**는 머신러닝 실험의 전체 라이프사이클을 관리하는 오픈소스 플랫폼입니다. 학습 실험에서 생성되는 하이퍼파라미터, 메트릭, 모델 아티팩트를 중앙에서 추적하고 비교할 수 있습니다.

### 실험 추적이 중요한 이유

VLA 모델의 fine-tuning에서는 다양한 하이퍼파라미터를 시도해야 합니다:

* 학습률 (learning rate)
* 배치 크기 (batch size)
* 데이터셋 구성
* 로봇 플랫폼 (embodiment)

각 실험마다:

* **손실값(loss)이 어떻게 감소하는가?**
* **최종 성능이 어떠한가?**
* **어떤 하이퍼파라미터가 최적인가?**

이러한 질문에 답하기 위해 **MLflow**를 사용하여 모든 실험을 중앙화된 대시보드에서 추적합니다.

### SageMaker Managed MLflow의 장점

AWS에서 제공하는 **SageMaker Managed MLflow**는:

* AWS 콘솔 인증 통합 (별도 계정 관리 불필요)
* 완전 관리형 (배포, 확장, 백업 자동 처리)
* HyperPod, SageMaker 학습, Batch 등 모든 AWS ML 서비스와 통합
* S3와 자동 동기화

***

## 5.1 MLflow 클라이언트 설치 및 연결

### Head Node에서 MLflow 설치

MLflow 설치는 `setup_groot_env.sh`에서 이미 완료되었습니다 (`mlflow`, `sagemaker-mlflow`, `boto3` 패키지 포함). MLflow 추적 URI를 설정하고 연결을 테스트합니다:

```bash
# 가상환경 활성화
source /fsx/envs/gr00t/bin/activate

# MLflow Tracking URI 설정 (ARN 형식 사용)
export MLFLOW_TRACKING_URI="arn:aws:sagemaker:us-east-1:<ACCOUNT_ID>:mlflow-tracking-server/hyperpod-<userId>-mlflow"

# 영구 저장
echo "export MLFLOW_TRACKING_URI=\"${MLFLOW_TRACKING_URI}\"" >> ~/.bashrc
source ~/.bashrc

# 연결 테스트
python3 -c "
import mlflow, os
mlflow.set_tracking_uri(os.environ['MLFLOW_TRACKING_URI'])
experiments = mlflow.search_experiments()
print(f'Connected! Found {len(experiments)} experiments.')
for e in experiments:
    print(f'  - {e.name}')
"
```

정상 출력:

```
Connected! Found 2 experiments.
  - Default
  - groot-finetune
```

{% hint style="info" %}
**MLflow URI 형식:**

* **Python SDK (학습 코드)**: ARN 형식 사용 — `arn:aws:sagemaker:us-east-1:<ACCOUNT_ID>:mlflow-tracking-server/hyperpod-<userId>-mlflow`
* **브라우저 UI**: URL 형식 — `https://us-east-1.experiments.sagemaker.aws/mlflow/hyperpod-<userId>-mlflow`

CDK 출력에서 `MlflowTrackingArn` (SDK용)과 `MlflowTrackingUri` (브라우저용)를 확인하세요.
{% endhint %}

***

## 5.2 MLflow UI 접속

### 브라우저에서 접속

CDK 출력의 `MLflowTrackingUri`를 브라우저에 복사합니다:

```
https://us-east-1.experiments.sagemaker.aws/mlflow/hyperpod-<userId>-mlflow
```

### MLflow UI 화면 구성

MLflow UI에 접속하면 다음과 같은 화면을 볼 수 있습니다:

**1. Experiments 목록 (좌측)**

* 모든 실험이 나열됨
* 각 실험 이름 클릭으로 상세 조회

**2. Runs 목록 (중앙)**

* 선택된 실험의 모든 Run (개별 학습 작업)
* 각 Run의 상태, 시작시간, 소요시간 표시

**3. Run 상세정보 (하단)**

* 하이퍼파라미터
* 메트릭 (손실값, 정확도 등)
* 아티팩트 (저장된 체크포인트)
* 태그

***

## 5.3 학습 코드에 MLflow 통합

학습 래퍼(`train_groot.py`)에는 MLflow 통합이 기본 포함되어 있습니다. `MLFLOW_TRACKING_URI` 환경변수만 설정하면 자동으로 하이퍼파라미터(`dataset_path`, `embodiment_tag`, `base_model`, `max_steps`, `global_batch_size`, `num_gpus`)와 학습 성공 여부가 기록됩니다.

상세 MLflow 사용 예시:

```bash
# MLflow 로깅 데모 실행
python3 /fsx/scratch/aws-physical-ai-recipes/training/hyperpod/mlflow/example_usage.py
```

***

## 5.4 실험 비교 (MLflow UI)

### 1. 실험 선택

MLflow UI의 좌측 패널에서 `groot-finetune` 실험을 클릭합니다.

현재 실험의 모든 Run이 중앙 패널에 표시됩니다:

| Run Name      | Started | Duration | Learning Rate | Batch Size | Final Loss |
| ------------- | ------- | -------- | ------------- | ---------- | ---------- |
| aloha-exp-001 | 10:23   | 45m      | 2e-5          | 32         | 0.124      |
| aloha-exp-002 | 10:25   | 42m      | 1e-5          | 32         | 0.131      |
| aloha-exp-003 | 11:15   | 48m      | 2e-5          | 64         | 0.118      |

### 2. 비교할 Run 선택

비교하고 싶은 Run들의 체크박스를 선택합니다.

### 3. 메트릭 비교

**Compare** 버튼을 클릭하면 메트릭 비교 화면으로 이동합니다:

```
손실값 (Loss) 비교 그래프
─────────────────────────────────────
Run 1: aloha-exp-001 (LR=2e-5, BS=32)   ─ 우상향
Run 2: aloha-exp-003 (LR=2e-5, BS=64)   ─ 우상향 (더 가파름)

결론: aloha-exp-001이 더 안정적인 학습 곡선을 보임
```

### 4. 최적 하이퍼파라미터 찾기

CLI를 사용하여 최고 성능 Run을 자동으로 찾을 수 있습니다:

```bash
python3 << 'EOF'
import mlflow
import os

# MLflow 설정
mlflow.set_tracking_uri(os.environ["MLFLOW_TRACKING_URI"])

# 실험 조회
runs = mlflow.search_runs(
    experiment_names=["groot-finetune"],
    order_by=["metrics.final_loss ASC"],  # 손실값 낮은 순서
    max_results=5,
)

print("=" * 80)
print("최고 성능 Top 5 Runs")
print("=" * 80)

for i, row in runs.iterrows():
    run_id = row["run_id"]
    run_name = row["tags.mlflow.runName"]
    loss = row["metrics.final_loss"]
    lr = row["params.learning_rate"]
    bs = row["params.batch_size"]
    
    print(f"{i+1}. {run_name}")
    print(f"   Final Loss: {loss:.4f}")
    print(f"   Learning Rate: {lr}, Batch Size: {bs}")
    print(f"   Run ID: {run_id}")
    print()
EOF
```

***

## 5.5 CLI로 실험 결과 조회

### 모든 Run 메트릭 조회

```bash
python3 << 'EOF'
import mlflow
import pandas as pd
import os

mlflow.set_tracking_uri(os.environ["MLFLOW_TRACKING_URI"])

# 실험의 모든 Run 조회
runs = mlflow.search_runs(
    experiment_names=["groot-finetune"],
    max_results=100,
)

# 관심 있는 컬럼만 선택
cols = [
    "tags.mlflow.runName",
    "params.learning_rate",
    "params.batch_size",
    "metrics.final_loss",
    "start_time",
    "status",
]

print("\n" + "=" * 100)
print("GR00T Fine-tuning 실험 결과")
print("=" * 100 + "\n")

# 선택한 컬럼만 출력
print(runs[[c for c in cols if c in runs.columns]].to_string(index=False))
EOF
```

출력 예:

```
100 rows × 6 columns
       tags.mlflow.runName params.learning_rate params.batch_size metrics.final_loss start_time status
0       aloha-exp-001             2e-05                 32              0.1240  ...  FINISHED
1       aloha-exp-002             1e-05                 32              0.1310  ...  FINISHED
2       aloha-exp-003             2e-05                 64              0.1180  ...  FINISHED
```

### 특정 Run의 상세 정보 조회

```bash
python3 << 'EOF'
import mlflow
import os

mlflow.set_tracking_uri(os.environ["MLFLOW_TRACKING_URI"])

# 특정 Run ID로 조회
run_id = "abc123def456"  # MLflow UI에서 복사
run = mlflow.get_run(run_id)

print("\n하이퍼파라미터:")
for key, value in run.data.params.items():
    print(f"  {key}: {value}")

print("\n최종 메트릭:")
for key, value in run.data.metrics.items():
    print(f"  {key}: {value:.4f}")

print("\n저장된 아티팩트:")
for artifact in mlflow.list_artifacts(run_id):
    print(f"  {artifact.path}")
EOF
```

***

## 5.6 MLflow에 기록되는 GR00T 메트릭

학습 코드가 다음과 같은 메트릭을 자동으로 생성합니다:

### 학습 메트릭 (Training Metrics)

| 메트릭                      | 설명             | 기록 빈도        |
| ------------------------ | -------------- | ------------ |
| `train/loss`             | 훈련 손실값         | 매 step       |
| `train/learning_rate`    | 현재 학습률         | 매 step       |
| `train/loss_action_head` | Action Head 손실 | 매 step       |
| `train/loss_vision`      | Vision 모듈 손실   | 매 step (선택적) |
| `train/step`             | 현재 Step        | 매 step       |

### 검증 메트릭 (Validation Metrics)

| 메트릭            | 설명     | 기록 빈도   |
| -------------- | ------ | ------- |
| `val/loss`     | 검증 손실값 | 매 epoch |
| `val/accuracy` | 검증 정확도 | 매 epoch |

### 최종 메트릭 (Final Metrics)

| 메트릭              | 설명          |
| ---------------- | ----------- |
| `final_loss`     | 최종 손실값      |
| `final_accuracy` | 최종 정확도      |
| `training_steps` | 총 학습 step 수 |

***

## 5.7 MLflow UI에서 메트릭 시각화

### 손실값 곡선 비교

MLflow UI에서:

1. 비교할 Run 선택
2. **Metrics** 탭 클릭
3. `train/loss` 선택
4. **Parallel Coordinates** 또는 **Scatter Plot** 선택

{% hint style="info" %}
**MLflow UI 화면 구성:**

좌측: Experiment 트리 중앙: Run 목록 우측: Run 상세정보 및 메트릭 그래프

하이퍼파라미터별 성능을 한눈에 비교할 수 있습니다.
{% endhint %}

***

## 5.8 MLflow 아티팩트 저장 및 관리

### 체크포인트 저장

```python
import mlflow

# 학습 완료 후
mlflow.log_artifact(
    local_path="/fsx/checkpoints/vla/groot-n1.7-aloha/checkpoint-final",
    artifact_path="model_checkpoints"
)

# 여러 파일 저장
mlflow.log_artifacts(
    local_dir="/fsx/checkpoints/vla/groot-n1.7-aloha",
    artifact_path="checkpoints"
)
```

### 저장된 아티팩트 다운로드

```python
import mlflow

# Run ID 지정
run_id = "abc123def456"

# 아티팩트 목록 확인
artifacts = mlflow.list_artifacts(run_id)
for artifact in artifacts:
    print(artifact.path)

# 특정 아티팩트 다운로드
mlflow.download_artifacts(
    run_id=run_id,
    artifact_path="model_checkpoints",
    dst_path="/local/path"
)
```

***

## 5.9 트러블슈팅

| 증상                   | 원인             | 해결 방법                          |
| -------------------- | -------------- | ------------------------------ |
| "401 Unauthorized"   | AWS 콘솔 미로그인    | AWS 콘솔 재로그인 후 MLflow UI 접속     |
| "Connection refused" | MLflow 서버 다운   | 몇 분 대기 후 재시도                   |
| 메트릭이 기록되지 않음         | MLflow URI 잘못됨 | `echo $MLFLOW_TRACKING_URI` 확인 |
| Run이 FAILED 상태       | 학습 오류          | CloudWatch 로그 확인               |
| 아티팩트 저장 실패           | S3 권한 부족       | IAM 역할 S3 권한 확인                |

***

## References

* [MLflow Documentation](https://mlflow.org/docs/)
* [AWS SageMaker Managed MLflow](https://docs.aws.amazon.com/sagemaker/latest/dg/mlflow.html)
* [MLflow Tracking API](https://mlflow.org/docs/latest/python_api/mlflow.html#mlflow.start_run)
* [GR00T Fine-tuning Metrics](https://github.com/NVIDIA/Isaac-GR00T)


---

# Agent Instructions: Querying This Documentation

If you need additional information that is not directly available in this page, you can query the documentation dynamically by asking a question.

Perform an HTTP GET request on the current page URL with the `ask` query parameter:

```
GET https://hi-space.gitbook.io/physical-ai-on-aws/physical-ai-on-aws-guide/hyperpod-distributed-training/6.-mlflow-tracking.md?ask=<question>
```

The question should be specific, self-contained, and written in natural language.
The response will contain a direct answer to the question and relevant excerpts and sources from the documentation.

Use this mechanism when the answer is not explicitly present in the current page, you need clarification or additional context, or you want to retrieve related documentation sections.