Orchestrate Fireworks AI with Apache Airflow®

Fireworks AI is a platform for running inference on AI models and fine-tuning open models. You can call models on shared serverless infrastructure, run them on dedicated GPU deployments, and fine-tune models on your own datasets. Apache Airflow® lets you orchestrate all of these actions as part of your Dags. In this guide you’ll learn:

How to connect Airflow to Fireworks AI.
How to run serverless inference with the Common AI provider.
How to run inference on a dedicated Fireworks AI deployment that scales to zero between Dag runs.
How to fine-tune a model on Fireworks AI.

Assumed knowledge

To get the most out of this guide, make sure you have an understanding of:

Basic Airflow concepts. See Introduction to Apache Airflow.
Basic use of the Common AI provider. See Orchestrate AI tasks with Apache Airflow and the Common AI provider.
Basic AI orchestration concepts. See AI orchestration overview.

Why use Fireworks AI with Airflow?

Using Fireworks AI together with Airflow gives you the following advantages:

Flexible scheduling: Orchestrate any Fireworks AI action on a time-based, data-aware, or event-driven schedule. See Schedule Dags in Airflow.
Durable execution with Airflow: Fine-tuning jobs often take several hours or longer. When you save the Fireworks AI job ID in the task state store, a task retry reconnects to the running job and doesn’t start a second one. See Durable execution with the task state store.
Deferrable operators: A deferrable operator defers waiting for a fine-tuning job to the triggerer, which frees the worker slot until the fine-tuning job finishes. See Deferrable operators.
Setup and teardown tasks: Scale a dedicated deployment up in a setup task and back down in a teardown task. As long as the setup task succeeded, the teardown task runs even when the inference task fails. See Use setup and teardown tasks in Airflow.
Dynamic task mapping: Run one inference task per file, record, or prompt, with the number of tasks determined at runtime. See Create dynamic Airflow tasks.
Common AI provider: The Common AI provider includes several operators and decorators (for example, @task.llm and @task.agent) that call Fireworks AI models and return structured output. This guide uses the Common AI provider to run serverless inference and to run inference on a dedicated deployment. See Common AI provider.

Ways to use Fireworks AI with Airflow

You can use Fireworks AI in Airflow in the following ways:

Call the Fireworks AI REST API or use the Fireworks AI Python SDK in a @task. This works for every Fireworks AI feature: inference, datasets, fine-tuning jobs, and deployments. To reuse the same logic across Dags, wrap it in a custom operator, which you can make durable, deferrable, or both.
Use the Common AI provider, which is based on PydanticAI, for inference. This works with serverless models and dedicated deployments, for both LLM calls and AI agents.

Airflow works with any harness that you can run in Python. If your harness, for example LangChain or CrewAI, supports Fireworks AI, you can run it inside a @task. See Start with the harness you already have.

Connect to Fireworks AI

How you provide your Fireworks AI API key depends on how you call Fireworks AI. The Common AI provider reads the key from an Airflow connection. When you use the Fireworks AI Python SDK or the API in a @task, you can read the key from the same connection or from the FIREWORKS_API_KEY environment variable.

Use a Fireworks AI API key directly

To use the Fireworks AI Python SDK or the API directly in a @task, set your API key and account ID as environment variables.

Create an API key in the Fireworks AI console.
Copy your Fireworks AI account ID.
If you want to use the Fireworks AI Python SDK, add it to your requirements.txt file.
fireworks-ai==<version>
Set the API key and the account ID as environment variables. If you run Airflow locally with the Astro CLI, add them to your .env file. On Astro, set them in the Astro Environment Manager.
FIREWORKS_API_KEY="<your-fireworks-api-key>" FIREWORKS_ACCOUNT_ID="<your-fireworks-account-id>"

Use the Common AI provider

The Common AI provider reads your API key and a default model from an Airflow connection of the pydanticai type.

Add the Common AI provider with the openai extra to your requirements.txt file. The Common AI provider uses PydanticAI to call models, and the Fireworks AI integration in PydanticAI uses the OpenAI client.
apache-airflow-providers-common-ai[openai]==<version>

Add the connection and your Fireworks AI account ID to your .env file. Use the pydanticai connection type, your Fireworks AI API key as the password, and a default model in the extra field in the format fireworks:<model>.

AIRFLOW_CONN_FIREWORKS_DEFAULT='{
    "conn_type": "pydanticai",
    "password": "<your-fireworks-api-key>",
    "extra": {
        "model": "fireworks:accounts/fireworks/models/gpt-oss-120b"
    }
}'
FIREWORKS_ACCOUNT_ID="<your-fireworks-account-id>"

PydanticAI sends requests for models with the fireworks: prefix to https://api.fireworks.ai/inference/v1, so you don’t need to set a host. You can override the default model on each task with the model_id parameter.

Run serverless inference

Serverless models run on infrastructure that Fireworks AI shares between customers. See the Fireworks AI model library for all available models.

LLM orchestration

The following Dag retrieves the user prompt with an upstream task, passes it to Fireworks AI, and returns the LLM output as a structured summary, which the downstream task prints.

For use cases and more information on how to orchestrate LLM calls with Airflow, see LLM orchestration with Apache Airflow®.

Common AI provider
Fireworks AI Python SDK

from pydantic import BaseModel

from airflow.sdk import chain, dag, task

LLM_CONN_ID = "fireworks_default"
MODEL_ID = "fireworks:accounts/fireworks/models/gpt-oss-120b"
SYSTEM_PROMPT = "You are a concise assistant. Keep the summary to one sentence."


class ToolSummary(BaseModel):
    name: str
    summary: str


@dag
def fireworks_ai_inference_llm():

    @task
    def get_user_prompt() -> str:
        return "What is Apache Airflow?"

    @task.llm(
        llm_conn_id=LLM_CONN_ID,
        model_id=MODEL_ID,
        system_prompt=SYSTEM_PROMPT,
        output_type=ToolSummary,
    )
    def ask_question(user_prompt: str):
        return user_prompt

    @task
    def print_response(response: ToolSummary):
        print(f"{response.name}: {response.summary}")

    _get_user_prompt = get_user_prompt()
    _ask_question = ask_question(user_prompt=_get_user_prompt)
    _print_response = print_response(response=_ask_question)

    chain(_get_user_prompt, _ask_question, _print_response)


fireworks_ai_inference_llm()

This Dag has three tasks:

get_user_prompt: Returns the user prompt. You can use any Python code to retrieve the prompt from any location.
ask_question: Uses @task.llm to send the prompt and a system prompt to gpt-oss-120b through the fireworks_default connection. The task returns the model response as a ToolSummary object.
print_response: Prints the name and summary from the model response. Airflow lets you use the model output in downstream tasks, for example to route data to different branches, classify records, or write the results to a database.

from fireworks import Fireworks
from pydantic import BaseModel

from airflow.sdk import chain, dag, task

MODEL_ID = "accounts/fireworks/models/gpt-oss-120b"
SYSTEM_PROMPT = "You are a concise assistant. Keep the summary to one sentence."


class ToolSummary(BaseModel):
    name: str
    summary: str


@dag
def fireworks_ai_inference_llm_sdk():

    @task
    def get_user_prompt() -> str:
        return "What is Apache Airflow?"

    @task
    def ask_question(user_prompt: str) -> dict:
        response = Fireworks().chat.completions.create(
            model=MODEL_ID,
            messages=[
                {
                    "role": "system",
                    "content": SYSTEM_PROMPT,
                },
                {"role": "user", "content": user_prompt},
            ],
            response_format={
                "type": "json_schema",
                "json_schema": {
                    "name": "ToolSummary",
                    "schema": ToolSummary.model_json_schema(),
                },
            },
        )
        return ToolSummary.model_validate_json(
            response.choices[0].message.content
        ).model_dump()

    @task
    def print_response(response: dict):
        print(f"{response['name']}: {response['summary']}")

    _get_user_prompt = get_user_prompt()
    _ask_question = ask_question(user_prompt=_get_user_prompt)
    _print_response = print_response(response=_ask_question)

    chain(_get_user_prompt, _ask_question, _print_response)


fireworks_ai_inference_llm_sdk()

This Dag has three tasks:

get_user_prompt: Returns the user prompt. You can use any Python code to retrieve the prompt from any location.
ask_question: Sends the prompt and a system prompt to gpt-oss-120b with the Fireworks AI Python SDK. The response_format parameter requests JSON that matches the ToolSummary schema, and the task validates the response with Pydantic and returns it as a dictionary.
print_response: Prints the name and summary from the model response. Airflow lets you use the model output in downstream tasks, for example to route data to different branches, classify records, or write the results to a database.

Agent orchestration

The following Dag retrieves the user prompt with an upstream task, passes it to an AI agent that runs on Fireworks AI, and returns the agent output as a structured result, which the downstream task prints. Agents can be given access to tools. Any Python function can be run as a tool by an agent.

For use cases and more information on how to orchestrate AI agents with Airflow, see Agent orchestration with Apache Airflow®.

Common AI provider
Fireworks AI Python SDK

import random

from pydantic import BaseModel
from pydantic_ai import FunctionToolset

from airflow.sdk import chain, dag, task

LLM_CONN_ID = "fireworks_default"
MODEL_ID = "fireworks:accounts/fireworks/models/gpt-oss-120b"
SYSTEM_PROMPT = "Always use the get_random_number tool when you need a random number."


class RandomNumberResult(BaseModel):
    number: int
    is_prime: bool


def get_random_number(min_value: int, max_value: int) -> int:
    """Return a random integer between min_value and max_value, inclusive."""
    return random.randint(min_value, max_value)


@dag
def fireworks_ai_inference_agent():

    @task
    def get_user_prompt() -> str:
        return "Pick a random number between 1 and 100 and tell me whether it is prime."

    @task.agent(
        llm_conn_id=LLM_CONN_ID,
        model_id=MODEL_ID,
        system_prompt=SYSTEM_PROMPT,
        output_type=RandomNumberResult,
        toolsets=[FunctionToolset([get_random_number])],
    )
    def pick_number(user_prompt: str):
        return user_prompt

    @task
    def print_response(response: RandomNumberResult):
        print(f"Number: {response.number}, prime: {response.is_prime}")

    _get_user_prompt = get_user_prompt()
    _pick_number = pick_number(user_prompt=_get_user_prompt)
    _print_response = print_response(response=_pick_number)

    chain(_get_user_prompt, _pick_number, _print_response)


fireworks_ai_inference_agent()

This Dag has three tasks:

get_user_prompt: Returns the user prompt. You can use any Python code to retrieve the prompt from any location.
pick_number: Uses @task.agent to run an agent that uses gpt-oss-120b. The get_random_number function is passed to the agent in a FunctionToolset with the toolsets parameter. See Toolsets for more information. The task returns the agent output as a RandomNumberResult object.
print_response: Prints the number and whether it’s prime. Airflow lets you use the agent output in downstream tasks, for example to route data to different branches, classify records, or write the results to a database.

import json
import random

from fireworks import Fireworks
from pydantic import BaseModel

from airflow.sdk import chain, dag, task

MODEL_ID = "accounts/fireworks/models/gpt-oss-120b"
SYSTEM_PROMPT = "Always use the get_random_number tool when you need a random number."
MAX_TURNS = 5


class RandomNumberResult(BaseModel):
    number: int
    is_prime: bool


def get_random_number(min_value: int, max_value: int) -> int:
    return random.randint(min_value, max_value)


TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "get_random_number",
            "description": "Return a random integer between min_value and max_value, inclusive.",
            "parameters": {
                "type": "object",
                "properties": {
                    "min_value": {"type": "integer"},
                    "max_value": {"type": "integer"},
                },
                "required": ["min_value", "max_value"],
            },
        },
    }
]


@dag
def fireworks_ai_inference_agent_sdk():

    @task
    def get_user_prompt() -> str:
        return "Pick a random number between 1 and 100 and tell me whether it is prime."

    @task
    def pick_number(user_prompt: str) -> dict:
        client = Fireworks()
        messages = [
            {
                "role": "system",
                "content": SYSTEM_PROMPT,
            },
            {"role": "user", "content": user_prompt},
        ]

        for _ in range(MAX_TURNS):
            response = client.chat.completions.create(
                model=MODEL_ID, messages=messages, tools=TOOLS
            )
            message = response.choices[0].message
            messages.append(message.model_dump(exclude_none=True))
            if not message.tool_calls:
                break

            for tool_call in message.tool_calls:
                arguments = json.loads(tool_call.function.arguments)
                print(f"Calling {tool_call.function.name} with {arguments}")
                messages.append(
                    {
                        "role": "tool",
                        "tool_call_id": tool_call.id,
                        "content": str(get_random_number(**arguments)),
                    }
                )
        else:
            raise RuntimeError(f"No final answer after {MAX_TURNS} turns")

        response = client.chat.completions.create(
            model=MODEL_ID,
            messages=messages,
            response_format={
                "type": "json_schema",
                "json_schema": {
                    "name": "RandomNumberResult",
                    "schema": RandomNumberResult.model_json_schema(),
                },
            },
        )
        return RandomNumberResult.model_validate_json(
            response.choices[0].message.content
        ).model_dump()

    @task
    def print_response(response: dict):
        print(f"Number: {response['number']}, prime: {response['is_prime']}")

    _get_user_prompt = get_user_prompt()
    _pick_number = pick_number(user_prompt=_get_user_prompt)
    _print_response = print_response(response=_pick_number)

    chain(_get_user_prompt, _pick_number, _print_response)


fireworks_ai_inference_agent_sdk()

This Dag has three tasks:

get_user_prompt: Returns the user prompt. You can use any Python code to retrieve the prompt from any location.
pick_number: Runs a tool-calling loop with the Fireworks AI Python SDK. In each turn, the task sends the conversation and the get_random_number tool definition to gpt-oss-120b, and runs the tool when the model calls it, for up to MAX_TURNS turns. A final request with response_format returns JSON that matches the RandomNumberResult schema, which the task validates with Pydantic and returns as a dictionary.
print_response: Prints the number and whether it’s prime. Airflow lets you use the agent output in downstream tasks, for example to route data to different branches, classify records, or write the results to a database.

Run inference on a dedicated deployment

A dedicated deployment runs a model on GPU capacity reserved for your Fireworks AI account. The following Dag scales a dedicated deployment up to one replica, runs inference on it, and scales the deployment back down to zero replicas at the end of the Dag run, even if inference fails. If the deployment doesn’t exist yet, the Dag creates it.

@task
Custom operator

import os

from fireworks import Fireworks, NotFoundError

from airflow.sdk import PokeReturnValue, chain, dag, setup, task, teardown

FIREWORKS_ACCOUNT_ID = os.getenv("FIREWORKS_ACCOUNT_ID")
BASE_MODEL = "accounts/fireworks/models/llama-v3p1-8b-instruct"
DEPLOYMENT_ID = "airflow-sdk-demo"
DEPLOYMENT_SHAPE = "accounts/fireworks/deploymentShapes/llama-v3p1-8b-instruct-minimal"
SCALE_TO_ZERO_WINDOW = "300s"
LLM_CONN_ID = "fireworks_default"
SYSTEM_PROMPT = "You are a concise assistant. Answer in one sentence."


@dag
def fireworks_ai_deployment_sdk():

    @setup
    def scale_up_deployment() -> str:
        client = Fireworks()
        try:
            client.deployments.get(DEPLOYMENT_ID)
        except NotFoundError:
            client.deployments.create(
                deployment_id=DEPLOYMENT_ID,
                base_model=BASE_MODEL,
                deployment_shape=DEPLOYMENT_SHAPE,
                min_replica_count=0,
                max_replica_count=1,
                autoscaling_policy={"scale_to_zero_window": SCALE_TO_ZERO_WINDOW},
            )
            print(f"Created deployment {DEPLOYMENT_ID}")

        client.deployments.scale(DEPLOYMENT_ID, replica_count=1)
        print(f"Scaling deployment {DEPLOYMENT_ID} to 1 replica")
        return DEPLOYMENT_ID

    @task.sensor(poke_interval=30, timeout=60 * 60, mode="reschedule")
    def wait_for_replica(deployment_id: str) -> PokeReturnValue:
        deployment = Fireworks().deployments.get(deployment_id)
        ready = (
            deployment.replica_stats.ready_replica_count
            if deployment.replica_stats
            else 0
        )
        print(f"Deployment state: {deployment.state}, ready replicas: {ready}")

        if deployment.state == "FAILED":
            raise RuntimeError(f"Deployment {deployment_id} failed")
        return PokeReturnValue(is_done=bool(ready))

    @task.llm(
        llm_conn_id=LLM_CONN_ID,
        model_id=f"fireworks:{BASE_MODEL}#accounts/{FIREWORKS_ACCOUNT_ID}/deployments/{DEPLOYMENT_ID}",
        system_prompt=SYSTEM_PROMPT,
    )
    def ask_question():
        return "What is Apache Airflow?"

    @teardown(on_failure_fail_dagrun=True)
    def scale_down_deployment(deployment_id: str):
        Fireworks().deployments.scale(deployment_id, replica_count=0)
        print(f"Scaled deployment {deployment_id} to 0 replicas")

    _scale_up_deployment = scale_up_deployment()
    _wait_for_replica = wait_for_replica(deployment_id=_scale_up_deployment)
    _ask_question = ask_question()
    _scale_down_deployment = scale_down_deployment(deployment_id=_scale_up_deployment)

    chain(
        _scale_up_deployment, _wait_for_replica, _ask_question, _scale_down_deployment
    )


fireworks_ai_deployment_sdk()

This Dag has four tasks:

scale_up_deployment: A setup task that scales the deployment in DEPLOYMENT_ID up to one replica. If the deployment doesn’t exist yet, the task creates it.
wait_for_replica: A sensor in reschedule mode that checks the deployment every 30 seconds until a replica is ready.
ask_question: Uses @task.llm to send a prompt to the deployment.
scale_down_deployment: A teardown task that scales the deployment to zero replicas, even if ask_question fails.

If the teardown task doesn’t run, Fireworks AI scales a deployment with a minimum of zero replicas down to zero after it receives no requests for the duration in scale_to_zero_window. The value must be at least 5 minutes, and the default is 1 hour.

Airflow UI graph view of a running fireworks_ai_deployment_sdk Dag run, where scale_up_deployment succeeded and the wait_for_replica sensor is up for reschedule while it waits for a ready replica

import os

from include.custom_operators.fireworks_deployment import (
    FireworksScaleDownDeploymentOperator,
    FireworksScaleUpDeploymentOperator,
)

from airflow.sdk import chain, dag, task

FIREWORKS_ACCOUNT_ID = os.getenv("FIREWORKS_ACCOUNT_ID")
BASE_MODEL = "accounts/fireworks/models/llama-v3p1-8b-instruct"
DEPLOYMENT_ID = "airflow-operator-demo"
DEPLOYMENT_SHAPE = "accounts/fireworks/deploymentShapes/llama-v3p1-8b-instruct-minimal"
LLM_CONN_ID = "fireworks_default"
SYSTEM_PROMPT = "You are a concise assistant. Answer in one sentence."


@dag
def fireworks_ai_deployment_custom_operator():

    _scale_up_deployment = FireworksScaleUpDeploymentOperator(
        task_id="scale_up_deployment",
        deployment_id=DEPLOYMENT_ID,
        base_model=BASE_MODEL,
        deployment_shape=DEPLOYMENT_SHAPE,
    )

    @task.llm(
        llm_conn_id=LLM_CONN_ID,
        model_id=f"fireworks:{BASE_MODEL}#accounts/{FIREWORKS_ACCOUNT_ID}/deployments/{DEPLOYMENT_ID}",
        system_prompt=SYSTEM_PROMPT,
    )
    def ask_question():
        return "What is Apache Airflow?"

    _ask_question = ask_question()

    _scale_down_deployment = FireworksScaleDownDeploymentOperator(
        task_id="scale_down_deployment",
        deployment_id=DEPLOYMENT_ID,
    )

    chain(
        _scale_up_deployment,
        _ask_question,
        _scale_down_deployment.as_teardown(
            setups=_scale_up_deployment, on_failure_fail_dagrun=True
        ),
    )


fireworks_ai_deployment_custom_operator()

This Dag has three tasks:

scale_up_deployment: Uses FireworksScaleUpDeploymentOperator as a setup task to create the deployment in DEPLOYMENT_ID if it doesn’t exist yet, scale it to one replica, and wait until the replica is ready. The operator is deferrable and durable. This means that the worker slot is released while the task waits for the deployment to be ready, and if the task fails and is retried, it reconnects to the deployment using the same deployment ID.
ask_question: Uses @task.llm to send a prompt to the deployment.
scale_down_deployment: Uses FireworksScaleDownDeploymentOperator as a teardown task to scale the deployment to zero replicas, even if ask_question fails.

scale_up_deployment scales the deployment up and waits for the replica in one task. If the deployment is scaled up but no replica becomes ready, the task fails, and the teardown task doesn’t run because its setup task failed. In that case, Fireworks AI scales the deployment to zero replicas after the duration in scale_to_zero_window. FireworksScaleUpDeploymentOperator sets it to 300 seconds unless you pass a different value. The Fireworks AI default is 1 hour.

Airflow UI graph view of a running fireworks_ai_deployment_custom_operator Dag run, where the scale_up_deployment task is deferred while it waits for a ready replica, and scale_down_deployment is the teardown task of scale_up_deployment

You can find the code of the two custom operators below:

from __future__ import annotations

import asyncio
import os
import time
from collections.abc import AsyncIterator, Sequence
from typing import Any

from fireworks import AsyncFireworks, Fireworks, NotFoundError

from airflow.sdk import BaseOperator, Context, ResumableJobMixin
from airflow.triggers.base import BaseTrigger, TriggerEvent


def _ready_replicas(deployment) -> int:
    return (
        deployment.replica_stats.ready_replica_count if deployment.replica_stats else 0
    )


class FireworksDeploymentReadyTrigger(BaseTrigger):
    """
    Wait asynchronously until a Fireworks AI deployment has a ready replica.

    :param deployment_id: ID of the deployment.
    :param account_id: Fireworks account ID.
    :param poll_interval: Seconds between status checks.
    """

    def __init__(
        self, deployment_id: str, account_id: str, poll_interval: float = 30.0
    ):
        super().__init__()
        self.deployment_id = deployment_id
        self.account_id = account_id
        self.poll_interval = poll_interval

    def serialize(self) -> tuple[str, dict[str, Any]]:
        return (
            "include.custom_operators.fireworks_deployment.FireworksDeploymentReadyTrigger",
            {
                "deployment_id": self.deployment_id,
                "account_id": self.account_id,
                "poll_interval": self.poll_interval,
            },
        )

    async def run(self) -> AsyncIterator[TriggerEvent]:
        client = AsyncFireworks(account_id=self.account_id)

        while True:
            deployment = await client.deployments.get(self.deployment_id)
            if _ready_replicas(deployment):
                yield TriggerEvent(
                    {"status": "success", "deployment_id": self.deployment_id}
                )
                return
            if deployment.state == "FAILED":
                message = deployment.status.message if deployment.status else ""
                yield TriggerEvent(
                    {
                        "status": "error",
                        "deployment_id": self.deployment_id,
                        "message": message,
                    }
                )
                return
            self.log.info(
                "Deployment %s is %s with no ready replica, checking again in %ss.",
                self.deployment_id,
                deployment.state,
                self.poll_interval,
            )
            await asyncio.sleep(self.poll_interval)


class FireworksScaleUpDeploymentOperator(ResumableJobMixin, BaseOperator):
    """
    Scale a Fireworks AI deployment to one replica and wait until the replica is ready.

    Creates the deployment with 0 to 1 replicas if it doesn't exist yet. Fireworks
    scales the deployment to 0 replicas after ``scale_to_zero_window`` without requests.
    The deployment ID is stored in the task state store, so a retry reconnects to the
    deployment instead of creating a second one.
    Returns the deployment resource name, e.g. ``accounts/<account>/deployments/<id>``.

    :param base_model: Model to deploy, e.g. ``accounts/fireworks/models/<model>``.
    :param deployment_shape: Deployment shape, e.g. ``accounts/fireworks/deploymentShapes/<shape>``.
    :param deployment_id: ID of the deployment. Fireworks generates an ID for a new
        deployment if not set.
    :param scale_to_zero_window: Time without requests after which Fireworks scales to 0 replicas.
    :param account_id: Fireworks account ID. Defaults to the ``FIREWORKS_ACCOUNT_ID`` environment variable.
    :param poll_interval: Seconds between status checks.
    :param deferrable: Wait for the replica in the triggerer instead of on the worker.
    :param durable: Store the deployment ID in the task state store so retries reconnect to it.
        Handled by ``ResumableJobMixin``, default ``True``.
    """

    template_fields: Sequence[str] = (
        "base_model",
        "deployment_shape",
        "deployment_id",
        "account_id",
    )
    external_id_key = "fireworks_deployment_id"

    def __init__(
        self,
        *,
        base_model: str,
        deployment_shape: str,
        deployment_id: str | None = None,
        scale_to_zero_window: str = "300s",
        account_id: str | None = None,
        poll_interval: int = 30,
        deferrable: bool = True,
        **kwargs,
    ) -> None:
        super().__init__(**kwargs)
        self.base_model = base_model
        self.deployment_shape = deployment_shape
        self.deployment_id = deployment_id
        self.scale_to_zero_window = scale_to_zero_window
        self.account_id = account_id or os.getenv("FIREWORKS_ACCOUNT_ID")
        self.poll_interval = poll_interval
        self.deferrable = deferrable

    def _client(self) -> Fireworks:
        return Fireworks(account_id=self.account_id)

    def execute(self, context: Context) -> str:
        return self.execute_resumable(context)

    def submit_job(self, context: Context) -> str:
        client = self._client()
        if self.deployment_id:
            try:
                client.deployments.get(self.deployment_id)
                return self.deployment_id
            except NotFoundError:
                pass

        create_kwargs = {
            "base_model": self.base_model,
            "deployment_shape": self.deployment_shape,
            "min_replica_count": 0,
            "max_replica_count": 1,
            "autoscaling_policy": {"scale_to_zero_window": self.scale_to_zero_window},
        }
        if self.deployment_id:
            create_kwargs["deployment_id"] = self.deployment_id

        deployment = client.deployments.create(**create_kwargs)
        self.log.info("Created deployment %s", deployment.name)
        return deployment.name.rsplit("/", 1)[-1]

    def get_job_status(self, external_id: str, context: Context) -> str:
        deployment = self._client().deployments.get(external_id)
        if _ready_replicas(deployment):
            return "READY"
        if deployment.state == "FAILED":
            return "FAILED"
        return "SCALING_UP"

    def is_job_active(self, status: str) -> bool:
        return status == "SCALING_UP"

    def is_job_succeeded(self, status: str) -> bool:
        return status == "READY"

    def poll_until_complete(self, external_id: str, context: Context) -> None:
        client = self._client()
        client.deployments.scale(external_id, replica_count=1)
        self.log.info("Scaling deployment %s to 1 replica", external_id)

        if self.deferrable:
            self.defer(
                trigger=FireworksDeploymentReadyTrigger(
                    deployment_id=external_id,
                    account_id=self.account_id,
                    poll_interval=self.poll_interval,
                ),
                method_name="execute_complete",
            )

        while True:
            status = self.get_job_status(external_id, context)
            if self.is_job_succeeded(status):
                return
            if status == "FAILED":
                raise RuntimeError(f"Deployment {external_id} failed")
            self.log.info(
                "Deployment %s has no ready replica, checking again in %ss.",
                external_id,
                self.poll_interval,
            )
            time.sleep(self.poll_interval)

    def get_job_result(self, external_id: str, context: Context) -> str:
        deployment = self._client().deployments.get(external_id)
        self.log.info("Deployment %s has a ready replica", deployment.name)
        return deployment.name

    def execute_complete(self, context: Context, event: dict[str, Any]) -> str:
        if event["status"] == "error":
            raise RuntimeError(
                f"Deployment {event['deployment_id']} failed: {event.get('message', '')}"
            )
        return self.get_job_result(event["deployment_id"], context)


class FireworksScaleDownDeploymentOperator(BaseOperator):
    """
    Scale a Fireworks AI deployment to 0 replicas.

    :param deployment_id: Deployment ID or resource name.
    :param account_id: Fireworks account ID. Defaults to the ``FIREWORKS_ACCOUNT_ID`` environment variable.
    """

    template_fields: Sequence[str] = ("deployment_id", "account_id")

    def __init__(
        self,
        *,
        deployment_id: str,
        account_id: str | None = None,
        **kwargs,
    ) -> None:
        super().__init__(**kwargs)
        self.deployment_id = deployment_id
        self.account_id = account_id or os.getenv("FIREWORKS_ACCOUNT_ID")

    def execute(self, context: Context) -> None:
        deployment_id = self.deployment_id.rsplit("/", 1)[-1]
        Fireworks(account_id=self.account_id).deployments.scale(
            deployment_id, replica_count=0
        )
        self.log.info("Scaled deployment %s to 0 replicas", deployment_id)

Fine-tune a model

Supervised fine-tuning on Fireworks AI takes a dataset in your account and a base model, and creates a new model in your account. A fine-tuning run in Airflow has three steps:

Upload the training data to Fireworks AI as a dataset.
Start a fine-tuning job.
Wait for the job to finish and save the name of the fine-tuned model, so that other tasks can use it, for example for inference.

To run inference on the fine-tuned model, deploy it on a dedicated deployment.

Data import and export

Fireworks AI stores training data as datasets in your account. Fine-tuning jobs for chat models expect a JSONL file with one conversation per line, in the OpenAI chat format.

Training data can include images as well as text. Fireworks AI expects each image as a base64-encoded data URI in a user message. See Vision training in the Fireworks AI documentation.

The following Dag uploads training data to a new Fireworks AI dataset with the Fireworks AI Python SDK.

import os
import uuid
from pathlib import Path

from fireworks import Fireworks

from airflow.sdk import dag, task

FIREWORKS_ACCOUNT_ID = os.getenv("FIREWORKS_ACCOUNT_ID")
DATASET_ID_PREFIX = "airflow-demo"
TRAINING_DATA_PATH = Path("include/receipt_training_data.jsonl")


@dag
def fireworks_ai_dataset_import():

    @task
    def upload_training_data() -> str:
        dataset_id = f"{DATASET_ID_PREFIX}-{uuid.uuid4().hex[:8]}"
        training_data = TRAINING_DATA_PATH.read_bytes()
        client = Fireworks()
        client.datasets.create(
            dataset_id=dataset_id,
            dataset={
                "user_uploaded": {},
                "format": "CHAT",
                "example_count": str(training_data.count(b"\n")),
            },
        )
        client.datasets.upload(
            dataset_id, file=(TRAINING_DATA_PATH.name, training_data)
        )
        print(f"Uploaded {TRAINING_DATA_PATH} to dataset {dataset_id}")
        return f"accounts/{FIREWORKS_ACCOUNT_ID}/datasets/{dataset_id}"

    upload_training_data()


fireworks_ai_dataset_import()

The upload_training_data task returns the dataset name, which a fine-tuning task can use as input. To export a dataset, for example to archive the exact training data of a model or to copy it to object storage, request signed download URLs with the SDK and download each file. The following Dag downloads one dataset to the include folder:

import os
from pathlib import Path

import requests
from fireworks import Fireworks

from airflow.sdk import chain, dag, task

FIREWORKS_ACCOUNT_ID = os.getenv("FIREWORKS_ACCOUNT_ID")
DATASET_ID = "onboarding-cord-receipts"
EXPORT_DIR = Path("include/exports")


@dag
def fireworks_ai_dataset_export():

    @task
    def get_download_urls() -> dict[str, str]:
        client = Fireworks()
        response = client.datasets.get_download_endpoint(DATASET_ID)
        return response.filename_to_signed_urls

    @task
    def download_files(signed_urls: dict[str, str]):
        target_dir = EXPORT_DIR / DATASET_ID
        target_dir.mkdir(parents=True, exist_ok=True)
        for filename, url in signed_urls.items():
            response = requests.get(url)
            response.raise_for_status()
            (target_dir / filename).write_bytes(response.content)
            print(
                f"Downloaded {filename} ({len(response.content)} bytes) to {target_dir}"
            )

    _get_download_urls = get_download_urls()
    _download_files = download_files(signed_urls=_get_download_urls)

    chain(_get_download_urls, _download_files)


fireworks_ai_dataset_export()

To export a dataset, you request signed download URLs for its files from Fireworks AI. The signed URLs expire, see Get Dataset Download Endpoint in the Fireworks AI documentation. Astronomer recommends requesting the signed URLs and downloading the files in the same Dag run.

Start the fine-tuning job and wait for it to finish

The following code starts a fine-tuning job on Fireworks AI and waits for it to finish.

@task
Custom operator

import os

from fireworks import Fireworks

from airflow.sdk import PokeReturnValue, task

FIREWORKS_ACCOUNT_ID = os.getenv("FIREWORKS_ACCOUNT_ID")
BASE_MODEL = "accounts/fireworks/models/muse-glimmer-30b"
DATASET_ID = "onboarding-cord-receipts"
LORA_RANK = 8
EPOCHS = 1


@task(retries=2)
def start_fine_tuning_job(task_state_store=None) -> str:
    client = Fireworks()
    job_id = task_state_store.get("fine_tuning_job_id")
    if job_id:
        state = client.supervised_fine_tuning_jobs.get(job_id).state
        if state not in (
            "JOB_STATE_FAILED",
            "JOB_STATE_CANCELLED",
            "JOB_STATE_EXPIRED",
        ):
            print(f"Reusing fine-tuning job {job_id} in state {state}")
            return job_id

    job = client.supervised_fine_tuning_jobs.create(
        dataset=f"accounts/{FIREWORKS_ACCOUNT_ID}/datasets/{DATASET_ID}",
        base_model=BASE_MODEL,
        lora_rank=LORA_RANK,
        epochs=EPOCHS,
    )
    job_id = job.name.rsplit("/", 1)[-1]
    task_state_store.set("fine_tuning_job_id", job_id)
    print(f"Started fine-tuning job: {job.name}")
    return job_id


@task.sensor(poke_interval=60, timeout=60 * 60 * 3, mode="reschedule")
def wait_for_fine_tuning_job(job_id: str) -> PokeReturnValue:
    job = Fireworks().supervised_fine_tuning_jobs.get(job_id)
    print(f"Fine-tuning job state: {job.state}")

    if job.state in (
        "JOB_STATE_FAILED",
        "JOB_STATE_CANCELLED",
        "JOB_STATE_EXPIRED",
    ):
        raise RuntimeError(f"Fine-tuning job ended in {job.state}: {job.status}")
    return PokeReturnValue(
        is_done=job.state == "JOB_STATE_COMPLETED",
        xcom_value=job.model_dump(mode="json"),
    )

start_fine_tuning_job: A durable task that starts the job and saves the job ID in the task state store. If the task fails after the job started and Airflow retries it, the retry reads the job ID and reuses the job, unless the job already failed, was cancelled, or expired.
wait_for_fine_tuning_job: A sensor in reschedule mode, which releases the worker slot between checks. It fails if the job fails and returns the job as a dictionary when the job is complete.

from __future__ import annotations

import asyncio
import os
import time
from collections.abc import AsyncIterator, Sequence
from typing import Any

from fireworks import AsyncFireworks, Fireworks

from airflow.sdk import BaseOperator, Context, ResumableJobMixin
from airflow.triggers.base import BaseTrigger, TriggerEvent

SUCCEEDED_STATES = {"JOB_STATE_COMPLETED"}
FAILED_STATES = {
    "JOB_STATE_FAILED",
    "JOB_STATE_CANCELLED",
    "JOB_STATE_CANCELLING",
    "JOB_STATE_EXPIRED",
    "JOB_STATE_EARLY_STOPPED",
    "JOB_STATE_DELETING",
    "JOB_STATE_DELETING_CLEANING_UP",
    "JOB_STATE_DELETED",
    "JOB_STATE_ARCHIVED",
}


class FireworksFineTuneTrigger(BaseTrigger):
    """
    Wait asynchronously for a Fireworks AI supervised fine-tuning job to finish.

    :param job_id: ID of the supervised fine-tuning job.
    :param account_id: Fireworks account ID.
    :param poll_interval: Seconds between status checks.
    """

    def __init__(self, job_id: str, account_id: str, poll_interval: float = 60.0):
        super().__init__()
        self.job_id = job_id
        self.account_id = account_id
        self.poll_interval = poll_interval

    def serialize(self) -> tuple[str, dict[str, Any]]:
        return (
            "include.custom_operators.fireworks_fine_tune.FireworksFineTuneTrigger",
            {
                "job_id": self.job_id,
                "account_id": self.account_id,
                "poll_interval": self.poll_interval,
            },
        )

    async def run(self) -> AsyncIterator[TriggerEvent]:
        client = AsyncFireworks(account_id=self.account_id)

        while True:
            job = await client.supervised_fine_tuning_jobs.get(self.job_id)
            if job.state in SUCCEEDED_STATES:
                yield TriggerEvent(
                    {"status": "success", "job_id": self.job_id, "state": job.state}
                )
                return
            if job.state in FAILED_STATES:
                message = job.status.message if job.status else ""
                yield TriggerEvent(
                    {
                        "status": "error",
                        "job_id": self.job_id,
                        "state": job.state,
                        "message": message,
                    }
                )
                return
            self.log.info(
                "Fine-tuning job %s is %s, checking again in %ss.",
                self.job_id,
                job.state,
                self.poll_interval,
            )
            await asyncio.sleep(self.poll_interval)


class FireworksFineTuneOperator(ResumableJobMixin, BaseOperator):
    """
    Run a supervised fine-tuning job on Fireworks AI and wait for it to finish.

    The job ID is stored in the task state store, so a retry reattaches to a job
    that is still running instead of starting a new one.

    :param dataset: Dataset resource name, e.g. ``accounts/<account>/datasets/<dataset>``.
    :param base_model: Model to fine-tune, e.g. ``accounts/fireworks/models/<model>``.
    :param output_model: Resource name for the fine-tuned model. Fireworks uses the job ID if not set.
    :param account_id: Fireworks account ID. Defaults to the ``FIREWORKS_ACCOUNT_ID`` environment variable.
    :param lora_rank: LoRA rank. Fireworks runs full-parameter training if this is 0 or unset.
    :param epochs: Number of training epochs.
    :param poll_interval: Seconds between status checks.
    :param deferrable: Wait for the job in the triggerer instead of on the worker.
    :param durable: Store the job ID in the task state store so retries reattach to the
        running job. Handled by ``ResumableJobMixin``, default ``True``.
    """

    template_fields: Sequence[str] = (
        "dataset",
        "base_model",
        "output_model",
        "account_id",
    )
    external_id_key = "fireworks_fine_tuning_job_id"

    def __init__(
        self,
        *,
        dataset: str,
        base_model: str,
        output_model: str | None = None,
        account_id: str | None = None,
        lora_rank: int = 8,
        epochs: int = 1,
        poll_interval: int = 60,
        deferrable: bool = True,
        **kwargs,
    ) -> None:
        super().__init__(**kwargs)
        self.dataset = dataset
        self.base_model = base_model
        self.output_model = output_model
        self.account_id = account_id or os.getenv("FIREWORKS_ACCOUNT_ID")
        self.lora_rank = lora_rank
        self.epochs = epochs
        self.poll_interval = poll_interval
        self.deferrable = deferrable

    def _client(self) -> Fireworks:
        return Fireworks(account_id=self.account_id)

    def execute(self, context: Context) -> dict:
        return self.execute_resumable(context)

    def submit_job(self, context: Context) -> str:
        create_kwargs = {
            "dataset": self.dataset,
            "base_model": self.base_model,
            "lora_rank": self.lora_rank,
            "epochs": self.epochs,
        }
        if self.output_model:
            create_kwargs["output_model"] = self.output_model

        job = self._client().supervised_fine_tuning_jobs.create(**create_kwargs)
        job_id = job.name.rsplit("/", 1)[-1]
        self.log.info("Started fine-tuning job %s", job.name)
        return job_id

    def get_job_status(self, external_id: str, context: Context) -> str:
        return self._client().supervised_fine_tuning_jobs.get(external_id).state

    def is_job_active(self, status: str) -> bool:
        return status not in SUCCEEDED_STATES and status not in FAILED_STATES

    def is_job_succeeded(self, status: str) -> bool:
        return status in SUCCEEDED_STATES

    def poll_until_complete(self, external_id: str, context: Context) -> None:
        if self.deferrable:
            self.defer(
                trigger=FireworksFineTuneTrigger(
                    job_id=external_id,
                    account_id=self.account_id,
                    poll_interval=self.poll_interval,
                ),
                method_name="execute_complete",
            )

        while True:
            status = self.get_job_status(external_id, context)
            if self.is_job_succeeded(status):
                return
            if status in FAILED_STATES:
                raise RuntimeError(f"Fine-tuning job {external_id} ended in {status}")
            self.log.info(
                "Fine-tuning job %s is %s, checking again in %ss.",
                external_id,
                status,
                self.poll_interval,
            )
            time.sleep(self.poll_interval)

    def get_job_result(self, external_id: str, context: Context) -> dict:
        job = self._client().supervised_fine_tuning_jobs.get(external_id)
        self.log.info("Fine-tuned model: %s", job.output_model)
        self.log.info("Base model: %s", job.base_model)
        self.log.info("Dataset: %s", job.dataset)
        if job.create_time and job.completed_time:
            self.log.info("Duration: %s", job.completed_time - job.create_time)
        return {
            "job_id": external_id,
            "job_name": job.name,
            "output_model": job.output_model,
            "base_model": job.base_model,
            "dataset": job.dataset,
            "state": job.state,
        }

    def execute_complete(self, context: Context, event: dict[str, Any]) -> dict:
        if event["status"] == "error":
            raise RuntimeError(
                f"Fine-tuning job {event['job_id']} ended in {event['state']}: {event.get('message', '')}"
            )
        return self.get_job_result(event["job_id"], context)

If several teams or Dags start fine-tuning jobs, you can put the start, wait, and result logic in one custom operator. Every Dag then starts jobs with the same defaults, logs the same information, and handles failed job states the same way.The FireworksFineTuneOperator combines two Airflow features:

Durable execution: The operator inherits from ResumableJobMixin. The mixin saves the job ID that submit_job returns in the task state store under the key in external_id_key. On a retry, it reads the ID, checks the job state with get_job_status, and reconnects to the job if is_job_active returns True. Durable execution is on by default. Set durable=False to start a new job on every try.
Deferrable execution: With deferrable=True, the default, poll_until_complete defers to FireworksFineTuneTrigger. The trigger checks the job state in the triggerer with the async Fireworks AI client, so the task doesn’t use a worker slot while the job runs. When the job finishes, the task resumes in execute_complete and returns the job as a dictionary. With deferrable=False, the operator checks the job state on the worker.

Overview

Get started

Airflow and AI

Airflow concepts

Dag authoring tools

Tutorials

Integrations & connections

Reference Architectures

Observability concepts

Glossary

Orchestrate Fireworks AI with Apache Airflow®

Assumed knowledge

Why use Fireworks AI with Airflow?

Ways to use Fireworks AI with Airflow

Connect to Fireworks AI

Use a Fireworks AI API key directly

Use the Common AI provider

Run serverless inference

LLM orchestration

Agent orchestration

Run inference on a dedicated deployment

Fine-tune a model

Data import and export

Start the fine-tuning job and wait for it to finish

​Assumed knowledge

​Why use Fireworks AI with Airflow?

​Ways to use Fireworks AI with Airflow

​Connect to Fireworks AI

​Use a Fireworks AI API key directly

​Use the Common AI provider

​Run serverless inference

​LLM orchestration

​Agent orchestration

​Run inference on a dedicated deployment

​Fine-tune a model

​Data import and export

​Start the fine-tuning job and wait for it to finish

Assumed knowledge

Why use Fireworks AI with Airflow?

Ways to use Fireworks AI with Airflow

Connect to Fireworks AI

Use a Fireworks AI API key directly

Use the Common AI provider

Run serverless inference

LLM orchestration

Agent orchestration

Run inference on a dedicated deployment

Fine-tune a model

Data import and export

Start the fine-tuning job and wait for it to finish