Tracer les appels Gemini

Le service RH de Valoria signale qu'une collaboratrice a attendu « une demi-minute » une réponse sur son matériel. Les journaux de l'application disent seulement 200 OK. Il manque la décomposition : combien de temps dans le chargement du prompt, combien dans Gemini, combien d'appels refusés avant le bon, combien de tokens. C'est ce que donne une trace.

Le mécanisme

Une trace est l'arbre des opérations d'une requête ; chaque nœud est un span avec ses entrées, sorties, durée, statut et attributs. MLflow en produit de deux façons, combinées ici :

• mlflow.gemini.autolog() instrumente le SDK google-genai : chaque generate_content devient un span avec le modèle, les tokens et un coût estimé.
• @mlflow.trace sur nos fonctions crée les spans parents, ce qui rattache les appels Gemini à la question qui les a provoqués.

flowchart TB
  Q["question_rh (racine)<br/>entrée : la question seule"] --> A["assistant_rh (CHAIN)<br/>tags : prompt_uri, prompt_version"]
  A --> G1["Models.generate_content<br/>429, 70 ms"]
  A --> G2["Models.generate_content<br/>429, 65 ms"]
  A --> G3["… 2 autres 429"]
  A --> G5["Models.generate_content<br/>OK, 2,82 s, 746 tokens"]

Les traces sont exportées vers le serveur en arrière-plan : l'appel de l'application ne les attend pas. Un test qui lit les traces juste après un appel doit donc demander la vidange de la file (flush=True).

La mise en œuvre

L'instrumentation tient en une ligne à la connexion :

def connect() -> None:
    mlflow.set_tracking_uri(os.environ.get("MLFLOW_TRACKING_URI", "http://localhost:5050"))
    mlflow.set_experiment(EXPERIMENT)
    mlflow.gemini.autolog()  # chaque appel google-genai devient une trace

Le client Gemini est créé une seule fois et gardé vivant, avec un délai maximal par appel. La reprise sur erreur reste dans notre code, si bien que chaque tentative apparaît comme un span :

@lru_cache(maxsize=1)
def client():
    from google import genai
    from google.genai import types
    # Sans délai maximal, une connexion bloquée gèle mlflow.genai.evaluate indéfiniment.
    return genai.Client(http_options=types.HttpOptions(timeout=TIMEOUT_MS))

def generate(prompt: str, llm=None, retries: int = 4) -> str:
    ...
        except (errors.APIError, httpx.TimeoutException) as e:
            code = getattr(e, "code", "timeout")
            if code not in (429, 500, 503, "timeout") or attempt == retries:
                raise
            wait = 2 ** (attempt + 1)
            log.warning("Gemini %s, nouvel essai dans %ss", code, wait)
            time.sleep(wait)

genai.Client() sans argument lit GEMINI_API_KEY, ou passe par Vertex AI si GOOGLE_GENAI_USE_VERTEXAI=TRUE (avec GOOGLE_CLOUD_PROJECT et GOOGLE_CLOUD_LOCATION) : le même code sert les deux modes.

Le test vérifie qu'une réponse est tracée et étiquetée, sans réseau :

out = answer("Combien de jours de télétravail ?", guide_text="§1 Télétravail : 2 jours.", llm=llm)
traces = mlflow.search_traces(return_type="list", flush=True)  # export asynchrone
assert traces and traces[0].info.tags["prompt_version"] == "1"

Résultat obtenu (2026-10-16)

La liste des traces après un make e2e : une ligne par question, entrée et sortie lisibles.

Le détail d'une question évaluée pendant une rafale de 429 sur le quota partagé : quatre appels refusés en quelques dizaines de millisecondes, puis la réponse en 2,82 s pour 746 tokens. La trace dure 33,3 s, presque entièrement en attentes de reprise (2 + 4 + 8 + 16 s).

Sur une autre question, les cinq tentatives ont échoué :

question_rh                SpanStatusCode.ERROR  ClientError: 429 RESOURCE_EXHAUSTED  30.6 s
assistant_rh               SpanStatusCode.ERROR  ClientError: 429 RESOURCE_EXHAUSTED  30.5 s
Models.generate_content    SpanStatusCode.ERROR  ClientError: 429 RESOURCE_EXHAUSTED   0.1 s   (x5)

mlflow.genai.evaluate a alors rejoué la ligne : le run d'évaluation de v2 compte 12 traces pour 10 questions, dont 2 en erreur.

Ce qui casse en production

• Données personnelles dans les traces. Les questions des collaborateurs (« je pars en congé maternité… ») sont stockées en clair, entrées et sorties comprises. Parade : durée de rétention définie, accès restreint au serveur, et masquage des champs sensibles avant export si le cas d'usage l'exige ; le RGPD s'applique aux traces comme au reste.
• Reprises invisibles. Une reprise faite par le client HTTP lui-même (HttpRetryOptions) n'apparaît que comme un appel long. Parade : reprise dans votre code quand vous voulez la voir ; sinon, comparez la durée du span à la latence normale.
• Traces perdues à l'arrêt. Un script court se termine avant l'export asynchrone. Symptôme : traces manquantes pour les derniers appels. Parade : MLflow vide la file à la sortie du processus (« Flushing the async trace logging queue before program exit ») ; pour un arrêt brutal (SIGKILL), rien ne la sauve.
• Volume. Une trace par appel, c'est 400 000 traces par mois pour 400 questions par jour et un facteur 30. Parade : échantillonnage en production (MLFLOW_TRACE_SAMPLING_RATIO), traces complètes en évaluation.

La règle de décision

Tracez tout en développement et en évaluation. En production, échantillonnez et gardez toutes les traces en erreur. Mettez la version du prompt et la version de l'application en tags de chaque trace : c'est la clé pour relier une plainte à une configuration.

Pratique

Exercice. Ajoutez un tag app_version (lu dans la variable d'environnement APP_VERSION, défaut dev) à chaque trace produite par answer(), et un test.

mlflow.update_current_trace(tags={"prompt_uri": prompt_uri, "prompt_version": str(prompt.version),
                                  "app_version": os.environ.get("APP_VERSION", "dev")})

def test_trace_has_app_version(tracking, monkeypatch):
    monkeypatch.setenv("APP_VERSION", "1.4.2")
    lab.register()
    answer("Q ?", guide_text="§1", llm=FakeLLM())
    assert mlflow.search_traces(return_type="list", flush=True)[0].info.tags["app_version"] == "1.4.2"

• J'active l'autolog Gemini et je trace mes propres fonctions pour obtenir un arbre lisible.
• Je garde la reprise sur erreur dans mon code pour qu'elle apparaisse dans les traces.
• J'étiquette chaque trace avec la version du prompt.
• Je sais que les traces contiennent des données personnelles et je fixe leur rétention.