Full SFT Customization

import os
os.environ["NHX_BASE_URL"] = "<YOUR_NHX_BASE_URL>"

import json
import os
from nemo_helix_plugin.client.client import NemoClient
from nemo_helix_plugin.client.errors import ConflictError
from nemo_helix_plugin.files.types import CreateFilesetRequest
from nemo_helix_plugin.models.types import CreateModelEntityRequest, UpdateModelEntityRequest
from nemo_helix_plugin.schema import SecretRef
from nemo_helix_plugin.secrets.types import HelixSecretCreateRequest
from nemo_customizer.sdk import Customization
from filesets import FilesetFileSystem
from nemo_helix_plugin.files.client import FilesClient
from pydantic import SecretStr
client = NemoClient.from_config()
customization = Customization.from_client(client)

from pathlib import Path
from datasets import load_dataset, DatasetDict
import json
# Load the SQuAD dataset from Hugging Face
print("Loading dataset rajpurkar/squad")
raw_dataset = load_dataset("rajpurkar/squad")
if not isinstance(raw_dataset, DatasetDict):
    raise ValueError("Dataset does not contain expected splits")
print("Loaded dataset")
# Configuration
VALIDATION_PROPORTION = 0.05
SEED = 1234
# For the purpose of this tutorial, we'll use a subset of the dataset
# The larger the datasets, the better the model will perform but longer the training will take
training_size = 3000
validation_size = 300
DATASET_PATH = Path("sft-dataset").absolute()
# Create directory if it doesn't exist
os.makedirs(DATASET_PATH, exist_ok=True)
# Get the train split and create a validation split from it
train_set = raw_dataset.get('train')
split_dataset = train_set.train_test_split(test_size=VALIDATION_PROPORTION, seed=SEED)
# Select subsets for the tutorial
train_ds = split_dataset['train'].select(range(min(training_size, len(split_dataset['train']))))
validation_ds = split_dataset['test'].select(range(min(validation_size, len(split_dataset['test']))))
# Convert SQuAD to chat format and save to JSONL
def convert_squad_to_chat_format(example):
    """Convert SQuAD to a chat `messages` row whose user turn matches the inference prompt."""
    user = f"Based on the following context, answer the question.\n\nContext: {example['context']}\n\nQuestion: {example['question']}"
    answer = example["answers"]["text"][0]  # Take the first answer
    return {"messages": [{"role": "user", "content": user}, {"role": "assistant", "content": answer}]}
# Save training data
with open(f"{DATASET_PATH}/training.jsonl", "w", encoding="utf-8") as f:
    for example in train_ds:
        converted = convert_squad_to_chat_format(example)
        f.write(json.dumps(converted) + "\n")
# Save validation data
with open(f"{DATASET_PATH}/validation.jsonl", "w", encoding="utf-8") as f:
    for example in validation_ds:
        converted = convert_squad_to_chat_format(example)
        f.write(json.dumps(converted) + "\n")
print(f"Saved training.jsonl with {len(train_ds)} rows")
print(f"Saved validation.jsonl with {len(validation_ds)} rows")
# Show a sample from the training data
print("\nSample from training data:")
with open(f"{DATASET_PATH}/training.jsonl", 'r') as f:
    first_line = f.readline()
    sample = json.loads(first_line)
    print(f"User: {sample['messages'][0]['content'][:200]}...")
    print(f"Assistant: {sample['messages'][1]['content']}")

# Create fileset to store SFT training data
DATASET_NAME = "sft-dataset"
try:
    client.files.create_fileset(
        body=CreateFilesetRequest(
            name=DATASET_NAME,
            description="SFT training data",
        ),
    ).data()
    print(f"Created fileset: {DATASET_NAME}")
except ConflictError:
    print(f"Fileset '{DATASET_NAME}' already exists, continuing...")
# Upload training data files individually to ensure correct structure
fs = FilesetFileSystem(client=FilesClient.from_client(client))
fs.put(
    lpath=f"{DATASET_PATH}/",  # Trailing slash uploads directory contents to fileset root
    rpath=f"default/{DATASET_NAME}/",
    recursive=True,
)
# Validate training data is uploaded correctly
print("Training data:")
print(json.dumps([f.model_dump() for f in client.files.list_files(name=DATASET_NAME).data().data], indent=2))

# Export the HF_TOKEN and NGC_API_KEY environment variables if they are not already set.
# This tutorial's default model (Qwen/Qwen3-0.6B) is public, so HF_TOKEN is optional.
HF_TOKEN = os.getenv("HF_TOKEN")
NGC_API_KEY = os.getenv("NGC_API_KEY")
def create_or_get_secret(name: str, value: str | None, label: str):
    if not value:
        print(f"{label} is not set - skipping setting secret")
        return None
    try:
        secret = client.secrets.create_secret(
            body=HelixSecretCreateRequest(name=name, value=SecretStr(value)),
        ).data()
        print(f"Created secret: {name}")
        return secret
    except ConflictError:
        print(f"Secret '{name}' already exists, continuing...")
        return client.secrets.get_secret(name=name).data()
# Create Hugging Face token secret
hf_secret = create_or_get_secret("hf-token", HF_TOKEN, "HF_TOKEN")
if hf_secret:
    print("HF_TOKEN secret:")
    print(hf_secret.model_dump_json(indent=2))
# Create NGC API key secret
# Uncomment the line below if you have NGC API Key and want to finetune NGC models
# ngc_api_key = create_or_get_secret("ngc-api-key", NGC_API_KEY, "NGC_API_KEY")

import time
# Create a fileset pointing to the desired Hugging Face model
from nemo_helix_plugin.files.storage_config import HuggingfaceStorageConfig
HF_REPO_ID = "Qwen/Qwen3-0.6B"
MODEL_NAME = "qwen3-0.6b"
try:
    base_model_fs = client.files.create_fileset(
        body=CreateFilesetRequest(
            name=MODEL_NAME,
            description="Qwen3 0.6B base model from Hugging Face",
            storage=HuggingfaceStorageConfig(
                # repo_id is the full model name from Hugging Face
                repo_id=HF_REPO_ID,
                repo_type="model",
                # we use the secret created in the previous step, if any
                token_secret=SecretRef(hf_secret.name) if hf_secret else None,
            ),
        ),
    ).data()
    print(f"Created base model fileset: {MODEL_NAME}")
except ConflictError:
    print(f"Base model fileset already exists. Skipping creation.")
    base_model_fs = client.files.get_fileset(name=MODEL_NAME).data()
# Create the Model Entity representation.
try:
    base_model = client.models.create_model(
        body=CreateModelEntityRequest(
            name=MODEL_NAME,
            fileset=f"default/{MODEL_NAME}",
        ),
    ).data()
    print(f"Created Model Entity: {MODEL_NAME}")
except ConflictError:
    print(f"Base model already exists. Updating fileset if different.")
    base_model = client.models.update_model(
        name=MODEL_NAME,
        body=UpdateModelEntityRequest(
            fileset=f"default/{MODEL_NAME}",
        ),
    ).data()
print(f"\nBase model fileset: fileset://default/{base_model.name}")
print("Base model fileset files list:")
print(json.dumps([f.model_dump() for f in client.files.list_files(name=MODEL_NAME).data().data], indent=2))
# Wait for ModelSpec to be populated from the checkpoint
print("\nWaiting for ModelSpec to be populated...")
SPEC_TIMEOUT_SECONDS = 120
spec_start = time.time()
while not base_model.spec:
    if time.time() - spec_start > SPEC_TIMEOUT_SECONDS:
        raise TimeoutError(f"ModelSpec not populated within {SPEC_TIMEOUT_SECONDS} seconds")
    time.sleep(2)
    base_model = client.models.get_model(name=MODEL_NAME).data()
print(f"ModelSpec populated: {base_model.spec}")

import uuid
from nemo_automodel_plugin.schema import AutomodelJobInput
job_suffix = uuid.uuid4().hex[:4]
JOB_NAME = f"my-sft-job-{job_suffix}"
OUTPUT_NAME = f"sft-model-{job_suffix}"
spec = AutomodelJobInput(
    model=f"default/{base_model.name}",
    dataset={"training": f"default/{DATASET_NAME}"},
    training={
        "training_type": "sft",
        "finetuning_type": "all_weights",
        "max_seq_length": 2048,
    },
    schedule={"epochs": 2},
    batch={"global_batch_size": 64, "micro_batch_size": 8},
    optimizer={"learning_rate": 1e-5, "weight_decay": 0.0},
    parallelism={
        "num_gpus_per_node": 1,
        "num_nodes": 1,
        "tensor_parallel_size": 1,
        "pipeline_parallel_size": 1,
    },
    output={"name": OUTPUT_NAME},
)
job = customization.automodel.jobs.create(
    spec=spec, workspace="default", name=JOB_NAME
)
print(f"Submitted job: {job.job.name}")
print(f"Output model: {OUTPUT_NAME}")

import time
# Poll the job until it reaches a terminal state
while True:
    status = client.jobs.get_job_status(name=job.job.name, workspace="default").data()
    print(f"Status: {status.status}")
    if status.status.lower() in ("completed", "failed", "cancelled", "error"):
        break
    time.sleep(15)
assert status.status.lower() == "completed"

# Validate model entity exists
model_entity = client.models.get_model(name=OUTPUT_NAME).data()
print(model_entity.model_dump_json(indent=2))

# Create deployment config
deploy_suffix = uuid.uuid4().hex[:4]
DEPLOYMENT_CONFIG_NAME = f"sft-model-deployment-cfg-{deploy_suffix}"
DEPLOYMENT_NAME = f"sft-model-deployment-{deploy_suffix}"
from nemo_helix_plugin.models.types import (
    ContainerExecutorConfig,
    CreateModelDeploymentConfigRequest,
    CreateModelDeploymentRequest,
    Engine,
    ModelDeploymentConfigModelSpec,
)
deployment_config = client.models.create_deployment_config(
    body=CreateModelDeploymentConfigRequest(
        name=DEPLOYMENT_CONFIG_NAME,
        engine=Engine.VLLM,
        model_spec=ModelDeploymentConfigModelSpec(
            model_namespace="default",
            model_name=OUTPUT_NAME,
        ),
        executor_config=ContainerExecutorConfig(
            gpu=1,
            image_name="docker.io/vllm/vllm-openai",
            image_tag="v0.22.1",
        ),
    ),
).data()
# Deploy model using deployment_config created above
deployment = client.models.create_deployment(
    body=CreateModelDeploymentRequest(
        name=DEPLOYMENT_NAME,
        config=deployment_config.name,
    ),
).data()
# Check deployment status
deployment_status = client.models.get_deployment(name=deployment.name).data()
print(f"Deployment name: {deployment.name}")
print(f"Deployment status: {deployment_status.status}")

from nemo_helix_plugin.models.types import (
    ContainerExecutorConfig,
    CreateModelDeploymentConfigRequest,
    Engine,
    ModelDeploymentConfigModelSpec,
)
client.models.create_deployment_config(
    body=CreateModelDeploymentConfigRequest(
        name="sft-model-config-multigpu",
        engine=Engine.VLLM,
        model_spec=ModelDeploymentConfigModelSpec(
            model_namespace="default",
            model_name=OUTPUT_NAME,
        ),
        executor_config=ContainerExecutorConfig(
            gpu=2,
            image_name="docker.io/vllm/vllm-openai",
            image_tag="v0.22.1",
        ),
    ),
).data()

import time
from IPython.display import clear_output
# Poll deployment status every 15 seconds until ready
TIMEOUT_MINUTES = 30
start_time = time.time()
timeout_seconds = TIMEOUT_MINUTES * 60
print(f"Monitoring deployment '{deployment.name}'...")
print(f"Timeout: {TIMEOUT_MINUTES} minutes\n")
while True:
    deployment_status = client.models.get_deployment(name=deployment.name).data()
    elapsed = time.time() - start_time
    elapsed_min = int(elapsed // 60)
    elapsed_sec = int(elapsed % 60)
    clear_output(wait=True)
    print(f"Deployment: {deployment.name}")
    print(f"Status: {deployment_status.status}")
    print(f"Elapsed time: {elapsed_min}m {elapsed_sec}s")
    # Check if deployment is ready
    if deployment_status.status == "READY":
        print("\nDeployment is ready!")
        if not client.models.wait_for_gateway(deployment.name, workspace="default", timeout=60):
            raise RuntimeError("Inference gateway did not become ready")
        break
    # Check for failure states
    if deployment_status.status in ("FAILED", "ERROR", "TERMINATED", "LOST"):
        raise RuntimeError(f"Deployment failed with status: {deployment_status.status}")
    # Check timeout
    if elapsed > timeout_seconds:
        raise TimeoutError(f"Deployment timeout after {TIMEOUT_MINUTES} minutes")
    time.sleep(15)

# Wait for deployment to be ready, then test
# Test the fine-tuned model with a question answering prompt
context = "The Apollo 11 mission was the first manned mission to land on the Moon. It was launched on July 16, 1969, and Neil Armstrong became the first person to walk on the lunar surface on July 20, 1969. Buzz Aldrin joined him shortly after, while Michael Collins remained in lunar orbit."
question = "Who was the first person to walk on the Moon?"
messages = [
    {"role": "user", "content": f"Based on the following context, answer the question.\n\nContext: {context}\n\nQuestion: {question}"}
]
from nemo_helix_plugin.inference_gateway.client import InferenceGatewayClient
from nemo_helix_plugin.inference_gateway.types import JsonBody
response = InferenceGatewayClient.from_client(client).provider_post(
    name=deployment.name,
    workspace="default",
    trailing_uri="v1/chat/completions",
    body=JsonBody(
        {
            "model": f"default/{OUTPUT_NAME}",
            "messages": messages,
            "temperature": 0,
            "max_tokens": 128,
            # Qwen3 reasons by default; turn it off for a short answer.
            "chat_template_kwargs": {"enable_thinking": False},
        }
    ),
).data()
print("=" * 60)
print("MODEL EVALUATION")
print("=" * 60)
print(f"Question: {question}")
print(f"Expected: Neil Armstrong")
print(f"Model output: {response['choices'][0]['message']['content']}")