From 6e503894eb7bcb31a598dfc4edf4ac92fbdbb1fb Mon Sep 17 00:00:00 2001 From: Yoav Katz Date: Thu, 16 Jul 2026 13:51:44 +0300 Subject: [PATCH 1/3] feat: add --experiment flag and delete-all-deployments script MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add --experiment NAME flag to deploy-agent.sh, deploy-benchmark.sh, deploy-and-evaluate.sh, and evaluate-benchmark.sh so parallel experiments get distinct pod/service names (e.g. exgentic-a2a-tool-calling-gsm8k-exp1 vs …-exp2). Reorder steps in deploy-agent.sh and deploy-benchmark.sh so resource limits and rollout stabilization happen before the health-check wait, avoiding a race where the pod restarts mid-check. Add delete-all-deployments.sh to bulk-delete all agent and benchmark deployments in a namespace via the Kagenti API. Co-Authored-By: Claude Sonnet 4.6 (1M context) Signed-off-by: Yoav Katz --- exgentic_a2a_runner/delete-all-deployments.sh | 283 ++++++++++++++++++ exgentic_a2a_runner/deploy-agent.sh | 89 +++--- exgentic_a2a_runner/deploy-and-evaluate.sh | 4 + exgentic_a2a_runner/deploy-benchmark.sh | 51 ++-- exgentic_a2a_runner/evaluate-benchmark.sh | 11 +- 5 files changed, 381 insertions(+), 57 deletions(-) create mode 100755 exgentic_a2a_runner/delete-all-deployments.sh diff --git a/exgentic_a2a_runner/delete-all-deployments.sh b/exgentic_a2a_runner/delete-all-deployments.sh new file mode 100755 index 0000000..3f32e93 --- /dev/null +++ b/exgentic_a2a_runner/delete-all-deployments.sh @@ -0,0 +1,283 @@ +#!/bin/bash +# Delete all agent and benchmark (tool) deployments from a Kagenti namespace via the API. +# Usage: ./delete-all-deployments.sh [OPTIONS] +# Example: ./delete-all-deployments.sh --openshift apps.mycluster.example.com +# Example: ./delete-all-deployments.sh --kind +# Example: ./delete-all-deployments.sh --dry-run + +set -e + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +ENV_FILE="$SCRIPT_DIR/.env" +if [ -f "$ENV_FILE" ]; then + while IFS= read -r line || [ -n "$line" ]; do + [[ -z "${line// }" || "$line" =~ ^[[:space:]]*# ]] && continue + line="${line#export }" + [[ "$line" != *=* ]] && continue + key="${line%%=*}" + val="${line#*=}" + if [[ "$val" =~ ^\"(.*)\"$ ]] || [[ "$val" =~ ^\'(.*)\'$ ]]; then + val="${BASH_REMATCH[1]}" + fi + if [ -z "${!key+x}" ]; then + export "$key=$val" + fi + done <"$ENV_FILE" +fi + +KEYCLOAK_USERNAME="admin" +KEYCLOAK_PASSWORD="${KEYCLOAK_PASSWORD:-unknown}" +NAMESPACE="team1" +CLUSTER_MODE="" +INGRESS_DOMAIN="" +DRY_RUN="false" +KUBECTL_BIN="${KUBECTL_BIN:-kubectl}" + +usage() { + cat < Target OpenShift cluster with given ingress domain + --in-cluster Run as a Kubernetes Job pod (uses in-cluster DNS) + +Options: + --namespace Kagenti namespace (default: team1) + --keycloak-user Keycloak username (default: admin) + --keycloak-pass Keycloak password (overrides KEYCLOAK_PASSWORD env) + --dry-run List what would be deleted without deleting anything + -h, --help Show this help message +EOF + exit 0 +} + +while [[ $# -gt 0 ]]; do + case $1 in + --kind) + CLUSTER_MODE="kind" + shift + ;; + --openshift) + CLUSTER_MODE="openshift" + INGRESS_DOMAIN="$2" + shift 2 + ;; + --in-cluster) + CLUSTER_MODE="in-cluster" + shift + ;; + --namespace) + NAMESPACE="$2" + shift 2 + ;; + --keycloak-user) + KEYCLOAK_USERNAME="$2" + shift 2 + ;; + --keycloak-pass) + KEYCLOAK_PASSWORD="$2" + shift 2 + ;; + --dry-run) + DRY_RUN="true" + shift + ;; + -h|--help) + usage + ;; + *) + echo "Error: Unknown argument: $1" + usage + ;; + esac +done + +export CLUSTER_MODE INGRESS_DOMAIN +# shellcheck source=libsh/urls.sh +source "$SCRIPT_DIR/libsh/urls.sh" +# shellcheck source=libsh/check-kubectl-context.sh +source "$SCRIPT_DIR/libsh/check-kubectl-context.sh" +check_kubectl_context + +KAGENTI_API="$(kagenti_api_url)" +KEYCLOAK_API="$(keycloak_api_url)" + +echo "==========================================" +echo " Delete All Deployments" +echo "==========================================" +echo " Namespace: $NAMESPACE" +echo " Kagenti API: $KAGENTI_API" +echo " Dry run: $DRY_RUN" +echo "==========================================" +echo "" + +# Step 1: Get Keycloak authentication token +echo "Step 1: Getting Keycloak authentication token..." + +if [ "$KEYCLOAK_PASSWORD" = "unknown" ]; then + echo "Step 1.5: Attempting to fetch Keycloak password from cluster..." + KAGENTI_PASSWORD=$("$KUBECTL_BIN" get secret kagenti-test-user -n keycloak -o jsonpath='{.data.password}' 2>/dev/null | base64 -d 2>/dev/null || echo "") + if [ -n "$KAGENTI_PASSWORD" ]; then + TEST_AUTH=$(curl -s -X POST "$KEYCLOAK_API/realms/kagenti/protocol/openid-connect/token" \ + -H "Content-Type: application/x-www-form-urlencoded" \ + -d "username=$KEYCLOAK_USERNAME" \ + -d "password=$KAGENTI_PASSWORD" \ + -d "grant_type=password" \ + -d "client_id=kagenti" 2>/dev/null || echo "") + if echo "$TEST_AUTH" | grep -q "access_token"; then + KEYCLOAK_PASSWORD="$KAGENTI_PASSWORD" + echo "✓ Fetched Keycloak password from cluster" + else + echo "⚠ Warning: Fetched password from cluster but authentication failed" + exit 1 + fi + else + TEST_AUTH=$(curl -s -X POST "$KEYCLOAK_API/realms/kagenti/protocol/openid-connect/token" \ + -H "Content-Type: application/x-www-form-urlencoded" \ + -d "username=$KEYCLOAK_USERNAME" \ + -d "password=admin" \ + -d "grant_type=password" \ + -d "client_id=kagenti" 2>/dev/null || echo "") + if echo "$TEST_AUTH" | grep -q "access_token"; then + KEYCLOAK_PASSWORD="admin" + echo "✓ Using default Keycloak password" + else + echo "Error: Could not determine Keycloak password. Use --keycloak-pass." + exit 1 + fi + fi +fi + +TOKEN_RESPONSE=$(curl -s -X POST "$KEYCLOAK_API/realms/kagenti/protocol/openid-connect/token" \ + -H "Content-Type: application/x-www-form-urlencoded" \ + -d "username=$KEYCLOAK_USERNAME" \ + -d "password=$KEYCLOAK_PASSWORD" \ + -d "grant_type=password" \ + -d "client_id=kagenti" || echo "TOKEN_ERROR") + +if [ "$TOKEN_RESPONSE" = "TOKEN_ERROR" ]; then + echo "Error: Failed to reach Keycloak at $KEYCLOAK_API" + exit 1 +fi + +ACCESS_TOKEN=$(echo "$TOKEN_RESPONSE" | grep -o '"access_token":"[^"]*"' | sed 's/"access_token":"\([^"]*\)"/\1/') +if [ -z "$ACCESS_TOKEN" ]; then + echo "Error: Failed to extract access token" + echo "Response: $TOKEN_RESPONSE" + exit 1 +fi +echo "✓ Authentication token obtained" +echo "" + +# Step 2: Verify Kagenti backend is accessible +echo "Step 2: Verifying Kagenti backend accessibility at $KAGENTI_API..." +if ! curl -s --max-time 10 "$KAGENTI_API/api/v1/namespaces" >/dev/null 2>&1; then + echo "Error: Kagenti backend is not accessible at $KAGENTI_API" + exit 1 +fi +echo "✓ Kagenti backend is accessible" +echo "" + +# Step 3: Discover deployments via kubectl +# Agents are named exgentic-a2a-*, tools/benchmarks are named exgentic-mcp-* +echo "Step 3: Listing deployments in namespace '$NAMESPACE'..." +ALL_DEPLOYMENTS=$("$KUBECTL_BIN" get deployments -n "$NAMESPACE" \ + --no-headers -o custom-columns=NAME:.metadata.name 2>/dev/null || true) + +AGENT_NAMES=$(echo "$ALL_DEPLOYMENTS" | grep "^exgentic-a2a-" || true) +TOOL_NAMES=$(echo "$ALL_DEPLOYMENTS" | grep "^exgentic-mcp-" || true) + +if [ -z "$AGENT_NAMES" ]; then + echo " No agents found." +else + echo " Found agents:" + echo "$AGENT_NAMES" | while read -r name; do + echo " - $name" + done +fi +echo "" + +echo "Step 4: Listing tools (benchmarks) in namespace '$NAMESPACE'..." +if [ -z "$TOOL_NAMES" ]; then + echo " No tools found." +else + echo " Found tools:" + echo "$TOOL_NAMES" | while read -r name; do + echo " - $name" + done +fi +echo "" + +if [ "$DRY_RUN" = "true" ]; then + echo "Dry-run mode — no deletions performed." + exit 0 +fi + +# wait_for_gone — polls until the record is 404 or 30s elapses. +wait_for_gone() { + local resource_type="$1" name="$2" api_path="$3" + local elapsed=0 max=30 code + while true; do + code=$(curl -s --max-time 5 -o /dev/null -w "%{http_code}" \ + "$KAGENTI_API$api_path" \ + -H "Authorization: Bearer $ACCESS_TOKEN") || code="000" + [ "$code" = "404" ] && echo " ✓ $resource_type '$name' confirmed gone" && return 0 + if [ "$elapsed" -ge "$max" ]; then + echo " ⚠ $resource_type '$name' still present after ${max}s — Kagenti cleanup stalled" + return 1 + fi + sleep 2 + elapsed=$((elapsed + 2)) + done +} + +# Step 5: Delete all agents +if [ -n "$AGENT_NAMES" ]; then + echo "Step 5: Deleting agents..." + echo "$AGENT_NAMES" | while read -r name; do + [ -z "$name" ] && continue + echo -n " Deleting agent '$name'... " + HTTP=$(curl -s --max-time 30 -w "%{http_code}" -o /tmp/kagenti_del_agent.txt \ + -X DELETE "$KAGENTI_API/api/v1/agents/$NAMESPACE/$name" \ + -H "Authorization: Bearer $ACCESS_TOKEN") || HTTP="000" + if [ "$HTTP" = "200" ] || [ "$HTTP" = "404" ]; then + echo "done (HTTP $HTTP)" + [ "$HTTP" = "200" ] && wait_for_gone "agent" "$name" "/api/v1/agents/$NAMESPACE/$name" + else + echo "FAILED (HTTP $HTTP)" + cat /tmp/kagenti_del_agent.txt + fi + done +else + echo "Step 5: No agents to delete." +fi +echo "" + +# Step 6: Delete all tools +if [ -n "$TOOL_NAMES" ]; then + echo "Step 6: Deleting tools (benchmarks)..." + echo "$TOOL_NAMES" | while read -r name; do + [ -z "$name" ] && continue + echo -n " Deleting tool '$name'... " + HTTP=$(curl -s --max-time 30 -w "%{http_code}" -o /tmp/kagenti_del_tool.txt \ + -X DELETE "$KAGENTI_API/api/v1/tools/$NAMESPACE/$name" \ + -H "Authorization: Bearer $ACCESS_TOKEN") || HTTP="000" + if [ "$HTTP" = "200" ] || [ "$HTTP" = "404" ]; then + echo "done (HTTP $HTTP)" + [ "$HTTP" = "200" ] && wait_for_gone "tool" "$name" "/api/v1/tools/$NAMESPACE/$name" + else + echo "FAILED (HTTP $HTTP)" + cat /tmp/kagenti_del_tool.txt + fi + done +else + echo "Step 6: No tools to delete." +fi +echo "" + +echo "==========================================" +echo " Done" +echo "==========================================" diff --git a/exgentic_a2a_runner/deploy-agent.sh b/exgentic_a2a_runner/deploy-agent.sh index 82c2e42..79d7743 100755 --- a/exgentic_a2a_runner/deploy-agent.sh +++ b/exgentic_a2a_runner/deploy-agent.sh @@ -39,6 +39,7 @@ KEYCLOAK_USERNAME="admin" KEYCLOAK_PASSWORD="${KEYCLOAK_PASSWORD:-unknown}" BENCHMARK_NAME="" AGENT_NAME_INPUT="" +EXPERIMENT_NAME="default" USE_MCP_GATEWAY="false" USE_LOCAL_IMAGE="false" CLUSTER_MODE="" @@ -62,6 +63,10 @@ while [[ $# -gt 0 ]]; do AGENT_NAME_INPUT="$2" shift 2 ;; + --experiment) + EXPERIMENT_NAME="$2" + shift 2 + ;; --model) MODEL_NAME="$2" shift 2 @@ -119,6 +124,7 @@ while [[ $# -gt 0 ]]; do echo " --agent NAME Agent name (e.g., tool_calling)" echo "" echo "Optional Arguments:" + echo " --experiment NAME Experiment name suffix appended to pod names (default: default)" echo " --model MODEL Model name (default: Azure/gpt-4.1)" echo " --keycloak-user USER Keycloak username (default: admin)" echo " --keycloak-pass PASS Keycloak password (auto-detected from cluster if not provided)" @@ -184,6 +190,11 @@ fi # Replace underscores with hyphens for Kubernetes compatibility AGENT_NAME="${FULL_AGENT_NAME}-${BENCHMARK_NAME}" AGENT_NAME="${AGENT_NAME//_/-}" +# Append experiment suffix when non-default so parallel experiments get distinct pod names +if [ -n "$EXPERIMENT_NAME" ] && [ "$EXPERIMENT_NAME" != "default" ]; then + EXPERIMENT_SUFFIX="${EXPERIMENT_NAME//_/-}" + AGENT_NAME="${AGENT_NAME}-${EXPERIMENT_SUFFIX}" +fi # Default to Exgentic registry, can be overridden with environment variable EXGENTIC_REGISTRY="${EXGENTIC_REGISTRY:-ghcr.io/exgentic}" @@ -191,6 +202,10 @@ IMAGE_TAG="${IMAGE_TAG:-latest}" REMOTE_IMAGE_NAME="${EXGENTIC_REGISTRY}/${FULL_AGENT_NAME}:${IMAGE_TAG}" TOOL_NAME="exgentic-mcp-${BENCHMARK_NAME}" +# Match the experiment suffix applied to the benchmark so MCP_URL points to the right service +if [ -n "$EXPERIMENT_NAME" ] && [ "$EXPERIMENT_NAME" != "default" ]; then + TOOL_NAME="${TOOL_NAME}-${EXPERIMENT_SUFFIX}" +fi NAMESPACE="team1" # Load shared URL helpers (kagenti_api_url, keycloak_api_url, agent_http_url, …) @@ -699,10 +714,44 @@ if [ "$CLUSTER_MODE" = "openshift" ]; then || echo "Warning: Could not patch route targetPort (route may not exist yet)" fi -# Step 10: Wait for agent to be ready. +# Step 10: Update openai-secret +echo "==========================================" +echo "Final Configuration" +echo "==========================================" +echo "" + +# Step 10.1: Update secrets +if [ "$CLUSTER_MODE" = "kind" ]; then + echo "Step 10.1: Updating secrets..." + "$SCRIPT_DIR/update-secrets.sh" --namespace "$NAMESPACE" +else + echo "Step 10.1: Updating secrets... (skipped — secrets are pre-provisioned on OpenShift/in-cluster)" +fi + +echo "" + +# Step 10.2/10.3: Set resource limits and wait for rollout (local/dev only). +if [ "$CLUSTER_MODE" = "in-cluster" ]; then + echo "Step 10.2: Setting resource limits... (skipped — kubectl not available in-cluster)" +else + echo "Step 10.2: Setting resource limits..." + kubectl set resources deployment/$AGENT_NAME -n $NAMESPACE \ + --limits=cpu=4,memory=2Gi \ + --requests=cpu=500m,memory=512Mi 2>/dev/null \ + && echo "✓ Agent resource limits set (CPU: 4 cores, Memory: 2Gi)" \ + || echo "Warning: Could not set resource limits" + echo "" + + echo "Step 10.3: Waiting for deployment to stabilize..." + kubectl rollout status deployment/$AGENT_NAME -n $NAMESPACE --timeout=120s + echo "✓ Deployment stable" +fi +echo "" + +# Step 11: Wait for agent to be ready after the rollout triggered by set resources. # Uses an HTTP health check (agent card endpoint) — kubectl is not available # inside the job container. -echo "Step 10: Waiting for agent to be ready..." +echo "Step 11: Waiting for agent to be ready..." AGENT_URL="$(agent_http_url "$AGENT_NAME" "$NAMESPACE")" echo " Agent URL: $AGENT_URL" @@ -740,40 +789,6 @@ fi echo "" -# Step 11: Update openai-secret -echo "==========================================" -echo "Final Configuration" -echo "==========================================" -echo "" - -# Step 11.1: Update secrets -if [ "$CLUSTER_MODE" = "kind" ]; then - echo "Step 11.1: Updating secrets..." - "$SCRIPT_DIR/update-secrets.sh" --namespace "$NAMESPACE" -else - echo "Step 11.1: Updating secrets... (skipped — secrets are pre-provisioned on OpenShift/in-cluster)" -fi - -echo "" - -# Step 11.2/11.3: Set resource limits and wait for rollout (local/dev only). -if [ "$CLUSTER_MODE" = "in-cluster" ]; then - echo "Step 11.2: Setting resource limits... (skipped — kubectl not available in-cluster)" -else - echo "Step 11.2: Setting resource limits..." - kubectl set resources deployment/$AGENT_NAME -n $NAMESPACE \ - --limits=cpu=4,memory=2Gi \ - --requests=cpu=500m,memory=512Mi 2>/dev/null \ - && echo "✓ Agent resource limits set (CPU: 4 cores, Memory: 2Gi)" \ - || echo "Warning: Could not set resource limits" - echo "" - - echo "Step 11.3: Waiting for deployment to stabilize..." - kubectl rollout status deployment/$AGENT_NAME -n $NAMESPACE --timeout=120s - echo "✓ Deployment stable" -fi -echo "" - # Step 11.4: Apply the AuthBridge plugin pipeline overlay (if any # plugin selector was supplied). The operator base config enables every # supported plugin; this overlay sets on_error: off on the ones we @@ -801,7 +816,7 @@ if [ "$AUTHBRIDGE_ENABLED" = "true" ]; then echo "" fi -# Step 12: Agent card already verified in Step 10; just show it. +# Step 12: Agent card already verified in Step 11; just show it. echo "Step 12: Agent card:" AGENT_HTTP_ROUTE_URL="$(agent_http_url "$AGENT_NAME" "$NAMESPACE")" CARD_RESPONSE=$(curl -s --max-time 5 "${AGENT_HTTP_ROUTE_URL}/.well-known/agent-card.json" 2>/dev/null || echo "") diff --git a/exgentic_a2a_runner/deploy-and-evaluate.sh b/exgentic_a2a_runner/deploy-and-evaluate.sh index 1b618fa..ed5dcd9 100755 --- a/exgentic_a2a_runner/deploy-and-evaluate.sh +++ b/exgentic_a2a_runner/deploy-and-evaluate.sh @@ -272,6 +272,7 @@ if [ "$DRY_RUN" = "true" ]; then BENCHMARK_CMD_DISPLAY=$(printf '%q ' \ "$SCRIPT_DIR/deploy-benchmark.sh" \ --benchmark "$BENCHMARK_NAME" \ + --experiment "$EXPERIMENT_NAME" \ --model "$MODEL_NAME" \ --keycloak-user "$KEYCLOAK_USERNAME" \ --keycloak-pass "$KEYCLOAK_PASSWORD" \ @@ -282,6 +283,7 @@ if [ "$DRY_RUN" = "true" ]; then echo "" else "$SCRIPT_DIR/deploy-benchmark.sh" --benchmark "$BENCHMARK_NAME" \ + --experiment "$EXPERIMENT_NAME" \ --model "$MODEL_NAME" \ --keycloak-user "$KEYCLOAK_USERNAME" \ --keycloak-pass "$KEYCLOAK_PASSWORD" \ @@ -306,6 +308,7 @@ if [ "$DRY_RUN" = "true" ]; then "$SCRIPT_DIR/deploy-agent.sh" \ --benchmark "$BENCHMARK_NAME" \ --agent "$AGENT_NAME" \ + --experiment "$EXPERIMENT_NAME" \ --model "$MODEL_NAME" \ --keycloak-user "$KEYCLOAK_USERNAME" \ --keycloak-pass "$KEYCLOAK_PASSWORD" \ @@ -319,6 +322,7 @@ if [ "$DRY_RUN" = "true" ]; then echo "" else "$SCRIPT_DIR/deploy-agent.sh" --benchmark "$BENCHMARK_NAME" --agent "$AGENT_NAME" \ + --experiment "$EXPERIMENT_NAME" \ --model "$MODEL_NAME" \ --keycloak-user "$KEYCLOAK_USERNAME" \ --keycloak-pass "$KEYCLOAK_PASSWORD" \ diff --git a/exgentic_a2a_runner/deploy-benchmark.sh b/exgentic_a2a_runner/deploy-benchmark.sh index 72624a6..30e68fb 100755 --- a/exgentic_a2a_runner/deploy-benchmark.sh +++ b/exgentic_a2a_runner/deploy-benchmark.sh @@ -12,6 +12,7 @@ MODEL_NAME="Azure/gpt-4.1" KEYCLOAK_USERNAME="admin" KEYCLOAK_PASSWORD="${KEYCLOAK_PASSWORD:-unknown}" BENCHMARK_NAME="" +EXPERIMENT_NAME="default" USE_MCP_GATEWAY="false" USE_LOCAL_IMAGE="false" CLUSTER_MODE="" @@ -24,6 +25,10 @@ while [[ $# -gt 0 ]]; do BENCHMARK_NAME="$2" shift 2 ;; + --experiment) + EXPERIMENT_NAME="$2" + shift 2 + ;; --model) MODEL_NAME="$2" shift 2 @@ -68,6 +73,7 @@ while [[ $# -gt 0 ]]; do echo " --benchmark NAME Benchmark name (e.g., gsm8k, tau2)" echo "" echo "Optional Arguments:" + echo " --experiment NAME Experiment name suffix appended to pod names (default: default)" echo " --model MODEL Model name (default: Azure/gpt-4.1)" echo " --action-timeout SECONDS Per-action step timeout in seconds (default: 30)" echo " --keycloak-user USER Keycloak username (default: admin)" @@ -123,6 +129,11 @@ EXGENTIC_REGISTRY="${EXGENTIC_REGISTRY:-ghcr.io/exgentic}" IMAGE_TAG="${IMAGE_TAG:-latest}" REMOTE_IMAGE_NAME="${EXGENTIC_REGISTRY}/exgentic-mcp-${BENCHMARK_NAME}:${IMAGE_TAG}" TOOL_NAME="exgentic-mcp-${BENCHMARK_NAME}" +# Append experiment suffix when non-default so parallel experiments get distinct pod names +if [ -n "$EXPERIMENT_NAME" ] && [ "$EXPERIMENT_NAME" != "default" ]; then + EXPERIMENT_SUFFIX="${EXPERIMENT_NAME//_/-}" + TOOL_NAME="${TOOL_NAME}-${EXPERIMENT_SUFFIX}" +fi NAMESPACE="team1" KAGENTI_API="$(kagenti_api_url)" KEYCLOAK_API="$(keycloak_api_url)" @@ -515,10 +526,29 @@ fi echo "" -# Step 11: Wait for MCP server to be ready. +# Step 11: Set resource limits (local/dev only — kubectl not available in-cluster). +if [ "$CLUSTER_MODE" = "in-cluster" ]; then + echo "Step 11: Setting resource limits... (skipped — kubectl not available in-cluster)" +else + echo "Step 11.1: Setting resource limits..." + kubectl set resources deployment/$TOOL_NAME -n $NAMESPACE \ + --limits=cpu=4,memory=4Gi \ + --requests=cpu=500m,memory=512Mi 2>/dev/null \ + && echo "✓ Benchmark resource limits set (CPU: 4 cores, Memory: 4Gi)" \ + || echo "Warning: Could not set resource limits" + echo "" + + echo "Step 11.2: Waiting for deployment to stabilize..." + kubectl rollout status deployment/$TOOL_NAME -n $NAMESPACE --timeout=120s + echo "✓ Deployment stable" +fi + +echo "" + +# Step 12: Wait for MCP server to be ready after the rollout triggered by set resources. # Uses an HTTP health check against the cluster-internal service URL — kubectl is # not available inside the job container. -echo "Step 11: Waiting for MCP server to be ready..." +echo "Step 12: Waiting for MCP server to be ready..." MCP_URL="$(tool_http_url "$TOOL_NAME" "$NAMESPACE")" echo " MCP URL: $MCP_URL" @@ -560,23 +590,6 @@ fi echo "" -# Step 12: Set resource limits (local/dev only — kubectl not available in-cluster). -if [ "$CLUSTER_MODE" = "in-cluster" ]; then - echo "Step 12: Setting resource limits... (skipped — kubectl not available in-cluster)" -else - echo "Step 12.1: Setting resource limits..." - kubectl set resources deployment/$TOOL_NAME -n $NAMESPACE \ - --limits=cpu=4,memory=4Gi \ - --requests=cpu=500m,memory=512Mi 2>/dev/null \ - && echo "✓ Benchmark resource limits set (CPU: 4 cores, Memory: 4Gi)" \ - || echo "Warning: Could not set resource limits" - echo "" - - echo "Step 12.2: Waiting for deployment to stabilize..." - kubectl rollout status deployment/$TOOL_NAME -n $NAMESPACE --timeout=120s - echo "✓ Deployment stable" -fi - # Step 14: Register with MCP Gateway (conditional) if [ "$USE_MCP_GATEWAY" = "true" ]; then echo "" diff --git a/exgentic_a2a_runner/evaluate-benchmark.sh b/exgentic_a2a_runner/evaluate-benchmark.sh index 033f899..bbe5aac 100755 --- a/exgentic_a2a_runner/evaluate-benchmark.sh +++ b/exgentic_a2a_runner/evaluate-benchmark.sh @@ -137,7 +137,16 @@ export AGENT_SERVICE="${FULL_AGENT_NAME}-${BENCHMARK_NAME}" AGENT_SERVICE="${AGENT_SERVICE//_/-}" # Set benchmark service name (override .env values) -export BENCHMARK_SERVICE="exgentic-mcp-${BENCHMARK_NAME}-mcp" +BENCHMARK_BASE="exgentic-mcp-${BENCHMARK_NAME}" + +# Append experiment suffix when non-default so service names match the deployed pods +if [ -n "$EXPERIMENT_NAME" ] && [ "$EXPERIMENT_NAME" != "default" ]; then + EXPERIMENT_SUFFIX="${EXPERIMENT_NAME//_/-}" + AGENT_SERVICE="${AGENT_SERVICE}-${EXPERIMENT_SUFFIX}" + BENCHMARK_BASE="${BENCHMARK_BASE}-${EXPERIMENT_SUFFIX}" +fi + +export BENCHMARK_SERVICE="${BENCHMARK_BASE}-mcp" # MCP Gateway configuration USE_MCP_GATEWAY="${USE_MCP_GATEWAY:-false}" From 27fa9c175b2ebef635df1b4563b4b2ee74254243 Mon Sep 17 00:00:00 2001 From: Yoav Katz Date: Sun, 19 Jul 2026 11:04:47 +0300 Subject: [PATCH 2/3] fix: harden MCP session cleanup and extend deploy timeout - runner: catch delete_session errors inside the MCP.DeleteSession child span so a residual cleanup failure (e.g. sidecar reaped after an agent timeout) never flips a completed session's span to ERROR - mcp_client: match benign "already gone" markers case-insensitively and recognize "no longer alive" as a clean no-op on session delete - deploy-benchmark: extend rollout stabilization timeout 120s -> 360s - analyze_traces: aggregate experiment comparison into a single table with dynamic column widths Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Yoav Katz --- exgentic_a2a_runner/analyze_traces.py | 198 +++++++----------- exgentic_a2a_runner/deploy-benchmark.sh | 2 +- .../exgentic_a2a_runner/mcp_client.py | 14 +- .../exgentic_a2a_runner/runner.py | 23 +- 4 files changed, 110 insertions(+), 127 deletions(-) diff --git a/exgentic_a2a_runner/analyze_traces.py b/exgentic_a2a_runner/analyze_traces.py index bbe957d..68d5cd3 100644 --- a/exgentic_a2a_runner/analyze_traces.py +++ b/exgentic_a2a_runner/analyze_traces.py @@ -589,137 +589,101 @@ def print_experiment_comparison(records: list[TraceRecord]) -> None: if not records: print("No traces to compare.") return - - # Get unique experiments + experiments = sorted(set(r.experiment_name for r in records)) - + if len(experiments) < 2: print(f"Only one experiment found: {experiments[0] if experiments else 'none'}") print("Need at least 2 experiments to compare.") return - + + exp_groups: dict[str, list[TraceRecord]] = defaultdict(list) + for r in records: + exp_groups[r.experiment_name].append(r) + + has_infra = any(t.has_infra for traces in exp_groups.values() for t in traces) + + col_w = max(20, max(len(e) for e in experiments) + 2) + print() print("=" * 140) print(f"Experiment Comparison: {' vs '.join(experiments)}") print("=" * 140) print() - - # Group by (agent, benchmark, model) and then by experiment - config_groups: dict[tuple, dict[str, list[TraceRecord]]] = defaultdict(lambda: defaultdict(list)) - - for r in records: - config_key = (r.agent_name, r.benchmark_name, r.model, r.num_parallel) - config_groups[config_key][r.experiment_name].append(r) - - # Print a table for each configuration - for config_key in sorted(config_groups.keys()): - agent, benchmark, model, num_parallel = config_key - exp_groups = config_groups[config_key] - - # Only show configs that have data for multiple experiments - if len(exp_groups) < 2: - continue - - print(f"\nAgent: {agent} | Benchmark: {benchmark} | Model: {model} | Parallel: {num_parallel}") - print("-" * 140) - - # Check if any traces have infra data - has_infra = any(t.has_infra for traces in exp_groups.values() for t in traces) - - # Build header - header = f"{'Metric':<35s}" - for exp in experiments: - if exp in exp_groups: - header += f" | {exp:>12s}" - print(header) - print("-" * len(header)) - - # Helper to print a metric row - def print_metric_row(label: str, metric_fn, fmt: str = ".2f") -> None: - row = f"{label:<35s}" - for exp in experiments: - if exp in exp_groups: - traces = exp_groups[exp] - values = [metric_fn(t) for t in traces] - avg_val = avg(values) - row += f" | {avg_val:>12{fmt}}" - else: - row += f" | {'N/A':>12s}" - print(row) - - # Count row - count_row = f"{'Traces (count)':<35s}" - for exp in experiments: - if exp in exp_groups: - count_row += f" | {len(exp_groups[exp]):>12d}" - else: - count_row += f" | {'N/A':>12s}" - print(count_row) - - # Evaluation success rate - eval_row = f"{'Eval Success Rate (%)':<35s}" - for exp in experiments: - if exp in exp_groups: - traces = exp_groups[exp] - eval_success = sum(1 for t in traces if t.evaluation_result is True) - rate = (eval_success / len(traces) * 100) if traces else 0 - eval_row += f" | {rate:>12.1f}" - else: - eval_row += f" | {'N/A':>12s}" - print(eval_row) - # --- Timing (absolute) --- - print_metric_row("Total Latency (s)", lambda t: t.total_latency_s) - print_metric_row("Session Creation (s)", lambda t: t.session_creation_s) - print_metric_row("Agent Call (s)", lambda t: t.agent_call_s) - print_metric_row(" Time to First Obs (s)", lambda t: t.time_to_first_obs_s) - print_metric_row(" LLM Calls (s)", lambda t: t.llm_after_obs_s) - print_metric_row(" Tool Calls (s)", lambda t: t.tool_total_s) - print_metric_row(" Overhead (s)", lambda t: t.overhead_s) - print_metric_row("Evaluation (s)", lambda t: t.evaluation_s) + header = f"{'Metric':<35s}" + for exp in experiments: + header += f" | {exp:>{col_w}s}" + print(header) + print("-" * len(header)) - # --- Timing (% of agent call) --- - sep = f"{'':35s}" + def print_metric_row(label: str, metric_fn, fmt: str = ".2f") -> None: + row = f"{label:<35s}" for exp in experiments: - if exp in exp_groups: - sep += f" | {'':>12s}" + traces = exp_groups[exp] + values = [metric_fn(t) for t in traces] + row += f" | {avg(values):>{col_w}{fmt}}" + print(row) + + count_row = f"{'Traces (count)':<35s}" + for exp in experiments: + count_row += f" | {len(exp_groups[exp]):>{col_w}d}" + print(count_row) + + eval_row = f"{'Eval Success Rate (%)':<35s}" + for exp in experiments: + traces = exp_groups[exp] + n = len(traces) + rate = sum(1 for t in traces if t.evaluation_result is True) / n * 100 if n else 0 + eval_row += f" | {rate:>{col_w}.1f}" + print(eval_row) + + print_metric_row("Total Latency (s)", lambda t: t.total_latency_s) + print_metric_row("Session Creation (s)", lambda t: t.session_creation_s) + print_metric_row("Agent Call (s)", lambda t: t.agent_call_s) + print_metric_row(" Time to First Obs (s)", lambda t: t.time_to_first_obs_s) + print_metric_row(" LLM Calls (s)", lambda t: t.llm_after_obs_s) + print_metric_row(" Tool Calls (s)", lambda t: t.tool_total_s) + print_metric_row(" Overhead (s)", lambda t: t.overhead_s) + print_metric_row("Evaluation (s)", lambda t: t.evaluation_s) + + sep = f"{'':35s}" + (f" | {'':>{col_w}s}" * len(experiments)) + print(sep) + print_metric_row("% Time before 1st Obs", + lambda t: (t.time_to_first_obs_s / t.agent_call_s * 100) if t.agent_call_s > 0 else 0, ".1f") + print_metric_row("% Time on LLM calls", + lambda t: (t.llm_after_obs_s / t.agent_call_s * 100) if t.agent_call_s > 0 else 0, ".1f") + print_metric_row("% Time on Tool calls", + lambda t: (t.tool_total_s / t.agent_call_s * 100) if t.agent_call_s > 0 else 0, ".1f") + print_metric_row("% Time overhead", + lambda t: (t.overhead_s / t.agent_call_s * 100) if t.agent_call_s > 0 else 0, ".1f") + + print(sep) + print_metric_row("LLM Calls (count)", lambda t: t.llm_count_after_obs, ".1f") + print_metric_row("Avg LLM Call Latency (s)", + lambda t: (t.llm_after_obs_s / t.llm_count_after_obs) if t.llm_count_after_obs > 0 else 0, ".3f") + print_metric_row("Tool Calls (count)", lambda t: t.tool_count, ".1f") + print_metric_row("Avg Tool Call Latency (s)", + lambda t: (t.tool_total_s / t.tool_count) if t.tool_count > 0 else 0, ".3f") + print_metric_row("LLM Input Tokens", lambda t: t.llm_input_tokens, ".0f") + print_metric_row("LLM Output Tokens", lambda t: t.llm_output_tokens, ".0f") + print_metric_row("LLM Total Tokens", lambda t: t.llm_input_tokens + t.llm_output_tokens, ".0f") + + if has_infra: print(sep) - print_metric_row("% Time before 1st Obs", - lambda t: (t.time_to_first_obs_s / t.agent_call_s * 100) if t.agent_call_s > 0 else 0, ".1f") - print_metric_row("% Time on LLM calls", - lambda t: (t.llm_after_obs_s / t.agent_call_s * 100) if t.agent_call_s > 0 else 0, ".1f") - print_metric_row("% Time on Tool calls", - lambda t: (t.tool_total_s / t.agent_call_s * 100) if t.agent_call_s > 0 else 0, ".1f") - print_metric_row("% Time overhead", - lambda t: (t.overhead_s / t.agent_call_s * 100) if t.agent_call_s > 0 else 0, ".1f") + print_metric_row("MCP CPU Utilization (%)", lambda t: t.mcp_cpu_utilization_pct if t.has_infra else 0, ".1f") + print_metric_row("MCP CPU Throttle (%)", lambda t: t.mcp_throttle_pct if t.has_infra else 0, ".1f") + print_metric_row("MCP Memory Max (MB)", lambda t: t.mcp_memory_max_mb if t.has_infra else 0, ".0f") + print_metric_row("MCP Memory Utilization (%)", lambda t: t.mcp_memory_utilization_pct if t.has_infra else 0, ".1f") + print_metric_row("MCP Network RX (MB)", lambda t: t.mcp_network_rx_mb if t.has_infra else 0, ".3f") + print_metric_row("MCP Network TX (MB)", lambda t: t.mcp_network_tx_mb if t.has_infra else 0, ".3f") + print_metric_row("A2A CPU Utilization (%)", lambda t: t.a2a_cpu_utilization_pct if t.has_infra else 0, ".1f") + print_metric_row("A2A CPU Throttle (%)", lambda t: t.a2a_throttle_pct if t.has_infra else 0, ".1f") + print_metric_row("A2A Memory Max (MB)", lambda t: t.a2a_memory_max_mb if t.has_infra else 0, ".0f") + print_metric_row("A2A Memory Utilization (%)", lambda t: t.a2a_memory_utilization_pct if t.has_infra else 0, ".1f") + print_metric_row("A2A Network RX (MB)", lambda t: t.a2a_network_rx_mb if t.has_infra else 0, ".3f") + print_metric_row("A2A Network TX (MB)", lambda t: t.a2a_network_tx_mb if t.has_infra else 0, ".3f") - # --- Agent metrics --- - print(sep) - print_metric_row("LLM Calls (count)", lambda t: t.llm_count_after_obs, ".1f") - print_metric_row("Avg LLM Call Latency (s)", lambda t: (t.llm_after_obs_s / max(t.llm_count_after_obs, 1)) if t.llm_count_after_obs > 0 else 0, ".3f") - print_metric_row("Tool Calls (count)", lambda t: t.tool_count, ".1f") - print_metric_row("Avg Tool Call Latency (s)", lambda t: (t.tool_total_s / max(t.tool_count, 1)) if t.tool_count > 0 else 0, ".3f") - print_metric_row("LLM Input Tokens", lambda t: t.llm_input_tokens, ".0f") - print_metric_row("LLM Output Tokens", lambda t: t.llm_output_tokens, ".0f") - print_metric_row("LLM Total Tokens", lambda t: t.llm_input_tokens + t.llm_output_tokens, ".0f") - - # --- Infrastructure metrics (if available) --- - if has_infra: - print(sep) - print_metric_row("MCP CPU Utilization (%)", lambda t: t.mcp_cpu_utilization_pct if t.has_infra else 0, ".1f") - print_metric_row("MCP CPU Throttle (%)", lambda t: t.mcp_throttle_pct if t.has_infra else 0, ".1f") - print_metric_row("MCP Memory Max (MB)", lambda t: t.mcp_memory_max_mb if t.has_infra else 0, ".0f") - print_metric_row("MCP Memory Utilization (%)", lambda t: t.mcp_memory_utilization_pct if t.has_infra else 0, ".1f") - print_metric_row("MCP Network RX (MB)", lambda t: t.mcp_network_rx_mb if t.has_infra else 0, ".3f") - print_metric_row("MCP Network TX (MB)", lambda t: t.mcp_network_tx_mb if t.has_infra else 0, ".3f") - - print_metric_row("A2A CPU Utilization (%)", lambda t: t.a2a_cpu_utilization_pct if t.has_infra else 0, ".1f") - print_metric_row("A2A CPU Throttle (%)", lambda t: t.a2a_throttle_pct if t.has_infra else 0, ".1f") - print_metric_row("A2A Memory Max (MB)", lambda t: t.a2a_memory_max_mb if t.has_infra else 0, ".0f") - print_metric_row("A2A Memory Utilization (%)", lambda t: t.a2a_memory_utilization_pct if t.has_infra else 0, ".1f") - print_metric_row("A2A Network RX (MB)", lambda t: t.a2a_network_rx_mb if t.has_infra else 0, ".3f") - print_metric_row("A2A Network TX (MB)", lambda t: t.a2a_network_tx_mb if t.has_infra else 0, ".3f") - print() @@ -738,7 +702,7 @@ def main() -> int: raw = json.load(sys.stdin) records = parse_traces(raw) - + if args.compare: print_experiment_comparison(records) else: diff --git a/exgentic_a2a_runner/deploy-benchmark.sh b/exgentic_a2a_runner/deploy-benchmark.sh index 30e68fb..ce53473 100755 --- a/exgentic_a2a_runner/deploy-benchmark.sh +++ b/exgentic_a2a_runner/deploy-benchmark.sh @@ -539,7 +539,7 @@ else echo "" echo "Step 11.2: Waiting for deployment to stabilize..." - kubectl rollout status deployment/$TOOL_NAME -n $NAMESPACE --timeout=120s + kubectl rollout status deployment/$TOOL_NAME -n $NAMESPACE --timeout=360s echo "✓ Deployment stable" fi diff --git a/exgentic_a2a_runner/exgentic_a2a_runner/mcp_client.py b/exgentic_a2a_runner/exgentic_a2a_runner/mcp_client.py index b18cd39..ec5c8be 100644 --- a/exgentic_a2a_runner/exgentic_a2a_runner/mcp_client.py +++ b/exgentic_a2a_runner/exgentic_a2a_runner/mcp_client.py @@ -43,6 +43,7 @@ def _tool_name(self, name: str) -> str: if self._tool_prefix: return f"{self._tool_prefix}{name}" return name + def _get_event_loop(self) -> asyncio.AbstractEventLoop: """Get or create thread-local event loop.""" if not hasattr(self._local, "loop"): @@ -155,6 +156,7 @@ def initialize(self) -> None: logger.info(f"Initializing MCP client for {self.mcp_url}") try: + async def _init(): session = await self._ensure_session() tools_result = await session.list_tools() @@ -284,8 +286,16 @@ def delete_session(self, session_id: str) -> None: status = result.get("status", "") if status != "success": error_msg = result.get("error", "") - if "client has been closed" in error_msg or "No session found" in error_msg: - logger.debug(f"Session {session_id} already cleaned up: {error_msg}") + # Markers indicating the session/process is already gone. Deleting an + # already-dead session (e.g. the MCP sidecar was reaped after an agent + # timeout) is a clean no-op, not a failure — treat it as such. + benign_markers = ( + "client has been closed", + "no session found", + "no longer alive", + ) + if any(marker in error_msg.lower() for marker in benign_markers): + logger.warning(f"Session {session_id} already cleaned up: {error_msg}") return raise RuntimeError(f"Session deletion failed: {result}") logger.info(f"Session {session_id} deleted") diff --git a/exgentic_a2a_runner/exgentic_a2a_runner/runner.py b/exgentic_a2a_runner/exgentic_a2a_runner/runner.py index 08d2826..0809ca8 100644 --- a/exgentic_a2a_runner/exgentic_a2a_runner/runner.py +++ b/exgentic_a2a_runner/exgentic_a2a_runner/runner.py @@ -331,9 +331,15 @@ def process_task(self, task_id: str) -> SessionResult: except Exception as e: logger.warning("Failed to collect infra metrics: %s", e) - # Delete session + # Delete session. Cleanup must not flip a completed session to + # ERROR: catch inside the child span so a residual cleanup error + # (e.g. sidecar already reaped after a timeout) never exits the + # span context nor propagates to the outer failure handler. with self.otel.child_span("MCP.DeleteSession"): - self.exgentic.delete_session(session_id) + try: + self.exgentic.delete_session(session_id) + except Exception as delete_error: + logger.warning(f"Failed to delete session {session_id}: {delete_error}") # Record success self.otel.record_success(span, evaluation_result) @@ -365,13 +371,16 @@ def process_task(self, task_id: str) -> SessionResult: logger.error(f"Session {session_id} failed: {error_type}: {error_msg}") - # Try to delete session even on failure (only if session was created) + # Try to delete session even on failure (only if session was created). + # Catch inside the child span so a cleanup failure never marks the + # MCP.DeleteSession span ERROR — the session's real failure is already + # recorded on the parent Agent.Session span below. if session_id and not session_id.startswith("pending-"): - try: - with self.otel.child_span("MCP.DeleteSession"): + with self.otel.child_span("MCP.DeleteSession"): + try: self.exgentic.delete_session(session_id) - except Exception as delete_error: - logger.warning(f"Failed to delete session {session_id}: {delete_error}") + except Exception as delete_error: + logger.warning(f"Failed to delete session {session_id}: {delete_error}") # Record failure self.otel.record_failure(span, e, error_type) From f2579bda0f8c9e4da74cf78256aa51cbadd91393 Mon Sep 17 00:00:00 2001 From: Yoav Katz Date: Sun, 19 Jul 2026 16:56:39 +0300 Subject: [PATCH 3/3] feat(analyze-run): support OpenShift MLflow and time-window trace fetch Rework analyze-run.sh to reach MLflow via kubectl port-forward for both --kind and --openshift, mirroring evaluate-benchmark.sh's OTEL collector approach. MLflow location, TLS, workspace, auth mode, and experiment id now default per cluster mode, with env/flag overrides. - Add --openshift DOMAIN / --kind cluster modes with per-mode defaults (namespace, service, port, TLS, workspace, auth mode, experiment id) - Add --auth-mode (secret | oc-token); oc-token uses `oc whoami -t` - Replace --limit trace count with --window (e.g. 3h, 90m, 2d); the downloader pages newest-first and stops at the first out-of-window trace - Skip trivially short traces (MIN_DURATION_S) before fetching their spans - Send x-mlflow-workspace header and allow insecure TLS for port-forwarded reencrypt HTTPS endpoints (RHOAI) - Verify kubectl context matches the cluster mode via libsh helper - Harden secret-mode token flow against set -e command-substitution aborts Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Yoav Katz --- exgentic_a2a_runner/analyze-run.sh | 382 ++++++++++++++---- exgentic_a2a_runner/download_mlflow_traces.py | 92 ++++- 2 files changed, 373 insertions(+), 101 deletions(-) diff --git a/exgentic_a2a_runner/analyze-run.sh b/exgentic_a2a_runner/analyze-run.sh index 9f76639..4ae29e1 100755 --- a/exgentic_a2a_runner/analyze-run.sh +++ b/exgentic_a2a_runner/analyze-run.sh @@ -4,6 +4,11 @@ # # Bash handles: connectivity, port-forwarding, OAuth token acquisition # Python handles: downloading traces, format transformation, and analysis (analyze_traces.py) +# +# MLflow access mirrors how evaluate-benchmark.sh reaches the OTEL collector: +# by default we kubectl port-forward svc/mlflow:5000 -> localhost:$MLFLOW_LOCAL_PORT +# for both --kind and --openshift, then talk to http://localhost:$MLFLOW_LOCAL_PORT. +# Pass -u/--url to skip the port-forward and hit a reachable MLflow URL directly. set -e @@ -11,34 +16,79 @@ SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" # Default values MLFLOW_URL="" -LIMIT=100 -AUTO_PORT_FORWARD="false" -MLFLOW_NAMESPACE="kagenti-system" -MLFLOW_SERVICE="mlflow" -MLFLOW_LOCAL_PORT="8080" +# Time window to fetch, e.g. 3h / 90m / 2d. Only traces newer than this are +# downloaded. Replaces the old trace-count limit. +WINDOW="${WINDOW:-3h}" +# MLflow location, TLS, workspace, and auth mode all default based on the +# cluster mode (--kind vs --openshift) — see the mode dispatch below. These +# start empty so we can tell "user/env supplied a value" from "apply the +# per-mode default". An env var of the same name still pre-seeds the value and +# takes precedence over the mode default; an explicit CLI flag overrides both. +MLFLOW_NAMESPACE="${MLFLOW_NAMESPACE:-}" +MLFLOW_SERVICE="${MLFLOW_SERVICE:-}" +MLFLOW_REMOTE_PORT="${MLFLOW_REMOTE_PORT:-}" +MLFLOW_LOCAL_PORT="${MLFLOW_LOCAL_PORT:-8080}" +MLFLOW_TLS="${MLFLOW_TLS:-}" +MLFLOW_WORKSPACE="${MLFLOW_WORKSPACE:-}" +AUTH_MODE="${AUTH_MODE:-}" KUBECTL_BIN="${KUBECTL_BIN:-kubectl}" -EXPERIMENT_ID="0" +# `whoami -t` is an OpenShift (oc) extension, not a kubectl subcommand, so the +# token command is separate from KUBECTL_BIN. Override with OC_BIN if needed. +OC_BIN="${OC_BIN:-oc}" +# Experiment ID also defaults per cluster mode (see the mode dispatch below): +# kind uses 0 (the Default experiment); OpenShift uses 1, since id 0 does not +# exist on RHOAI. Starts empty so we can tell "user set it" from "use default". +EXPERIMENT_ID="${EXPERIMENT_ID:-}" EXPERIMENT_FILTER="" COMPARE_EXPERIMENTS="" +CLUSTER_MODE="" +INGRESS_DOMAIN="" usage() { cat << EOF Usage: $0 [OPTIONS] Options: - -u, --url URL MLflow REST API base URL (default: http://mlflow.localtest.me:8080) - -l, --limit NUM Limit number of traces to download (default: 100) + -u, --url URL MLflow REST API base URL. If omitted, port-forward svc/mlflow from the cluster. + -w, --window DURATION Fetch traces from the last DURATION, e.g. 3h, 90m, 2d (default: 3h) -e, --experiment NAME Filter traces by experiment name attribute -c, --compare EXP1,EXP2 Compare two experiments (comma-separated) - --experiment-id ID MLflow experiment ID to query (default: 0) - -f, --forward Auto port-forward MLflow from kind cluster if not accessible + --experiment-id ID MLflow experiment ID to query (default: per cluster mode, see below) + --kind Target a local Kind cluster (default) + --openshift DOMAIN Target an OpenShift cluster with the given ingress domain + --mlflow-namespace NS Namespace of the MLflow service + --mlflow-service NAME Name of the MLflow service + --mlflow-port PORT Remote MLflow service port to forward + --mlflow-tls MLflow serves HTTPS on the forwarded port + --mlflow-workspace NAME Send x-mlflow-workspace header + --auth-mode MODE Token source: secret (kagenti oauth secret) or oc-token (oc whoami -t) -h, --help Show this help message +The MLflow location, TLS, workspace, auth mode, and experiment id all DEFAULT +from the cluster mode, so a plain --kind or --openshift needs no other flags: + + --kind --openshift + namespace kagenti-system redhat-ods-applications + service mlflow mlflow + remote port 5000 8443 + tls off (http) on (https) + workspace (none) team1 + auth mode secret oc-token + experiment id 0 1 + +Any of the above flags (or the matching env var) overrides its per-mode +default; env vars are also honored over the default. By default (no -u/--url) +the script port-forwards the MLflow service to localhost:${MLFLOW_LOCAL_PORT} for both +modes, matching how evaluate-benchmark.sh reaches the OTEL collector. + Examples: - $0 -f -l 50 + $0 --window 1h $0 --experiment baseline $0 --compare baseline,test1 - $0 -u http://mlflow.localtest.me:8080 -l 200 + $0 --kind --window 6h + $0 --openshift apps.mycluster.example.com + $0 --openshift apps.mycluster.example.com --experiment-id 3 --compare baseline,test1 + $0 -u http://mlflow.localtest.me:8080 --window 2d EOF exit 1 } @@ -46,25 +96,130 @@ EOF while [[ $# -gt 0 ]]; do case $1 in -u|--url) MLFLOW_URL="$2"; shift 2 ;; - -l|--limit) LIMIT="$2"; shift 2 ;; + -w|--window) WINDOW="$2"; shift 2 ;; -e|--experiment) EXPERIMENT_FILTER="$2"; shift 2 ;; -c|--compare) COMPARE_EXPERIMENTS="$2"; shift 2 ;; --experiment-id) EXPERIMENT_ID="$2"; shift 2 ;; - -f|--forward) AUTO_PORT_FORWARD="true"; shift ;; + --mlflow-namespace) MLFLOW_NAMESPACE="$2"; shift 2 ;; + --mlflow-service) MLFLOW_SERVICE="$2"; shift 2 ;; + --mlflow-port) MLFLOW_REMOTE_PORT="$2"; shift 2 ;; + --mlflow-tls) MLFLOW_TLS="true"; shift ;; + --mlflow-workspace) MLFLOW_WORKSPACE="$2"; shift 2 ;; + --auth-mode) AUTH_MODE="$2"; shift 2 ;; + --kind) CLUSTER_MODE="kind"; shift ;; + --openshift) + CLUSTER_MODE="openshift" + if [ $# -lt 2 ]; then + echo "Error: --openshift requires an ingress domain argument" + usage + fi + INGRESS_DOMAIN="$2" + shift 2 + ;; -h|--help) usage ;; *) echo "Unknown option: $1"; usage ;; esac done -# Set default URL if not provided +# Default to kind when no cluster mode is given +if [ -z "$CLUSTER_MODE" ]; then + CLUSTER_MODE="kind" +fi + +# Validate cluster mode and its arguments, and fill in per-mode MLflow defaults. +# Every assignment uses ${VAR:-default} so an env var or explicit CLI flag that +# already set the value wins; only unset values fall back to the mode default. +case "$CLUSTER_MODE" in + kind) + # kagenti's kind MLflow: HTTP on port 5000 in kagenti-system, no + # workspace header, client-credentials secret flow for auth. + MLFLOW_NAMESPACE="${MLFLOW_NAMESPACE:-kagenti-system}" + MLFLOW_SERVICE="${MLFLOW_SERVICE:-mlflow}" + MLFLOW_REMOTE_PORT="${MLFLOW_REMOTE_PORT:-5000}" + MLFLOW_TLS="${MLFLOW_TLS:-false}" + # MLFLOW_WORKSPACE left empty: kind MLflow needs no workspace header. + AUTH_MODE="${AUTH_MODE:-secret}" + EXPERIMENT_ID="${EXPERIMENT_ID:-0}" + ;; + openshift) + if [ -z "$INGRESS_DOMAIN" ]; then + echo "Error: --openshift requires an ingress domain argument" + usage + fi + # RHOAI-managed MLflow: HTTPS on port 8443 in redhat-ods-applications, + # behind an oauth-proxy that accepts the logged-in user token, and it + # requires an x-mlflow-workspace header (team1). + MLFLOW_NAMESPACE="${MLFLOW_NAMESPACE:-redhat-ods-applications}" + MLFLOW_SERVICE="${MLFLOW_SERVICE:-mlflow}" + MLFLOW_REMOTE_PORT="${MLFLOW_REMOTE_PORT:-8443}" + MLFLOW_TLS="${MLFLOW_TLS:-true}" + MLFLOW_WORKSPACE="${MLFLOW_WORKSPACE:-team1}" + AUTH_MODE="${AUTH_MODE:-oc-token}" + # Experiment 0 does not exist on RHOAI; default to the lowest real id. + EXPERIMENT_ID="${EXPERIMENT_ID:-1}" + ;; + *) + echo "Error: unsupported cluster mode '${CLUSTER_MODE}'. Use --kind or --openshift DOMAIN." + exit 1 + ;; +esac + +case "$AUTH_MODE" in + secret|oc-token) ;; + *) + echo "Error: unsupported --auth-mode '${AUTH_MODE}'. Use secret or oc-token." + exit 1 + ;; +esac + +# Decide how to reach MLflow. An explicit -u/--url is used as-is and skips the +# port-forward; otherwise we port-forward svc/mlflow and talk to it on localhost +# (same approach evaluate-benchmark.sh uses for the OTEL collector). +USE_PORT_FORWARD="false" if [ -z "$MLFLOW_URL" ]; then - MLFLOW_URL="http://mlflow.localtest.me:8080" + USE_PORT_FORWARD="true" + if [ "$MLFLOW_TLS" = "true" ]; then + MLFLOW_URL="https://localhost:${MLFLOW_LOCAL_PORT}" + else + MLFLOW_URL="http://localhost:${MLFLOW_LOCAL_PORT}" + fi +fi + +# Parse the time window (e.g. 3h, 90m, 2d, or a bare number = hours) into +# milliseconds for the Python downloader. +parse_window_ms() { + local w="$1" num unit + if [[ "$w" =~ ^([0-9]+)([hmd]?)$ ]]; then + num="${BASH_REMATCH[1]}" + unit="${BASH_REMATCH[2]:-h}" + case "$unit" in + h) echo $(( num * 3600 * 1000 )) ;; + m) echo $(( num * 60 * 1000 )) ;; + d) echo $(( num * 86400 * 1000 )) ;; + esac + return 0 + fi + return 1 +} + +if ! WINDOW_MS=$(parse_window_ms "$WINDOW"); then + echo "Error: invalid --window '$WINDOW'. Use e.g. 3h, 90m, 2d." + exit 1 fi echo "=== MLflow Trace Analysis ===" -echo "MLflow URL: $MLFLOW_URL" +echo "Cluster mode: $CLUSTER_MODE" +if [ "$USE_PORT_FORWARD" = "true" ]; then + echo "MLflow URL: $MLFLOW_URL (via port-forward svc/${MLFLOW_SERVICE})" +else + echo "MLflow URL: $MLFLOW_URL (direct)" +fi echo "Experiment ID: $EXPERIMENT_ID" -echo "Limit: $LIMIT" +echo "Auth mode: $AUTH_MODE" +if [ -n "$MLFLOW_WORKSPACE" ]; then + echo "Workspace: $MLFLOW_WORKSPACE" +fi +echo "Window: $WINDOW" if [ -n "$EXPERIMENT_FILTER" ]; then echo "Experiment Filter: $EXPERIMENT_FILTER" fi @@ -73,24 +228,75 @@ if [ -n "$COMPARE_EXPERIMENTS" ]; then fi echo "" +# --- Verify kubectl points at the cluster matching CLUSTER_MODE --- +# The port-forward and OAuth steps below read a secret, exec into the MLflow +# pod, and forward its service, so the active kubectl context must match the +# requested mode. Catching a mismatch here gives a clear error up front. +export CLUSTER_MODE INGRESS_DOMAIN KUBECTL_BIN +# shellcheck source=libsh/check-kubectl-context.sh +source "$SCRIPT_DIR/libsh/check-kubectl-context.sh" +check_kubectl_context +echo "" + # --- Helper functions --- OAUTH_TOKEN="" +PF_MLFLOW_PID="" -get_oauth_token() { - echo "Obtaining OAuth token from cluster..." +# Port-forward the MLflow service (remote port $MLFLOW_REMOTE_PORT) to +# localhost:$MLFLOW_LOCAL_PORT. Mirrors the OTEL collector port-forward in +# evaluate-benchmark.sh; used for both kind and openshift when no explicit +# -u/--url was given. +setup_port_forward() { + echo "Starting port-forward for MLflow (${MLFLOW_NAMESPACE}/svc/${MLFLOW_SERVICE}:${MLFLOW_REMOTE_PORT} -> localhost:${MLFLOW_LOCAL_PORT})..." + + echo "Checking if MLflow pod is ready..." + if ! "$KUBECTL_BIN" wait --for=condition=ready pod -l app=mlflow -n "$MLFLOW_NAMESPACE" --timeout=30s >/dev/null 2>&1; then + echo "Error: MLflow pod (label app=mlflow) is not ready in namespace $MLFLOW_NAMESPACE" + return 1 + fi - local secret_json - secret_json=$("$KUBECTL_BIN" get secret mlflow-oauth-secret -n "$MLFLOW_NAMESPACE" -o json 2>/dev/null) - if [ $? -ne 0 ] || [ -z "$secret_json" ]; then + "$KUBECTL_BIN" port-forward -n "$MLFLOW_NAMESPACE" "svc/${MLFLOW_SERVICE}" "${MLFLOW_LOCAL_PORT}:${MLFLOW_REMOTE_PORT}" >/dev/null 2>&1 & + PF_MLFLOW_PID=$! + sleep 3 + + if ! ps -p "$PF_MLFLOW_PID" > /dev/null; then + echo "Error: MLflow port-forward failed to start" + return 1 + fi + + echo "✓ MLflow port-forward established (PID: $PF_MLFLOW_PID)" + return 0 +} + +cleanup_port_forward() { + if [ -n "$PF_MLFLOW_PID" ]; then + echo "" + echo "Stopping MLflow port-forward (PID: $PF_MLFLOW_PID)..." + kill "$PF_MLFLOW_PID" 2>/dev/null || true + fi +} + +# secret mode: kagenti's client-credentials flow. Reads mlflow-oauth-secret and +# execs into the MLflow pod to exchange it for an access token. +get_token_from_secret() { + echo "Obtaining OAuth token via mlflow-oauth-secret..." + + # Note: under `set -e`, a failing command substitution aborts the script + # before the following `if` can run. Capture status explicitly so the + # error message below actually prints instead of the script dying silently. + local secret_json secret_status + secret_json=$("$KUBECTL_BIN" get secret mlflow-oauth-secret -n "$MLFLOW_NAMESPACE" -o json 2>/dev/null) && secret_status=0 || secret_status=$? + if [ "$secret_status" -ne 0 ] || [ -z "$secret_json" ]; then echo "Error: Could not read mlflow-oauth-secret from namespace $MLFLOW_NAMESPACE" + echo "Hint: confirm the secret exists on the current cluster ($($KUBECTL_BIN config current-context 2>/dev/null))" return 1 fi local client_id client_secret token_url - client_id=$(echo "$secret_json" | jq -r '.data["OIDC_CLIENT_ID"]' | base64 -d) - client_secret=$(echo "$secret_json" | jq -r '.data["OIDC_CLIENT_SECRET"]' | base64 -d) - token_url=$(echo "$secret_json" | jq -r '.data["OIDC_TOKEN_URL"]' | base64 -d) + client_id=$(echo "$secret_json" | jq -r '.data["OIDC_CLIENT_ID"]' | base64 -d) || true + client_secret=$(echo "$secret_json" | jq -r '.data["OIDC_CLIENT_SECRET"]' | base64 -d) || true + token_url=$(echo "$secret_json" | jq -r '.data["OIDC_TOKEN_URL"]' | base64 -d) || true if [ -z "$client_id" ] || [ -z "$client_secret" ] || [ -z "$token_url" ]; then echo "Error: Could not extract OAuth credentials from secret" @@ -98,7 +304,7 @@ get_oauth_token() { fi local mlflow_pod - mlflow_pod=$("$KUBECTL_BIN" get pod -n "$MLFLOW_NAMESPACE" -l app=mlflow -o jsonpath='{.items[0].metadata.name}' 2>/dev/null) + mlflow_pod=$("$KUBECTL_BIN" get pod -n "$MLFLOW_NAMESPACE" -l app=mlflow -o jsonpath='{.items[0].metadata.name}' 2>/dev/null) || true if [ -z "$mlflow_pod" ]; then echo "Error: Could not find MLflow pod" return 1 @@ -116,9 +322,9 @@ data = urllib.parse.urlencode({ req = urllib.request.Request('${token_url}', data=data, headers={'Content-Type': 'application/x-www-form-urlencoded'}) resp = urllib.request.urlopen(req) print(resp.read().decode()) -" 2>/dev/null) +" 2>/dev/null) || true - OAUTH_TOKEN=$(echo "$token_response" | jq -r '.access_token' 2>/dev/null) + OAUTH_TOKEN=$(echo "$token_response" | jq -r '.access_token' 2>/dev/null) || true if [ -z "$OAUTH_TOKEN" ] || [ "$OAUTH_TOKEN" = "null" ]; then echo "Error: Could not obtain OAuth token" echo "Response: $token_response" @@ -128,80 +334,67 @@ print(resp.read().decode()) echo "✓ OAuth token obtained" } -setup_port_forward() { - echo "Setting up port forwarding to MLflow in kind cluster..." - - if ! command -v "$KUBECTL_BIN" &> /dev/null; then - echo "Error: $KUBECTL_BIN is not installed or not in PATH"; return 1 - fi - if ! CURRENT_CONTEXT=$("$KUBECTL_BIN" config current-context 2>/dev/null); then - echo "Error: Unable to determine current kubectl context"; return 1 - fi - if [ "$CURRENT_CONTEXT" != "kind-kagenti" ]; then - echo "Warning: Not connected to kind-kagenti cluster (current: $CURRENT_CONTEXT)"; return 1 +# oc-token mode: use the logged-in user token, which the RHOAI mlflow-oauth-proxy +# accepts as a bearer token (mirrors the collector's serviceaccount-token auth). +get_token_from_oc() { + if ! command -v "$OC_BIN" >/dev/null 2>&1; then + echo "Error: '$OC_BIN' not found; the oc-token auth mode needs the OpenShift CLI" + echo "Hint: install oc, or set OC_BIN to its path" + return 1 fi - - echo "Checking if MLflow pod is ready..." - if ! "$KUBECTL_BIN" wait --for=condition=ready pod -l app=mlflow -n $MLFLOW_NAMESPACE --timeout=30s >/dev/null 2>&1; then - echo "Error: MLflow pod is not ready in cluster"; return 1 + echo "Obtaining bearer token via '$OC_BIN whoami -t'..." + OAUTH_TOKEN=$("$OC_BIN" whoami -t 2>/dev/null) || true + if [ -z "$OAUTH_TOKEN" ]; then + echo "Error: Could not obtain a user token from '$OC_BIN whoami -t'" + echo "Hint: log in first (e.g. 'oc login ...') so a bearer token is available" + return 1 fi + echo "✓ Bearer token obtained" +} - echo "Cleaning up existing port-forward on port ${MLFLOW_LOCAL_PORT}..." - lsof -ti:${MLFLOW_LOCAL_PORT} | xargs kill -9 2>/dev/null || true - sleep 2 - - echo "Starting port-forward: localhost:${MLFLOW_LOCAL_PORT} -> ${MLFLOW_SERVICE}.${MLFLOW_NAMESPACE}:5000" - "$KUBECTL_BIN" port-forward -n $MLFLOW_NAMESPACE svc/$MLFLOW_SERVICE ${MLFLOW_LOCAL_PORT}:5000 >/dev/null 2>&1 & - PF_MLFLOW_PID=$! - sleep 3 - - if curl -s --max-time 2 -o /dev/null -w "%{http_code}" "${MLFLOW_URL}/health" 2>/dev/null | grep -q "200\|404"; then - echo "✓ Port-forward established successfully" - else - echo "Warning: Port-forward started but MLflow may not be responding yet" - fi - return 0 +# Dispatch token acquisition based on the resolved auth mode. +get_oauth_token() { + case "$AUTH_MODE" in + secret) get_token_from_secret ;; + oc-token) get_token_from_oc ;; + *) echo "Error: unsupported auth mode '$AUTH_MODE'"; return 1 ;; + esac } -cleanup_port_forward() { - if [ -n "$PF_MLFLOW_PID" ]; then - echo "" - echo "Cleaning up port-forward (PID: $PF_MLFLOW_PID)..." - kill $PF_MLFLOW_PID 2>/dev/null || true +# --- Step 1: Port-forward (if needed) and test connectivity --- + +if [ "$USE_PORT_FORWARD" = "true" ]; then + if ! setup_port_forward; then + echo "Error: Failed to set up MLflow port-forward" + exit 1 fi -} + trap cleanup_port_forward EXIT + echo "" +fi -# --- Step 1: Test connectivity / setup port-forward --- +# A reencrypt/self-signed TLS endpoint on localhost won't pass cert +# verification, so allow insecure TLS for the health check when --mlflow-tls +# is set. (This only affects the bash health probe; the Python downloader has +# its own connection handling — see note below.) +CURL_TLS_OPTS=() +if [ "$MLFLOW_TLS" = "true" ]; then + CURL_TLS_OPTS=(-k) +fi echo "Connecting to MLflow..." set +e -HEALTH_CHECK=$(curl -s --max-time 5 -o /dev/null -w "%{http_code}" "${MLFLOW_URL}/health" 2>&1) +HEALTH_CHECK=$(curl -s "${CURL_TLS_OPTS[@]}" --max-time 5 -o /dev/null -w "%{http_code}" "${MLFLOW_URL}/health" 2>&1) CURL_EXIT=$? set -e if [[ $CURL_EXIT -ne 0 ]] || [[ "$HEALTH_CHECK" == "000" ]]; then - if [ "$AUTO_PORT_FORWARD" = "true" ]; then - echo "MLflow not accessible locally, attempting to port-forward from kind cluster..." - echo "" - if setup_port_forward; then - trap cleanup_port_forward EXIT - echo "" - echo "Retrying MLflow connection..." - set +e - HEALTH_CHECK=$(curl -s --max-time 5 -o /dev/null -w "%{http_code}" "${MLFLOW_URL}/health" 2>&1) - CURL_EXIT=$? - set -e - if [[ $CURL_EXIT -ne 0 ]] || [[ "$HEALTH_CHECK" == "000" ]]; then - echo "Error: Still unable to connect to MLflow after port-forwarding"; exit 1 - fi - else - echo "Error: Failed to setup port-forward to MLflow"; exit 1 - fi + echo "Error: Failed to connect to MLflow at $MLFLOW_URL" + if [ "$USE_PORT_FORWARD" = "true" ]; then + echo "The port-forward to svc/${MLFLOW_SERVICE} started but MLflow is not responding." else - echo "Error: Failed to connect to MLflow at $MLFLOW_URL" - echo "Use --forward flag to auto port-forward from kind cluster" - exit 1 + echo "Check that the URL passed via -u/--url points to a reachable MLflow instance." fi + exit 1 fi echo "✓ Connected to MLflow" @@ -209,12 +402,23 @@ echo "" # --- Step 2: Obtain OAuth token --- -get_oauth_token +if ! get_oauth_token; then + echo "Error: Failed to obtain OAuth token; cannot download traces" + exit 1 +fi echo "" # --- Step 3: Download traces, transform, and pipe to analyze_traces.py --- -export MLFLOW_URL OAUTH_TOKEN EXPERIMENT_ID LIMIT EXPERIMENT_FILTER COMPARE_EXPERIMENTS +# MLFLOW_WORKSPACE: sent as the x-mlflow-workspace header (required by RHOAI). +# MLFLOW_INSECURE_TLS: skip cert verification for the port-forwarded HTTPS +# endpoint (reencrypt cert won't validate against localhost). +MLFLOW_INSECURE_TLS="false" +if [ "$MLFLOW_TLS" = "true" ]; then + MLFLOW_INSECURE_TLS="true" +fi +export MLFLOW_URL OAUTH_TOKEN EXPERIMENT_ID WINDOW_MS EXPERIMENT_FILTER COMPARE_EXPERIMENTS +export MLFLOW_WORKSPACE MLFLOW_INSECURE_TLS PYTHON_ARGS="" if [ -n "$COMPARE_EXPERIMENTS" ]; then diff --git a/exgentic_a2a_runner/download_mlflow_traces.py b/exgentic_a2a_runner/download_mlflow_traces.py index 79ae5a3..cf3fd42 100644 --- a/exgentic_a2a_runner/download_mlflow_traces.py +++ b/exgentic_a2a_runner/download_mlflow_traces.py @@ -5,9 +5,17 @@ MLFLOW_URL - MLflow base URL (e.g. http://localhost:8081) OAUTH_TOKEN - Bearer token for auth EXPERIMENT_ID - MLflow experiment ID (default: 0) - LIMIT - Max traces to fetch (default: 100) + WINDOW_MS - Only fetch traces newer than this many milliseconds + ago (default: 10800000 = 3h). The listing is paged + newest-first and stops at the first older trace. EXPERIMENT_FILTER - Filter by experiment_name attribute COMPARE_EXPERIMENTS - Comma-separated experiment names to include + MLFLOW_WORKSPACE - Workspace context; sent as x-mlflow-workspace header + when set (required by RHOAI-managed MLflow) + MLFLOW_INSECURE_TLS - "true" to skip TLS cert verification (needed when + reaching a reencrypt HTTPS endpoint via port-forward) + MIN_DURATION_S - Skip traces shorter than this (seconds, default 0.1) + at the listing stage; their spans are never fetched. Outputs JSON to stdout: {"traces": [{"traceId": ..., "spans": [...]}]} """ @@ -16,7 +24,9 @@ import json import os +import ssl import sys +import time import urllib.request import urllib.error from datetime import datetime, timezone @@ -24,41 +34,94 @@ MLFLOW_URL = os.environ["MLFLOW_URL"] TOKEN = os.environ["OAUTH_TOKEN"] EXPERIMENT_ID = os.environ.get("EXPERIMENT_ID", "0") -LIMIT = int(os.environ.get("LIMIT", "100")) +# Time window: fetch traces from the last WINDOW_MS milliseconds (default 3h). +WINDOW_MS = int(os.environ.get("WINDOW_MS", str(3 * 3600 * 1000))) EXPERIMENT_FILTER = os.environ.get("EXPERIMENT_FILTER", "") COMPARE_EXPERIMENTS = os.environ.get("COMPARE_EXPERIMENTS", "") +MLFLOW_WORKSPACE = os.environ.get("MLFLOW_WORKSPACE", "") +MLFLOW_INSECURE_TLS = os.environ.get("MLFLOW_INSECURE_TLS", "false").lower() == "true" +# Skip trivially short traces (agent-card probes, health checks) that have no +# real request. Filtered from the listing before spans are fetched. In seconds. +MIN_DURATION_S = float(os.environ.get("MIN_DURATION_S", "0.1")) + +CUTOFF_MS = int(time.time() * 1000) - WINDOW_MS + +# When talking to a port-forwarded reencrypt HTTPS endpoint the cert won't +# validate against localhost, so allow opting out of verification. +_SSL_CONTEXT = ssl._create_unverified_context() if MLFLOW_INSECURE_TLS else None def mlflow_get(path: str) -> dict: url = f"{MLFLOW_URL}{path}" - req = urllib.request.Request(url, headers={ + headers = { "Authorization": f"Bearer {TOKEN}", "Content-Type": "application/json", - }) - with urllib.request.urlopen(req, timeout=30) as resp: + } + # RHOAI MLflow requires a workspace context on every request. + if MLFLOW_WORKSPACE: + headers["x-mlflow-workspace"] = MLFLOW_WORKSPACE + req = urllib.request.Request(url, headers=headers) + with urllib.request.urlopen(req, timeout=30, context=_SSL_CONTEXT) as resp: return json.loads(resp.read()) +def _is_long_enough(trace: dict) -> bool: + """True if the trace ran at least MIN_DURATION_S. + + Filters out trivially short traces (agent-card probes / health checks with a + null request) at the listing stage, so their spans are never fetched. A + trace with no reported duration is kept — we can't prove it is short. + """ + ms = trace.get("execution_time_ms") + if ms is None: + return True + return (ms / 1000.0) >= MIN_DURATION_S + + def fetch_trace_list() -> list[dict]: - """Fetch trace listing with pagination.""" - all_traces: list[dict] = [] + """Fetch traces within the time window, dropping too-short traces. + + The listing is paged newest-first (descending timestamp_ms), so we stop as + soon as we reach a trace older than CUTOFF_MS — older pages are never + requested. Traces shorter than MIN_DURATION_S are filtered here, before any + span download. + """ + kept: list[dict] = [] + skipped_short = 0 page_token = None - page_size = min(LIMIT, 500) + page_size = 500 + reached_cutoff = False - while len(all_traces) < LIMIT: + while not reached_cutoff: url = f"/api/2.0/mlflow/traces?experiment_ids={EXPERIMENT_ID}&max_results={page_size}" if page_token: url += f"&page_token={page_token}" response = mlflow_get(url) traces = response.get("traces", []) - all_traces.extend(traces) + for t in traces: + ts = t.get("timestamp_ms") + # Newest-first ordering: once older than the window, we are done. + if ts is not None and ts < CUTOFF_MS: + reached_cutoff = True + break + if _is_long_enough(t): + kept.append(t) + else: + skipped_short += 1 page_token = response.get("next_page_token") if not page_token or not traces: break - return all_traces[:LIMIT] + if skipped_short: + print( + f"Skipped {skipped_short} traces shorter than {MIN_DURATION_S}s " + f"(not fetched)", + file=sys.stderr, + ) + + return kept def transform_spans(raw_spans: list[dict]) -> list[dict]: @@ -136,7 +199,12 @@ def transform_spans(raw_spans: list[dict]) -> list[dict]: def main() -> int: - print(f"Fetching traces from MLflow (experiment_id={EXPERIMENT_ID}, limit={LIMIT})...", file=sys.stderr) + window_h = WINDOW_MS / 3600000.0 + print( + f"Fetching traces from MLflow (experiment_id={EXPERIMENT_ID}, " + f"window={window_h:g}h)...", + file=sys.stderr, + ) all_traces = fetch_trace_list() real_traces = [t for t in all_traces if t.get("status") in ("OK", "ERROR")]