> ## Documentation Index > Fetch the complete documentation index at: https://astronomer.io/docs/llms.txt > Use this file to discover all available pages before exploring further. # Orchestrate Fireworks AI with Apache Airflow® > Learn how to use Fireworks AI with Airflow for serverless inference, dedicated deployments, and fine-tuning. [Fireworks AI](https://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®](https://airflow.apache.org/) 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](https://astronomer.io/docs/airflow/intro-to-airflow.md). * Basic use of the Common AI provider. See [Orchestrate AI tasks with Apache Airflow and the Common AI provider](https://astronomer.io/docs/airflow/airflow-common-ai-provider.md). * Basic AI orchestration concepts. See [AI orchestration overview](https://astronomer.io/docs/airflow/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](https://astronomer.io/docs/airflow/scheduling-in-airflow.md). * **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](https://astronomer.io/docs/airflow/airflow-task-state-store.md). * **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](https://astronomer.io/docs/airflow/deferrable-operators.md). * **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](https://astronomer.io/docs/airflow/airflow-setup-teardown.md). * **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](https://astronomer.io/docs/airflow/dynamic-tasks.md). * **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](#run-serverless-inference) and to [run inference on a dedicated deployment](#run-inference-on-a-dedicated-deployment). See [Common AI provider](https://astronomer.io/docs/airflow/airflow-common-ai-provider.md). ## Ways to use Fireworks AI with Airflow You can use Fireworks AI in Airflow in the following ways: * Call the [Fireworks AI REST API](https://docs.fireworks.ai/) or use the [Fireworks AI Python SDK](https://pypi.org/project/fireworks-ai/) 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](#start-the-fine-tuning-job-and-wait-for-it-to-finish), which you can make durable, deferrable, or both. * Use the [Common AI provider](https://astronomer.io/docs/airflow/airflow-common-ai-provider.md), which is based on [PydanticAI](https://pydantic.dev/docs/ai/overview/), for inference. This works with serverless models and dedicated deployments, for both LLM calls and AI agents. > [!TIP] > > 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](https://astronomer.io/docs/airflow/ai-orchestration-llm.md#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. 1. Create an API key in the [Fireworks AI console](https://fireworks.ai/). 2. Copy your Fireworks AI account ID. 3. If you want to use the Fireworks AI Python SDK, add it to your `requirements.txt` file. ```text wrap theme={null} fireworks-ai== ``` 4. 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](https://astronomer.io/docs/astro/environment-manager-overview.md). ```text wrap theme={null} FIREWORKS_API_KEY="" 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. 1. 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. ```text wrap theme={null} apache-airflow-providers-common-ai[openai]== ``` 2. 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:`. ```text wrap theme={null} AIRFLOW_CONN_FIREWORKS_DEFAULT='{ "conn_type": "pydanticai", "password": "", "extra": { "model": "fireworks:accounts/fireworks/models/gpt-oss-120b" } }' FIREWORKS_ACCOUNT_ID="" ``` > [!NOTE] > > 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](https://fireworks.ai/models) 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. > [!TIP] > > For use cases and more information on how to orchestrate LLM calls with Airflow, see [LLM orchestration with Apache Airflow®](https://astronomer.io/docs/airflow/ai-orchestration-llm.md). > [tab] Common AI provider > > ```python expandable wrap theme={null} > 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. > > ![Airflow UI graph view of a successful run of the fireworks\_ai\_inference\_llm Dag with the tasks get\_user\_prompt, ask\_question, which uses @task.llm, and print\_response](https://mintcdn.com/astronomer/tUqEvVpMUNByhW9D/images/img/guides/3-4_fireworks_ai_inference_llm_common_ai.png?fit=max\&auto=format\&n=tUqEvVpMUNByhW9D\&q=85\&s=6d4b047c24492f3d55f32bc2f4d56bfc) > [tab] Fireworks AI Python SDK > > ```python expandable wrap theme={null} > 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. > > ![Airflow UI graph view of a successful run of the fireworks\_ai\_inference\_llm\_sdk Dag with the tasks get\_user\_prompt, ask\_question, and print\_response](https://mintcdn.com/astronomer/tUqEvVpMUNByhW9D/images/img/guides/3-4_fireworks_ai_inference_llm_sdk.png?fit=max\&auto=format\&n=tUqEvVpMUNByhW9D\&q=85\&s=553580eb50916b26946ddde1774a74d5) ### 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. > [!TIP] > > For use cases and more information on how to orchestrate AI agents with Airflow, see [Agent orchestration with Apache Airflow®](https://astronomer.io/docs/airflow/ai-orchestration-agent.md). > [tab] Common AI provider > > ```python expandable wrap theme={null} > 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](https://astronomer.io/docs/airflow/airflow-common-ai-provider.md#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. > > ![Airflow UI graph view of a successful run of the fireworks\_ai\_inference\_agent Dag with the tasks get\_user\_prompt, pick\_number, which uses @task.agent, and print\_response](https://mintcdn.com/astronomer/tUqEvVpMUNByhW9D/images/img/guides/3-4_fireworks_ai_inference_agent_common_ai.png?fit=max\&auto=format\&n=tUqEvVpMUNByhW9D\&q=85\&s=7b8c4bce47a0736db0e5cb912fd77ce7) > [tab] Fireworks AI Python SDK > > ```python expandable wrap theme={null} > 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. > > ![Airflow UI graph view of a successful run of the fireworks\_ai\_inference\_agent\_sdk Dag with the tasks get\_user\_prompt, pick\_number, and print\_response](https://mintcdn.com/astronomer/tUqEvVpMUNByhW9D/images/img/guides/3-4_fireworks_ai_inference_agent_sdk.png?fit=max\&auto=format\&n=tUqEvVpMUNByhW9D\&q=85\&s=d7dea5f0364d4e2f237cf017117ead13) ## 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. > [tab] @task > > ```python expandable wrap theme={null} > 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](https://astronomer.io/docs/airflow/airflow-setup-teardown.md) 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](https://astronomer.io/docs/airflow/airflow-setup-teardown.md) that scales the deployment to zero replicas, even if `ask_question` fails. > > > [!NOTE] > > > > 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](https://mintcdn.com/astronomer/tUqEvVpMUNByhW9D/images/img/guides/3-4_fireworks_ai_deployment_sdk.png?fit=max\&auto=format\&n=tUqEvVpMUNByhW9D\&q=85\&s=5cc3a3ea0570c280b945d8bfa1454551) > [tab] Custom operator > > ```python expandable wrap theme={null} > 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](https://astronomer.io/docs/airflow/airflow-setup-teardown.md) 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](https://astronomer.io/docs/airflow/deferrable-operators.md) and [durable](https://astronomer.io/docs/airflow/airflow-task-state-store.md). 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](https://astronomer.io/docs/airflow/airflow-setup-teardown.md) to scale the deployment to zero replicas, even if `ask_question` fails. > > > [!NOTE] > > > > `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](https://mintcdn.com/astronomer/tUqEvVpMUNByhW9D/images/img/guides/3-4_fireworks_ai_deployment_custom_operator.png?fit=max\&auto=format\&n=tUqEvVpMUNByhW9D\&q=85\&s=1354daa96885944ac0f92a8098f3a65f) > > You can find the code of the two [custom operators](https://astronomer.io/docs/airflow/airflow-importing-custom-hooks-operators.md) below: > > ```python expandable wrap theme={null} > 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//deployments/``. > > :param base_model: Model to deploy, e.g. ``accounts/fireworks/models/``. > :param deployment_shape: Deployment shape, e.g. ``accounts/fireworks/deploymentShapes/``. > :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: 1. Upload the training data to Fireworks AI as a dataset. 2. Start a fine-tuning job. 3. 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](#run-inference-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. > [!NOTE] > > 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](https://docs.fireworks.ai/fine-tuning/fine-tuning-models#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. ```python expandable wrap theme={null} 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: ```python expandable wrap theme={null} 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() ``` > [!NOTE] > > To export a dataset, you request signed download URLs for its files from Fireworks AI. The signed URLs expire, see [Get Dataset Download Endpoint](https://docs.fireworks.ai/api-reference/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. > [tab] @task > > ```python expandable wrap theme={null} > 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](https://astronomer.io/docs/airflow/airflow-task-state-store.md) 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. > [tab] Custom operator > > ```python expandable wrap theme={null} > 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//datasets/``. > :param base_model: Model to fine-tune, e.g. ``accounts/fireworks/models/``. > :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](https://astronomer.io/docs/airflow/airflow-task-state-store.md)**: The operator inherits from [`ResumableJobMixin`](https://astronomer.io/docs/airflow/airflow-task-state-store#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](https://astronomer.io/docs/airflow/deferrable-operators.md)**: 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.