Привет! Я Саша Абакумов, DevOps-инженер в KTS.

Нашей команде часто приходится поднимать инфраструктуру под ML-проекты. Со временем число ML-инженеров и разработчиков на таких проектах росло, и логиниться в каждый по отдельности становилось все больнее. Чтобы упростить коллегам жизнь, мы интегрировали Single Sign-On (SSO) в стек одного из наших проектов, состоящий из JupyterHub, Airflow и MLflow.

SSO позволяет единообразно аутентифицироваться во всех инструментах под одной учетной записью. Помимо очевидного удобства, нам это также дало возможность централизованно управлять доступом и внедрить RBAC — сопоставление ролей в инструментах с группами или ролями в IdP.

В качестве инструмента для реализации SSO я использовал OIDC-провайдер Keycloak, наверняка многим хорошо знакомый. Ниже я расскажу о том, как с его помощью настроить SSO для JupyterHub, MLflow и Airflow (все компоненты разворачиваются с помощью Helm-чартов).

Оглавление

• JupyterHub

• Передача токена MLflow в ноутбук

• MLflow

  • Интерфейс и управление доступами

• Airflow

• Выводы

JupyterHub

JupyterHub изначально поддерживает OAuth-аутентификацию через OAuthenticator. Его можно настроить на внешний IdP, такой как GitHub или Keycloak, для аутентификации пользователей. В нашем случае мы используем GenericOAuthenticator с Keycloak. Основные шаги интеграции JupyterHub с Keycloak:

1. Сначала в Keycloak создаётся OAuth-клиент для JupyterHub. Вы получите client_id, client_secret, настроите Redirect URI (например, https://jupyterhub.example.com/hub/oauth_callback) и разрешенные URL.

2. Далее в Helm-чарт JupyterHub (Zero to JupyterHub) указываем класс аутентификации GenericOAuthenticator и настраиваем его параметры.

Пример конфигурации JupyterHub в Helm values:

hub:
  config:
    JupyterHub:
      authenticator_class: generic-oauth
    GenericOAuthenticator:
      client_id: "jupyterhub-client"  
      client_secret: "KEYCLOAK_CLIENT_SECRET"  
      oauth_callback_url: "https://jupyterhub.example.com/hub/oauth_callback"
      authorize_url: "https://keycloak.example.com/realms/MLRealm/protocol/openid-connect/auth"
      token_url: "https://keycloak.example.com/realms/MLRealm/protocol/openid-connect/token"
      userdata_url: "https://keycloak.example.com/realms/MLRealm/protocol/openid-connect/userinfo"
      username_claim: email
      scope:
        - openid
        - email
        - profile
      enable_auth_state: true
      allow_all: true

Здесь JupyterHub будет перенаправлять пользователя на Keycloak для входа и п��лучать от него токен. Параметр enable_auth_state: true включает сохранение расширенного состояния аутентификации (например, OIDC-токена) в базе JupyterHub (подробнее см. документацию). Это важно для передачи токена другим сервисам.

1. После такой настройки при попытке логина JupyterHub отправит пользователя на страницу Keycloak. Успешный вход вернет пользователя обратно, JupyterHub создаст учетную запись (при включенном allow_all: true допускаются все авторизованные Keycloak-пользователи) и сохранит токен Keycloak во внутренней базе (Auth State).

Передача токена MLflow в ноутбук

Одно из требований, с которыми я работал над настройкой SSO — обеспечить автоматическую авторизацию пользователя в MLflow, когда он работает в Jupyter Notebook. Мы реализовали это через pre-spawn hook JupyterHub. Этот хук выполняется перед запуском отдельного сервера (ноутбука) пользователя:

async def pre_spawn_hook(spawner):
    # Получаем сохраненный auth_state с Keycloak-токеном
    auth_state = await spawner.user.get_auth_state()
    oidc_token = auth_state["access_token"]
    username = auth_state["oauth_user"]["email"]

    # Задаем переменные окружения для MLflow в контейнере пользователя
    spawner.environment["MLFLOW_TRACKING_USERNAME"] = username
    spawner.environment["MLFLOW_TRACKING_URI"] = "https://mlflow.example.com"

    # Запрашиваем персональный токен доступа MLflow через API
    headers = {"Authorization": f"Bearer {oidc_token}"}
    url = "https://mlflow.example.com/api/2.0/mlflow/users/access-token"
    async with aiohttp.ClientSession() as session:
        async with session.get(url, headers=headers) as resp:
            if resp.status != 200:
                body = await resp.text()
                raise Exception(f"Failed to fetch MLflow token: {resp.status} {body}")
            token_json = await resp.json()
            spawner.environment["MLFLOW_TRACKING_PASSWORD"] = token_json["token"]

c.KubeSpawner.pre_spawn_hook = pre_spawn_hook

Как это работает:

1. При запуске пользовательского ноутбука JupyterHub извлекает OIDC Access Token (JWT) выданный Keycloak (oidc_token). Этот токен позволяет от имени пользователя обращаться к сервисам, доверяющим тому же Keycloak.

2. Далее хук устанавливает переменные окружения MLFLOW_TRACKING_URI (URL сервера MLflow) и MLFLOW_TRACKING_USERNAME (например, email пользователя).

3. Затем выполняется запрос к API MLflow:/api/2.0/mlflow/users/access-token – специальный эндпоинт, возвращающий персональный токен MLflow для данного пользователя. Мы передаем Keycloak-токен в заголовке Authorization. Если MLflow успешно распознал пользователя, он возвращает уникальный token – эквивалент пароля (API Token) для MLflow. Этот токен сохраняется в MLFLOW_TRACKING_PASSWORD в окружении ноутбука.

В результате пользователь автоматически получает настроенные переменные внутри Jupyter Notebook для MLflow Tracking. Любые вызовы mlflow (Python-клиент) будут использовать эти учетные данные. Таким образом, ML-инженер может сразу начинать логировать эксперименты в MLflow, не выполняя отдельный вход в MLflow вручную.

MLflow

По умолчанию MLflow поддерживает только простую базовую аутентификацию (username/password). Чтобы интегрировать его с SSO, мы воспользовались опенсорсным плагином mlflow-oidc-auth. Этот плагин добавляет поддержку OpenID Connect для MLflow, включая SSO и централизованное управление пользователями через провайдера OIDC. Также он обеспечивает аутентификацию через OIDC для UI и API, управление пользователями и группами и разграничение доступа по ролям/группам.

Шаги интеграции:

1. Первым делом мы собрали custom Docker-образ MLflow, включив в него плагин (можно установить через pip install mlflow-oidc-auth[full]). Запуск MLflow сервера производится с параметром --app-name oidc-auth (это указание использовать плагин). В Helm values переопределяем команду запуска сервера, добавив этот аргумент.

2. Далее мы настроили переменные окружения, через которые конфигурируется плагин. Вот основные из них:

   • OIDC_CLIENT_ID — ID OIDC-клиента для MLflow, как настроено в Keycloak (например, "mlflow-client");

   • OIDC_CLIENT_SECRET — секрет этого клиента;

   • OIDC_DISCOVERY_URL — URL OpenID Connect Discovery для Keycloak. Обычно это https://<Keycloak>/realms/<Realm>/.well-known/openid-configuration. По этому URL плагин узнает адреса авторизации, токен-эндпоинта, публичные ключи и пр.;

   • OIDC_REDIRECT_URI — куда Keycloak перенаправит браузер после успешного входа. Должен совпадать с настроенным redirect URI клиента. Например: https://mlflow.example.com/callback;

   • OIDC_PROVIDER_DISPLAY_NAME — название провайдера для отображения на кнопке входа (например, "Login with Keycloak" или название вашей SSO-системы);

   • OIDC_SCOPE — скоуп OIDC. В нашем случае это "openid email profile groups", так как нам нужно получать базовые данные профиля и группы пользователя;

   • OIDC_GROUP_NAME — группа пользователей, которой разрешен вход в MLflow. Мы указали здесь имя группы Keycloak, объединяющей всех пользователей ML-стека (например, MLOPS_RO). Плагин будет пропускать только тех, у кого в токене присутствует эта группа;

   • OIDC_ADMIN_GROUP_NAME — имя группы, члены которой получают администраторские права в MLflow (например, MLOPS_RW для команды с правами записи/управления). Эти пользователи смогут удалять эксперименты, менять права и т.д.;

   • DEFAULT_MLFLOW_PERMISSION — уровень доступа по умолчанию для пользователей из обычной группы (не админов);

   • OIDC_USERS_DB_URI — строка подключения к базе данных для хранения пользователей/токенов MLflow (плагин ведет свою таблицу пользователей и привязанные токены).

Что происходит после настройки:

1. При попытке доступа к интерфейсу MLflow пользователь увидит кнопку входа через Keycloak.

2. При нажатии на кнопку входа выполняется стандартный OIDC-флоу: редирект на Keycloak, ввод логина, возврат на MLflow. В этот момент плагин проверит, состоит ли пользователь в разрешенной группе (OIDC_GROUP_NAME):

   • Если нет, вход будет запрещен.

   • Если да, плагин создаст (или обновит) запись пользователя во внутренней базе MLflow и применит ему права: участнику группы по умолчанию — права, заданные DEFAULT_MLFLOW_PERMISSION, участнику admin-группы — права администратора.

Интерфейс и управление доступами

При попытке доступа к интерфейсу MLflow пользователь увидит кнопку входа через Keycloak:

При нажатии выполняется стандартный OIDC-флоу — пользователь перенаправляется на страницу входа Keycloak:

После успешной аутентификации плагин проверяет группы пользователя и применяет соответствующие права доступа.

Плагин предоставляет веб-интерфейс для управления доступами (MLflow Permission Manager), где администраторы могут:

• управлять правами пользователей на уровне экспериментов, моделей и промптов;

• создавать персональные access key для программного доступа к MLflow API;

• назначать групповые права доступа.

Для входа в интерфейс управления используется кнопка "Login to MLflow" в главном меню. Пользователи могут самостоятельно создавать access key через кнопку "Create access key" для использования в скриптах и CI/CD пайплайнах.

Важно, что плагин поддерживает и работу API: MLflow-клиент может использовать Basic Auth с email/токеном (именно такой токен мы получали в pre-spawn hook JupyterHub). Плагин признает Basic Auth, сравнивая имя пользователя и токен с записями в своей базе. Соответственно, API-запросы с токеном от Keycloak успешно авторизуются.

Airflow

Для Apache Airflow в версии 3.x интеграция с SSO также возможна (здесь и далее примеры для официального Apache Airflow Helm chart), хотя из коробки Airflow не имел готового провайдера для OIDC. Для интеграции мы задействовали механизм Flask AppBuilder (FAB) Security. По сути, это внутренняя система Airflow для аутентификации, которую можно расширить под OAuth.

Благодаря использованию Keycloak мы получаем не только единый вход, но и возможность маппить роли Keycloak на роли Airflow. Таким образом, управлять пользователями в Airflow вручную нам также больше не приходится.

Подготовка Airflow к OAuth:

1. Создаем кастомный Docker-образ Airflow на базе официального. В него устанавливаем необходимые компоненты:

   • Провайдер FAB: apache-airflow-providers-fab.

   • Библиотека OAuth для Flask AppBuilder: flask-appbuilder[oauth] (она привносит поддержку OAuth2/OIDC в FAB).

   • Библиотеки для OIDC: authlib (используется FAB для OAuth), python-jose[cryptography] или PyJWT для декодирования JWT, если нужно.

2. Далее задаем необходимую конфигурацию Airflow (config.py). В Helm values в секции apiServer.apiServerConfig настраиваем следующие параметры безопасности:

   • Тип аутентификации: AUTH_TYPE = AUTH_OAUTH (переводим FAB в режим внешнего OAuth).

   • Разрешить саморегистрацию: AUTH_USER_REGISTRATION = True — новые пользователи автоматически создаются при первом входе.

   • Синхронизация ролей: AUTH_ROLES_SYNC_AT_LOGIN = True — при каждом логине обновлять роли пользователя согласно текущим данным (важно, если групповые роли могли измениться).

   • Роль по умолчанию для новых пользователей: AUTH_USER_REGISTRATION_ROLE = "Viewer" (гость/чтение — минимальные права).

3. Следующим шагом настраиваем маппинг ролей через AUTH_ROLES_MAPPING. Здесь мы сопоставляем названия групп/ролей Keycloak с ролями Airflow. Например:
   AUTH_ROLES_MAPPING = {
    "MLOPS_RW": ["Admin"],  # группа Keycloak -> роль Airflow
    "MLOPS_RO": ["User"],   # (прочие пользователи получат роль Viewer по умолчанию)
}
   Если в токене пользователя есть группа MLOPS_RW, он получит роль Admin в Airflow; если есть MLOPS_RO — роль User (обычный пользователь с доступом к своим дагам). Все остальные, не имеющие этих групп, по умолчанию останутся с ролью Viewer.

4. Задаем необходимую конфигурацию OAuth-провайдера. В переменной OAUTH_PROVIDERS описываем наш Keycloak:
   OAUTH_PROVIDERS = [{
    "name": "keycloak",
    "icon": "Keycloak",
    "token_key": "access_token",
    "remote_app": {
        "client_id": "airflow-client",
        "client_secret": os.getenv("KC_CLIENT_SECRET"),
        "server_metadata_url": "https://keycloak.example.com/realms/MLRealm/.well-known/openid-configuration",
        "api_base_url": "https://keycloak.example.com/realms/MLRealm/protocol/openid-connect",
        "access_token_url": "https://keycloak.example.com/realms/MLRealm/protocol/openid-connect/token",
        "authorize_url": "https://keycloak.example.com/realms/MLRealm/protocol/openid-connect/auth",
        "request_token_url": None,
        "client_kwargs": {"scope": "openid email profile groups"}
    },
}]

   Здесь мы использовали server_metadata_url, чтобы Airflow сам подтянул нужные эндпоинты и публичный ключ Realm’а. client_id и client_secret соответствуют настроенному в Keycloak клиенту Airflow.

5. Настраиваем получение публичного ключа. Поскольку Airflow (Flask AppBuilder) самостоятельно не валидирует токены по сигнатуре, мы можем заранее получить публичный ключ Keycloak и использовать его для декодирования JWT. Например:
   
   import requests, jwtfrom base64 import b64decodefrom cryptography.hazmat.primitives import serialization

OIDC_ISSUER = "https://keycloak.example.com/realms/MLRealm"resp = requests.get(OIDC_ISSUER, timeout=5)public_key_der = b64decode(resp.json()["public_key"])PUBLIC_KEY = serialization.load_der_public_key(public_key_der)
   
   Эндпоинты Keycloak Realm (через /.well-known/openid-configuration или прямым GET на URL Realm) обычно возвращают ключ в поле "public_key", который мы декодируем из Base64. Этот ключ нужен для проверки JWT.

6. Добавляем кастомный Security Manager. Для этого мы сначала определяем класс, наследующий стандартный FAB SecurityManager (в Airflow 3 он называется FabAirflowSecurityManagerOverride), а затем переопределяем в нем метод get_oauth_user_info(self, provider, response). Здесь реализуется извлечение информации о пользователе из токена Keycloak:
   class KeycloakSecurityManager(FabAirflowSecurityManagerOverride):
    def get_oauth_user_info(self, provider, response):
        if provider != "keycloak":
            return {}
        token = response["access_token"]
        # Декодируем JWT без проверки аудитории
        payload = jwt.decode(token, PUBLIC_KEY, algorithms=["RS256", "HS256"], options={"verify_aud": False})
        # Извлекаем username и email
        username = payload.get("preferred_username")
        user_info = {
            "username": username,
            "email": payload.get("email"),
            "first_name": payload.get("given_name"),
            "last_name": payload.get("family_name"),
            # Соберем все группы/роли пользователя
            "role_keys": []
        }
        # Пример: если используем realm roles:
        realm_roles = payload.get("realm_access", {}).get("roles", [])
        if realm_roles:
            user_info["role_keys"].extend(realm_roles)
        # Пример: если группы приходят в claim "groups":
        groups = payload.get("groups", [])
        if groups:
            # Keycloak может включать путь группы, берем только имя последней подгруппы
            clean_names = [g.split("/")[-1] for g in groups]
            user_info["role_keys"].extend(clean_names)
        return user_info
SECURITY_MANAGER_CLASS = KeycloakSecurityManager
   
   В этой реализации мы берем из токена preferred_username (в Keycloak обычно совпадает с email prefix или UID пользователя) в качестве username в Airflow. Далее собираем список меток пользователя: realm roles и группы (в зависимости от того, как настроен Keycloak). Все найденные ключи помещаем в role_keys. Airflow затем сопоставит эти ключи с AUTH_ROLES_MAPPING, то есть присвоит пользователю роли согласно нашей таблице соответствий. Например, если в role_keys есть "MLOPS_RW", то пользователь получит роль Admin в Airflow.

7. Наконец, убеждаемся, что Airflow использует FabAuthManager вместо базового (В Airflow 3.0 по умолчанию стоял SimpleAuthManager, не поддерживающий OAuth). В airflow.cfg (или overrideConfiguration в helm) устанавливаем:

   [core]
auth_manager = airflow.providers.fab.auth_manager.fab_auth_manager.FabAuthManager
   
   После этого API-сервер Airflow будет запускаться с нашей конфигурацией.

Авторизация в Airflow будет выглядеть так:

1. На странице входа Airflow появится опция «Login with Keycloak».

2. При нажатии пользователь будет перенаправлен на Keycloak и выполнит вход, после чего его редиректнет обратно.

3. По возвращении Airflow создаст его учетную запись (если ее нет) и назначит роли.

При каждом входе роли Airflow автоматически маппятся на актуальное состояние групп в Keycloak. Благодаря этому управление доступом становится централизованным: достаточно в Keycloak добавить человека в нужную группу, и при следующем входе в Airflow у него будут соответствующие права.

Выводы

1. Мы получили единое пространство аутентификации для всех компонентов ML-стека. Пользователи используют корпоративный логин (Keycloak) для доступа к JupyterHub, MLflow и Airflow.

2. Интеграция JupyterHub с MLflow позволяет ML-инженерам отслеживать эксперименты в MLflow прямо из ноутбуков без повторной авторизации.

3. Администраторы управляют доступами в одном месте — через группы и роли в Keycloak. Это:

   а) повышает безопасность (можно применять политики MFA, автоматическое отключение доступа) и снижает шанс рассинхронизации прав;

   б) добавляет масштабируемости — добавление новой роли или изменения прав делается через конфигурацию, без переработки кода. MLflow-плагин поддерживает управление на уровне экспериментов и моделей, а Airflow — на уровне ролей RBAC, что покрывает требования большинства команд.

В итоге, SSO с Keycloak значительно упростил жизнь нашим ML-инженерам и админам. Новый сотрудник теперь проходит единую процедуру подключения: его учетная запись появляется в Keycloak, добавляется в нужные группы, и он сразу получает доступ ко всем инструментам ML-платформы с соответствующими разрешениями.

Надеюсь, этот опыт пригодится и вам. А узнать больше о том, как мы работаем с инфраструктурой, вы можете в других статьях нашего блога:

• Тонкости обновления драйверов NVIDIA в Yandex Managed Kubernetes

• BuildKit в Kubernetes: мануал по быстрой и автомасштабируемой сборке проектов

• Infrastructure as Code на практике: как мы рефакторили сложный Ansible-репозиторий

• Firezone, или как спрятать свою инфраструктуру от посторонних глаз