aoai_finetune.py

# Copyright (c) Microsoft. All rights reserved.

"""The Azure OpenAI fine-tuning algorithm implementation."""

import asyncio
import copy
import json
import logging
import os
import random
import subprocess
import tempfile
import time
from typing import Any, Dict, List, Optional, Sequence, Tuple

import requests
from openai import OpenAI

from agentlightning.adapter.messages import OpenAIMessages, TraceToMessages
from agentlightning.algorithm import Algorithm
from agentlightning.algorithm.utils import batch_iter_over_dataset
from agentlightning.reward import find_final_reward
from agentlightning.types import LLM, RolloutMode, TaskInput

logger = logging.getLogger("agentlightning.aoai")

ROLLOUT_IDLE_SLEEP_SECONDS = 5.0
FILE_STATUS_POLL_INTERVAL = 10
FINETUNE_JOB_POLL_INTERVAL = 60


class AzureOpenAIFinetune(Algorithm):
    """Coordinate iterative fine-tuning runs for an Azure OpenAI deployment.

    The algorithm batches rollouts, extracts the recorded traces, converts them into JSONL records
    that comply with Azure OpenAI fine-tuning, and optionally redeploys the resulting checkpoint so
    subsequent rollouts evaluate the newest model revision.
    """

    def __init__(
        self,
        base_deployment_name: str,
        finetuned_deployment_name: str,
        base_model_name: str,
        *,
        finetune_every_n_rollouts: int = 32,
        azure_openai_endpoint: Optional[str] = None,
        azure_openai_api_key: Optional[str] = None,
        azure_openai_api_version: Optional[str] = None,
        subscription_id: Optional[str] = None,
        resource_group: Optional[str] = None,
        resource_name: Optional[str] = None,
        seed: int = 42,
        n_iterations: int = 3,
        finetune_epochs: int = 1,
        finetune_batch_size: int = 2,
        finetune_learning_rate: float = 1.0,
        max_deployments: int = 2,
        data_filter_ratio: float = 0.5,
    ) -> None:
        """Create a fine-tuning workflow tied to an Azure OpenAI endpoint.

        Args:
            base_deployment_name: Deployment used as the base model for the first fine-tuning job.
            deployment_name: Deployment that should serve the fine-tuned weights after each round.
                Currently, this name is only used as a prefix for the actual deployment created after
                each fine-tuning job, because multiple versions cannot be assigned to the same deployment.
            base_model_name: On Azure, deployments are instantiated from base models
                (e.g., "gpt-4.1-mini" deployment is created from "gpt-4.1-mini-2025-04-14").
                This name is used to identify the latter name when launching fine-tuning jobs.
            finetune_every_n_rollouts: Number of rollouts grouped together before launching a job.
                We don't recommend setting this value too low as fine-tuning jobs have a minimum rows requirement.
            azure_openai_endpoint: Azure OpenAI endpoint (e.g. `https://{resource}.openai.azure.com`).
            azure_openai_api_key: API key with access to the Azure OpenAI resource.
            azure_openai_api_version: API version to use when talking to Azure OpenAI.
            subscription_id: Azure subscription that owns the OpenAI resource (used for deployment).
            resource_group: Resource group of the target Azure OpenAI resource.
            resource_name: Azure OpenAI resource name, usually the Azure OpenAI resource name.
            seed: Random seed forwarded to the fine-tuning job for reproducibility.
            n_iterations: Number of algorithm iterations (fine-tune → deploy → evaluate).
            finetune_epochs: Number of epochs per fine-tuning job (not the number of epochs to go through `train_dataset`).
            finetune_batch_size: Batch size to use for the fine-tuning job.
            finetune_learning_rate: Learning rate to use for the fine-tuning job.
            max_deployments: Maximum number of deployments to keep active; older ones are deleted.
                Use this to avoid hitting the capacity limit on Azure service.
            data_filter_ratio: Fraction of high-reward examples to keep when preparing JSONL data.
        """
        super().__init__()

        self.azure_openai_endpoint = azure_openai_endpoint or os.getenv("AZURE_OPENAI_ENDPOINT", "")
        if not self.azure_openai_endpoint:
            raise ValueError("Azure OpenAI endpoint must be provided via parameter or AZURE_OPENAI_ENDPOINT env var")

        self.azure_openai_api_key = azure_openai_api_key or os.getenv("AZURE_OPENAI_API_KEY", "")
        if not self.azure_openai_api_key:
            raise ValueError("Azure OpenAI API key must be provided via parameter or AZURE_OPENAI_API_KEY env var")

        self.azure_openai_api_version = azure_openai_api_version or os.getenv("AZURE_OPENAI_API_VERSION", "")
        if not self.azure_openai_api_version:
            raise ValueError(
                "Azure OpenAI API version must be provided via parameter or AZURE_OPENAI_API_VERSION env var"
            )

        self.subscription_id = subscription_id or os.getenv("AZURE_SUBSCRIPTION_ID", "")
        if not self.subscription_id:
            raise ValueError("Azure subscription ID must be provided via parameter or AZURE_SUBSCRIPTION_ID env var")

        self.resource_group = resource_group or os.getenv("AZURE_RESOURCE_GROUP", "")
        if not self.resource_group:
            raise ValueError("Azure resource group must be provided via parameter or AZURE_RESOURCE_GROUP env var")

        self.resource_name = resource_name or os.getenv("AZURE_RESOURCE_NAME", "")
        if not self.resource_name:
            raise ValueError("Azure resource name must be provided via parameter or AZURE_RESOURCE_NAME env var")

        self.base_deployment_name = base_deployment_name
        self.finetuned_deployment_name = finetuned_deployment_name
        self.base_model_name = base_model_name

        self.finetune_every_n_rollouts = finetune_every_n_rollouts
        self.seed = seed
        self.n_iterations = n_iterations
        self.finetune_epochs = finetune_epochs
        self.finetune_batch_size = finetune_batch_size
        self.finetune_learning_rate = finetune_learning_rate
        self.max_deployments = max_deployments
        self.data_filter_ratio = data_filter_ratio

        self.openai_client = OpenAI(
            api_key=self.azure_openai_api_key,
            base_url=self.azure_openai_endpoint,
        )

        # Tracks the deployments created. They can be deleted later if needed.
        self._created_deployments: List[str] = []
        self._log_prefix: str = ""

    async def run(  # type: ignore
        self,
        train_dataset: Optional[List[TaskInput]] = None,
        val_dataset: Optional[List[TaskInput]] = None,
    ) -> None:
        """
        Run the training loop.

        Args:
            train_dataset: Optional training dataset
            val_dataset: Optional validation dataset
        """
        if train_dataset is None or val_dataset is None:
            raise ValueError("Both train_dataset and val_dataset must be provided")

        resources: LLM = LLM(endpoint=self.azure_openai_endpoint, model=self.base_deployment_name)
        store = self.get_store()

        # This tracks the model name used in training
        # It's different from the deployment name which used for inference
        training_model_name: str = self.base_model_name

        data_iterator = batch_iter_over_dataset(train_dataset, self.finetune_every_n_rollouts)
        for i_iteration in range(self.n_iterations):
            self._log_prefix = f"[AOAI FT {i_iteration + 1}/{self.n_iterations}] "
            # (1) Fetch the next batch of tasks to process
            tasks = next(data_iterator)
            self._log_info(f"[Stage 1] Starting fine-tuning iteration with {len(tasks)} tasks...")

            # (2) Update the current active LLM deployment address
            await store.add_resources({"main_llm": resources})
            self._log_info(f"[Stage 2] Using model deployment: {resources.model}")

            # (3) Spawn and wait for the rollouts to complete
            messages_group, reward_group = await self.batch_rollout_and_collect_data(tasks, "train")
            self._log_info(f"[Stage 3] Completed rollouts for {len(tasks)} tasks.")

            # (4) Filter the data based on rewards
            training_data = await self.prepare_data_for_training(messages_group, reward_group, "train")
            self._log_info(f"[Stage 4] Prepared {len(training_data)} training examples after filtering.")

            # (5) Perform fine-tuning
            self._log_info(f"[Stage 5] Starting fine-tuning for model {training_model_name}...")
            training_model_name = self.finetune(training_data, training_model_name, i_iteration)
            self._log_info(f"[Stage 5] Fine-tuning completed. Updated training model base name: {training_model_name}")

            # (6) Deploy the fine-tuned model
            self._log_info(f"[Stage 6] Deploying fine-tuned model...")
            resources = self.deploy_finetuned_model(training_model_name, i_iteration + 1)
            self._log_info(f"[Stage 6] Deployment completed. Updated resources to: {resources}")

            # (7) Evaluate on validation dataset
            self._log_info(f"[Stage 7] Evaluating on validation dataset...")
            _, val_reward_group = await self.batch_rollout_and_collect_data(val_dataset, "val")
            self._log_info(
                f"[Stage 7] Evaluation completed. Average reward: {sum(val_reward_group) / len(val_reward_group):.4f}"
            )

    async def batch_rollout_and_collect_data(
        self,
        tasks: Sequence[TaskInput],
        rollout_mode: RolloutMode = "train",
    ) -> Tuple[List[OpenAIMessages], List[float]]:
        """Launch rollouts for a batch of tasks and aggregate their traces.

        Each task is executed concurrently and the resulting spans are converted into OpenAI-style
        chat messages. Rewards from the traces are preserved so downstream filtering can prefer the
        highest quality examples.

        Args:
            tasks: Rollout payloads collected from the dataset.
            rollout_mode: Semantic label that differentiates training from validation passes.

        Returns:
            Tuple containing the flattened list of OpenAI messages and the aligned list of rewards.
        """
        if not tasks:
            return [], []

        results = await asyncio.gather(*(self.rollout_and_collect_data(task, mode=rollout_mode) for task in tasks))

        messages_group: List[OpenAIMessages] = []
        reward_group: List[float] = []

        for messages_list, reward in results:
            if not messages_list:
                continue
            messages_group.extend(messages_list)
            # Duplicate the reward for each message set produced by the rollout
            reward_group.extend([reward] * len(messages_list))

        return messages_group, reward_group

    async def rollout_and_collect_data(self, task: TaskInput, mode: RolloutMode) -> Tuple[List[OpenAIMessages], float]:
        """Execute a single rollout, returning OpenAI messages together with the final reward.

        The method waits for the rollout to enter a terminal state, retrieves the recorded spans,
        converts them into OpenAI chat messages using the configured trace adapter, and extracts the
        reward emitted by the runner.

        Args:
            task: Rollout payload to enqueue in the store.
            mode: Execution mode to annotate the rollout (`"train"`, `"val"` or `"test"`).

        Returns:
            A tuple containing the list of OpenAI messages reconstructed from the trace and the
            numeric reward associated with the rollout. Rewards default to `0.0` when not found.
        """
        store = self.get_store()
        rollout = await store.enqueue_rollout(input=task, mode=mode)
        rollout_id = rollout.rollout_id

        self._log_debug("Waiting for rollout %s to finish in mode=%s", rollout_id, mode)

        while True:
            completed = await store.wait_for_rollouts(rollout_ids=[rollout_id], timeout=0.0)
            if completed:
                finished = completed[0]
                if finished.status != "succeeded":
                    self._log_error(f"Rollout {rollout_id} finished with status {finished.status}. Skipping.")
                break
            await asyncio.sleep(ROLLOUT_IDLE_SLEEP_SECONDS)

        spans = await store.query_spans(rollout_id=rollout_id, attempt_id="latest")

        try:
            adapter = self.get_adapter()
        except ValueError:
            adapter = TraceToMessages()
            self.set_adapter(adapter)
        if not isinstance(adapter, TraceToMessages):
            raise RuntimeError(
                "The adapter is configured but not a TraceToMessages adapter. "
                "AzureOpenAIFinetune requires a TraceToMessages adapter. Please set that in Trainer."
            )

        messages_list = adapter.adapt(spans)
        if not messages_list:
            self._log_error(f"Rollout {rollout_id} produced no OpenAI messages for training.")

        # NOTE: Patch the messages list for AOAI requirements
        # This should ideally be merged into message adapter
        for messages in messages_list:
            for message in messages["messages"]:
                if "content" in message and message["content"] is None:
                    message.pop("content")

        reward = find_final_reward(spans)
        if reward is None:
            self._log_error(f"Rollout {rollout_id} produced no reward; defaulting to 0.0.")
            reward = 0.0

        self._log_info("Rollout %s produced %d message set(s) with reward %.3f", rollout_id, len(messages_list), reward)

        return messages_list, reward

    async def prepare_data_for_training(
        self,
        messages_group: List[OpenAIMessages],
        reward_group: List[float],
        split: RolloutMode,
    ) -> List[Dict[str, Any]]:
        """Combine rollouts and rewards into JSONL training payloads.

        Args:
            messages_group: Flattened list of OpenAI message dictionaries.
            reward_group: Rewards aligned with `messages_group` entries.
            split: Dataset split that produced the examples (e.g., `"train"` or `"val"`).

        Returns:
            JSON-serializable dictionaries ready to be written into a fine-tuning file.
        """
        if len(messages_group) != len(reward_group):
            raise ValueError("Mismatch between number of message entries and reward entries.")

        tagged_examples: List[Dict[str, Any]] = []
        for idx, (messages, reward) in enumerate(zip(messages_group, reward_group)):
            example: Dict[str, Any] = {
                "messages": messages["messages"],
                "metadata": {"split": split, "rollout_index": idx},
                "reward": reward,
                "reward_jitter": random.uniform(0, 1),
            }
            if messages.get("tools"):
                example["tools"] = messages["tools"]
            tagged_examples.append(example)

        self._log_info(
            "Collected %d candidate example(s) for split=%s before filtering (ratio=%.2f).",
            len(tagged_examples),
            split,
            self.data_filter_ratio,
        )

        filtered_examples = self._filter_training_data(tagged_examples)

        self._log_info("Keeping %d example(s) for fine-tuning after reward-based filtering.", len(filtered_examples))

        return filtered_examples

    def finetune(self, training_data: List[Dict[str, Any]], base_model: str, iteration_idx: int) -> str:
        """Launch a fine-tuning job on Azure OpenAI using the supplied dataset.

        Args:
            training_data: JSONL-ready records that describe the conversation transcripts.
            iteration_idx: Current iteration index.

        Returns:
            Identifier of the fine-tuned model produced by Azure OpenAI.
        """
        if not training_data:
            raise ValueError("Training data must not be empty before launching fine-tuning.")
        if not self.openai_client:
            raise RuntimeError("Azure OpenAI client is not initialized; cannot fine-tune.")

        next_iteration = iteration_idx + 1

        train_file_path: Optional[str] = None
        try:
            with tempfile.NamedTemporaryFile(
                mode="w", prefix=f"{base_model}_{iteration_idx:02d}_", suffix=".jsonl", delete=False
            ) as handle:
                for record in training_data:
                    handle.write(json.dumps(record) + "\n")
                train_file_path = handle.name

            self._log_info(
                "Prepared temporary training file %s with %d example(s).", train_file_path, len(training_data)
            )

            with open(train_file_path, "rb") as file_handle:
                training_response = self.openai_client.files.create(file=file_handle, purpose="fine-tune")
            train_file_id = training_response.id
            self._log_info("Uploaded training file to Azure OpenAI (file_id=%s).", train_file_id)

            self._wait_for_file_processed(train_file_id)

            job = self.openai_client.fine_tuning.jobs.create(
                training_file=train_file_id,
                model=base_model,
                seed=self.seed,
                method={
                    "type": "supervised",
                    "supervised": {
                        "hyperparameters": {
                            "batch_size": self.finetune_batch_size,
                            "learning_rate_multiplier": self.finetune_learning_rate,
                            "n_epochs": self.finetune_epochs,
                        }
                    },
                },
                # TODO: continuously adding suffix will make model names very long after a few iterations
                # investigate if we can just specify the fine-tuned model name directly
                suffix=f"v{next_iteration:02d}",
                # NOTE: https://learn.microsoft.com/en-us/azure/ai-foundry/openai/how-to/fine-tuning
                # Other options are "GlobalStandard" and "Standard"
                extra_body={"trainingType": "GlobalStandard"},
            )
            job_id = job.id
            self._log_info("Fine-tuning job %s created for base model %s.", job_id, base_model)

            fine_tuned_model = self._wait_for_finetuning(job_id)
            if not fine_tuned_model:
                raise RuntimeError(f"Fine-tuning job {job_id} finished without producing a model id.")

            self._log_info("Fine-tuning job %s succeeded with new model id %s.", job_id, fine_tuned_model)

            return fine_tuned_model

        finally:
            if train_file_path and os.path.exists(train_file_path):
                try:
                    os.unlink(train_file_path)
                except OSError:
                    self._log_warning("Failed to remove temporary training file %s.", train_file_path)

    def deploy_finetuned_model(self, finetuned_model_id: str, iteration_idx: int) -> LLM:
        """Deploy the fine-tuned checkpoint and return an `LLM` resource descriptor.

        Args:
            finetuned_model_id: Identifier returned by the fine-tuning job.
            iteration_idx: Current iteration index.

        Returns:
            `LLM` resource pointing to either the Azure deployment or the direct model id.
        """
        if not finetuned_model_id:
            raise ValueError("finetuned_model_id must be a non-empty string.")

        while len(self._created_deployments) >= self.max_deployments:
            self._log_warning(
                "Maximum number of deployments reached (%d). Cleaning up old deployments.", self.max_deployments
            )
            oldest_deployment = self._created_deployments.pop(0)
            self._log_info("Deleting old deployment %s.", oldest_deployment)
            self._delete_deployment(oldest_deployment)

        if self.subscription_id and self.resource_group and self.resource_name:
            # version should be like this: str(iteration_idx)
            # Because of this issue: {"code":"ModelUpgradeNotSupported","message":"Model updates are not supported for finetuned model deployments."}
            # We need to concatenate the version to the model name
            # and version is always "1"
            deployment_name = f"{self.finetuned_deployment_name}_v{iteration_idx:02d}"
            self._deploy_model(finetuned_model_id, deployment_name, "1")
            self._wait_for_deployment_ready(deployment_name, "1")
            self._created_deployments.append(deployment_name)
            self._log_info(
                "Deployed fine-tuned model %s to deployment %s. We now have %d active deployments.",
                finetuned_model_id,
                deployment_name,
                len(self._created_deployments),
            )
        else:
            raise RuntimeError("Azure deployment parameters missing; using fine-tuned model id directly.")

        return LLM(endpoint=self.azure_openai_endpoint, model=deployment_name, api_key=self.azure_openai_api_key)

    def cleanup_deployments(self) -> None:
        """Delete all deployments created by this algorithm instance."""
        for deployment_name in self._created_deployments:
            self._log_info("Cleaning up deployment %s.", deployment_name)
            self._delete_deployment(deployment_name)
        self._created_deployments = []

    def _filter_training_data(self, data: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """Select the top-performing examples and strip reward metadata.

        Args:
            data: Candidate training examples carrying a temporary `reward` key.

        Returns:
            List of examples suitable for JSONL serialization (without the `reward` field).
        """
        if not data:
            return []

        if self.data_filter_ratio >= 1.0:
            selected = data
        else:
            sorted_data = sorted(data, key=lambda x: (x.get("reward", 0.0), x.get("reward_jitter", 0.0)), reverse=True)
            keep_count = max(1, int(len(sorted_data) * self.data_filter_ratio))
            selected = sorted_data[:keep_count]

        self._log_debug("Filtering kept %d/%d example(s).", len(selected), len(data))

        filtered: List[Dict[str, Any]] = []
        for entry in selected:
            entry_copy = copy.deepcopy(entry)
            entry_copy.pop("reward", None)
            entry_copy.pop("reward_jitter", None)
            entry_copy.pop("metadata", None)
            filtered.append(entry_copy)

        return filtered

    def _wait_for_file_processed(self, file_id: str, interval: int = FILE_STATUS_POLL_INTERVAL) -> None:
        """Poll the uploaded training file until Azure marks it as processed.

        Args:
            file_id: Identifier returned by `files.create`.
            interval: Number of seconds to wait between polling attempts.
        """
        self._log_info("Waiting for training file %s to reach the processed state.", file_id)
        while True:
            file_info = self.openai_client.files.retrieve(file_id)
            status = getattr(file_info, "status", None)
            self._log_debug("Training file %s status: %s", file_id, status)

            if status == "processed":
                return
            if status == "failed":
                raise RuntimeError(f"Azure OpenAI reported a failure while processing file {file_id}.")

            time.sleep(interval)

    def _wait_for_finetuning(self, job_id: str, interval: int = FINETUNE_JOB_POLL_INTERVAL) -> str:
        """Poll the fine-tuning job until a terminal status is reached.

        Args:
            job_id: Identifier of the fine-tuning job to monitor.
            interval: Number of seconds between polling attempts.

        Returns:
            The identifier of the fine-tuned model when successful.
            Otherwise, raise an exception.
        """

        self._log_info("Waiting for fine-tuning job %s to complete.", job_id)

        while True:
            job = self.openai_client.fine_tuning.jobs.retrieve(job_id)
            self._log_debug("Fine-tuning job %s status: %s", job_id, job.status)

            if job.status == "succeeded":
                if job.fine_tuned_model:
                    return job.fine_tuned_model
                else:
                    raise RuntimeError(f"Fine-tuning job {job_id} succeeded but no model id was returned: {job}")
            if job.status in {"failed", "cancelled"}:
                raise RuntimeError(f"Fine-tuning job {job_id} ended with status {job.status}.")

            time.sleep(interval)

    def _deploy_model(self, model_name: str, deployment_name: str, version: str) -> None:
        """Deploy the fine-tuned model using Azure's control plane REST API.

        Args:
            model_name: Fine-tuned (training) model identifier returned by Azure OpenAI.
            deployment_name: Name of the deployment to update.
            version: Version string to stamp on the deployment update.
        """
        token = self._get_azure_token()

        request_url = (
            f"https://management.azure.com/subscriptions/{self.subscription_id}"
            f"/resourceGroups/{self.resource_group}"
            f"/providers/Microsoft.CognitiveServices/accounts/{self.resource_name}"
            f"/deployments/{deployment_name}"
        )

        headers = {
            "Authorization": f"Bearer {token}",
            "Content-Type": "application/json",
        }

        # Follows the setup in https://github.com/azure-ai-foundry/fine-tuning/blob/047fd230a77e327e75d4bc41403ee8e7bff4de9e/Demos/DistillingSarcasm/sarcasm.ipynb
        deploy_data = {
            "sku": {"name": "DeveloperTier", "capacity": 250},
            "properties": {
                "model": {
                    "format": "OpenAI",
                    "name": model_name,
                    "version": version,
                }
            },
        }

        self._log_info("Deploying model %s (version %s) to deployment %s.", model_name, version, deployment_name)

        response = requests.put(
            request_url,
            params={"api-version": "2025-06-01"},
            headers=headers,
            data=json.dumps(deploy_data),
            timeout=180,
        )

        if response.status_code < 400:
            self._log_info("Deployment %s updated successfully.", deployment_name)
        else:
            self._log_error("Deployment failed: %s %s", response.status_code, response.text)

    def _wait_for_deployment_ready(self, deployment_name: str, version: str, interval: int = 30) -> None:
        """Poll the deployment status until it is marked as ready.

        Args:
            deployment_name: Name of the deployment to monitor.
            interval: Number of seconds between polling attempts.
        """
        self._log_info("Waiting for deployment %s to become ready.", deployment_name)
        while True:
            request_url = (
                f"https://management.azure.com/subscriptions/{self.subscription_id}"
                f"/resourceGroups/{self.resource_group}"
                f"/providers/Microsoft.CognitiveServices/accounts/{self.resource_name}"
                f"/deployments/{deployment_name}"
            )

            token = self._get_azure_token()
            headers = {
                "Authorization": f"Bearer {token}",
                "Content-Type": "application/json",
            }

            response = requests.get(
                request_url,
                params={"api-version": "2025-06-01"},
                headers=headers,
                timeout=60,
            )

            if response.status_code >= 400:
                self._log_error(
                    "Failed to query deployment status. Retry later: %s, %s", response.status_code, response.text
                )
            else:
                deployment_info = response.json()
                properties = deployment_info.get("properties", {})
                model_info = properties.get("model", {})
                provisioning_state = properties.get("provisioningState")
                self._log_info(
                    "Waiting for deployment to be ready. Current provisioning state of %s: %s",
                    deployment_name,
                    provisioning_state,
                )

                if provisioning_state == "Succeeded":
                    version_found = model_info.get("version")
                    if version_found == version:
                        self._log_info("Deployment %s is ready with version %s.", deployment_name, version)
                        return
                    else:
                        self._log_warning(
                            "Deployment succeeded, but version mismatch: expected %s, got %s. Try again later.",
                            version,
                            version_found,
                        )
                elif provisioning_state == "Cancelled" or provisioning_state == "Failed":
                    raise RuntimeError(f"Deployment {deployment_name} failed with state {provisioning_state}.")
                else:
                    # Just wait and poll again
                    self._log_debug(
                        "Deployment %s not ready yet. Current state: %s", deployment_name, provisioning_state
                    )

            time.sleep(interval)

    def _delete_deployment(self, deployment_name: str) -> None:
        """Delete a specific deployment in Azure OpenAI.

        Args:
            deployment_name: Name of the deployment to delete.
        """
        token = self._get_azure_token()
        request_url = (
            f"https://management.azure.com/subscriptions/{self.subscription_id}"
            f"/resourceGroups/{self.resource_group}"
            f"/providers/Microsoft.CognitiveServices/accounts/{self.resource_name}"
            f"/deployments/{deployment_name}"
        )

        headers = {
            "Authorization": f"Bearer {token}",
            "Content-Type": "application/json",
        }

        self._log_info("Deleting deployment %s...", deployment_name)

        response = requests.delete(
            request_url,
            params={"api-version": "2025-06-01"},
            headers=headers,
            timeout=60,
        )

        if response.status_code in (200, 202, 204):
            self._log_info("Deployment %s deleted successfully.", deployment_name)
        else:
            self._log_error(
                "Failed to delete deployment %s: %s %s",
                deployment_name,
                response.status_code,
                response.text,
            )

    def _get_azure_token(self) -> str:
        """Request an Azure management token via the Azure CLI.

        Returns:
            Bearer token that authorizes calls to the Azure management plane.
        """
        cmd = [
            "az",
            "account",
            "get-access-token",
            "--resource",
            "https://management.azure.com",
            "--query",
            "accessToken",
            "-o",
            "tsv",
        ]
        try:
            token = subprocess.check_output(cmd, text=True).strip()
        except subprocess.CalledProcessError as exc:
            raise ValueError("Azure CLI command failed. Could not fetch token from Azure CLI.") from exc
        if token:
            return token
        else:
            raise ValueError("Could not fetch token from Azure CLI.")

    # Logging helpers

    def _log_info(self, message: str, *args: Any, **kwargs: Any) -> None:
        logger.info(f"{self._log_prefix}{message}", *args, **kwargs)

    def _log_debug(self, message: str, *args: Any, **kwargs: Any) -> None:
        logger.debug(f"{self._log_prefix}{message}", *args, **kwargs)

    def _log_warning(self, message: str, *args: Any, **kwargs: Any) -> None:
        logger.warning(f"{self._log_prefix}{message}", *args, **kwargs)

    def _log_error(self, message: str, *args: Any, **kwargs: Any) -> None:
        logger.error(f"{self._log_prefix}{message}", *args, **kwargs)