MLFlow – wprowadzenie

MLFlow to popularna biblioteka służąca do śledzenia wyników pracy funkcji. Oryginalnie została zaplanowana do śledzenia działania funkcji Machine Learning, stąd jej nazwa, ale z czasem została zaadoptowana do AI i generalnie może służyć do śledzenia czegoklowiek. MLFlow jest rekomendowany do używania w Databricks.

MLFlow może być zainstalowany lokalnie. Tutaj tworzymy środowisko wirtualne, aktywujemy je, instalujemy mlflow i przy okazji openai. Potem startujemy serwer mlflow:

python -m venv venv
venv/scripst/activate.ps1
pip install --upgrade mlflow openai>=1.0.0
mlflow server

Serwer po uruchomieniu nasłuchuje domyślnie na lokalnym komputerze, o czym informuje output tego polecenia:

INFO: Uvicorn running on http://127.0.0.1:5000 (Press CTRL+C to quit)

Od tej pory do serwera można się podłączyć z aplikacji i zacząć coś sobie logować. Zazwyczaj MLFlow służy do zbierania trace ze skomplikowanych aplikacji np. korzystających z OpenAI, stąd i ten przykład do tego nawiązuje:

import mlflow
from openai import OpenAI
import time

# Specify the tracking URI for the MLflow server.
mlflow.set_tracking_uri("http://localhost:5000")

# Specify the experiment you just created for your LLM application or AI agent.
mlflow.set_experiment("DEMO-01")

@mlflow.trace
def transform_parameter(input_param: str) -> str:
  time.sleep(2)
  return input_param.upper()

@mlflow.trace
def chat_completion(message: list[dict], user_id: str, session_id: str):
    """Process a chat message with user and session tracking."""

    # Add user and session context to the current trace
    mlflow.update_current_trace(
        metadata={
            "mlflow.trace.user": user_id,  # Links trace to specific user
            "mlflow.trace.session": session_id,  # Groups trace with conversation
            "ui_theme_preference": "dark",
            "country": "poland",
        }
    )

    # Your chat logic here
    ret = transform_parameter(message[0]['input'])
    return ret
    # return generate_response(message)

chat_completion(
  message = [{'input':'hello'},{'output':'what do you want?'}],
  user_id = 'wloczykij',
  session_id = '013'
)

Mamy tutaj przykład niby-to-chata z OpenAI, bo tego OpenAI na tym etapie usunąłem, bo przecież nie każdy ma API key pozwalający wykonać zapytania, do których są potrzebne tokeny. Najważniejsze rzeczy do zaobserwowania:

• Importujemy mlflow (a oprócz tego też inne potrzebne moduły)
• Przy pomocy mlflow.set_tracking_uri łączymy się do uruchomionego serwera MLFlow
• Na tym serwerze łączymy się do eksperymentu „DEMO-01”. Gdyby go nie było, to zostanie on automatycznie utworzony. Eksperyment to najwyższy poziom w MLFlow, pod nim w strukturze znajduje się run a poniżej trace.
• Dalej możemy mieć dziesiątki funkcji. Jeśli te funkcje mają jakieś specjalne znaczenie dla twojego programu, to należy je zonaczyć dekoratorem @mlflow.trace. Dzieki temu w logach zbieranych przez MLFlow moment uruchomienia tej funkcji, jej parametry i zwracany wynik będą wyraźnie widoczne
• Przydatne może być wywołanie mlflow.update_current_trace(…), które pozwala dodać pewne metadane do zbieranego trace. Część z tych metadanych jest natywnie obsługiwana przez MLFlow, jak np widoczne tu mlflow.trace.user czy mlflow.trace.session. Jeśli jednak chcesz dodać własne metadane opisujące sposób w jaki ta funkcja została uruchomiona, jej środowisko itp, to możesz to zrobić. Tutaj dodaliśmy tak: ui_theme_preference czy country.

Ponadto ten kod coś robi. W naszym przypadku – same głupie rzeczy, bo zmienia teksty z małych liter na wielkie litery, ale śmiało w tym miejscu można by właśnie wywoływać API OpenAI, uruchamiać uczenie maszynowe, czy pobierać i przetwarzać dane. Cokolwiek!

Kod uruchomisz mniej więcej tak:

$env:OPENAI_API_KEY="sk-proj-Bx***"
python ./demo-01.py

Definiowanie klucza w postaci zmiennej środowiskowej uznajmy na tym etapie za opcjonalne, ale przyda się, jeśli doczytasz ten wpis do końca.

Po kilku uruchomieniach można zajrzeć do tego co zapisało się pod adresem, na którym nasłuchuje MLFLOW, a tutaj powinno się udać zobaczyć listę eksperymentów, a wśród nich ten jeden nasz:

Jeśli klikniemy na eksperymencie, to zobaczymy overview, który pokazuje podstawowe informacje o eksperymencie:

Przechodząc do Traces zobaczymy wszystkie zarejestrowane uruchomienia, które można filtrować między innymi w oparciu o wcześniej zarejestrowane metadane

A wybierając pojedyncze uruchomienie zobaczymy szczegóły trace. W szczególności, jeśli jedna funkcja wywoływała drugą to zobaczymy informacje zarówno o tej funkcji zewnętrznej jak i o tej wewnętrznej

Oprócz samodzielnego oznaczania funkcji przez @mlflow.trace można też zdać się na automat i w kodzie, po podpięciu się do ekspertymentu wywołać mlflow.openai.autolog(). O tak:

mlflow.set_experiment(„DEMO-03”)
mlflow.openai.autolog()

Wtedy logowanie może być zbierane automatycznie. Nie musisz oznaczać funkcji przez @mlflow.trace a działanie i tak zostanie zalogowane, ale… jest haczyk. W takim przypadku kod musi wywoływać np. funkcje OpenAi. Ten Automat nie zadziała dla naszych własnych funkcji. Jeśli więc np. masz u siebie jeszcze coś w tym stylu:

client = OpenAI()
# The trace of the following is sent to the MLflow server.
client.chat.completions.create(
    model="o4-mini",
    messages=[
        {"role": "system", "content": "You are a helpful weather assistant."},
        {"role": "user", "content": "What's the weather like in Seattle?"},
    ],
)

To MLFLow to złapie i zaloguje – nawet jeśli nie masz OpenAI API Key, lub jest on nieważny lub brakło tokenów – wtedy w logach znajdzie się po prostu zgłoszony bład – i super, bo przecież właśnie o to chodziło żeby takie sytuacje wychwycić

Komentarze są wyłączone