Современные платформы для машинного обучения (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 каждый их этих компонентов требует отдельной аутентификации. На практике это приводит к путанице из-за множества учетных записей для каждого инструмента. И, как следствие, к вытекающим из этого рискам информационной безопасности.

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

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

  2. Вопросы информационной безопасности. Обилие логинов и паролей создает риски компрометации учетных записей. Использование 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:

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 \ l
      iblzma-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.

$ 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 
Key	Value
--- -----
access_token_ttl 24h 
assignments	[allow_all]
client_id hmXyMbH4tIResWptajk2QwgX5Fd6R7dk 
client_secret hvo_secret_mWDcX0C91i2H8wGGMnq7n8t4s5NXpILDu1t8irSTE5EGauiwhkCaP8Ics38CNMvM
client_type	confidential 
id_token_ttl 24h
key	default
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.

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 
Key	Value
--- -----
access_token_ttl 24h 
assignments	[allow_all]
client_id HC53gCO2rob89DrpvrJi32mPJlefkNza 
client_secret hvo_secret_7J8DMIxBijR0E1WVJYwl1y9UcnnxJVrohRPvJh4Lg1JqeBYcrS7XAmvP456ya84p
client_type	confidential 
id_token_ttl 24h 
key	default
redirect_uris [https://airflow.example.com/oauth-authorized/sso]

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

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

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_7J8DMIxBijR0E1WVJYwl1y9UcnnxJVrohRPvJh4Lg1JqeBYcrS7XAmvP456ya84p",
       "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 = StarVaultSecurityManager

Этот файл необходимо примонтировать в 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.

$ 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_qod6qYjwy16oeYZcvz0fbU1B2pTJ0skPR9dCWZ45ZNG1ib1BeGUNKAmCQeTFfQR1
client_type confidential
id_token_ttl 24h
key default
redirect_uris [https://jupyterhub.example.com/hub/oauth_callback]

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

# Указание нужного варианта настройки авторизации
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. Получили возможность настройки ролевой модели для доступа к конечным сервисам при помощи StarVault. 

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

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

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

Комментарии (0)