Поделиться через

Трассировка MLflow для агентов

  • Статья

Внимание

Эта функция предоставляется в режиме общедоступной предварительной версии.

В этой статье описывается трассировка MLflow в Databricks и его использование для добавления наблюдаемости в созданные приложения ИИ.

Что такое трассировка MLflow?

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

Трассировка MLflow интегрирована с инструментами и инфраструктурой Databricks, что позволяет хранить и отображать трассировки в записных книжках Databricks или пользовательском интерфейсе эксперимента MLflow.

Встроенная трассировка фиксирует подробные сведения на каждом этапе в приложении генеративного ИИ

Зачем использовать трассировку MLflow?

Трассировка MLflow обеспечивает ряд преимуществ.

  • Просмотрите интерактивную визуализацию трассировки и используйте средство исследования для диагностики проблем.
  • Убедитесь, что шаблоны запросов и ограничители дают разумные результаты.
  • Анализ задержки различных платформ, моделей и размеров блоков.
  • Оцените затраты на программное приложение, анализируя количество используемых токенов в разных моделях.
  • Создайте эталонные наборы данных для оценки производительности различных версий.
  • Храните трассировки из конечных точек рабочей модели для отладки проблем и выполнения автономной проверки и оценки.

Добавление трассировок в агент

Трассировка MLflow поддерживает три метода добавления трассировок в созданные приложения ИИ. Дополнительные сведения о API см. в документации по MLflow.

API Рекомендуемый вариант использования Description
автоматическое логирование в MLflow Разработка с интегрированными библиотеками GenAI Автоматическое ведение журнала автоматически регистрирует трассировки для поддерживаемых платформ с открытым кодом, таких как LangChain, LlamaIndex и OpenAI.
API Fluent Настраиваемый агент с Pyfunc API с низким кодом для добавления трассировок, не беспокоясь об управлении структурой дерева трассировки. MLflow определяет соответствующие связи между родительскими и дочерними элементами автоматически с помощью стека Python.
API клиента MLflow Расширенные варианты использования, такие как многопоточное использование MLflowClient предоставляет детализированные, потокобезопасные API для расширенных вариантов использования. Необходимо вручную управлять отношениями между родительскими и дочерними элементами. Это позволяет лучше контролировать жизненный цикл трассировки, особенно для вариантов использования с несколькими потоками.

Установка трассировки MLflow

Трассировка MLflow доступна в MLflow версии 2.13.0 и выше, которая предварительно установлена в <DBR< 15.4 LTS ML и выше. При необходимости установите MLflow со следующим кодом:

%pip install mlflow>=2.13.0 -qqqU
%restart_python

Кроме того, можно установить последнюю версию databricks-agents, которая включает совместимую версию MLflow:

%pip install databricks-agents

Использование автологгинга, чтобы добавить трассировки вашим агентам

Если библиотека GenAI поддерживает трассировку, например LangChain или OpenAI, включите автологирование, добавив mlflow.<library>.autolog() в код. Например:

mlflow.langchain.autolog()

Примечание.

По состоянию на Databricks Runtime 15.4 LTS ML трассировка MLflow включена по умолчанию в записных книжках. Чтобы отключить трассировку, например, с LangChain, можно выполнить mlflow.langchain.autolog(log_traces=False) в ноутбуке.

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

Используйте Fluent API для ручного добавления трассировок в агент

API Fluent в MLflow автоматически создают иерархии трассировки на основе потока выполнения кода.

Декорируйте функцию

Используйте декоратор @mlflow.trace для создания диапазона для области украшенной функции.

Объект MLflow Span упорядочивает шаги трассировки. Спаны собирают информацию об отдельных операциях или шагах в рабочем процессе, а именно, таких как вызовы API или запросы векторного хранилища.

Диапазон начинается при вызове функции и заканчивается при возвращении. MLflow записывает входные и выходные данные функции и все исключения, вызванные функцией.

Например, следующий код создает диапазон с именем my_function, который записывает входные аргументы x и y и выходные данные.

@mlflow.trace(name="agent", span_type="TYPE", attributes={"key": "value"})
def my_function(x, y):
    return x + y

Используйте диспетчер контекстов трассировки

Если вы хотите создать диапазон для произвольного блока кода, а не только функции, можно использовать mlflow.start_span() в качестве диспетчера контекстов, который упаковывает блок кода. Диапазон начинается при вводе контекста и заканчивается при выходе контекста. Входные и выходные параметры диапазона должны быть предоставлены вручную с помощью методов установки объекта диапазона, предоставленного контекстным менеджером.

with mlflow.start_span("my_span") as span:
    span.set_inputs({"x": x, "y": y})
    result = x + y
    span.set_outputs(result)
    span.set_attribute("key", "value")

Обернуть внешнюю функцию

Чтобы отслеживать функции внешней библиотеки, оберните функцию с помощью mlflow.trace.

from sklearn.metrics import accuracy_score

y_pred = [0, 2, 1, 3]
y_true = [0, 1, 2, 3]

traced_accuracy_score = mlflow.trace(accuracy_score)
traced_accuracy_score(y_true, y_pred)

### Fluent API example

The following example shows how to use the Fluent APIs `mlflow.trace` and `mlflow.start_span` to trace the `quickstart-agent`:

```python
import mlflow
from mlflow.deployments import get_deploy_client

class QAChain(mlflow.pyfunc.PythonModel):
    def __init__(self):
        self.client = get_deploy_client("databricks")

    @mlflow.trace(name="quickstart-agent")
    def predict(self, model_input, system_prompt, params):
        messages = [
                {
                    "role": "system",
                    "content": system_prompt,
                },
                {
                    "role": "user",
                    "content":  model_input[0]["query"]
                }
          ]

        traced_predict = mlflow.trace(self.client.predict)
        output = traced_predict(
            endpoint=params["model_name"],
            inputs={
                "temperature": params["temperature"],
                "max_tokens": params["max_tokens"],
                "messages": messages,
            },
        )

        with mlflow.start_span(name="_final_answer") as span:
          # Initiate another span generation
            span.set_inputs({"query": model_input[0]["query"]})

            answer = output["choices"][0]["message"]["content"]

            span.set_outputs({"generated_text": answer})
            # Attributes computed at runtime can be set using the set_attributes() method.
            span.set_attributes({
              "model_name": params["model_name"],
                        "prompt_tokens": output["usage"]["prompt_tokens"],
                        "completion_tokens": output["usage"]["completion_tokens"],
                        "total_tokens": output["usage"]["total_tokens"]
                    })
              return answer

После добавления трассировки запустите функцию. Следующий текст продолжает пример с функцией predict() из предыдущего раздела. Трассировки автоматически отображаются при запуске метода вызова predict().

SYSTEM_PROMPT = """
You are an assistant for Databricks users. You answer Python, coding, SQL, data engineering, spark, data science, DW and platform, API, or infrastructure administration questions related to Databricks. If the question is unrelated to one of these topics, kindly decline to answer. If you don't know the answer, say that you don't know; don't try to make up an answer. Keep the answer as concise as possible. Use the following pieces of context to answer the question at the end:
"""

model = QAChain()

prediction = model.predict(
  [
      {"query": "What is in MLflow 5.0"},
  ],
  SYSTEM_PROMPT,
  {
    # Using Databricks Foundation Model for easier testing, feel free to replace it.
    "model_name": "databricks-dbrx-instruct",
    "temperature": 0.1,
    "max_tokens": 1000,
  }
)

API клиента MLflow

MlflowClient предоставляет детализированные, потокобезопасные API для запуска и завершения трассировки, управления диапазонами и задания полей диапазона. Он обеспечивает полный контроль над жизненным циклом и структурой трассировки. Эти API полезны, если api Fluent недостаточно для ваших требований, например многопоточных приложений и обратных вызовов.

Ниже приведены шаги по созданию полной трассировки с помощью клиента MLflow.

  1. Создайте экземпляр MLflowClient с помощью client = MlflowClient().

  2. Запустите трассировку с помощью метода client.start_trace(). Это инициирует контекст трассировки, запускает абсолютный корневой диапазон и возвращает объект корневого диапазона. Этот метод должен выполняться перед API start_span().

    1. Задайте атрибуты, входные данные и выходные данные трассировки в client.start_trace().

    Примечание.

    В API Fluent нет эквивалента методу start_trace(). Это связано с тем, что API Fluent автоматически инициализирует контекст трассировки и определяет, является ли он корневым спаном на основе управляемого состояния.

  3. API start_trace() возвращает диапазон. Получите идентификатор запроса, уникальный идентификатор трассировки, также называемый trace_id, и идентификатор возвращаемого диапазона с помощью span.request_id и span.span_id.

  4. Запустите дочерний диапазон с помощью client.start_span(request_id, parent_id=span_id), чтобы задать атрибуты, входные данные и выходные данные для диапазона.

    1. Для этого метода требуются request_id и parent_id, чтобы связать отрезок с правильной позицией в иерархии трассировки. Он возвращает другой объект диапазона.

  5. Завершение дочернего диапазона путем вызова client.end_span(request_id, span_id).

  6. Повторите шаги 3 – 5 для всех дочерних диапазонов, которые вы хотите создать.

  7. После завершения всех дочерних диапазонов вызовите client.end_trace(request_id), чтобы закрыть трассировку и записать ее.

from mlflow.client import MlflowClient

mlflow_client = MlflowClient()

root_span = mlflow_client.start_trace(
  name="simple-rag-agent",
  inputs={
          "query": "Demo",
          "model_name": "DBRX",
          "temperature": 0,
          "max_tokens": 200
         }
  )

request_id = root_span.request_id

# Retrieve documents that are similar to the query
similarity_search_input = dict(query_text="demo", num_results=3)

span_ss = mlflow_client.start_span(
      "search",
      # Specify request_id and parent_id to create the span at the right position in the trace
        request_id=request_id,
        parent_id=root_span.span_id,
        inputs=similarity_search_input
  )
retrieved = ["Test Result"]

# You must explicitly end the span
mlflow_client.end_span(request_id, span_id=span_ss.span_id, outputs=retrieved)

root_span.end_trace(request_id, outputs={"output": retrieved})

Просмотр трассировок

Чтобы просмотреть следы после запуска агента, используйте один из следующих вариантов:

  • Визуализация трассировки отображается встроенным в выходных данных ячейки.
  • Трассировки записываются в эксперимент MLflow. Вы можете просмотреть и выполнить поиск по полному списку исторических следов на вкладке Трейсы на странице Эксперимент. Когда агент выполняется в активном сеансе MLflow, трассы отображаются на странице Run.
  • Программное извлечение трассировок с помощью API search_traces().

Использование трассировки MLflow в рабочей среде

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

Чтобы включить трассировку MLflow для конечной точки обслуживания, необходимо установить переменную окружения ENABLE_MLFLOW_TRACING в конфигурации конечной точки на True. Сведения о развертывании конечной точки с настраиваемыми переменными среды см. в статье Добавление переменных среды обычного текста. Если вы развернули агент с помощью API deploy(), трассировки автоматически записываются в таблицу вывода. См. раздел "Развертывание агента для создания приложения ИИ".

Примечание.

Запись трассировок в таблицу вывода выполняется асинхронно, поэтому она не добавляет те же издержки, что и в среде записной книжки во время разработки. Однако это может по-прежнему привести к некоторой нагрузке на скорость отклика конечной точки, особенно если размер трассировки для каждого запроса вывода велик. Databricks не гарантирует никаких соглашений об уровне обслуживания (SLA) для фактической задержки в конечной точке модели, так как она сильно зависит от среды и реализации модели. Databricks рекомендует протестировать производительность конечной точки и получить аналитические сведения о затратах трассировки перед развертыванием в рабочем приложении.

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

Размер трассировки на запрос Влияние на задержку (мс)
~10 КБ ~ 1 мс
~ 1 МБ 50 ~ 100 мс
10 МБ 150 мс ~

Ограничения

  • Трассировка MLflow доступна в записных книжках Databricks, заданиях записных книжек и службе моделей.

Автологирование LangChain может не поддерживать все API прогнозирования LangChain. Для получения полного списка поддерживаемых API см. документацию по MLflow .