MNIST 예제 전체
계약 3가지를 그대로 구현한 작동하는 최소 학습 컨테이너(Dockerfile, train.py, predictor.py, experiment.json)
계약을 그대로 구현한 작동하는 최소 예제입니다. 아래 목록의 앞 네 파일을 복사해 학습 코드만 바꾸면 됩니다. 마지막 파일은 로컬 검증용입니다. 아래 코드는 실제로 이미지로 빌드해 로컬 MLflow 와 플랫폼 양쪽에서 돌려 확인했습니다.
mnist/ ├── Dockerfile # 이미지 정의. ENTRYPOINT 는 그냥 train.py ├── train.py # 학습 진입점. 계약 3가지(지표 · 모델 · 종료 코드)를 구현 ├── predictor.py # 서빙 래퍼 + 모델 구조. 레지스트리에 등록돼 mlflow models serve 로 뜨는 파일 ├── experiment.json # 플랫폼이 넣어 주는 설정 파일의 예 (로컬 검증용) └── make_sample_data.py # 로컬 검증 전용. 플랫폼이 만들어 주는 데이터 디렉터리를 흉내 낸다
계약과 코드의 대응
| 계약 | 이 예제의 구현 |
|---|---|
| ① 지표를 MLflow 로 | mlflow.log_metrics({...}, step=epoch) + log_params({"epochs": …}) |
② log_model() 한 번 | train.py 끝의 mlflow.pyfunc.log_model(...) (signature · pip_requirements · metadata 포함) |
| ③ 종료 코드 | main() 이 0 을 돌려준다. 실패는 예외로 죽어 0 이 아닌 코드 |
하지 않는 것: set_tracking_uri() · set_experiment() · run 만들기 · 데이터 내려받기 · register_model(). 전부 플랫폼이 합니다.
predictor.py 가 모델 구조를 갖는 이유
서빙 컨테이너에는 학습 코드가 없습니다. mlflow models serve 는 등록된 모델만 로드하므로 모델 구조 · 전처리 · 추론이 predictor.py 한 파일 안에서 끝나야 합니다. 그래서 의존 방향이 train.py → predictor.py 입니다. 반대로 두면 학습은 잘 되는데 서빙에서만 ModuleNotFoundError 가 납니다.
전처리(preprocess)를 한 곳에 두는 것도 같은 이유입니다. 학습과 서빙의 전처리가 갈리면 "학습 정확도는 높은데 서빙 예측만 이상한", 가장 찾기 어려운 버그가 됩니다.
Dockerfile
# MNIST 학습 컨테이너: 계약 샘플
#
# 특별한 것이 없다는 점이 핵심이다: 플랫폼이 제공하는 스크립트를 넣지 않고,
# ENTRYPOINT 는 그냥 학습 스크립트를 가리킨다. 플랫폼이 이 ENTRYPOINT 를 그대로 실행한다.
FROM python:3.11-slim
# CPU 전용 wheel. 이미지가 ~1.5GB 로 끝난다.
# GPU 학습이면 이 줄 대신 `pip install torch mlflow==3.13.0 pillow pandas` 를 쓴다(6~8GB).
RUN pip install --no-cache-dir \
--index-url https://download.pytorch.org/whl/cpu torch==2.12.0 \
&& pip install --no-cache-dir mlflow==3.13.0 pillow pandas
# train.py 는 predictor.py 를 import 하므로 둘은 같은 디렉터리에 있어야 한다.
COPY train.py predictor.py /app/
# 없으면 학습 진행 로그가 버퍼링되어 화면에 실시간으로 안 뜬다.
ENV PYTHONUNBUFFERED=1
# 플랫폼이 마운트하는 위치. 컨테이너 안에서 상대 경로를 쓸 일이 있으면 여기가 기준이 된다.
WORKDIR /geo
ENTRYPOINT ["python", "/app/train.py"]train.py
# -*- coding: utf-8 -*-
"""MNIST 학습: 학습 컨테이너 계약 구현 샘플.
이 파일이 계약을 위해 하는 일은 **셋뿐**이다:
1. 지표를 MLflow 로 기록한다 (`step` = epoch, 총 epoch 은 param `epochs`)
2. 모델을 `mlflow.pyfunc.log_model()` 로 한 번 로깅한다
3. 성공하면 0, 실패하면 0 이외로 끝낸다
하지 **않는** 일: MLflow 주소·토큰 설정, run 생성, 실험 이름 지정, 데이터셋 다운로드,
모델 레지스트리 등록. 전부 플랫폼이 처리한다.
"""
import json
import os
import random
import signal
import sys
from pathlib import Path
import mlflow
import pandas as pd
import torch
import torch.nn.functional as F
from mlflow.models import infer_signature
from PIL import Image
from torch.utils.data import DataLoader, Dataset
# 서빙 래퍼에서 가져온다. 모델 구조와 전처리는 학습·서빙이 같아야 한다.
from predictor import CLASS_NAMES, SmallCNN, preprocess
# ── 플랫폼이 넣어 주는 것 ──────────────────────────────────────────────────
DATA_DIR = Path(os.environ.get("GEO_DATA_DIR", "/geo/dataset"))
WORK_DIR = Path(os.environ.get("GEO_WORK_DIR", "/geo/work"))
CONFIG_PATH = Path(os.environ.get("GEO_CONFIG", "/geo/experiment.json"))
DEVICE = "cpu" if os.environ.get("GEO_PARAM_DEVICE", "cpu") == "cpu" else "cuda:0"
# SIGTERM 은 "중단 요청"이다. 핸들러에서 죽지 말고 플래그만 세운 뒤
# 안전한 지점(에폭 경계)에서 체크포인트를 저장하고 나간다. 종료코드는 무엇이든 된다.
stop_requested = False
def _on_sigterm(*_):
global stop_requested
stop_requested = True
print("-> SIGTERM: 이번 에폭까지만 하고 저장 후 종료합니다", flush=True)
def hyperparameter(name: str, default):
"""``GEO_HP_<이름>`` 환경변수를 기본값의 타입으로 읽는다."""
raw = os.environ.get(f"GEO_HP_{name.upper()}")
if raw is None or raw == "":
return default
return type(default)(raw)
def load_config() -> dict:
"""``experiment.json`` 을 읽는다. 없어도 동작해야 한다(로컬 검증 편의)."""
try:
return json.loads(CONFIG_PATH.read_text(encoding="utf-8"))
except (OSError, ValueError):
return {}
class ManifestDataset(Dataset):
"""``manifest.json`` 이 가리키는 이미지 파일을 그대로 읽는 데이터셋.
데이터는 플랫폼이 학습 시작 **전에** 로컬 디렉터리로 내려 둔다. 다운로드도,
스토리지 자격증명도 필요 없다. 파일을 열기만 하면 된다.
라벨은 각 항목의 ``meta.label`` 에서 온다(분류 데이터셋의 경우).
"""
def __init__(self, items: list[dict]) -> None:
self.items = items
def __len__(self) -> int:
return len(self.items)
def __getitem__(self, index: int):
item = self.items[index]
image = Image.open(DATA_DIR / item["path"])
# 전처리는 predictor.preprocess 하나만 쓴다. 학습과 서빙이 갈리면
# 학습 정확도는 좋은데 서빙 예측만 틀리는, 가장 찾기 어려운 버그가 된다.
tensor = preprocess(image)[0]
return tensor, CLASS_NAMES.index(str(item["meta"]["label"]))
def load_manifest() -> list[dict]:
"""이미지 파일 목록(라벨 있는 것만)."""
manifest = json.loads((DATA_DIR / "manifest.json").read_text(encoding="utf-8"))
items = [
f
for f in manifest.get("files", [])
if f.get("kind") == "image" and (f.get("meta") or {}).get("label") is not None
]
if not items:
raise SystemExit("manifest.json 에 라벨이 있는 이미지가 없습니다")
return items
def split_items(items: list[dict], config: dict) -> tuple[list[dict], list[dict]]:
"""``experiment.json`` 의 지시대로 train/val 을 나눈다.
**분할은 이미지의 책임이다.** 플랫폼은 파일을 나눠 주지 않고 비율만 알려 준다.
"""
spec = config.get("split") or {}
percent = int(spec.get("train_percent", 80))
shuffled = list(items)
random.Random(int(spec.get("seed", 0))).shuffle(shuffled)
cut = max(1, len(shuffled) * percent // 100)
return shuffled[:cut], shuffled[cut:] or shuffled[cut - 1 :]
def run_epoch(model, loader, optimizer=None) -> tuple[float, float]:
"""(평균 loss, 정확도). ``optimizer`` 가 없으면 평가 모드."""
training = optimizer is not None
model.train(training)
total_loss, correct, seen = 0.0, 0, 0
with torch.set_grad_enabled(training):
for images, labels in loader:
images, labels = images.to(DEVICE), labels.to(DEVICE)
logits = model(images)
loss = F.cross_entropy(logits, labels)
if training:
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_loss += loss.item() * len(labels)
correct += int((logits.argmax(dim=-1) == labels).sum())
seen += len(labels)
return total_loss / seen, correct / seen
def build_signature(model, sample_item: dict):
"""(signature, input_example): 서빙 API 계약을 실제 입력 1건으로 만든다.
컬럼 이름·타입이 그대로 추론 요청 형식이 되므로, 서빙에서 받을 형태
(base64 이미지 문자열)를 여기서 그대로 만들어야 한다. ``input_example`` 은
로깅 시 MLflow 가 예측을 한 번 돌려 보게 해서, 서빙에서야 드러날 오류를
학습 단계에서 잡아 준다.
"""
import base64
encoded = base64.b64encode((DATA_DIR / sample_item["path"]).read_bytes()).decode()
model_input = pd.DataFrame({"image_b64": [encoded]})
with torch.no_grad():
probs = model.cpu()(preprocess(Image.open(DATA_DIR / sample_item["path"])))
probs = probs.softmax(dim=-1)[0]
index = int(probs.argmax())
model_output = pd.DataFrame(
[{"label": CLASS_NAMES[index], "confidence": float(probs[index])}]
)
return infer_signature(model_input, model_output), model_input
def main() -> int:
signal.signal(signal.SIGTERM, _on_sigterm)
WORK_DIR.mkdir(parents=True, exist_ok=True)
config = load_config()
epochs = hyperparameter("epochs", 5)
batch_size = hyperparameter("batch", 64)
learning_rate = hyperparameter("lr", 0.001)
train_items, val_items = split_items(load_manifest(), config)
print(f"-> train {len(train_items)} / val {len(val_items)} @ {DEVICE}", flush=True)
train_loader = DataLoader(
ManifestDataset(train_items), batch_size=batch_size, shuffle=True
)
val_loader = DataLoader(ManifestDataset(val_items), batch_size=batch_size)
model = SmallCNN(len(CLASS_NAMES)).to(DEVICE)
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
checkpoint = WORK_DIR / "best.pt"
best_accuracy = 0.0
# 접속 정보(URI·토큰)와 run 은 환경변수로 이미 세팅되어 있다. set_tracking_uri()
# 나 set_experiment() 를 부르면 안 된다. 실험이 어긋나 실행이 즉시 죽는다.
with mlflow.start_run():
# 총 epoch 수는 param 으로. 학습 화면의 진행률이 이 값으로 계산된다.
mlflow.log_params(
{"epochs": epochs, "batch": batch_size, "lr": learning_rate}
)
for epoch in range(epochs):
train_loss, train_accuracy = run_epoch(model, train_loader, optimizer)
_, val_accuracy = run_epoch(model, val_loader)
# 지표의 단일 소스는 MLflow 다. step 은 epoch 이고, 화면 x축이 이 값이다.
mlflow.log_metrics(
{
"train/loss": train_loss,
"train/accuracy": train_accuracy,
"eval/accuracy": val_accuracy,
},
step=epoch,
)
print(
f"-> epoch {epoch + 1}/{epochs} "
f"loss={train_loss:.4f} val_acc={val_accuracy:.4f}",
flush=True,
)
if val_accuracy >= best_accuracy:
best_accuracy = val_accuracy
torch.save(model.state_dict(), checkpoint)
if stop_requested: # 중단 요청. 저장은 이미 끝났다
break
if not checkpoint.exists(): # 한 에폭도 못 돌았다면 남길 모델이 없다
raise SystemExit("체크포인트가 없습니다 — 학습이 진행되지 않았습니다")
model.load_state_dict(torch.load(checkpoint, map_location=DEVICE))
signature, input_example = build_signature(model, val_items[0])
# 모델은 반드시 MLflow 모델로 남긴다. 생 .pt 만 남기면 MLmodel 이 없어
# `mlflow models serve` 로 뜨지 않고, 서빙 이미지도 만들 수 없다.
mlflow.pyfunc.log_model(
name="model",
python_model=str(Path(__file__).resolve().parent / "predictor.py"),
artifacts={"weights": str(checkpoint)},
signature=signature,
input_example=input_example,
pip_requirements=[
f"torch=={torch.__version__.split('+')[0]}",
"pillow",
"pandas",
"mlflow==3.13.0", # 플랫폼 고정 버전
],
metadata={
"input_kind": "image_b64", # 추론 콘솔의 입력 위젯을 정한다
"class_names": {str(i): name for i, name in enumerate(CLASS_NAMES)},
},
)
# 레지스트리 등록(register_model)은 부르지 않는다. 플랫폼이 한다.
mlflow.log_metric("eval/best_accuracy", best_accuracy)
print(f"-> done. best val accuracy = {best_accuracy:.4f}", flush=True)
return 0
if __name__ == "__main__":
sys.exit(main())predictor.py
# -*- coding: utf-8 -*-
"""서빙 래퍼: 모델 레지스트리에 등록되어 `mlflow models serve` 로 뜨는 파일.
**서빙 컨테이너에는 학습 코드(`train.py`)가 없다.** 그래서 이 파일 하나가
모델 구조 정의 + 가중치 로드 + 추론까지 스스로 끝내야 한다(models-from-code).
`train.py` 가 이 파일에서 ``SmallCNN`` 을 import 하는 방향이지 그 반대가 아닌 것도
그래서다. 반대로 두면 서빙 쪽에서 학습 코드가 필요해진다.
입력/출력 계약(= `signature`):
- 입력: ``image_b64`` 컬럼 하나. PNG/JPEG 바이트를 base64 로 인코딩한 문자열.
640x640 텐서를 JSON 숫자로 풀면 요청이 수십 MB 가 되므로 이미지 모델은 base64 를 쓴다.
- 출력: ``label``(문자열), ``confidence``(0~1 실수) 두 컬럼.
"""
import base64
import io
import mlflow
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image
#: 인덱스 → 표시 이름. MNIST 는 인덱스와 숫자가 같지만, 일반적으로는
#: 학습에 쓴 클래스 순서를 그대로 적어야 한다.
CLASS_NAMES = [str(i) for i in range(10)]
#: 학습·서빙이 같은 전처리를 써야 한다. 여기 상수를 바꾸면 train.py 도 같이 바뀐다.
IMAGE_SIZE = 28
NORM_MEAN = 0.1307
NORM_STD = 0.3081
class SmallCNN(nn.Module):
"""28x28 흑백 입력용 소형 CNN (샘플이므로 구조 자체는 중요하지 않다)."""
def __init__(self, num_classes: int = 10) -> None:
super().__init__()
self.conv1 = nn.Conv2d(1, 16, 3, padding=1)
self.conv2 = nn.Conv2d(16, 32, 3, padding=1)
self.fc1 = nn.Linear(32 * 7 * 7, 128)
self.fc2 = nn.Linear(128, num_classes)
def forward(self, x: torch.Tensor) -> torch.Tensor:
x = F.max_pool2d(F.relu(self.conv1(x)), 2)
x = F.max_pool2d(F.relu(self.conv2(x)), 2)
x = x.flatten(1)
x = F.relu(self.fc1(x))
return self.fc2(x)
def preprocess(image: Image.Image) -> torch.Tensor:
"""PIL 이미지 → (1, 1, 28, 28) 텐서. 학습과 **동일한** 전처리."""
image = image.convert("L").resize((IMAGE_SIZE, IMAGE_SIZE))
buffer = bytearray(image.tobytes()) # 쓰기 가능한 사본(frombuffer 요구사항)
tensor = torch.frombuffer(buffer, dtype=torch.uint8).float() / 255.0
tensor = tensor.reshape(1, 1, IMAGE_SIZE, IMAGE_SIZE)
return (tensor - NORM_MEAN) / NORM_STD
class MnistPredictor(mlflow.pyfunc.PythonModel):
"""스코어링 서버가 로드하는 모델."""
def load_context(self, context):
"""컨테이너 기동 시 1회. 가중치는 ``artifacts`` 로 함께 실려 온다."""
self.model = SmallCNN(len(CLASS_NAMES))
state = torch.load(context.artifacts["weights"], map_location="cpu")
self.model.load_state_dict(state)
self.model.eval()
def predict(self, context, model_input, params=None):
"""요청 1건당 DataFrame 한 장이 들어오고 한 장을 돌려준다.
반환값은 **JSON 직렬화 가능**해야 한다(DataFrame / ndarray / list / dict).
로깅할 때 "Add type hints to the `predict` method" 경고가 뜨는데 무시해도 된다.
MLflow 의 타입 힌트 기반 검증은 ``list[...]`` 형태만 지원하고, 이 모델은
DataFrame 계약(`signature`)을 쓰므로 힌트를 붙이면 오히려 경고가 늘어난다.
"""
rows = []
for encoded in model_input["image_b64"]:
image = Image.open(io.BytesIO(base64.b64decode(encoded)))
with torch.no_grad():
probs = self.model(preprocess(image)).softmax(dim=-1)[0]
index = int(probs.argmax())
rows.append(
{"label": CLASS_NAMES[index], "confidence": float(probs[index])}
)
return pd.DataFrame(rows)
# models-from-code 의 마지막 줄. 이 호출이 없으면 로드되지 않는다.
mlflow.models.set_model(MnistPredictor())CLASS_NAMES 의 순서가 곧 모델 출력 인덱스의 뜻입니다. 학습에 쓴 순서와 반드시 같아야 합니다.
experiment.json
플랫폼에서는 자동으로 만들어집니다. 로컬 검증 때 /geo/experiment.json 에 마운트하면 split 등을 읽어 갑니다(없어도 기본값으로 돕니다).
{
"experiment_id": "exp-local-mnist",
"name": "mnist-local-20260801-1030",
"tenant": "DEMO",
"created_at": "2026-08-01T10:30:12Z",
"task": "classification",
"framework": "mnist-example",
"model": "small-cnn",
"classes": ["0", "1", "2", "3", "4", "5", "6", "7", "8", "9"],
"split": { "method": "random", "train_percent": 80, "seed": 0 },
"hyperparameters": { "epochs": 5, "batch": 64, "lr": 0.001 },
"device": "cpu",
"gpu": 0,
"evaluation": { "benchmark": true, "speed_test": false },
"export": { "onnx": false, "tensorrt": false },
"serving": { "runtime": "cpu" },
"register": { "enabled": true, "model_name": "mnist" },
"pretrained": { "source": "catalog" },
"dataset": {
"id": "ds-local-mnist",
"name": "MNIST (로컬 검증용)",
"file_count": 2000,
"path": "/geo/dataset",
"manifest": "/geo/dataset/manifest.json"
},
"paths": { "data_dir": "/geo/dataset", "work_dir": "/geo/work" },
"mlflow": { "experiment_name": "mnist-local-20260801-1030", "run_id": "0000000000000000000000000000abcd" }
}make_sample_data.py (로컬 검증 전용)
# -*- coding: utf-8 -*-
"""로컬 검증용 데이터 생성기. **플랫폼에서는 필요 없습니다**.
실제 실행에서는 플랫폼이 `GEO_DATA_DIR` 에 데이터셋을 미리 내려놓고
`manifest.json` 도 함께 씁니다. 이 스크립트는 그 상태를 로컬에서 흉내 내
전달 전 자가 검증을 할 수 있게 합니다.
python make_sample_data.py ./sample-data # torchvision MNIST (권장)
python make_sample_data.py ./sample-data --synthetic # 네트워크 없이 합성 이미지
"""
import argparse
import json
import random
from pathlib import Path
from PIL import Image, ImageDraw
def write_manifest(root: Path, records: list[tuple[str, str]]) -> None:
"""플랫폼과 같은 형식의 ``manifest.json`` 을 쓴다."""
manifest = {
"dataset": {"id": "ds-local-mnist", "name": "MNIST (로컬 검증용)"},
"files": [
{"path": path, "kind": "image", "meta": {"label": label}}
for path, label in records
],
}
(root / "manifest.json").write_text(
json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8"
)
print(f"-> {root/'manifest.json'} ({len(records)} files)")
def from_torchvision(root: Path, count: int) -> list[tuple[str, str]]:
from torchvision import datasets
dataset = datasets.MNIST(root=str(root / ".cache"), train=True, download=True)
records = []
for index in range(min(count, len(dataset))):
image, label = dataset[index]
name = f"images/{index:05d}.png"
(root / "images").mkdir(parents=True, exist_ok=True)
image.save(root / name)
records.append((name, str(label)))
return records
def synthetic(root: Path, count: int) -> list[tuple[str, str]]:
"""숫자를 그려 넣은 28x28 이미지. 학습이 도는지 확인하는 용도."""
(root / "images").mkdir(parents=True, exist_ok=True)
rng = random.Random(0)
records = []
for index in range(count):
label = rng.randrange(10)
image = Image.new("L", (28, 28), color=0)
draw = ImageDraw.Draw(image)
draw.text((9, 8), str(label), fill=255)
name = f"images/{index:05d}.png"
image.save(root / name)
records.append((name, str(label)))
return records
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("output", type=Path, help="GEO_DATA_DIR 로 쓸 디렉터리")
parser.add_argument("--count", type=int, default=2000)
parser.add_argument(
"--synthetic", action="store_true", help="torchvision 없이 합성 이미지 사용"
)
args = parser.parse_args()
args.output.mkdir(parents=True, exist_ok=True)
if args.synthetic:
records = synthetic(args.output, args.count)
else:
records = from_torchvision(args.output, args.count)
write_manifest(args.output, records)
if __name__ == "__main__":
main()다음: 로컬에서 먼저 검증하기