用 Flask 代理打通外部访问 Amazon SageMaker MLflow 的 REST 通道

外部调用方 → HTTPS → API Gateway / ALB → Flask Proxy → SageMaker MLflow Server

# app.py — Flask MLflow Proxy
import os
import requests
from flask import Flask, request, jsonify, Response

app = Flask(__name__)

# SageMaker MLflow 内网地址，从环境变量读取
MLFLOW_INTERNAL_URL = os.environ.get(
    "MLFLOW_INTERNAL_URL", "https://mlflow-server.internal:5000"
)

# 简易 token 校验（生产环境建议换成 IAM 签名或 OAuth）
API_TOKEN = os.environ.get("PROXY_API_TOKEN", "change-me-in-production")


def verify_token():
    """校验请求头中的 Authorization token"""
    auth = request.headers.get("Authorization", "")
    if auth != f"Bearer {API_TOKEN}":
        return jsonify({"error": "unauthorized"}), 401
    return None


@app.route("/api/2.0/mlflow/<path:subpath>", methods=["GET", "POST", "PUT", "DELETE"])
def proxy_mlflow(subpath):
    # 1. 鉴权
    err = verify_token()
    if err is not None:
        return err

    # 2. 构造转发目标 URL
    target_url = f"{MLFLOW_INTERNAL_URL}/api/2.0/mlflow/{subpath}"

    # 3. 复制请求头，去掉 host 和 authorization（避免泄露到内网）
    fwd_headers = {
        k: v for k, v in request.headers
        if k.lower() not in ("host", "authorization")
    }

    # 4. 转发请求
    resp = requests.request(
        method=request.method,
        url=target_url,
        headers=fwd_headers,
        data=request.get_data(),
        params=request.args,
        timeout=30,
        verify=False,  # 内网自签证书场景；生产环境应配 CA
    )

    # 5. 透传响应
    excluded_headers = {"content-encoding", "transfer-encoding"}
    response_headers = {
        k: v for k, v in resp.headers.items() if k.lower() not in excluded_headers
    }
    return Response(resp.content, status=resp.status_code, headers=response_headers)


@app.route("/health", methods=["GET"])
def health():
    return jsonify({"status": "ok"})


if __name__ == "__main__":
    app.run(host="0.0.0.0", port=8080)

# 安装依赖
pip install flask requests

# 设置环境变量
export MLFLOW_INTERNAL_URL="https://your-sagemaker-mlflow.internal:5000"
export PROXY_API_TOKEN="my-secret-token"

# 启动代理
python app.py

curl -X POST \
  http://localhost:8080/api/2.0/mlflow/experiments/create \
  -H "Authorization: Bearer my-secret-token" \
  -H "Content-Type: application/json" \
  -d '{"name": "proxy-test-experiment"}'

curl -X GET \
  "http://localhost:8080/api/2.0/mlflow/experiments/search?max_results=10" \
  -H "Authorization: Bearer my-secret-token"

外部 → Amazon API Gateway (IAM/Auth) → ALB → ECS/Fargate (Flask Proxy) → SageMaker MLflow

# ecs_service.tf — 核心部分
resource "aws_ecs_task_definition" "mlflow_proxy" {
  family                   = "mlflow-proxy"
  network_mode             = "awsvpc"
  requires_compatibilities = ["FARGATE"]
  cpu                      = "256"
  memory                   = "512"

  container_definitions = jsonencode([
    {
      name      = "proxy"
      image     = "your-registry/mlflow-proxy:latest"
      essential = true
      portMappings = [{ containerPort = 8080, protocol = "tcp" }]
      environment = [
        { name = "MLFLOW_INTERNAL_URL", value = var.mlflow_internal_url },
        { name = "PROXY_API_TOKEN",     value = var.proxy_api_token },
      ]
      logConfiguration = {
        logDriver = "awslogs"
        options = {
          "awslogs-group" = aws_cloudwatch_log_group.mlflow_proxy.name
          "awslogs-region" = var.region
        }
      }
    }
  ])
}

resource "aws_ecs_service" "mlflow_proxy" {
  name            = "mlflow-proxy"
  cluster         = aws_ecs_cluster.main.id
  task_definition = aws_ecs_task_definition.mlflow_proxy.arn
  launch_type     = "FARGATE"
  desired_count   = 2

  network_configuration {
    subnets         = var.private_subnet_ids   # 和 MLflow Server 同 VPC
    security_groups = [aws_security_group.proxy.id]
  }

  load_balancer {
    target_group_arn = aws_lb_target_group.proxy.arn
    container_name   = "proxy"
    container_port   = 8080
  }
}