Безопасный HTTPS-доступ к SageMaker MLflow без SDK: прокси на Flask + IAM

MLflow в SageMaker — штука удобная, пока не упираешься в корпоративную безопасность. Типичная картина: у вас есть managed MLflow Tracking Server, запущенный через CreateMlflowTrackingServer, вы работаете из Jupyter Notebook или с удаленной машины, но ваш сетевой админ заворачивает все порты, кроме 443, и требует HTTPS. Вдобавок SDK (boto3 + MLflow-клиент) часто блокируется политиками — либо из-за установки зависимостей, либо из-за необходимости прямого доступа к API с подписью SigV4, которую SDK делает за кулисами. Решение: построить легковесный REST API прокси на Flask, который берет на себя аутентификацию через IAM и выставляет наружу чистый HTTPS. Без SDK, без головной боли.

Почему SDK — зло (или хотя бы не везде)

SDK удобен, когда у вас есть полный контроль над средой. Но в реальном мире часто:

• Клиент — микросервис на Go или Rust, где тащить boto3 тяжело.
• CI/CD пайплайн запускается в контейнере без доступа к AWS-метаданным.
• Команда использует MLflow Tracking API через curl для быстрых экспериментов, но сталкивается с ошибкой SignatureDoesNotMatch (кто не проклинал ручную подпись SigV4?).

Именно в таких случаях прокси-сервер на Flask — спасательный круг. Вы один раз настраиваете подпись запросов к SageMaker MLflow API внутри безопасного сегмента (например, в VPC), а наружу торчите только чистый REST с авторизацией по ключу или JWT. Я уже писал про OpenAI-совместимое API без SigV4 — подход похож, но тут мы идем глубже: строим полноценный reverse proxy с кастомной аутентификацией.

Архитектура решения

Наша прокси (назовем ее mlflow-proxy) работает так:

1. Принимает внешние HTTPS-запросы на /api/2.0/mlflow/...
2. Проверяет API-ключ (или JWT) в заголовке X-API-Key.
3. Подписывает исходный URL и тело запроса по правилам SigV4, используя IAM-роль (или временные credentials из STS).
4. Проксирует подписанный запрос к эндпоинту SageMaker MLflow Tracking Server (который всегда работает через HTTPS).
5. Возвращает ответ клиенту.

По желанию можно генерировать presigned URL (срок жизни 1 час) для прямого доступа к определённым ресурсам (например, для загрузки артефактов). Это удобно, когда нельзя проксировать большие файлы через свой сервер.

Пошагово: пишем прокси

1 Настройка IAM-роли

Создайте IAM-роль (например, MLflowProxyRole) с политикой:

{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Effect": "Allow",
      "Action": [
        "sagemaker-mlflow:InvokeMlflowTrackingServer",
        "sagemaker-mlflow:GetMlflowTrackingServer",
        "sagemaker-mlflow:ListMlflowTrackingServers"
      ],
      "Resource": "arn:aws:sagemaker:*:ACCOUNT_ID:mlflow-tracking-server/*"
    },
    {
      "Effect": "Allow",
      "Action": [
        "s3:PutObject",
        "s3:GetObject"
      ],
      "Resource": "arn:aws:s3:::BUCKET_NAME/*"
    }
  ]
}

Роль привяжите к инстансу EC2, ECS-таске или лямбде, где будет крутиться прокси. Если запускаете на своей машине — используйте aws configure с профилем, имеющим права на sts:AssumeRole.

2 Flask-приложение (версия 3.1)

Устанавливаем Flask (на момент 28.05.2026 актуальна версия 3.1.2), requests и boto3 (1.38+). Код приложения:

import os
import json
import hmac
import hashlib
import base64
from datetime import datetime, timezone
from flask import Flask, request, jsonify, Response
from requests import request as http_request
from requests_auth_aws_sigv4 import AWSSigV4  # pip install requests-aws-sigv4

app = Flask(__name__)

MLFLOW_SERVER_URL = os.environ["MLFLOW_TRACKING_SERVER_ARN"]  # или URL эндпоинта
API_KEY = os.environ["PROXY_API_KEY"]
REGION = os.environ.get("AWS_REGION", "us-east-1")

@app.before_request
def check_auth():
    if request.headers.get("X-API-Key") != API_KEY:
        return jsonify({"error": "Unauthorized"}), 401

@app.route("/presign-upload", methods=["GET"])
def presign_url():
    """Генерация presigned URL для S3 (загрузка артефактов)"""
    s3_client = boto3.client("s3", region_name=REGION)
    bucket = request.args.get("bucket")
    key = request.args.get("key")
    if not bucket or not key:
        return jsonify({"error": "Missing bucket or key"}), 400
    url = s3_client.generate_presigned_url(
        ClientMethod="put_object",
        Params={"Bucket": bucket, "Key": key},
        ExpiresIn=3600
    )
    return jsonify({"url": url})

@app.route("/api/2.0/mlflow/", methods=["GET", "POST"])
def proxy(subpath):
    target_url = f"{MLFLOW_SERVER_URL}/api/2.0/mlflow/{subpath}"
    method = request.method
    headers = {k: v for k, v in request.headers if k.lower() not in ["host", "x-api-key", "content-length"]}
    data = request.get_data()
    
    # Подписываем запрос SigV4
    auth = AWSSigV4("sagemaker-mlflow", region=REGION)
    signed = auth(method, target_url, data=data, headers=headers)
    
    # Выполняем подписанный запрос
    resp = http_request(method, target_url, data=data, headers=signed.headers, verify=True)
    return Response(resp.content, status=resp.status_code, headers=dict(resp.headers))

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=8080, ssl_context=("/certs/cert.pem", "/certs/key.pem"))

Важный нюанс: библиотека requests-aws-sigv4 (версия 1.2.0) автоматически вычисляет заголовок Authorization, X-Amz-Date и X-Amz-Security-Token (если используются временные креды). Без неё пришлось бы вручную считать HMAC — делать этого не советую, легко ошибиться. Альтернатива — встроенный aws_sigv4 из boto3, но он работает на уровне сессии, а не отдельных запросов.

3 Запускаем с HTTPS

Никакого HTTP наружу — только HTTPS. Для локального тестирования сгенерируйте самоподписанные сертификаты (openssl), для продакшена используйте AWS Certificate Manager (ACM) на ALB или NGINX. Рекомендуемая конфигурация:

• Запускаете прокси за внутренним Application Load Balancer.
• ALB терминирует TLS (сертификат из ACM), прокси работает по HTTP:8080.
• У ALB включите WAF для дополнительной защиты.

Так мы уменьшаем поверхность атаки и упрощаем ротацию сертификатов. Подобную архитектуру я описывал в статье про serverless conversational AI agent — там тоже использовал внутренний прокси для аутентификации.

Как это выглядит со стороны клиента?

Клиент (curl, Postman, код на Python без SDK) шлёт запросы на публичный эндпоинт прокси:

curl -X POST "https://mlflow-proxy.example.com/api/2.0/mlflow/runs/create" \
  -H "X-API-Key: secret-key-123" \
  -H "Content-Type: application/json" \
  -d '{"experiment_id": "1", "run_name": "test-run"}'

Прокси отвечает так же, как и сам MLflow Server — никаких дополнительных врапперов. Если нужно загрузить артефакт, клиент сначала получает presigned URL через /presign-upload, а потом пишет напрямую в S3. Это снижает нагрузку на прокси и ускоряет передачу.

Нюансы и частые грабли

Развёртывание: ECS Fargate vs Lambda

Для небольшой нагрузки (<5 запросов/сек) подойдёт AWS Lambda с помощью контейнерного рантайма (образ на базе python:3.13-slim). Flask в Lambda запускается через aws-lambda-python с адаптером Mangum. Пример:

from mangum import Mangum
handler = Mangum(app)

Для высоконагруженных систем используйте ECS Fargate с автоматическим масштабированием. Я рекомендую ECS — меньше cold start, проще логи через FireLens. В статье ModelOps без головной боли описан схожий подход с S3-шаблонами — там тоже используется лёгкий прокси-слой.

FAQ

Можно ли использовать этот прокси для MLflow на собственном сервере (не SageMaker)?

Да, просто замените MLFLOW_SERVER_URL на адрес вашего сервера. SigV4 отключается, остаётся только проверка ключа. Но смысл тогда теряется — проще поставить nginx с basic auth.

Как обновлять версию MLflow на SageMaker?

SageMaker автоматически обновляет managed Tracking Server до новой версии MLflow (по состоянию на 28.05.2026 — MLflow 2.18). Прокси не требует изменений, так как API обратно совместим.

Что если утечёт API-ключ?

Смените PROXY_API_KEY в переменных окружения и перезапустите сервис. Также настройте CloudWatch Logs для мониторинга подозрительной активности. Утечка ключа не даст доступ к MLflow напрямую — только к прокси, который ограничен вашей IAM-ролью.

Вместо заключения: куда движемся?

С ростом требований к безопасности такие прокси станут стандартом — уже сейчас многие компании отказываются от прямых SDK-вызовов в пользу «чистых» REST с кастомной авторизацией. Особенно это касается гибридных сред, где часть пайплайнов крутится on-prem, а MLflow — в облаке. Я предсказываю, что в ближайшие полгода AWS выпустит официальный механизм для «API Gateway перед MLflow», но пока учимся делать всё руками — это даёт полный контроль. Кстати, в наших сквозных трекингах на DVC и MLflow такой прокси уже используется в продакшне — полет нормальный.

Подписаться на канал