# finetune_pipeline.py
from kfp import dsl, compiler, kubernetes
BASE_MODEL = "Qwen/Qwen3-0.6B"
REGISTERED_NAME = "qwen3-0.6b-sft"
WORKSPACE = "finetune" # MLflow workspace (= namespace)
EXPERIMENT = "qwen3-0.6b-daily-sft"
PVC_NAME = "finetune-shared"
PVC_MOUNT = "/mnt/shared"
MLFLOW_URI = "http://mlflow-tracking-server.kubeflow:5000"
# Pre-built runtime image with torch (CUDA 12.6) + HF + trl + mlflow + kserve
# + kubernetes baked in. No pip install runs at pod start, so the same tag
# works on air-gapped clusters once it's mirrored into an internal registry.
RUNTIME_IMAGE = "docker.io/alaudadockerhub/finetune-pipeline-runtime-cu126-amd64:v0.1.0"
# ---------------------------------------------------------------------------
# 1. Fine-tune with the HF Trainer, autolog to MLflow, register the model.
# ---------------------------------------------------------------------------
@dsl.component(base_image=RUNTIME_IMAGE)
def fine_tune(
workspace: str,
experiment: str,
base_model: str,
registered_name: str,
pvc_mount: str,
run_id: str,
# Empty = generate a tiny synthetic instruction corpus in-process (no
# network egress; smoke-shape only, not a meaningful fine-tune). Non-empty
# = a local JSONL path (e.g. "/mnt/shared/datasets/alpaca.jsonl") or an
# "s3://…" URI to a JSONL file with a `text` column. See the "Using a
# real dataset" section for how to prepare the file.
dataset_path: str = "",
num_train_epochs: int = 1,
learning_rate: float = 2e-4,
per_device_train_batch_size: int = 2,
max_samples: int = 512,
) -> str:
"""Full SFT of `base_model` on a small synthetic corpus (or a JSONL file
at `dataset_path`). Logs metrics + registers the fine-tuned model.
Returns the model version as a string."""
import operator, os, random, mlflow, mlflow.transformers
from datasets import Dataset, load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments
from trl import SFTTrainer
# 1) MLflow (auth via MLFLOW_TRACKING_TOKEN injected from the Secret).
mlflow.set_tracking_uri(MLFLOW_URI)
mlflow.set_workspace(workspace)
mlflow.set_experiment(experiment)
mlflow.transformers.autolog(log_models=False) # models are logged manually below
# 2) Data. Synthetic by default so the pipeline's smoke run needs no
# network egress and no dataset pre-staging. Swap in a real dataset by
# passing `dataset_path` (JSONL file with a `text` column) — either a
# PVC path or an S3 URI. See "Using a real dataset" below.
if dataset_path.startswith("s3://"):
ds = load_dataset(
"json", data_files=dataset_path, split=f"train[:{max_samples}]",
# storage_options / credentials come from AWS_* env; see the docs.
)
elif dataset_path:
ds = load_dataset(
"json", data_files=dataset_path, split=f"train[:{max_samples}]",
)
else:
rng = random.Random(42)
ops = [("+", operator.add), ("-", operator.sub), ("*", operator.mul)]
rows = []
for _ in range(max_samples):
a, b = rng.randint(0, 99), rng.randint(0, 99)
sym, op = rng.choice(ops)
rows.append({"text":
f"### Instruction:\nWhat is {a} {sym} {b}?\n"
f"### Response:\nThe answer is {op(a, b)}."
})
ds = Dataset.from_list(rows)
# 3) Model & tokenizer.
tokenizer = AutoTokenizer.from_pretrained(base_model)
model = AutoModelForCausalLM.from_pretrained(base_model)
output_dir = f"{pvc_mount}/models/{registered_name}/{run_id}"
# 4) Train. `report_to="mlflow"` streams loss / eval_loss / lr per step.
args = TrainingArguments(
output_dir=output_dir,
num_train_epochs=num_train_epochs,
learning_rate=learning_rate,
per_device_train_batch_size=per_device_train_batch_size,
logging_steps=10,
save_strategy="no",
report_to="mlflow",
run_name=f"train-{run_id}",
)
with mlflow.start_run(run_name=f"pipeline-{run_id}") as parent:
mlflow.set_tag("pipeline_run_id", run_id)
mlflow.set_tag("base_model", base_model)
with mlflow.start_run(run_name=f"train-{run_id}", nested=True):
# TRL 0.12+ renamed the `tokenizer` kwarg to `processing_class`.
trainer = SFTTrainer(model=model, args=args, train_dataset=ds,
processing_class=tokenizer)
trainer.train()
trainer.save_model(output_dir)
# 5) Register the fine-tuned model.
# The artifact upload uses the MLflow plugin's artifact store (S3).
# KServe then pulls the model back from that S3 URI in the deploy step.
# `task="text-generation"` is required because MLflow cannot infer the
# task from a locally-loaded (non-Hub) model.
mv = mlflow.transformers.log_model(
transformers_model={"model": trainer.model, "tokenizer": tokenizer},
artifact_path="model",
registered_model_name=registered_name,
task="text-generation",
)
# 6) Record the PVC path on the model version so downstream steps
# can serve it from the shared PVC without another S3 round-trip.
client = mlflow.MlflowClient()
version = client.get_latest_versions(registered_name, stages=["None"])[0].version
client.set_model_version_tag(registered_name, version, "pvc_path", output_dir)
client.set_model_version_tag(registered_name, version, "pipeline_run_id", run_id)
print(f"registered {registered_name} version {version} @ {output_dir}")
return version
# ---------------------------------------------------------------------------
# 2. Deploy the just-registered model as a temporary KServe InferenceService.
# ---------------------------------------------------------------------------
@dsl.component(base_image=RUNTIME_IMAGE)
def deploy_for_evaluation(
workspace: str,
registered_name: str,
model_version: str,
pvc_name: str,
pvc_mount: str,
run_id: str,
) -> str:
"""Create an InferenceService that serves the model straight off the PVC.
Returns the InferenceService name."""
import time, mlflow
from kubernetes import client, config
from kserve import KServeClient, constants
from kserve import (V1beta1InferenceService, V1beta1InferenceServiceSpec,
V1beta1PredictorSpec, V1beta1ModelSpec, V1beta1ModelFormat)
# Resolve the PVC path from the tag set by fine_tune().
mlflow.set_tracking_uri(MLFLOW_URI); mlflow.set_workspace(workspace)
tags = mlflow.MlflowClient().get_model_version(registered_name, model_version).tags
pvc_path = tags["pvc_path"]
isvc_name = f"{registered_name}-eval-{run_id[-8:]}"
config.load_incluster_config()
ns = open("/var/run/secrets/kubernetes.io/serviceaccount/namespace").read().strip()
isvc = V1beta1InferenceService(
api_version=constants.KSERVE_GROUP + "/v1beta1", kind=constants.KSERVE_KIND,
metadata=client.V1ObjectMeta(name=isvc_name, namespace=ns, labels={
"app.kubernetes.io/part-of": "finetune-pipeline",
"pipeline-run-id": run_id,
}),
spec=V1beta1InferenceServiceSpec(predictor=V1beta1PredictorSpec(
model=V1beta1ModelSpec(
model_format=V1beta1ModelFormat(name="huggingface"),
runtime="kserve-huggingfaceserver",
storage_uri=f"pvc://{pvc_name}{pvc_path.removeprefix(pvc_mount)}",
resources=client.V1ResourceRequirements(
limits={"nvidia.com/gpu": "1", "memory": "16Gi"},
requests={"nvidia.com/gpu": "1", "memory": "8Gi"},
),
),
)),
)
KServeClient().create(isvc)
# Wait for READY.
deadline = time.time() + 15 * 60
while time.time() < deadline:
got = KServeClient().get(isvc_name, namespace=ns)
cond = {c["type"]: c["status"] for c in (got.get("status", {}).get("conditions") or [])}
if cond.get("Ready") == "True":
print(f"InferenceService {isvc_name} is Ready")
return isvc_name
time.sleep(10)
raise TimeoutError(f"{isvc_name} did not become Ready in 15 min")
# ---------------------------------------------------------------------------
# 3. Evaluate the InferenceService with a TrustyAI LMEvalJob, log to MLflow.
# ---------------------------------------------------------------------------
@dsl.component(base_image=RUNTIME_IMAGE)
def evaluate(
workspace: str,
experiment: str,
registered_name: str,
model_version: str,
isvc_name: str,
run_id: str,
tasks: list = ["arc_easy", "hellaswag"],
limit: str = "50",
) -> dict:
"""Create an LMEvalJob against the InferenceService, wait for it to
Complete, log the metrics into the parent MLflow run + as a nested run,
and tag the model version with the primary accuracy."""
import json, time, mlflow
from kubernetes import client, config
config.load_incluster_config()
ns = open("/var/run/secrets/kubernetes.io/serviceaccount/namespace").read().strip()
api = client.CustomObjectsApi()
job_name = f"eval-{isvc_name}"
body = dict(
apiVersion="trustyai.opendatahub.io/v1alpha1", kind="LMEvalJob",
metadata=dict(name=job_name, labels={"pipeline-run-id": run_id}),
spec=dict(
model="local-completions",
modelArgs=[
dict(name="model", value=isvc_name),
dict(name="base_url",
value=f"http://{isvc_name}-predictor.{ns}.svc/v1/completions"),
dict(name="num_concurrent", value="1"),
dict(name="max_retries", value="3"),
dict(name="tokenized_requests", value="True"),
dict(name="tokenizer", value="Qwen/Qwen3-0.6B"),
],
taskList=dict(taskNames=list(tasks)),
allowOnline=True, allowCodeExecution=False,
batchSize="1", limit=limit, logSamples=False,
chatTemplate=dict(enabled=False),
outputs=dict(pvcManaged=dict(size="100Mi")),
# The ta-lmes-job image runs as UID 65532; a fresh PVC is root-owned,
# so the driver's `open(stdout.log)` fails with permission denied
# unless fsGroup gives the pod's group write access to the mount.
pod=dict(securityContext=dict(fsGroup=65532)),
),
)
api.create_namespaced_custom_object(
"trustyai.opendatahub.io", "v1alpha1", ns, "lmevaljobs", body)
# Wait for a terminal state. `state` reaches `Complete` on both success
# and failure — the reason field is what actually distinguishes them.
deadline = time.time() + 60 * 60
while time.time() < deadline:
got = api.get_namespaced_custom_object(
"trustyai.opendatahub.io", "v1alpha1", ns, "lmevaljobs", job_name)
status = got.get("status") or {}
state, reason = status.get("state"), status.get("reason")
if state == "Complete" and reason == "Succeeded":
break
if state in {"Cancelled", "Failed"} or (state == "Complete" and reason in {"Failed", "Cancelled"}):
raise RuntimeError(
f"LMEvalJob {job_name} ended: state={state} reason={reason} "
f"message={status.get('message')}")
time.sleep(20)
else:
raise TimeoutError(f"LMEvalJob {job_name} did not finish in 1 h")
results = json.loads(got["status"]["results"])["results"]
# Log each task's metrics back to MLflow.
mlflow.set_tracking_uri(MLFLOW_URI)
mlflow.set_workspace(workspace)
mlflow.set_experiment(experiment)
parent = mlflow.search_runs(filter_string=f"tags.pipeline_run_id = '{run_id}'",
order_by=["attributes.start_time DESC"], max_results=1)
parent_run_id = parent.iloc[0]["run_id"]
flat = {}
with mlflow.start_run(run_id=parent_run_id):
with mlflow.start_run(run_name=f"eval-{run_id}", nested=True):
for task, metrics in results.items():
for k, v in metrics.items():
if isinstance(v, (int, float)) and "stderr" not in k:
name = f"{task}/{k.replace(',', '_')}"
mlflow.log_metric(name, float(v))
flat[name] = float(v)
mlflow.set_tag("model_version", model_version)
mlflow.set_tag("isvc_name", isvc_name)
# Tag the primary accuracy on the model version so it shows up on the
# Model Registry page next to the version number.
primary = next((k for k in flat if k.endswith("/acc_none")), None)
if primary is not None:
mlflow.MlflowClient().set_model_version_tag(
registered_name, model_version, "eval_acc", f"{flat[primary]:.4f}")
return flat
# ---------------------------------------------------------------------------
# 4. Cleanup: delete the temporary InferenceService and LMEvalJob.
# ---------------------------------------------------------------------------
@dsl.component(base_image=RUNTIME_IMAGE)
def cleanup(isvc_name: str):
from kubernetes import client, config
config.load_incluster_config()
ns = open("/var/run/secrets/kubernetes.io/serviceaccount/namespace").read().strip()
api = client.CustomObjectsApi()
for gv, plural, name in [
("serving.kserve.io/v1beta1", "inferenceservices", isvc_name),
("trustyai.opendatahub.io/v1alpha1", "lmevaljobs", f"eval-{isvc_name}"),
]:
group, version = gv.split("/")
try:
api.delete_namespaced_custom_object(group, version, ns, plural, name)
except Exception as exc:
print(f"delete {plural}/{name} skipped: {exc}")
# ---------------------------------------------------------------------------
# Pipeline: wire the four components together.
# ---------------------------------------------------------------------------
@dsl.pipeline(name="qwen3-sft-mlflow-trustyai",
description="Daily SFT + MLflow tracking + TrustyAI evaluation")
def sft_mlflow_trustyai(
workspace: str = WORKSPACE,
experiment: str = EXPERIMENT,
base_model: str = BASE_MODEL,
registered_name: str = REGISTERED_NAME,
# Default = synthetic in-process data (offline smoke). Override with a
# JSONL path on the shared PVC or an s3:// URI to fine-tune on a real
# dataset — see the "Using a real dataset" section.
dataset_path: str = "",
num_train_epochs: int = 1,
learning_rate: float = 2e-4,
max_samples: int = 512,
eval_limit: str = "50",
):
train = fine_tune(
workspace=workspace, experiment=experiment,
base_model=base_model, registered_name=registered_name,
pvc_mount=PVC_MOUNT, run_id=dsl.PIPELINE_JOB_ID_PLACEHOLDER,
dataset_path=dataset_path, num_train_epochs=num_train_epochs,
learning_rate=learning_rate, max_samples=max_samples,
).set_accelerator_type("nvidia.com/gpu").set_accelerator_limit(1)
train.set_memory_limit("32Gi").set_cpu_limit("8")
deploy = deploy_for_evaluation(
workspace=workspace, registered_name=registered_name,
model_version=train.output, pvc_name=PVC_NAME, pvc_mount=PVC_MOUNT,
run_id=dsl.PIPELINE_JOB_ID_PLACEHOLDER,
)
ev = evaluate(
workspace=workspace, experiment=experiment,
registered_name=registered_name, model_version=train.output,
isvc_name=deploy.output, run_id=dsl.PIPELINE_JOB_ID_PLACEHOLDER,
)
with dsl.ExitHandler(cleanup(isvc_name=deploy.output)):
ev # cleanup runs whether evaluate() succeeds or not
# Mount the shared PVC on training + deploy.
for task in (train, deploy):
kubernetes.mount_pvc(task, pvc_name=PVC_NAME, mount_path=PVC_MOUNT)
# Inject the Dex id token into every component that talks to MLflow.
for task in (train, deploy, ev):
kubernetes.use_secret_as_env(
task, secret_name="mlflow-token",
secret_key_to_env={"token": "MLFLOW_TRACKING_TOKEN"})
compiler.Compiler().compile(sft_mlflow_trustyai, "pipeline.yaml")
