Treating LLM Inference Like a First-Class MLOps Citizen: MLflow + Azure ML + OpenAI

How to bring experiment tracking discipline to generative AI — with token-level observability baked in

The Problem Nobody Talks About in GenAI

When you train a traditional ML model, the MLOps toolbox is well-established: log your hyperparameters, track your loss curves, version your artifacts, compare runs. Tools like MLflow were built for exactly this.

But the moment you switch to a generative AI workload — calling GPT-4 in a loop, chaining prompts, building a chatbot — most teams throw that discipline out the window. API calls happen in isolation, token costs are a mystery until the bill arrives, and debugging a bad response means staring at logs hoping to spot the culprit prompt.

This post walks through a pattern that closes that gap: connecting MLflow to Azure Machine Learning as a remote tracking server, then instrumenting every OpenAI API call to log params, metrics, and latency — automatically, per request.

The Architecture in One Sentence

Your Python notebook calls OpenAI → measures latency and counts tokens via TikToken → logs everything to MLflow → which persists to Azure ML Studio for visualization and comparison.

No custom dashboards. No manual exports. Just structured, reproducible observability.

Step 1: Stand Up the Azure ML Workspace

Before writing a line of Python, you need an Azure ML workspace. This becomes the remote backend for MLflow — think of it as your experiment database in the cloud.

In the Azure Portal:

1. Search for Azure Machine Learning → Create
2. Name it something meaningful (mlops-ws-mlflow-demo)
3. Assign it to a resource group (rg-mlops-demo)
4. Region: East US (or closest to you)
5. Hit Review + Create — takes about 2–3 minutes

Once provisioned, go to Overview and copy the MLflow tracking URI. It looks like:

azureml://eastus.api.azureml.ms/mlflow/v1.0/subscriptions/<sub-id>/resourceGroups/<rg>/providers/Microsoft.MachineLearningServices/workspaces/<workspace>

This single URI is how your local notebook talks to Azure’s tracking backend.

Step 2: Install Dependencies

pip install azure-ai-ml azureml-mlflow azure-identity openai tiktoken colorama

Four packages do the heavy lifting:

Package Role azure-ai-ml Azure ML Python SDK — fetches workspace config azureml-mlflow MLflow plugin that speaks the Azure ML tracking protocol azure-identity Handles browser-based auth to your Azure subscription tiktoken OpenAI's official tokenizer — exact token counts, no guessing

Step 3: Authenticate and Connect MLflow to Azure

from azure.identity import InteractiveBrowserCredential
from azure.ai.ml import MLClient
import mlflow

credential = InteractiveBrowserCredential()

ml_client = MLClient(
    credential=credential,
    subscription_id="<your-subscription-id>",
    resource_group_name="rg-mlops-demo",
    workspace_name="mlops-ws-mlflow-demo"
)

# Pull the MLflow URI directly from the workspace object
mlflow_tracking_uri = ml_client.workspaces.get(
    ml_client.workspace_name
).mlflow_tracking_uri

mlflow.set_tracking_uri(mlflow_tracking_uri)
mlflow.set_experiment("genai_mlflow_demo")

InteractiveBrowserCredential() opens a browser tab for Microsoft login — no service principals, no secrets to manage locally. Clean and sufficient for development.

After this cell runs, every mlflow.log_* call in your session will write to Azure ML Studio, not your local filesystem.

Step 4: Why TikToken (and Not Just Trusting the API Response)

Before building the core function, a word on token counting.

OpenAI’s API response does include token usage in response.usage. So why use TikToken separately?

Three reasons:

1. Pre-request cost estimation. You can count tokens before sending the request and decide whether to truncate a prompt that would exceed budget thresholds.
2. Conversation-level accounting. In a multi-turn chat, the full conversation history is sent every time. TikToken lets you track the growing conversation_tokens count — the real cost driver in long sessions.
3. Exact match to billing. TikToken uses the same cl100k_base encoding as gpt-3.5-turbo and gpt-4. Your counts match what OpenAI charges — no estimates, no surprises.

import tiktoken as tk

def count_tokens(string: str, encoding_name: str = "cl100k_base") -> int:
    encoding = tk.get_encoding(encoding_name)
    return len(encoding.encode(string, disallowed_special=()))

Quick sanity check on token counts for familiar phrases:

Text                                                              Tokens
------------------------------------------------------------------------
Can you answer medical question?                                       6
Is ADHD that serious?                                                  5
Tell me a detailed story about a boy...                               15
What is the meaning of life, the universe, and everything in it?      15

Intuitive and exact — “Is ADHD that serious?” is 5 tokens, not 5 words (the ? merges with the adjacent token in GPT's vocabulary).

Step 5: The Instrumented generate_text() Function

This is the core of the pattern. Every OpenAI call goes through this wrapper — which measures latency, counts tokens, and logs everything to MLflow before returning the response.

import time
import mlflow
from openai import OpenAI

client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

MODEL             = "gpt-3.5-turbo"
TEMPERATURE       = 0.7
TOP_P             = 1
FREQUENCY_PENALTY = 0
PRESENCE_PENALTY  = 0
MAX_TOKENS        = 800

def generate_text(conversation: list, max_tokens: int = MAX_TOKENS) -> str:
    start_time = time.time()

    response = client.chat.completions.create(
        model=MODEL,
        messages=conversation,
        temperature=TEMPERATURE,
        max_tokens=max_tokens,
        top_p=TOP_P,
        frequency_penalty=FREQUENCY_PENALTY,
        presence_penalty=PRESENCE_PENALTY
    )

    latency          = time.time() - start_time
    message_response = response.choices[0].message.content

    # Token accounting
    prompt_tokens       = count_tokens(conversation[-1]['content'])
    conversation_tokens = count_tokens(str(conversation))
    completion_tokens   = count_tokens(message_response)

    # Log metrics to Azure MLflow
    mlflow.log_metrics({
        "request_latency":     latency,
        "prompt_tokens":       prompt_tokens,
        "completion_tokens":   completion_tokens,
        "conversation_tokens": conversation_tokens,
        "request_count":       1
    })

    # Log params to Azure MLflow
    mlflow.log_params({
        "model":             MODEL,
        "temperature":       TEMPERATURE,
        "top_p":             TOP_P,
        "frequency_penalty": FREQUENCY_PENALTY,
        "presence_penalty":  PRESENCE_PENALTY,
        "max_tokens":        max_tokens
    })

    return message_response

What gets logged per call:

Type Key What it tells you Param model Which GPT variant was used Param temperature Sampling randomness Param max_tokens Response ceiling Metric request_latency Wall-clock seconds for the API round-trip Metric prompt_tokens Tokens in the latest user message Metric completion_tokens Tokens in the model's response Metric conversation_tokens Cumulative context window usage Metric request_count Step counter for the run

Step 6: The Load Test Loop

With the wrapper in place, running a 10-iteration conversational load test is straightforward:

test_inputs = [
    "Hello, how are you?",
    "What is the capital of France?",
    "Tell me a dad joke",
    "Tell me a short story",
    "What is the meaning of life?",
    "What is the largest mammal?",
    "What is the square root of 144?",
    "Give me a productivity tip",
    "Explain quantum computing in one sentence",
    "What is machine learning?"
]

with mlflow.start_run():
    conversation = [
        {"role": "system", "content": "You are a helpful assistant."},
    ]

    for i in range(10):
        user_input = random.choice(test_inputs)
        conversation.append({"role": "user", "content": user_input})

        ai_output = generate_text(conversation, MAX_TOKENS)
        conversation.append({"role": "assistant", "content": ai_output})

        print(f"Iteration {i+1}/10")
        print(f"User: {user_input}")
        print(f"AI:   {ai_output[:100]}...")

The with mlflow.start_run(): context manager groups all 10 iterations into a single MLflow run. Each call to generate_text() logs a new step within that run — which is what produces the time-series charts in Azure ML Studio.

Step 7: What You See in Azure ML Studio

Once the loop completes, Azure ML Studio shows two completed runs under the genai_mlflow_demo experiment.

The Overview tab shows the logged params at a glance: model: gpt-3.5-turbo, temperature: 0.7, max_tokens: 800 — exactly what was configured in code.

The Metrics tab renders all five tracked signals as interactive time-series charts across the 10 steps:

Reading the charts:

• **conversation_tokens** grows monotonically — this is the compound cost of maintaining conversation history. By step 9, the model is processing 657 tokens of context to answer each new question, even if that question is just "What is the capital of France?"
• **completion_tokens** spikes at steps 1 and 7 — longer questions ("Tell me a short story", "What is the meaning of life?") elicit more verbose responses.
• **request_latency** oscillates between 0.5s and 2.8s — correlating with response length, not prompt complexity.
• **prompt_tokens** stays flat around 5–7 — individual user messages are short; the cost is in the accumulated context.
• **request_count** is constant at 1 per step — confirming the logging step counter works correctly.

These patterns are immediately actionable. If you’re running a production chatbot, conversation_tokens growing to 657 by turn 9 is a signal to implement context summarization or a sliding window before costs compound further.

The Metrics Defined

Metric Definition Completion tokens Tokens generated by the model in the response Conversation tokens Total tokens in the full conversation context sent per request Prompt tokens Tokens in the most recent user message only Request latency Wall-clock seconds for one API round-trip Request count Incremental counter — one per call

Why This Matters Beyond the Demo

This pattern implements something that most GenAI teams skip: inference-time MLOps.

Traditional MLOps focuses on training — experiment tracking, model versioning, deployment pipelines. But for LLM applications, the model is fixed. What you’re engineering is the prompt strategy, context management, and cost profile — and those need the same rigor as model training.

By treating every OpenAI call as a logged, versioned, comparable experiment you get:

• Cost attribution — know exactly which prompt templates or conversation patterns are expensive
• Latency benchmarking — compare model versions or temperature settings head-to-head
• Reproducibility — params are logged alongside outputs; you can reconstruct any run
• Governance — centralized tracking under Azure ML’s enterprise security and RBAC model
• Prompt versioning — extend log_params to include system prompt hashes and iterate systematically

Conclusion

MLflow was designed for traditional ML training loops. But its primitives — params, metrics, runs, experiments — map cleanly onto LLM inference behavior. Connecting it to Azure ML as a remote backend gives you persistence, visualization, and enterprise-grade access control with minimal additional code.

The key insight: every API call is an experiment. Treat it that way, and the path from prototype chatbot to production GenAI system becomes observable, auditable, and — critically — debuggable when it inevitably behaves unexpectedly.

As generative AI moves deeper into enterprise workflows, this pattern scales naturally into full GenAI-Ops practices: prompt versioning, cost attribution, performance benchmarking, and governance under enterprise security standards. MLflow paired with Azure ML provides a robust foundation for building responsible, observable, and production-ready AI systems.

The full notebook is available on GitHub. If you found this useful, follow for more posts on applied MLOps and generative AI engineering.

Tags: MLflow Azure Machine Learning OpenAI MLOps GenAI LLMOps Python Observability