Implementation:Dagster io Dagster OpenAI Fine Tuning Pattern
| Property | Value |
|---|---|
| Type | Implementation |
| Category | AI, NLP, Fine_Tuning |
| Repository | Dagster_io_Dagster |
| Implements | Principle:Dagster_io_Dagster_LLM_Fine_Tuning_Orchestration |
Overview
Concrete implementation pattern for OpenAI fine-tuning workflows using the dagster-openai resource and Dagster asset checks.
Description
This implementation demonstrates how to orchestrate an OpenAI fine-tuning workflow as Dagster assets. The OpenAIResource wraps the OpenAI client, the fine-tuning asset manages file uploads and asynchronous job polling, and an asset check validates the fine-tuned model by comparing its accuracy against the base model on a holdout dataset.
Usage
Define training and validation data assets in JSONL format, then wire the fine_tuned_model asset with an OpenAIResource. The asset handles file upload, job creation, and polling. Attach an asset check to validate model quality after fine-tuning completes.
Code Reference
Source Location
examples/docs_projects/project_llm_fine_tune/src/project_llm_fine_tune/defs/assets.py:L340-514
Signature/Pattern
Fine-tuning asset with polling:
from dagster_openai import OpenAIResource
import dagster as dg
MODEL_NAME = "gpt-4o-mini-2024-07-18"
@dg.asset(kinds={"openai"}, group_name="training")
def fine_tuned_model(context: dg.AssetExecutionContext, openai: OpenAIResource, training_file, validation_file) -> str:
client = openai.get_client(context)
# Upload training files
train_file = client.files.create(file=open(training_file, "rb"), purpose="fine-tune")
val_file = client.files.create(file=open(validation_file, "rb"), purpose="fine-tune")
# Create fine-tuning job
job = client.fine_tuning.jobs.create(
training_file=train_file.id,
validation_file=val_file.id,
model=MODEL_NAME,
suffix="goodreads",
)
# Poll for completion
while True:
status = client.fine_tuning.jobs.retrieve(job.id)
if status.status == "succeeded":
break
time.sleep(30)
model_name = status.fine_tuned_model
context.add_output_metadata({"model_name": model_name})
return model_name
Model validation via asset check:
@dg.asset_check(
asset=fine_tuned_model,
additional_ins={"data": dg.AssetIn("enriched_graphic_novels")},
)
def validate_fine_tuned_model(context, openai: OpenAIResource, data: pd.DataFrame) -> dg.AssetCheckResult:
client = openai.get_client(context)
# Compare fine-tuned vs base model accuracy
ft_accuracy = evaluate_model(client, context.get_input("fine_tuned_model"), data)
base_accuracy = evaluate_model(client, MODEL_NAME, data)
return dg.AssetCheckResult(
passed=ft_accuracy >= base_accuracy,
severity=dg.AssetCheckSeverity.WARN,
metadata={"fine_tuned_accuracy": ft_accuracy, "base_accuracy": base_accuracy},
)
Import
from dagster_openai import OpenAIResource
import dagster as dg
I/O Contract
| Direction | Name | Type | Description |
|---|---|---|---|
| Input | training_file | str (path) | Path to JSONL training data file |
| Input | validation_file | str (path) | Path to JSONL validation data file |
| Input | openai | OpenAIResource | Dagster resource wrapping the OpenAI client |
| Input | data (for check) | pd.DataFrame | Holdout dataset for model evaluation |
| Output | fine_tuned_model | str | Name/ID of the fine-tuned model |
| Output | validation result | AssetCheckResult | Pass/fail with accuracy metadata |
Usage Examples
Defining the fine-tuning pipeline:
import dagster as dg
from dagster_openai import OpenAIResource
defs = dg.Definitions(
assets=[training_file, validation_file, fine_tuned_model],
asset_checks=[validate_fine_tuned_model],
resources={
"openai": OpenAIResource(api_key=dg.EnvVar("OPENAI_API_KEY")),
},
)
Monitoring fine-tuning progress:
The asset logs metadata including the fine-tuned model name once the job completes. The asset check runs automatically after materialization and reports accuracy metrics in the Dagster UI.