# 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"
# ---------------------------------------------------------------------------
# 1. Fine-tune with the HF Trainer, autolog to MLflow, register the model.
# ---------------------------------------------------------------------------
@dsl.component(
base_image="python:3.11-slim",
packages_to_install=[
"mlflow>=3.10",
"transformers>=4.44",
"datasets>=2.20",
"accelerate>=0.34",
"trl>=0.10",
"torch>=2.3",
],
)
def fine_tune(
workspace: str,
experiment: str,
base_model: str,
registered_name: str,
pvc_mount: str,
run_id: str,
dataset_id: str = "tatsu-lab/alpaca",
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 `dataset_id`. Logs metrics + registers the
fine-tuned model. Returns the model version string (an int as str)."""
import os, mlflow, mlflow.transformers
from datasets import 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.
ds = load_dataset(dataset_id, split=f"train[:{max_samples}]")
# 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):
trainer = SFTTrainer(model=model, args=args, train_dataset=ds,
tokenizer=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.
mv = mlflow.transformers.log_model(
transformers_model={"model": trainer.model, "tokenizer": tokenizer},
artifact_path="model",
registered_model_name=registered_name,
)
# 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="python:3.11-slim",
packages_to_install=["kserve>=0.13", "kubernetes>=29", "mlflow>=3.10"],
)
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="python:3.11-slim",
packages_to_install=["kubernetes>=29", "mlflow>=3.10"],
)
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="python:3.11-slim",
packages_to_install=["kubernetes>=29"],
)
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,
dataset_id: str = "tatsu-lab/alpaca",
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_id=dataset_id, 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")