Train a Custom Model in Azure ML and Call It from a Foundry Agent

Foundry's model catalog is broad, yet some jobs sit outside it. Imagine predicting churn from a company's own customer records, or scoring fraud against its own transaction history. Those are conventional machine learning tasks on private tabular data, and Microsoft's answer for them is Azure Machine Learning, a different product with a separate SDK, a workspace of its own and its own way of deploying models.

The build below covers the full path: train a model with the Azure ML SDK, register it, serve it from a managed endpoint and then expose that endpoint to a Foundry agent as a function tool. The agent can then call the custom model mid-conversation exactly as it would call any other tool.

Where the two platforms meet

The two products run on separate runtimes but belong to one ecosystem, and Microsoft has worked to keep the seam between them thin. Knowing exactly where it lies makes the build easier to follow. Four pieces are involved:

• Command job. Training runs here. The SDK is pointed at a script, a compute target and some inputs, and it executes the code as a tracked, reproducible job.
• Model registry. A trained model turns into a named, versioned artifact that exists apart from the job that produced it, so one version can be promoted to production without repeating the training run.
• Managed Online Endpoint. This serves predictions from a registered model. It's real infrastructure, always on or autoscaled, with its own URL, authentication and cost, and it lives inside the Azure ML workspace, not Foundry.
• Foundry FunctionTool. This is the bridge. Foundry neither knows nor cares that the function behind the tool calls an Azure ML endpoint. To the agent it's a name, a schema and a result.

Prerequisites

1. An Azure ML workspace and a Foundry project. They can sit in the same or different resource groups, because nothing here requires shared infrastructure.
2. Python 3.9 or newer with both SDKs:

pip install azure-ai-ml azure-ai-projects azure-identity

An MLClient authenticates and targets the workspace:

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>",
)

Training as a command job

The training step is a command job. It takes a training script, in this case an ordinary scikit-learn classifier, and executes 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)

The call to ml_client.jobs.stream waits and prints logs until training finishes. Prefer it in any script that will really run. A job that fails silently in the background is an unpleasant way to discover that a pipeline is broken.

Registering the model

Registration converts the job's output into a named, versioned artifact that can be referenced independently:

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.",
    )
)

Choosing MLFLOW_MODEL as the type is deliberate. If the training script logs the model through MLflow autologging, managed endpoints can deploy it using a built-in scoring container, with no custom score.py inference script. That saves real time, and the training script is best designed around it from the beginning.

Deploying behind an endpoint

An endpoint plus a deployment puts the model online:

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()

The deployment name blue isn't mandatory, but keep it, because it sets up the habit worth having at the first new model version. Create a green deployment next to blue, split traffic between them and shift fully once the new version earns trust. Overwriting blue and hoping for the best is the risky alternative.

Before wiring anything else, get a real prediction directly:

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)

Debugging a broken endpoint through an agent's function-calling loop is much harder than debugging it directly, so confirm it works first.

Connecting the endpoint to a Foundry agent

The agent never calls Azure ML itself. Application code does, when the agent asks for it. The setup declares a FunctionTool, creates the agent and defines the function that sends the HTTP request to the endpoint:

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,
    ),
)

Then a conversation runs, the model's function call is intercepted and the result is returned:

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)

The model decides whether a question warrants calling get_churn_score. When it does, application code makes the HTTP call to Azure ML, and the output returns to the same conversation so the model can finish its answer. Nothing in the agent definition or the Responses API calls depends on what the function does behind the scenes. It could query a database, invoke a Lambda function or reach a different Azure ML workspace, as it does here.

Two resources, two auth boundaries

One point stays hidden in the SDK docs. The managed endpoint in Azure ML and the project endpoint in Foundry are two resources behind two separate authentication walls, and nothing in this setup joins them. Azure ML has its own key or managed identity for the workspace, while Foundry holds a different credential. Declaring a FunctionTool establishes no trust between them; it only gives the model a name and schema to ask your code for something. All authentication toward Azure ML happens inside your own function, exactly as it would when calling any other outside API.

That's a strength. The Azure ML side can be owned, secured and rotated independently of Foundry, even by another team, and neither side needs write access to the other's resources.

Production notes

• A function call expires after 10 minutes. Output must go back to the conversation before the window closes, or the run fails. If the HTTP call to Azure ML is slow, treat that as a hard deadline and budget for it.
• Move to a managed identity beyond the prototype. A key stored in an environment variable is fine locally and a liability in anything unattended. Grant the identity running the tool code the AzureML Data Scientist role, or a narrower custom role scoped to this endpoint and not the whole workspace.
• Roll out new models gradually. Launch the new deployment next to the existing one and move traffic over in steps. That's the sole way to spot a bad retrain before every real request depends on it.
• Leave strict=True on. Strict schemas stop the model sending incomplete or malformed arguments to get_churn_score, which would otherwise blow up inside your function and not before the call.
• Idle compute costs money. A Managed Online Endpoint running one instance all day is billed even with no traffic, which a serverless model call isn't. For low or spiky volume, look at scale-to-zero or batch endpoints.
• Drift is silent. Nothing here retrains when real-world data moves away from the training distribution. Log the inputs and outputs of each get_churn_score call somewhere reviewable, and revisit training data on a schedule, not only after predictions visibly worsen.

Picking the right platform

For most generative AI work, Foundry's catalog is the right tool, and teams without proprietary training data may never need more. Once the question becomes a prediction about your own customers from your own history, Azure ML is the better home. Otherwise a language model gets stretched over a task it wasn't designed for, or serving infrastructure gets built by hand. The glue is small: one training job, one registered model, one managed endpoint and a single FunctionTool. The platforms stay separate, each doing the thing it's good at.