Single Sign‑On для MLflow, Jupyterhub и Airflow: OIDC без костылей

Современные платформы для машинного обучения (ML)  — это комплексные системы. В их состав входит множество разнообразных инструментов — от средств обработки данных до систем развертывания моделей. А по мере увеличения масштаба и сложности таких платформ на первый план выходит вопрос эффективного управления доступом и безопасностью. Решить его можно, внедрив технологию Single Sign-On (SSO), которая позволяет пользователям получать доступ сразу ко всем компонентам платформы. 

Меня зовут Дмитрий Матушкин, я инженер платформы Nova Container Platfrom в Orion soft. В этой статье мы подробно рассмотрим процесс внедрения и настройки StarVault (аналог HashiCorp Vault, но все действия похожи на те, что нужно произвести в Vault) с использованием технологии OpenID Connect (OIDC) в качестве единой точки входа для популярных компонентов ML-платформы: MLflow, Airflow и JupyterHub.   

Все данные сервисы будут развернуты в кластере Kubernetes. Для удобства развертывания и настройки ванильного кластера я буду использовать решение Nova Container Platform, которое позволяет получить готовый кластер за 10 минут. Также будем считать, что в StarVault уже создан OIDC provider, например, с названием «some_provider».

Почему именно SSO?

Почему же работать c ИИ-фреймворками без SSO сегодня очень сложно? Есть три основных причины:

1. Путаница и снижение эффективности. SSO дает единую точку входа для пользователей множества инструментов и сервисов. Среды для экспериментов, сервисы для управления жизненным циклом моделей, платформы для управления обработкой данных и т.д – без SSO каждый их этих компонентов требует отдельной аутентификации. На практике это приводит к путанице из-за множества учетных записей для каждого инструмента. И, как следствие, к вытекающим из этого рискам информационной безопасности.

2. Проблемы с разграничением доступа. В отличие от плоской модели работы с учетными записями SSO позволяет управлять доступом централизованно. Это особенно важно в командах, состоящих из сотрудников с разными обязанностями. Например, вы можете разрешить специалистам из группы Data Scientist запускать эксперименты, но изолировать от них production-окружение. MlOps инженерам вы позволяете деплоить модели, но не допускаете их к исправлению raw-данных. 

   Разграничение доступа в зависимости от ролей позволяет централизованно настраивать и применять права доступа для всех членов команды без лишних сложностей.

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

Внедряем SSO

Но давайте разберемся, как именно реализовать SSO на практике. Чтобы дальше было проще работать, мы развернули нужные сервисы в кластере Kubernetes на базе Nova Container Platform. Также мы предварительно создали в StarVault OIDC-provider с названием «some_provider», который в вашем случае, разумеется, будет называться по-другому.

Особенности настройки SSO для MLflow

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

На мой взгляд проще всего решить задачу настройки SSO в MLFlow с помощью плагина mlflow-oidc-auth, который основан на базе стандартного плагина обычной авторизации по логину и паролю basic-auth. Для работы mlflow-oidc-auth требуются 2 базы данных в PostgreSQL (в них хранятся метаданные и параметры доступа пользователей). Если вы все сделаете правильно, он позволит использовать для авторизации OpenID Connect (OIDC).

Надо учитывать, что данный плагин не установлен в базовом образе MLflow. И поэтому для работы с ним нужно пересобрать образ для добавления SSO-функционала. Для этого мы использовали следующий Dockerfile:

```Dockerfile FROM python:3.13.4 AS foundation LABEL maintainer="OrionSoft" WORKDIR /mlflow-build/ COPY pyproject.toml poetry.toml poetry.lock LICENSE README.md ./ COPY mlflowstack ./mlflowstack  RUN ln -s /usr/bin/dpkg-split /usr/sbin/dpkg-split \ && ln -s /usr/bin/dpkg-deb /usr/sbin/dpkg-deb \ && ln -s /bin/rm /usr/sbin/rm \ && ln -s /bin/tar /usr/sbin/tar  RUN apt-get update && \ apt-get install -y --no-install-recommends \ make \ build-essential \ libssl-dev \ zlib1g-dev \ libbz2-dev \ libreadline-dev \ libsqlite3-dev \ wget \ curl \ libncursesw5-dev \ xz-utils \ tk-dev \ libxml2-dev \ libxmlsec1-dev \ libffi-dev \ liblzma-dev && \ apt-get clean && \ rm -rf /var/lib/apt/lists/* /var/cache/* /var/log/* /tmp/* /var/tmp/*   RUN python -m pip install --upgrade pip --no-cache-dir && \ pip install poetry wheel --no-cache-dir RUN poetry build WORKDIR /mlflow/ RUN python -m venv .venv && \ . .venv/bin/activate && \ pip install /mlflow-build/dist/mlflowstack-1.0-py3-none-any.whl FROM python:3.13.4-slim LABEL maintainer="OrionSoft"  RUN groupadd -r -g 1001 mlflow && useradd -r -u 1001 -g mlflow -m -d /home/mlflow mlflow WORKDIR /mlflow/ RUN chown -R mlflow:mlflow /mlflow  COPY --from=foundation --chown=mlflow:mlflow /mlflow/.venv /mlflow/.venv ENV PATH=/mlflow/.venv/bin:$PATH ENV PYTHONUNBUFFERED=1  USER mlflow  CMD ["mlflow", "server", "--backend-store-uri", "sqlite:///mlflow.sqlite", "--default- artifact-root", "./mlruns", "--host=0.0.0.0", "--port=5000"] ```

В файле pyproject.toml была указана зависимость от нашего нового плагина «mlflow- oidc-auth (==5.0.1)» и на выходе был получен образ MLflow с поддержкой OIDC.

Следующим шагом нужно настроить интеграцию между MLflow и StarVault. Для этого в StarVault необходимо создать client application для MLflow, который будет работать с нашим провайдером OIDC, а также определить в нем поля Redirect URI и Assigments.

```bash   $ starvault write identity/oidc/client/mlflow \ redirect_uris="https://mlflow.example.com" \ assignments="allow_all" Success! Data written to: identity/oidc/client/mlflow $ starvault read identity/oidc/client/mlflow KeyValue  access_token_ttl24h assignments[allow_all] client_idhmXyMbH4tIResWptajk2QwgX5Fd6R7dk client_secret hvo_secret_mWDcX0C91i2H8wGGMnq7n8t4s5NXpILDu1t8irSTE5EGauiwhkCaP 8Ics38CNMvM client_typeconfidential id_token_ttl24h keydefault redirect_uris[https://mlflow.example.com]  ```

Для корректной работы с OIDC-провайдером необходимо создать следующие scope: groups, email и name. И если scope groups остается на ваше усмотрение, последние два являются обязательными, поскольку OIDC-плагин для MLflow использует их для определения почты и отображения имени в веб-интерфейсе.

```Часть Python кода oidc плагина def handle_user_and_group_management(token) -> list[str]: """Handle user and group management based on the token. Returns list of error messages or empty list.""" errors = [] email = token["userinfo"].get("email") or token["userinfo"].get("preferred_username") display_name = token["userinfo"].get("name") if not email: errors.append("User profile error: No email provided in OIDC userinfo.") if not display_name: errors.append("User profile error: No display name provided in OIDC userinfo.") if errors: return errors ... ```

Перейдем к настройке MLflow. Чтобы наша авторизация работала, необходимо создать ConfigMap с необходимыми переменными окружения для подключения к StarVault.  Эти значения должны быть загружены в переменные окружения пода с MLflow.

```yaml apiVersion: v1 kind: ConfigMap metadata: name: mlflow-env-configmap namespace: mlflow labels: app: mlflow data: OIDC_REDIRECT_URI: "https://mlflow.example.com/callback" OIDC_PROVIDER_TYPE: "oidc" OIDC_PROVIDER_DISPLAY_NAME: "sso" # отображаемое имя в веб интерфейсе OIDC_SCOPE: "openid email name groups" OIDC_GROUP_NAME: "mlflow-access" OIDC_ADMIN_GROUP_NAME: "mlflow-admins" DEFAULT_MLFLOW_PERMISSION: "MANAGE" LOG_LEVEL: "INFO" OIDC_USERS_DB_URI: "postgresql://admin:admin@psql- cls.postgresql.svc:5432/mlflow_users" # строка для подключения к базе для хранения данных пользователей SECRET_KEY: "dbAtlCg3GNY3lIjebcYM7QpsNJMEIJrH" OIDC_DISCOVERY_URL: "https://starvault.example.com/v1/identity/oidc/provider/some_provider/.well-known/ openid-configuration" OIDC_CLIENT_SECRET: "hvo_secret_mWDcX0C91i2H8wGGMnq7n8t4s5NXpILDu1t8irSTE5EGauiwhkCa P8Ics38CNMvM" OIDC_CLIENT_ID: "hmXyMbH4tIResWptajk2QwgX5Fd6R7dk" ```

Если все описанные выше действия были проделаны правильно, то после захода в веб интерфейс MLflow откроется следующая страница:

А после успешной авторизации вы попадете на домашнюю страницу MLflow.

Чтобы продолжить работу с Mlflow из консоли, например, вести трекинг экспериментов, необходимо получить токен текущего пользователя. Для этого нужно нажать кнопку «Create 

Для дальнейшей работы с Mlflow из консоли, например, для трекинга экспериментов, нужно получить токен для текущего пользователя. Для этого необходимо нажать кнопку «Create access key», после чего откроется следующая форма:

Все! Интеграция между StarVault и Mlflow настроена успешно!

Особенности настройки SSO для Airflow

Airflow использует для авторизации Flask AppBuilder (FAB) auth manager, который поддерживает несколько методов авторизации, в том числе и OAuth, но StarVault и OIDC по умолчанию в нем нет. Но зато Airflow позволяет настроить аутентификацию через кастомного провайдера методом создания своего собственного класса с наследованием от системного класса FabAirflowSecurityManagerOverride.

Интеграцию между StarVault и Airflow легче всего реализовать именно таким способом. Для этого в StarVault создается client application для Airflow, который будет работать с OIDC провайдером. В нем задаются такие поля, как Redirect URI и Assigments.

$ starvault write identity/oidc/client/airflow redirect_uris="https://airflow.example.com/oauth-authorized/sso" assignments="allow_all" Success! Data written to: identity/oidc/client/airflow  $ starvault read identity/oidc/client/airflow KeyValue  access_token_ttl24h assignments[allow_all] client_idHC53gCO2rob89DrpvrJi32mPJlefkNza client_secret hvo_secret_7J8DMIxBijR0E1WVJYwl1y9UcnnxJVrohRPvJh4Lg1JqeBYcrS7XAm vP456ya84p client_typeconfidential id_token_ttl24h   keydefault redirect_uris[https://airflow.example.com/oauth-authorized/sso]

Чтобы настроить Airflow в нем нужно создать ConfigMap для компонента веб-сервера и описать в нем класс, реализующий нужную нам логику работы.

Важно! В качестве имени в поле «name» у OAuth провайдера необходимо указать сегмент пути (в нашем случае «sso») после сегмента «oauth-authorized», заданного в поле Redirect URI.

```yaml apiVersion: v1 kind: ConfigMap metadata:  name: airflow-webserver-config  namespace: airflow  labels:  app: airflow  instance: webserver data:  webserver_config.ctmpl: |- from airflow.providers.fab.auth_manager.security_manager.override import FabAirflowSecurityManagerOverride from flask_appbuilder.security.manager import AUTH_OAUTH from typing import Any, List, Union import requests AUTH_TYPE = AUTH_OAUTH # Задание нужно AUTH_ROLES_SYNC_AT_LOGIN = True AUTH_USER_REGISTRATION = True # Позволяет пользователям, которые не созданы в БД FAB, зарегистрироваться # Задание соответствия между ролями, возвращаемыми провайдером авторизации и заданными в FAB AUTH_ROLES_MAPPING = { "Viewer": ["Viewer"], "Admin": ["Admin"], } # Задание StarVault в качестве OAuth провайдера. Значения большинства полей можно получить из эндпоинта "https://starvault.example.com/v1/identity/oidc/provider/some_provider/.well-known/ openid-configuration". OAUTH_PROVIDERS = [ {  "name": "sso",  "icon": "fa-sign-in",  "token_key": "access_token",  "remote_app": {  "client_id": "HC53gCO2rob89DrpvrJi32mPJlefkNza",  "client_secret": "hvo_secret_7J8DMIxBijR0E1WVJYwl1y9UcnnxJVrohRPvJh4Lg1JqeBYcrS7XAm vP456ya84p",  "client_kwargs": {  "scope": "openid email groups",  "token_endpoint_auth_method": "client_secret_post",  },  "server_metadata_url": "https://starvault.example.com/v1/identity/oidc/provider/some_provider/.well-known/ openid-configuration",  "api_base_url": "https://starvault.example.com/v1/identity/oidc/provider/some_provider",  "access_token_url": "https://starvault.example.com/v1/identity/oidc/provider/some_provider/token",  "authorize_url": "https://starvault.example.com/ui/vault/identity/oidc/provider/some_provider/authorize",  "jwks_uri": None,  }, }, ] # Создание собственного класса для реализации логики авторизации class StarVaultSecurityManager(FabAirflowSecurityManagerOverride): def get_oauth_user_info(self, provider: str, resp: Any) -> dict[str, Union[str, list[str]]]:  if provider == "sso":  remote = self.appbuilder.sm.oauth_remotes[provider]  # Получение токена для отправки запросов в StarVault  access_token = resp.get('access_token')  # Формирование URL и заголовков для получения информации о пользователе  userinfo_url = f"{remote.api_base_url}/userinfo"  headers = {  "Authorization": f"Bearer {access_token}",  "Accept": "application/json"  }  # Отправка запроса на получение информации о пользователе  response = requests.get(userinfo_url, headers=headers)  data = response.json()  # возращаем почту и роль (в данном случае значение роли совпадает с названием группы)  return {  "email": data.get("email", ""),  "role_keys": data.get("groups", []),  }   # Указание использования собственного класса для авторизации SECURITY_MANAGER_CLASS = VaultSecurityManager ```

Этот файл необходимо примонтировать в pod веб-сервера, используя путь «/opt/airflow/webserver_config.py«. 

Если вам необходимо определить другой путь для данного файла (например, когда происходят конфликты при монтировании других томов), то в файле «airflow.cfg» в поле «[webserver]» необходимо будет отдельно указать параметр с требуемым путем, например «config_file =/opt/airflow/webserver/webserver_config.py«.

Если все описанные выше действия были проделаны правильно, то при входе на веб-интерфейс Airflow откроется следующая страница:

Интеграция между StarVault и Airflow успешно настроена!

Особенности настройки SSO для Jupyterhub

Jupyterhub – это многопользовательский сервер с возможностью создания и запуска Jupyter Notebooks из одного интерфейса. Для авторизации Jupyterhub использует встроенный модуль oauthenticator, который поддерживает интеграции как для заранее приготовленных платформ (например, GitLab или Google), так и с помощью GenericAuthenticator, который позволяет настроить подключение к любому провайдеру.

Чтобы настроить интеграцию между StarVault и Jupyterhub в StarVault нужно создать client application для Jupyterhub, который будет работать с провайдером OIDC, и определить в нем такие поля как Redirect URI и Assigments.

```bash $ starvault write identity/oidc/client/jupyterhub redirect_uris="https://jupyterhub.example.com/hub/oauth_callback" assignments="allow_all" Success! Data written to: identity/oidc/client/jupyterhub $ starvault read identity/oidc/client/jupyterhub Key Value --- ----- access_token_ttl 24h assignments [allow_all] client_id vuuYyveqysCc5BqmbfFiUJ9naGs2M4kc client_secret hvo_secret_qod6qYjwy16oeYZcvz0fbU1B2pTJ0skPR9dCWZ45ZNG1ib1BeGUNK AmCQeTFfQR1 client_type confidential id_token_ttl 24h key default redirect_uris [https://jupyterhub.example.com/hub/oauth_callback] ```

Для настройки Jupyterhub добавляем в файл jupyterhub_config.py следующие строки:

```yaml # Указание нужного варианта настройки авторизации c.JupyterHub.authenticator_class = "generic-oauth" # Задание полей о клиенте OIDC c.GenericOAuthenticator.client_id = "vuuYyveqysCc5BqmbfFiUJ9naGs2M4kc" c.GenericOAuthenticator.client_secret = "hvo_secret_qod6qYjwy16oeYZcvz0fbU1B2pTJ0skPR9dCWZ45ZNG1ib1BeGUN KAmCQeTFfQR1" # Задание информации о провайдере. Значения данных полей можно получить из эндпоинта "https://starvault.example.com/v1/identity/oidc/provider/some_provider/.well-known/ openid-configuration". c.GenericOAuthenticator.authorize_url = "https://starvault.example.com/ui/vault/identity/oidc/provider/some_provider/authorize" c.GenericOAuthenticator.token_url = "https://starvault.example.com/v1/identity/oidc/provider/some_provider/token" c.GenericOAuthenticator.userdata_url = "https://starvault.example.com/v1/identity/oidc/provider/some_provider/userinfo" # Настройка информации о пользователе c.GenericOAuthenticator.scope = ["openid", "email", "groups"] # Указываем, что в качестве username следует использовать email c.GenericOAuthenticator.username_claim = "email" # Задаем соответствие между полями для определения группы пользователя c.GenericOAuthenticator.auth_state_groups_key = "oauth_user.groups" # Настройка авторизации c.GenericOAuthenticator.allowed_groups = {"jupyterhub_users"} c.GenericOAuthenticator.admin_groups = {"jupyterhub_admins"} ```

Если все описанные выше действия были проделаны правильно, то после захода в веб интерфейс Jupyterhub откроется следующая страница:

Заключение

Время подводить итоги. Что мы сделали: 

1. Настроили единую точку входа для популярных ML сервисов, тем самым упростили жизнь ML-инженерам.

2. Получили возможность настройки ролевой модели для доступа к конечным сервисам при помощи Vault. 

3. Настроили централизованное управление доступом на базе StarVault к конечным сервисам.

Если у вас есть вопросы по настройке авторизации такого типа или имеется собственный опыт внедрения SSO в ML, обязательно пишите – обсудим это в комментариях. Также пишите, на какую тему стоит написать следующую статью, связанную с ML/AI 🙂

Бодрой всем нам SSO-авторизации!