DEV Community

Cover image for Building and Serving a Custom Model with Azure ML, Then Wiring It Into a Foundry Agent
Jubin Soni
Jubin Soni Subscriber

Posted on

Building and Serving a Custom Model with Azure ML, Then Wiring It Into a Foundry Agent

The Foundry model catalog covers a lot of ground, but it doesn't cover everything. If you need a model trained on your own proprietary, tabular data, a churn predictor built on your actual customer history, a fraud score trained on your actual transaction patterns, that's not a model catalog problem. That's a real machine learning problem, and on Microsoft's stack it belongs to Azure Machine Learning, a genuinely separate platform from Foundry with its own SDK, its own workspace concept, and its own deployment model.

This is a hands-on build of the whole path: train a model with Azure ML's SDK, register it, deploy it behind a managed endpoint, and then wire that endpoint into a Foundry agent as a function tool, so a conversational agent can call your custom model mid-conversation the same way it would call any other tool.

The mental model first

Azure ML and Foundry don't share a runtime. They share an ecosystem, and a genuine amount of engineering effort has gone into making the seam between them small, but it's worth being precise about where that seam actually is:

  • A command job is how training happens. You point Azure ML's SDK at a script, a compute target, and a set of inputs, and it runs your training code as a tracked, reproducible job.
  • The model registry is where a trained model becomes a named, versioned artifact, independent of the job that produced it. This is what lets you promote a specific version to production without recreating the training run.
  • A Managed Online Endpoint is how a registered model actually serves predictions. It's a real, always-on (or autoscaling) piece of infrastructure with its own URL, its own auth, and its own cost, running inside the Azure ML workspace, not inside Foundry.
  • A Foundry agent's FunctionTool is the bridge. Foundry doesn't know or care that the function it's calling happens to invoke an Azure ML endpoint. From the agent's perspective, it's just a tool with a name, a schema, and a result.

From a training job to a tool call inside an agent turn

Prerequisites

  1. An Azure ML workspace and a Foundry project, in the same or different resource groups, it doesn't matter which, since nothing about this integration requires them to share infrastructure.
  2. Python 3.9+ with both SDKs installed.
pip install azure-ai-ml azure-ai-projects azure-identity
Enter fullscreen mode Exit fullscreen mode
from azure.ai.ml import MLClient
from azure.identity import DefaultAzureCredential

ml_client = MLClient(
    DefaultAzureCredential(),
    subscription_id="<subscription-id>",
    resource_group_name="<resource-group>",
    workspace_name="<aml-workspace-name>",
)
Enter fullscreen mode Exit fullscreen mode

Step 1: train the model

A command job wraps a training script, here a plain scikit-learn classifier, and runs it on managed compute.

from azure.ai.ml import command, Input, Output

train_job = command(
    code="./src",
    command="python train.py --data ${{inputs.training_data}} --model_output ${{outputs.model_output}}",
    inputs={"training_data": Input(type="uri_folder", path="azureml://datastores/workspaceblobstore/paths/churn-training/")},
    outputs={"model_output": Output(type="uri_folder")},
    environment="azureml://registries/azureml/environments/sklearn-1.5/labels/latest",
    compute="cpu-cluster",
    display_name="churn-model-training",
)

returned_job = ml_client.jobs.create_or_update(train_job)
ml_client.jobs.stream(returned_job.name)
Enter fullscreen mode Exit fullscreen mode

ml_client.jobs.stream blocks and prints logs until the job finishes, which is worth doing in any script you're actually going to run rather than fire-and-forget, since a training job failing silently in the background is a bad way to find out your pipeline is broken.

Step 2: register the trained model

Registration turns the job's output into a named, versioned artifact you can reference independently of the job.

from azure.ai.ml.entities import Model
from azure.ai.ml.constants import AssetTypes

model = ml_client.models.create_or_update(
    Model(
        path=f"azureml://jobs/{returned_job.name}/outputs/model_output",
        name="churn-classifier",
        type=AssetTypes.MLFLOW_MODEL,
        description="Customer churn classifier, trained on 18 months of account history.",
    )
)
Enter fullscreen mode Exit fullscreen mode

Using MLFLOW_MODEL as the type here isn't incidental. If your training script logs the model with MLflow's autologging, Azure ML's managed endpoints can deploy it with a built-in scoring container, no custom score.py inference script required. That's a real time saver worth designing your training script around from the start rather than discovering after the fact.

Step 3: deploy it behind a managed endpoint

from azure.ai.ml.entities import ManagedOnlineEndpoint, ManagedOnlineDeployment

endpoint = ManagedOnlineEndpoint(name="churn-endpoint", auth_mode="key")
ml_client.online_endpoints.begin_create_or_update(endpoint).result()

deployment = ManagedOnlineDeployment(
    name="blue",
    endpoint_name="churn-endpoint",
    model=model,
    instance_type="Standard_DS3_v2",
    instance_count=1,
)
ml_client.online_deployments.begin_create_or_update(deployment).result()

endpoint.traffic = {"blue": 100}
ml_client.online_endpoints.begin_create_or_update(endpoint).result()
Enter fullscreen mode Exit fullscreen mode

The blue deployment name isn't a convention you have to follow, but it's worth keeping, since it sets up the pattern you'll want the first time you deploy a new model version: create a green deployment alongside blue, split traffic between them, and shift fully once you trust the new version, rather than replacing blue outright and hoping.

Step 4: confirm it works before anything else touches it

import json

test_input = {"input_data": {"columns": ["tenure_months", "monthly_spend", "support_tickets"], "data": [[14, 89.50, 3]]}}

response = ml_client.online_endpoints.invoke(
    endpoint_name="churn-endpoint",
    request_file=None,
    deployment_name="blue",
    input_data=json.dumps(test_input),
)
print(response)
Enter fullscreen mode Exit fullscreen mode

Get a real prediction back here before wiring anything else into it. Debugging a broken endpoint through the extra layer of an agent's function-calling loop is meaningfully harder than debugging it directly.

Step 5: wrap the endpoint as a Foundry agent function tool

This is the actual bridge. The Foundry agent doesn't call Azure ML directly, your application code does, in response to the agent asking for it.

import os
import requests
from azure.ai.projects import AIProjectClient
from azure.ai.projects.models import PromptAgentDefinition, Tool, FunctionTool
from azure.identity import DefaultAzureCredential

def get_churn_score(tenure_months: int, monthly_spend: float, support_tickets: int) -> dict:
    payload = {"input_data": {"columns": ["tenure_months", "monthly_spend", "support_tickets"], "data": [[tenure_months, monthly_spend, support_tickets]]}}
    resp = requests.post(
        "https://churn-endpoint.<region>.inference.ml.azure.com/score",
        headers={"Authorization": f"Bearer {os.environ['AML_ENDPOINT_KEY']}", "Content-Type": "application/json"},
        json=payload,
        timeout=10,
    )
    resp.raise_for_status()
    return {"churn_probability": resp.json()[0]}

func_tool = FunctionTool(
    name="get_churn_score",
    description="Predict churn probability for a customer given tenure, spend, and support ticket history.",
    parameters={
        "type": "object",
        "properties": {
            "tenure_months": {"type": "integer", "description": "How many months the customer has been active."},
            "monthly_spend": {"type": "number", "description": "Average monthly spend in dollars."},
            "support_tickets": {"type": "integer", "description": "Number of support tickets in the last 90 days."},
        },
        "required": ["tenure_months", "monthly_spend", "support_tickets"],
        "additionalProperties": False,
    },
    strict=True,
)

project = AIProjectClient(endpoint=os.environ["FOUNDRY_PROJECT_ENDPOINT"], credential=DefaultAzureCredential())
tools: list[Tool] = [func_tool]

agent = project.agents.create_version(
    agent_name="retention-agent",
    definition=PromptAgentDefinition(
        model="gpt-4.1-mini",
        instructions="Help the team assess churn risk. Call get_churn_score whenever specific customer numbers are provided.",
        tools=tools,
    ),
)
Enter fullscreen mode Exit fullscreen mode

Step 6: run it end to end

import json
from openai.types.responses.response_input_param import FunctionCallOutput

openai_client = project.get_openai_client()
conversation = openai_client.conversations.create()

response = openai_client.responses.create(
    input="A customer's been with us 14 months, spends about $90/month, and filed 3 tickets recently. Churn risk?",
    conversation=conversation.id,
    extra_body={"agent_reference": {"name": agent.name, "type": "agent_reference"}},
)

for item in response.output:
    if item.type == "function_call" and item.name == "get_churn_score":
        result = get_churn_score(**json.loads(item.arguments))
        follow_up = openai_client.responses.create(
            input=[FunctionCallOutput(type="function_call_output", call_id=item.call_id, output=json.dumps(result))],
            conversation=conversation.id,
            extra_body={"agent_reference": {"name": agent.name, "type": "agent_reference"}},
        )
        print(follow_up.output_text)
Enter fullscreen mode Exit fullscreen mode

The model decides whether the question warrants calling get_churn_score at all, your code executes the actual HTTP call to Azure ML when it does, and the result flows back into the same conversation for the model to finish answering with. Nothing about the agent definition or the Responses API calls changes based on what the function actually does behind the scenes, whether it hits a database, a Lambda function, or, as here, a completely separate Azure ML workspace.

Where this fits in the bigger picture

Two workspaces, two endpoints, one function call in between

Worth being explicit about something the SDKs don't make obvious on their own: Azure ML's managed endpoint and the Foundry project endpoint are two different resources with two different auth boundaries, and nothing about this integration merges them. The Azure ML endpoint has its own key or managed identity, scoped to that workspace. The Foundry project has its own credential, scoped separately. FunctionTool doesn't create a trust relationship between the two, it's just a name and a schema the model uses to ask your code to do something. The actual authentication to Azure ML happens entirely in your own function implementation, same as it would if that function called any other external API.

That's a feature, not a limitation. It means the Azure ML side of this can be owned, secured, and rotated independently of the Foundry side, by a different team if that's how your organization is structured, without either side needing write access to the other's resources.

Production considerations before you commit

  • Process expires 10 minutes after a function call is issued. Submit the tool's output back to the conversation before that window closes, or the run fails. If get_churn_score's HTTP call to Azure ML is slow, that's a hard deadline, not a soft one, budget for it.
  • Prefer a managed identity over a static endpoint key once this is more than a prototype. A key baked into an environment variable is fine for local testing and a real liability in anything that runs unattended. Grant the identity running your function tool code the AzureML Data Scientist or a more narrowly scoped custom role against just this endpoint, not the whole workspace.
  • Blue/green the endpoint, don't overwrite it. Standing up a new deployment alongside the old one and shifting traffic gradually is the only way to catch a regression in a newly trained model before it's serving 100% of real requests.
  • Treat strict=True on the FunctionTool schema as a correctness feature, not boilerplate. It's what keeps the model from calling get_churn_score with a malformed or partial argument set that would otherwise fail inside your function rather than being caught before the call.
  • Online endpoint compute is billed whether or not it's actively scoring. Unlike a serverless model call, a Managed Online Endpoint with instance_count=1 running around the clock costs money at idle. If call volume is low and bursty, look at scale-to-zero options or batch endpoints instead of defaulting to always-on instance count 1.
  • Model drift doesn't announce itself. Nothing in this pipeline retrains automatically when the real-world distribution shifts away from what the model was trained on. Log the inputs and outputs of every get_churn_score call somewhere you can actually review, and revisit training data on a real cadence, not only when someone notices predictions have gotten worse.

Where this leaves you

Foundry's model catalog is the right tool for the overwhelming majority of generative AI work, and for teams that never need a model trained on their own proprietary data, it's genuinely the whole story. But the moment the problem is "predict something specific about our own customers, from our own historical data," that's Azure ML's job, and pretending otherwise means either forcing a language model to approximate a task it was never built for, or building your own training and serving infrastructure by hand. The actual integration effort here is small: a training job, a registered model, a managed endpoint, and one FunctionTool definition. The two platforms don't need to merge for that to work. They just both need to keep doing the one thing each is actually good at.

References

  1. Microsoft Learn. "Train models with the Python SDK v2." Azure Machine Learning. learn.microsoft.com/en-us/azure/machine-learning/how-to-train-model?view=azureml-api-2
  2. Microsoft Learn. "Deploy machine learning models to online endpoints." Azure Machine Learning. learn.microsoft.com/en-us/azure/machine-learning/how-to-deploy-online-endpoints?view=azureml-api-2
  3. Microsoft Learn. "Use function calling with Microsoft Foundry agents." learn.microsoft.com/en-us/azure/foundry/agents/how-to/tools/function-calling
  4. Microsoft Learn. "Explore Microsoft Foundry Models in Azure Machine Learning." learn.microsoft.com/en-us/azure/machine-learning/foundry-models-overview?view=azureml-api-2
  5. Microsoft Learn. "Get started with Microsoft Foundry SDKs and endpoints." learn.microsoft.com/en-us/azure/foundry/how-to/develop/sdk-overview

Top comments (0)