Protokollieren und Registrieren von KI-Agents

Artikel
10/30/2024

Wichtig

Dieses Feature befindet sich in der Public Preview.

Protokollieren von KI-Agenten mithilfe von Mosaik AI Agent Framework. Das Protokollieren eines Agents ist die Grundlage des Entwicklungsprozesses. Die Protokollierung erfasst einen "Zeitpunkt" des Codes und der Konfiguration des Agents, damit Sie die Qualität der Konfiguration auswerten können.

Anforderungen

Erstellen Sie einen KI-Agent, bevor Sie ihn protokollieren.

Codebasierte Protokollierung im Vergleich zur serialisierungsbasierten Protokollierung

Sie können die codebasierte MLflow-Protokollierung oder die serialisierungsbasierte MLflow-Protokollierung verwenden. Databricks empfiehlt, die codebasierte Protokollierung zu verwenden.

Codebasierte MLflow-Protokollierung: Der Code der Kette wird als Python-Datei erfasst. Die Python-Umgebung wird als Liste der Pakete erfasst. Wenn die Kette bereitgestellt wird, wird die Python-Umgebung wiederhergestellt, und der Code der Kette wird ausgeführt, um die Kette in den Arbeitsspeicher zu laden, damit sie aufgerufen werden kann, wenn der Endpunkt aufgerufen wird.

Serialisierungsbasierte MLflow-Protokollierung: Der Code und der aktuelle Zustand der Kette in der Python-Umgebung wird auf den Datenträger serialisiert, häufig mithilfe von Bibliotheken wie pickle oder joblib. Wenn die Kette bereitgestellt wird, wird die Python-Umgebung wiederhergestellt, und das serialisierte Objekt wird in den Arbeitsspeicher geladen, sodass es aufgerufen werden kann, wenn der Endpunkt aufgerufen wird.

In der Tabelle sind die Vor- und Nachteile der einzelnen Methoden aufgeführt.

Methode	Vorteile	Nachteile
Codebasierte MLflow-Protokollierung	– Überwindet inhärente Einschränkungen der Serialisierung, die von vielen beliebten GenAI-Bibliotheken nicht unterstützt wird. – Speichert eine Kopie des ursprünglichen Codes für späteren Verweis. – Sie müssen Ihren Code nicht in ein einzelnes Objekt umstrukturieren, das serialisiert werden kann.	`log_model(...)` muss aus einem anderen Notizbuch als dem Code der Kette (als Treibernotizbuch bezeichnet) aufgerufen werden.
Serialisierungsbasierte MLflow-Protokollierung	`log_model(...)` kann aus demselben Notizbuch aufgerufen werden, in dem das Modell definiert ist.	– Originalcode ist nicht verfügbar. – Alle Bibliotheken und Objekte, die in der Kette verwendet werden, müssen die Serialisierung unterstützen.

Für die codebasierte Protokollierung muss sich der Code, der Ihren Agent oder Ihre Kette protokolliert, in einem separaten Notizbuch von Ihrem Kettencode befinden. Dieses Notizbuch wird als Treibernotizbuch bezeichnet. Ein Beispielnotizbuch finden Sie unter Beispielnotizbücher.

Codebasierte Protokollierung mit LangChain

Erstellen Sie ein Notizbuch oder eine Python-Datei mit Ihrem Code. Für dieses Beispiel wird das Notizbuch oder die Datei chain.py genannt. Das Notizbuch oder die Datei muss eine LangChain-Kette enthalten, die hier als lc_chain bezeichnet wird.
Fügen Sie mlflow.models.set_model(lc_chain) zum Notizbuch oder zur Datei hinzu.
Erstellen Sie ein neues Notizbuch, das als Treibernotizbuch dienen soll (in diesem Beispiel driver.py genannt).
Verwenden Sie im Treibernotebook mlflow.lang_chain.log_model(lc_model=”/path/to/chain.py”), um chain.py auszuführen, und protokollieren Sie die Ergebnisse in einem MLflow-Modell.
Bereitstellen des Modells. Siehe Bereitstellen eines Agents für generative KI-Anwendung. Die Bereitstellung Ihres Agents hängt möglicherweise von anderen Databricks-Ressourcen ab, z. B. von einem Vektorsuchindex und einem Modellbereitstellungsendpunkten. Für LangChain-Agents:
- MLflow log_model leitet die von der Kette benötigten Abhängigkeiten ab und protokolliert sie in der MLmodel-Datei im protokollierten Modellartefakt. Ab Mlflow Version 2.17.0 können Sie diese abgeleiteten Abhängigkeiten überschreiben. Siehe Angeben von Ressourcen für PyFunc- oder LangChain-Agent.
- Während der Bereitstellung erstellt databricks.agents.deploy automatisch die M2M OAuth-Token, die erforderlich sind, um auf diese abgeleiteten Ressourcenabhängigkeiten zuzugreifen und mit ihnen zu kommunizieren.
Wenn die Bereitstellungsumgebung geladen wird, wird chain.py ausgeführt.
Wenn eine Dienstanforderung erhalten wird, wird lc_chain.invoke(...) aufgerufen.


import mlflow

code_path = "/Workspace/Users/first.last/chain.py"
config_path = "/Workspace/Users/first.last/config.yml"

input_example = {
    "messages": [
        {
            "role": "user",
            "content": "What is Retrieval-augmented Generation?",
        }
    ]
}

# example using LangChain
with mlflow.start_run():
  logged_chain_info = mlflow.langchain.log_model(
    lc_model=code_path,
    model_config=config_path, # If you specify this parameter, this is the configuration that is used for training the model. The development_config is overwritten.
    artifact_path="chain", # This string is used as the path inside the MLflow model where artifacts are stored
    input_example=input_example, # Must be a valid input to your chain
    example_no_conversion=True, # Required
  )

print(f"MLflow Run: {logged_chain_info.run_id}")
print(f"Model URI: {logged_chain_info.model_uri}")

# To verify that the model has been logged correctly, load the chain and call `invoke`:
model = mlflow.langchain.load_model(logged_chain_info.model_uri)
model.invoke(example)

Codebasierte Protokollierung mit PyFunc

Erstellen Sie ein Notizbuch oder eine Python-Datei mit Ihrem Code. Für dieses Beispiel wird das Notizbuch oder die Datei chain.py genannt. Das Notizbuch oder die Datei muss eine PyFunc-Klasse enthalten, die hier als PyFuncClass bezeichnet wird.
Fügen Sie mlflow.models.set_model(PyFuncClass) zum Notizbuch oder zur Datei hinzu.
Erstellen Sie ein neues Notizbuch, das als Treibernotizbuch dienen soll (in diesem Beispiel driver.py genannt).
Verwenden Sie im Treibernotebook mlflow.pyfunc.log_model(python_model=”/path/to/chain.py”, resources=”/path/to/resources.yaml”), um chain.py auszuführen, und protokollieren Sie die Ergebnisse in einem MLflow-Modell. Der resources Parameter deklariert alle Ressourcen, die für das Modell erforderlich sind, z. B. einen Vektorsuchindex oder einen Endpunkt, der ein Foundation-Modell dient. Eine Beispielressourcendatei für PyFunc finden Sie unter "Angeben von Ressourcen für PyFunc" oder "LangChain".For an example resources file for PyFunc, see Specify resources for PyFunc or LangChain agent.
Bereitstellen des Modells. Siehe Bereitstellen eines Agents für generative KI-Anwendung.
Wenn die Bereitstellungsumgebung geladen wird, wird chain.py ausgeführt.
Wenn eine Dienstanforderung erhalten wird, wird PyFuncClass.predict(...) aufgerufen.

import mlflow

code_path = "/Workspace/Users/first.last/chain.py"
config_path = "/Workspace/Users/first.last/config.yml"

input_example = {
    "messages": [
        {
            "role": "user",
            "content": "What is Retrieval-augmented Generation?",
        }
    ]
}

# example using PyFunc model

resources_path = "/Workspace/Users/first.last/resources.yml"

with mlflow.start_run():
  logged_chain_info = mlflow.pyfunc.log_model(
    python_model=chain_notebook_path,
    artifact_path="chain",
    input_example=input_example,
    resources=resources_path,
    example_no_conversion=True,
  )

print(f"MLflow Run: {logged_chain_info.run_id}")
print(f"Model URI: {logged_chain_info.model_uri}")

# To verify that the model has been logged correctly, load the chain and call `invoke`:
model = mlflow.pyfunc.load_model(logged_chain_info.model_uri)
model.invoke(example)

Angeben von Ressourcen für PyFunc- oder LangChain-Agent

Sie können Ressourcen angeben, die für das Modell erforderlich sind, z. B. einen Vektorsuchindex und einen Bereitstellungsendpunkt.

Bei LangChain werden Ressourcen mithilfe eines Best-Effort-Ansatzes automatisch erkannt und mit dem Modell protokolliert. Ab MLflow Version 2.17.0 können Sie diese automatisch abgeleiteten Ressourcen mithilfe von Code außer Kraft setzen, der dem unten gezeigten ähnelt. Dies wird für Produktionsanwendungsfälle empfohlen, da Sie sicherstellen können, dass Agents mit den erforderlichen Abhängigkeiten protokolliert werden.

Beim Bereitstellen eines pyfunc-Agents müssen Sie alle Ressourcenabhängigkeiten des bereitgestellten Agents manuell hinzufügen. Ein M2M OAuth-Token mit Zugriff auf alle angegebenen Ressourcen im resources-Parameter wird für den bereitgestellten Agent erstellt.

Hinweis

Sie können die Ressourcen außer Kraft setzen, für die Ihr Endpunkt über Berechtigungen verfügt, indem Sie die Ressourcen beim Protokollieren der Kette manuell angeben.

Der folgende Code gibt Abhängigkeiten mithilfe des resources Parameters an.

import mlflow
from mlflow.models.resources import (
    DatabricksFunction,
    DatabricksServingEndpoint,
    DatabricksSQLWarehouse,
    DatabricksVectorSearchIndex,
)

with mlflow.start_run():
  logged_chain_info = mlflow.pyfunc.log_model(
    python_model=chain_notebook_path,
    artifact_path="chain",
    input_example=input_example,
    example_no_conversion=True,
    resources=[
      DatabricksServingEndpoint(endpoint_name="databricks-mixtral-8x7b-instruct"),
      DatabricksServingEndpoint(endpoint_name="databricks-bge-large-en"),
      DatabricksVectorSearchIndex(index_name="prod.agents.databricks_docs_index"),
      DatabricksSQLWarehouse(warehouse_id="your_warehouse_id"),
      DatabricksFunction(function_name="ml.tools.python_exec"),
    ]
  )

Sie können auch Ressourcen hinzufügen, indem Sie sie in einer resources.yaml-Datei angeben. Sie können im resources-Parameter auf diesen Dateipfad verweisen. Ein M2M OAuth-Token mit Zugriff auf alle angegebenen Ressourcen in resources.yaml wird für den bereitgestellten Agent erstellt.

Im Folgenden finden Sie eine Beispieldatei resources.yaml, die Modellbereitstellungsendpunkte und einen Vektorsuchindex definiert.


api_version: "1"
databricks:
  vector_search_index:
    - name: "catalog.schema.my_vs_index"
  serving_endpoint:
    - name: databricks-dbrx-instruct
    - name: databricks-bge-large-en

Registrieren der Kette im Unity-Katalog

Bevor Sie die Kette bereitstellen, müssen Sie die Kette im Unity-Katalog registrieren. Wenn Sie die Kette registrieren, wird sie als Modell im Unity-Katalog verpackt, und Sie können Unity-Katalogberechtigungen für die Autorisierung für Ressourcen in der Kette verwenden.

import mlflow

mlflow.set_registry_uri("databricks-uc")

catalog_name = "test_catalog"
schema_name = "schema"
model_name = "chain_name"

model_name = catalog_name + "." + schema_name + "." + model_name
uc_model_info = mlflow.register_model(model_uri=logged_chain_info.model_uri, name=model_name)

Freigeben über