From 15a7d0b5cf67624438800374506dd365943b8156 Mon Sep 17 00:00:00 2001 From: Sam Oluwalana Date: Tue, 21 Jul 2026 14:18:44 -0600 Subject: [PATCH 01/24] docs(customizer): align guides with current APIs and workflows Resolve accumulated Customizer documentation drift across API examples, tutorials, model references, and Fern rendering. - align FileSet, ModelSpec, job, metrics, and inference examples with current schemas - add tutorial failure guards and correct deployment readiness handling - update GPT-OSS, storage, Hugging Face, Unsloth, and CUDA 13 guidance - fix internal links, terminology, JSON examples, titles, and stale snippets - synchronize source notebooks and generated Fern artifacts - add validation for missing NotebookViewer artifacts Validation: - make docs-check - npm run broken-links Signed-off-by: Sam Oluwalana --- docs/AGENTS.md | 4 +- docs/customizer/about.mdx | 22 +- docs/customizer/index.mdx | 10 +- .../manage-customization-jobs/create-job.mdx | 10 +- .../get-job-status.mdx | 6 +- .../hyperparameters.mdx | 35 +- .../manage-customization-jobs/index.mdx | 2 +- .../list-active-jobs.mdx | 19 +- .../manage-model-entities/create-fileset.mdx | 24 +- .../create-model-entity.mdx | 4 +- .../manage-model-entities/index.mdx | 10 +- docs/customizer/models/data-format.mdx | 57 +- docs/customizer/models/embedding.mdx | 15 +- docs/customizer/models/gpt-oss.mdx | 22 +- docs/customizer/models/index.mdx | 6 +- docs/customizer/models/llama-nemotron.mdx | 8 +- docs/customizer/models/llama.mdx | 6 +- docs/customizer/models/mistral.mdx | 6 +- docs/customizer/models/phi.mdx | 2 +- docs/customizer/models/qwen.mdx | 4 +- .../output/chat-basic-format-example.jsonl | 1 - .../output/chat-expanded-format-example.json | 14 - .../output/chat-thinking-off-example.jsonl | 5 - .../output/chat-thinking-on-example.jsonl | 5 - .../chat-tool-calling-basic-example.jsonl | 1 - .../chat-tool-calling-expanded-example.json | 35 - .../output/completion-format-example.jsonl | 1 - .../_snippets/output/config-list-example.json | 95 - .../output/job-creation-response.json | 25 - .../output/lora-hyperparameters-example.json | 13 - .../output/lora-training-data-example.jsonl | 2 - .../output/version-error-example.json | 1 - .../distillation-customization-job.ipynb | 1860 ++++++++-------- .../distillation-customization-job.mdx | 24 +- .../tutorials/dpo-customization-job.ipynb | 1035 ++++----- .../embedding-customization-job.ipynb | 66 +- .../tutorials/embedding-customization-job.mdx | 60 +- .../tutorials/format-training-dataset.mdx | 6 +- docs/customizer/tutorials/import-hf-model.mdx | 31 +- .../tutorials/lora-customization-job.ipynb | 15 +- .../tutorials/lora-customization-job.mdx | 13 +- docs/customizer/tutorials/metrics.mdx | 16 +- .../tutorials/optimize-throughput.ipynb | 1952 +++++++++-------- .../tutorials/optimize-throughput.mdx | 80 +- .../tutorials/sft-customization-job.ipynb | 40 +- .../tutorials/sft-customization-job.mdx | 35 +- .../understand-configurations-and-models.mdx | 70 +- docs/fern/README.md | 6 +- .../distillation-customization-job.json | 16 +- .../distillation-customization-job.ts | 16 +- .../notebooks/dpo-customization-job.json | 155 ++ .../notebooks/dpo-customization-job.ts | 159 ++ .../embedding-customization-job.json | 40 +- .../notebooks/embedding-customization-job.ts | 40 +- .../notebooks/lora-customization-job.json | 16 +- .../notebooks/lora-customization-job.ts | 16 +- .../notebooks/optimize-throughput.json | 44 +- .../notebooks/optimize-throughput.ts | 44 +- .../notebooks/sft-customization-job.json | 32 +- .../notebooks/sft-customization-job.ts | 32 +- docs/fern/package.json | 3 +- docs/fern/scripts/README.md | 3 + docs/fern/scripts/ipynb-to-mdx.py | 6 + .../fern/scripts/validate-notebook-viewer.mjs | 126 ++ docs/requirements.mdx | 4 +- docs/support-matrix.mdx | 6 +- docs/troubleshooting/customizer.mdx | 13 +- 67 files changed, 3452 insertions(+), 3098 deletions(-) delete mode 100644 docs/customizer/tutorials/_snippets/output/chat-basic-format-example.jsonl delete mode 100644 docs/customizer/tutorials/_snippets/output/chat-expanded-format-example.json delete mode 100644 docs/customizer/tutorials/_snippets/output/chat-thinking-off-example.jsonl delete mode 100644 docs/customizer/tutorials/_snippets/output/chat-thinking-on-example.jsonl delete mode 100644 docs/customizer/tutorials/_snippets/output/chat-tool-calling-basic-example.jsonl delete mode 100644 docs/customizer/tutorials/_snippets/output/chat-tool-calling-expanded-example.json delete mode 100644 docs/customizer/tutorials/_snippets/output/completion-format-example.jsonl delete mode 100644 docs/customizer/tutorials/_snippets/output/config-list-example.json delete mode 100644 docs/customizer/tutorials/_snippets/output/job-creation-response.json delete mode 100644 docs/customizer/tutorials/_snippets/output/lora-hyperparameters-example.json delete mode 100644 docs/customizer/tutorials/_snippets/output/lora-training-data-example.jsonl delete mode 100644 docs/customizer/tutorials/_snippets/output/version-error-example.json create mode 100644 docs/fern/components/notebooks/dpo-customization-job.json create mode 100644 docs/fern/components/notebooks/dpo-customization-job.ts create mode 100644 docs/fern/scripts/validate-notebook-viewer.mjs diff --git a/docs/AGENTS.md b/docs/AGENTS.md index 120fc034b4..a5a83d7239 100644 --- a/docs/AGENTS.md +++ b/docs/AGENTS.md @@ -11,7 +11,7 @@ Run these from the repo root (they wrap `cd docs/fern && npm run …`): | `make docs-deps` | Install docs tooling (first run on a machine) | | `make docs` | Local dev server (live preview) | | `make docs-watch` | Local dev server plus repo-level watcher for `docs/**` changes outside `docs/fern/` | -| `make docs-check` | `fern check` + MDX validation + gated-link check (what CI runs) | +| `make docs-check` | `fern check` + MDX validation + NotebookViewer artifact validation + gated-link check (what CI runs) | | `make docs-check-python-snippets DOCS_PATH=...` | Syntax-check and type-check Python fenced snippets in one doc | | `make docs-run-notebook DOCS_PATH=...` | Execute the source notebook for one Fern `.mdx`/`.ipynb` doc using `nemo-nb` markers | | `make docs-broken-links` | Report broken links | @@ -24,7 +24,7 @@ Use `make docs` when you are only editing `docs/fern/` config. Use `make docs-wa ## Rules that bite if you miss them - **Navigation is the build.** Fern only builds pages listed in `docs/fern/versions/latest.yml`. A `.mdx` not in the nav is **not built** (404, not indexed) — that is how unready features are gated. Do **not** use `hidden: true` for gating (it still builds/serves the page). -- **Gated (unready) features** stay in the repo but out of the nav: `auth/`, `customizer/`, `safe-synthesizer/`, `evaluator/benchmarks/`, and a few individual pages. Ready-to-paste nav blocks for re-publishing are in `docs/fern/gated-nav.yml`. To publish one: move its block into `latest.yml`, re-add inbound links, run `make docs-check && make docs-broken-links`. +- **Publication state is nav-derived.** Do not maintain or rely on a hard-coded list of gated directories. Check `docs/fern/versions/latest.yml`: listed pages are published, and omitted pages are gated. `docs/fern/gated-nav.yml` contains reference blocks for some gated features. To publish one: move its block into `latest.yml`, re-add inbound links, run `make docs-check && make docs-broken-links`. - **Don't link into gated pages.** A link from a published page into a gated page is a dead link. `make docs-check` fails on it; `make docs-fix-links` delinks it to plain text. (Replaces the old MkDocs `hide_unready_docs` auto-delinking.) - **Internal links** use canonical nav URLs like `/documentation/get-started/core-concepts/workspaces`, not relative `.md`/source paths. `make docs-broken-links` is the check. - **No `{{variable}}` substitutions.** Fern has no substitution step; product names are inlined as literal text. (Prompt-template tokens like `` `{{input}}` `` inside backticks are real content — leave them.) diff --git a/docs/customizer/about.mdx b/docs/customizer/about.mdx index b7cf983f24..e46563195f 100644 --- a/docs/customizer/about.mdx +++ b/docs/customizer/about.mdx @@ -16,28 +16,28 @@ Full SFT models require a NIM deployment to serve inference. The Deployment Mana | Deployment Mode | Image Type | Weight Loading | Best For | |-----------------|------------|----------------|----------| -| **Multi-LLM** (Default) | Generic multi-model NIM | On-the-fly download via Files service | Any HF model, custom fine-tuned models, development | +| **Multi-LLM** (Default) | Generic multi-model NIM | On-the-fly download via Files service | Supported Hugging Face architectures, custom fine-tuned models, development | | **Model-Specific NIM** | Dedicated model image | Pre-download via model puller | Production, optimized performance and latency | -- **Multi-LLM Image**: Can deploy any HuggingFace-compatible model, providing maximum flexibility for custom fine-tuned models. Does not guarantee optimized inference performance. +- **Multi-LLM Image**: Can deploy Hugging Face checkpoints whose architectures are supported by the image's inference engine, providing flexibility for custom fine-tuned models. Importing a checkpoint does not guarantee training or deployment compatibility; for example, Automodel LoRA does not support Conv1D-based architectures. Does not guarantee optimized inference performance. - **Model-Specific NIM**: Provides optimized inference performance and latency through model-specific optimizations. Recommended for production deployments where performance is critical. ## Parameter-Efficient Fine-Tuning -Parameter-Efficient Fine-Tuning (PEFT) methods enable efficient model customization by training a small number of parameters while keeping the base model frozen. For example, when customizing LLaMa 3.3 70B: +Parameter-Efficient Fine-Tuning (PEFT) methods enable efficient model customization by training a small number of parameters while keeping the base model frozen. For example, when customizing Llama 3.3 70B: -- **Traditional SFT**: Trains and stores ~40 GB per task. -- **PEFT**: Trains and stores only a few MB per task while maintaining comparable performance. +- **Traditional SFT**: Produces a full BF16 checkpoint of approximately 140 GB per task. During training, budget free disk space separately for the base checkpoint, intermediate checkpoint, and final output—approximately 3× the downloaded base checkpoint size. +- **PEFT**: Produces an adapter that is typically approximately 100–500 MB per task while maintaining comparable performance. During training, budget approximately 1.5× the downloaded base checkpoint size. ```mermaid --- caption: Traditional Fine-Tuning --- flowchart TD - T1[Task 1] --> M1[LLaMa 3.3 - 70B] - T2[Task 2] --> M2[LLaMa 3.3 - 70B] - T3[Task 3] --> M3[LLaMa 3.3 - 70B] + T1[Task 1] --> M1[Llama 3.3 - 70B] + T2[Task 2] --> M2[Llama 3.3 - 70B] + T3[Task 3] --> M3[Llama 3.3 - 70B] style M1 fill:#B8D5F2 style M2 fill:#B8D5F2 @@ -47,13 +47,13 @@ flowchart TD ```mermaid --- -caption: Parameter-Efficient Fine Tuning +caption: Parameter-Efficient Fine-Tuning --- flowchart TD P1[Task 1] --> A1[LoRA 113M] P2[Task 2] --> A2[LoRA 20M] P3[Task 3] --> A3[LoRA 10M] - A1 & A2 & A3 --> M[LLaMa 3.3 - 70B] + A1 & A2 & A3 --> M[Llama 3.3 - 70B] style A1 fill:#B8D5F2 style A2 fill:#B8D5F2 @@ -215,7 +215,7 @@ Common hyperparameters you'll tune include: -NeMo Customizer offers **two training backends** — Automodel (multi-GPU) and Unsloth (single-GPU, quantized) — and each accepts its own job configuration. The exact field names, defaults, and available knobs differ between them. For the full per-backend hyperparameter reference, see [Training Configuration](/documentation/customizer-reference/manage-customization-jobs/training-configuration). +NeMo Customizer offers **two training backends** — Automodel (multi-GPU) and Unsloth (single-GPU, with optional quantized loading for LoRA) — and each accepts its own job configuration. Unsloth full-weight training requires unquantized model loading. The exact field names, defaults, and available knobs differ between them. For the full per-backend hyperparameter reference, see [Training Configuration](/documentation/customizer-reference/manage-customization-jobs/training-configuration). diff --git a/docs/customizer/index.mdx b/docs/customizer/index.mdx index 488229a01d..4b405ec48b 100644 --- a/docs/customizer/index.mdx +++ b/docs/customizer/index.mdx @@ -68,7 +68,7 @@ View the available Phi models from Microsoft, designed for strong reasoning capa -View the available GPT-OSS models supported for Full SFT customization. +View the available GPT-OSS models supported for Full SFT and LoRA customization. @@ -118,21 +118,21 @@ Learn how to format datasets for different model types. datasets chat-models completion-models - + Learn how to start a LoRA customization job using a custom dataset. nemo-customizer - + Learn how to start a SFT customization job using a custom dataset. nemo-customizer - + Learn how to compress a larger teacher model into a smaller student model. @@ -146,7 +146,7 @@ Learn how to check job metrics using MLFlow or Weights & Biases. nemo-customizer mlflow wandb - + Learn how to optimize the token-per-GPU throughput for a LoRA optimization job. diff --git a/docs/customizer/manage-customization-jobs/create-job.mdx b/docs/customizer/manage-customization-jobs/create-job.mdx index 90d3940d7d..441cf038ee 100644 --- a/docs/customizer/manage-customization-jobs/create-job.mdx +++ b/docs/customizer/manage-customization-jobs/create-job.mdx @@ -9,7 +9,7 @@ Customization jobs are submitted to one of two backends. Choose the backend that | Backend | Best for | Methods | |---------|----------|---------| | **Automodel** (default) | Production fine-tuning, larger models, multi-GPU scaling | SFT, distillation; LoRA, merged-LoRA, or full-weight | -| **Unsloth** | Memory-constrained single-GPU LoRA | SFT; LoRA or full-weight, with 4-bit / 8-bit loading | +| **Unsloth** | Memory-constrained single-GPU training | SFT; LoRA with optional 4-bit / 8-bit loading, or unquantized full-weight | ## Prerequisites @@ -19,7 +19,7 @@ Before you can create a customization job, make sure that you have: - Created a [FileSet and Model Entity](/documentation/customizer-reference/manage-model-entities/overview) for your base model. - [Uploaded a dataset](/documentation/get-started/core-concepts/manage-files) as a FileSet. - Determined the [training configuration](/documentation/customizer-reference/manage-customization-jobs/training-configuration) you want to use for the customization job. -- Verified that the platform has sufficient storage for the job. Full SFT jobs require approximately 3× the base model size in free disk space; LoRA jobs require approximately 1.5×. See [ft-tut-understand-models](/documentation/customizer-reference/tutorials/understanding-models-and-training) for details. If you are also deploying the model from a base checkpoint fileset, plan for ~2.5× model size overall for LoRA. +- Verified that the platform has sufficient storage for the job. Budget against the downloaded base checkpoint size: Full SFT jobs require approximately 3× in free disk space, and LoRA jobs require approximately 1.5×. See [ft-tut-understand-models](/documentation/customizer-reference/tutorials/understanding-models-and-training) for details. Include any retained deployment copies separately. - Set the `NMP_BASE_URL` environment variable to your NeMo Platform endpoint. ```bash @@ -67,13 +67,15 @@ print(f"Submitted job: {job.job.name}") print(f"Job status: {job.job.status}") ``` +The response preserves the explicit `name`. If you omit `name`, the platform generates a backend-prefixed job name. + :open: ```json { - "name": "automodel-a1b2c3d4e5f6", + "name": "my-lora-job", "workspace": "default", "id": "platform-job-2k8i3i1HqJHHPVB5M6Bk9Z", "status": "queued", @@ -104,7 +106,7 @@ print(f"Job status: {job.job.status}") ## Submit an Unsloth Job -The Unsloth backend runs on a single GPU and supports 4-bit / 8-bit quantized loading. Build a `UnslothJobInput` spec and submit it to the `unsloth` backend. Note that Unsloth uses its own field names (`model.name`, `dataset.path`, `batch.per_device_train_batch_size`). +The Unsloth backend runs on a single GPU and supports 4-bit / 8-bit quantized loading for LoRA. Full-weight training requires `model.load_in_4bit=false` and `model.load_in_8bit=false`. Build a `UnslothJobInput` spec and submit it to the `unsloth` backend. Note that Unsloth uses its own field names (`model.name`, `dataset.path`, `batch.per_device_train_batch_size`). ```python import os diff --git a/docs/customizer/manage-customization-jobs/get-job-status.mdx b/docs/customizer/manage-customization-jobs/get-job-status.mdx index 0ca20c7e42..b7f852f887 100644 --- a/docs/customizer/manage-customization-jobs/get-job-status.mdx +++ b/docs/customizer/manage-customization-jobs/get-job-status.mdx @@ -10,7 +10,7 @@ Get detailed execution status for a customization job, including step-by-step pr This endpoint provides granular execution details including: - **Step-level status**: `model-and-dataset-download` → `training` → `model-upload` → `model-entity-creation` -- **Training metrics**: `step`, `epoch`, `loss`, `lr` (learning rate), `grad_norm`, `val_loss` +- **Training metrics**: `step`, `epoch`, `train_loss`, `lr` (learning rate), `grad_norm`, `val_loss` - **Progress tracking**: `downloaded_files`, `uploaded_bytes`, `progress_pct` To list jobs or get job definitions (model entity, hyperparameters, spec), use [List Active Jobs](/documentation/customizer-reference/manage-customization-jobs/list-active-jobs) instead. @@ -146,7 +146,7 @@ curl -X GET \ "num_epochs": 2, "step": 8, "epoch": 1, - "loss": 2.8918895721435547, + "train_loss": 2.8918895721435547, "lr": 4.9101714686276044e-05, "grad_norm": 26.0 } @@ -222,7 +222,7 @@ curl -X GET \ "num_epochs": 2, "step": 94, "epoch": 2, - "loss": 0.3437718152999878, + "train_loss": 0.3437718152999878, "lr": 5.000000000000001e-07, "grad_norm": 20.125, "val_loss": 0.5527229905128479, diff --git a/docs/customizer/manage-customization-jobs/hyperparameters.mdx b/docs/customizer/manage-customization-jobs/hyperparameters.mdx index 760d6493bf..3265c29a1e 100644 --- a/docs/customizer/manage-customization-jobs/hyperparameters.mdx +++ b/docs/customizer/manage-customization-jobs/hyperparameters.mdx @@ -14,7 +14,7 @@ NeMo Customizer ships **two training backends**, and each accepts its own job co | Backend | Best for | Training methods | Hardware | |---------|----------|------------------|----------| | **Automodel** (default) | Production fine-tuning, larger models, multi-GPU scaling | SFT, distillation; LoRA, merged-LoRA, or full-weight | Single- or multi-GPU (tensor / pipeline / context / expert parallel) | -| **Unsloth** | Memory-constrained single-GPU LoRA | SFT; LoRA or full-weight | Single GPU (4-bit / 8-bit quantization) | +| **Unsloth** | Memory-constrained single-GPU training | SFT; LoRA or full-weight | Single GPU; optional 4-bit / 8-bit loading for LoRA, unquantized loading for full-weight | @@ -93,7 +93,7 @@ The `parallelism` block scales Automodel training across GPUs and nodes. | `parallelism.tensor_parallel_size` | GPUs for tensor parallelism (splits layers across GPUs for large models) | `1` | | `parallelism.pipeline_parallel_size` | GPUs for pipeline parallelism (splits model stages across GPUs) | `1` | | `parallelism.context_parallel_size` | GPUs for context parallelism (for very long sequences) | `1` | -| `parallelism.expert_parallel_size` | Expert parallelism for MoE models; must divide the number of experts | `null` | +| `parallelism.expert_parallel_size` | Expert parallelism for MoE models; must divide the number of experts. Leave unset for non-MoE models | `null` | @@ -102,7 +102,7 @@ The `parallelism` block scales Automodel training across GPUs and nodes. - `total_gpus = num_gpus_per_node × num_nodes`. - `total_gpus` must be divisible by `tensor_parallel_size × pipeline_parallel_size × context_parallel_size`. - `data_parallel_size` is derived as `total_gpus / (TP × PP × CP)`, and `global_batch_size` must be divisible by `micro_batch_size × data_parallel_size`. -- For MoE models, tensor parallelism must be `1` when `expert_parallel_size > 1`. +- For MoE models, when `expert_parallel_size` is set: the number of experts must be divisible by `expert_parallel_size`, `(data_parallel_size × context_parallel_size)` must be divisible by `expert_parallel_size`, and `tensor_parallel_size` must be `1` when `expert_parallel_size > 1`. @@ -131,7 +131,7 @@ When `training.training_type` is `"distillation"`, the following additional fiel ## Unsloth Configuration -An Unsloth job is configured with the following top-level sections: `model`, `dataset`, `training`, `schedule`, `batch`, `optimizer`, `hardware`, `output`, and (optionally) `integrations`. Unsloth runs on a **single GPU** and supports 4-bit / 8-bit quantized loading. +An Unsloth job is configured with the following top-level sections: `model`, `dataset`, `training`, `schedule`, `batch`, `optimizer`, `hardware`, `output`, and (optionally) `integrations`. Unsloth runs on a **single GPU** and supports 4-bit / 8-bit quantized loading for LoRA. Full-weight training must load the model without quantization. ### Model @@ -143,6 +143,8 @@ An Unsloth job is configured with the following top-level sections: `model`, `da | `model.load_in_8bit` | Load the base model in 8-bit | `false` | | `model.dtype` | Compute dtype (`auto`, `bfloat16`, `float16`, `float32`) | `auto` | | `model.trust_remote_code` | Allow custom model code from the checkpoint | `false` | +| `model.device_map` | Device placement forwarded to Unsloth. Accepts `auto`, `balanced`, `sequential`, a device index, or a custom map. `null` pins the model to the single visible GPU | `null` | +| `model.rope_scaling` | RoPE scaling configuration for long-context extension, such as `{"type": "linear", "factor": 2.0}` | `null` | @@ -166,7 +168,7 @@ Full-weight training (`training.finetuning_type: "all_weights"`) cannot be combi |-----------|--------|-------------|---------| | `training.training_type` | `sft` | Training method | `sft` | | `training.finetuning_type` | `lora`, `all_weights` | Adapter regime. `lora` trains an adapter; `all_weights` performs full-weight training | `lora` | -| `training.lora` | `{ rank, alpha, dropout, target_modules, bias, use_rslora, random_state }` | LoRA configuration (auto-filled with defaults when `finetuning_type` is `lora`) | *(see below)* | +| `training.lora` | `LoRAParams` object | LoRA configuration (auto-filled with defaults when `finetuning_type` is `lora`) | *(see below)* | | `training.use_gradient_checkpointing` | `unsloth`, `true`, `false` | Gradient checkpointing mode. `unsloth` uses Unsloth's optimized implementation | `unsloth` | LoRA parameters (`training.lora`): @@ -180,6 +182,12 @@ LoRA parameters (`training.lora`): | `bias` | Bias training mode (`none`, `all`, `lora_only`) | `none` | | `use_rslora` | Use rank-stabilized LoRA | `false` | | `random_state` | LoRA initialization seed | `3407` | +| `use_dora` | Use weight-decomposed LoRA (DoRA). Can improve quality at low ranks with additional training overhead | `false` | +| `loftq_config` | LoftQ initialization configuration for quantized base models | `null` | +| `modules_to_save` | Additional non-LoRA modules to train and save in full, such as `embed_tokens` or `lm_head` | `null` | +| `layers_to_transform` | Layer index or list of layer indexes to receive LoRA; `null` applies LoRA to all layers | `null` | +| `layer_replication` | Layer-replication ranges, such as `[[0, 16], [8, 24]]` | `null` | +| `init_lora_weights` | LoRA initialization: `true`, `false`, `gaussian`, `pissa`, `olora`, or `loftq` | `true` | ### Schedule @@ -194,6 +202,7 @@ LoRA parameters (`training.lora`): | `schedule.save_steps` | Checkpoint cadence (steps) | `null` | | `schedule.eval_steps` | Evaluation cadence (steps) | `null` | | `schedule.seed` | Random seed | `3407` | +| `schedule.lr_scheduler_kwargs` | Additional scheduler arguments, such as `{"num_cycles": 3}` for `cosine_with_restarts` | `null` | ### Batch @@ -209,6 +218,12 @@ LoRA parameters (`training.lora`): | `optimizer.learning_rate` | Step size for weight updates | `2e-4` | | `optimizer.weight_decay` | L2 regularization strength | `0.0` | | `optimizer.optim` | Optimizer (`adamw_torch`, `adamw_torch_fused`, `adamw_8bit`, `paged_adamw_8bit`, `sgd`). 8-bit optimizers reduce optimizer-state memory | `adamw_8bit` | +| `optimizer.adam_beta1` | Adam/AdamW first-moment decay | `0.9` | +| `optimizer.adam_beta2` | Adam/AdamW second-moment decay | `0.999` | +| `optimizer.adam_epsilon` | Adam/AdamW epsilon for numerical stability | `1e-8` | +| `optimizer.max_grad_norm` | Maximum gradient norm for clipping | `1.0` | +| `optimizer.label_smoothing_factor` | Cross-entropy label smoothing factor; `0.0` disables smoothing | `0.0` | +| `optimizer.neftune_noise_alpha` | NEFTune embedding-noise alpha; `null` disables NEFTune | `null` | ### Hardware @@ -217,9 +232,15 @@ LoRA parameters (`training.lora`): | `hardware.gpus` | Comma-separated GPU indices (`0` or `0,1`) for `CUDA_VISIBLE_DEVICES` (selection, not reservation) | `null` | | `hardware.precision` | Mixed-precision dtype (`bf16`, `fp16`). `bf16` recommended for Ampere+ | `bf16` | -### Output (save method) +### Output -Unsloth's output `save_method` controls the saved checkpoint shape: +| Parameter | Description | Default | +|-----------|-------------|---------| +| `output.name` | Output Model Entity or adapter name | Auto-generated from the job name | +| `output.description` | Optional description for the generated artifact | `null` | +| `output.save_method` | Saved checkpoint shape (see below) | `lora` | + +The `output.save_method` field accepts: | `save_method` | Result | |---------------|--------| diff --git a/docs/customizer/manage-customization-jobs/index.mdx b/docs/customizer/manage-customization-jobs/index.mdx index f01cfe9349..88d8fe9619 100644 --- a/docs/customizer/manage-customization-jobs/index.mdx +++ b/docs/customizer/manage-customization-jobs/index.mdx @@ -7,7 +7,7 @@ Use customization jobs to fine-tune a [model](/documentation/customizer-referenc ## How It Works -A customization job references a **Model Entity** that contains the base model checkpoint, and is submitted to one of two backends — **automodel** (default, multi-GPU) or **unsloth** (single-GPU, quantized). The job then runs on the platform's GPU cluster. When training completes: +A customization job references a **Model Entity** that contains the base model checkpoint, and is submitted to one of two backends — **automodel** (default, multi-GPU) or **unsloth** (single-GPU, with optional quantized loading for LoRA). The job then runs on the platform's GPU cluster. When training completes: - **LoRA jobs**: Create an **Adapter** attached to the original Model Entity. Adapters can be auto-deployed to NIMs. - **Full fine-tuning jobs**: Create a **new Model Entity** with the customized weights, linked to the base model. diff --git a/docs/customizer/manage-customization-jobs/list-active-jobs.mdx b/docs/customizer/manage-customization-jobs/list-active-jobs.mdx index ef8ddb3603..e993f1cf80 100644 --- a/docs/customizer/manage-customization-jobs/list-active-jobs.mdx +++ b/docs/customizer/manage-customization-jobs/list-active-jobs.mdx @@ -3,7 +3,7 @@ title: "List Active Jobs" description: "" --- -List customization jobs and their high-level status. Customization jobs run on the platform's Jobs service, so you list them through that service and filter by `source` to scope the results to a backend (`automodel` or `unsloth`). Each entry includes the job definition (model, dataset, training configuration) and overall status. +List active customization jobs and their high-level status. Customization jobs run on the platform's Jobs service, so you list them through that service and filter by `source` and `status`. The `source` scopes results to a backend (`automodel` or `unsloth`), while `status="active"` excludes completed, failed, and cancelled jobs. Each entry includes the job definition (model, dataset, training configuration) and overall status. @@ -25,7 +25,7 @@ export NMP_BASE_URL="https://your-nmp-base-url" ## To List Active Customization Jobs -Use the SDK to list jobs, filtering by `source` to scope the results to a customization backend: +Use the SDK to list jobs, filtering by `source` to scope the results to a customization backend and by `status` to return only active jobs: ```python import os @@ -37,10 +37,13 @@ client = NeMoPlatform( workspace="default", ) -# List automodel customization jobs +# List active automodel customization jobs jobs = client.jobs.list( workspace="default", - filter={"source": "automodel"}, # Use "unsloth" for the Unsloth backend + filter={ + "source": "automodel", # Use "unsloth" for the Unsloth backend + "status": "active", + }, page=1, page_size=10, sort="created_at", @@ -56,7 +59,8 @@ filtered_jobs = client.jobs.list( workspace="default", filter={ "source": "automodel", - "status": "active", # Filter by job status + "status": "active", + "project": "my-finetuning-project", }, sort="-created_at", # Sort by created_at descending ) @@ -118,7 +122,10 @@ for job in filtered_jobs.data: "total_results": 1 }, "sort": "created_at", - "filter": {}, + "filter": { + "source": "automodel", + "status": "active" + }, "search": {} } ``` diff --git a/docs/customizer/manage-model-entities/create-fileset.mdx b/docs/customizer/manage-model-entities/create-fileset.mdx index 410cea2757..2d629b9ad0 100644 --- a/docs/customizer/manage-model-entities/create-fileset.mdx +++ b/docs/customizer/manage-model-entities/create-fileset.mdx @@ -8,7 +8,7 @@ Create a FileSet containing your base model checkpoint before creating a Model E ## Prerequisites - Obtained the base URL of your NeMo Platform. -- For gated or private HuggingFace models: Created a secret with your HF token. Refer to [Manage Secrets](/documentation/get-started/core-concepts/manage-secrets). +- For gated or private Hugging Face models: Created a secret with your HF token. Refer to [Manage Secrets](/documentation/get-started/core-concepts/manage-secrets). - Set the `NMP_BASE_URL` environment variable. ```bash @@ -17,9 +17,9 @@ export NMP_BASE_URL="https://your-nemo-platform-url" --- -## From HuggingFace Hub +## From Hugging Face Hub -The most common method is downloading directly from HuggingFace. The example below uses [Qwen/Qwen3-1.7B](https://huggingface.co/Qwen/Qwen3-1.7B), a public model that requires no token: +The most common method is downloading directly from Hugging Face. The example below uses [Qwen/Qwen3-1.7B](https://huggingface.co/Qwen/Qwen3-1.7B), a public model that requires no token: ```python import os @@ -34,12 +34,12 @@ client = NeMoPlatform( HF_REPO_ID = "Qwen/Qwen3-1.7B" MODEL_NAME = "qwen3-1.7b" -# Create FileSet from HuggingFace +# Create FileSet from Hugging Face try: fileset = client.files.filesets.create( workspace="default", name=MODEL_NAME, - description="Qwen3 1.7B from HuggingFace", + description="Qwen3 1.7B from Hugging Face", purpose="model", storage=HuggingfaceStorageConfigParam( type="huggingface", @@ -57,10 +57,10 @@ print(f"FileSet ready: {fileset.name}") -Gated models (such as Llama) require a HuggingFace token. To use one: +Gated models (such as Llama) require a Hugging Face token. To use one: -1. Accept the model license on the HuggingFace model page. -2. Create a HuggingFace token with read access. +1. Accept the model license on the Hugging Face model page. +2. Create a Hugging Face token with read access. 3. Store the token as a secret in the platform (see [Manage Secrets](/documentation/get-started/core-concepts/manage-secrets)), then pass it as `token_secret` in the storage config: ```python @@ -90,7 +90,7 @@ client = NeMoPlatform( ) MODEL_NAME = "nemotron-mini-4b" -NGC_RESOURCE = "nemotron-mini-4b-instruct" +NGC_TARGET = "nemotron-mini-4b-instruct" NGC_ORG = "nvidia" NGC_TEAM = "nemo" NGC_VERSION = "1.0" @@ -122,7 +122,8 @@ try: type="ngc", org=NGC_ORG, team=NGC_TEAM, - resource=NGC_RESOURCE, # NGC resource name + target=NGC_TARGET, # NGC asset name + target_type="resource", version=NGC_VERSION, api_key_secret=ngc_secret.name, ), @@ -141,10 +142,11 @@ Files are downloaded in the background after you create a FileSet. Check the sta ```python # List files in the FileSet -files = client.files.list( +response = client.files.list( workspace="default", fileset="qwen3-1.7b", ) +files = response.data print(f"Files in FileSet ({len(files)} total):") for f in files[:10]: # Show first 10 diff --git a/docs/customizer/manage-model-entities/create-model-entity.mdx b/docs/customizer/manage-model-entities/create-model-entity.mdx index d06767e59a..a08bad2a64 100644 --- a/docs/customizer/manage-model-entities/create-model-entity.mdx +++ b/docs/customizer/manage-model-entities/create-model-entity.mdx @@ -96,9 +96,9 @@ print(f" Attention Heads: {model.spec.num_attention_heads}") "hidden_size": 2048, "num_layers": 28, "num_attention_heads": 16, - "num_key_value_heads": 8, + "num_kv_heads": 8, "vocab_size": 151936, - "max_sequence_length": 40960 + "context_size": 40960 } } ``` diff --git a/docs/customizer/manage-model-entities/index.mdx b/docs/customizer/manage-model-entities/index.mdx index 3fbe6ec2cd..0a747b3585 100644 --- a/docs/customizer/manage-model-entities/index.mdx +++ b/docs/customizer/manage-model-entities/index.mdx @@ -11,7 +11,7 @@ Before running a customization job, you need to set up a **Model Entity** that p -Create a FileSet containing your base model checkpoint from HuggingFace, NGC, or local storage. +Create a FileSet containing your base model checkpoint from Hugging Face, NGC, or local storage. @@ -35,7 +35,7 @@ A **FileSet** is a collection of files managed by the platform. For customizatio - Tokenizer files (`tokenizer.json`, `tokenizer_config.json`, and so on) FileSets can be populated from: -- **HuggingFace Hub** - Download directly from HF repositories +- **Hugging Face Hub** - Download directly from HF repositories - **NGC** - Download from NVIDIA NGC catalogs - **Local upload** - Upload files from your local machine @@ -56,7 +56,7 @@ Complete example of setting up a model for customization: -**HuggingFace Token**: If downloading from a gated HuggingFace repository (like Llama models), you will need to create a secret containing your HuggingFace API token first. Refer to [Manage Secrets](/documentation/get-started/core-concepts/manage-secrets) for instructions. +**Hugging Face Token**: If downloading from a gated Hugging Face repository (like Llama models), you will need to create a secret containing your Hugging Face API token first. Refer to [Manage Secrets](/documentation/get-started/core-concepts/manage-secrets) for instructions. ```python @@ -71,12 +71,12 @@ client = NeMoPlatform( workspace="default", ) -# Step 1: Create FileSet from HuggingFace +# Step 1: Create FileSet from Hugging Face try: fileset = client.files.filesets.create( workspace="default", name="qwen3-1.7b", - description="Qwen3 1.7B base model from HuggingFace", + description="Qwen3 1.7B base model from Hugging Face", storage=HuggingfaceStorageConfigParam( type="huggingface", repo_id="Qwen/Qwen3-1.7B", diff --git a/docs/customizer/models/data-format.mdx b/docs/customizer/models/data-format.mdx index dd0b3aba2b..5140cac824 100644 --- a/docs/customizer/models/data-format.mdx +++ b/docs/customizer/models/data-format.mdx @@ -9,7 +9,7 @@ Use the following guidelines to prepare your training dataset for the supported ## Dataset Preparation Guidelines - **File Format**: Save your training data as `.jsonl` files (one JSON object per line). -- **Validation**: Each record is automatically validated against the appropriate schema when training begins. The required format depends on `training.type` (for example, `sft`) specified in your job configuration. +- **Validation**: Each record is automatically validated against the appropriate schema when training begins. The required format depends on `training.training_type` (for example, `sft`) specified in your job configuration. For dataset creation tutorials, refer to [Format Training Dataset](/documentation/customizer-reference/tutorials/format-training-dataset). @@ -51,20 +51,27 @@ Each line in your JSONL file must contain a JSON object with these required fiel Each line in your JSONL file must contain a JSON object with these required fields: - **`messages`** (array of objects): The messages in the conversation. - - **`role`** (string): The role of the message. - - **`content`** (string): The content of the message. + - **`role`** (string): The role of the message. + - **`content`** (string): The content of the message. #### Example Dataset Entry -``` +```json { - "messages": [ - { - "role": "system", - "content": "You are an email writing assistant. Please help people write cogent emails." - }, - ... - ] + "messages": [ + { + "role": "system", + "content": "You are an email writing assistant. Please help people write cogent emails." + }, + { + "role": "user", + "content": "Write a concise follow-up after a project review." + }, + { + "role": "assistant", + "content": "Thank you for reviewing the project. Please let me know if you have any additional feedback." + } + ] } ``` @@ -158,10 +165,10 @@ Each line in your JSONL file must contain a JSON object with these required fiel #### Example Dataset Entry -``` +```json { - "prompt": "your string", - "completion": "your expected response" + "prompt": "What is the capital of France?", + "completion": "Paris." } ``` @@ -173,17 +180,23 @@ Each line in your JSONL file must contain a JSON object with these required fiel - **`system`** (string): The system message that defines the assistant's role or behavior. - **`conversations`** (array of objects): The conversation turns between user and assistant. - - **`from`** (string): The role of the message sender ("User" or "Assistant"). - - **`value`** (string): The content of the message. + - **`from`** (string): The role of the message sender (`User` or `Assistant`). + - **`value`** (string): The content of the message. #### Example Dataset Entry -``` +```json { - "system": "you are a robot", - "conversations": [ - {"from": "User", "value": "Choose a number that is greater than 0 and less than 2\n"}, - {"from": "Assistant", "value": "1"} - ] + "system": "You are a helpful assistant.", + "conversations": [ + { + "from": "User", + "value": "Choose a number that is greater than 0 and less than 2." + }, + { + "from": "Assistant", + "value": "1" + } + ] } ``` diff --git a/docs/customizer/models/embedding.mdx b/docs/customizer/models/embedding.mdx index 14ea7c522d..4abe495c94 100644 --- a/docs/customizer/models/embedding.mdx +++ b/docs/customizer/models/embedding.mdx @@ -20,7 +20,7 @@ This page provides detailed technical specifications for the embedding model fam | Training Data | Semi-supervised pre-training on 12M samples and fine-tuning on 1M samples from public QA datasets with commercial licenses | | License | [NVIDIA Open Model License](https://www.nvidia.com/en-us/agreements/enterprise-software/nvidia-open-model-license/), [Llama 3.2 Community License](https://www.llama.com/llama3_2/license/) | | Default Name | nvidia/llama-nemotron-embed-1b-v2 | -| HuggingFace | [nvidia/llama-nemotron-embed-1b-v2](https://huggingface.co/nvidia/llama-nemotron-embed-1b-v2) | +| Hugging Face | [nvidia/llama-nemotron-embed-1b-v2](https://huggingface.co/nvidia/llama-nemotron-embed-1b-v2) | | NIM | [nvidia/llama-nemotron-embed-1b-v2](https://catalog.ngc.nvidia.com/orgs/nim/teams/nvidia/containers/llama-nemotron-embed-1b-v2) | ### Model Entity Configuration @@ -59,9 +59,9 @@ The following table provides conservative hyperparameter defaults specifically o | Parameter | API Field Name | Type | Description | Recommended Value | | --- | --- | --- | --- | --- | -| Learning Rate | `learning_rate` | `number` | Step size for updating model parameters. Lower values help prevent overfitting in embedding models. | `5e-6` | -| Weight Decay | `weight_decay` | `number` | Regularization parameter to prevent overfitting by penalizing large weights. | `0.01` | -| Number of Epochs | `epochs` | `integer` | Number of complete passes through the training dataset. Limited to prevent overfitting. | `1` | +| Learning Rate | `optimizer.learning_rate` | `number` | Step size for updating model parameters. Lower values help prevent overfitting in embedding models. | `5e-6` | +| Weight Decay | `optimizer.weight_decay` | `number` | Regularization parameter to prevent overfitting by penalizing large weights. | `0.01` | +| Number of Epochs | `schedule.epochs` | `integer` | Number of complete passes through the training dataset. Limited to prevent overfitting. | `1` | | Training Data Size | N/A | N/A | Number of training examples to prevent overfitting while maintaining model performance. | `5,000-10,000` examples | NVIDIA recommends evaluating fine-tuned embedding models against the baseline to detect overfitting and potential performance degradation. @@ -79,8 +79,9 @@ This model supports inference deployment through NVIDIA Inference Microservices 1. **Deploy the model**: Create a ModelDeploymentConfig and ModelDeployment to deploy your fine-tuned model. See [about](/documentation/models-and-inference) for details. 2. **Access through Inference Gateway**: The Inference Gateway provides unified access to all deployed models via three routing patterns: - - **Model Entity routing**: `/v2/workspaces/{workspace}/inference/gateway/model/{name}/-/v1/embeddings` - - **Provider routing**: `/v2/workspaces/{workspace}/inference/gateway/provider/{deployment}/-/v1/embeddings` + - **Model Entity routing**: `/apis/inference-gateway/v2/workspaces/{workspace}/model/{name}/-/v1/embeddings` + - **Provider routing**: `/apis/inference-gateway/v2/workspaces/{workspace}/provider/{deployment}/-/v1/embeddings` + - **OpenAI routing**: `/apis/inference-gateway/v2/workspaces/{workspace}/openai/-/v1/embeddings` (specify the model in the request body) ```python import os @@ -132,6 +133,6 @@ embedding = response["data"][0]["embedding"] print(f"Embedding dimension: {len(embedding)}") ``` -For detailed fine-tuning instructions, refer to the [Embedding Customization tutorial](../tutorials/embedding-customization-job.ipynb). +For detailed fine-tuning instructions, refer to the [Embedding Customization tutorial](/documentation/customizer-reference/tutorials/embedding-customization-job). For more information about formatting training datasets for the embedding model, refer to [Dataset Format Requirements](/documentation/customizer-reference/models/dataset-format). diff --git a/docs/customizer/models/gpt-oss.mdx b/docs/customizer/models/gpt-oss.mdx index 2d1b4e2eb6..20cffd397e 100644 --- a/docs/customizer/models/gpt-oss.mdx +++ b/docs/customizer/models/gpt-oss.mdx @@ -8,7 +8,7 @@ This page provides detailed technical specifications for the OpenAI GPT-OSS mode ## Before You Start -These models require a HuggingFace token to download. Create a secret with your HuggingFace API key, then create a FileSet and Model Entity referencing the model. See [index](/documentation/customizer-reference/manage-model-entities/overview) for setup instructions. +These models require a Hugging Face token to download. Create a secret with your Hugging Face API key, then create a FileSet and Model Entity referencing the model. See [index](/documentation/customizer-reference/manage-model-entities/overview) for setup instructions. --- @@ -19,18 +19,24 @@ These models require a HuggingFace token to download. Create a secret with your | Creator | OpenAI | | Architecture | Mixture of Experts (MoE) Transformer | | Description | GPT-OSS 20B provides lower latency for local or specialized use cases, featuring full chain-of-thought reasoning and agentic capabilities. | -| Max I/O Tokens | Not specified | +| Context Length | 131,072 tokens | | Parameters | 21B parameters (3.6B active parameters) | | Training Data | Trained on harmony response format | -| Memory Requirements | Runs within 32GB of memory with BFloat16 quantization | +| Checkpoint Quantization | MXFP4 quantization of the MoE weights | +| Inference Memory | The official checkpoint can run within 16GB of memory; Customizer training requires the GPU configurations below | | Default Name | openai/gpt-oss-20b | -| HuggingFace | [openai/gpt-oss-20b](https://huggingface.co/openai/gpt-oss-20b) | +| Hugging Face | [openai/gpt-oss-20b](https://huggingface.co/openai/gpt-oss-20b) | ### Training Options (20B) -- **LoRA**: 4x 80GB GPU, tensor parallel size 1, expert parallel size 4, pipeline parallel size 1 -- **Full SFT**: 8x 80GB GPU, tensor parallel size 1, expert parallel size 8, pipeline parallel size 1 -- Sequence Packing: Not supported +The Automodel contract matrix validates these configurations: + +| Fine-Tuning | Dataset Format | GPUs | Sequence Packing | +| --- | --- | --- | --- | +| LoRA | Prompt-completion | 1 | Supported | +| Full SFT | Prompt-completion or chat | 8 | Not supported in the tested configurations | + +For the tested Full SFT configuration, tensor and pipeline parallel sizes are `1` and expert parallel size is `8`. The tested LoRA configuration uses one GPU without expert parallelism. Default training max sequence length: 4096. @@ -81,6 +87,6 @@ GPT-OSS models use the harmony response format and require this format for prope -Sequence packing is not supported for GPT-OSS models in NeMo Customizer. +For GPT-OSS 20B, sequence packing is supported by the tested Automodel LoRA configuration with prompt-completion data. It is not supported by the tested Full SFT or chat configurations. diff --git a/docs/customizer/models/index.mdx b/docs/customizer/models/index.mdx index 8ea46b6871..25e86ca846 100644 --- a/docs/customizer/models/index.mdx +++ b/docs/customizer/models/index.mdx @@ -14,7 +14,7 @@ For fine-tuning and deployment tutorials, see the [Tutorials](/documentation/cus ## Before You Start -If downloading models hosted on Hugging Face, create a secret with your HuggingFace API key, then create a FileSet and Model Entity referencing the model. See [index](/documentation/customizer-reference/manage-model-entities/overview) for setup instructions. +If downloading models hosted on Hugging Face, create a secret with your Hugging Face API key, then create a FileSet and Model Entity referencing the model. See [index](/documentation/customizer-reference/manage-model-entities/overview) for setup instructions. --- @@ -63,7 +63,7 @@ View the available Mistral models, including Mistral and Ministral variants for ## Tested Models -The following table lists models that NVIDIA tested and their available features. This is a list of *known-good* combinations, not a list of limits: NeMo Customizer can fine-tune many models and regimes beyond those listed, including additional Hugging Face checkpoints, other fine-tuning regimes (LoRA, merged-LoRA, full-weight, distillation), and either training backend (Automodel or Unsloth). Models and regimes outside this table may work but have not been formally validated. +The following table lists models that NVIDIA tested and their available features. This is a list of *known-good* combinations, not a list of limits: NeMo Customizer can fine-tune additional Hugging Face checkpoints and regimes when their architectures are supported by the selected training backend. Compatibility varies by architecture and fine-tuning method—for example, Automodel LoRA does not support Conv1D-based models. Models and regimes outside this table may work but have not been formally validated; test them before relying on them in production. For detailed technical specifications of each model such as architecture, parameters, and token limits, refer to the [model family](#model-families) pages. @@ -81,7 +81,7 @@ The following models support both chat and completion model training. | [nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16](https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16) | No | Full SFT, LoRA | No | Supported (only Full SFT) | Yes | | [nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16](https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16) | No | LoRA | No | Supported | Yes | | [microsoft/phi-4](https://huggingface.co/microsoft/phi-4) | No | Full SFT, LoRA | No | Supported | No | -| [openai/gpt-oss-20b](https://huggingface.co/openai/gpt-oss-20b) | Yes | Full SFT, LoRA | No | Supported | Yes | +| [openai/gpt-oss-20b](https://huggingface.co/openai/gpt-oss-20b) | Yes | Full SFT, LoRA | LoRA with prompt-completion data | Supported | Yes | | [Qwen/Qwen2.5-1.5B-Instruct](https://huggingface.co/Qwen/Qwen2.5-1.5B-Instruct) | No | Full SFT, LoRA | No| Supported | Yes | | [Qwen/Qwen3-0.6B](https://huggingface.co/Qwen/Qwen3-0.6B) | No | Full SFT, LoRA | No | Supported | Yes | | [mistralai/Mistral-7B-Instruct-v0.3](https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.3) | No | Full SFT, LoRA | No | Supported | No | diff --git a/docs/customizer/models/llama-nemotron.mdx b/docs/customizer/models/llama-nemotron.mdx index 091b420d7c..ccf85a78d3 100644 --- a/docs/customizer/models/llama-nemotron.mdx +++ b/docs/customizer/models/llama-nemotron.mdx @@ -17,7 +17,7 @@ This page provides detailed technical specifications for the Nemotron model fami | Parameters | 8 billion | | Training Data | Not specified | | Default Name | nvidia/Llama-3.1-Nemotron-Nano-8B-v1 | -| HuggingFace | [nvidia/Llama-3.1-Nemotron-Nano-8B-v1](https://huggingface.co/nvidia/Llama-3.1-Nemotron-Nano-8B-v1) | +| Hugging Face | [nvidia/Llama-3.1-Nemotron-Nano-8B-v1](https://huggingface.co/nvidia/Llama-3.1-Nemotron-Nano-8B-v1) | | NIM | [nvidia/llama-3.1-nemotron-nano-8b-v1](https://catalog.ngc.nvidia.com/orgs/nim/teams/nvidia/containers/llama-3.1-nemotron-nano-8b-v1?version=1.8.4) | ### Training Options @@ -47,7 +47,7 @@ This page provides detailed technical specifications for the Nemotron model fami | Max I/O Tokens | 4096 | | Parameters | 9 billion | | Default Name | nvidia/NVIDIA-Nemotron-Nano-9B-v2 | -| HuggingFace | [nvidia/NVIDIA-Nemotron-Nano-9B-v2](https://huggingface.co/nvidia/NVIDIA-Nemotron-Nano-9B-v2) | +| Hugging Face | [nvidia/NVIDIA-Nemotron-Nano-9B-v2](https://huggingface.co/nvidia/NVIDIA-Nemotron-Nano-9B-v2) | | NIM | [NVIDIA-Nemotron-Nano-9B-v2](https://catalog.ngc.nvidia.com/orgs/nim/teams/nvidia/containers/nvidia-nemotron-nano-9b-v2?version=latest) | ### Training Options @@ -79,7 +79,7 @@ This page provides detailed technical specifications for the Nemotron model fami | MoE Configuration | 128 experts + 1 shared expert, 6 experts activated per token | | Supported Languages | English, German, Spanish, French, Italian, Japanese | | Default Name | nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 | -| HuggingFace | [nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16](https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16) | +| Hugging Face | [nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16](https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16) | | NIM | [Nemotron-3-Nano-30B-A3B](https://catalog.ngc.nvidia.com/orgs/nim/teams/nvidia/containers/nemotron-3-nano?version=2.0.1) | ### Training Options @@ -115,7 +115,7 @@ Deployment for LoRA using NIM is not supported for this model. | Max I/O Tokens | 4096 | | Parameters | 120B total (12B active) | | Default Name | nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16 | -| HuggingFace | [nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16](https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16) | +| Hugging Face | [nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16](https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16) | ### Training Options diff --git a/docs/customizer/models/llama.mdx b/docs/customizer/models/llama.mdx index 31573aa778..a4fee46942 100644 --- a/docs/customizer/models/llama.mdx +++ b/docs/customizer/models/llama.mdx @@ -17,7 +17,7 @@ This page provides detailed technical specifications for the Llama model family | Parameters | 3 billion | | Training Data | 15+ trillion tokens (up to 2024) | | Default Name | meta-llama/Llama-3.2-3B-Instruct | -| HuggingFace | [meta-llama/Llama-3.2-3B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-3B-Instruct) | +| Hugging Face | [meta-llama/Llama-3.2-3B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-3B-Instruct) | ### Training Options @@ -46,7 +46,7 @@ This page provides detailed technical specifications for the Llama model family | Parameters | 1 billion | | Training Data | 15+ trillion tokens (up to 2024) | | Default Name | meta-llama/Llama-3.2-1B-Instruct | -| HuggingFace | [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct) | +| Hugging Face | [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct) | ### Training Options @@ -76,7 +76,7 @@ This page provides detailed technical specifications for the Llama model family | Parameters | 8 billion | | Training Data | 15 trillion tokens (up to December 2023) | | Default Name | meta-llama/Llama-3.1-8B-Instruct | -| HuggingFace | [meta-llama/Llama-3.1-8B-Instruct](https://huggingface.co/meta-llama/Llama-3.1-8B-Instruct) | +| Hugging Face | [meta-llama/Llama-3.1-8B-Instruct](https://huggingface.co/meta-llama/Llama-3.1-8B-Instruct) | ### Training Options diff --git a/docs/customizer/models/mistral.mdx b/docs/customizer/models/mistral.mdx index dc3850f9c7..9d6cbcb716 100644 --- a/docs/customizer/models/mistral.mdx +++ b/docs/customizer/models/mistral.mdx @@ -17,7 +17,7 @@ This page provides detailed technical specifications for the Mistral model famil | Parameters | 7 billion | | Training Data | Not specified | | Default Name | mistralai/Mistral-7B-Instruct-v0.3 | -| HuggingFace | [mistralai/Mistral-7B-Instruct-v0.3](https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.3) | +| Hugging Face | [mistralai/Mistral-7B-Instruct-v0.3](https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.3) | ### Training Options @@ -46,7 +46,7 @@ This page provides detailed technical specifications for the Mistral model famil | Parameters | 3 billion | | Training Data | Not specified | | Default Name | mistralai/Ministral-3-3B-Instruct-2512 | -| HuggingFace | [mistralai/Ministral-3-3B-Instruct-2512](https://huggingface.co/mistralai/Ministral-3-3B-Instruct-2512) | +| Hugging Face | [mistralai/Ministral-3-3B-Instruct-2512](https://huggingface.co/mistralai/Ministral-3-3B-Instruct-2512) | ### Training Options @@ -69,7 +69,7 @@ Deployment using NIM is not supported for this model. | Parameters | 3 billion | | Training Data | Not specified | | Default Name | mistralai/Ministral-3-3B-Reasoning-2512 | -| HuggingFace | [mistralai/Ministral-3-3B-Reasoning-2512](https://huggingface.co/mistralai/Ministral-3-3B-Reasoning-2512) | +| Hugging Face | [mistralai/Ministral-3-3B-Reasoning-2512](https://huggingface.co/mistralai/Ministral-3-3B-Reasoning-2512) | ### Training Options diff --git a/docs/customizer/models/phi.mdx b/docs/customizer/models/phi.mdx index caeaf972a1..ad498c654e 100644 --- a/docs/customizer/models/phi.mdx +++ b/docs/customizer/models/phi.mdx @@ -16,7 +16,7 @@ This page provides detailed technical specifications for the Phi model family su | Parameters | 14 billion | | Training Data | High-quality data with emphasis on reasoning and code | | Default Name | microsoft/phi-4 | -| HuggingFace | [microsoft/phi-4](https://huggingface.co/microsoft/phi-4) | +| Hugging Face | [microsoft/phi-4](https://huggingface.co/microsoft/phi-4) | ### Training Options diff --git a/docs/customizer/models/qwen.mdx b/docs/customizer/models/qwen.mdx index 3d414a3d3d..473002a750 100644 --- a/docs/customizer/models/qwen.mdx +++ b/docs/customizer/models/qwen.mdx @@ -17,7 +17,7 @@ This page provides detailed technical specifications for the Qwen model family s | Parameters | 1.5 billion | | Training Data | Not specified | | Default Name | Qwen/Qwen2.5-1.5B-Instruct | -| HuggingFace | [Qwen/Qwen2.5-1.5B-Instruct](https://huggingface.co/Qwen/Qwen2.5-1.5B-Instruct) | +| Hugging Face | [Qwen/Qwen2.5-1.5B-Instruct](https://huggingface.co/Qwen/Qwen2.5-1.5B-Instruct) | ### Training Options @@ -46,7 +46,7 @@ This page provides detailed technical specifications for the Qwen model family s | Parameters | 0.6 billion | | Training Data | Not specified | | Default Name | Qwen/Qwen3-0.6B | -| HuggingFace | [Qwen/Qwen3-0.6B](https://huggingface.co/Qwen/Qwen3-0.6B) | +| Hugging Face | [Qwen/Qwen3-0.6B](https://huggingface.co/Qwen/Qwen3-0.6B) | ### Training Options diff --git a/docs/customizer/tutorials/_snippets/output/chat-basic-format-example.jsonl b/docs/customizer/tutorials/_snippets/output/chat-basic-format-example.jsonl deleted file mode 100644 index 3d63e84cd3..0000000000 --- a/docs/customizer/tutorials/_snippets/output/chat-basic-format-example.jsonl +++ /dev/null @@ -1 +0,0 @@ -{"messages": [{"role": "system","content": ""}, {"role": "user","content": ""}, {"role": "assistant","content": ""}]} diff --git a/docs/customizer/tutorials/_snippets/output/chat-expanded-format-example.json b/docs/customizer/tutorials/_snippets/output/chat-expanded-format-example.json deleted file mode 100644 index c98ab1de5b..0000000000 --- a/docs/customizer/tutorials/_snippets/output/chat-expanded-format-example.json +++ /dev/null @@ -1,14 +0,0 @@ -{ - "messages": [ - { - "role": "system", - "content": "" - }, { - "role": "user", - "content": "" - }, { - "role": "assistant", - "content": "" - } - ] -} diff --git a/docs/customizer/tutorials/_snippets/output/chat-thinking-off-example.jsonl b/docs/customizer/tutorials/_snippets/output/chat-thinking-off-example.jsonl deleted file mode 100644 index 7cf16c2f8f..0000000000 --- a/docs/customizer/tutorials/_snippets/output/chat-thinking-off-example.jsonl +++ /dev/null @@ -1,5 +0,0 @@ -{"messages": [ - {"role": "system", "content": "detailed thinking off"}, - {"role": "user", "content": "What is 2 + 2?"}, - {"role": "assistant", "content": "4"} -]} diff --git a/docs/customizer/tutorials/_snippets/output/chat-thinking-on-example.jsonl b/docs/customizer/tutorials/_snippets/output/chat-thinking-on-example.jsonl deleted file mode 100644 index 637a0e338e..0000000000 --- a/docs/customizer/tutorials/_snippets/output/chat-thinking-on-example.jsonl +++ /dev/null @@ -1,5 +0,0 @@ -{"messages": [ - {"role": "system", "content": "detailed thinking on"}, - {"role": "user", "content": "What is 2 + 2?"}, - {"role": "assistant", "content": "To solve 2 + 2, add 2 and 2 together. The answer is 4."} -]} diff --git a/docs/customizer/tutorials/_snippets/output/chat-tool-calling-basic-example.jsonl b/docs/customizer/tutorials/_snippets/output/chat-tool-calling-basic-example.jsonl deleted file mode 100644 index 7a1e354929..0000000000 --- a/docs/customizer/tutorials/_snippets/output/chat-tool-calling-basic-example.jsonl +++ /dev/null @@ -1 +0,0 @@ -{"messages": [{"role": "user","content": ""},{"role": "assistant","content": "","tool_calls": [{"type": "function","function": {"name": "fibonacci","arguments": {"n": 20}}}]}],"tools": [{"type": "function","function": {"name": "fibonacci","description": "Calculates the nth Fibonacci number.","parameters": {"type": "object","properties": {"n": {"description": "The position of the Fibonacci number.","type": "integer"}}}}}]} diff --git a/docs/customizer/tutorials/_snippets/output/chat-tool-calling-expanded-example.json b/docs/customizer/tutorials/_snippets/output/chat-tool-calling-expanded-example.json deleted file mode 100644 index 2bf7f0d13b..0000000000 --- a/docs/customizer/tutorials/_snippets/output/chat-tool-calling-expanded-example.json +++ /dev/null @@ -1,35 +0,0 @@ -{ - "messages": [ - { - "role": "user", - "content": "" - }, - { - "role": "assistant", - "content": "", - "tool_calls": [{ - "type": "function", - "function": { - "name": "fibonacci", - "arguments": {"n": 20} - } - }] - } - ], - "tools": [{ - "type": "function", - "function": { - "name": "fibonacci", - "description": "Calculates the nth Fibonacci number.", - "parameters": { - "type": "object", - "properties": { - "n": { - "description": "The position of the Fibonacci number.", - "type": "integer" - } - } - } - } - }] -} diff --git a/docs/customizer/tutorials/_snippets/output/completion-format-example.jsonl b/docs/customizer/tutorials/_snippets/output/completion-format-example.jsonl deleted file mode 100644 index 9fc99eeb16..0000000000 --- a/docs/customizer/tutorials/_snippets/output/completion-format-example.jsonl +++ /dev/null @@ -1 +0,0 @@ -{"prompt": "Hello", "completion": " world."} diff --git a/docs/customizer/tutorials/_snippets/output/config-list-example.json b/docs/customizer/tutorials/_snippets/output/config-list-example.json deleted file mode 100644 index 25bc371e54..0000000000 --- a/docs/customizer/tutorials/_snippets/output/config-list-example.json +++ /dev/null @@ -1,95 +0,0 @@ -{ - "object": "list", - "data": [ - { - "name": "meta/llama-3.2-1b-instruct@v1.0.0+80GB", - "namespace": "default", - "dataset_schemas": [ - { - "title": "Newline-Delimited JSON File", - "type": "array", - "items": { - "description": "Schema for Supervised Fine-Tuning (SFT) training data items.", - "properties": { - "prompt": { - "description": "The prompt for the entry", - "title": "Prompt", - "type": "string" - }, - "completion": { - "description": "The completion to train on", - "title": "Completion", - "type": "string" - } - }, - "required": ["prompt", "completion"], - "title": "SFTDatasetItemSchema", - "type": "object" - } - } - ], - "training_options": [ - { - "training_type": "sft", - "finetuning_type": "lora", - "num_gpus": 1, - "num_nodes": 1, - "tensor_parallel_size": 1, - "use_sequence_parallel": false - }, - { - "training_type": "sft", - "finetuning_type": "all_weights", - "num_gpus": 1, - "num_nodes": 1, - "tensor_parallel_size": 1, - "use_sequence_parallel": false - } - ] - }, - { - "name": "nvidia/llama-3.2-nv-embedqa-1b@v2+80GB", - "namespace": "nvidia", - "dataset_schemas": [ - { - "title": "Newline-Delimited JSON File", - "type": "array", - "items": { - "description": "Schema for embedding training data items.", - "properties": { - "query": { - "description": "The query to use as an anchor", - "title": "Query", - "type": "string" - }, - "pos_doc": { - "description": "A document that should match positively with the anchor", - "title": "Positive Document", - "type": "string" - }, - "neg_doc": { - "description": "Documents that should not match with the anchor", - "title": "Negative Documents", - "type": "array", - "items": {"type": "string"} - } - }, - "required": ["query", "pos_doc", "neg_doc"], - "title": "EmbeddingDatasetItemSchema", - "type": "object" - } - } - ], - "training_options": [ - { - "training_type": "sft", - "finetuning_type": "lora_merged", - "num_gpus": 1, - "num_nodes": 1, - "tensor_parallel_size": 1, - "use_sequence_parallel": false - } - ] - } - ] -} diff --git a/docs/customizer/tutorials/_snippets/output/job-creation-response.json b/docs/customizer/tutorials/_snippets/output/job-creation-response.json deleted file mode 100644 index 11c3c6741f..0000000000 --- a/docs/customizer/tutorials/_snippets/output/job-creation-response.json +++ /dev/null @@ -1,25 +0,0 @@ -{ - "id": "cust-Pi95UoDbNcqwgkruAB8LY6", - "created_at": "2025-02-19T20:10:06.278132", - "updated_at": "2025-02-19T20:10:06.278133", - "workspace": "default", - "spec": { - "model": "default/llama-3.2-1b-instruct", - "dataset": "fileset://default/test-dataset", - "training": { - "type": "sft", - "peft": { - "type": "lora", - "rank": 8, - "alpha": 32, - "dropout": 0.01 - }, - "batch_size": 16, - "epochs": 10, - "learning_rate": 0.0001 - }, - "output": {"name": "llama-3.2-1b-instruct-test-dataset-a1b2c3d4e5f6", "type": "adapter", "fileset": "llama-3.2-1b-instruct-test-dataset-a1b2c3d4e5f6"} - }, - "status": "created", - "custom_fields": {} -} diff --git a/docs/customizer/tutorials/_snippets/output/lora-hyperparameters-example.json b/docs/customizer/tutorials/_snippets/output/lora-hyperparameters-example.json deleted file mode 100644 index 114aa5c34d..0000000000 --- a/docs/customizer/tutorials/_snippets/output/lora-hyperparameters-example.json +++ /dev/null @@ -1,13 +0,0 @@ -{ - "training": { - "type": "sft", - "peft": { - "type": "lora", - "rank": 8, - "dropout": 0.01 - }, - "epochs": 10, - "batch_size": 16, - "learning_rate": 0.0001 - } -} diff --git a/docs/customizer/tutorials/_snippets/output/lora-training-data-example.jsonl b/docs/customizer/tutorials/_snippets/output/lora-training-data-example.jsonl deleted file mode 100644 index e0fb6a2dc9..0000000000 --- a/docs/customizer/tutorials/_snippets/output/lora-training-data-example.jsonl +++ /dev/null @@ -1,2 +0,0 @@ -{"prompt": "You are an email writting assistant, please help people write cogent emails.\n\nTell John I finally got the PS5 with 1TB from Target and will be home at 8 pm so we can play Call of Duty. State Jason told me it was fun.", "completion": "Dear John, I got it! I managed to find the new PS5 at Target. I will be home by 8pm so we can play Call of Duty on it. It actually is a lot of fun (or so has Jason told me). It is the 1tb one so I hope it serves us well. See you later. Jackson Lee"} -{"prompt": "You are an email writting assistant, please help people write cogent emails.\n\nCheck in with the horseback riding facility if my daughter can go for a trial ride to see if she likes it. Ask if the horses are gentle and what the capacity is for each lesson.","completion": "Hello! My 9-year-old is interested in horseback riding lessons. She's only ridden ponies at the beach, but had taken up a kean interest in horses lately. How many children are in each lesson group? Are the horses gentle? Would she be able to try a lesson before committing to a full six-month contract? Thank you for replying at your earliest convenience. Heather"} diff --git a/docs/customizer/tutorials/_snippets/output/version-error-example.json b/docs/customizer/tutorials/_snippets/output/version-error-example.json deleted file mode 100644 index 9dc8710d4f..0000000000 --- a/docs/customizer/tutorials/_snippets/output/version-error-example.json +++ /dev/null @@ -1 +0,0 @@ -{ "detail": "Version is not specified in the config URN: meta/llama-3.2-1b-instruct" } diff --git a/docs/customizer/tutorials/distillation-customization-job.ipynb b/docs/customizer/tutorials/distillation-customization-job.ipynb index 00d6515395..af0e2830a6 100644 --- a/docs/customizer/tutorials/distillation-customization-job.ipynb +++ b/docs/customizer/tutorials/distillation-customization-job.ipynb @@ -1,930 +1,932 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "\n", - "\n", - "\n", - "# Knowledge Distillation Customization\n", - "\n", - "Learn how to train a smaller student model to mimic a larger teacher model using knowledge distillation (KD).\n", - "\n", - "## About\n", - "\n", - "Knowledge Distillation transfers knowledge from a large **teacher** model to a smaller **student** model. During training, the student learns to match the teacher's output probability distribution, producing a compact model that retains much of the teacher's capability.\n", - "\n", - "**What you can achieve with KD:**\n", - "\n", - "- **Compress models:** Distill a 3B model into a 1B model for faster inference and lower deployment costs\n", - "- **Reduce latency:** Deploy a smaller model that responds faster while preserving quality\n", - "- **Lower resource requirements:** Serve a distilled model on fewer GPUs\n", - "\n", - "### KD vs SFT: Understanding the Trade-offs\n", - "\n", - "| Aspect | Full SFT | Knowledge Distillation |\n", - "| --- | --- | --- |\n", - "| **Training signal** | Ground-truth labels only | Teacher's soft probability distribution + labels |\n", - "| **Knowledge source** | Dataset examples | Teacher model's learned representations |\n", - "| **Output model size** | Same as input model | Typically a smaller student model |\n", - "| **GPU requirements** | Needs to fit one model | Needs to fit both teacher and student in memory |\n", - "| **Best for** | Domain adaptation, new knowledge injection | Model compression, latency reduction |\n", - "\n", - "### Key Parameters\n", - "\n", - "| Parameter | Default | Description |\n", - "| --- | --- | --- |\n", - "| `teacher_model` | *(required)* | Teacher model entity URN (e.g., `default/llama-3-2-3b-teacher`) |\n", - "| `teacher_precision` | `bf16` | Precision for the frozen teacher (`bf16`, `fp16`, `fp32`). Lower = less memory |\n", - "| `distillation_ratio` | `0.5` | Balance between CE loss and KD loss. `0.0` = CE only, `1.0` = KD only |\n", - "| `distillation_temperature` | `1.0` | Softmax temperature. Higher = softer distributions, more knowledge transfer |\n", - "\n", - "### Workflow Overview\n", - "\n", - "This tutorial follows a complete distillation pipeline:\n", - "\n", - "1. **Fine-tune the teacher** (SFT on the task dataset) so it learns the domain\n", - "2. **Establish a baseline** by deploying the base student model and measuring ROUGE scores\n", - "3. **Distill into the student** using the fine-tuned teacher's soft targets\n", - "4. **Evaluate the distilled student** and compare ROUGE scores against the baseline\n", - "\n", - "**When to choose KD:**\n", - "\n", - "- You have a high-quality large model and want a smaller, faster version\n", - "- Deployment latency or cost is a constraint\n", - "- The teacher and student share the same vocabulary (e.g., both are Llama models)\n", - "\n", - "**When to choose SFT instead:** Refer to the [Full SFT tutorial](./sft-customization-job) when you want to train a model directly on labeled data without a teacher." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Prerequisites\n", - "\n", - "Before starting this tutorial, ensure you have:\n", - "\n", - "1. **Completed the [Quickstart](../../get-started/quickstart.md)** to install and deploy NeMo Platform locally\n", - "2. **Installed the Python SDK** (PyPI wrapper: `pip install \"nemo-platform[all]\"`; source checkout: run `make bootstrap` from the repository root)\n", - "3. **Installed evaluation dependencies:**\n", - "\n", - "```sh\n", - "pip install evaluate rouge_score datasets\n", - "```" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Quick Start\n", - "\n", - "### 1. Initialize SDK\n", - "\n", - "The SDK needs to know your NeMo Platform server URL. By default, `http://localhost:8080` is used in accordance with the [Quickstart](../../get-started/quickstart.md) guide. If NeMo Platform is running at a custom location, you can override the URL by setting the `NMP_BASE_URL` environment variable:\n", - "\n", - "```sh\n", - "export NMP_BASE_URL=\n", - "```" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "import json\n", - "import os\n", - "import time\n", - "import uuid\n", - "from pathlib import Path\n", - "\n", - "from nemo_platform import NeMoPlatform, ConflictError\n", - "\n", - "NMP_BASE_URL = os.environ.get(\"NMP_BASE_URL\", \"http://localhost:8080\")\n", - "client = NeMoPlatform(\n", - " base_url=NMP_BASE_URL,\n", - " workspace=\"default\"\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### 2. Prepare Dataset\n", - "\n", - "Knowledge distillation uses the same dataset formats as SFT. We use the SQuAD dataset for both teacher training and distillation so that the teacher first learns the task, then transfers that knowledge to the student.\n", - "\n", - "We also hold out a small **test split** for ROUGE evaluation at the end." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from datasets import load_dataset, DatasetDict\n", - "\n", - "print(\"Loading dataset rajpurkar/squad\")\n", - "raw_dataset = load_dataset(\"rajpurkar/squad\")\n", - "if not isinstance(raw_dataset, DatasetDict):\n", - " raise ValueError(\"Dataset does not contain expected splits\")\n", - "\n", - "print(\"Loaded dataset\")\n", - "\n", - "SEED = 1234\n", - "TRAINING_SIZE = 3000\n", - "VALIDATION_SIZE = 300\n", - "TEST_SIZE = 100\n", - "DATASET_PATH = Path(\"kd-dataset\").absolute()\n", - "\n", - "os.makedirs(DATASET_PATH, exist_ok=True)\n", - "\n", - "train_set = raw_dataset.get('train')\n", - "split = train_set.train_test_split(test_size=0.05, seed=SEED)\n", - "\n", - "train_ds = split['train'].select(range(min(TRAINING_SIZE, len(split['train']))))\n", - "val_ds = split['test'].select(range(min(VALIDATION_SIZE, len(split['test']))))\n", - "test_ds = split['test'].select(range(VALIDATION_SIZE, min(VALIDATION_SIZE + TEST_SIZE, len(split['test']))))\n", - "\n", - "\n", - "def convert_squad(example):\n", - " \"\"\"Convert SQuAD format to prompt/completion format.\"\"\"\n", - " prompt = f\"Context: {example['context']} Question: {example['question']} Answer:\"\n", - " completion = example[\"answers\"][\"text\"][0]\n", - " return {\"prompt\": prompt, \"completion\": completion}\n", - "\n", - "\n", - "def write_jsonl(dataset, path):\n", - " with open(path, \"w\", encoding=\"utf-8\") as f:\n", - " for example in dataset:\n", - " f.write(json.dumps(convert_squad(example)) + \"\\n\")\n", - "\n", - "\n", - "def write_test_jsonl(dataset, path):\n", - " \"\"\"Save test split with raw context/question for chat-style evaluation.\"\"\"\n", - " with open(path, \"w\", encoding=\"utf-8\") as f:\n", - " for example in dataset:\n", - " f.write(json.dumps({\n", - " \"context\": example[\"context\"],\n", - " \"question\": example[\"question\"],\n", - " \"completion\": example[\"answers\"][\"text\"][0],\n", - " }) + \"\\n\")\n", - "\n", - "\n", - "write_jsonl(train_ds, f\"{DATASET_PATH}/training.jsonl\")\n", - "write_jsonl(val_ds, f\"{DATASET_PATH}/validation.jsonl\")\n", - "write_test_jsonl(test_ds, f\"{DATASET_PATH}/testing.jsonl\")\n", - "\n", - "print(f\"Training: {len(train_ds)} rows\")\n", - "print(f\"Validation: {len(val_ds)} rows\")\n", - "print(f\"Test: {len(test_ds)} rows\")\n", - "\n", - "with open(f\"{DATASET_PATH}/training.jsonl\", 'r') as f:\n", - " sample = json.loads(f.readline())\n", - " print(f\"\\nSample prompt: {sample['prompt'][:150]}...\")\n", - " print(f\"Sample completion: {sample['completion']}\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "DATASET_NAME = \"kd-dataset\"\n", - "\n", - "try:\n", - " client.files.filesets.create(\n", - " workspace=\"default\",\n", - " name=DATASET_NAME,\n", - " description=\"Knowledge distillation training data\"\n", - " )\n", - " print(f\"Created fileset: {DATASET_NAME}\")\n", - "except ConflictError:\n", - " print(f\"Fileset '{DATASET_NAME}' already exists, continuing...\")\n", - "\n", - "client.files.upload(\n", - " local_path=f\"{DATASET_PATH}/\",\n", - " remote_path=\"\",\n", - " fileset=DATASET_NAME,\n", - " workspace=\"default\"\n", - ")\n", - "\n", - "print(\"Uploaded files:\")\n", - "print(json.dumps([f.model_dump() for f in client.files.list(fileset=DATASET_NAME, workspace=\"default\").data], indent=2))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### 3. Secrets Setup\n", - "\n", - "In this tutorial we use two Llama 3.2 Instruct models from HuggingFace:\n", - "- **Teacher:** [meta-llama/Llama-3.2-3B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-3B-Instruct) (3B parameters)\n", - "- **Student:** [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct) (1B parameters)\n", - "\n", - "Both models share the same tokenizer/vocabulary (required for knowledge distillation) and include a chat template for deployment with `/chat/completions`.\n", - "\n", - "**HuggingFace Authentication:**\n", - "- For gated models (Llama, Gemma), you must provide a HuggingFace token via the `token_secret` parameter\n", - "- Get your token from [HuggingFace Settings](https://huggingface.co/settings/tokens) (requires Read access)\n", - "- Accept the model's terms on the HuggingFace model page before using it:\n", - " - [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct)\n", - " - [meta-llama/Llama-3.2-3B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-3B-Instruct)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "HF_TOKEN = os.getenv(\"HF_TOKEN\")\n", - "\n", - "\n", - "def create_or_get_secret(name: str, value: str | None, label: str):\n", - " if not value:\n", - " raise ValueError(f\"{label} is not set\")\n", - " try:\n", - " secret = client.secrets.create(\n", - " name=name,\n", - " workspace=\"default\",\n", - " value=value,\n", - " )\n", - " print(f\"Created secret: {name}\")\n", - " return secret\n", - " except ConflictError:\n", - " print(f\"Secret '{name}' already exists, continuing...\")\n", - " return client.secrets.retrieve(name=name, workspace=\"default\")\n", - "\n", - "\n", - "hf_secret = create_or_get_secret(\"hf-token\", HF_TOKEN, \"HF_TOKEN\")\n", - "print(\"HF_TOKEN secret:\")\n", - "print(hf_secret.model_dump_json(indent=2))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### 4. Create Model FileSets and Model Entities\n", - "\n", - "Knowledge distillation requires **two** model entities:\n", - "1. **Student model** — the smaller model that will be trained ([meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct))\n", - "2. **Teacher model** — the larger model that provides soft targets ([meta-llama/Llama-3.2-3B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-3B-Instruct))\n", - "\n", - "Using the Instruct variants ensures the output model includes a chat template, which is required for the `/chat/completions` inference endpoint." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from nemo_platform.types.files import HuggingfaceStorageConfigParam\n", - "\n", - "SPEC_TIMEOUT_SECONDS = 120\n", - "\n", - "\n", - "def create_model(hf_repo: str, model_name: str, description: str):\n", - " \"\"\"Create a fileset + model entity and wait for ModelSpec.\"\"\"\n", - " try:\n", - " client.files.filesets.create(\n", - " workspace=\"default\",\n", - " name=model_name,\n", - " description=description,\n", - " storage=HuggingfaceStorageConfigParam(\n", - " type=\"huggingface\",\n", - " repo_id=hf_repo,\n", - " repo_type=\"model\",\n", - " token_secret=hf_secret.name\n", - " )\n", - " )\n", - " print(f\"Created fileset: {model_name}\")\n", - " except ConflictError:\n", - " print(f\"Fileset '{model_name}' already exists.\")\n", - "\n", - " try:\n", - " model = client.models.create(\n", - " workspace=\"default\",\n", - " name=model_name,\n", - " fileset=f\"default/{model_name}\",\n", - " )\n", - " print(f\"Created Model Entity: {model_name}\")\n", - " except ConflictError:\n", - " print(f\"Model '{model_name}' already exists. Updating fileset.\")\n", - " model = client.models.update(\n", - " workspace=\"default\",\n", - " name=model_name,\n", - " fileset=f\"default/{model_name}\",\n", - " )\n", - "\n", - " print(f\"Waiting for ModelSpec on {model_name}...\")\n", - " spec_start = time.time()\n", - " while not model.spec:\n", - " if time.time() - spec_start > SPEC_TIMEOUT_SECONDS:\n", - " raise TimeoutError(f\"ModelSpec not populated within {SPEC_TIMEOUT_SECONDS}s\")\n", - " time.sleep(2)\n", - " model = client.models.retrieve(workspace=\"default\", name=model_name)\n", - " print(f\"ModelSpec populated: {model.spec}\")\n", - " return model\n", - "\n", - "\n", - "student_model = create_model(\n", - " hf_repo=\"meta-llama/Llama-3.2-1B-Instruct\",\n", - " model_name=\"llama-3-2-1b-student\",\n", - " description=\"Llama 3.2 1B Instruct student model\",\n", - ")\n", - "\n", - "print()\n", - "\n", - "teacher_model = create_model(\n", - " hf_repo=\"meta-llama/Llama-3.2-3B-Instruct\",\n", - " model_name=\"llama-3-2-3b-teacher\",\n", - " description=\"Llama 3.2 3B Instruct teacher model\",\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "---\n", - "\n", - "## Phase 1: Fine-Tune the Teacher\n", - "\n", - "### 5. Train Teacher with Full SFT\n", - "\n", - "For best distillation results, fine-tune the teacher on the **same dataset** that will be used for distillation. This ensures the teacher has learned the task-specific knowledge that the student will inherit.\n", - "\n", - "We train the 3B Instruct model with Full SFT on the SQuAD dataset." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from nemo_automodel_plugin.schema import AutomodelJobInput\n", - "\n", - "job_suffix = uuid.uuid4().hex[:4]\n", - "\n", - "TEACHER_JOB_NAME = f\"teacher-sft-job-{job_suffix}\"\n", - "TEACHER_OUTPUT_NAME = f\"teacher-model-{job_suffix}\"\n", - "\n", - "teacher_spec = AutomodelJobInput(\n", - " model=f\"default/{teacher_model.name}\",\n", - " dataset={\"training\": f\"default/{DATASET_NAME}\"},\n", - " training={\n", - " \"training_type\": \"sft\",\n", - " \"finetuning_type\": \"all_weights\",\n", - " \"max_seq_length\": 2048,\n", - " },\n", - " schedule={\"epochs\": 1},\n", - " batch={\"global_batch_size\": 64, \"micro_batch_size\": 1},\n", - " optimizer={\"learning_rate\": 5e-5},\n", - " parallelism={\"num_gpus_per_node\": 1},\n", - " output={\"name\": TEACHER_OUTPUT_NAME},\n", - ")\n", - "\n", - "teacher_job = client.customization.automodel.jobs.create(\n", - " spec=teacher_spec, workspace=\"default\", name=TEACHER_JOB_NAME\n", - ")\n", - "\n", - "TRAINED_TEACHER_NAME = TEACHER_OUTPUT_NAME\n", - "print(f\"Teacher training job: {teacher_job.job.name}\")\n", - "print(f\"Output teacher model: {TRAINED_TEACHER_NAME}\")\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from IPython.display import clear_output\n", - "\n", - "\n", - "def wait_for_job(job_name: str):\n", - " \"\"\"Poll job status until completion.\"\"\"\n", - " while True:\n", - " status = client.jobs.get_status(name=job_name, workspace=\"default\")\n", - " clear_output(wait=True)\n", - " print(f\"Job: {job_name}\")\n", - " print(f\"Status: {status.status}\")\n", - "\n", - " for job_step in status.steps or []:\n", - " if job_step.name == \"training\":\n", - " for task in job_step.tasks or []:\n", - " details = task.status_details or {}\n", - " step = details.get(\"step\")\n", - " max_steps = details.get(\"max_steps\")\n", - " if step is not None and max_steps is not None:\n", - " print(f\"Progress: Step {step}/{max_steps} ({step / max_steps * 100:.1f}%)\")\n", - " phase = details.get(\"phase\")\n", - " if phase:\n", - " print(f\"Phase: {phase}\")\n", - " break\n", - " break\n", - "\n", - " if status.status in (\"completed\", \"failed\", \"cancelled\", \"error\"):\n", - " print(f\"\\nJob finished: {status.status}\")\n", - " return status\n", - "\n", - " time.sleep(10)\n", - "\n", - "\n", - "teacher_status = wait_for_job(TEACHER_JOB_NAME)\n", - "assert teacher_status.status == \"completed\"" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "---\n", - "\n", - "## Phase 2: Establish Baseline (Base Student)\n", - "\n", - "### 6. Deploy the Base Student Model\n", - "\n", - "Before distillation, deploy the base student model (1B Instruct, without any fine-tuning) to establish a baseline ROUGE score. After distillation, we compare the distilled student against this baseline to measure improvement." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "baseline_suffix = uuid.uuid4().hex[:4]\n", - "BASELINE_DEPLOYMENT_CONFIG = f\"baseline-student-cfg-{baseline_suffix}\"\n", - "BASELINE_DEPLOYMENT_NAME = f\"baseline-student-{baseline_suffix}\"\n", - "\n", - "baseline_deployment_config = client.inference.deployment_configs.create(\n", - " workspace=\"default\",\n", - " name=BASELINE_DEPLOYMENT_CONFIG,\n", - " engine=\"vllm\",\n", - " model_spec={\n", - " \"model_namespace\": \"default\",\n", - " \"model_name\": student_model.name,\n", - " },\n", - " executor_config={\n", - " \"gpu\": 1,\n", - " \"image_name\": \"vllm/vllm-openai\",\n", - " \"image_tag\": \"v0.22.1\",\n", - " },\n", - ")\n", - "\n", - "baseline_deployment = client.inference.deployments.create(\n", - " workspace=\"default\",\n", - " name=BASELINE_DEPLOYMENT_NAME,\n", - " config=baseline_deployment_config.name\n", - ")\n", - "\n", - "print(f\"Baseline student deployment: {baseline_deployment.name}\")\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "def wait_for_deployment(deployment_name: str, timeout_minutes: int = 30):\n", - " \"\"\"Poll deployment until ready.\"\"\"\n", - " start = time.time()\n", - " timeout = timeout_minutes * 60\n", - " while True:\n", - " dep = client.inference.deployments.retrieve(name=deployment_name, workspace=\"default\")\n", - " elapsed = time.time() - start\n", - " clear_output(wait=True)\n", - " print(f\"Deployment: {deployment_name}\")\n", - " print(f\"Status: {dep.status}\")\n", - " print(f\"Elapsed: {int(elapsed // 60)}m {int(elapsed % 60)}s\")\n", - "\n", - " if dep.status == \"READY\":\n", - " print(\"\\nDeployment is ready!\")\n", - " return dep\n", - " if dep.status in (\"FAILED\", \"ERROR\", \"TERMINATED\", \"LOST\"):\n", - " print(f\"\\nDeployment failed: {dep.status}\")\n", - " return dep\n", - " if elapsed > timeout:\n", - " print(f\"\\nTimeout ({timeout_minutes}m). Check status manually.\")\n", - " return dep\n", - " time.sleep(15)\n", - "\n", - "\n", - "dep_status = wait_for_deployment(BASELINE_DEPLOYMENT_NAME)\n", - "assert dep_status.status == \"READY\"" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### 7. Generate Baseline Predictions on Test Set" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "with open(f\"{DATASET_PATH}/testing.jsonl\", \"r\", encoding=\"utf-8\") as f:\n", - " test_data = [json.loads(line) for line in f]\n", - "\n", - "contexts = [row[\"context\"] for row in test_data]\n", - "questions = [row[\"question\"] for row in test_data]\n", - "reference_completions = [row[\"completion\"] for row in test_data]\n", - "\n", - "print(f\"Test samples: {len(contexts)}\")\n", - "print(f\"Sample context: {contexts[0]}\")\n", - "print(f\"Sample question: {questions[0]}\")\n", - "print(f\"Sample reference: {reference_completions[0]}\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "def generate_completions(\n", - " deployment_name: str,\n", - " output_model_name: str,\n", - " contexts: list[str],\n", - " questions: list[str],\n", - ") -> list[str]:\n", - " \"\"\"Generate completions for a list of context/question pairs using a deployed model.\"\"\"\n", - " completions = []\n", - " for context, question in zip(contexts, questions):\n", - " messages = [\n", - " {\n", - " \"role\": \"user\",\n", - " \"content\": f\"Based on the following context, answer the question.\\n\\nContext: {context}\\n\\nQuestion: {question}\",\n", - " }\n", - " ]\n", - " response = client.inference.gateway.provider.post(\n", - " \"v1/chat/completions\",\n", - " name=deployment_name,\n", - " workspace=\"default\",\n", - " body={\n", - " \"model\": f\"default/{output_model_name}\",\n", - " \"messages\": messages,\n", - " \"temperature\": 0,\n", - " \"max_tokens\": 128,\n", - " }\n", - " )\n", - " completions.append(response[\"choices\"][0][\"message\"][\"content\"])\n", - " return completions\n", - "\n", - "\n", - "print(\"Generating baseline (base student) predictions...\")\n", - "baseline_completions = generate_completions(BASELINE_DEPLOYMENT_NAME, student_model.name, contexts, questions)\n", - "print(f\"Generated {len(baseline_completions)} baseline predictions\")\n", - "print(f\"\\nSample baseline output: {baseline_completions[0]}\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### 8. Delete Baseline Deployment\n", - "\n", - "Delete the baseline student deployment to free GPU resources for the distillation training job and subsequent distilled model deployment." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "client.inference.deployments.delete(name=BASELINE_DEPLOYMENT_NAME, workspace=\"default\")\n", - "print(f\"Deleted baseline deployment: {BASELINE_DEPLOYMENT_NAME}\")\n", - "\n", - "if not client.models.wait_for_status(\n", - " deployment_name=BASELINE_DEPLOYMENT_NAME,\n", - " desired_status=\"DELETED\",\n", - " workspace=\"default\",\n", - " timeout=600,\n", - "):\n", - " raise TimeoutError(\n", - " f\"Deployment {BASELINE_DEPLOYMENT_NAME} was not deleted within timeout\"\n", - " )\n", - "\n", - "client.inference.deployment_configs.delete(name=BASELINE_DEPLOYMENT_CONFIG, workspace=\"default\")\n", - "print(f\"Deleted baseline deployment config: {BASELINE_DEPLOYMENT_CONFIG}\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "---\n", - "\n", - "## Phase 3: Distill into Student\n", - "\n", - "### 9. Create Knowledge Distillation Job\n", - "\n", - "Now create a distillation job that trains the 1B student using the **fine-tuned** 3B teacher's output distribution. The `model` field specifies the student, and `teacher_model` references the trained teacher model entity from Phase 1." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "**GPU Requirements:**\n", - "\n", - "KD requires loading both student and teacher models, so plan GPU memory accordingly:\n", - "- 1B student + 3B teacher: 1 GPU (24GB+ VRAM each)\n", - "- 3B student + 8B teacher: 4 GPUs\n", - "- 8B student + 70B teacher: 8+ GPUs\n", - "\n", - "Use `teacher_precision=\"bf16\"` (default) to reduce teacher memory footprint." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from nemo_automodel_plugin.schema import AutomodelJobInput\n", - "\n", - "KD_JOB_NAME = f\"my-kd-job-{job_suffix}\"\n", - "KD_OUTPUT_NAME = f\"kd-student-{job_suffix}\"\n", - "\n", - "kd_spec = AutomodelJobInput(\n", - " model=f\"default/{student_model.name}\",\n", - " dataset={\"training\": f\"default/{DATASET_NAME}\"},\n", - " training={\n", - " \"training_type\": \"distillation\",\n", - " \"finetuning_type\": \"all_weights\",\n", - " \"teacher_model\": f\"default/{TRAINED_TEACHER_NAME}\",\n", - " \"teacher_precision\": \"bf16\",\n", - " \"distillation_ratio\": 0.5,\n", - " \"distillation_temperature\": 2.0,\n", - " \"max_seq_length\": 2048,\n", - " },\n", - " schedule={\"epochs\": 1},\n", - " batch={\"global_batch_size\": 64, \"micro_batch_size\": 1},\n", - " optimizer={\"learning_rate\": 5e-5},\n", - " parallelism={\"num_gpus_per_node\": 1},\n", - " output={\"name\": KD_OUTPUT_NAME},\n", - ")\n", - "\n", - "kd_job = client.customization.automodel.jobs.create(\n", - " spec=kd_spec, workspace=\"default\", name=KD_JOB_NAME\n", - ")\n", - "\n", - "DISTILLED_STUDENT_NAME = KD_OUTPUT_NAME\n", - "print(f\"Distillation job: {kd_job.job.name}\")\n", - "print(f\"Output student model: {DISTILLED_STUDENT_NAME}\")\n" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### 10. Track Distillation Progress" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "kd_status = wait_for_job(KD_JOB_NAME)\n", - "assert kd_status.status == \"completed\"" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "---\n", - "\n", - "## Phase 4: Evaluate the Distilled Student Model\n", - "\n", - "### 11. Deploy the Distilled Student Model\n", - "\n", - "The output model has the same architecture as the 1B student—only its weights have been updated via distillation. It requires just 1 GPU to deploy." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "deploy_suffix_2 = uuid.uuid4().hex[:4]\n", - "STUDENT_DEPLOYMENT_CONFIG = f\"kd-student-deploy-cfg-{deploy_suffix_2}\"\n", - "STUDENT_DEPLOYMENT_NAME = f\"kd-student-deploy-{deploy_suffix_2}\"\n", - "\n", - "student_deployment_config = client.inference.deployment_configs.create(\n", - " workspace=\"default\",\n", - " name=STUDENT_DEPLOYMENT_CONFIG,\n", - " engine=\"vllm\",\n", - " model_spec={\n", - " \"model_namespace\": \"default\",\n", - " \"model_name\": DISTILLED_STUDENT_NAME,\n", - " },\n", - " executor_config={\n", - " \"gpu\": 1,\n", - " \"image_name\": \"vllm/vllm-openai\",\n", - " \"image_tag\": \"v0.22.1\",\n", - " },\n", - ")\n", - "\n", - "student_deployment = client.inference.deployments.create(\n", - " workspace=\"default\",\n", - " name=STUDENT_DEPLOYMENT_NAME,\n", - " config=student_deployment_config.name\n", - ")\n", - "\n", - "print(f\"Student deployment: {student_deployment.name}\")\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "wait_for_deployment(STUDENT_DEPLOYMENT_NAME)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### 12. Generate Student Predictions on Test Set" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "print(\"Generating distilled student predictions...\")\n", - "student_completions = generate_completions(STUDENT_DEPLOYMENT_NAME, DISTILLED_STUDENT_NAME, contexts, questions)\n", - "print(f\"Generated {len(student_completions)} student predictions\")\n", - "print(f\"\\nSample student output: {student_completions[0]}\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### 13. Compute ROUGE Scores\n", - "\n", - "Compare the base student (before distillation) and the distilled student against the ground-truth reference completions using ROUGE metrics." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "import evaluate\n", - "\n", - "rouge = evaluate.load(\"rouge\")\n", - "\n", - "baseline_scores = rouge.compute(predictions=baseline_completions, references=reference_completions)\n", - "student_scores = rouge.compute(predictions=student_completions, references=reference_completions)\n", - "\n", - "metrics = list(baseline_scores.keys())\n", - "header = f\"{'Model':<35} \" + \" \".join(f\"{m:>10}\" for m in metrics)\n", - "separator = \"-\" * len(header)\n", - "\n", - "print(\"=\" * 60)\n", - "print(\"ROUGE SCORE COMPARISON\")\n", - "print(\"=\" * 60)\n", - "print(header)\n", - "print(separator)\n", - "print(f\"{'Base Student (1B, no training)':<35} \" + \" \".join(f\"{baseline_scores[m]:>10.4f}\" for m in metrics))\n", - "print(f\"{'Distilled Student (1B, KD)':<35} \" + \" \".join(f\"{student_scores[m]:>10.4f}\" for m in metrics))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "print(\"Sample predictions (first 3):\\n\")\n", - "for i in range(min(3, len(contexts))):\n", - " print(f\"--- Sample {i + 1} ---\")\n", - " print(f\"Question: {questions[i]}\")\n", - " print(f\"Reference: {reference_completions[i]}\")\n", - " print(f\"Baseline: {baseline_completions[i][:200]}\")\n", - " print(f\"Distilled: {student_completions[i][:200]}\")\n", - " print()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "**Interpreting ROUGE Scores:**\n", - "\n", - "| Metric | Measures |\n", - "|--------|----------|\n", - "| **ROUGE-1** | Unigram overlap between prediction and reference |\n", - "| **ROUGE-2** | Bigram overlap (captures phrase-level similarity) |\n", - "| **ROUGE-L** | Longest common subsequence (captures sentence structure) |\n", - "| **ROUGE-Lsum** | ROUGE-L computed over full summaries |\n", - "\n", - "**What to expect:**\n", - "- The base student (1B, no training) provides a lower bound since it has not seen the task data\n", - "- The distilled student (1B, KD) should significantly outperform the base student, demonstrating the knowledge transferred from the 3B teacher\n", - "- If the distilled student scores are not much higher than the baseline, try increasing `distillation_temperature`, adjusting `distillation_ratio`, or training for more epochs\n", - "\n", - "---\n", - "\n", - "## Hyperparameters\n", - "\n", - "For detailed information on all available hyperparameters, recommended values, and tuning guidance, refer to the [Hyperparameter Reference](../manage-customization-jobs/hyperparameters.md).\n", - "\n", - "---\n", - "\n", - "## Troubleshooting\n", - "\n", - "**Job fails during model download:**\n", - "- Verify authentication secrets are configured (refer to [Managing Secrets](../../get-started/concepts/manage-secrets.md))\n", - "- For gated HuggingFace models (Llama, Gemma), accept the license on the model page\n", - "- Check both `model` (student) and `teacher_model` URNs are correct\n", - "- Ensure both model entities exist: `client.models.retrieve(name=..., workspace=\"default\")`\n", - "\n", - "**Job fails with OOM (Out of Memory) error:**\n", - "\n", - "KD loads both models, so OOM is more likely than with SFT:\n", - "1. **First try:** Use `teacher_precision=\"bf16\"` to reduce teacher memory\n", - "2. **Still OOM:** Reduce `micro_batch_size` to 1\n", - "3. **Still OOM:** Reduce `global_batch_size` and `max_seq_length`\n", - "4. **Last resort:** Increase `num_gpus_per_node`\n", - "\n", - "**No chat template / `/chat/completions` fails:**\n", - "- Use Instruct model variants (e.g., `Llama-3.2-1B-Instruct`) instead of base models (`Llama-3.2-1B`). Base models do not include a chat template in their tokenizer, so the output model will also lack one.\n", - "\n", - "**Distilled model quality is poor:**\n", - "- Increase `distillation_temperature` (try 2.0–5.0) to transfer more nuanced knowledge\n", - "- Adjust `distillation_ratio`—if dataset labels are high-quality, lower the ratio; if the teacher is strong, raise it\n", - "- Increase `epochs` or `max_steps` for more training\n", - "- Verify teacher and student share the same vocabulary\n", - "\n", - "**Vocabulary mismatch error:**\n", - "- Teacher and student must use the same tokenizer. Use models from the same family (e.g., Llama 3.2 1B Instruct + Llama 3.2 3B Instruct)\n", - "\n", - "**Deployment fails:**\n", - "- Verify output model exists: `client.models.retrieve(name=DISTILLED_STUDENT_NAME, workspace=\"default\")`\n", - "- Check deployment logs: `client.inference.deployments.get_logs(name=deployment.name, workspace=\"default\")`\n", - "- The distilled model has the same size as the student, so GPU requirements match the student model\n", - "\n", - "\n", - "## Next Steps\n", - "\n", - "- [Monitor training metrics](fine-tune-metrics) in detail\n", - "- [Evaluate your fine-tuned model](../../evaluator/index) using the Evaluator service\n", - "- Learn about [LoRA customization](./lora-customization-job) for resource-efficient fine-tuning\n", - "- Learn about [Full SFT](./sft-customization-job) for direct supervised fine-tuning" - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": ".venv", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.11.14" - } - }, - "nbformat": 4, - "nbformat_minor": 4 -} + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "\n", + "\n", + "\n", + "# Knowledge Distillation Customization\n", + "\n", + "Learn how to train a smaller student model to mimic a larger teacher model using knowledge distillation (KD).\n", + "\n", + "## About\n", + "\n", + "Knowledge Distillation transfers knowledge from a large **teacher** model to a smaller **student** model. During training, the student learns to match the teacher's output probability distribution, producing a compact model that retains much of the teacher's capability.\n", + "\n", + "**What you can achieve with KD:**\n", + "\n", + "- **Compress models:** Distill a 3B model into a 1B model for faster inference and lower deployment costs\n", + "- **Reduce latency:** Deploy a smaller model that responds faster while preserving quality\n", + "- **Lower resource requirements:** Serve a distilled model on fewer GPUs\n", + "\n", + "### KD vs SFT: Understanding the Trade-offs\n", + "\n", + "| Aspect | Full SFT | Knowledge Distillation |\n", + "| --- | --- | --- |\n", + "| **Training signal** | Ground-truth labels only | Teacher's soft probability distribution + labels |\n", + "| **Knowledge source** | Dataset examples | Teacher model's learned representations |\n", + "| **Output model size** | Same as input model | Typically a smaller student model |\n", + "| **GPU requirements** | Needs to fit one model | Needs to fit both teacher and student in memory |\n", + "| **Best for** | Domain adaptation, new knowledge injection | Model compression, latency reduction |\n", + "\n", + "### Key Parameters\n", + "\n", + "| Parameter | Default | Description |\n", + "| --- | --- | --- |\n", + "| `teacher_model` | *(required)* | Teacher model entity URN (e.g., `default/llama-3-2-3b-teacher`) |\n", + "| `teacher_precision` | `bf16` | Precision for the frozen teacher (`bf16`, `fp16`, `fp32`). Lower = less memory |\n", + "| `distillation_ratio` | `0.5` | Balance between CE loss and KD loss. `0.0` = CE only, `1.0` = KD only |\n", + "| `distillation_temperature` | `1.0` | Softmax temperature. Higher = softer distributions, more knowledge transfer |\n", + "\n", + "### Workflow Overview\n", + "\n", + "This tutorial follows a complete distillation pipeline:\n", + "\n", + "1. **Fine-tune the teacher** (SFT on the task dataset) so it learns the domain\n", + "2. **Establish a baseline** by deploying the base student model and measuring ROUGE scores\n", + "3. **Distill into the student** using the fine-tuned teacher's soft targets\n", + "4. **Evaluate the distilled student** and compare ROUGE scores against the baseline\n", + "\n", + "**When to choose KD:**\n", + "\n", + "- You have a high-quality large model and want a smaller, faster version\n", + "- Deployment latency or cost is a constraint\n", + "- The teacher and student share the same vocabulary (e.g., both are Llama models)\n", + "\n", + "**When to choose SFT instead:** Refer to the [Full SFT tutorial](./sft-customization-job) when you want to train a model directly on labeled data without a teacher." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Prerequisites\n", + "\n", + "Before starting this tutorial, ensure you have:\n", + "\n", + "1. **Completed the [Quickstart](../../get-started/quickstart.md)** to install and deploy NeMo Platform locally\n", + "2. **Installed the Python SDK** (PyPI wrapper: `pip install \"nemo-platform[all]\"`; source checkout: run `make bootstrap` from the repository root)\n", + "3. **Installed evaluation dependencies:**\n", + "\n", + "```sh\n", + "pip install evaluate rouge_score datasets\n", + "```\n", + "\n", + "4. **At least one GPU with CUDA 13+**" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Quick Start\n", + "\n", + "### 1. Initialize SDK\n", + "\n", + "The SDK needs to know your NeMo Platform server URL. By default, `http://localhost:8080` is used in accordance with the [Quickstart](../../get-started/quickstart.md) guide. If NeMo Platform is running at a custom location, you can override the URL by setting the `NMP_BASE_URL` environment variable:\n", + "\n", + "```sh\n", + "export NMP_BASE_URL=\n", + "```" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "import json\n", + "import os\n", + "import time\n", + "import uuid\n", + "from pathlib import Path\n", + "\n", + "from nemo_platform import NeMoPlatform, ConflictError\n", + "\n", + "NMP_BASE_URL = os.environ.get(\"NMP_BASE_URL\", \"http://localhost:8080\")\n", + "client = NeMoPlatform(\n", + " base_url=NMP_BASE_URL,\n", + " workspace=\"default\"\n", + ")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 2. Prepare Dataset\n", + "\n", + "Knowledge distillation uses the same dataset formats as SFT. We use the SQuAD dataset for both teacher training and distillation so that the teacher first learns the task, then transfers that knowledge to the student.\n", + "\n", + "We also hold out a small **test split** for ROUGE evaluation at the end." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from datasets import load_dataset, DatasetDict\n", + "\n", + "print(\"Loading dataset rajpurkar/squad\")\n", + "raw_dataset = load_dataset(\"rajpurkar/squad\")\n", + "if not isinstance(raw_dataset, DatasetDict):\n", + " raise ValueError(\"Dataset does not contain expected splits\")\n", + "\n", + "print(\"Loaded dataset\")\n", + "\n", + "SEED = 1234\n", + "TRAINING_SIZE = 3000\n", + "VALIDATION_SIZE = 300\n", + "TEST_SIZE = 100\n", + "DATASET_PATH = Path(\"kd-dataset\").absolute()\n", + "\n", + "os.makedirs(DATASET_PATH, exist_ok=True)\n", + "\n", + "train_set = raw_dataset.get('train')\n", + "split = train_set.train_test_split(test_size=0.05, seed=SEED)\n", + "\n", + "train_ds = split['train'].select(range(min(TRAINING_SIZE, len(split['train']))))\n", + "val_ds = split['test'].select(range(min(VALIDATION_SIZE, len(split['test']))))\n", + "test_ds = split['test'].select(range(VALIDATION_SIZE, min(VALIDATION_SIZE + TEST_SIZE, len(split['test']))))\n", + "\n", + "\n", + "def convert_squad(example):\n", + " \"\"\"Convert SQuAD format to prompt/completion format.\"\"\"\n", + " prompt = f\"Context: {example['context']} Question: {example['question']} Answer:\"\n", + " completion = example[\"answers\"][\"text\"][0]\n", + " return {\"prompt\": prompt, \"completion\": completion}\n", + "\n", + "\n", + "def write_jsonl(dataset, path):\n", + " with open(path, \"w\", encoding=\"utf-8\") as f:\n", + " for example in dataset:\n", + " f.write(json.dumps(convert_squad(example)) + \"\\n\")\n", + "\n", + "\n", + "def write_test_jsonl(dataset, path):\n", + " \"\"\"Save test split with raw context/question for chat-style evaluation.\"\"\"\n", + " with open(path, \"w\", encoding=\"utf-8\") as f:\n", + " for example in dataset:\n", + " f.write(json.dumps({\n", + " \"context\": example[\"context\"],\n", + " \"question\": example[\"question\"],\n", + " \"completion\": example[\"answers\"][\"text\"][0],\n", + " }) + \"\\n\")\n", + "\n", + "\n", + "write_jsonl(train_ds, f\"{DATASET_PATH}/training.jsonl\")\n", + "write_jsonl(val_ds, f\"{DATASET_PATH}/validation.jsonl\")\n", + "write_test_jsonl(test_ds, f\"{DATASET_PATH}/testing.jsonl\")\n", + "\n", + "print(f\"Training: {len(train_ds)} rows\")\n", + "print(f\"Validation: {len(val_ds)} rows\")\n", + "print(f\"Test: {len(test_ds)} rows\")\n", + "\n", + "with open(f\"{DATASET_PATH}/training.jsonl\", 'r') as f:\n", + " sample = json.loads(f.readline())\n", + " print(f\"\\nSample prompt: {sample['prompt'][:150]}...\")\n", + " print(f\"Sample completion: {sample['completion']}\")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "DATASET_NAME = \"kd-dataset\"\n", + "\n", + "try:\n", + " client.files.filesets.create(\n", + " workspace=\"default\",\n", + " name=DATASET_NAME,\n", + " description=\"Knowledge distillation training data\"\n", + " )\n", + " print(f\"Created fileset: {DATASET_NAME}\")\n", + "except ConflictError:\n", + " print(f\"Fileset '{DATASET_NAME}' already exists, continuing...\")\n", + "\n", + "client.files.upload(\n", + " local_path=f\"{DATASET_PATH}/\",\n", + " remote_path=\"\",\n", + " fileset=DATASET_NAME,\n", + " workspace=\"default\"\n", + ")\n", + "\n", + "print(\"Uploaded files:\")\n", + "print(json.dumps([f.model_dump() for f in client.files.list(fileset=DATASET_NAME, workspace=\"default\").data], indent=2))" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 3. Secrets Setup\n", + "\n", + "In this tutorial we use two Llama 3.2 Instruct models from Hugging Face:\n", + "- **Teacher:** [meta-llama/Llama-3.2-3B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-3B-Instruct) (3B parameters)\n", + "- **Student:** [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct) (1B parameters)\n", + "\n", + "Both models share the same tokenizer/vocabulary (required for knowledge distillation) and include a chat template for deployment with `/chat/completions`.\n", + "\n", + "**Hugging Face Authentication:**\n", + "- For gated models (Llama, Gemma), you must provide a Hugging Face token via the `token_secret` parameter\n", + "- Get your token from [Hugging Face Settings](https://huggingface.co/settings/tokens) (requires Read access)\n", + "- Accept the model's terms on the Hugging Face model page before using it:\n", + " - [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct)\n", + " - [meta-llama/Llama-3.2-3B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-3B-Instruct)" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "HF_TOKEN = os.getenv(\"HF_TOKEN\")\n", + "\n", + "\n", + "def create_or_get_secret(name: str, value: str | None, label: str):\n", + " if not value:\n", + " raise ValueError(f\"{label} is not set\")\n", + " try:\n", + " secret = client.secrets.create(\n", + " name=name,\n", + " workspace=\"default\",\n", + " value=value,\n", + " )\n", + " print(f\"Created secret: {name}\")\n", + " return secret\n", + " except ConflictError:\n", + " print(f\"Secret '{name}' already exists, continuing...\")\n", + " return client.secrets.retrieve(name=name, workspace=\"default\")\n", + "\n", + "\n", + "hf_secret = create_or_get_secret(\"hf-token\", HF_TOKEN, \"HF_TOKEN\")\n", + "print(\"HF_TOKEN secret:\")\n", + "print(hf_secret.model_dump_json(indent=2))" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 4. Create Model FileSets and Model Entities\n", + "\n", + "Knowledge distillation requires **two** model entities:\n", + "1. **Student model** — the smaller model that will be trained ([meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct))\n", + "2. **Teacher model** — the larger model that provides soft targets ([meta-llama/Llama-3.2-3B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-3B-Instruct))\n", + "\n", + "Using the Instruct variants ensures the output model includes a chat template, which is required for the `/chat/completions` inference endpoint." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from nemo_platform.types.files import HuggingfaceStorageConfigParam\n", + "\n", + "SPEC_TIMEOUT_SECONDS = 120\n", + "\n", + "\n", + "def create_model(hf_repo: str, model_name: str, description: str):\n", + " \"\"\"Create a fileset + model entity and wait for ModelSpec.\"\"\"\n", + " try:\n", + " client.files.filesets.create(\n", + " workspace=\"default\",\n", + " name=model_name,\n", + " description=description,\n", + " storage=HuggingfaceStorageConfigParam(\n", + " type=\"huggingface\",\n", + " repo_id=hf_repo,\n", + " repo_type=\"model\",\n", + " token_secret=hf_secret.name\n", + " )\n", + " )\n", + " print(f\"Created fileset: {model_name}\")\n", + " except ConflictError:\n", + " print(f\"Fileset '{model_name}' already exists.\")\n", + "\n", + " try:\n", + " model = client.models.create(\n", + " workspace=\"default\",\n", + " name=model_name,\n", + " fileset=f\"default/{model_name}\",\n", + " )\n", + " print(f\"Created Model Entity: {model_name}\")\n", + " except ConflictError:\n", + " print(f\"Model '{model_name}' already exists. Updating fileset.\")\n", + " model = client.models.update(\n", + " workspace=\"default\",\n", + " name=model_name,\n", + " fileset=f\"default/{model_name}\",\n", + " )\n", + "\n", + " print(f\"Waiting for ModelSpec on {model_name}...\")\n", + " spec_start = time.time()\n", + " while not model.spec:\n", + " if time.time() - spec_start > SPEC_TIMEOUT_SECONDS:\n", + " raise TimeoutError(f\"ModelSpec not populated within {SPEC_TIMEOUT_SECONDS}s\")\n", + " time.sleep(2)\n", + " model = client.models.retrieve(workspace=\"default\", name=model_name)\n", + " print(f\"ModelSpec populated: {model.spec}\")\n", + " return model\n", + "\n", + "\n", + "student_model = create_model(\n", + " hf_repo=\"meta-llama/Llama-3.2-1B-Instruct\",\n", + " model_name=\"llama-3-2-1b-student\",\n", + " description=\"Llama 3.2 1B Instruct student model\",\n", + ")\n", + "\n", + "print()\n", + "\n", + "teacher_model = create_model(\n", + " hf_repo=\"meta-llama/Llama-3.2-3B-Instruct\",\n", + " model_name=\"llama-3-2-3b-teacher\",\n", + " description=\"Llama 3.2 3B Instruct teacher model\",\n", + ")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "---\n", + "\n", + "## Phase 1: Fine-Tune the Teacher\n", + "\n", + "### 5. Train Teacher with Full SFT\n", + "\n", + "For best distillation results, fine-tune the teacher on the **same dataset** that will be used for distillation. This ensures the teacher has learned the task-specific knowledge that the student will inherit.\n", + "\n", + "We train the 3B Instruct model with Full SFT on the SQuAD dataset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from nemo_automodel_plugin.schema import AutomodelJobInput\n", + "\n", + "job_suffix = uuid.uuid4().hex[:4]\n", + "\n", + "TEACHER_JOB_NAME = f\"teacher-sft-job-{job_suffix}\"\n", + "TEACHER_OUTPUT_NAME = f\"teacher-model-{job_suffix}\"\n", + "\n", + "teacher_spec = AutomodelJobInput(\n", + " model=f\"default/{teacher_model.name}\",\n", + " dataset={\"training\": f\"default/{DATASET_NAME}\"},\n", + " training={\n", + " \"training_type\": \"sft\",\n", + " \"finetuning_type\": \"all_weights\",\n", + " \"max_seq_length\": 2048,\n", + " },\n", + " schedule={\"epochs\": 1},\n", + " batch={\"global_batch_size\": 64, \"micro_batch_size\": 1},\n", + " optimizer={\"learning_rate\": 5e-5},\n", + " parallelism={\"num_gpus_per_node\": 1},\n", + " output={\"name\": TEACHER_OUTPUT_NAME},\n", + ")\n", + "\n", + "teacher_job = client.customization.automodel.jobs.create(\n", + " spec=teacher_spec, workspace=\"default\", name=TEACHER_JOB_NAME\n", + ")\n", + "\n", + "TRAINED_TEACHER_NAME = TEACHER_OUTPUT_NAME\n", + "print(f\"Teacher training job: {teacher_job.job.name}\")\n", + "print(f\"Output teacher model: {TRAINED_TEACHER_NAME}\")\n" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from IPython.display import clear_output\n", + "\n", + "\n", + "def wait_for_job(job_name: str):\n", + " \"\"\"Poll job status until completion.\"\"\"\n", + " while True:\n", + " status = client.jobs.get_status(name=job_name, workspace=\"default\")\n", + " clear_output(wait=True)\n", + " print(f\"Job: {job_name}\")\n", + " print(f\"Status: {status.status}\")\n", + "\n", + " for job_step in status.steps or []:\n", + " if job_step.name == \"training\":\n", + " for task in job_step.tasks or []:\n", + " details = task.status_details or {}\n", + " step = details.get(\"step\")\n", + " max_steps = details.get(\"max_steps\")\n", + " if step is not None and max_steps is not None:\n", + " print(f\"Progress: Step {step}/{max_steps} ({step / max_steps * 100:.1f}%)\")\n", + " phase = details.get(\"phase\")\n", + " if phase:\n", + " print(f\"Phase: {phase}\")\n", + " break\n", + " break\n", + "\n", + " if status.status in (\"completed\", \"failed\", \"cancelled\", \"error\"):\n", + " print(f\"\\nJob finished: {status.status}\")\n", + " return status\n", + "\n", + " time.sleep(10)\n", + "\n", + "\n", + "teacher_status = wait_for_job(TEACHER_JOB_NAME)\n", + "assert teacher_status.status == \"completed\"" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "---\n", + "\n", + "## Phase 2: Establish Baseline (Base Student)\n", + "\n", + "### 6. Deploy the Base Student Model\n", + "\n", + "Before distillation, deploy the base student model (1B Instruct, without any fine-tuning) to establish a baseline ROUGE score. After distillation, we compare the distilled student against this baseline to measure improvement." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "baseline_suffix = uuid.uuid4().hex[:4]\n", + "BASELINE_DEPLOYMENT_CONFIG = f\"baseline-student-cfg-{baseline_suffix}\"\n", + "BASELINE_DEPLOYMENT_NAME = f\"baseline-student-{baseline_suffix}\"\n", + "\n", + "baseline_deployment_config = client.inference.deployment_configs.create(\n", + " workspace=\"default\",\n", + " name=BASELINE_DEPLOYMENT_CONFIG,\n", + " engine=\"vllm\",\n", + " model_spec={\n", + " \"model_namespace\": \"default\",\n", + " \"model_name\": student_model.name,\n", + " },\n", + " executor_config={\n", + " \"gpu\": 1,\n", + " \"image_name\": \"vllm/vllm-openai\",\n", + " \"image_tag\": \"v0.22.1\",\n", + " },\n", + ")\n", + "\n", + "baseline_deployment = client.inference.deployments.create(\n", + " workspace=\"default\",\n", + " name=BASELINE_DEPLOYMENT_NAME,\n", + " config=baseline_deployment_config.name\n", + ")\n", + "\n", + "print(f\"Baseline student deployment: {baseline_deployment.name}\")\n" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "def wait_for_deployment(deployment_name: str, timeout_minutes: int = 30):\n", + " \"\"\"Poll deployment until ready.\"\"\"\n", + " start = time.time()\n", + " timeout = timeout_minutes * 60\n", + " while True:\n", + " dep = client.inference.deployments.retrieve(name=deployment_name, workspace=\"default\")\n", + " elapsed = time.time() - start\n", + " clear_output(wait=True)\n", + " print(f\"Deployment: {deployment_name}\")\n", + " print(f\"Status: {dep.status}\")\n", + " print(f\"Elapsed: {int(elapsed // 60)}m {int(elapsed % 60)}s\")\n", + "\n", + " if dep.status == \"READY\":\n", + " print(\"\\nDeployment is ready!\")\n", + " if not client.models.wait_for_gateway(deployment_name, workspace=\"default\", timeout=60):\n", + " raise RuntimeError(\"Inference gateway did not become ready\")\n", + " return dep\n", + " if dep.status in (\"FAILED\", \"ERROR\", \"TERMINATED\", \"LOST\"):\n", + " raise RuntimeError(f\"Deployment failed with status: {dep.status}\")\n", + " if elapsed > timeout:\n", + " raise TimeoutError(f\"Deployment timeout after {timeout_minutes} minutes\")\n", + " time.sleep(15)\n", + "\n", + "\n", + "dep_status = wait_for_deployment(BASELINE_DEPLOYMENT_NAME)\n", + "assert dep_status.status == \"READY\"" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 7. Generate Baseline Predictions on Test Set" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "with open(f\"{DATASET_PATH}/testing.jsonl\", \"r\", encoding=\"utf-8\") as f:\n", + " test_data = [json.loads(line) for line in f]\n", + "\n", + "contexts = [row[\"context\"] for row in test_data]\n", + "questions = [row[\"question\"] for row in test_data]\n", + "reference_completions = [row[\"completion\"] for row in test_data]\n", + "\n", + "print(f\"Test samples: {len(contexts)}\")\n", + "print(f\"Sample context: {contexts[0]}\")\n", + "print(f\"Sample question: {questions[0]}\")\n", + "print(f\"Sample reference: {reference_completions[0]}\")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "def generate_completions(\n", + " deployment_name: str,\n", + " output_model_name: str,\n", + " contexts: list[str],\n", + " questions: list[str],\n", + ") -> list[str]:\n", + " \"\"\"Generate completions for a list of context/question pairs using a deployed model.\"\"\"\n", + " completions = []\n", + " for context, question in zip(contexts, questions):\n", + " messages = [\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": f\"Based on the following context, answer the question.\\n\\nContext: {context}\\n\\nQuestion: {question}\",\n", + " }\n", + " ]\n", + " response = client.inference.gateway.provider.post(\n", + " \"v1/chat/completions\",\n", + " name=deployment_name,\n", + " workspace=\"default\",\n", + " body={\n", + " \"model\": f\"default/{output_model_name}\",\n", + " \"messages\": messages,\n", + " \"temperature\": 0,\n", + " \"max_tokens\": 128,\n", + " }\n", + " )\n", + " completions.append(response[\"choices\"][0][\"message\"][\"content\"])\n", + " return completions\n", + "\n", + "\n", + "print(\"Generating baseline (base student) predictions...\")\n", + "baseline_completions = generate_completions(BASELINE_DEPLOYMENT_NAME, student_model.name, contexts, questions)\n", + "print(f\"Generated {len(baseline_completions)} baseline predictions\")\n", + "print(f\"\\nSample baseline output: {baseline_completions[0]}\")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 8. Delete Baseline Deployment\n", + "\n", + "Delete the baseline student deployment to free GPU resources for the distillation training job and subsequent distilled model deployment." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "client.inference.deployments.delete(name=BASELINE_DEPLOYMENT_NAME, workspace=\"default\")\n", + "print(f\"Deleted baseline deployment: {BASELINE_DEPLOYMENT_NAME}\")\n", + "\n", + "if not client.models.wait_for_status(\n", + " deployment_name=BASELINE_DEPLOYMENT_NAME,\n", + " desired_status=\"DELETED\",\n", + " workspace=\"default\",\n", + " timeout=600,\n", + "):\n", + " raise TimeoutError(\n", + " f\"Deployment {BASELINE_DEPLOYMENT_NAME} was not deleted within timeout\"\n", + " )\n", + "\n", + "client.inference.deployment_configs.delete(name=BASELINE_DEPLOYMENT_CONFIG, workspace=\"default\")\n", + "print(f\"Deleted baseline deployment config: {BASELINE_DEPLOYMENT_CONFIG}\")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "---\n", + "\n", + "## Phase 3: Distill into Student\n", + "\n", + "### 9. Create Knowledge Distillation Job\n", + "\n", + "Now create a distillation job that trains the 1B student using the **fine-tuned** 3B teacher's output distribution. The `model` field specifies the student, and `teacher_model` references the trained teacher model entity from Phase 1." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "**GPU Requirements:**\n", + "\n", + "KD requires loading both student and teacher models, so plan GPU memory accordingly:\n", + "- 1B student + 3B teacher: 1 GPU (24GB+ VRAM each)\n", + "- 3B student + 8B teacher: 4 GPUs\n", + "- 8B student + 70B teacher: 8+ GPUs\n", + "\n", + "Use `teacher_precision=\"bf16\"` (default) to reduce teacher memory footprint." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from nemo_automodel_plugin.schema import AutomodelJobInput\n", + "\n", + "KD_JOB_NAME = f\"my-kd-job-{job_suffix}\"\n", + "KD_OUTPUT_NAME = f\"kd-student-{job_suffix}\"\n", + "\n", + "kd_spec = AutomodelJobInput(\n", + " model=f\"default/{student_model.name}\",\n", + " dataset={\"training\": f\"default/{DATASET_NAME}\"},\n", + " training={\n", + " \"training_type\": \"distillation\",\n", + " \"finetuning_type\": \"all_weights\",\n", + " \"teacher_model\": f\"default/{TRAINED_TEACHER_NAME}\",\n", + " \"teacher_precision\": \"bf16\",\n", + " \"distillation_ratio\": 0.5,\n", + " \"distillation_temperature\": 2.0,\n", + " \"max_seq_length\": 2048,\n", + " },\n", + " schedule={\"epochs\": 1},\n", + " batch={\"global_batch_size\": 64, \"micro_batch_size\": 1},\n", + " optimizer={\"learning_rate\": 5e-5},\n", + " parallelism={\"num_gpus_per_node\": 1},\n", + " output={\"name\": KD_OUTPUT_NAME},\n", + ")\n", + "\n", + "kd_job = client.customization.automodel.jobs.create(\n", + " spec=kd_spec, workspace=\"default\", name=KD_JOB_NAME\n", + ")\n", + "\n", + "DISTILLED_STUDENT_NAME = KD_OUTPUT_NAME\n", + "print(f\"Distillation job: {kd_job.job.name}\")\n", + "print(f\"Output student model: {DISTILLED_STUDENT_NAME}\")\n" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 10. Track Distillation Progress" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "kd_status = wait_for_job(KD_JOB_NAME)\n", + "assert kd_status.status == \"completed\"" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "---\n", + "\n", + "## Phase 4: Evaluate the Distilled Student Model\n", + "\n", + "### 11. Deploy the Distilled Student Model\n", + "\n", + "The output model has the same architecture as the 1B student—only its weights have been updated via distillation. It requires just 1 GPU to deploy." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "deploy_suffix_2 = uuid.uuid4().hex[:4]\n", + "STUDENT_DEPLOYMENT_CONFIG = f\"kd-student-deploy-cfg-{deploy_suffix_2}\"\n", + "STUDENT_DEPLOYMENT_NAME = f\"kd-student-deploy-{deploy_suffix_2}\"\n", + "\n", + "student_deployment_config = client.inference.deployment_configs.create(\n", + " workspace=\"default\",\n", + " name=STUDENT_DEPLOYMENT_CONFIG,\n", + " engine=\"vllm\",\n", + " model_spec={\n", + " \"model_namespace\": \"default\",\n", + " \"model_name\": DISTILLED_STUDENT_NAME,\n", + " },\n", + " executor_config={\n", + " \"gpu\": 1,\n", + " \"image_name\": \"vllm/vllm-openai\",\n", + " \"image_tag\": \"v0.22.1\",\n", + " },\n", + ")\n", + "\n", + "student_deployment = client.inference.deployments.create(\n", + " workspace=\"default\",\n", + " name=STUDENT_DEPLOYMENT_NAME,\n", + " config=student_deployment_config.name\n", + ")\n", + "\n", + "print(f\"Student deployment: {student_deployment.name}\")\n" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "wait_for_deployment(STUDENT_DEPLOYMENT_NAME)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 12. Generate Student Predictions on Test Set" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\"Generating distilled student predictions...\")\n", + "student_completions = generate_completions(STUDENT_DEPLOYMENT_NAME, DISTILLED_STUDENT_NAME, contexts, questions)\n", + "print(f\"Generated {len(student_completions)} student predictions\")\n", + "print(f\"\\nSample student output: {student_completions[0]}\")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 13. Compute ROUGE Scores\n", + "\n", + "Compare the base student (before distillation) and the distilled student against the ground-truth reference completions using ROUGE metrics." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "import evaluate\n", + "\n", + "rouge = evaluate.load(\"rouge\")\n", + "\n", + "baseline_scores = rouge.compute(predictions=baseline_completions, references=reference_completions)\n", + "student_scores = rouge.compute(predictions=student_completions, references=reference_completions)\n", + "\n", + "metrics = list(baseline_scores.keys())\n", + "header = f\"{'Model':<35} \" + \" \".join(f\"{m:>10}\" for m in metrics)\n", + "separator = \"-\" * len(header)\n", + "\n", + "print(\"=\" * 60)\n", + "print(\"ROUGE SCORE COMPARISON\")\n", + "print(\"=\" * 60)\n", + "print(header)\n", + "print(separator)\n", + "print(f\"{'Base Student (1B, no training)':<35} \" + \" \".join(f\"{baseline_scores[m]:>10.4f}\" for m in metrics))\n", + "print(f\"{'Distilled Student (1B, KD)':<35} \" + \" \".join(f\"{student_scores[m]:>10.4f}\" for m in metrics))" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "print(\"Sample predictions (first 3):\\n\")\n", + "for i in range(min(3, len(contexts))):\n", + " print(f\"--- Sample {i + 1} ---\")\n", + " print(f\"Question: {questions[i]}\")\n", + " print(f\"Reference: {reference_completions[i]}\")\n", + " print(f\"Baseline: {baseline_completions[i][:200]}\")\n", + " print(f\"Distilled: {student_completions[i][:200]}\")\n", + " print()" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "**Interpreting ROUGE Scores:**\n", + "\n", + "| Metric | Measures |\n", + "|--------|----------|\n", + "| **ROUGE-1** | Unigram overlap between prediction and reference |\n", + "| **ROUGE-2** | Bigram overlap (captures phrase-level similarity) |\n", + "| **ROUGE-L** | Longest common subsequence (captures sentence structure) |\n", + "| **ROUGE-Lsum** | ROUGE-L computed over full summaries |\n", + "\n", + "**What to expect:**\n", + "- The base student (1B, no training) provides a lower bound since it has not seen the task data\n", + "- The distilled student (1B, KD) should significantly outperform the base student, demonstrating the knowledge transferred from the 3B teacher\n", + "- If the distilled student scores are not much higher than the baseline, try increasing `distillation_temperature`, adjusting `distillation_ratio`, or training for more epochs\n", + "\n", + "---\n", + "\n", + "## Hyperparameters\n", + "\n", + "For detailed information on all available hyperparameters, recommended values, and tuning guidance, refer to the [Hyperparameter Reference](../manage-customization-jobs/hyperparameters.md).\n", + "\n", + "---\n", + "\n", + "## Troubleshooting\n", + "\n", + "**Job fails during model download:**\n", + "- Verify authentication secrets are configured (refer to [Managing Secrets](../../get-started/concepts/manage-secrets.md))\n", + "- For gated Hugging Face models (Llama, Gemma), accept the license on the model page\n", + "- Check both `model` (student) and `teacher_model` URNs are correct\n", + "- Ensure both model entities exist: `client.models.retrieve(name=..., workspace=\"default\")`\n", + "\n", + "**Job fails with OOM (Out of Memory) error:**\n", + "\n", + "KD loads both models, so OOM is more likely than with SFT:\n", + "1. **First try:** Use `teacher_precision=\"bf16\"` to reduce teacher memory\n", + "2. **Still OOM:** Reduce `micro_batch_size` to 1\n", + "3. **Still OOM:** Reduce `global_batch_size` and `max_seq_length`\n", + "4. **Last resort:** Increase `num_gpus_per_node`\n", + "\n", + "**No chat template / `/chat/completions` fails:**\n", + "- Use Instruct model variants (e.g., `Llama-3.2-1B-Instruct`) instead of base models (`Llama-3.2-1B`). Base models do not include a chat template in their tokenizer, so the output model will also lack one.\n", + "\n", + "**Distilled model quality is poor:**\n", + "- Increase `distillation_temperature` (try 2.0–5.0) to transfer more nuanced knowledge\n", + "- Adjust `distillation_ratio`—if dataset labels are high-quality, lower the ratio; if the teacher is strong, raise it\n", + "- Increase `epochs` or `max_steps` for more training\n", + "- Verify teacher and student share the same vocabulary\n", + "\n", + "**Vocabulary mismatch error:**\n", + "- Teacher and student must use the same tokenizer. Use models from the same family (e.g., Llama 3.2 1B Instruct + Llama 3.2 3B Instruct)\n", + "\n", + "**Deployment fails:**\n", + "- Verify output model exists: `client.models.retrieve(name=DISTILLED_STUDENT_NAME, workspace=\"default\")`\n", + "- Check deployment logs: `client.inference.deployments.get_logs(name=deployment.name, workspace=\"default\")`\n", + "- The distilled model has the same size as the student, so GPU requirements match the student model\n", + "\n", + "\n", + "## Next Steps\n", + "\n", + "- [Monitor training metrics](fine-tune-metrics) in detail\n", + "- [Evaluate your fine-tuned model](../../evaluator/index) using the Evaluator service\n", + "- Learn about [LoRA customization](./lora-customization-job) for resource-efficient fine-tuning\n", + "- Learn about [Full SFT](./sft-customization-job) for direct supervised fine-tuning" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": ".venv", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.14" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} \ No newline at end of file diff --git a/docs/customizer/tutorials/distillation-customization-job.mdx b/docs/customizer/tutorials/distillation-customization-job.mdx index 8da102ae5f..3c78c4cae3 100644 --- a/docs/customizer/tutorials/distillation-customization-job.mdx +++ b/docs/customizer/tutorials/distillation-customization-job.mdx @@ -5,8 +5,6 @@ description: "" [Run in Google Colab](https://colab.research.google.com/github/NVIDIA-NeMo/nemo-platform/blob/main/docs/customizer/tutorials/distillation-customization-job.ipynb) -# Knowledge Distillation Customization - Learn how to train a smaller student model to mimic a larger teacher model using knowledge distillation (KD). ## About @@ -67,6 +65,8 @@ Before starting this tutorial, ensure you have: pip install evaluate rouge_score datasets ``` +4. **At least one GPU with CUDA 13+** + ## Quick Start ### 1. Initialize SDK @@ -189,16 +189,16 @@ print(json.dumps([f.model_dump() for f in client.files.list(fileset=DATASET_NAME ### 3. Secrets Setup -In this tutorial we use two Llama 3.2 Instruct models from HuggingFace: +In this tutorial we use two Llama 3.2 Instruct models from Hugging Face: - **Teacher:** [meta-llama/Llama-3.2-3B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-3B-Instruct) (3B parameters) - **Student:** [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct) (1B parameters) Both models share the same tokenizer/vocabulary (required for knowledge distillation) and include a chat template for deployment with `/chat/completions`. -**HuggingFace Authentication:** -- For gated models (Llama, Gemma), you must provide a HuggingFace token via the `token_secret` parameter -- Get your token from [HuggingFace Settings](https://huggingface.co/settings/tokens) (requires Read access) -- Accept the model's terms on the HuggingFace model page before using it: +**Hugging Face Authentication:** +- For gated models (Llama, Gemma), you must provide a Hugging Face token via the `token_secret` parameter +- Get your token from [Hugging Face Settings](https://huggingface.co/settings/tokens) (requires Read access) +- Accept the model's terms on the Hugging Face model page before using it: - [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct) - [meta-llama/Llama-3.2-3B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-3B-Instruct) @@ -433,13 +433,13 @@ def wait_for_deployment(deployment_name: str, timeout_minutes: int = 30): if dep.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") return dep if dep.status in ("FAILED", "ERROR", "TERMINATED", "LOST"): - print(f"\nDeployment failed: {dep.status}") - return dep + raise RuntimeError(f"Deployment failed with status: {dep.status}") if elapsed > timeout: - print(f"\nTimeout ({timeout_minutes}m). Check status manually.") - return dep + raise TimeoutError(f"Deployment timeout after {timeout_minutes} minutes") time.sleep(15) @@ -694,7 +694,7 @@ For detailed information on all available hyperparameters, recommended values, a **Job fails during model download:** - Verify authentication secrets are configured (refer to [Managing Secrets](/documentation/get-started/core-concepts/manage-secrets)) -- For gated HuggingFace models (Llama, Gemma), accept the license on the model page +- For gated Hugging Face models (Llama, Gemma), accept the license on the model page - Check both `model` (student) and `teacher_model` URNs are correct - Ensure both model entities exist: `client.models.retrieve(name=..., workspace="default")` diff --git a/docs/customizer/tutorials/dpo-customization-job.ipynb b/docs/customizer/tutorials/dpo-customization-job.ipynb index d9a56d4f64..6cd9dd4110 100644 --- a/docs/customizer/tutorials/dpo-customization-job.ipynb +++ b/docs/customizer/tutorials/dpo-customization-job.ipynb @@ -1,515 +1,524 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "\n", - "\n", - "\n", - "# DPO Model Customization Job\n", - "\n", - "Learn how to use the NeMo Platform to align a model with **DPO** (Direct Preference Optimization) on a preference dataset. For each prompt, DPO trains on a *chosen* (preferred) and a *rejected* response so the model prefers the chosen style — no separate reward model required.\n", - "\n", - "This tutorial uses the `rl` customization backend (powered by [NVIDIA NeMo-RL](https://github.com/NVIDIA-NeMo/RL)), which runs DPO on a **Ray** cluster. Unlike the [SFT](./sft-customization-job) and [LoRA](./lora-customization-job) tutorials (Docker GPU jobs), `rl` requires a **Kubernetes-backed** NeMo Platform. DPO here is **full-weight** (no LoRA/adapter); the output is a full model entity.\n", - "\n", - "**Time to complete:** approximately 45-60 minutes. Job duration increases with model and dataset size." - ] + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "\n", + "\n", + "\n", + "# DPO Model Customization Job\n", + "\n", + "Learn how to use the NeMo Platform to align a model with **DPO** (Direct Preference Optimization) on a preference dataset. For each prompt, DPO trains on a *chosen* (preferred) and a *rejected* response so the model prefers the chosen style — no separate reward model required.\n", + "\n", + "This tutorial uses the `rl` customization backend (powered by [NVIDIA NeMo-RL](https://github.com/NVIDIA-NeMo/RL)), which runs DPO on a **Ray** cluster. Unlike the [SFT](/documentation/customizer-reference/tutorials/sft-customization-job) and [LoRA](/documentation/customizer-reference/tutorials/lora-customization-job) tutorials (Docker GPU jobs), `rl` requires a **Kubernetes-backed** NeMo Platform. DPO here is **full-weight** (no LoRA/adapter); the output is a full model entity.\n", + "\n", + "**Time to complete:** approximately 45-60 minutes. Job duration increases with model and dataset size." + ], + "id": "40916acf" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Prerequisites\n", + "\n", + "Before starting this tutorial, ensure you have:\n", + "\n", + "1. **Completed the [Quickstart](/documentation/get-started)** to install the NeMo Platform and Python SDK.\n", + "2. **Installed the Python SDK** (PyPI wrapper: `pip install \"nemo-platform[all]\"`; source checkout: run `make bootstrap` from the repository root).\n", + "3. **Installed the `datasets` package**: `pip install datasets`.\n", + "4. **A platform configured with `platform.runtime: kubernetes`.** The `rl` (DPO) backend provisions a Ray cluster and has **no local Docker fallback** — `submit` fails fast on a Docker-runtime platform. Multi-node jobs (`parallelism.num_nodes > 1`) additionally require the platform-side `NMP_RL_MULTINODE_SHARED_STORAGE_PATH`.\n", + "5. **A Hugging Face token** with access to the gated base model (this tutorial uses `meta-llama/Llama-3.2-1B-Instruct`). Export it as `HF_TOKEN`.\n", + "6. **At least one GPU with CUDA 13+** and a GPU execution profile (`nemo jobs list-execution-profiles`)." + ], + "id": "15435f3c" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Quick Start\n", + "\n", + "### 1. Initialize the SDK\n", + "\n", + "The SDK needs your NeMo Platform server URL. By default `http://localhost:8080` is used; set `NMP_BASE_URL` to override:\n", + "\n", + "```sh\n", + "export NMP_BASE_URL=\n", + "```" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "import json\n", + "import os\n", + "import time\n", + "import uuid\n", + "from pathlib import Path\n", + "from nemo_platform import NeMoPlatform, ConflictError\n", + "from nemo_platform.types.secrets import PlatformSecretResponse\n", + "from nemo_platform.types.files import HuggingfaceStorageConfigParam\n", + "from nemo_rl_plugin.schema import RlJobInput\n", + "\n", + "\n", + "def max_wait_time_checker(seconds: int, label: str = \"\"):\n", + " \"\"\"Return a check() that raises TimeoutError once `seconds` have elapsed.\"\"\"\n", + " start = time.time()\n", + "\n", + " def check():\n", + " if time.time() - start > seconds:\n", + " raise TimeoutError(f\"{label} took longer than {seconds} seconds\")\n", + "\n", + " return check\n", + "\n", + "\n", + "NMP_BASE_URL = os.environ.get(\"NMP_BASE_URL\", \"http://localhost:8080\")\n", + "sdk = NeMoPlatform(base_url=NMP_BASE_URL, workspace=\"default\")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 2. Prepare the Preference Dataset\n", + "\n", + "DPO trains on **preference data**. The `rl` backend takes a **single** dataset fileset that holds both `training.jsonl` and `validation.jsonl`, and auto-detects the row schema from the first line. Three preference formats are supported (see the platform's `BinaryPreferenceDatasetItemSchema` / `HelpSteer3DatasetItemSchema` / `Tulu3PreferenceDatasetItemSchema`):" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "#### Binary Preference Format\n", + "\n", + "Simple `prompt` / `chosen` / `rejected` (the `prompt` may be a string or a list of chat messages):\n", + "\n", + "```json\n", + "{\"prompt\": \"What is the capital of France?\", \"chosen\": \"The capital of France is Paris.\", \"rejected\": \"I'm not sure.\"}\n", + "```" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "#### HelpSteer3 Format (used here)\n", + "\n", + "A conversation `context` (string or chat messages), two candidate `response1` / `response2`, and a signed `overall_preference` in -3..3 — **negative** means response 1 is preferred, **positive** means response 2, **0** is a tie. This is the **raw** schema of `nvidia/HelpSteer3`, so no conversion is needed:\n", + "\n", + "```json\n", + "{\"context\": [{\"role\": \"user\", \"content\": \"Explain how to use git rebase\"}], \"response1\": \"...\", \"response2\": \"...\", \"overall_preference\": -2}\n", + "```" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "#### Tulu3 Preference Format\n", + "\n", + "Full chat conversations for both the chosen and rejected branches (each a list of messages ending with the assistant turn):\n", + "\n", + "```json\n", + "{\"chosen\": [{\"role\": \"user\", \"content\": \"...\"}, {\"role\": \"assistant\", \"content\": \"preferred\"}], \"rejected\": [{\"role\": \"user\", \"content\": \"...\"}, {\"role\": \"assistant\", \"content\": \"dispreferred\"}]}\n", + "```" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "#### Download nvidia/HelpSteer3\n", + "\n", + "We use [nvidia/HelpSteer3](https://huggingface.co/datasets/nvidia/HelpSteer3) (the `preference` subset), NVIDIA's open preference dataset. It ships native `train` and `validation` splits and matches the HelpSteer3 schema above, so we upload the rows **as-is** — the platform's `HelpSteer3Dataset` loader handles the `overall_preference` semantics (including ties) at training time." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from datasets import load_dataset, Dataset\n", + "\n", + "print(\"Loading dataset nvidia/HelpSteer3 (preference subset)\")\n", + "ds = load_dataset(\"nvidia/HelpSteer3\", \"preference\")\n", + "\n", + "# Small subsets keep the tutorial fast; larger sets train better but take longer.\n", + "training_size = 3000\n", + "validation_size = 300\n", + "DATASET_NAME = \"dpo-dataset\"\n", + "DATASET_PATH = Path(\"dpo-dataset\").absolute()\n", + "os.makedirs(DATASET_PATH, exist_ok=True)\n", + "\n", + "train_dataset = ds[\"train\"]\n", + "validation_dataset = ds[\"validation\"]\n", + "assert isinstance(train_dataset, Dataset) and isinstance(validation_dataset, Dataset)\n", + "\n", + "# Save raw HelpSteer3 rows directly — no conversion. The platform detects the\n", + "# HelpSteer3 schema from the row keys (context / response1 / response2 / overall_preference).\n", + "train_dataset.select(range(training_size)).to_json(f\"{DATASET_PATH}/training.jsonl\")\n", + "validation_dataset.select(range(validation_size)).to_json(f\"{DATASET_PATH}/validation.jsonl\")\n", + "\n", + "print(f\"Saved training.jsonl ({training_size} rows) and validation.jsonl ({validation_size} rows)\")\n", + "with open(f\"{DATASET_PATH}/training.jsonl\") as f:\n", + " sample = json.loads(f.readline())\n", + "print(\"Sample keys:\", sorted(sample.keys()))\n", + "print(\"overall_preference:\", sample[\"overall_preference\"])" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 3. Create FileSet and Upload Preference Data\n", + "\n", + "Upload both JSONL files to a single FileSet so the DPO job can read them." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "try:\n", + " sdk.files.filesets.create(workspace=\"default\", name=DATASET_NAME, description=\"DPO preference data\")\n", + " print(f\"Created fileset: {DATASET_NAME}\")\n", + "except ConflictError:\n", + " print(f\"Fileset '{DATASET_NAME}' already exists, continuing...\")\n", + "\n", + "sdk.files.upload(local_path=DATASET_PATH, remote_path=\"\", fileset=DATASET_NAME, workspace=\"default\")\n", + "\n", + "print(\"Preference data:\")\n", + "print(json.dumps([f.model_dump() for f in sdk.files.list(fileset=DATASET_NAME, workspace=\"default\").data], indent=2, default=str))" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 4. Secrets Setup\n", + "\n", + "The base model (`meta-llama/Llama-3.2-1B-Instruct`) is gated, so store your Hugging Face token as a platform secret named `hf-token` and reference it on the model fileset." + ], + "id": "7f8bde21" + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "HF_TOKEN = os.getenv(\"HF_TOKEN\")\n", + "\n", + "def create_or_get_secret(name: str, value: str | None, label: str) -> PlatformSecretResponse | None:\n", + " if not value:\n", + " print(f\"{label} is not set - skipping secret (gated model downloads will fail without it)\")\n", + " return None\n", + " try:\n", + " secret = sdk.secrets.create(name=name, workspace=\"default\", value=value)\n", + " print(f\"Created secret: {name}\")\n", + " return secret\n", + " except ConflictError:\n", + " print(f\"Secret '{name}' already exists, continuing...\")\n", + " return sdk.secrets.retrieve(name=name, workspace=\"default\")\n", + "\n", + "\n", + "hf_secret = create_or_get_secret(\"hf-token\", HF_TOKEN, \"HF_TOKEN\")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 5. Create Base Model FileSet and Model Entity\n", + "\n", + "DPO starts from an instruction-tuned base model. The model entity's spec is inferred asynchronously after creation." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "HF_REPO_ID = \"meta-llama/Llama-3.2-1B-Instruct\"\n", + "MODEL_NAME = \"llama-3-2-1b-instruct\"\n", + "\n", + "storage = HuggingfaceStorageConfigParam(type=\"huggingface\", repo_id=HF_REPO_ID, repo_type=\"model\")\n", + "if hf_secret:\n", + " storage[\"token_secret\"] = hf_secret.name\n", + "\n", + "try:\n", + " base_model_fs = sdk.files.filesets.create(\n", + " workspace=\"default\", name=MODEL_NAME, description=\"Llama 3.2 1B Instruct base model\", storage=storage\n", + " )\n", + " print(f\"Created base model fileset: {MODEL_NAME}\")\n", + "except ConflictError:\n", + " base_model_fs = sdk.files.filesets.retrieve(workspace=\"default\", name=MODEL_NAME)\n", + " print(\"Base model fileset already exists.\")\n", + "\n", + "try:\n", + " base_model = sdk.models.create(workspace=\"default\", name=MODEL_NAME, fileset=f\"default/{MODEL_NAME}\")\n", + "except ConflictError:\n", + " base_model = sdk.models.retrieve(workspace=\"default\", name=MODEL_NAME)\n", + "\n", + "print(f\"Base model fileset: fileset://default/{base_model.name}\")\n", + "\n", + "# Wait for the ModelSpec to be inferred from the checkpoint.\n", + "check = max_wait_time_checker(600, \"Model spec\")\n", + "while not base_model.spec:\n", + " check()\n", + " time.sleep(10)\n", + " base_model = sdk.models.retrieve(workspace=\"default\", name=MODEL_NAME)\n", + "print(\"Model spec ready\")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 6. Create the DPO Customization Job\n", + "\n", + "Submit a DPO job to the `rl` backend with `RlJobInput`. Note the DPO-specific shape:\n", + "\n", + "- `model` is a string ref to the model entity; `dataset` is a **single** string ref to the preference fileset (holding both files).\n", + "- The training method is `{\"type\": \"dpo\", ...}` — full-weight, no `finetuning_type`/LoRA.\n", + "- `ref_policy_kl_penalty` is **β** (DPO paper): how strongly the policy stays tied to the reference model.\n", + "- `rl` auto-generates the job id (`rl-`); read it back from the response.\n", + "\n", + "Other configurable knobs: `optimizer_type`, `adam_eps`, `activation_checkpointing`, `keep_top_k`, `val_at_end`, `preference_loss_weight`, `sft_loss_weight`. Run `nemo customization rl explain` for the live schema." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "job_suffix = uuid.uuid4().hex[:8]\n", + "OUTPUT_NAME = f\"llama-3-2-1b-dpo-{job_suffix}\"\n", + "\n", + "spec = RlJobInput(\n", + " model=f\"default/{base_model.name}\",\n", + " dataset=f\"default/{DATASET_NAME}\",\n", + " training={\n", + " \"type\": \"dpo\",\n", + " \"epochs\": 1,\n", + " \"batch_size\": 16,\n", + " \"micro_batch_size\": 1,\n", + " \"learning_rate\": 5e-6,\n", + " \"max_seq_length\": 4096,\n", + " \"ref_policy_kl_penalty\": 0.1,\n", + " \"parallelism\": {\n", + " \"num_nodes\": 1,\n", + " \"num_gpus_per_node\": 1,\n", + " \"tensor_parallel_size\": 1,\n", + " \"pipeline_parallel_size\": 1,\n", + " },\n", + " },\n", + " output={\"name\": OUTPUT_NAME},\n", + ")\n", + "\n", + "# `rl` auto-generates the job id (rl-); do not pass name=.\n", + "job = sdk.customization.rl.jobs.create(spec=spec, workspace=\"default\")\n", + "print(f\"Job ID: {job.job.name}\")\n", + "print(f\"Output model: {OUTPUT_NAME}\")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 7. Track Training Progress\n", + "\n", + "The DPO job runs four steps: download -> **dpo-training** (Ray) -> upload -> model-entity. We poll the top-level job status and surface the training step's progress." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from IPython.display import clear_output\n", + "\n", + "check = max_wait_time_checker(7200, \"DPO job\")\n", + "while True:\n", + " check()\n", + " status = sdk.jobs.get_status(name=job.job.name, workspace=\"default\")\n", + " clear_output(wait=True)\n", + " print(f\"Job Status: {status.status}\")\n", + "\n", + " step = max_steps = phase = None\n", + " for job_step in status.steps or []:\n", + " if job_step.name == \"dpo-training\":\n", + " for task in job_step.tasks or []:\n", + " d = task.status_details or {}\n", + " step, max_steps, phase = d.get(\"step\"), d.get(\"max_steps\"), d.get(\"phase\")\n", + " break\n", + " break\n", + " if step is not None and max_steps:\n", + " print(f\"Training: Step {step}/{max_steps} ({100 * step / max_steps:.1f}%)\")\n", + " if phase:\n", + " print(f\"Phase: {phase}\")\n", + "\n", + " if status.status in (\"completed\", \"failed\", \"cancelled\", \"error\"):\n", + " print(f\"\\nJob finished: {status.status}\")\n", + " break\n", + " time.sleep(15)\n", + "\n", + "assert status.status == \"completed\"" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "**Interpreting DPO training metrics** (in `status_details.metrics`):\n", + "\n", + "- **`loss`** — the DPO loss; should trend down as the policy learns to separate chosen from rejected.\n", + "- **Reward margin** (chosen minus rejected reward) — should trend **up**: the model increasingly prefers chosen responses.\n", + "- **Validation `loss`** — watch for divergence from training loss (overfitting). Raise `ref_policy_kl_penalty` (β) or add `sft_loss_weight` if the policy drifts too far from the reference." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 8. Validate the Output Model\n", + "\n", + "DPO produces a **full-weight model entity** (not an adapter). Confirm it was registered." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "model_entity = sdk.models.retrieve(workspace=\"default\", name=OUTPUT_NAME)\n", + "print(model_entity.model_dump_json(indent=2))" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 9. Deploy and Evaluate (optional)\n", + "\n", + "The DPO output is a full model, so it deploys like any full-weight checkpoint (see the [Full SFT](/documentation/customizer-reference/tutorials/sft-customization-job) tutorial for details). We deploy with vLLM and send a chat completion." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "deploy_suffix = uuid.uuid4().hex[:8]\n", + "DEPLOYMENT_CONFIG_NAME = f\"dpo-deployment-cfg-{deploy_suffix}\"\n", + "DEPLOYMENT_NAME = f\"dpo-deployment-{deploy_suffix}\"\n", + "\n", + "deployment_config = sdk.inference.deployment_configs.create(\n", + " workspace=\"default\",\n", + " name=DEPLOYMENT_CONFIG_NAME,\n", + " engine=\"vllm\",\n", + " model_spec={\"model_namespace\": \"default\", \"model_name\": OUTPUT_NAME},\n", + " executor_config={\"gpu\": 1, \"image_name\": \"vllm/vllm-openai\", \"image_tag\": \"v0.22.1\"},\n", + ")\n", + "\n", + "deployment = sdk.inference.deployments.create(\n", + " workspace=\"default\", name=DEPLOYMENT_NAME, config=deployment_config.name\n", + ")\n", + "print(f\"Deployment name: {deployment.name}\")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "check = max_wait_time_checker(1800, \"Deployment\")\n", + "while True:\n", + " check()\n", + " deployment_status = sdk.inference.deployments.retrieve(name=deployment.name, workspace=\"default\")\n", + " clear_output(wait=True)\n", + " print(f\"Deployment status: {deployment_status.status}\")\n", + " deployment_state = str(deployment_status.status).lower()\n", + " if deployment_state in (\"ready\", \"running\"):\n", + " if not sdk.models.wait_for_gateway(deployment.name, workspace=\"default\", timeout=60):\n", + " raise RuntimeError(\"Inference gateway did not become ready\")\n", + " break\n", + " if deployment_state in (\"failed\", \"error\", \"terminated\", \"lost\"):\n", + " raise RuntimeError(f\"Deployment failed with status: {deployment_status.status}\")\n", + " time.sleep(15)" + ], + "execution_count": null, + "outputs": [], + "id": "a5811863" + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "messages = [\n", + " {\"role\": \"system\", \"content\": \"You are a helpful assistant.\"},\n", + " {\"role\": \"user\", \"content\": \"Write a short, friendly email to a colleague asking to reschedule our meeting to Thursday.\"},\n", + "]\n", + "\n", + "response = sdk.inference.gateway.provider.post(\n", + " \"v1/chat/completions\",\n", + " name=deployment.name,\n", + " workspace=\"default\",\n", + " body={\"model\": f\"default/{OUTPUT_NAME}\", \"messages\": messages, \"temperature\": 0.7, \"max_tokens\": 256},\n", + ")\n", + "print(\"Model output:\\n\")\n", + "print(response[\"choices\"][0][\"message\"][\"content\"])" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Conclusion\n", + "\n", + "You aligned a base model with **DPO** on the NeMo Platform using the `rl` backend:\n", + "\n", + "- Uploaded a HelpSteer3 preference dataset **as-is** (the platform detects the schema natively).\n", + "- Submitted a full-weight DPO job that ran on a Ray cluster via the Kubernetes executor.\n", + "- Registered the output as a full model entity and (optionally) deployed it for inference.\n", + "\n", + "**Next steps:** tune the alignment strength with `ref_policy_kl_penalty` (β), add `sft_loss_weight` to anchor the policy to the chosen responses, enable `activation_checkpointing` for memory headroom, or scale up with `parallelism`. See the `nemo-customizer` skill's `references/hyperparameters.md` (section NeMo-RL (DPO)) for the full knob reference." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": ".venv", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.14" + } }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Prerequisites\n", - "\n", - "Before starting this tutorial, ensure you have:\n", - "\n", - "1. **Completed the [Quickstart](../../get-started/quickstart.md)** to install the NeMo Platform and Python SDK.\n", - "2. **Installed the Python SDK** (PyPI wrapper: `pip install \"nemo-platform[all]\"`; source checkout: run `make bootstrap` from the repository root).\n", - "3. **Installed the `datasets` package**: `pip install datasets`.\n", - "4. **A platform configured with `platform.runtime: kubernetes`.** The `rl` (DPO) backend provisions a Ray cluster and has **no local Docker fallback** — `submit` fails fast on a Docker-runtime platform. Multi-node jobs (`parallelism.num_nodes > 1`) additionally require the platform-side `NMP_RL_MULTINODE_SHARED_STORAGE_PATH`.\n", - "5. **A HuggingFace token** with access to the gated base model (this tutorial uses `meta-llama/Llama-3.2-1B-Instruct`). Export it as `HF_TOKEN`.\n", - "6. **At least one GPU with CUDA 12.8+** and a GPU execution profile (`nemo jobs list-execution-profiles`)." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Quick Start\n", - "\n", - "### 1. Initialize the SDK\n", - "\n", - "The SDK needs your NeMo Platform server URL. By default `http://localhost:8080` is used; set `NMP_BASE_URL` to override:\n", - "\n", - "```sh\n", - "export NMP_BASE_URL=\n", - "```" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "import json\n", - "import os\n", - "import time\n", - "import uuid\n", - "from pathlib import Path\n", - "from nemo_platform import NeMoPlatform, ConflictError\n", - "from nemo_platform.types.secrets import PlatformSecretResponse\n", - "from nemo_platform.types.files import HuggingfaceStorageConfigParam\n", - "from nemo_rl_plugin.schema import RlJobInput\n", - "\n", - "\n", - "def max_wait_time_checker(seconds: int, label: str = \"\"):\n", - " \"\"\"Return a check() that raises TimeoutError once `seconds` have elapsed.\"\"\"\n", - " start = time.time()\n", - "\n", - " def check():\n", - " if time.time() - start > seconds:\n", - " raise TimeoutError(f\"{label} took longer than {seconds} seconds\")\n", - "\n", - " return check\n", - "\n", - "\n", - "NMP_BASE_URL = os.environ.get(\"NMP_BASE_URL\", \"http://localhost:8080\")\n", - "sdk = NeMoPlatform(base_url=NMP_BASE_URL, workspace=\"default\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### 2. Prepare the Preference Dataset\n", - "\n", - "DPO trains on **preference data**. The `rl` backend takes a **single** dataset fileset that holds both `training.jsonl` and `validation.jsonl`, and auto-detects the row schema from the first line. Three preference formats are supported (see the platform's `BinaryPreferenceDatasetItemSchema` / `HelpSteer3DatasetItemSchema` / `Tulu3PreferenceDatasetItemSchema`):" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "#### Binary Preference Format\n", - "\n", - "Simple `prompt` / `chosen` / `rejected` (the `prompt` may be a string or a list of chat messages):\n", - "\n", - "```json\n", - "{\"prompt\": \"What is the capital of France?\", \"chosen\": \"The capital of France is Paris.\", \"rejected\": \"I'm not sure.\"}\n", - "```" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "#### HelpSteer3 Format (used here)\n", - "\n", - "A conversation `context` (string or chat messages), two candidate `response1` / `response2`, and a signed `overall_preference` in -3..3 — **negative** means response 1 is preferred, **positive** means response 2, **0** is a tie. This is the **raw** schema of `nvidia/HelpSteer3`, so no conversion is needed:\n", - "\n", - "```json\n", - "{\"context\": [{\"role\": \"user\", \"content\": \"Explain how to use git rebase\"}], \"response1\": \"...\", \"response2\": \"...\", \"overall_preference\": -2}\n", - "```" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "#### Tulu3 Preference Format\n", - "\n", - "Full chat conversations for both the chosen and rejected branches (each a list of messages ending with the assistant turn):\n", - "\n", - "```json\n", - "{\"chosen\": [{\"role\": \"user\", \"content\": \"...\"}, {\"role\": \"assistant\", \"content\": \"preferred\"}], \"rejected\": [{\"role\": \"user\", \"content\": \"...\"}, {\"role\": \"assistant\", \"content\": \"dispreferred\"}]}\n", - "```" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "#### Download nvidia/HelpSteer3\n", - "\n", - "We use [nvidia/HelpSteer3](https://huggingface.co/datasets/nvidia/HelpSteer3) (the `preference` subset), NVIDIA's open preference dataset. It ships native `train` and `validation` splits and matches the HelpSteer3 schema above, so we upload the rows **as-is** — the platform's `HelpSteer3Dataset` loader handles the `overall_preference` semantics (including ties) at training time." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from datasets import load_dataset, Dataset\n", - "\n", - "print(\"Loading dataset nvidia/HelpSteer3 (preference subset)\")\n", - "ds = load_dataset(\"nvidia/HelpSteer3\", \"preference\")\n", - "\n", - "# Small subsets keep the tutorial fast; larger sets train better but take longer.\n", - "training_size = 3000\n", - "validation_size = 300\n", - "DATASET_NAME = \"dpo-dataset\"\n", - "DATASET_PATH = Path(\"dpo-dataset\").absolute()\n", - "os.makedirs(DATASET_PATH, exist_ok=True)\n", - "\n", - "train_dataset = ds[\"train\"]\n", - "validation_dataset = ds[\"validation\"]\n", - "assert isinstance(train_dataset, Dataset) and isinstance(validation_dataset, Dataset)\n", - "\n", - "# Save raw HelpSteer3 rows directly — no conversion. The platform detects the\n", - "# HelpSteer3 schema from the row keys (context / response1 / response2 / overall_preference).\n", - "train_dataset.select(range(training_size)).to_json(f\"{DATASET_PATH}/training.jsonl\")\n", - "validation_dataset.select(range(validation_size)).to_json(f\"{DATASET_PATH}/validation.jsonl\")\n", - "\n", - "print(f\"Saved training.jsonl ({training_size} rows) and validation.jsonl ({validation_size} rows)\")\n", - "with open(f\"{DATASET_PATH}/training.jsonl\") as f:\n", - " sample = json.loads(f.readline())\n", - "print(\"Sample keys:\", sorted(sample.keys()))\n", - "print(\"overall_preference:\", sample[\"overall_preference\"])" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### 3. Create FileSet and Upload Preference Data\n", - "\n", - "Upload both JSONL files to a single FileSet so the DPO job can read them." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "try:\n", - " sdk.files.filesets.create(workspace=\"default\", name=DATASET_NAME, description=\"DPO preference data\")\n", - " print(f\"Created fileset: {DATASET_NAME}\")\n", - "except ConflictError:\n", - " print(f\"Fileset '{DATASET_NAME}' already exists, continuing...\")\n", - "\n", - "sdk.files.upload(local_path=DATASET_PATH, remote_path=\"\", fileset=DATASET_NAME, workspace=\"default\")\n", - "\n", - "print(\"Preference data:\")\n", - "print(json.dumps([f.model_dump() for f in sdk.files.list(fileset=DATASET_NAME, workspace=\"default\").data], indent=2, default=str))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### 4. Secrets Setup\n", - "\n", - "The base model (`meta-llama/Llama-3.2-1B-Instruct`) is gated, so store your HuggingFace token as a platform secret named `hf-token` and reference it on the model fileset." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "HF_TOKEN = os.getenv(\"HF_TOKEN\")\n", - "\n", - "def create_or_get_secret(name: str, value: str | None, label: str) -> PlatformSecretResponse | None:\n", - " if not value:\n", - " print(f\"{label} is not set - skipping secret (gated model downloads will fail without it)\")\n", - " return None\n", - " try:\n", - " secret = sdk.secrets.create(name=name, workspace=\"default\", value=value)\n", - " print(f\"Created secret: {name}\")\n", - " return secret\n", - " except ConflictError:\n", - " print(f\"Secret '{name}' already exists, continuing...\")\n", - " return sdk.secrets.retrieve(name=name, workspace=\"default\")\n", - "\n", - "\n", - "hf_secret = create_or_get_secret(\"hf-token\", HF_TOKEN, \"HF_TOKEN\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### 5. Create Base Model FileSet and Model Entity\n", - "\n", - "DPO starts from an instruction-tuned base model. The model entity's spec is inferred asynchronously after creation." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "HF_REPO_ID = \"meta-llama/Llama-3.2-1B-Instruct\"\n", - "MODEL_NAME = \"llama-3-2-1b-instruct\"\n", - "\n", - "storage = HuggingfaceStorageConfigParam(type=\"huggingface\", repo_id=HF_REPO_ID, repo_type=\"model\")\n", - "if hf_secret:\n", - " storage[\"token_secret\"] = hf_secret.name\n", - "\n", - "try:\n", - " base_model_fs = sdk.files.filesets.create(\n", - " workspace=\"default\", name=MODEL_NAME, description=\"Llama 3.2 1B Instruct base model\", storage=storage\n", - " )\n", - " print(f\"Created base model fileset: {MODEL_NAME}\")\n", - "except ConflictError:\n", - " base_model_fs = sdk.files.filesets.retrieve(workspace=\"default\", name=MODEL_NAME)\n", - " print(\"Base model fileset already exists.\")\n", - "\n", - "try:\n", - " base_model = sdk.models.create(workspace=\"default\", name=MODEL_NAME, fileset=f\"default/{MODEL_NAME}\")\n", - "except ConflictError:\n", - " base_model = sdk.models.retrieve(workspace=\"default\", name=MODEL_NAME)\n", - "\n", - "print(f\"Base model fileset: fileset://default/{base_model.name}\")\n", - "\n", - "# Wait for the ModelSpec to be inferred from the checkpoint.\n", - "check = max_wait_time_checker(600, \"Model spec\")\n", - "while not base_model.spec:\n", - " check()\n", - " time.sleep(10)\n", - " base_model = sdk.models.retrieve(workspace=\"default\", name=MODEL_NAME)\n", - "print(\"Model spec ready\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### 6. Create the DPO Customization Job\n", - "\n", - "Submit a DPO job to the `rl` backend with `RlJobInput`. Note the DPO-specific shape:\n", - "\n", - "- `model` is a string ref to the model entity; `dataset` is a **single** string ref to the preference fileset (holding both files).\n", - "- The training method is `{\"type\": \"dpo\", ...}` — full-weight, no `finetuning_type`/LoRA.\n", - "- `ref_policy_kl_penalty` is **β** (DPO paper): how strongly the policy stays tied to the reference model.\n", - "- `rl` auto-generates the job id (`rl-`); read it back from the response.\n", - "\n", - "Other configurable knobs: `optimizer_type`, `adam_eps`, `activation_checkpointing`, `keep_top_k`, `val_at_end`, `preference_loss_weight`, `sft_loss_weight`. Run `nemo customization rl explain` for the live schema." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "job_suffix = uuid.uuid4().hex[:8]\n", - "OUTPUT_NAME = f\"llama-3-2-1b-dpo-{job_suffix}\"\n", - "\n", - "spec = RlJobInput(\n", - " model=f\"default/{base_model.name}\",\n", - " dataset=f\"default/{DATASET_NAME}\",\n", - " training={\n", - " \"type\": \"dpo\",\n", - " \"epochs\": 1,\n", - " \"batch_size\": 16,\n", - " \"micro_batch_size\": 1,\n", - " \"learning_rate\": 5e-6,\n", - " \"max_seq_length\": 4096,\n", - " \"ref_policy_kl_penalty\": 0.1,\n", - " \"parallelism\": {\n", - " \"num_nodes\": 1,\n", - " \"num_gpus_per_node\": 1,\n", - " \"tensor_parallel_size\": 1,\n", - " \"pipeline_parallel_size\": 1,\n", - " },\n", - " },\n", - " output={\"name\": OUTPUT_NAME},\n", - ")\n", - "\n", - "# `rl` auto-generates the job id (rl-); do not pass name=.\n", - "job = sdk.customization.rl.jobs.create(spec=spec, workspace=\"default\")\n", - "print(f\"Job ID: {job.job.name}\")\n", - "print(f\"Output model: {OUTPUT_NAME}\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### 7. Track Training Progress\n", - "\n", - "The DPO job runs four steps: download -> **dpo-training** (Ray) -> upload -> model-entity. We poll the top-level job status and surface the training step's progress." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from IPython.display import clear_output\n", - "\n", - "check = max_wait_time_checker(7200, \"DPO job\")\n", - "while True:\n", - " check()\n", - " status = sdk.jobs.get_status(name=job.job.name, workspace=\"default\")\n", - " clear_output(wait=True)\n", - " print(f\"Job Status: {status.status}\")\n", - "\n", - " step = max_steps = phase = None\n", - " for job_step in status.steps or []:\n", - " if job_step.name == \"dpo-training\":\n", - " for task in job_step.tasks or []:\n", - " d = task.status_details or {}\n", - " step, max_steps, phase = d.get(\"step\"), d.get(\"max_steps\"), d.get(\"phase\")\n", - " break\n", - " break\n", - " if step is not None and max_steps:\n", - " print(f\"Training: Step {step}/{max_steps} ({100 * step / max_steps:.1f}%)\")\n", - " if phase:\n", - " print(f\"Phase: {phase}\")\n", - "\n", - " if status.status in (\"completed\", \"failed\", \"cancelled\", \"error\"):\n", - " print(f\"\\nJob finished: {status.status}\")\n", - " break\n", - " time.sleep(15)\n", - "\n", - "assert status.status == \"completed\"" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "**Interpreting DPO training metrics** (in `status_details.metrics`):\n", - "\n", - "- **`loss`** — the DPO loss; should trend down as the policy learns to separate chosen from rejected.\n", - "- **Reward margin** (chosen minus rejected reward) — should trend **up**: the model increasingly prefers chosen responses.\n", - "- **Validation `loss`** — watch for divergence from training loss (overfitting). Raise `ref_policy_kl_penalty` (β) or add `sft_loss_weight` if the policy drifts too far from the reference." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### 8. Validate the Output Model\n", - "\n", - "DPO produces a **full-weight model entity** (not an adapter). Confirm it was registered." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "model_entity = sdk.models.retrieve(workspace=\"default\", name=OUTPUT_NAME)\n", - "print(model_entity.model_dump_json(indent=2))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### 9. Deploy and Evaluate (optional)\n", - "\n", - "The DPO output is a full model, so it deploys like any full-weight checkpoint (see the [Full SFT](./sft-customization-job) tutorial for details). We deploy with vLLM and send a chat completion." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "deploy_suffix = uuid.uuid4().hex[:8]\n", - "DEPLOYMENT_CONFIG_NAME = f\"dpo-deployment-cfg-{deploy_suffix}\"\n", - "DEPLOYMENT_NAME = f\"dpo-deployment-{deploy_suffix}\"\n", - "\n", - "deployment_config = sdk.inference.deployment_configs.create(\n", - " workspace=\"default\",\n", - " name=DEPLOYMENT_CONFIG_NAME,\n", - " engine=\"vllm\",\n", - " model_spec={\"model_namespace\": \"default\", \"model_name\": OUTPUT_NAME},\n", - " executor_config={\"gpu\": 1, \"image_name\": \"vllm/vllm-openai\", \"image_tag\": \"v0.22.1\"},\n", - ")\n", - "\n", - "deployment = sdk.inference.deployments.create(\n", - " workspace=\"default\", name=DEPLOYMENT_NAME, config=deployment_config.name\n", - ")\n", - "print(f\"Deployment name: {deployment.name}\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "check = max_wait_time_checker(1800, \"Deployment\")\n", - "while True:\n", - " check()\n", - " deployment_status = sdk.inference.deployments.retrieve(name=deployment.name, workspace=\"default\")\n", - " clear_output(wait=True)\n", - " print(f\"Deployment status: {deployment_status.status}\")\n", - " if str(deployment_status.status).lower() in (\"ready\", \"running\", \"failed\", \"error\"):\n", - " break\n", - " time.sleep(15)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "messages = [\n", - " {\"role\": \"system\", \"content\": \"You are a helpful assistant.\"},\n", - " {\"role\": \"user\", \"content\": \"Write a short, friendly email to a colleague asking to reschedule our meeting to Thursday.\"},\n", - "]\n", - "\n", - "response = sdk.inference.gateway.provider.post(\n", - " \"v1/chat/completions\",\n", - " name=deployment.name,\n", - " workspace=\"default\",\n", - " body={\"model\": f\"default/{OUTPUT_NAME}\", \"messages\": messages, \"temperature\": 0.7, \"max_tokens\": 256},\n", - ")\n", - "print(\"Model output:\\n\")\n", - "print(response[\"choices\"][0][\"message\"][\"content\"])" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Conclusion\n", - "\n", - "You aligned a base model with **DPO** on the NeMo Platform using the `rl` backend:\n", - "\n", - "- Uploaded a HelpSteer3 preference dataset **as-is** (the platform detects the schema natively).\n", - "- Submitted a full-weight DPO job that ran on a Ray cluster via the Kubernetes executor.\n", - "- Registered the output as a full model entity and (optionally) deployed it for inference.\n", - "\n", - "**Next steps:** tune the alignment strength with `ref_policy_kl_penalty` (β), add `sft_loss_weight` to anchor the policy to the chosen responses, enable `activation_checkpointing` for memory headroom, or scale up with `parallelism`. See the `nemo-customizer` skill's `references/hyperparameters.md` (section NeMo-RL (DPO)) for the full knob reference." - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": ".venv", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.11.14" - } - }, - "nbformat": 4, - "nbformat_minor": 5 -} + "nbformat": 4, + "nbformat_minor": 5 +} \ No newline at end of file diff --git a/docs/customizer/tutorials/embedding-customization-job.ipynb b/docs/customizer/tutorials/embedding-customization-job.ipynb index ee2332e9f4..ee6ebc3796 100644 --- a/docs/customizer/tutorials/embedding-customization-job.ipynb +++ b/docs/customizer/tutorials/embedding-customization-job.ipynb @@ -50,8 +50,10 @@ "\n", "1. **Completed the [Quickstart](../../get-started/quickstart.md)** to install and deploy NeMo Platform locally\n", "2. **Installed the Python SDK** (PyPI wrapper: `pip install \"nemo-platform[all]\"`; source checkout: run `make bootstrap` from the repository root)\n", - "3. **HuggingFace token** with read access to download the SPECTER dataset (get one at [huggingface.co/settings/tokens](https://huggingface.co/settings/tokens))\n", - "4. **NGC API key** to pull NIM container images from nvcr.io (get one at [ngc.nvidia.com](https://ngc.nvidia.com/) → Setup → Generate API Key)" + "3. **NGC API key** to pull NIM container images from nvcr.io (get one at [ngc.nvidia.com](https://ngc.nvidia.com/) → Setup → Generate API Key)\n", + "4. **At least one GPU with CUDA 13+**\n", + "\n", + "The SPECTER dataset and the tutorial's base model are public and do not require a Hugging Face token. If you substitute a gated or private model, provide a token with read access." ] }, { @@ -312,7 +314,7 @@ "source": [ "### 3. Prepare Dataset\n", "\n", - "Use the [SPECTER dataset](https://huggingface.co/datasets/embedding-data/SPECTER) from HuggingFace, a collection of scientific paper triplets where papers that cite each other are considered related.\n", + "Use the [SPECTER dataset](https://huggingface.co/datasets/embedding-data/SPECTER) from Hugging Face, a collection of scientific paper triplets where papers that cite each other are considered related.\n", "\n", "**Dataset structure:**\n", "- ~684K scientific paper triplets (this tutorial uses 10%)\n", @@ -344,12 +346,6 @@ "from datasets import load_dataset\n", "import json\n", "\n", - "# HuggingFace token for dataset access\n", - "HF_TOKEN = os.environ.get(\"HF_TOKEN\")\n", - "if not HF_TOKEN:\n", - " raise ValueError(\"HF_TOKEN environment variable is required. Get one at https://huggingface.co/settings/tokens\")\n", - "os.environ[\"HF_TOKEN\"] = HF_TOKEN\n", - "\n", "# Configuration\n", "DATASET_SIZE = 3000 # Number of triplets (increase for better results, max ~684K)\n", "VALIDATION_SPLIT = 0.05 # 5% held out for validation\n", @@ -439,11 +435,11 @@ "Configure authentication for accessing base models:\n", "\n", "- **NGC models** (`ngc://` URIs): Requires NGC API key\n", - "- **HuggingFace models** (`hf://` URIs): Requires HF token for gated/private models\n", + "- **Hugging Face models** (`hf://` URIs): Requires HF token for gated/private models\n", "\n", "Get your credentials:\n", "- [NGC API Key](https://ngc.nvidia.com/) (Setup → Generate API Key)\n", - "- [HuggingFace Token](https://huggingface.co/settings/tokens) (Create token with Read access)\n", + "- [Hugging Face Token](https://huggingface.co/settings/tokens) (Optional; needed only for a gated/private replacement model)\n", "\n", "---\n", "\n", @@ -477,9 +473,10 @@ " return client.secrets.retrieve(name=name, workspace=\"default\")\n", "\n", "\n", - "# Create HuggingFace token secret (for downloading model from HF during training)\n", - "hf_secret = create_or_get_secret(\"hf-token\", HF_TOKEN, \"HF_TOKEN\")\n", - "print(f\"HF_TOKEN secret: {hf_secret.name}\")\n", + "# Public Hugging Face models need no token. Create a secret only when HF_TOKEN is set.\n", + "hf_secret = create_or_get_secret(\"hf-token\", HF_TOKEN, \"HF_TOKEN\") if HF_TOKEN else None\n", + "if hf_secret:\n", + " print(f\"HF_TOKEN secret: {hf_secret.name}\")\n", "\n", "# NGC secret was already created in baseline step (Step 2), or use the platform default\n", "if \"NGC_SECRET_NAME\" not in globals():\n", @@ -495,7 +492,7 @@ "source": [ "### 7. Create Base Model FileSet and Model Entity\n", "\n", - "Create a fileset pointing to the [nvidia/llama-nemotron-embed-1b-v2](https://huggingface.co/nvidia/llama-nemotron-embed-1b-v2) embedding model from HuggingFace, then create a Model Entity that references this fileset. Model downloading will take place at training time." + "Create a fileset pointing to the [nvidia/llama-nemotron-embed-1b-v2](https://huggingface.co/nvidia/llama-nemotron-embed-1b-v2) embedding model from Hugging Face, then create a Model Entity that references this fileset. Model downloading will take place at training time." ] }, { @@ -508,20 +505,20 @@ "HF_REPO_ID = \"nvidia/llama-nemotron-embed-1b-v2\"\n", "MODEL_NAME = \"nv-nemotron-embed-1b-base\"\n", "\n", - "# Ensure you have a HuggingFace token secret created\n", + "storage = HuggingfaceStorageConfigParam(\n", + " type=\"huggingface\",\n", + " repo_id=HF_REPO_ID,\n", + " repo_type=\"model\",\n", + ")\n", + "if hf_secret:\n", + " storage[\"token_secret\"] = hf_secret.name\n", + "\n", "try:\n", " base_model_fs = client.files.filesets.create(\n", " workspace=\"default\",\n", " name=MODEL_NAME,\n", " description=\"NVIDIA Llama Nemotron Embed 1B v2 embedding model\",\n", - " storage=HuggingfaceStorageConfigParam(\n", - " type=\"huggingface\",\n", - " # repo_id is the full model name from Hugging Face\n", - " repo_id=HF_REPO_ID,\n", - " repo_type=\"model\",\n", - " # we use the secret created in the previous step\n", - " token_secret=hf_secret.name\n", - " )\n", + " storage=storage,\n", " )\n", "except ConflictError as e:\n", " print(f\"Base model fileset already exists. Skipping creation.\")\n", @@ -595,7 +592,7 @@ " \"training_type\": \"sft\",\n", " \"finetuning_type\": \"lora_merged\",\n", " \"lora\": {\"rank\": 16, \"alpha\": 32},\n", - " \"max_seq_length\": MAX_SEQ_LENGTH,\n", + " \"max_seq_length\": 512,\n", "}\n", "```\n" ] @@ -693,7 +690,10 @@ " print(f\"\\nJob finished with status: {status.status}\")\n", " break\n", " \n", - " time.sleep(10)" + " time.sleep(10)\n", + "\n", + "if status.status != \"completed\":\n", + " raise RuntimeError(f\"Training job finished with status: {status.status}\")" ], "execution_count": null, "outputs": [] @@ -918,7 +918,7 @@ "\n", "**Benchmark Evaluation**\n", "\n", - "For systematic evaluation, use the NeMo Evaluator service with retrieval benchmarks like SciDocs, BEIR, or MTEB. Refer to the [Evaluator documentation](../../evaluator/index.md) for details.\n", + "For systematic evaluation of end-to-end retrieval quality in a RAG pipeline, use the NeMo Evaluator [RAG metrics](../../evaluator/metrics/rag.md) (RAGAS `context_recall`, `context_precision`, and `context_relevance`).\n", "\n", "---\n", "\n", @@ -930,10 +930,10 @@ "\n", "| Parameter | Recommended | Notes |\n", "|-----------|-------------|-------|\n", - "| `learning_rate` | 1e-6 to 5e-6 | Lower than standard SFT |\n", - "| `batch_size` | 128-256 | Larger batches improve contrastive learning |\n", - "| `max_seq_length` | 512 | Typical for embedding models |\n", - "| `epochs` | 1-3 | Start small, increase if needed |\n", + "| `optimizer.learning_rate` | 1e-6 to 5e-6 | Lower than standard SFT |\n", + "| `batch.global_batch_size` | 128-256 | Larger batches improve contrastive learning |\n", + "| `training.max_seq_length` | 512 | Typical for embedding models |\n", + "| `schedule.epochs` | 1-3 | Start small, increase if needed |\n", "\n", "---\n", "\n", @@ -959,7 +959,7 @@ "## Next Steps\n", "\n", "- [Monitor training metrics](../manage-customization-jobs/get-job-status.md) in detail\n", - "- [Evaluate your model](../../evaluator/index.md) with retrieval benchmarks\n", + "- [Evaluate your model](../../evaluator/metrics/rag.md) with RAG metrics\n", "- Integrate the fine-tuned embedding model into your RAG pipeline\n", "- Scale up training with the full SPECTER dataset (~684K triplets) for better results" ] @@ -986,4 +986,4 @@ }, "nbformat": 4, "nbformat_minor": 2 -} +} \ No newline at end of file diff --git a/docs/customizer/tutorials/embedding-customization-job.mdx b/docs/customizer/tutorials/embedding-customization-job.mdx index 222dbf3ccb..a47770e29b 100644 --- a/docs/customizer/tutorials/embedding-customization-job.mdx +++ b/docs/customizer/tutorials/embedding-customization-job.mdx @@ -5,8 +5,6 @@ description: "" [Run in Google Colab](https://colab.research.google.com/github/NVIDIA-NeMo/nemo-platform/blob/main/docs/customizer/tutorials/embedding-customization-job.ipynb) -# Embedding Model Customization - Learn how to fine-tune an embedding model to improve retrieval accuracy for your specific domain. ## About @@ -43,8 +41,10 @@ Before starting this tutorial, ensure you have: 1. **Completed the [Quickstart](/documentation/get-started)** to install and deploy NeMo Platform locally 2. **Installed the Python SDK** (PyPI wrapper: `pip install "nemo-platform[all]"`; source checkout: run `make bootstrap` from the repository root) -3. **HuggingFace token** with read access to download the SPECTER dataset (get one at [huggingface.co/settings/tokens](https://huggingface.co/settings/tokens)) -4. **NGC API key** to pull NIM container images from nvcr.io (get one at [ngc.nvidia.com](https://ngc.nvidia.com/) → Setup → Generate API Key) +3. **NGC API key** to pull NIM container images from nvcr.io (get one at [ngc.nvidia.com](https://ngc.nvidia.com/) → Setup → Generate API Key) +4. **At least one GPU with CUDA 13+** + +The SPECTER dataset and the tutorial's base model are public and do not require a Hugging Face token. If you substitute a gated or private model, provide a token with read access. ## Quick Start @@ -256,7 +256,7 @@ print("GPU freed. Proceed to fine-tune and improve these rankings.") ### 3. Prepare Dataset -Use the [SPECTER dataset](https://huggingface.co/datasets/embedding-data/SPECTER) from HuggingFace, a collection of scientific paper triplets where papers that cite each other are considered related. +Use the [SPECTER dataset](https://huggingface.co/datasets/embedding-data/SPECTER) from Hugging Face, a collection of scientific paper triplets where papers that cite each other are considered related. **Dataset structure:** - ~684K scientific paper triplets (this tutorial uses 10%) @@ -279,12 +279,6 @@ from pathlib import Path from datasets import load_dataset import json -# HuggingFace token for dataset access -HF_TOKEN = os.environ.get("HF_TOKEN") -if not HF_TOKEN: - raise ValueError("HF_TOKEN environment variable is required. Get one at https://huggingface.co/settings/tokens") -os.environ["HF_TOKEN"] = HF_TOKEN - # Configuration DATASET_SIZE = 3000 # Number of triplets (increase for better results, max ~684K) VALIDATION_SPLIT = 0.05 # 5% held out for validation @@ -358,11 +352,11 @@ print(json.dumps([f.model_dump() for f in client.files.list(fileset=DATASET_NAME Configure authentication for accessing base models: - **NGC models** (`ngc://` URIs): Requires NGC API key -- **HuggingFace models** (`hf://` URIs): Requires HF token for gated/private models +- **Hugging Face models** (`hf://` URIs): Requires HF token for gated/private models Get your credentials: - [NGC API Key](https://ngc.nvidia.com/) (Setup → Generate API Key) -- [HuggingFace Token](https://huggingface.co/settings/tokens) (Create token with Read access) +- [Hugging Face Token](https://huggingface.co/settings/tokens) (Optional; needed only for a gated/private replacement model) --- @@ -392,9 +386,10 @@ def create_or_get_secret(name: str, value: str | None, label: str): return client.secrets.retrieve(name=name, workspace="default") -# Create HuggingFace token secret (for downloading model from HF during training) -hf_secret = create_or_get_secret("hf-token", HF_TOKEN, "HF_TOKEN") -print(f"HF_TOKEN secret: {hf_secret.name}") +# Public Hugging Face models need no token. Create a secret only when HF_TOKEN is set. +hf_secret = create_or_get_secret("hf-token", HF_TOKEN, "HF_TOKEN") if HF_TOKEN else None +if hf_secret: + print(f"HF_TOKEN secret: {hf_secret.name}") # NGC secret was already created in baseline step (Step 2), or use the platform default if "NGC_SECRET_NAME" not in globals(): @@ -404,7 +399,7 @@ print(f"NGC_API_KEY secret: {NGC_SECRET_NAME}") ### 7. Create Base Model FileSet and Model Entity -Create a fileset pointing to the [nvidia/llama-nemotron-embed-1b-v2](https://huggingface.co/nvidia/llama-nemotron-embed-1b-v2) embedding model from HuggingFace, then create a Model Entity that references this fileset. Model downloading will take place at training time. +Create a fileset pointing to the [nvidia/llama-nemotron-embed-1b-v2](https://huggingface.co/nvidia/llama-nemotron-embed-1b-v2) embedding model from Hugging Face, then create a Model Entity that references this fileset. Model downloading will take place at training time. ```python import time @@ -413,20 +408,20 @@ from nemo_platform.types.files import HuggingfaceStorageConfigParam HF_REPO_ID = "nvidia/llama-nemotron-embed-1b-v2" MODEL_NAME = "nv-nemotron-embed-1b-base" -# Ensure you have a HuggingFace token secret created +storage = HuggingfaceStorageConfigParam( + type="huggingface", + repo_id=HF_REPO_ID, + repo_type="model", +) +if hf_secret: + storage["token_secret"] = hf_secret.name + try: base_model_fs = client.files.filesets.create( workspace="default", name=MODEL_NAME, description="NVIDIA Llama Nemotron Embed 1B v2 embedding model", - storage=HuggingfaceStorageConfigParam( - type="huggingface", - # 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 - token_secret=hf_secret.name - ) + storage=storage, ) except ConflictError as e: print(f"Base model fileset already exists. Skipping creation.") @@ -494,7 +489,7 @@ training={ "training_type": "sft", "finetuning_type": "lora_merged", "lora": {"rank": 16, "alpha": 32}, - "max_seq_length": MAX_SEQ_LENGTH, + "max_seq_length": 512, } ``` @@ -580,6 +575,9 @@ while True: break time.sleep(10) + +if status.status != "completed": + raise RuntimeError(f"Training job finished with status: {status.status}") ``` **Interpreting Embedding Training Metrics:** @@ -770,10 +768,10 @@ For detailed information on all available hyperparameters, recommended values, a | Parameter | Recommended | Notes | |-----------|-------------|-------| -| `learning_rate` | 1e-6 to 5e-6 | Lower than standard SFT | -| `batch_size` | 128-256 | Larger batches improve contrastive learning | -| `max_seq_length` | 512 | Typical for embedding models | -| `epochs` | 1-3 | Start small, increase if needed | +| `optimizer.learning_rate` | 1e-6 to 5e-6 | Lower than standard SFT | +| `batch.global_batch_size` | 128-256 | Larger batches improve contrastive learning | +| `training.max_seq_length` | 512 | Typical for embedding models | +| `schedule.epochs` | 1-3 | Start small, increase if needed | --- diff --git a/docs/customizer/tutorials/format-training-dataset.mdx b/docs/customizer/tutorials/format-training-dataset.mdx index b08d157467..b60e8f88c1 100644 --- a/docs/customizer/tutorials/format-training-dataset.mdx +++ b/docs/customizer/tutorials/format-training-dataset.mdx @@ -288,7 +288,7 @@ if model.spec: print(f"Is Chat Model: {model.spec.is_chat}") print(f"Family: {model.spec.family}") print(f"Parameters: {model.spec.base_num_parameters:,}") - print(f"Max Sequence Length: {model.spec.max_sequence_length}") + print(f"Context Size: {model.spec.context_size}") ``` ### Chat with the Model @@ -334,8 +334,8 @@ print(f"Response: {response.choices[0].message.content}") Now that you know how to format your training datasets, you can proceed with creating customization jobs: -- [Start a LoRA Model Customization Job](./lora-customization-job.ipynb) - For parameter-efficient fine-tuning -- [Start a Full SFT Customization Job](./sft-customization-job.ipynb) - For full model fine-tuning +- [Start a LoRA Model Customization Job](/documentation/customizer-reference/tutorials/lora-customization-job) - For parameter-efficient fine-tuning +- [Start a Full SFT Customization Job](/documentation/customizer-reference/tutorials/sft-customization-job) - For full model fine-tuning --- diff --git a/docs/customizer/tutorials/import-hf-model.mdx b/docs/customizer/tutorials/import-hf-model.mdx index 8160f4a531..b75bf1be5f 100644 --- a/docs/customizer/tutorials/import-hf-model.mdx +++ b/docs/customizer/tutorials/import-hf-model.mdx @@ -1,10 +1,10 @@ --- -title: "Import and Fine-Tune Private HuggingFace Models" +title: "Import and Fine-Tune Private Hugging Face Models" description: "" --- -Use this tutorial to learn how to import a private HuggingFace model into NeMo Customizer, fine-tune it with LoRA, and deploy it for inference. +Use this tutorial to learn how to import a private Hugging Face model into NeMo Customizer, fine-tune it with LoRA, and deploy it for inference. ## Prerequisites @@ -14,8 +14,8 @@ Use this tutorial to learn how to import a private HuggingFace model into NeMo C - Completed the [Quickstart](/documentation/get-started) to install and deploy NeMo Platform locally. - Installed the Python SDK and any tutorial packages you need in your environment. -- A HuggingFace token with access to the private or gated model repository. -- A HuggingFace model with a compatible architecture. This tutorial uses `google/gemma-2-2b-it` as an example, but success depends on architectural compatibility. +- A Hugging Face token with access to the private or gated model repository. +- A Hugging Face model with a compatible architecture. This tutorial uses `google/gemma-2-2b-it` as an example, but success depends on architectural compatibility. - Sufficient GPU memory for the model and LoRA training job. @@ -100,9 +100,9 @@ NMP_BASE_URL = os.environ.get("NMP_BASE_URL", "http://localhost:8080") client = NeMoPlatform(base_url=NMP_BASE_URL, workspace="default") ``` -### 2. Store the HuggingFace Token +### 2. Store the Hugging Face Token -Private and gated HuggingFace repositories require a token. Store it as a NeMo Platform secret and reference that secret from the HuggingFace fileset. +Private and gated Hugging Face repositories require a token. Store it as a NeMo Platform secret and reference that secret from the Hugging Face fileset. ```python def create_or_get_secret(name: str, value: str | None, label: str) -> PlatformSecretResponse: @@ -126,7 +126,7 @@ hf_secret = create_or_get_secret("hf-token", os.getenv("HF_TOKEN"), "HF_TOKEN") ### 3. Create a Model FileSet and Model Entity -Create a HuggingFace-backed fileset for the private model, then register a Model Entity that points to that fileset. Model files are downloaded by the platform when training or deployment needs them. +Create a Hugging Face-backed fileset for the private model, then register a Model Entity that points to that fileset. Model files are downloaded by the platform when training or deployment needs them. ```python HF_REPO_ID = "google/gemma-2-2b-it" @@ -142,18 +142,18 @@ try: base_model_fs = client.files.filesets.create( workspace="default", name=MODEL_NAME, - description=f"Private HuggingFace model {HF_REPO_ID}", + description=f"Private Hugging Face model {HF_REPO_ID}", storage=model_storage, cache=True, ) print(f"Created model fileset: {base_model_fs.name}") except ConflictError: - print(f"Model fileset '{MODEL_NAME}' already exists, refreshing HuggingFace settings...") + print(f"Model fileset '{MODEL_NAME}' already exists, refreshing Hugging Face settings...") client.files.filesets.delete(workspace="default", name=MODEL_NAME) base_model_fs = client.files.filesets.create( workspace="default", name=MODEL_NAME, - description=f"Private HuggingFace model {HF_REPO_ID}", + description=f"Private Hugging Face model {HF_REPO_ID}", storage=model_storage, cache=True, ) @@ -258,7 +258,7 @@ try: client.files.filesets.create( workspace="default", name=DATASET_NAME, - description="Private HuggingFace model LoRA training data", + description="Private Hugging Face model LoRA training data", cache=True, ) print(f"Created dataset fileset: {DATASET_NAME}") @@ -442,8 +442,10 @@ def chat(model_id: str): ) -base_response = chat(f"default/{MODEL_NAME}") -lora_response = chat(f"default--{OUTPUT_NAME}") +BASE_INFERENCE_MODEL_NAME = f"default/{MODEL_NAME}" +INFERENCE_MODEL_NAME = f"default--{OUTPUT_NAME}" +base_response = chat(BASE_INFERENCE_MODEL_NAME) +lora_response = chat(INFERENCE_MODEL_NAME) print("Base model response:") print(base_response["choices"][0]["message"]["content"]) @@ -457,11 +459,12 @@ print(lora_response["choices"][0]["message"]["content"]) ```bash export OUTPUT_NAME="" +export INFERENCE_MODEL_NAME="default--${OUTPUT_NAME}" curl -s "${NMP_BASE_URL}/apis/inference-gateway/v2/workspaces/default/openai/-/v1/chat/completions" \ -H 'Content-Type: application/json' \ -d '{ - "model": "'${OUTPUT_NAME}'", + "model": "'${INFERENCE_MODEL_NAME}'", "messages": [ {"role": "user", "content": "Can you summarize what LoRA fine-tuning does?"} ], diff --git a/docs/customizer/tutorials/lora-customization-job.ipynb b/docs/customizer/tutorials/lora-customization-job.ipynb index 6bb42e9947..cf8f0667b0 100644 --- a/docs/customizer/tutorials/lora-customization-job.ipynb +++ b/docs/customizer/tutorials/lora-customization-job.ipynb @@ -27,7 +27,8 @@ "1. **Completed the [Quickstart](../../get-started/quickstart.md)** to install and deploy NeMo Platform locally\n", "2. **Installed the Python SDK** (PyPI wrapper: `pip install \"nemo-platform[all]\"`; source checkout: run `make bootstrap` from the repository root)\n", "3. **Installed the `datasets` package** for loading SQuAD: `pip install datasets`\n", - "4. **At least one GPU with CUDA 12.8+**\n" + "4. **At least one GPU with CUDA 13+**\n", + "" ] }, { @@ -213,9 +214,9 @@ "source": [ "### 4. Secrets Setup\n", "\n", - "For Huggingface models that require authentication, create a secret with your HF token. Get a token from [Huggingface Settings](https://huggingface.co/settings/tokens) and accept the model terms.\n", + "For Hugging Face models that require authentication, create a secret with your HF token. Get a token from [Hugging Face Settings](https://huggingface.co/settings/tokens) and accept the model terms.\n", "\n", - "This is generally true for LLaMa based models (e.g. [Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct)).\n", + "This is generally true for Llama-based models (for example, [Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct)).\n", "\n", "```sh\n", "export HF_TOKEN=\n", @@ -276,7 +277,7 @@ " base_model_fs = client.files.filesets.create(\n", " workspace=\"default\",\n", " name=MODEL_NAME,\n", - " description=\"Qwen3 0.6b base model from Huggingface\",\n", + " description=\"Qwen3 0.6b base model from Hugging Face\",\n", " storage=storage,\n", " cache=True,\n", " )\n", @@ -517,12 +518,13 @@ "messages = [\n", " {\"role\": \"user\", \"content\": f\"Based on the following context, answer the question.\\n\\nContext: {context}\\n\\nQuestion: {question}\"}\n", "]\n", + "INFERENCE_MODEL_NAME = f\"default--{OUTPUT_NAME}\"\n", "response = client.inference.gateway.provider.post(\n", " \"v1/chat/completions\",\n", " name=deployment_name,\n", " workspace=\"default\",\n", " body={\n", - " \"model\": OUTPUT_NAME,\n", + " \"model\": INFERENCE_MODEL_NAME,\n", " \"messages\": messages,\n", " \"temperature\": 0,\n", " \"max_tokens\": 256,\n", @@ -550,8 +552,7 @@ "\n", "- [Monitor training metrics](fine-tune-metrics) in detail\n", "- [Evaluate your fine-tuned model](../../evaluator/index) using the Evaluator service\n", - "- Try [Full SFT](./sft-customization-job) for other customization options\n", - "" + "- Try [Full SFT](./sft-customization-job) for other customization options\n" ] } ], diff --git a/docs/customizer/tutorials/lora-customization-job.mdx b/docs/customizer/tutorials/lora-customization-job.mdx index d7434b20a7..4cf7f208fa 100644 --- a/docs/customizer/tutorials/lora-customization-job.mdx +++ b/docs/customizer/tutorials/lora-customization-job.mdx @@ -5,8 +5,6 @@ description: "" [Run in Google Colab](https://colab.research.google.com/github/NVIDIA-NeMo/nemo-platform/blob/main/docs/customizer/tutorials/lora-customization-job.ipynb) -# LoRA Model Customization Job - Learn how to use the NeMo Platform to create a LoRA (Low-Rank Adaptation) customization job using a custom dataset. In this tutorial we use LoRA to fine-tune a **question-answering model** from the SQuAD dataset. LoRA is a parameter-efficient fine-tuning method that requires fewer computational resources than full fine-tuning. If you need full model fine-tuning instead, see the [Full SFT Customization Job](/documentation/customizer-reference/tutorials/sft-customization-job) tutorial. @@ -20,7 +18,7 @@ Before starting this tutorial, ensure you have: 1. **Completed the [Quickstart](/documentation/get-started)** to install and deploy NeMo Platform locally 2. **Installed the Python SDK** (PyPI wrapper: `pip install "nemo-platform[all]"`; source checkout: run `make bootstrap` from the repository root) 3. **Installed the `datasets` package** for loading SQuAD: `pip install datasets` -4. **At least one GPU with CUDA 12.8+** +4. **At least one GPU with CUDA 13+** ## Quick Start @@ -164,9 +162,9 @@ print(client.files.list(fileset=DATASET_NAME, workspace="default")) ### 4. Secrets Setup -For Huggingface models that require authentication, create a secret with your HF token. Get a token from [Huggingface Settings](https://huggingface.co/settings/tokens) and accept the model terms. +For Hugging Face models that require authentication, create a secret with your HF token. Get a token from [Hugging Face Settings](https://huggingface.co/settings/tokens) and accept the model terms. -This is generally true for LLaMa based models (e.g. [Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct)). +This is generally true for Llama-based models (for example, [Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct)). ```sh export HF_TOKEN= @@ -213,7 +211,7 @@ try: base_model_fs = client.files.filesets.create( workspace="default", name=MODEL_NAME, - description="Qwen3 0.6b base model from Huggingface", + description="Qwen3 0.6b base model from Hugging Face", storage=storage, cache=True, ) @@ -408,12 +406,13 @@ 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}"} ] +INFERENCE_MODEL_NAME = f"default--{OUTPUT_NAME}" response = client.inference.gateway.provider.post( "v1/chat/completions", name=deployment_name, workspace="default", body={ - "model": OUTPUT_NAME, + "model": INFERENCE_MODEL_NAME, "messages": messages, "temperature": 0, "max_tokens": 256, diff --git a/docs/customizer/tutorials/metrics.mdx b/docs/customizer/tutorials/metrics.mdx index b857a39caa..a7794d62ca 100644 --- a/docs/customizer/tutorials/metrics.mdx +++ b/docs/customizer/tutorials/metrics.mdx @@ -22,8 +22,8 @@ The time to complete this tutorial is approximately 10 minutes. ### Tutorial-Specific Prerequisites -- Completed customization job with a valid ID -- (Optional) Access to NeMo with MLflow tracking enabled +- Completed customization job with a valid job name +- (Optional) A job created with `spec.integrations.mlflow` and access to its configured MLflow tracking server ## Available Metrics @@ -62,7 +62,7 @@ for step in status.steps or []: print(f"Training Phase: {details.get('phase')}") print(f"Step: {details.get('step')}/{details.get('max_steps')}") print(f"Epoch: {details.get('epoch')}/{details.get('num_epochs')}") - print(f"Training Loss: {details.get('loss')}") + print(f"Training Loss: {details.get('train_loss')}") print(f"Validation Loss: {details.get('val_loss')}") print(f"Learning Rate: {details.get('lr')}") print(f"Gradient Norm: {details.get('grad_norm')}") @@ -72,16 +72,16 @@ The response includes training progress and metrics including loss, learning rat ### Using MLflow -If your deployment has MLflow tracking enabled: +If your customization job was created with an `integrations.mlflow` configuration (see [MLflow Integration](/documentation/customizer-reference/manage-customization-jobs/customization-job-reference#mlflow-integration)): -1. Access the MLflow UI at your cluster's MLflow tracking URL -2. Locate your experiment by the output model name -3. Find the run using your customization job ID +1. Access the MLflow UI at the configured `tracking_uri` +2. Locate the configured `experiment_name` (defaults to the output model name) +3. Find the configured run `name` (defaults to the customization job ID) 4. View detailed metrics, including training and validation loss curves, under the "Metrics" tab -MLflow integration is configured at the cluster level. Contact your administrator if you need access to the MLflow UI or if MLflow tracking is not enabled for your deployment. +MLflow tracking is requested per job through `spec.integrations.mlflow`; it is not automatically enabled for every job in a cluster. The tracking server can be selected with `tracking_uri` in the job spec or the platform-side `MLFLOW_TRACKING_URI` environment variable. Contact your administrator if you need access to that server. diff --git a/docs/customizer/tutorials/optimize-throughput.ipynb b/docs/customizer/tutorials/optimize-throughput.ipynb index a4b6171f91..771a3fe5ef 100644 --- a/docs/customizer/tutorials/optimize-throughput.ipynb +++ b/docs/customizer/tutorials/optimize-throughput.ipynb @@ -1,973 +1,985 @@ { - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "\n", - "\n", - "\n", - "# Optimize for Tokens/GPU Throughput\n", - "\n", - "## About\n", - "Learn how to use the NeMo Platform Customizer to create a [LoRA](nemo-ms-about-concepts-customization) (Low-Rank Adaptation) customization job optimized for higher tokens/GPU throughput and lower runtime. \n", - "\n", - "**In this tutorial, you will:**\n", - "1. Fine-tune [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct) on the SQuAD dataset using LoRA, with [sequence packing](nemo-ms-about-concepts-customization) enabled for one run and disabled for another.\n", - "2. Compare training runtime, GPU utilization, and memory allocation between the two runs.\n", - "3. Verify that validation loss remains comparable, confirming that sequence packing improves throughput without sacrificing model quality.\n", - "\n", - "> **Note:** While this tutorial demonstrates sequence packing with LoRA, the optimization is also available for all_weights (full) SFT customization jobs." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Prerequisites\n", - "\n", - "Before starting this tutorial, ensure you have:\n", - "\n", - "1. **Completed the [Quickstart](../../get-started/quickstart.md)** to install and deploy NeMo Platform locally\n", - "2. **Installed the Python SDK** (PyPI wrapper: `pip install \"nemo-platform[all]\"`; source checkout: run `make bootstrap` from the repository root)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Quick Start\n", - "\n", - "### 1. Initialize SDK\n", - "\n", - "The SDK needs to know your NeMo Platform server URL. By default, `http://localhost:8080` is used in accordance with the [Quickstart](../../get-started/quickstart.md) guide. If NeMo Platform is running at a custom location, you can override the URL by setting the `NMP_BASE_URL` environment variable:\n", - "\n", - "```sh\n", - "export NMP_BASE_URL=\n", - "```" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "import json\n", - "import os\n", - "from nemo_platform import NeMoPlatform, ConflictError\n", - "\n", - "NMP_BASE_URL = os.environ.get(\"NMP_BASE_URL\", \"http://localhost:8080\")\n", - "client = NeMoPlatform(\n", - " base_url=NMP_BASE_URL,\n", - " workspace=\"default\"\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### 2. Create Dataset FileSet and Upload Training Data" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Install additional dependencies if they are not installed in your Python environment.\n", - "\n", - "The cell below automatically detects your environment and uses:\n", - "- `uv pip install` if you're in a uv-managed virtual environment\n", - "- `pip install` otherwise\n", - "\n", - "Required packages:\n", - "- `datasets` - Download the public [rajpurkar/squad](https://huggingface.co/datasets/rajpurkar/squad) dataset\n", - "- `pandas` - Compare job results in table format\n", - "- `matplotlib` - Plot live training metrics (loss curves, GPU utilization)\n", - "- `nvidia-ml-py` - Collect GPU VRAM and compute utilization metrics during training" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "vscode": { - "languageId": "shellscript" + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "\n", + "\n", + "\n", + "# Optimize for Tokens/GPU Throughput\n", + "\n", + "## About\n", + "Learn how to use the NeMo Platform Customizer to create a [LoRA](nemo-ms-about-concepts-customization) (Low-Rank Adaptation) customization job optimized for higher tokens/GPU throughput and lower runtime. \n", + "\n", + "**In this tutorial, you will:**\n", + "1. Fine-tune [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct) on the SQuAD dataset using LoRA, with [sequence packing](nemo-ms-about-concepts-customization) enabled for one run and disabled for another.\n", + "2. Compare training runtime, GPU utilization, and memory allocation between the two runs.\n", + "3. Verify that validation loss remains comparable, confirming that sequence packing improves throughput without sacrificing model quality.\n", + "\n", + "> **Note:** While this tutorial demonstrates sequence packing with LoRA, the optimization is also available for all_weights (full) SFT customization jobs." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Prerequisites\n", + "\n", + "Before starting this tutorial, ensure you have:\n", + "\n", + "1. **Completed the [Quickstart](../../get-started/quickstart.md)** to install and deploy NeMo Platform locally\n", + "2. **Installed the Python SDK** (PyPI wrapper: `pip install \"nemo-platform[all]\"`; source checkout: run `make bootstrap` from the repository root)\n", + "3. **At least one GPU with CUDA 13+**" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Quick Start\n", + "\n", + "### 1. Initialize SDK\n", + "\n", + "The SDK needs to know your NeMo Platform server URL. By default, `http://localhost:8080` is used in accordance with the [Quickstart](../../get-started/quickstart.md) guide. If NeMo Platform is running at a custom location, you can override the URL by setting the `NMP_BASE_URL` environment variable:\n", + "\n", + "```sh\n", + "export NMP_BASE_URL=\n", + "```" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "import json\n", + "import os\n", + "from nemo_platform import NeMoPlatform, ConflictError\n", + "\n", + "NMP_BASE_URL = os.environ.get(\"NMP_BASE_URL\", \"http://localhost:8080\")\n", + "client = NeMoPlatform(\n", + " base_url=NMP_BASE_URL,\n", + " workspace=\"default\"\n", + ")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 2. Create Dataset FileSet and Upload Training Data" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Install additional dependencies if they are not installed in your Python environment.\n", + "\n", + "The cell below automatically detects your environment and uses:\n", + "- `uv pip install` if you're in a uv-managed virtual environment\n", + "- `pip install` otherwise\n", + "\n", + "Required packages:\n", + "- `datasets` - Download the public [rajpurkar/squad](https://huggingface.co/datasets/rajpurkar/squad) dataset\n", + "- `pandas` - Compare job results in table format\n", + "- `matplotlib` - Plot live training metrics (loss curves, GPU utilization)\n", + "- `nvidia-ml-py` - Collect GPU VRAM and compute utilization metrics during training" + ] + }, + { + "cell_type": "code", + "metadata": { + "vscode": { + "languageId": "shellscript" + } + }, + "source": [ + "if command -v uv >/dev/null 2>&1 && [ -n \"$VIRTUAL_ENV\" ]; then\n", + " uv pip install datasets pandas matplotlib nvidia-ml-py\n", + "else\n", + " pip install datasets pandas matplotlib nvidia-ml-py\n", + "fi" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "#### Download rajpurkar/squad Dataset\n", + "\n", + "SQuAD (Stanford Question Answering Dataset) is a reading comprehension dataset consisting of questions posed on Wikipedia articles, where the answer is a segment of text from the corresponding passage." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "import json\n", + "import os\n", + "from pathlib import Path\n", + "from datasets import load_dataset, Dataset, DatasetDict\n", + "\n", + "# Configuration\n", + "SEED = 1234\n", + "DATASET_NAME = \"sft-dataset\"\n", + "\n", + "# Convert SQuAD format to prompt/completion format and save to JSONL\n", + "def convert_squad_to_sft_format(example):\n", + " \"\"\"Convert SQuAD format to prompt/completion format for SFT training.\"\"\"\n", + " prompt = f\"Context: {example['context']} Question: {example['question']} Answer:\"\n", + " completion = example[\"answers\"][\"text\"][0] # Take the first answer\n", + " return {\"prompt\": prompt, \"completion\": completion}\n", + "\n", + "# Load the SQuAD dataset from Hugging Face\n", + "print(\"Loading dataset rajpurkar/squad\")\n", + "ds = load_dataset(\"rajpurkar/squad\")\n", + "if not isinstance(ds, DatasetDict):\n", + " raise ValueError(\"Dataset does not contain expected splits\")\n", + "\n", + "print(\"Loaded dataset\")\n", + "\n", + "# For the purpose of this tutorial, we'll use a subset of the dataset\n", + "# We use a reduced dataset size (3000 training/300 validation samples) to keep tutorial runtime manageable\n", + "# while still demonstrating the performance benefits of sequence packing. The larger the dataset,\n", + "# the better the model will perform but the longer the training will take.\n", + "training_size = 3000\n", + "validation_size = 300\n", + "DATASET_PATH = Path(DATASET_NAME).absolute()\n", + "\n", + "# Get training split and verify it's a Dataset (not IterableDataset)\n", + "train_dataset = ds[\"train\"]\n", + "validation_dataset = ds[\"validation\"]\n", + "assert isinstance(train_dataset, Dataset), \"Expected Dataset type\"\n", + "assert isinstance(validation_dataset, Dataset), \"Expected Dataset type\"\n", + "\n", + "# Select subsets and save to JSONL files\n", + "training_ds = train_dataset.select(range(training_size))\n", + "validation_ds = validation_dataset.select(range(validation_size))\n", + "\n", + "# Transform to SFT format (prompt/completion)\n", + "training_ds = training_ds.map(convert_squad_to_sft_format, remove_columns=training_ds.column_names)\n", + "validation_ds = validation_ds.map(convert_squad_to_sft_format, remove_columns=validation_ds.column_names)\n", + "\n", + "# Create directory if it doesn't exist\n", + "# Note: This will create a local 'sft-dataset/' directory with training.jsonl and validation.jsonl files\n", + "os.makedirs(DATASET_PATH, exist_ok=True)\n", + "\n", + "# Save subsets to JSONL files\n", + "training_ds.to_json(f\"{DATASET_PATH}/training.jsonl\")\n", + "validation_ds.to_json(f\"{DATASET_PATH}/validation.jsonl\")\n", + "\n", + "print(f\"Saved training.jsonl with {len(training_ds)} rows\")\n", + "print(f\"Saved validation.jsonl with {len(validation_ds)} rows\")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Create fileset to store SFT training data\n", + "\n", + "try:\n", + " client.files.filesets.create(\n", + " workspace=\"default\",\n", + " name=DATASET_NAME,\n", + " description=\"SFT training data\"\n", + " )\n", + " print(f\"Created fileset: {DATASET_NAME}\")\n", + "except ConflictError:\n", + " print(f\"Fileset '{DATASET_NAME}' already exists, continuing...\")\n", + "\n", + "# Upload training data files individually to ensure correct structure\n", + "client.files.upload(\n", + " local_path=f\"{DATASET_PATH}/\", # Trailing slash uploads directory contents to fileset root\n", + " remote_path=\"\",\n", + " fileset=DATASET_NAME,\n", + " workspace=\"default\"\n", + ")\n", + "\n", + "# Validate training data is uploaded correctly\n", + "print(\"Training data:\")\n", + "print(json.dumps([f.model_dump() for f in client.files.list(fileset=DATASET_NAME, workspace=\"default\").data], indent=2))" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 3. Secrets Setup\n", + "\n", + "If you plan to use NGC or Hugging Face models, you will need to configure authentication:\n", + "\n", + "- **NGC models** (`ngc://` URIs): Requires NGC API key\n", + "- **Hugging Face models** (`hf://` URIs): Requires HF token for gated/private models\n", + "\n", + "\n", + "Configure these as secrets in your platform. Refer to [Managing Secrets](../../get-started/concepts/manage-secrets.md) for detailed instructions.\n", + "\n", + "Get your credentials to access base models:\n", + "- [NGC API Key](https://ngc.nvidia.com/) (Setup → Generate API Key)\n", + "- [Hugging Face Token](https://huggingface.co/settings/tokens) (Create token with Read access)\n", + "\n", + "\n", + "---\n", + "\n", + "#### Quick Setup Example\n", + "\n", + "This tutorial uses the [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main) model from Hugging Face. Ensure that you have sufficient permissions to download the model. If you cannot access the files on the [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main) Hugging Face page, request access.\n", + "\n", + "**Hugging Face Authentication:**\n", + "- For gated models (Llama, Gemma), you must provide a Hugging Face token via the `token_secret` parameter\n", + "- Get your token from [Hugging Face Settings](https://huggingface.co/settings/tokens) (requires Read access)\n", + "- Accept the model's terms on the Hugging Face model page before using it. Example: [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main)\n", + "- For public models, you can omit the `token_secret` parameter when creating a fileset for the model in the next step." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Export the HF_TOKEN and NGC_API_KEY environment variables if they are not already set\n", + "HF_TOKEN = os.getenv(\"HF_TOKEN\")\n", + "NGC_API_KEY = os.getenv(\"NGC_API_KEY\")\n", + "\n", + "\n", + "def create_or_get_secret(name: str, value: str | None, label: str):\n", + " if not value:\n", + " raise ValueError(f\"{label} environment variable is not set. Set it and try again.\")\n", + " try:\n", + " secret = client.secrets.create(\n", + " name=name,\n", + " workspace=\"default\",\n", + " value=value,\n", + " )\n", + " print(f\"Created secret: {name}\")\n", + " return secret\n", + " except ConflictError:\n", + " print(f\"Secret '{name}' already exists, continuing...\")\n", + " return client.secrets.retrieve(name=name, workspace=\"default\")\n", + "\n", + "\n", + "# Create Hugging Face token secret\n", + "hf_secret = create_or_get_secret(\"hf-token\", HF_TOKEN, \"HF_TOKEN\")\n", + "print(\"HF_TOKEN secret:\")\n", + "print(hf_secret.model_dump_json(indent=2))\n", + "\n", + "# Create NGC API key secret\n", + "# Uncomment the line below if you have NGC API Key and want to finetune NGC models\n", + "# ngc_api_key = create_or_get_secret(\"ngc-api-key\", NGC_API_KEY, \"NGC_API_KEY\")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 4. Create Base Model FileSet\n", + "\n", + "Create a fileset pointing to the [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main) model on Hugging Face. This step creates a pointer to the model on Hugging Face and does not download it. The model is downloaded at job creation time.\n", + "\n", + "Note: for public models, you can omit the `token_secret` parameter when creating a model fileset." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "import time\n", + "from nemo_platform.types.files import HuggingfaceStorageConfigParam\n", + "\n", + "HF_REPO_ID = \"meta-llama/Llama-3.2-1B-Instruct\"\n", + "MODEL_NAME = \"llama-3-2-1b-base\"\n", + "\n", + "# Ensure you have a Hugging Face token secret created\n", + "# Create a fileset pointing to the desired Hugging Face model\n", + "try:\n", + " base_model_fs = client.files.filesets.create(\n", + " workspace=\"default\",\n", + " name=MODEL_NAME,\n", + " description=\"Llama 3.2 1B base model from Hugging Face\",\n", + " storage=HuggingfaceStorageConfigParam(\n", + " type=\"huggingface\",\n", + " # repo_id is the full model name from Hugging Face\n", + " repo_id=HF_REPO_ID,\n", + " repo_type=\"model\",\n", + " # we use the secret created in the previous step\n", + " token_secret=hf_secret.name\n", + " )\n", + " )\n", + " print(f\"Created base model fileset: {MODEL_NAME}\")\n", + "except ConflictError:\n", + " print(f\"Base model fileset already exists. Skipping creation.\")\n", + " base_model_fs = client.files.filesets.retrieve(\n", + " workspace=\"default\",\n", + " name=MODEL_NAME,\n", + " )\n", + "\n", + "# Create the Model Entity representation.\n", + "try:\n", + " base_model = client.models.create(\n", + " workspace=\"default\",\n", + " name=MODEL_NAME,\n", + " fileset=f\"default/{MODEL_NAME}\",\n", + " )\n", + " print(f\"Created Model Entity: {MODEL_NAME}\")\n", + "except ConflictError:\n", + " print(f\"Base model already exists. Updating fileset if different.\")\n", + " base_model = client.models.update(\n", + " workspace=\"default\",\n", + " name=MODEL_NAME,\n", + " fileset=f\"default/{MODEL_NAME}\",\n", + " )\n", + "\n", + "print(f\"\\nBase model fileset: fileset://default/{base_model.name}\")\n", + "print(\"Base model fileset files list:\")\n", + "print(json.dumps([f.model_dump() for f in client.files.list(fileset=MODEL_NAME, workspace=\"default\").data], indent=2))\n", + "\n", + "# Wait for ModelSpec to be populated from the checkpoint\n", + "print(\"\\nWaiting for ModelSpec to be populated...\")\n", + "SPEC_TIMEOUT_SECONDS = 120\n", + "spec_start = time.time()\n", + "while not base_model.spec:\n", + " if time.time() - spec_start > SPEC_TIMEOUT_SECONDS:\n", + " raise TimeoutError(f\"ModelSpec not populated within {SPEC_TIMEOUT_SECONDS} seconds\")\n", + " time.sleep(2)\n", + " base_model = client.models.retrieve(\n", + " workspace=\"default\",\n", + " name=MODEL_NAME,\n", + " )\n", + "\n", + "print(f\"ModelSpec populated: {base_model.spec}\")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 5. Create LoRA Job with Sequence Packing\n", + "Create a LoRA customization job with **sequence packing** enabled via `AutomodelJobInput` (`batch.sequence_packing=True`)." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "import uuid\n", + "from nemo_automodel_plugin.schema import AutomodelJobInput\n", + "\n", + "SEQUENCE_PACKING_ENABLED = True\n", + "\n", + "job_suffix = uuid.uuid4().hex[:4]\n", + "JOB_NAME = f\"packing-job-{job_suffix}\"\n", + "PACK_OUTPUT_NAME = f\"packing-out-{job_suffix}\"\n", + "\n", + "spec = AutomodelJobInput(\n", + " model=f\"default/{base_model.name}\",\n", + " dataset={\"training\": f\"default/{DATASET_NAME}\"},\n", + " training={\n", + " \"training_type\": \"sft\",\n", + " \"finetuning_type\": \"lora\",\n", + " \"max_seq_length\": 4096,\n", + " },\n", + " schedule={\"epochs\": 1, \"val_check_interval\": 0.1},\n", + " batch={\n", + " \"global_batch_size\": 64,\n", + " \"micro_batch_size\": 1,\n", + " \"sequence_packing\": SEQUENCE_PACKING_ENABLED,\n", + " },\n", + " optimizer={\"learning_rate\": 5e-5},\n", + " parallelism={\"num_gpus_per_node\": 1},\n", + " output={\"name\": PACK_OUTPUT_NAME},\n", + ")\n", + "\n", + "job_with_sequence_packing = client.customization.automodel.jobs.create(\n", + " spec=spec, workspace=\"default\", name=JOB_NAME\n", + ")\n", + "\n", + "print(f\"Submitted job: {job_with_sequence_packing.job.name}\")\n", + "print(f\"Output adapter: {PACK_OUTPUT_NAME}\")\n" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 6. Track Fine-Tuning Progress\n", + "\n", + "A training job contains multiple steps: \n", + "- Model and dataset downloading\n", + "- Fine-tuning where LoRA adapter weights are trained\n", + "- Creating a fileset entry for the fine-tuned model\n", + "- Fine-tuned weights uploading\n", + "\n", + "The elapsed time printed below reflects progress of the entire job. We compare the time taken by the fine-tuning step for both jobs in the last section of this tutorial." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "#### Define Helper Functions" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "\n", + "# Helpers to draw GPU VRAM Utilization and Validation Loss\n", + "import matplotlib.pyplot as plt\n", + "try:\n", + " import pynvml\n", + " _PYNVML_AVAILABLE = True\n", + "except ImportError:\n", + " _PYNVML_AVAILABLE = False\n", + " print(\"Note: Install nvidia-ml-py ('pip install nvidia-ml-py' or 'uv pip install nvidia-ml-py') to enable live GPU metrics.\")\n", + "\n", + "# ---------------------------------------------------------------------------\n", + "# GPU metrics collection (nvidia-ml-py; import name is pynvml)\n", + "# ---------------------------------------------------------------------------\n", + "\n", + "def _get_gpu_snapshot() -> tuple[list[float], list[float]]:\n", + " \"\"\"Return (vram_usage_pcts, compute_util_pcts) for each GPU.\"\"\"\n", + " if not _PYNVML_AVAILABLE:\n", + " return [], []\n", + " pynvml.nvmlInit()\n", + " try:\n", + " vram, util = [], []\n", + " for i in range(pynvml.nvmlDeviceGetCount()):\n", + " h = pynvml.nvmlDeviceGetHandleByIndex(i)\n", + " mem = pynvml.nvmlDeviceGetMemoryInfo(h)\n", + " rates = pynvml.nvmlDeviceGetUtilizationRates(h)\n", + " vram.append(int(mem.used) / int(mem.total) * 100)\n", + " util.append(float(rates.gpu))\n", + " return vram, util\n", + " finally:\n", + " pynvml.nvmlShutdown()\n", + "\n", + "\n", + "# ---------------------------------------------------------------------------\n", + "# Dashboard drawing helpers\n", + "# ---------------------------------------------------------------------------\n", + "\n", + "_PALETTE = {\n", + " \"val_loss\": \"#E74C3C\",\n", + " \"train_loss\": \"#F39C12\",\n", + " \"vram\": [\"#3498DB\", \"#9B59B6\", \"#1ABC9C\", \"#E67E22\"],\n", + " \"util\": [\"#2ECC71\", \"#E74C3C\", \"#3498DB\", \"#F1C40F\"],\n", + " \"grid\": \"#ECECEC\",\n", + " \"title\": \"#2C3E50\",\n", + " \"subtitle\": \"#7F8C8D\",\n", + " \"spine\": \"#CCCCCC\",\n", + " \"tick\": \"#666666\",\n", + "}\n", + "\n", + "\n", + "def _style_axis(ax):\n", + " \"\"\"Apply shared cosmetic styling to a subplot axis.\"\"\"\n", + " ax.set_facecolor(\"white\")\n", + " ax.grid(True, alpha=0.4, color=_PALETTE[\"grid\"], linewidth=0.8)\n", + " for spine in (\"top\", \"right\"):\n", + " ax.spines[spine].set_visible(False)\n", + " ax.spines[\"left\"].set_color(_PALETTE[\"spine\"])\n", + " ax.spines[\"bottom\"].set_color(_PALETTE[\"spine\"])\n", + " ax.tick_params(colors=_PALETTE[\"tick\"], labelsize=9)\n", + "\n", + "\n", + "def _plot_line(ax, xs, ys, color, label, fill=True):\n", + " \"\"\"Plot a time series, gracefully skipping None values.\"\"\"\n", + " pts = [(x, y) for x, y in zip(xs, ys) if y is not None]\n", + " if not pts:\n", + " return\n", + " px, py = zip(*pts)\n", + " ax.plot(\n", + " px, py, color=color, linewidth=2.2,\n", + " marker=\"o\", markersize=4,\n", + " markerfacecolor=\"white\", markeredgewidth=1.8, markeredgecolor=color,\n", + " label=label, zorder=3,\n", + " )\n", + " if fill:\n", + " ax.fill_between(px, py, alpha=0.08, color=color)\n", + "\n", + "\n", + "def _plot_gpu_panel(ax, xs, history, colors, fallback_label):\n", + " \"\"\"Plot per-GPU time series with area fill.\"\"\"\n", + " if not history or not history[0]:\n", + " ax.text(\n", + " 0.5, 0.5, \"No GPU data\", transform=ax.transAxes,\n", + " ha=\"center\", va=\"center\", fontsize=11, color=\"#AAAAAA\",\n", + " )\n", + " return\n", + " n_gpus = max(len(snap) for snap in history)\n", + " for g in range(n_gpus):\n", + " vals = [snap[g] if g < len(snap) else 0 for snap in history]\n", + " c = colors[g % len(colors)]\n", + " label = f\"GPU {g}\" if n_gpus > 1 else fallback_label\n", + " ax.plot(xs[: len(vals)], vals, color=c, linewidth=2, label=label)\n", + " ax.fill_between(xs[: len(vals)], vals, alpha=0.08, color=c)\n", + " if n_gpus > 1:\n", + " ax.legend(fontsize=9, framealpha=0.9, edgecolor=\"#DDD\")\n", + "\n", + "\n", + "def _draw_dashboard(\n", + " elapsed_mins, val_losses, train_losses,\n", + " vram_history, util_history,\n", + " job_name, status_str, step_str, elapsed_str,\n", + "):\n", + " \"\"\"Render a live 1x3 training dashboard.\"\"\"\n", + " fig, axes = plt.subplots(1, 3, figsize=(20, 5.5))\n", + " fig.patch.set_facecolor(\"#FAFBFC\")\n", + "\n", + " fig.suptitle(\n", + " job_name, fontsize=15, fontweight=\"bold\",\n", + " color=_PALETTE[\"title\"], y=1.10,\n", + " )\n", + " fig.text(\n", + " 0.5, 1.01,\n", + " f\"{status_str} | {step_str} | {elapsed_str}\",\n", + " ha=\"center\", fontsize=13, color=_PALETTE[\"subtitle\"],\n", + " )\n", + "\n", + " for ax in axes:\n", + " _style_axis(ax)\n", + "\n", + " # -- Panel 1: Loss curves --\n", + " _plot_line(axes[0], elapsed_mins, val_losses, _PALETTE[\"val_loss\"], \"Val Loss\", fill=True)\n", + " _plot_line(axes[0], elapsed_mins, train_losses, _PALETTE[\"train_loss\"], \"Train Loss\", fill=False)\n", + " axes[0].set_title(\"Train/Validation Loss\", fontsize=13, fontweight=\"bold\", color=_PALETTE[\"title\"], pad=12)\n", + " axes[0].set_xlabel(\"Time (min)\", fontsize=10, color=\"#666\")\n", + " axes[0].set_ylabel(\"Loss\", fontsize=10, color=\"#666\")\n", + " if any(v is not None for v in val_losses + train_losses):\n", + " axes[0].legend(fontsize=9, framealpha=0.9, edgecolor=\"#DDD\")\n", + "\n", + " # -- Panel 2: GPU VRAM usage --\n", + " _plot_gpu_panel(axes[1], elapsed_mins, vram_history, _PALETTE[\"vram\"], \"VRAM\")\n", + " axes[1].set_title(\"GPU VRAM Usage\", fontsize=13, fontweight=\"bold\", color=_PALETTE[\"title\"], pad=12)\n", + " axes[1].set_xlabel(\"Time (min)\", fontsize=10, color=\"#666\")\n", + " axes[1].set_ylabel(\"Usage (%)\", fontsize=10, color=\"#666\")\n", + " axes[1].set_ylim(-2, 105)\n", + "\n", + " # -- Panel 3: GPU utilization --\n", + " _plot_gpu_panel(axes[2], elapsed_mins, util_history, _PALETTE[\"util\"], \"Utilization\")\n", + " axes[2].set_title(\"GPU Utilization\", fontsize=13, fontweight=\"bold\", color=_PALETTE[\"title\"], pad=12)\n", + " axes[2].set_xlabel(\"Time (min)\", fontsize=10, color=\"#666\")\n", + " axes[2].set_ylabel(\"Utilization (%)\", fontsize=10, color=\"#666\")\n", + " axes[2].set_ylim(-2, 105)\n", + "\n", + " plt.tight_layout(rect=[0, 0, 1, 0.98])\n", + " plt.show()" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "#### Monitor the Job Until Completion\n", + "\n", + "The cell below polls the job status every 10 seconds and renders a live dashboard with validation loss, GPU VRAM usage, and GPU utilization charts. The charts appear empty at first while the model and dataset download; training metrics and GPU activity populate after the fine-tuning step begins.\n", + "\n", + "> **Note:** This is additional code. You can also use the Weights & Biases or MLflow integrations." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "import time\n", + "from typing import cast\n", + "from IPython.display import clear_output\n", + "from nemo_platform.types.shared import PlatformJobStatusResponse\n", + "\n", + "# Timeout set to 30 minutes to accommodate typical LoRA training duration for this dataset size.\n", + "# Actual training time will vary based on hardware, model size, and dataset complexity.\n", + "TIMEOUT_SECONDS = 30 * 60 # 30 minutes\n", + "VAL_LOSS_KEY = \"val_loss\"\n", + "TRAIN_LOSS_KEY = \"train_loss\"\n", + "\n", + "\n", + "def get_training_metric(\n", + " status: PlatformJobStatusResponse,\n", + " metric_key: str,\n", + ") -> float | None:\n", + " \"\"\"Return a metric reported by a task in the training step.\"\"\"\n", + " for job_step in status.steps or []:\n", + " if job_step.name == \"training\":\n", + " for task in job_step.tasks or []:\n", + " value = (task.status_details or {}).get(metric_key)\n", + " if value is not None:\n", + " return float(value)\n", + " return None\n", + "\n", + "\n", + "# ---------------------------------------------------------------------------\n", + "# Job polling with live dashboard\n", + "# ---------------------------------------------------------------------------\n", + "\n", + "def wait_for_job(\n", + " workspace: str,\n", + " job_name: str,\n", + " timeout: int = TIMEOUT_SECONDS,\n", + " poll_interval: int = 10,\n", + " val_loss_key: str = VAL_LOSS_KEY,\n", + " train_loss_key: str = TRAIN_LOSS_KEY,\n", + ") -> PlatformJobStatusResponse:\n", + " \"\"\"\n", + " Poll job status until completed, failed, cancelled, or timeout.\n", + " Displays a live dashboard with loss curves and GPU metrics.\n", + "\n", + " Args:\n", + " workspace: The workspace where the job is running.\n", + " job_name: The name of the job to monitor.\n", + " timeout: Maximum time to wait in seconds (default: 30 minutes).\n", + " poll_interval: Time between status checks in seconds (default: 10).\n", + "\n", + " Returns:\n", + " The final job status response.\n", + " \"\"\"\n", + " start_time = time.time()\n", + "\n", + " # Time-series accumulators required for plotting\n", + " elapsed_mins: list[float] = []\n", + " val_losses: list[float | None] = []\n", + " train_losses: list[float | None] = []\n", + " vram_history: list[list[float]] = []\n", + " util_history: list[list[float]] = []\n", + "\n", + " while True:\n", + " elapsed = time.time() - start_time\n", + " elapsed_min = elapsed / 60\n", + "\n", + " # Check for timeout\n", + " if elapsed > timeout:\n", + " error_message = f\"Timeout reached after {elapsed_min:.1f} minutes\"\n", + " print(f\"\\n{error_message}\")\n", + " print(\"Job did not complete within the timeout period.\")\n", + " raise Exception(error_message)\n", + "\n", + " status = client.jobs.get_status(name=job_name, workspace=workspace)\n", + "\n", + " # -- Extract training progress from nested steps structure --\n", + " step: int | None = None\n", + " max_steps: int | None = None\n", + " training_phase: str | None = None\n", + " val_loss: float | None = None\n", + " train_loss: float | None = None\n", + " current_step_name: str | None = None\n", + " current_step_phase: str | None = None\n", + "\n", + " for job_step in status.steps or []:\n", + " # Track the current active step name and phase for progress display\n", + " if job_step.tasks:\n", + " task = job_step.tasks[0]\n", + " td = task.status_details or {}\n", + " phase = cast(str, td.get(\"phase\", \"\"))\n", + " # Update current step if it's active or pending (not completed)\n", + " if job_step.status in (\"active\", \"pending\"):\n", + " current_step_name = job_step.name\n", + " current_step_phase = phase or \"started\"\n", + "\n", + " if job_step.name == \"training\":\n", + " for task in job_step.tasks or []:\n", + " td = task.status_details or {}\n", + " step = cast(int, td[\"step\"]) if \"step\" in td else None\n", + " max_steps = cast(int, td[\"max_steps\"]) if \"max_steps\" in td else None\n", + " training_phase = cast(str, td[\"phase\"]) if \"phase\" in td else None\n", + " raw_val_loss = td.get(val_loss_key)\n", + " val_loss = float(raw_val_loss) if raw_val_loss is not None else None\n", + " raw_train_loss = td.get(train_loss_key)\n", + " train_loss = float(raw_train_loss) if raw_train_loss is not None else None\n", + " break\n", + " break\n", + "\n", + " if val_loss is None:\n", + " raw_val_loss = (status.status_details or {}).get(val_loss_key)\n", + " val_loss = float(raw_val_loss) if raw_val_loss is not None else None\n", + " if train_loss is None:\n", + " raw_train_loss = (status.status_details or {}).get(train_loss_key)\n", + " train_loss = float(raw_train_loss) if raw_train_loss is not None else None\n", + "\n", + " # -- Collect GPU snapshot --\n", + " vram_pcts, util_pcts = _get_gpu_snapshot()\n", + "\n", + " # -- Append to accumulators used for the plots --\n", + " elapsed_mins.append(elapsed_min)\n", + " val_losses.append(val_loss)\n", + " train_losses.append(train_loss)\n", + " vram_history.append(vram_pcts)\n", + " util_history.append(util_pcts)\n", + "\n", + " # -- Build status strings --\n", + " status_str = f\"Status: {status.status}\"\n", + " if step is not None and max_steps is not None:\n", + " pct = step / max_steps * 100\n", + " step_str = f\"Step {step}/{max_steps} ({pct:.0f}%)\"\n", + " if training_phase:\n", + " step_str += f\" - {training_phase}\"\n", + " else:\n", + " if current_step_name and current_step_phase:\n", + " step_str = f\"{current_step_name} - {current_step_phase}\"\n", + " elif current_step_name:\n", + " step_str = f\"{current_step_name}\"\n", + " else:\n", + " step_str = \"Waiting for training to start...\"\n", + " elapsed_str = f\"Elapsed: {elapsed_min:.1f} min\"\n", + "\n", + " # -- Redraw dashboard --\n", + " clear_output(wait=True)\n", + " _draw_dashboard(\n", + " elapsed_mins, val_losses, train_losses,\n", + " vram_history, util_history,\n", + " job_name, status_str, step_str, elapsed_str,\n", + " )\n", + "\n", + " # -- Check terminal conditions --\n", + " if status.status.lower() == \"completed\":\n", + " # Redraw dashboard one final time with \"completed\" status\n", + " status_str = f\"Status: {status.status}\"\n", + " if step is not None and max_steps is not None:\n", + " step_str = f\"Step {max_steps}/{max_steps} (100%)\"\n", + " clear_output(wait=True)\n", + " _draw_dashboard(\n", + " elapsed_mins, val_losses, train_losses,\n", + " vram_history, util_history,\n", + " job_name, status_str, step_str, elapsed_str,\n", + " )\n", + " print(f\"\\nJob completed in {elapsed_min:.1f} minutes ({elapsed:.0f}s)\")\n", + " return status\n", + " elif status.status.lower() in (\"failed\", \"cancelled\", \"error\"):\n", + " print(f\"\\nJob finished with status: {status.status}\")\n", + " print(f\"Total time elapsed: {elapsed_min:.1f} minutes ({elapsed:.0f}s)\")\n", + "\n", + " # Print error details from the job level\n", + " if status.error_details:\n", + " error_msg = status.error_details.get(\"message\", \"\")\n", + " if error_msg:\n", + " print(f\"\\nError: {error_msg}\")\n", + "\n", + " # Find and print error details from the failed step/task\n", + " for job_step in status.steps or []:\n", + " if job_step.status == \"error\":\n", + " print(f\"\\nFailed step: {job_step.name}\")\n", + " if job_step.error_details:\n", + " step_error = job_step.error_details.get(\"message\", \"\")\n", + " if step_error:\n", + " print(f\"Step error: {step_error}\")\n", + " # Get error_stack from the failed task\n", + " for task in job_step.tasks or []:\n", + " if task.status == \"error\" and hasattr(task, \"error_stack\") and task.error_stack:\n", + " print(f\"\\nError stack trace:\\n{task.error_stack}\")\n", + " elif task.status == \"error\" and task.error_details:\n", + " task_error = task.error_details.get(\"message\", \"\")\n", + " if task_error:\n", + " print(f\"Task error: {task_error}\")\n", + " break\n", + "\n", + " raise Exception(f\"Job finished with status: {status.status}\")\n", + "\n", + " time.sleep(poll_interval)\n", + "\n", + "\n", + "# Wait for the job to complete\n", + "job_with_sequence_packing_status = wait_for_job(\n", + " workspace=\"default\",\n", + " job_name=job_with_sequence_packing.job.name,\n", + " timeout=TIMEOUT_SECONDS,\n", + ")\n", + "\n", + "packed_val_loss = get_training_metric(job_with_sequence_packing_status, VAL_LOSS_KEY)\n", + "if packed_val_loss is not None:\n", + " print(f\"Validation loss: {packed_val_loss:.2f}\")\n", + "else:\n", + " print(\"Validation loss: not reported in job status\")\n" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 7. Create LoRA Job without Sequence Packing\n", + "Create a second Automodel LoRA job with `batch.sequence_packing=False` for comparison." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "import uuid\n", + "from nemo_automodel_plugin.schema import AutomodelJobInput\n", + "\n", + "job_suffix = uuid.uuid4().hex[:4]\n", + "JOB_NAME = f\"no-packing-job-{job_suffix}\"\n", + "NO_PACK_OUTPUT_NAME = f\"no-packing-out-{job_suffix}\"\n", + "\n", + "spec = AutomodelJobInput(\n", + " model=f\"default/{base_model.name}\",\n", + " dataset={\"training\": f\"default/{DATASET_NAME}\"},\n", + " training={\n", + " \"training_type\": \"sft\",\n", + " \"finetuning_type\": \"lora\",\n", + " \"max_seq_length\": 4096,\n", + " },\n", + " schedule={\"epochs\": 1, \"val_check_interval\": 0.1},\n", + " batch={\n", + " \"global_batch_size\": 64,\n", + " \"micro_batch_size\": 1,\n", + " \"sequence_packing\": False,\n", + " },\n", + " optimizer={\"learning_rate\": 5e-5},\n", + " parallelism={\"num_gpus_per_node\": 1},\n", + " output={\"name\": NO_PACK_OUTPUT_NAME},\n", + ")\n", + "\n", + "job_without_sequence_packing = client.customization.automodel.jobs.create(\n", + " spec=spec, workspace=\"default\", name=JOB_NAME\n", + ")\n", + "\n", + "print(f\"Submitted job: {job_without_sequence_packing.job.name}\")\n", + "print(f\"Output adapter: {NO_PACK_OUTPUT_NAME}\")\n" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 8. Track Fine-Tuning Progress for Job without Sequence Packing" + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Wait for the training step to complete\n", + "job_without_sequence_packing_status = wait_for_job(\n", + " workspace=\"default\",\n", + " job_name=job_without_sequence_packing.job.name,\n", + " timeout=TIMEOUT_SECONDS\n", + ")\n", + "\n", + "no_pack_val_loss = get_training_metric(job_without_sequence_packing_status, VAL_LOSS_KEY)\n", + "if no_pack_val_loss is not None:\n", + " print(f\"Validation loss: {no_pack_val_loss:.2f}\")\n", + "else:\n", + " print(\"Validation loss: not reported in job status\")" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 9. Compare Results\n", + "- Time to complete training should be significantly lower for the job that used sequence packing.\n", + "- The expected validation loss for both jobs should be similar.\n", + "- Sequence packed version should have a higher GPU utilization and higher GPU Memory Allocation." + ] + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "from nemo_platform.types.jobs import PlatformJobStep\n", + "import pandas as pd\n", + "\n", + "STEP_NAME = \"training\"\n", + "\n", + "def get_elapsed_time(step: PlatformJobStep) -> float:\n", + " \"\"\"Calculate elapsed time in seconds from step's created_at to updated_at.\"\"\"\n", + " if step.created_at is None or step.updated_at is None:\n", + " raise ValueError(\"Training step timestamps are unavailable\")\n", + " return (step.updated_at - step.created_at).total_seconds()\n", + "\n", + "step_with_sequence_packing = client.jobs.steps.retrieve(\n", + " name=STEP_NAME,\n", + " workspace=\"default\",\n", + " job=job_with_sequence_packing.job.name,\n", + ")\n", + "\n", + "step_without_sequence_packing = client.jobs.steps.retrieve(\n", + " name=STEP_NAME,\n", + " workspace=\"default\",\n", + " job=job_without_sequence_packing.job.name,\n", + ")\n", + "\n", + "time_to_complete_with_sequence_packing = get_elapsed_time(step_with_sequence_packing)\n", + "time_to_complete_without_sequence_packing = get_elapsed_time(step_without_sequence_packing)\n", + "\n", + "# Display results as a table\n", + "results_df = pd.DataFrame({\n", + " \"Seq Packing Enabled\": [True, False],\n", + " \"Val Loss\": [packed_val_loss, no_pack_val_loss],\n", + " \"Training Step Time, sec\": [\n", + " time_to_complete_with_sequence_packing,\n", + " time_to_complete_without_sequence_packing\n", + " ]\n", + "})\n", + "\n", + "results_df.style.format({\"Val Loss\": \"{:.2f}\", \"Training Step Time, sec\": \"{:.0f}\"}).hide(axis='index')" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "#### Examples of Validation Loss\n", + "\n", + "The expected validation loss curves should match closely for both jobs.\n", + "![Validation loss comparison chart showing similar convergence patterns between sequence-packed and non-packed training runs over training steps](../_images/packed_vs_not_packed_val_loss.png)\n", + "\n", + "Sequence packed version should complete significantly faster.\n", + "![Runtime comparison chart demonstrating significantly reduced training time for sequence-packed job compared to non-packed baseline](../_images/runtime.png)\n", + "\n", + "#### GPU Utilization\n", + "Sequence packed version should have a higher GPU utilization.\n", + "![GPU utilization chart showing higher and more consistent GPU usage with sequence packing enabled throughout the training process](../_images/gpu_utilization.png)\n", + "\n", + "#### GPU Memory Allocation\n", + "Sequence packed version should have a higher GPU Memory Allocation.\n", + "![GPU memory allocation chart illustrating increased memory utilization efficiency with sequence packing enabled](../_images/gpu_memory.png)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": ".venv", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.14" } - }, - "outputs": [], - "source": [ - "if command -v uv >/dev/null 2>&1 && [ -n \"$VIRTUAL_ENV\" ]; then\n", - " uv pip install datasets pandas matplotlib nvidia-ml-py\n", - "else\n", - " pip install datasets pandas matplotlib nvidia-ml-py\n", - "fi" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "#### Download rajpurkar/squad Dataset\n", - "\n", - "SQuAD (Stanford Question Answering Dataset) is a reading comprehension dataset consisting of questions posed on Wikipedia articles, where the answer is a segment of text from the corresponding passage." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "import json\n", - "import os\n", - "from pathlib import Path\n", - "from datasets import load_dataset, Dataset, DatasetDict\n", - "\n", - "# Configuration\n", - "SEED = 1234\n", - "DATASET_NAME = \"sft-dataset\"\n", - "\n", - "# Convert SQuAD format to prompt/completion format and save to JSONL\n", - "def convert_squad_to_sft_format(example):\n", - " \"\"\"Convert SQuAD format to prompt/completion format for SFT training.\"\"\"\n", - " prompt = f\"Context: {example['context']} Question: {example['question']} Answer:\"\n", - " completion = example[\"answers\"][\"text\"][0] # Take the first answer\n", - " return {\"prompt\": prompt, \"completion\": completion}\n", - "\n", - "# Load the SQuAD dataset from Hugging Face\n", - "print(\"Loading dataset rajpurkar/squad\")\n", - "ds = load_dataset(\"rajpurkar/squad\")\n", - "if not isinstance(ds, DatasetDict):\n", - " raise ValueError(\"Dataset does not contain expected splits\")\n", - "\n", - "print(\"Loaded dataset\")\n", - "\n", - "# For the purpose of this tutorial, we'll use a subset of the dataset\n", - "# We use a reduced dataset size (3000 training/300 validation samples) to keep tutorial runtime manageable\n", - "# while still demonstrating the performance benefits of sequence packing. The larger the dataset,\n", - "# the better the model will perform but the longer the training will take.\n", - "training_size = 3000\n", - "validation_size = 300\n", - "DATASET_PATH = Path(DATASET_NAME).absolute()\n", - "\n", - "# Get training split and verify it's a Dataset (not IterableDataset)\n", - "train_dataset = ds[\"train\"]\n", - "validation_dataset = ds[\"validation\"]\n", - "assert isinstance(train_dataset, Dataset), \"Expected Dataset type\"\n", - "assert isinstance(validation_dataset, Dataset), \"Expected Dataset type\"\n", - "\n", - "# Select subsets and save to JSONL files\n", - "training_ds = train_dataset.select(range(training_size))\n", - "validation_ds = validation_dataset.select(range(validation_size))\n", - "\n", - "# Transform to SFT format (prompt/completion)\n", - "training_ds = training_ds.map(convert_squad_to_sft_format, remove_columns=training_ds.column_names)\n", - "validation_ds = validation_ds.map(convert_squad_to_sft_format, remove_columns=validation_ds.column_names)\n", - "\n", - "# Create directory if it doesn't exist\n", - "# Note: This will create a local 'sft-dataset/' directory with training.jsonl and validation.jsonl files\n", - "os.makedirs(DATASET_PATH, exist_ok=True)\n", - "\n", - "# Save subsets to JSONL files\n", - "training_ds.to_json(f\"{DATASET_PATH}/training.jsonl\")\n", - "validation_ds.to_json(f\"{DATASET_PATH}/validation.jsonl\")\n", - "\n", - "print(f\"Saved training.jsonl with {len(training_ds)} rows\")\n", - "print(f\"Saved validation.jsonl with {len(validation_ds)} rows\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Create fileset to store SFT training data\n", - "\n", - "try:\n", - " client.files.filesets.create(\n", - " workspace=\"default\",\n", - " name=DATASET_NAME,\n", - " description=\"SFT training data\"\n", - " )\n", - " print(f\"Created fileset: {DATASET_NAME}\")\n", - "except ConflictError:\n", - " print(f\"Fileset '{DATASET_NAME}' already exists, continuing...\")\n", - "\n", - "# Upload training data files individually to ensure correct structure\n", - "client.files.upload(\n", - " local_path=f\"{DATASET_PATH}/\", # Trailing slash uploads directory contents to fileset root\n", - " remote_path=\"\",\n", - " fileset=DATASET_NAME,\n", - " workspace=\"default\"\n", - ")\n", - "\n", - "# Validate training data is uploaded correctly\n", - "print(\"Training data:\")\n", - "print(json.dumps([f.model_dump() for f in client.files.list(fileset=DATASET_NAME, workspace=\"default\").data], indent=2))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### 3. Secrets Setup\n", - "\n", - "If you plan to use NGC or HuggingFace models, you will need to configure authentication:\n", - "\n", - "- **NGC models** (`ngc://` URIs): Requires NGC API key\n", - "- **HuggingFace models** (`hf://` URIs): Requires HF token for gated/private models\n", - "\n", - "\n", - "Configure these as secrets in your platform. Refer to [Managing Secrets](../../get-started/concepts/manage-secrets.md) for detailed instructions.\n", - "\n", - "Get your credentials to access base models:\n", - "- [NGC API Key](https://ngc.nvidia.com/) (Setup → Generate API Key)\n", - "- [HuggingFace Token](https://huggingface.co/settings/tokens) (Create token with Read access)\n", - "\n", - "\n", - "---\n", - "\n", - "#### Quick Setup Example\n", - "\n", - "This tutorial uses the [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main) model from HuggingFace. Ensure that you have sufficient permissions to download the model. If you cannot access the files on the [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main) Hugging Face page, request access.\n", - "\n", - "**HuggingFace Authentication:**\n", - "- For gated models (Llama, Gemma), you must provide a HuggingFace token via the `token_secret` parameter\n", - "- Get your token from [HuggingFace Settings](https://huggingface.co/settings/tokens) (requires Read access)\n", - "- Accept the model's terms on the HuggingFace model page before using it. Example: [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main)\n", - "- For public models, you can omit the `token_secret` parameter when creating a fileset for the model in the next step." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Export the HF_TOKEN and NGC_API_KEY environment variables if they are not already set\n", - "HF_TOKEN = os.getenv(\"HF_TOKEN\")\n", - "NGC_API_KEY = os.getenv(\"NGC_API_KEY\")\n", - "\n", - "\n", - "def create_or_get_secret(name: str, value: str | None, label: str):\n", - " if not value:\n", - " raise ValueError(f\"{label} environment variable is not set. Set it and try again.\")\n", - " try:\n", - " secret = client.secrets.create(\n", - " name=name,\n", - " workspace=\"default\",\n", - " value=value,\n", - " )\n", - " print(f\"Created secret: {name}\")\n", - " return secret\n", - " except ConflictError:\n", - " print(f\"Secret '{name}' already exists, continuing...\")\n", - " return client.secrets.retrieve(name=name, workspace=\"default\")\n", - "\n", - "\n", - "# Create HuggingFace token secret\n", - "hf_secret = create_or_get_secret(\"hf-token\", HF_TOKEN, \"HF_TOKEN\")\n", - "print(\"HF_TOKEN secret:\")\n", - "print(hf_secret.model_dump_json(indent=2))\n", - "\n", - "# Create NGC API key secret\n", - "# Uncomment the line below if you have NGC API Key and want to finetune NGC models\n", - "# ngc_api_key = create_or_get_secret(\"ngc-api-key\", NGC_API_KEY, \"NGC_API_KEY\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### 4. Create Base Model FileSet\n", - "\n", - "Create a fileset pointing to the [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main) model on HuggingFace. This step creates a pointer to the model on Hugging Face and does not download it. The model is downloaded at job creation time.\n", - "\n", - "Note: for public models, you can omit the `token_secret` parameter when creating a model fileset." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "import time\n", - "from nemo_platform.types.files import HuggingfaceStorageConfigParam\n", - "\n", - "HF_REPO_ID = \"meta-llama/Llama-3.2-1B-Instruct\"\n", - "MODEL_NAME = \"llama-3-2-1b-base\"\n", - "\n", - "# Ensure you have a HuggingFace token secret created\n", - "# Create a fileset pointing to the desired HuggingFace model\n", - "try:\n", - " base_model_fs = client.files.filesets.create(\n", - " workspace=\"default\",\n", - " name=MODEL_NAME,\n", - " description=\"Llama 3.2 1B base model from HuggingFace\",\n", - " storage=HuggingfaceStorageConfigParam(\n", - " type=\"huggingface\",\n", - " # repo_id is the full model name from Hugging Face\n", - " repo_id=HF_REPO_ID,\n", - " repo_type=\"model\",\n", - " # we use the secret created in the previous step\n", - " token_secret=hf_secret.name\n", - " )\n", - " )\n", - " print(f\"Created base model fileset: {MODEL_NAME}\")\n", - "except ConflictError:\n", - " print(f\"Base model fileset already exists. Skipping creation.\")\n", - " base_model_fs = client.files.filesets.retrieve(\n", - " workspace=\"default\",\n", - " name=MODEL_NAME,\n", - " )\n", - "\n", - "# Create the Model Entity representation.\n", - "try:\n", - " base_model = client.models.create(\n", - " workspace=\"default\",\n", - " name=MODEL_NAME,\n", - " fileset=f\"default/{MODEL_NAME}\",\n", - " )\n", - " print(f\"Created Model Entity: {MODEL_NAME}\")\n", - "except ConflictError:\n", - " print(f\"Base model already exists. Updating fileset if different.\")\n", - " base_model = client.models.update(\n", - " workspace=\"default\",\n", - " name=MODEL_NAME,\n", - " fileset=f\"default/{MODEL_NAME}\",\n", - " )\n", - "\n", - "print(f\"\\nBase model fileset: fileset://default/{base_model.name}\")\n", - "print(\"Base model fileset files list:\")\n", - "print(json.dumps([f.model_dump() for f in client.files.list(fileset=MODEL_NAME, workspace=\"default\").data], indent=2))\n", - "\n", - "# Wait for ModelSpec to be populated from the checkpoint\n", - "print(\"\\nWaiting for ModelSpec to be populated...\")\n", - "SPEC_TIMEOUT_SECONDS = 120\n", - "spec_start = time.time()\n", - "while not base_model.spec:\n", - " if time.time() - spec_start > SPEC_TIMEOUT_SECONDS:\n", - " raise TimeoutError(f\"ModelSpec not populated within {SPEC_TIMEOUT_SECONDS} seconds\")\n", - " time.sleep(2)\n", - " base_model = client.models.retrieve(\n", - " workspace=\"default\",\n", - " name=MODEL_NAME,\n", - " )\n", - "\n", - "print(f\"ModelSpec populated: {base_model.spec}\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### 5. Create LoRA Job with Sequence Packing\n", - "Create a LoRA customization job with **sequence packing** enabled via `AutomodelJobInput` (`batch.sequence_packing=True`)." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "import uuid\n", - "from nemo_automodel_plugin.schema import AutomodelJobInput\n", - "\n", - "SEQUENCE_PACKING_ENABLED = True\n", - "\n", - "job_suffix = uuid.uuid4().hex[:4]\n", - "JOB_NAME = f\"packing-job-{job_suffix}\"\n", - "PACK_OUTPUT_NAME = f\"packing-out-{job_suffix}\"\n", - "\n", - "spec = AutomodelJobInput(\n", - " model=f\"default/{base_model.name}\",\n", - " dataset={\"training\": f\"default/{DATASET_NAME}\"},\n", - " training={\n", - " \"training_type\": \"sft\",\n", - " \"finetuning_type\": \"lora\",\n", - " \"max_seq_length\": 4096,\n", - " },\n", - " schedule={\"epochs\": 1, \"val_check_interval\": 0.1},\n", - " batch={\n", - " \"global_batch_size\": 64,\n", - " \"micro_batch_size\": 1,\n", - " \"sequence_packing\": SEQUENCE_PACKING_ENABLED,\n", - " },\n", - " optimizer={\"learning_rate\": 5e-5},\n", - " parallelism={\"num_gpus_per_node\": 1},\n", - " output={\"name\": PACK_OUTPUT_NAME},\n", - ")\n", - "\n", - "job_with_sequence_packing = client.customization.automodel.jobs.create(\n", - " spec=spec, workspace=\"default\", name=JOB_NAME\n", - ")\n", - "\n", - "print(f\"Submitted job: {job_with_sequence_packing.job.name}\")\n", - "print(f\"Output adapter: {PACK_OUTPUT_NAME}\")\n" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### 6. Track Finetuning Progress\n", - "\n", - "A training job contains multiple steps: \n", - "- Model and dataset downloading\n", - "- Finetuning where LoRA adapter weights are trained\n", - "- Creating a fileset entry for the finetuned model\n", - "- Finetuned weights uploading\n", - "\n", - "The elapsed time printed below reflects progress of the entire job. We compare the time taken by the finetuning step for both jobs in the last section of this tutorial." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "#### Define Helper Functions" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "\n", - "# Helpers to draw GPU VRAM Utilization and Validation Loss\n", - "import matplotlib.pyplot as plt\n", - "try:\n", - " import pynvml\n", - " _PYNVML_AVAILABLE = True\n", - "except ImportError:\n", - " _PYNVML_AVAILABLE = False\n", - " print(\"Note: Install nvidia-ml-py ('pip install nvidia-ml-py' or 'uv pip install nvidia-ml-py') to enable live GPU metrics.\")\n", - "\n", - "# ---------------------------------------------------------------------------\n", - "# GPU metrics collection (nvidia-ml-py; import name is pynvml)\n", - "# ---------------------------------------------------------------------------\n", - "\n", - "def _get_gpu_snapshot() -> tuple[list[float], list[float]]:\n", - " \"\"\"Return (vram_usage_pcts, compute_util_pcts) for each GPU.\"\"\"\n", - " if not _PYNVML_AVAILABLE:\n", - " return [], []\n", - " pynvml.nvmlInit()\n", - " try:\n", - " vram, util = [], []\n", - " for i in range(pynvml.nvmlDeviceGetCount()):\n", - " h = pynvml.nvmlDeviceGetHandleByIndex(i)\n", - " mem = pynvml.nvmlDeviceGetMemoryInfo(h)\n", - " rates = pynvml.nvmlDeviceGetUtilizationRates(h)\n", - " vram.append(int(mem.used) / int(mem.total) * 100)\n", - " util.append(float(rates.gpu))\n", - " return vram, util\n", - " finally:\n", - " pynvml.nvmlShutdown()\n", - "\n", - "\n", - "# ---------------------------------------------------------------------------\n", - "# Dashboard drawing helpers\n", - "# ---------------------------------------------------------------------------\n", - "\n", - "_PALETTE = {\n", - " \"val_loss\": \"#E74C3C\",\n", - " \"train_loss\": \"#F39C12\",\n", - " \"vram\": [\"#3498DB\", \"#9B59B6\", \"#1ABC9C\", \"#E67E22\"],\n", - " \"util\": [\"#2ECC71\", \"#E74C3C\", \"#3498DB\", \"#F1C40F\"],\n", - " \"grid\": \"#ECECEC\",\n", - " \"title\": \"#2C3E50\",\n", - " \"subtitle\": \"#7F8C8D\",\n", - " \"spine\": \"#CCCCCC\",\n", - " \"tick\": \"#666666\",\n", - "}\n", - "\n", - "\n", - "def _style_axis(ax):\n", - " \"\"\"Apply shared cosmetic styling to a subplot axis.\"\"\"\n", - " ax.set_facecolor(\"white\")\n", - " ax.grid(True, alpha=0.4, color=_PALETTE[\"grid\"], linewidth=0.8)\n", - " for spine in (\"top\", \"right\"):\n", - " ax.spines[spine].set_visible(False)\n", - " ax.spines[\"left\"].set_color(_PALETTE[\"spine\"])\n", - " ax.spines[\"bottom\"].set_color(_PALETTE[\"spine\"])\n", - " ax.tick_params(colors=_PALETTE[\"tick\"], labelsize=9)\n", - "\n", - "\n", - "def _plot_line(ax, xs, ys, color, label, fill=True):\n", - " \"\"\"Plot a time series, gracefully skipping None values.\"\"\"\n", - " pts = [(x, y) for x, y in zip(xs, ys) if y is not None]\n", - " if not pts:\n", - " return\n", - " px, py = zip(*pts)\n", - " ax.plot(\n", - " px, py, color=color, linewidth=2.2,\n", - " marker=\"o\", markersize=4,\n", - " markerfacecolor=\"white\", markeredgewidth=1.8, markeredgecolor=color,\n", - " label=label, zorder=3,\n", - " )\n", - " if fill:\n", - " ax.fill_between(px, py, alpha=0.08, color=color)\n", - "\n", - "\n", - "def _plot_gpu_panel(ax, xs, history, colors, fallback_label):\n", - " \"\"\"Plot per-GPU time series with area fill.\"\"\"\n", - " if not history or not history[0]:\n", - " ax.text(\n", - " 0.5, 0.5, \"No GPU data\", transform=ax.transAxes,\n", - " ha=\"center\", va=\"center\", fontsize=11, color=\"#AAAAAA\",\n", - " )\n", - " return\n", - " n_gpus = max(len(snap) for snap in history)\n", - " for g in range(n_gpus):\n", - " vals = [snap[g] if g < len(snap) else 0 for snap in history]\n", - " c = colors[g % len(colors)]\n", - " label = f\"GPU {g}\" if n_gpus > 1 else fallback_label\n", - " ax.plot(xs[: len(vals)], vals, color=c, linewidth=2, label=label)\n", - " ax.fill_between(xs[: len(vals)], vals, alpha=0.08, color=c)\n", - " if n_gpus > 1:\n", - " ax.legend(fontsize=9, framealpha=0.9, edgecolor=\"#DDD\")\n", - "\n", - "\n", - "def _draw_dashboard(\n", - " elapsed_mins, val_losses, train_losses,\n", - " vram_history, util_history,\n", - " job_name, status_str, step_str, elapsed_str,\n", - "):\n", - " \"\"\"Render a live 1x3 training dashboard.\"\"\"\n", - " fig, axes = plt.subplots(1, 3, figsize=(20, 5.5))\n", - " fig.patch.set_facecolor(\"#FAFBFC\")\n", - "\n", - " fig.suptitle(\n", - " job_name, fontsize=15, fontweight=\"bold\",\n", - " color=_PALETTE[\"title\"], y=1.10,\n", - " )\n", - " fig.text(\n", - " 0.5, 1.01,\n", - " f\"{status_str} | {step_str} | {elapsed_str}\",\n", - " ha=\"center\", fontsize=13, color=_PALETTE[\"subtitle\"],\n", - " )\n", - "\n", - " for ax in axes:\n", - " _style_axis(ax)\n", - "\n", - " # -- Panel 1: Loss curves --\n", - " _plot_line(axes[0], elapsed_mins, val_losses, _PALETTE[\"val_loss\"], \"Val Loss\", fill=True)\n", - " _plot_line(axes[0], elapsed_mins, train_losses, _PALETTE[\"train_loss\"], \"Train Loss\", fill=False)\n", - " axes[0].set_title(\"Train/Validation Loss\", fontsize=13, fontweight=\"bold\", color=_PALETTE[\"title\"], pad=12)\n", - " axes[0].set_xlabel(\"Time (min)\", fontsize=10, color=\"#666\")\n", - " axes[0].set_ylabel(\"Loss\", fontsize=10, color=\"#666\")\n", - " if any(v is not None for v in val_losses + train_losses):\n", - " axes[0].legend(fontsize=9, framealpha=0.9, edgecolor=\"#DDD\")\n", - "\n", - " # -- Panel 2: GPU VRAM usage --\n", - " _plot_gpu_panel(axes[1], elapsed_mins, vram_history, _PALETTE[\"vram\"], \"VRAM\")\n", - " axes[1].set_title(\"GPU VRAM Usage\", fontsize=13, fontweight=\"bold\", color=_PALETTE[\"title\"], pad=12)\n", - " axes[1].set_xlabel(\"Time (min)\", fontsize=10, color=\"#666\")\n", - " axes[1].set_ylabel(\"Usage (%)\", fontsize=10, color=\"#666\")\n", - " axes[1].set_ylim(-2, 105)\n", - "\n", - " # -- Panel 3: GPU utilization --\n", - " _plot_gpu_panel(axes[2], elapsed_mins, util_history, _PALETTE[\"util\"], \"Utilization\")\n", - " axes[2].set_title(\"GPU Utilization\", fontsize=13, fontweight=\"bold\", color=_PALETTE[\"title\"], pad=12)\n", - " axes[2].set_xlabel(\"Time (min)\", fontsize=10, color=\"#666\")\n", - " axes[2].set_ylabel(\"Utilization (%)\", fontsize=10, color=\"#666\")\n", - " axes[2].set_ylim(-2, 105)\n", - "\n", - " plt.tight_layout(rect=[0, 0, 1, 0.98])\n", - " plt.show()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "#### Monitor the Job Until Completion\n", - "\n", - "The cell below polls the job status every 10 seconds and renders a live dashboard with validation loss, GPU VRAM usage, and GPU utilization charts. The charts appear empty at first while the model and dataset download; training metrics and GPU activity populate after the finetuning step begins.\n", - "\n", - "> **Note:** This is additional code. You can also use the Weights & Biases or MLflow integrations." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "import time\n", - "from typing import cast\n", - "from IPython.display import clear_output\n", - "from nemo_platform.types.shared import PlatformJobStatusResponse\n", - "\n", - "# Timeout set to 30 minutes to accommodate typical LoRA training duration for this dataset size.\n", - "# Actual training time will vary based on hardware, model size, and dataset complexity.\n", - "TIMEOUT_SECONDS = 30 * 60 # 30 minutes\n", - "VAL_LOSS_KEY = \"val_loss\"\n", - "TRAIN_LOSS_KEY = \"loss\"\n", - "\n", - "# ---------------------------------------------------------------------------\n", - "# Job polling with live dashboard\n", - "# ---------------------------------------------------------------------------\n", - "\n", - "def wait_for_job(\n", - " workspace: str,\n", - " job_name: str,\n", - " timeout: int = TIMEOUT_SECONDS,\n", - " poll_interval: int = 10,\n", - " val_loss_key: str = VAL_LOSS_KEY,\n", - " train_loss_key: str = TRAIN_LOSS_KEY,\n", - ") -> PlatformJobStatusResponse:\n", - " \"\"\"\n", - " Poll job status until completed, failed, cancelled, or timeout.\n", - " Displays a live dashboard with loss curves and GPU metrics.\n", - "\n", - " Args:\n", - " workspace: The workspace where the job is running.\n", - " job_name: The name of the job to monitor.\n", - " timeout: Maximum time to wait in seconds (default: 30 minutes).\n", - " poll_interval: Time between status checks in seconds (default: 10).\n", - "\n", - " Returns:\n", - " The final job status response.\n", - " \"\"\"\n", - " start_time = time.time()\n", - "\n", - " # Time-series accumulators required for plotting\n", - " elapsed_mins: list[float] = []\n", - " val_losses: list[float | None] = []\n", - " train_losses: list[float | None] = []\n", - " vram_history: list[list[float]] = []\n", - " util_history: list[list[float]] = []\n", - "\n", - " while True:\n", - " elapsed = time.time() - start_time\n", - " elapsed_min = elapsed / 60\n", - "\n", - " # Check for timeout\n", - " if elapsed > timeout:\n", - " error_message = f\"Timeout reached after {elapsed_min:.1f} minutes\"\n", - " print(f\"\\n{error_message}\")\n", - " print(\"Job did not complete within the timeout period.\")\n", - " raise Exception(error_message)\n", - "\n", - " status = client.jobs.get_status(name=job_name, workspace=workspace)\n", - "\n", - " # -- Extract training progress from nested steps structure --\n", - " step: int | None = None\n", - " max_steps: int | None = None\n", - " training_phase: str | None = None\n", - " val_loss: float | None = None\n", - " train_loss: float | None = None\n", - " current_step_name: str | None = None\n", - " current_step_phase: str | None = None\n", - "\n", - " for job_step in status.steps or []:\n", - " # Track the current active step name and phase for progress display\n", - " if job_step.tasks:\n", - " task = job_step.tasks[0]\n", - " td = task.status_details or {}\n", - " phase = cast(str, td.get(\"phase\", \"\"))\n", - " # Update current step if it's active or pending (not completed)\n", - " if job_step.status in (\"active\", \"pending\"):\n", - " current_step_name = job_step.name\n", - " current_step_phase = phase or \"started\"\n", - "\n", - " if job_step.name == \"training\":\n", - " for task in job_step.tasks or []:\n", - " td = task.status_details or {}\n", - " step = cast(int, td[\"step\"]) if \"step\" in td else None\n", - " max_steps = cast(int, td[\"max_steps\"]) if \"max_steps\" in td else None\n", - " training_phase = cast(str, td[\"phase\"]) if \"phase\" in td else None\n", - " raw_val_loss = td.get(val_loss_key)\n", - " val_loss = float(raw_val_loss) if raw_val_loss is not None else None\n", - " raw_train_loss = td.get(train_loss_key)\n", - " train_loss = float(raw_train_loss) if raw_train_loss is not None else None\n", - " break\n", - " break\n", - "\n", - " if val_loss is None:\n", - " raw_val_loss = (status.status_details or {}).get(val_loss_key)\n", - " val_loss = float(raw_val_loss) if raw_val_loss is not None else None\n", - " if train_loss is None:\n", - " raw_train_loss = (status.status_details or {}).get(train_loss_key)\n", - " train_loss = float(raw_train_loss) if raw_train_loss is not None else None\n", - "\n", - " # -- Collect GPU snapshot --\n", - " vram_pcts, util_pcts = _get_gpu_snapshot()\n", - "\n", - " # -- Append to accumulators used for the plots --\n", - " elapsed_mins.append(elapsed_min)\n", - " val_losses.append(val_loss)\n", - " train_losses.append(train_loss)\n", - " vram_history.append(vram_pcts)\n", - " util_history.append(util_pcts)\n", - "\n", - " # -- Build status strings --\n", - " status_str = f\"Status: {status.status}\"\n", - " if step is not None and max_steps is not None:\n", - " pct = step / max_steps * 100\n", - " step_str = f\"Step {step}/{max_steps} ({pct:.0f}%)\"\n", - " if training_phase:\n", - " step_str += f\" - {training_phase}\"\n", - " else:\n", - " if current_step_name and current_step_phase:\n", - " step_str = f\"{current_step_name} - {current_step_phase}\"\n", - " elif current_step_name:\n", - " step_str = f\"{current_step_name}\"\n", - " else:\n", - " step_str = \"Waiting for training to start...\"\n", - " elapsed_str = f\"Elapsed: {elapsed_min:.1f} min\"\n", - "\n", - " # -- Redraw dashboard --\n", - " clear_output(wait=True)\n", - " _draw_dashboard(\n", - " elapsed_mins, val_losses, train_losses,\n", - " vram_history, util_history,\n", - " job_name, status_str, step_str, elapsed_str,\n", - " )\n", - "\n", - " # -- Check terminal conditions --\n", - " if status.status.lower() == \"completed\":\n", - " # Redraw dashboard one final time with \"completed\" status\n", - " status_str = f\"Status: {status.status}\"\n", - " if step is not None and max_steps is not None:\n", - " step_str = f\"Step {max_steps}/{max_steps} (100%)\"\n", - " clear_output(wait=True)\n", - " _draw_dashboard(\n", - " elapsed_mins, val_losses, train_losses,\n", - " vram_history, util_history,\n", - " job_name, status_str, step_str, elapsed_str,\n", - " )\n", - " print(f\"\\nJob completed in {elapsed_min:.1f} minutes ({elapsed:.0f}s)\")\n", - " return status\n", - " elif status.status.lower() in (\"failed\", \"cancelled\", \"error\"):\n", - " print(f\"\\nJob finished with status: {status.status}\")\n", - " print(f\"Total time elapsed: {elapsed_min:.1f} minutes ({elapsed:.0f}s)\")\n", - "\n", - " # Print error details from the job level\n", - " if status.error_details:\n", - " error_msg = status.error_details.get(\"message\", \"\")\n", - " if error_msg:\n", - " print(f\"\\nError: {error_msg}\")\n", - "\n", - " # Find and print error details from the failed step/task\n", - " for job_step in status.steps or []:\n", - " if job_step.status == \"error\":\n", - " print(f\"\\nFailed step: {job_step.name}\")\n", - " if job_step.error_details:\n", - " step_error = job_step.error_details.get(\"message\", \"\")\n", - " if step_error:\n", - " print(f\"Step error: {step_error}\")\n", - " # Get error_stack from the failed task\n", - " for task in job_step.tasks or []:\n", - " if task.status == \"error\" and hasattr(task, \"error_stack\") and task.error_stack:\n", - " print(f\"\\nError stack trace:\\n{task.error_stack}\")\n", - " elif task.status == \"error\" and task.error_details:\n", - " task_error = task.error_details.get(\"message\", \"\")\n", - " if task_error:\n", - " print(f\"Task error: {task_error}\")\n", - " break\n", - "\n", - " raise Exception(f\"Job finished with status: {status.status}\")\n", - "\n", - " time.sleep(poll_interval)\n", - "\n", - "\n", - "# Wait for the job to complete\n", - "job_with_sequence_packing_status = wait_for_job(\n", - " workspace=\"default\",\n", - " job_name=job_with_sequence_packing.job.name,\n", - " timeout=TIMEOUT_SECONDS,\n", - ")\n", - "\n", - "packed_val_loss = (job_with_sequence_packing_status.status_details or {}).get(\"val_loss\")\n", - "if packed_val_loss is not None:\n", - " print(f\"Validation loss: {float(packed_val_loss):.2f}\")\n", - "else:\n", - " print(\"Validation loss: not reported in job status\")\n" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### 7. Create LoRA Job without Sequence Packing\n", - "Create a second Automodel LoRA job with `batch.sequence_packing=False` for comparison." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "import uuid\n", - "from nemo_automodel_plugin.schema import AutomodelJobInput\n", - "\n", - "job_suffix = uuid.uuid4().hex[:4]\n", - "JOB_NAME = f\"no-packing-job-{job_suffix}\"\n", - "NO_PACK_OUTPUT_NAME = f\"no-packing-out-{job_suffix}\"\n", - "\n", - "spec = AutomodelJobInput(\n", - " model=f\"default/{base_model.name}\",\n", - " dataset={\"training\": f\"default/{DATASET_NAME}\"},\n", - " training={\n", - " \"training_type\": \"sft\",\n", - " \"finetuning_type\": \"lora\",\n", - " \"max_seq_length\": 4096,\n", - " },\n", - " schedule={\"epochs\": 1, \"val_check_interval\": 0.1},\n", - " batch={\n", - " \"global_batch_size\": 64,\n", - " \"micro_batch_size\": 1,\n", - " \"sequence_packing\": False,\n", - " },\n", - " optimizer={\"learning_rate\": 5e-5},\n", - " parallelism={\"num_gpus_per_node\": 1},\n", - " output={\"name\": NO_PACK_OUTPUT_NAME},\n", - ")\n", - "\n", - "job_without_sequence_packing = client.customization.automodel.jobs.create(\n", - " spec=spec, workspace=\"default\", name=JOB_NAME\n", - ")\n", - "\n", - "print(f\"Submitted job: {job_without_sequence_packing.job.name}\")\n", - "print(f\"Output adapter: {NO_PACK_OUTPUT_NAME}\")\n" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### 8. Track Finetuning Progress for Job without Sequence Packing" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Wait for the training step to complete\n", - "job_without_sequence_packing_status = wait_for_job(\n", - " workspace=\"default\",\n", - " job_name=job_without_sequence_packing.job.name,\n", - " timeout=TIMEOUT_SECONDS\n", - ")\n", - "\n", - "no_pack_val_loss = (job_without_sequence_packing_status.status_details or {}).get(\"val_loss\")\n", - "if no_pack_val_loss is not None:\n", - " print(f\"Validation loss: {float(no_pack_val_loss):.2f}\")\n", - "else:\n", - " print(\"Validation loss: not reported in job status\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### 9. Compare Results\n", - "- Time to complete training should be significantly lower for the job that used sequence packing.\n", - "- The expected validation loss for both jobs should be similar.\n", - "- Sequence packed version should have a higher GPU utilization and higher GPU Memory Allocation." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from nemo_platform.types.jobs import PlatformJobStep\n", - "from datetime import datetime\n", - "import pandas as pd\n", - "\n", - "STEP_NAME = \"training\"\n", - "\n", - "def get_elapsed_time(step: PlatformJobStep) -> float:\n", - " \"\"\"Calculate elapsed time in seconds from step's created_at to updated_at.\"\"\"\n", - " created_at = datetime.fromisoformat(step.created_at.replace(\"Z\", \"+00:00\"))\n", - " updated_at = datetime.fromisoformat(step.updated_at.replace(\"Z\", \"+00:00\"))\n", - " return (updated_at - created_at).total_seconds()\n", - "\n", - "step_with_sequence_packing = client.jobs.steps.retrieve(\n", - " name=STEP_NAME,\n", - " workspace=\"default\",\n", - " job=job_with_sequence_packing.job.name,\n", - ")\n", - "\n", - "step_without_sequence_packing = client.jobs.steps.retrieve(\n", - " name=STEP_NAME,\n", - " workspace=\"default\",\n", - " job=job_without_sequence_packing.job.name,\n", - ")\n", - "\n", - "time_to_complete_with_sequence_packing = get_elapsed_time(step_with_sequence_packing)\n", - "time_to_complete_without_sequence_packing = get_elapsed_time(step_without_sequence_packing)\n", - "\n", - "# Display results as a table\n", - "results_df = pd.DataFrame({\n", - " \"Seq Packing Enabled\": [True, False],\n", - " \"Val Loss\": [\n", - " (job_with_sequence_packing_status.status_details or {}).get(\"val_loss\"),\n", - " (job_without_sequence_packing_status.status_details or {}).get(\"val_loss\"),\n", - " ],\n", - " \"Training Step Time, sec\": [\n", - " time_to_complete_with_sequence_packing,\n", - " time_to_complete_without_sequence_packing\n", - " ]\n", - "})\n", - "\n", - "results_df.style.format({\"Val Loss\": \"{:.2f}\", \"Training Step Time, sec\": \"{:.0f}\"}).hide(axis='index')" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "#### Examples of Validation Loss\n", - "\n", - "The expected validation loss curves should match closely for both jobs.\n", - "![Validation loss comparison chart showing similar convergence patterns between sequence-packed and non-packed training runs over training steps](../_images/packed_vs_not_packed_val_loss.png)\n", - "\n", - "Sequence packed version should complete significantly faster.\n", - "![Runtime comparison chart demonstrating significantly reduced training time for sequence-packed job compared to non-packed baseline](../_images/runtime.png)\n", - "\n", - "#### GPU Utilization\n", - "Sequence packed version should have a higher GPU utilization.\n", - "![GPU utilization chart showing higher and more consistent GPU usage with sequence packing enabled throughout the training process](../_images/gpu_utilization.png)\n", - "\n", - "#### GPU Memory Allocation\n", - "Sequence packed version should have a higher GPU Memory Allocation.\n", - "![GPU memory allocation chart illustrating increased memory utilization efficiency with sequence packing enabled](../_images/gpu_memory.png)" - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": ".venv", - "language": "python", - "name": "python3" }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.11.14" - } - }, - "nbformat": 4, - "nbformat_minor": 2 -} + "nbformat": 4, + "nbformat_minor": 2 +} \ No newline at end of file diff --git a/docs/customizer/tutorials/optimize-throughput.mdx b/docs/customizer/tutorials/optimize-throughput.mdx index fa87e0fea4..c06bbf193f 100644 --- a/docs/customizer/tutorials/optimize-throughput.mdx +++ b/docs/customizer/tutorials/optimize-throughput.mdx @@ -5,8 +5,6 @@ description: "" [Run in Google Colab](https://colab.research.google.com/github/NVIDIA-NeMo/nemo-platform/blob/main/docs/customizer/tutorials/optimize-throughput.ipynb) -# Optimize for Tokens/GPU Throughput - ## About Learn how to use the NeMo Platform Customizer to create a [LoRA](/documentation/customizer-reference/customization-concepts#nemo-ms-about-concepts-customization) (Low-Rank Adaptation) customization job optimized for higher tokens/GPU throughput and lower runtime. @@ -23,6 +21,7 @@ Before starting this tutorial, ensure you have: 1. **Completed the [Quickstart](/documentation/get-started)** to install and deploy NeMo Platform locally 2. **Installed the Python SDK** (PyPI wrapper: `pip install "nemo-platform[all]"`; source checkout: run `make bootstrap` from the repository root) +3. **At least one GPU with CUDA 13+** ## Quick Start @@ -159,29 +158,29 @@ print(json.dumps([f.model_dump() for f in client.files.list(fileset=DATASET_NAME ### 3. Secrets Setup -If you plan to use NGC or HuggingFace models, you will need to configure authentication: +If you plan to use NGC or Hugging Face models, you will need to configure authentication: - **NGC models** (`ngc://` URIs): Requires NGC API key -- **HuggingFace models** (`hf://` URIs): Requires HF token for gated/private models +- **Hugging Face models** (`hf://` URIs): Requires HF token for gated/private models Configure these as secrets in your platform. Refer to [Managing Secrets](/documentation/get-started/core-concepts/manage-secrets) for detailed instructions. Get your credentials to access base models: - [NGC API Key](https://ngc.nvidia.com/) (Setup → Generate API Key) -- [HuggingFace Token](https://huggingface.co/settings/tokens) (Create token with Read access) +- [Hugging Face Token](https://huggingface.co/settings/tokens) (Create token with Read access) --- #### Quick Setup Example -This tutorial uses the [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main) model from HuggingFace. Ensure that you have sufficient permissions to download the model. If you cannot access the files on the [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main) Hugging Face page, request access. +This tutorial uses the [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main) model from Hugging Face. Ensure that you have sufficient permissions to download the model. If you cannot access the files on the [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main) Hugging Face page, request access. -**HuggingFace Authentication:** -- For gated models (Llama, Gemma), you must provide a HuggingFace token via the `token_secret` parameter -- Get your token from [HuggingFace Settings](https://huggingface.co/settings/tokens) (requires Read access) -- Accept the model's terms on the HuggingFace model page before using it. Example: [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main) +**Hugging Face Authentication:** +- For gated models (Llama, Gemma), you must provide a Hugging Face token via the `token_secret` parameter +- Get your token from [Hugging Face Settings](https://huggingface.co/settings/tokens) (requires Read access) +- Accept the model's terms on the Hugging Face model page before using it. Example: [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main) - For public models, you can omit the `token_secret` parameter when creating a fileset for the model in the next step. ```python @@ -206,7 +205,7 @@ def create_or_get_secret(name: str, value: str | None, label: str): return client.secrets.retrieve(name=name, workspace="default") -# Create HuggingFace token secret +# Create Hugging Face token secret hf_secret = create_or_get_secret("hf-token", HF_TOKEN, "HF_TOKEN") print("HF_TOKEN secret:") print(hf_secret.model_dump_json(indent=2)) @@ -218,7 +217,7 @@ print(hf_secret.model_dump_json(indent=2)) ### 4. Create Base Model FileSet -Create a fileset pointing to the [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main) model on HuggingFace. This step creates a pointer to the model on Hugging Face and does not download it. The model is downloaded at job creation time. +Create a fileset pointing to the [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main) model on Hugging Face. This step creates a pointer to the model on Hugging Face and does not download it. The model is downloaded at job creation time. Note: for public models, you can omit the `token_secret` parameter when creating a model fileset. @@ -229,13 +228,13 @@ from nemo_platform.types.files import HuggingfaceStorageConfigParam HF_REPO_ID = "meta-llama/Llama-3.2-1B-Instruct" MODEL_NAME = "llama-3-2-1b-base" -# Ensure you have a HuggingFace token secret created -# Create a fileset pointing to the desired HuggingFace model +# Ensure you have a Hugging Face token secret created +# Create a fileset pointing to the desired Hugging Face model try: base_model_fs = client.files.filesets.create( workspace="default", name=MODEL_NAME, - description="Llama 3.2 1B base model from HuggingFace", + description="Llama 3.2 1B base model from Hugging Face", storage=HuggingfaceStorageConfigParam( type="huggingface", # repo_id is the full model name from Hugging Face @@ -330,15 +329,15 @@ print(f"Output adapter: {PACK_OUTPUT_NAME}") ``` -### 6. Track Finetuning Progress +### 6. Track Fine-Tuning Progress A training job contains multiple steps: - Model and dataset downloading -- Finetuning where LoRA adapter weights are trained -- Creating a fileset entry for the finetuned model -- Finetuned weights uploading +- Fine-tuning where LoRA adapter weights are trained +- Creating a fileset entry for the fine-tuned model +- Fine-tuned weights uploading -The elapsed time printed below reflects progress of the entire job. We compare the time taken by the finetuning step for both jobs in the last section of this tutorial. +The elapsed time printed below reflects progress of the entire job. We compare the time taken by the fine-tuning step for both jobs in the last section of this tutorial. #### Define Helper Functions @@ -489,7 +488,7 @@ def _draw_dashboard( #### Monitor the Job Until Completion -The cell below polls the job status every 10 seconds and renders a live dashboard with validation loss, GPU VRAM usage, and GPU utilization charts. The charts appear empty at first while the model and dataset download; training metrics and GPU activity populate after the finetuning step begins. +The cell below polls the job status every 10 seconds and renders a live dashboard with validation loss, GPU VRAM usage, and GPU utilization charts. The charts appear empty at first while the model and dataset download; training metrics and GPU activity populate after the fine-tuning step begins. > **Note:** This is additional code. You can also use the Weights & Biases or MLflow integrations. @@ -503,7 +502,22 @@ from nemo_platform.types.shared import PlatformJobStatusResponse # Actual training time will vary based on hardware, model size, and dataset complexity. TIMEOUT_SECONDS = 30 * 60 # 30 minutes VAL_LOSS_KEY = "val_loss" -TRAIN_LOSS_KEY = "loss" +TRAIN_LOSS_KEY = "train_loss" + + +def get_training_metric( + status: PlatformJobStatusResponse, + metric_key: str, +) -> float | None: + """Return a metric reported by a task in the training step.""" + for job_step in status.steps or []: + if job_step.name == "training": + for task in job_step.tasks or []: + value = (task.status_details or {}).get(metric_key) + if value is not None: + return float(value) + return None + # --------------------------------------------------------------------------- # Job polling with live dashboard @@ -680,9 +694,9 @@ job_with_sequence_packing_status = wait_for_job( timeout=TIMEOUT_SECONDS, ) -packed_val_loss = (job_with_sequence_packing_status.status_details or {}).get("val_loss") +packed_val_loss = get_training_metric(job_with_sequence_packing_status, VAL_LOSS_KEY) if packed_val_loss is not None: - print(f"Validation loss: {float(packed_val_loss):.2f}") + print(f"Validation loss: {packed_val_loss:.2f}") else: print("Validation loss: not reported in job status") @@ -727,7 +741,7 @@ print(f"Output adapter: {NO_PACK_OUTPUT_NAME}") ``` -### 8. Track Finetuning Progress for Job without Sequence Packing +### 8. Track Fine-Tuning Progress for Job without Sequence Packing ```python # Wait for the training step to complete @@ -737,9 +751,9 @@ job_without_sequence_packing_status = wait_for_job( timeout=TIMEOUT_SECONDS ) -no_pack_val_loss = (job_without_sequence_packing_status.status_details or {}).get("val_loss") +no_pack_val_loss = get_training_metric(job_without_sequence_packing_status, VAL_LOSS_KEY) if no_pack_val_loss is not None: - print(f"Validation loss: {float(no_pack_val_loss):.2f}") + print(f"Validation loss: {no_pack_val_loss:.2f}") else: print("Validation loss: not reported in job status") ``` @@ -751,16 +765,15 @@ else: ```python from nemo_platform.types.jobs import PlatformJobStep -from datetime import datetime import pandas as pd STEP_NAME = "training" def get_elapsed_time(step: PlatformJobStep) -> float: """Calculate elapsed time in seconds from step's created_at to updated_at.""" - created_at = datetime.fromisoformat(step.created_at.replace("Z", "+00:00")) - updated_at = datetime.fromisoformat(step.updated_at.replace("Z", "+00:00")) - return (updated_at - created_at).total_seconds() + if step.created_at is None or step.updated_at is None: + raise ValueError("Training step timestamps are unavailable") + return (step.updated_at - step.created_at).total_seconds() step_with_sequence_packing = client.jobs.steps.retrieve( name=STEP_NAME, @@ -780,10 +793,7 @@ time_to_complete_without_sequence_packing = get_elapsed_time(step_without_sequen # Display results as a table results_df = pd.DataFrame({ "Seq Packing Enabled": [True, False], - "Val Loss": [ - (job_with_sequence_packing_status.status_details or {}).get("val_loss"), - (job_without_sequence_packing_status.status_details or {}).get("val_loss"), - ], + "Val Loss": [packed_val_loss, no_pack_val_loss], "Training Step Time, sec": [ time_to_complete_with_sequence_packing, time_to_complete_without_sequence_packing diff --git a/docs/customizer/tutorials/sft-customization-job.ipynb b/docs/customizer/tutorials/sft-customization-job.ipynb index d59af1f9a4..8a6416ebf6 100644 --- a/docs/customizer/tutorials/sft-customization-job.ipynb +++ b/docs/customizer/tutorials/sft-customization-job.ipynb @@ -29,7 +29,7 @@ "- ✅ Can fundamentally change model behavior\n", "- ✅ Best for significant domain shifts or specialized tasks\n", "- ❌ Requires substantial GPU resources (4-8x more than LoRA)\n", - "- ❌ Produces full model weights (~140GB for Llama 70B)\n", + "- ❌ Produces a full BF16 checkpoint (~140 GB for Llama 70B); peak job disk usage can reach approximately 3× the downloaded base checkpoint size\n", "- ❌ Longer training time\n", "\n", "**LoRA** trains only ~1% of weights by adding thin matrices to existing weights:\n", @@ -58,7 +58,8 @@ "Before starting this tutorial, ensure you have:\n", "\n", "1. **Completed the [Quickstart](../../get-started/quickstart.md)** to install and deploy NeMo Platform locally\n", - "2. **Installed the Python SDK** (PyPI wrapper: `pip install \"nemo-platform[all]\"`; source checkout: run `make bootstrap` from the repository root)" + "2. **Installed the Python SDK** (PyPI wrapper: `pip install \"nemo-platform[all]\"`; source checkout: run `make bootstrap` from the repository root)\n", + "3. **At least one GPU with CUDA 13+**" ] }, { @@ -306,29 +307,29 @@ "source": [ "### 4. Secrets Setup\n", "\n", - "If you plan to use NGC or HuggingFace models, you will need to configure authentication:\n", + "If you plan to use NGC or Hugging Face models, you will need to configure authentication:\n", "\n", "- **NGC models** (`ngc://` URIs): Requires NGC API key\n", - "- **HuggingFace models** (`hf://` URIs): Requires HF token for gated/private models\n", + "- **Hugging Face models** (`hf://` URIs): Requires HF token for gated/private models\n", "\n", "\n", "Configure these as secrets in your platform. Refer to [Managing Secrets](../../get-started/concepts/manage-secrets.md) for detailed instructions.\n", "\n", "Get your credentials to access base models:\n", "- [NGC API Key](https://ngc.nvidia.com/) (Setup → Generate API Key)\n", - "- [HuggingFace Token](https://huggingface.co/settings/tokens) (Create token with Read access)\n", + "- [Hugging Face Token](https://huggingface.co/settings/tokens) (Create token with Read access)\n", "\n", "\n", "---\n", "\n", "#### Quick Setup Example\n", "\n", - "In this tutorial we are going to work with [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main) model from HuggingFace. Ensure that you have sufficient permissions to download the model. If you cannot access the files on the [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main) Hugging Face page, request access\n", + "In this tutorial we are going to work with the [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main) model from Hugging Face. Ensure that you have sufficient permissions to download the model. If you cannot access the files on the [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main) Hugging Face page, request access.\n", "\n", - "**HuggingFace Authentication:**\n", - "- For gated models (Llama, Gemma), you must provide a HuggingFace token via the `token_secret` parameter\n", - "- Get your token from [HuggingFace Settings](https://huggingface.co/settings/tokens) (requires Read access)\n", - "- Accept the model's terms on the HuggingFace model page before using it. Example: [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main)\n", + "**Hugging Face Authentication:**\n", + "- For gated models (Llama, Gemma), you must provide a Hugging Face token via the `token_secret` parameter\n", + "- Get your token from [Hugging Face Settings](https://huggingface.co/settings/tokens) (requires Read access)\n", + "- Accept the model's terms on the Hugging Face model page before using it. Example: [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main)\n", "- For public models, you can omit the `token_secret` parameter when creating a fileset for model in the next step" ] }, @@ -357,7 +358,7 @@ " return client.secrets.retrieve(name=name, workspace=\"default\")\n", "\n", "\n", - "# Create HuggingFace token secret\n", + "# Create Hugging Face token secret\n", "hf_secret = create_or_get_secret(\"hf-token\", HF_TOKEN, \"HF_TOKEN\")\n", "print(\"HF_TOKEN secret:\")\n", "print(hf_secret.model_dump_json(indent=2))\n", @@ -375,7 +376,7 @@ "source": [ "### 5. Create Base Model FileSet and Model Entity\n", "\n", - "Create a fileset pointing to [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main) model in HuggingFace that we will train with SFT. Then create a Model Entity that references this fileset. Model downloading will take place at training time.\n", + "Create a fileset pointing to the [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main) model in Hugging Face that we will train with SFT. Then create a Model Entity that references this fileset. Model downloading will take place at training time.\n", "\n", "Note: for public models, you can omit the `token_secret` parameter when creating a model fileset." ] @@ -386,18 +387,18 @@ "source": [ "import time\n", "\n", - "# Create a fileset pointing to the desired HuggingFace model\n", + "# Create a fileset pointing to the desired Hugging Face model\n", "from nemo_platform.types.files import HuggingfaceStorageConfigParam\n", "\n", "HF_REPO_ID = \"meta-llama/Llama-3.2-1B-Instruct\"\n", "MODEL_NAME = \"llama-3-2-1b-base\"\n", "\n", - "# Ensure you have a HuggingFace token secret created\n", + "# Ensure you have a Hugging Face token secret created\n", "try:\n", " base_model_fs = client.files.filesets.create(\n", " workspace=\"default\",\n", " name=MODEL_NAME,\n", - " description=\"Llama 3.2 1B base model from HuggingFace\",\n", + " description=\"Llama 3.2 1B base model from Hugging Face\",\n", " storage=HuggingfaceStorageConfigParam(\n", " type=\"huggingface\",\n", " # repo_id is the full model name from Hugging Face\n", @@ -457,7 +458,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "### 6. Create SFT Finetuning Job\n", + "### 6. Create SFT Fine-Tuning Job\n", "Create a customization job to fine-tune all model weights using the **Automodel** backend and `AutomodelJobInput`." ] }, @@ -689,8 +690,7 @@ ")\n", "```\n", "\n", - "**Single-Node Constraint:** Model deployments are limited to a single node. The maximum `gpu` value depends on the total GPUs available on a single node in your cluster. Multi-node deployments are not supported.\n", - "" + "**Single-Node Constraint:** Model deployments are limited to a single node. The maximum `gpu` value depends on the total GPUs available on a single node in your cluster. Multi-node deployments are not supported.\n" ] }, { @@ -827,7 +827,7 @@ "\n", "**Job fails during model download:**\n", "- Verify authentication secrets are configured (refer to [Managing Secrets](../../get-started/concepts/manage-secrets.md))\n", - "- For gated HuggingFace models (Llama, Gemma), accept the license on the model page (for example, [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct))\n", + "- For gated Hugging Face models (Llama, Gemma), accept the license on the model page (for example, [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct))\n", "- Confirm the model fileset uses `token_secret=hf_secret.name` for gated models\n", "- Check `AutomodelJobInput` references use the `workspace/name` format: `model=f\"default/{MODEL_NAME}\"` and `dataset={\"training\": f\"default/{DATASET_NAME}\"}` (for example, `default/llama-3-2-1b-base`, `default/sft-dataset`)\n", "- Verify the model entity points at the fileset: `fileset=f\"default/{MODEL_NAME}\"`\n", @@ -885,4 +885,4 @@ }, "nbformat": 4, "nbformat_minor": 2 -} +} \ No newline at end of file diff --git a/docs/customizer/tutorials/sft-customization-job.mdx b/docs/customizer/tutorials/sft-customization-job.mdx index 044abd8e01..24a2038d26 100644 --- a/docs/customizer/tutorials/sft-customization-job.mdx +++ b/docs/customizer/tutorials/sft-customization-job.mdx @@ -5,8 +5,6 @@ description: "" [Run in Google Colab](https://colab.research.google.com/github/NVIDIA-NeMo/nemo-platform/blob/main/docs/customizer/tutorials/sft-customization-job.ipynb) -# Full SFT Customization - Learn how to fine-tune all model weights using supervised fine-tuning (SFT) to customize LLM behavior for your specific tasks. ## About @@ -27,7 +25,7 @@ Supervised Fine-Tuning (SFT) customizes model behavior, injects new knowledge, a - ✅ Can fundamentally change model behavior - ✅ Best for significant domain shifts or specialized tasks - ❌ Requires substantial GPU resources (4-8x more than LoRA) -- ❌ Produces full model weights (~140GB for Llama 70B) +- ❌ Produces a full BF16 checkpoint (~140 GB for Llama 70B); peak job disk usage can reach approximately 3× the downloaded base checkpoint size - ❌ Longer training time **LoRA** trains only ~1% of weights by adding thin matrices to existing weights: @@ -52,6 +50,7 @@ Before starting this tutorial, ensure you have: 1. **Completed the [Quickstart](/documentation/get-started)** to install and deploy NeMo Platform locally 2. **Installed the Python SDK** (PyPI wrapper: `pip install "nemo-platform[all]"`; source checkout: run `make bootstrap` from the repository root) +3. **At least one GPU with CUDA 13+** ## Quick Start @@ -224,29 +223,29 @@ print(json.dumps([f.model_dump() for f in client.files.list(fileset=DATASET_NAME ### 4. Secrets Setup -If you plan to use NGC or HuggingFace models, you will need to configure authentication: +If you plan to use NGC or Hugging Face models, you will need to configure authentication: - **NGC models** (`ngc://` URIs): Requires NGC API key -- **HuggingFace models** (`hf://` URIs): Requires HF token for gated/private models +- **Hugging Face models** (`hf://` URIs): Requires HF token for gated/private models Configure these as secrets in your platform. Refer to [Managing Secrets](/documentation/get-started/core-concepts/manage-secrets) for detailed instructions. Get your credentials to access base models: - [NGC API Key](https://ngc.nvidia.com/) (Setup → Generate API Key) -- [HuggingFace Token](https://huggingface.co/settings/tokens) (Create token with Read access) +- [Hugging Face Token](https://huggingface.co/settings/tokens) (Create token with Read access) --- #### Quick Setup Example -In this tutorial we are going to work with [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main) model from HuggingFace. Ensure that you have sufficient permissions to download the model. If you cannot access the files on the [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main) Hugging Face page, request access. +In this tutorial we are going to work with the [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main) model from Hugging Face. Ensure that you have sufficient permissions to download the model. If you cannot access the files on the [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main) Hugging Face page, request access. -**HuggingFace Authentication:** -- For gated models (Llama, Gemma), you must provide a HuggingFace token via the `token_secret` parameter -- Get your token from [HuggingFace Settings](https://huggingface.co/settings/tokens) (requires Read access) -- Accept the model's terms on the HuggingFace model page before using it. Example: [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main) +**Hugging Face Authentication:** +- For gated models (Llama, Gemma), you must provide a Hugging Face token via the `token_secret` parameter +- Get your token from [Hugging Face Settings](https://huggingface.co/settings/tokens) (requires Read access) +- Accept the model's terms on the Hugging Face model page before using it. Example: [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main) - For public models, you can omit the `token_secret` parameter when creating a fileset for model in the next step ```python @@ -271,7 +270,7 @@ def create_or_get_secret(name: str, value: str | None, label: str): return client.secrets.retrieve(name=name, workspace="default") -# Create HuggingFace token secret +# Create Hugging Face token secret hf_secret = create_or_get_secret("hf-token", HF_TOKEN, "HF_TOKEN") print("HF_TOKEN secret:") print(hf_secret.model_dump_json(indent=2)) @@ -283,25 +282,25 @@ print(hf_secret.model_dump_json(indent=2)) ### 5. Create Base Model FileSet and Model Entity -Create a fileset pointing to [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main) model in HuggingFace that we will train with SFT. Then create a Model Entity that references this fileset. Model downloading will take place at training time. +Create a fileset pointing to the [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main) model in Hugging Face that we will train with SFT. Then create a Model Entity that references this fileset. Model downloading will take place at training time. Note: for public models, you can omit the `token_secret` parameter when creating a model fileset. ```python import time -# Create a fileset pointing to the desired HuggingFace model +# Create a fileset pointing to the desired Hugging Face model from nemo_platform.types.files import HuggingfaceStorageConfigParam HF_REPO_ID = "meta-llama/Llama-3.2-1B-Instruct" MODEL_NAME = "llama-3-2-1b-base" -# Ensure you have a HuggingFace token secret created +# Ensure you have a Hugging Face token secret created try: base_model_fs = client.files.filesets.create( workspace="default", name=MODEL_NAME, - description="Llama 3.2 1B base model from HuggingFace", + description="Llama 3.2 1B base model from Hugging Face", storage=HuggingfaceStorageConfigParam( type="huggingface", # repo_id is the full model name from Hugging Face @@ -355,7 +354,7 @@ while not base_model.spec: print(f"ModelSpec populated: {base_model.spec}") ``` -### 6. Create SFT Finetuning Job +### 6. Create SFT Fine-Tuning Job Create a customization job to fine-tune all model weights using the **Automodel** backend and `AutomodelJobInput`. **GPU Requirements:** @@ -655,7 +654,7 @@ For detailed information on all available hyperparameters, recommended values, a **Job fails during model download:** - Verify authentication secrets are configured (refer to [Managing Secrets](/documentation/get-started/core-concepts/manage-secrets)) -- For gated HuggingFace models (Llama, Gemma), accept the license on the model page (for example, [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct)) +- For gated Hugging Face models (Llama, Gemma), accept the license on the model page (for example, [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct)) - Confirm the model fileset uses `token_secret=hf_secret.name` for gated models - Check `AutomodelJobInput` references use the `workspace/name` format: `model=f"default/{MODEL_NAME}"` and `dataset={"training": f"default/{DATASET_NAME}"}` (for example, `default/llama-3-2-1b-base`, `default/sft-dataset`) - Verify the model entity points at the fileset: `fileset=f"default/{MODEL_NAME}"` diff --git a/docs/customizer/tutorials/understand-configurations-and-models.mdx b/docs/customizer/tutorials/understand-configurations-and-models.mdx index e9b9aa725d..1d36654a17 100644 --- a/docs/customizer/tutorials/understand-configurations-and-models.mdx +++ b/docs/customizer/tutorials/understand-configurations-and-models.mdx @@ -6,7 +6,7 @@ description: "" Learn the fundamentals of how NeMo Customizer works to make informed decisions about your fine-tuning projects. This tutorial covers how models are organized, how adapters attach to base models, training types and GPU requirements, and how to choose the right approach for your use case. -Understanding these basics will help you navigate the fine-tuning process more effectively and avoid common issues. If you're ready to start fine-tuning immediately, you can jump to [SFT Customization Job](sft-customization-job.ipynb) after completing this tutorial. +Understanding these basics will help you navigate the fine-tuning process more effectively and avoid common issues. If you're ready to start fine-tuning immediately, you can jump to [SFT Customization Job](/documentation/customizer-reference/tutorials/sft-customization-job) after completing this tutorial. @@ -42,7 +42,7 @@ An **Adapter** is a set of parameter-efficient fine-tuning weights (like LoRA) t - Are **nested within** the parent Model Entity - Are **enabled** for inference by default post training - Have their own FileSet for storing the adapter weights -- Track metadata like finetuning type, rank, and alpha values +- Track metadata like fine-tuning type, rank, and alpha values ### What is a FileSet? @@ -72,7 +72,7 @@ flowchart LR **1. Create a FileSet for your base model** -Upload your model checkpoint files (from HuggingFace, NGC, or local storage) to a FileSet: +Upload your model checkpoint files (from Hugging Face, NGC, or local storage) to a FileSet: ```python import os @@ -84,7 +84,7 @@ client = NeMoPlatform( workspace="default", ) -# Create a FileSet from HuggingFace +# Create a FileSet from Hugging Face fileset = client.files.filesets.create( workspace="default", name="llama-3-2-1b", @@ -93,7 +93,7 @@ fileset = client.files.filesets.create( type="huggingface", repo_id="meta-llama/Llama-3.2-1B-Instruct", repo_type="model", - token_secret="my-hf-token", # Secret containing HuggingFace token + token_secret="my-hf-token", # Secret containing Hugging Face token ), ) ``` @@ -240,12 +240,12 @@ model = client.models.adapters.create( ### Storage Requirements -Customization jobs consume disk space on the platform's shared persistent volume for model files, finetuning checkpoints, and the final output artifact. Required space depends on the training type: +Customization jobs consume disk space on the platform's shared persistent volume for model files, fine-tuning checkpoints, and the final output artifact. Required space depends on the training type: | Training Type | Approximate Disk Usage | Notes | |---------------|----------------------|-------| -| LoRA | ~1.5× base model size | Stores base model + small adapter weights | -| Full SFT | ~3× base model size | Stores base model + full checkpoint + output model | +| LoRA | ~1.5× downloaded base checkpoint size | Stores base model + small adapter weights | +| Full SFT | ~3× downloaded base checkpoint size | Stores base model + intermediate checkpoint + full output model | @@ -253,7 +253,7 @@ These estimates cover model weights only and do not include training dataset siz If the platform disk fills during a job, the job fails with an I/O error and the job service may return a ``500`` status when you retrieve logs. -Ensure your platform's shared persistent volume has at least **3× the base model size** +Ensure your platform's shared persistent volume has at least **3× the downloaded base checkpoint size** of free space before starting a full SFT job, or **1.5×** for LoRA jobs. @@ -261,15 +261,17 @@ For troubleshooting disk-related failures, see [customizer](/documentation/refer ### Parallelism Parameters Explained -Parallelism is configured via `training.parallelism`. These parameters control how training workloads are distributed across GPUs: +Parallelism is configured via the top-level Automodel `parallelism` block (for example, `parallelism={"num_gpus_per_node": 1}`). These parameters control how training workloads are distributed across GPUs: | Parameter | Description | Default | |-----------|-------------|---------| -| `tensor_parallel_size` | Number of GPUs to distribute each layer's parameters across | 1 | -| `pipeline_parallel_size` | Number of GPUs to distribute layers across sequentially | 1 | -| `context_parallel_size` | Number of GPUs to distribute sequence context across | 1 | -| `sequence_parallel` | Enable sequence parallelism to distribute activation memory along the sequence dimension | `false` | -| `expert_parallel_size` | Number of GPUs to distribute MoE experts across (MoE models only) | 1 | +| `parallelism.num_nodes` | Number of training nodes | `1` | +| `parallelism.num_gpus_per_node` | GPUs per node | `1` | +| `parallelism.tensor_parallel_size` | Number of GPUs to distribute each layer's parameters across | `1` | +| `parallelism.pipeline_parallel_size` | Number of GPUs to distribute layers across sequentially | `1` | +| `parallelism.context_parallel_size` | Number of GPUs to distribute sequence context across | `1` | +| `parallelism.sequence_parallel` | Enable sequence parallelism to distribute activation memory along the sequence dimension | `false` | +| `parallelism.expert_parallel_size` | Number of GPUs to distribute MoE experts across (MoE models only) | `null` | `data_parallel_size` is automatically derived as `total_gpus / (TP × PP × CP)` and is not set directly. @@ -277,16 +279,18 @@ Parallelism is configured via `training.parallelism`. These parameters control h **Recommended parallelism for Experts (MoE) Models**: -The `expert_parallel_size` parameter is used to parallelize a Mixture of Experts (MoE) model's experts across GPUs. For non-MoE models, this parameter is ignored. A model's model card will indicate if it is a Mixture of Experts model and specifies its number of experts. +The `parallelism.expert_parallel_size` parameter parallelizes a Mixture of Experts (MoE) model's experts across GPUs. For non-MoE models, leave it unset (`null`). A model's model card indicates whether it is a Mixture of Experts model and how many experts it has. -The number of experts in the model must be divisible by `expert_parallel_size`. For example, if a model has 8 experts, setting `expert_parallel_size=4` results in each GPU processing 2 experts. +When you set `expert_parallel_size`: +- The number of experts in the model must be divisible by `expert_parallel_size`. For example, if a model has 8 experts, `expert_parallel_size=4` gives each GPU 2 experts. +- `(data_parallel_size × context_parallel_size)` must be divisible by `expert_parallel_size`. +- When `expert_parallel_size > 1`, `tensor_parallel_size` must be `1`. -Also, the value of `expert_parallel_size` must evenly divide the derived `data_parallel_size`, which is automatically calculated as `data_parallel_size = total GPUs / (tensor_parallel_size × pipeline_parallel_size × context_parallel_size)`. - -For example, with 8 total GPUs, `tensor_parallel_size=2`, and `pipeline_parallel_size=1`: +For example, with 8 total GPUs, `tensor_parallel_size=2`, `pipeline_parallel_size=1`, and `context_parallel_size=1`: - Derived `data_parallel_size = 8 / (2 × 1 × 1) = 4` -- Valid `expert_parallel_size` values: `1`, `2`, or `4` (must evenly divide 4) -- Invalid `expert_parallel_size` value: `3` (does not evenly divide 4) +- `data_parallel_size × context_parallel_size = 4` +- Valid `expert_parallel_size` values: `1`, `2`, or `4` +- Invalid `expert_parallel_size` value: `3` (does not divide 4) ### Resource Allocation Rules @@ -359,7 +363,7 @@ flowchart TD | **Llama Models** | General-purpose language models excellent for instruction following, conversation, and text generation tasks | `llama-3.1-8b-instruct`, `llama-3.2-1b-instruct` | | **Llama Nemotron Models** | NVIDIA's specialized variants optimized for specific use cases with enhanced reasoning capabilities | Various Nano and Super variants | | **Phi Models** | Microsoft's efficient models designed for strong reasoning with optimized deployment characteristics | Phi model family configurations | -| **GPT-OSS Models** | Open-source GPT-based models supporting Full SFT customization workflows | Various GPT-OSS configurations | +| **GPT-OSS Models** | Open-weight reasoning models with tested Full SFT and LoRA configurations | `openai/gpt-oss-20b` | ### Specialized Models @@ -370,12 +374,12 @@ flowchart TD ### Importing Custom Models -You can import any HuggingFace-compatible model: +You can import a Hugging Face checkpoint into a FileSet and register it as a Model Entity: ```python from nemo_platform.types.files import HuggingfaceStorageConfigParam -# Create FileSet from HuggingFace +# Create FileSet from Hugging Face fileset = client.files.filesets.create( workspace="default", name="my-custom-model", @@ -393,7 +397,13 @@ model = client.models.create( ) ``` -For detailed guidance, see [Import HuggingFace Model](/documentation/customizer-reference/tutorials/import-hugging-face-models). +For detailed guidance, see [Import Hugging Face Model](/documentation/customizer-reference/tutorials/import-hugging-face-models). + + + +Importing a checkpoint does not guarantee that every training or deployment backend supports its architecture. In particular, Automodel LoRA does not support Conv1D-based architectures such as older GPT-2 variants. Confirm the model and fine-tuning regime in the [Tested Models](/documentation/customizer-reference/models/model-catalog) table, and review the import tutorial's known architecture limitations before submitting a job. + + --- @@ -408,19 +418,19 @@ Now that you understand how Model Entities and Adapters work, you're ready to pr Learn how to prepare your data for fine-tuning. - + Create a parameter-efficient LoRA adapter. - + Use full supervised fine-tuning for maximum performance. -Import and fine-tune private HuggingFace models. +Import and fine-tune private Hugging Face models. @@ -437,7 +447,7 @@ Import and fine-tune private HuggingFace models. ✅ **Full SFT training** creates a new Model Entity with full weights ✅ **Adapters are enabled by default** and automatically loaded by NIMs serving the base model ✅ **GPU requirements** vary significantly between LoRA and full fine-tuning -✅ **Custom HuggingFace models** can be imported via FileSet + Model Entity +✅ **Custom Hugging Face models** can be imported via FileSet + Model Entity ### Quick Reference Commands diff --git a/docs/fern/README.md b/docs/fern/README.md index bc339e89ee..115c1b2e3e 100644 --- a/docs/fern/README.md +++ b/docs/fern/README.md @@ -19,7 +19,7 @@ From the repo root (these wrap `cd docs/fern && npm run …`): ```bash make docs-deps # one-time: install docs/fern tooling (needed for MDX validation) make docs-login # one-time per machine: Fern CLI auth for the nvidia org -make docs-check # validate: fern check + MDX validation + gated-link check (what CI runs) +make docs-check # validate: fern check + MDX + NotebookViewer artifacts + gated links make docs # start local preview (prints a localhost URL) make docs-watch # start local preview plus a repo-level watcher for docs/** changes ``` @@ -93,7 +93,7 @@ Fern groups endpoints by their OpenAPI tag in the sidebar (Customizer, Evaluator Some features are not shipped yet and must be **fully excluded from the build** — not just hidden from the sidebar. Fern's `hidden: true` still builds and serves the page (reachable by direct URL and indexable), so it is **not** used for this. Instead, the gated pages are simply **left out of `versions/latest.yml`**: Fern only builds pages referenced in the navigation, so an omitted page is never built (it 404s and is not indexed). This matches the old MkDocs `hide_unready_docs` hook, which dropped the same files from the build. -The gated `.mdx` files stay in the repo so they remain maintained. The gated trees today are: `auth/`, `customizer/`, `safe-synthesizer/`, `evaluator/benchmarks/`, plus individual pages (`evaluator/metrics/{job-management,results}`, `run-inference/tutorials/deploy-models`, `example-applications/`, `troubleshooting/{cluster-setup,customizer}`, `get-started/quickstart`). +Gated `.mdx` files stay in the repo so they remain maintained. Do not keep a separate list of gated directories in contributor docs: publication state is derived from `versions/latest.yml`. A page listed there is published; an omitted page is gated. Inbound links from visible pages into gated pages are **delinked to plain text** (not rewritten URLs), since the target is not built — otherwise they would be broken links. @@ -114,7 +114,7 @@ One difference from the old MkDocs hook: that hook ran at build time and kept th | Workflow | Trigger | Purpose | | --- | --- | --- | -| `fern-docs-ci.yaml` | `pull_request` touching `docs/**` | `npm run check` (fern check + MDX + gated-link check) and `npm run broken-links` | +| `fern-docs-ci.yaml` | `pull_request` touching `docs/**` | `npm run check` (fern check + MDX + NotebookViewer artifacts + gated links) and `npm run broken-links` | | `fern-docs-preview-build.yaml` | `pull_request` touching `docs/**` | Upload PR `docs/` sources as an artifact (no secrets — fork-safe) | | `fern-docs-preview-comment.yaml` | successful preview build (`workflow_run`) | Build a Fern preview with `DOCS_FERN_TOKEN` and post/update the PR comment | | `publish-fern-docs.yaml` | push to `main` touching `docs/**`, `docs/v*` tag, or manual dispatch | Publish the Fern docs site | diff --git a/docs/fern/components/notebooks/distillation-customization-job.json b/docs/fern/components/notebooks/distillation-customization-job.json index 969cdfa323..c25043b5b4 100644 --- a/docs/fern/components/notebooks/distillation-customization-job.json +++ b/docs/fern/components/notebooks/distillation-customization-job.json @@ -7,8 +7,8 @@ }, { "type": "markdown", - "source": "## Prerequisites\n\nBefore starting this tutorial, ensure you have:\n\n1. **Completed the [Quickstart](../../get-started/quickstart.md)** to install and deploy NeMo Platform locally\n2. **Installed the Python SDK** (PyPI wrapper: `pip install \"nemo-platform[all]\"`; source checkout: run `make bootstrap` from the repository root)\n3. **Installed evaluation dependencies:**\n\n```sh\npip install evaluate rouge_score datasets\n```", - "source_html": "

Prerequisites

\n

Before starting this tutorial, ensure you have:

\n
    \n
  1. Completed the Quickstart to install and deploy NeMo Platform locally
    2. \n
  2. Installed the Python SDK (PyPI wrapper: pip install "nemo-platform[all]"; source checkout: run make bootstrap from the repository root)
    3. \n
  3. Installed evaluation dependencies:
    4. \n
\n
pip install evaluate rouge_score datasets\n
\n" + "source": "## Prerequisites\n\nBefore starting this tutorial, ensure you have:\n\n1. **Completed the [Quickstart](../../get-started/quickstart.md)** to install and deploy NeMo Platform locally\n2. **Installed the Python SDK** (PyPI wrapper: `pip install \"nemo-platform[all]\"`; source checkout: run `make bootstrap` from the repository root)\n3. **Installed evaluation dependencies:**\n\n```sh\npip install evaluate rouge_score datasets\n```\n\n4. **At least one GPU with CUDA 13+**", + "source_html": "

Prerequisites

\n

Before starting this tutorial, ensure you have:

\n
    \n
  1. Completed the Quickstart to install and deploy NeMo Platform locally
    2. \n
  2. Installed the Python SDK (PyPI wrapper: pip install "nemo-platform[all]"; source checkout: run make bootstrap from the repository root)
    3. \n
  3. Installed evaluation dependencies:
    4. \n
\n
pip install evaluate rouge_score datasets\n
\n
    \n
  1. At least one GPU with CUDA 13+
    2. \n
\n" }, { "type": "markdown", @@ -40,8 +40,8 @@ }, { "type": "markdown", - "source": "### 3. Secrets Setup\n\nIn this tutorial we use two Llama 3.2 Instruct models from HuggingFace:\n- **Teacher:** [meta-llama/Llama-3.2-3B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-3B-Instruct) (3B parameters)\n- **Student:** [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct) (1B parameters)\n\nBoth models share the same tokenizer/vocabulary (required for knowledge distillation) and include a chat template for deployment with `/chat/completions`.\n\n**HuggingFace Authentication:**\n- For gated models (Llama, Gemma), you must provide a HuggingFace token via the `token_secret` parameter\n- Get your token from [HuggingFace Settings](https://huggingface.co/settings/tokens) (requires Read access)\n- Accept the model's terms on the HuggingFace model page before using it:\n - [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct)\n - [meta-llama/Llama-3.2-3B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-3B-Instruct)", - "source_html": "

3. Secrets Setup

\n

In this tutorial we use two Llama 3.2 Instruct models from HuggingFace:

\n\n

Both models share the same tokenizer/vocabulary (required for knowledge distillation) and include a chat template for deployment with /chat/completions.

\n

HuggingFace Authentication:

\n\n" + "source": "### 3. Secrets Setup\n\nIn this tutorial we use two Llama 3.2 Instruct models from Hugging Face:\n- **Teacher:** [meta-llama/Llama-3.2-3B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-3B-Instruct) (3B parameters)\n- **Student:** [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct) (1B parameters)\n\nBoth models share the same tokenizer/vocabulary (required for knowledge distillation) and include a chat template for deployment with `/chat/completions`.\n\n**Hugging Face Authentication:**\n- For gated models (Llama, Gemma), you must provide a Hugging Face token via the `token_secret` parameter\n- Get your token from [Hugging Face Settings](https://huggingface.co/settings/tokens) (requires Read access)\n- Accept the model's terms on the Hugging Face model page before using it:\n - [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct)\n - [meta-llama/Llama-3.2-3B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-3B-Instruct)", + "source_html": "

3. Secrets Setup

\n

In this tutorial we use two Llama 3.2 Instruct models from Hugging Face:

\n\n

Both models share the same tokenizer/vocabulary (required for knowledge distillation) and include a chat template for deployment with /chat/completions.

\n

Hugging Face Authentication:

\n\n" }, { "type": "code", @@ -90,9 +90,9 @@ }, { "type": "code", - "source": "def wait_for_deployment(deployment_name: str, timeout_minutes: int = 30):\n \"\"\"Poll deployment until ready.\"\"\"\n start = time.time()\n timeout = timeout_minutes * 60\n while True:\n dep = client.inference.deployments.retrieve(name=deployment_name, workspace=\"default\")\n elapsed = time.time() - start\n clear_output(wait=True)\n print(f\"Deployment: {deployment_name}\")\n print(f\"Status: {dep.status}\")\n print(f\"Elapsed: {int(elapsed // 60)}m {int(elapsed % 60)}s\")\n\n if dep.status == \"READY\":\n print(\"\\nDeployment is ready!\")\n return dep\n if dep.status in (\"FAILED\", \"ERROR\", \"TERMINATED\", \"LOST\"):\n print(f\"\\nDeployment failed: {dep.status}\")\n return dep\n if elapsed > timeout:\n print(f\"\\nTimeout ({timeout_minutes}m). Check status manually.\")\n return dep\n time.sleep(15)\n\n\ndep_status = wait_for_deployment(BASELINE_DEPLOYMENT_NAME)\nassert dep_status.status == \"READY\"", + "source": "def wait_for_deployment(deployment_name: str, timeout_minutes: int = 30):\n \"\"\"Poll deployment until ready.\"\"\"\n start = time.time()\n timeout = timeout_minutes * 60\n while True:\n dep = client.inference.deployments.retrieve(name=deployment_name, workspace=\"default\")\n elapsed = time.time() - start\n clear_output(wait=True)\n print(f\"Deployment: {deployment_name}\")\n print(f\"Status: {dep.status}\")\n print(f\"Elapsed: {int(elapsed // 60)}m {int(elapsed % 60)}s\")\n\n if dep.status == \"READY\":\n print(\"\\nDeployment is ready!\")\n if not client.models.wait_for_gateway(deployment_name, workspace=\"default\", timeout=60):\n raise RuntimeError(\"Inference gateway did not become ready\")\n return dep\n if dep.status in (\"FAILED\", \"ERROR\", \"TERMINATED\", \"LOST\"):\n raise RuntimeError(f\"Deployment failed with status: {dep.status}\")\n if elapsed > timeout:\n raise TimeoutError(f\"Deployment timeout after {timeout_minutes} minutes\")\n time.sleep(15)\n\n\ndep_status = wait_for_deployment(BASELINE_DEPLOYMENT_NAME)\nassert dep_status.status == \"READY\"", "language": "python", - "source_html": "def wait_for_deployment(deployment_name: str, timeout_minutes: int = 30):\n """Poll deployment until ready."""\n start = time.time()\n timeout = timeout_minutes * 60\n while True:\n dep = client.inference.deployments.retrieve(name=deployment_name, workspace="default")\n elapsed = time.time() - start\n clear_output(wait=True)\n print(f"Deployment: {deployment_name}")\n print(f"Status: {dep.status}")\n print(f"Elapsed: {int(elapsed // 60)}m {int(elapsed % 60)}s")\n\n if dep.status == "READY":\n print("\\nDeployment is ready!")\n return dep\n if dep.status in ("FAILED", "ERROR", "TERMINATED", "LOST"):\n print(f"\\nDeployment failed: {dep.status}")\n return dep\n if elapsed > timeout:\n print(f"\\nTimeout ({timeout_minutes}m). Check status manually.")\n return dep\n time.sleep(15)\n\n\ndep_status = wait_for_deployment(BASELINE_DEPLOYMENT_NAME)\nassert dep_status.status == "READY"\n" + "source_html": "def wait_for_deployment(deployment_name: str, timeout_minutes: int = 30):\n """Poll deployment until ready."""\n start = time.time()\n timeout = timeout_minutes * 60\n while True:\n dep = client.inference.deployments.retrieve(name=deployment_name, workspace="default")\n elapsed = time.time() - start\n clear_output(wait=True)\n print(f"Deployment: {deployment_name}")\n print(f"Status: {dep.status}")\n print(f"Elapsed: {int(elapsed // 60)}m {int(elapsed % 60)}s")\n\n if dep.status == "READY":\n print("\\nDeployment is ready!")\n if not client.models.wait_for_gateway(deployment_name, workspace="default", timeout=60):\n raise RuntimeError("Inference gateway did not become ready")\n return dep\n if dep.status in ("FAILED", "ERROR", "TERMINATED", "LOST"):\n raise RuntimeError(f"Deployment failed with status: {dep.status}")\n if elapsed > timeout:\n raise TimeoutError(f"Deployment timeout after {timeout_minutes} minutes")\n time.sleep(15)\n\n\ndep_status = wait_for_deployment(BASELINE_DEPLOYMENT_NAME)\nassert dep_status.status == "READY"\n" }, { "type": "markdown", @@ -196,8 +196,8 @@ }, { "type": "markdown", - "source": "**Interpreting ROUGE Scores:**\n\n| Metric | Measures |\n|--------|----------|\n| **ROUGE-1** | Unigram overlap between prediction and reference |\n| **ROUGE-2** | Bigram overlap (captures phrase-level similarity) |\n| **ROUGE-L** | Longest common subsequence (captures sentence structure) |\n| **ROUGE-Lsum** | ROUGE-L computed over full summaries |\n\n**What to expect:**\n- The base student (1B, no training) provides a lower bound since it has not seen the task data\n- The distilled student (1B, KD) should significantly outperform the base student, demonstrating the knowledge transferred from the 3B teacher\n- If the distilled student scores are not much higher than the baseline, try increasing `distillation_temperature`, adjusting `distillation_ratio`, or training for more epochs\n\n---\n\n## Hyperparameters\n\nFor detailed information on all available hyperparameters, recommended values, and tuning guidance, refer to the [Hyperparameter Reference](../manage-customization-jobs/hyperparameters.md).\n\n---\n\n## Troubleshooting\n\n**Job fails during model download:**\n- Verify authentication secrets are configured (refer to [Managing Secrets](../../get-started/concepts/manage-secrets.md))\n- For gated HuggingFace models (Llama, Gemma), accept the license on the model page\n- Check both `model` (student) and `teacher_model` URNs are correct\n- Ensure both model entities exist: `client.models.retrieve(name=..., workspace=\"default\")`\n\n**Job fails with OOM (Out of Memory) error:**\n\nKD loads both models, so OOM is more likely than with SFT:\n1. **First try:** Use `teacher_precision=\"bf16\"` to reduce teacher memory\n2. **Still OOM:** Reduce `micro_batch_size` to 1\n3. **Still OOM:** Reduce `global_batch_size` and `max_seq_length`\n4. **Last resort:** Increase `num_gpus_per_node`\n\n**No chat template / `/chat/completions` fails:**\n- Use Instruct model variants (e.g., `Llama-3.2-1B-Instruct`) instead of base models (`Llama-3.2-1B`). Base models do not include a chat template in their tokenizer, so the output model will also lack one.\n\n**Distilled model quality is poor:**\n- Increase `distillation_temperature` (try 2.0–5.0) to transfer more nuanced knowledge\n- Adjust `distillation_ratio`—if dataset labels are high-quality, lower the ratio; if the teacher is strong, raise it\n- Increase `epochs` or `max_steps` for more training\n- Verify teacher and student share the same vocabulary\n\n**Vocabulary mismatch error:**\n- Teacher and student must use the same tokenizer. Use models from the same family (e.g., Llama 3.2 1B Instruct + Llama 3.2 3B Instruct)\n\n**Deployment fails:**\n- Verify output model exists: `client.models.retrieve(name=DISTILLED_STUDENT_NAME, workspace=\"default\")`\n- Check deployment logs: `client.inference.deployments.get_logs(name=deployment.name, workspace=\"default\")`\n- The distilled model has the same size as the student, so GPU requirements match the student model\n\n\n## Next Steps\n\n- [Monitor training metrics](fine-tune-metrics) in detail\n- [Evaluate your fine-tuned model](../../evaluator/index) using the Evaluator service\n- Learn about [LoRA customization](./lora-customization-job) for resource-efficient fine-tuning\n- Learn about [Full SFT](./sft-customization-job) for direct supervised fine-tuning", - "source_html": "

Interpreting ROUGE Scores:

\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n
MetricMeasures
ROUGE-1Unigram overlap between prediction and reference
ROUGE-2Bigram overlap (captures phrase-level similarity)
ROUGE-LLongest common subsequence (captures sentence structure)
ROUGE-LsumROUGE-L computed over full summaries
\n

What to expect:

\n
    \n
  • The base student (1B, no training) provides a lower bound since it has not seen the task data
    • \n
  • The distilled student (1B, KD) should significantly outperform the base student, demonstrating the knowledge transferred from the 3B teacher
    • \n
  • If the distilled student scores are not much higher than the baseline, try increasing distillation_temperature, adjusting distillation_ratio, or training for more epochs
    • \n
\n
\n

Hyperparameters

\n

For detailed information on all available hyperparameters, recommended values, and tuning guidance, refer to the Hyperparameter Reference.

\n
\n

Troubleshooting

\n

Job fails during model download:

\n
    \n
  • Verify authentication secrets are configured (refer to Managing Secrets)
    • \n
  • For gated HuggingFace models (Llama, Gemma), accept the license on the model page
    • \n
  • Check both model (student) and teacher_model URNs are correct
    • \n
  • Ensure both model entities exist: client.models.retrieve(name=..., workspace="default")
    • \n
\n

Job fails with OOM (Out of Memory) error:

\n

KD loads both models, so OOM is more likely than with SFT:

\n
    \n
  1. First try: Use teacher_precision="bf16" to reduce teacher memory
    2. \n
  2. Still OOM: Reduce micro_batch_size to 1
    3. \n
  3. Still OOM: Reduce global_batch_size and max_seq_length
    4. \n
  4. Last resort: Increase num_gpus_per_node
    5. \n
\n

No chat template / /chat/completions fails:

\n
    \n
  • Use Instruct model variants (e.g., Llama-3.2-1B-Instruct) instead of base models (Llama-3.2-1B). Base models do not include a chat template in their tokenizer, so the output model will also lack one.
    • \n
\n

Distilled model quality is poor:

\n
    \n
  • Increase distillation_temperature (try 2.0–5.0) to transfer more nuanced knowledge
    • \n
  • Adjust distillation_ratio—if dataset labels are high-quality, lower the ratio; if the teacher is strong, raise it
    • \n
  • Increase epochs or max_steps for more training
    • \n
  • Verify teacher and student share the same vocabulary
    • \n
\n

Vocabulary mismatch error:

\n
    \n
  • Teacher and student must use the same tokenizer. Use models from the same family (e.g., Llama 3.2 1B Instruct + Llama 3.2 3B Instruct)
    • \n
\n

Deployment fails:

\n
    \n
  • Verify output model exists: client.models.retrieve(name=DISTILLED_STUDENT_NAME, workspace="default")
    • \n
  • Check deployment logs: client.inference.deployments.get_logs(name=deployment.name, workspace="default")
    • \n
  • The distilled model has the same size as the student, so GPU requirements match the student model
    • \n
\n

Next Steps

\n\n" + "source": "**Interpreting ROUGE Scores:**\n\n| Metric | Measures |\n|--------|----------|\n| **ROUGE-1** | Unigram overlap between prediction and reference |\n| **ROUGE-2** | Bigram overlap (captures phrase-level similarity) |\n| **ROUGE-L** | Longest common subsequence (captures sentence structure) |\n| **ROUGE-Lsum** | ROUGE-L computed over full summaries |\n\n**What to expect:**\n- The base student (1B, no training) provides a lower bound since it has not seen the task data\n- The distilled student (1B, KD) should significantly outperform the base student, demonstrating the knowledge transferred from the 3B teacher\n- If the distilled student scores are not much higher than the baseline, try increasing `distillation_temperature`, adjusting `distillation_ratio`, or training for more epochs\n\n---\n\n## Hyperparameters\n\nFor detailed information on all available hyperparameters, recommended values, and tuning guidance, refer to the [Hyperparameter Reference](../manage-customization-jobs/hyperparameters.md).\n\n---\n\n## Troubleshooting\n\n**Job fails during model download:**\n- Verify authentication secrets are configured (refer to [Managing Secrets](../../get-started/concepts/manage-secrets.md))\n- For gated Hugging Face models (Llama, Gemma), accept the license on the model page\n- Check both `model` (student) and `teacher_model` URNs are correct\n- Ensure both model entities exist: `client.models.retrieve(name=..., workspace=\"default\")`\n\n**Job fails with OOM (Out of Memory) error:**\n\nKD loads both models, so OOM is more likely than with SFT:\n1. **First try:** Use `teacher_precision=\"bf16\"` to reduce teacher memory\n2. **Still OOM:** Reduce `micro_batch_size` to 1\n3. **Still OOM:** Reduce `global_batch_size` and `max_seq_length`\n4. **Last resort:** Increase `num_gpus_per_node`\n\n**No chat template / `/chat/completions` fails:**\n- Use Instruct model variants (e.g., `Llama-3.2-1B-Instruct`) instead of base models (`Llama-3.2-1B`). Base models do not include a chat template in their tokenizer, so the output model will also lack one.\n\n**Distilled model quality is poor:**\n- Increase `distillation_temperature` (try 2.0–5.0) to transfer more nuanced knowledge\n- Adjust `distillation_ratio`—if dataset labels are high-quality, lower the ratio; if the teacher is strong, raise it\n- Increase `epochs` or `max_steps` for more training\n- Verify teacher and student share the same vocabulary\n\n**Vocabulary mismatch error:**\n- Teacher and student must use the same tokenizer. Use models from the same family (e.g., Llama 3.2 1B Instruct + Llama 3.2 3B Instruct)\n\n**Deployment fails:**\n- Verify output model exists: `client.models.retrieve(name=DISTILLED_STUDENT_NAME, workspace=\"default\")`\n- Check deployment logs: `client.inference.deployments.get_logs(name=deployment.name, workspace=\"default\")`\n- The distilled model has the same size as the student, so GPU requirements match the student model\n\n\n## Next Steps\n\n- [Monitor training metrics](fine-tune-metrics) in detail\n- [Evaluate your fine-tuned model](../../evaluator/index) using the Evaluator service\n- Learn about [LoRA customization](./lora-customization-job) for resource-efficient fine-tuning\n- Learn about [Full SFT](./sft-customization-job) for direct supervised fine-tuning", + "source_html": "

Interpreting ROUGE Scores:

\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n
MetricMeasures
ROUGE-1Unigram overlap between prediction and reference
ROUGE-2Bigram overlap (captures phrase-level similarity)
ROUGE-LLongest common subsequence (captures sentence structure)
ROUGE-LsumROUGE-L computed over full summaries
\n

What to expect:

\n
    \n
  • The base student (1B, no training) provides a lower bound since it has not seen the task data
    • \n
  • The distilled student (1B, KD) should significantly outperform the base student, demonstrating the knowledge transferred from the 3B teacher
    • \n
  • If the distilled student scores are not much higher than the baseline, try increasing distillation_temperature, adjusting distillation_ratio, or training for more epochs
    • \n
\n
\n

Hyperparameters

\n

For detailed information on all available hyperparameters, recommended values, and tuning guidance, refer to the Hyperparameter Reference.

\n
\n

Troubleshooting

\n

Job fails during model download:

\n
    \n
  • Verify authentication secrets are configured (refer to Managing Secrets)
    • \n
  • For gated Hugging Face models (Llama, Gemma), accept the license on the model page
    • \n
  • Check both model (student) and teacher_model URNs are correct
    • \n
  • Ensure both model entities exist: client.models.retrieve(name=..., workspace="default")
    • \n
\n

Job fails with OOM (Out of Memory) error:

\n

KD loads both models, so OOM is more likely than with SFT:

\n
    \n
  1. First try: Use teacher_precision="bf16" to reduce teacher memory
    2. \n
  2. Still OOM: Reduce micro_batch_size to 1
    3. \n
  3. Still OOM: Reduce global_batch_size and max_seq_length
    4. \n
  4. Last resort: Increase num_gpus_per_node
    5. \n
\n

No chat template / /chat/completions fails:

\n
    \n
  • Use Instruct model variants (e.g., Llama-3.2-1B-Instruct) instead of base models (Llama-3.2-1B). Base models do not include a chat template in their tokenizer, so the output model will also lack one.
    • \n
\n

Distilled model quality is poor:

\n
    \n
  • Increase distillation_temperature (try 2.0–5.0) to transfer more nuanced knowledge
    • \n
  • Adjust distillation_ratio—if dataset labels are high-quality, lower the ratio; if the teacher is strong, raise it
    • \n
  • Increase epochs or max_steps for more training
    • \n
  • Verify teacher and student share the same vocabulary
    • \n
\n

Vocabulary mismatch error:

\n
    \n
  • Teacher and student must use the same tokenizer. Use models from the same family (e.g., Llama 3.2 1B Instruct + Llama 3.2 3B Instruct)
    • \n
\n

Deployment fails:

\n
    \n
  • Verify output model exists: client.models.retrieve(name=DISTILLED_STUDENT_NAME, workspace="default")
    • \n
  • Check deployment logs: client.inference.deployments.get_logs(name=deployment.name, workspace="default")
    • \n
  • The distilled model has the same size as the student, so GPU requirements match the student model
    • \n
\n

Next Steps

\n\n" } ] } \ No newline at end of file diff --git a/docs/fern/components/notebooks/distillation-customization-job.ts b/docs/fern/components/notebooks/distillation-customization-job.ts index 5190b821cd..d33d81b925 100644 --- a/docs/fern/components/notebooks/distillation-customization-job.ts +++ b/docs/fern/components/notebooks/distillation-customization-job.ts @@ -12,8 +12,8 @@ export default { cells: [ }, { "type": "markdown", - "source": "## Prerequisites\n\nBefore starting this tutorial, ensure you have:\n\n1. **Completed the [Quickstart](../../get-started/quickstart.md)** to install and deploy NeMo Platform locally\n2. **Installed the Python SDK** (PyPI wrapper: `pip install \"nemo-platform[all]\"`; source checkout: run `make bootstrap` from the repository root)\n3. **Installed evaluation dependencies:**\n\n```sh\npip install evaluate rouge_score datasets\n```", - "source_html": "

Prerequisites

\n

Before starting this tutorial, ensure you have:

\n
    \n
  1. Completed the Quickstart to install and deploy NeMo Platform locally
    2. \n
  2. Installed the Python SDK (PyPI wrapper: pip install "nemo-platform[all]"; source checkout: run make bootstrap from the repository root)
    3. \n
  3. Installed evaluation dependencies:
    4. \n
\n
pip install evaluate rouge_score datasets\n
\n" + "source": "## Prerequisites\n\nBefore starting this tutorial, ensure you have:\n\n1. **Completed the [Quickstart](../../get-started/quickstart.md)** to install and deploy NeMo Platform locally\n2. **Installed the Python SDK** (PyPI wrapper: `pip install \"nemo-platform[all]\"`; source checkout: run `make bootstrap` from the repository root)\n3. **Installed evaluation dependencies:**\n\n```sh\npip install evaluate rouge_score datasets\n```\n\n4. **At least one GPU with CUDA 13+**", + "source_html": "

Prerequisites

\n

Before starting this tutorial, ensure you have:

\n
    \n
  1. Completed the Quickstart to install and deploy NeMo Platform locally
    2. \n
  2. Installed the Python SDK (PyPI wrapper: pip install "nemo-platform[all]"; source checkout: run make bootstrap from the repository root)
    3. \n
  3. Installed evaluation dependencies:
    4. \n
\n
pip install evaluate rouge_score datasets\n
\n
    \n
  1. At least one GPU with CUDA 13+
    2. \n
\n" }, { "type": "markdown", @@ -45,8 +45,8 @@ export default { cells: [ }, { "type": "markdown", - "source": "### 3. Secrets Setup\n\nIn this tutorial we use two Llama 3.2 Instruct models from HuggingFace:\n- **Teacher:** [meta-llama/Llama-3.2-3B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-3B-Instruct) (3B parameters)\n- **Student:** [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct) (1B parameters)\n\nBoth models share the same tokenizer/vocabulary (required for knowledge distillation) and include a chat template for deployment with `/chat/completions`.\n\n**HuggingFace Authentication:**\n- For gated models (Llama, Gemma), you must provide a HuggingFace token via the `token_secret` parameter\n- Get your token from [HuggingFace Settings](https://huggingface.co/settings/tokens) (requires Read access)\n- Accept the model's terms on the HuggingFace model page before using it:\n - [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct)\n - [meta-llama/Llama-3.2-3B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-3B-Instruct)", - "source_html": "

3. Secrets Setup

\n

In this tutorial we use two Llama 3.2 Instruct models from HuggingFace:

\n\n

Both models share the same tokenizer/vocabulary (required for knowledge distillation) and include a chat template for deployment with /chat/completions.

\n

HuggingFace Authentication:

\n\n" + "source": "### 3. Secrets Setup\n\nIn this tutorial we use two Llama 3.2 Instruct models from Hugging Face:\n- **Teacher:** [meta-llama/Llama-3.2-3B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-3B-Instruct) (3B parameters)\n- **Student:** [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct) (1B parameters)\n\nBoth models share the same tokenizer/vocabulary (required for knowledge distillation) and include a chat template for deployment with `/chat/completions`.\n\n**Hugging Face Authentication:**\n- For gated models (Llama, Gemma), you must provide a Hugging Face token via the `token_secret` parameter\n- Get your token from [Hugging Face Settings](https://huggingface.co/settings/tokens) (requires Read access)\n- Accept the model's terms on the Hugging Face model page before using it:\n - [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct)\n - [meta-llama/Llama-3.2-3B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-3B-Instruct)", + "source_html": "

3. Secrets Setup

\n

In this tutorial we use two Llama 3.2 Instruct models from Hugging Face:

\n\n

Both models share the same tokenizer/vocabulary (required for knowledge distillation) and include a chat template for deployment with /chat/completions.

\n

Hugging Face Authentication:

\n\n" }, { "type": "code", @@ -95,9 +95,9 @@ export default { cells: [ }, { "type": "code", - "source": "def wait_for_deployment(deployment_name: str, timeout_minutes: int = 30):\n \"\"\"Poll deployment until ready.\"\"\"\n start = time.time()\n timeout = timeout_minutes * 60\n while True:\n dep = client.inference.deployments.retrieve(name=deployment_name, workspace=\"default\")\n elapsed = time.time() - start\n clear_output(wait=True)\n print(f\"Deployment: {deployment_name}\")\n print(f\"Status: {dep.status}\")\n print(f\"Elapsed: {int(elapsed // 60)}m {int(elapsed % 60)}s\")\n\n if dep.status == \"READY\":\n print(\"\\nDeployment is ready!\")\n return dep\n if dep.status in (\"FAILED\", \"ERROR\", \"TERMINATED\", \"LOST\"):\n print(f\"\\nDeployment failed: {dep.status}\")\n return dep\n if elapsed > timeout:\n print(f\"\\nTimeout ({timeout_minutes}m). Check status manually.\")\n return dep\n time.sleep(15)\n\n\ndep_status = wait_for_deployment(BASELINE_DEPLOYMENT_NAME)\nassert dep_status.status == \"READY\"", + "source": "def wait_for_deployment(deployment_name: str, timeout_minutes: int = 30):\n \"\"\"Poll deployment until ready.\"\"\"\n start = time.time()\n timeout = timeout_minutes * 60\n while True:\n dep = client.inference.deployments.retrieve(name=deployment_name, workspace=\"default\")\n elapsed = time.time() - start\n clear_output(wait=True)\n print(f\"Deployment: {deployment_name}\")\n print(f\"Status: {dep.status}\")\n print(f\"Elapsed: {int(elapsed // 60)}m {int(elapsed % 60)}s\")\n\n if dep.status == \"READY\":\n print(\"\\nDeployment is ready!\")\n if not client.models.wait_for_gateway(deployment_name, workspace=\"default\", timeout=60):\n raise RuntimeError(\"Inference gateway did not become ready\")\n return dep\n if dep.status in (\"FAILED\", \"ERROR\", \"TERMINATED\", \"LOST\"):\n raise RuntimeError(f\"Deployment failed with status: {dep.status}\")\n if elapsed > timeout:\n raise TimeoutError(f\"Deployment timeout after {timeout_minutes} minutes\")\n time.sleep(15)\n\n\ndep_status = wait_for_deployment(BASELINE_DEPLOYMENT_NAME)\nassert dep_status.status == \"READY\"", "language": "python", - "source_html": "def wait_for_deployment(deployment_name: str, timeout_minutes: int = 30):\n """Poll deployment until ready."""\n start = time.time()\n timeout = timeout_minutes * 60\n while True:\n dep = client.inference.deployments.retrieve(name=deployment_name, workspace="default")\n elapsed = time.time() - start\n clear_output(wait=True)\n print(f"Deployment: {deployment_name}")\n print(f"Status: {dep.status}")\n print(f"Elapsed: {int(elapsed // 60)}m {int(elapsed % 60)}s")\n\n if dep.status == "READY":\n print("\\nDeployment is ready!")\n return dep\n if dep.status in ("FAILED", "ERROR", "TERMINATED", "LOST"):\n print(f"\\nDeployment failed: {dep.status}")\n return dep\n if elapsed > timeout:\n print(f"\\nTimeout ({timeout_minutes}m). Check status manually.")\n return dep\n time.sleep(15)\n\n\ndep_status = wait_for_deployment(BASELINE_DEPLOYMENT_NAME)\nassert dep_status.status == "READY"\n" + "source_html": "def wait_for_deployment(deployment_name: str, timeout_minutes: int = 30):\n """Poll deployment until ready."""\n start = time.time()\n timeout = timeout_minutes * 60\n while True:\n dep = client.inference.deployments.retrieve(name=deployment_name, workspace="default")\n elapsed = time.time() - start\n clear_output(wait=True)\n print(f"Deployment: {deployment_name}")\n print(f"Status: {dep.status}")\n print(f"Elapsed: {int(elapsed // 60)}m {int(elapsed % 60)}s")\n\n if dep.status == "READY":\n print("\\nDeployment is ready!")\n if not client.models.wait_for_gateway(deployment_name, workspace="default", timeout=60):\n raise RuntimeError("Inference gateway did not become ready")\n return dep\n if dep.status in ("FAILED", "ERROR", "TERMINATED", "LOST"):\n raise RuntimeError(f"Deployment failed with status: {dep.status}")\n if elapsed > timeout:\n raise TimeoutError(f"Deployment timeout after {timeout_minutes} minutes")\n time.sleep(15)\n\n\ndep_status = wait_for_deployment(BASELINE_DEPLOYMENT_NAME)\nassert dep_status.status == "READY"\n" }, { "type": "markdown", @@ -201,7 +201,7 @@ export default { cells: [ }, { "type": "markdown", - "source": "**Interpreting ROUGE Scores:**\n\n| Metric | Measures |\n|--------|----------|\n| **ROUGE-1** | Unigram overlap between prediction and reference |\n| **ROUGE-2** | Bigram overlap (captures phrase-level similarity) |\n| **ROUGE-L** | Longest common subsequence (captures sentence structure) |\n| **ROUGE-Lsum** | ROUGE-L computed over full summaries |\n\n**What to expect:**\n- The base student (1B, no training) provides a lower bound since it has not seen the task data\n- The distilled student (1B, KD) should significantly outperform the base student, demonstrating the knowledge transferred from the 3B teacher\n- If the distilled student scores are not much higher than the baseline, try increasing `distillation_temperature`, adjusting `distillation_ratio`, or training for more epochs\n\n---\n\n## Hyperparameters\n\nFor detailed information on all available hyperparameters, recommended values, and tuning guidance, refer to the [Hyperparameter Reference](../manage-customization-jobs/hyperparameters.md).\n\n---\n\n## Troubleshooting\n\n**Job fails during model download:**\n- Verify authentication secrets are configured (refer to [Managing Secrets](../../get-started/concepts/manage-secrets.md))\n- For gated HuggingFace models (Llama, Gemma), accept the license on the model page\n- Check both `model` (student) and `teacher_model` URNs are correct\n- Ensure both model entities exist: `client.models.retrieve(name=..., workspace=\"default\")`\n\n**Job fails with OOM (Out of Memory) error:**\n\nKD loads both models, so OOM is more likely than with SFT:\n1. **First try:** Use `teacher_precision=\"bf16\"` to reduce teacher memory\n2. **Still OOM:** Reduce `micro_batch_size` to 1\n3. **Still OOM:** Reduce `global_batch_size` and `max_seq_length`\n4. **Last resort:** Increase `num_gpus_per_node`\n\n**No chat template / `/chat/completions` fails:**\n- Use Instruct model variants (e.g., `Llama-3.2-1B-Instruct`) instead of base models (`Llama-3.2-1B`). Base models do not include a chat template in their tokenizer, so the output model will also lack one.\n\n**Distilled model quality is poor:**\n- Increase `distillation_temperature` (try 2.0–5.0) to transfer more nuanced knowledge\n- Adjust `distillation_ratio`—if dataset labels are high-quality, lower the ratio; if the teacher is strong, raise it\n- Increase `epochs` or `max_steps` for more training\n- Verify teacher and student share the same vocabulary\n\n**Vocabulary mismatch error:**\n- Teacher and student must use the same tokenizer. Use models from the same family (e.g., Llama 3.2 1B Instruct + Llama 3.2 3B Instruct)\n\n**Deployment fails:**\n- Verify output model exists: `client.models.retrieve(name=DISTILLED_STUDENT_NAME, workspace=\"default\")`\n- Check deployment logs: `client.inference.deployments.get_logs(name=deployment.name, workspace=\"default\")`\n- The distilled model has the same size as the student, so GPU requirements match the student model\n\n\n## Next Steps\n\n- [Monitor training metrics](fine-tune-metrics) in detail\n- [Evaluate your fine-tuned model](../../evaluator/index) using the Evaluator service\n- Learn about [LoRA customization](./lora-customization-job) for resource-efficient fine-tuning\n- Learn about [Full SFT](./sft-customization-job) for direct supervised fine-tuning", - "source_html": "

Interpreting ROUGE Scores:

\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n
MetricMeasures
ROUGE-1Unigram overlap between prediction and reference
ROUGE-2Bigram overlap (captures phrase-level similarity)
ROUGE-LLongest common subsequence (captures sentence structure)
ROUGE-LsumROUGE-L computed over full summaries
\n

What to expect:

\n
    \n
  • The base student (1B, no training) provides a lower bound since it has not seen the task data
    • \n
  • The distilled student (1B, KD) should significantly outperform the base student, demonstrating the knowledge transferred from the 3B teacher
    • \n
  • If the distilled student scores are not much higher than the baseline, try increasing distillation_temperature, adjusting distillation_ratio, or training for more epochs
    • \n
\n
\n

Hyperparameters

\n

For detailed information on all available hyperparameters, recommended values, and tuning guidance, refer to the Hyperparameter Reference.

\n
\n

Troubleshooting

\n

Job fails during model download:

\n
    \n
  • Verify authentication secrets are configured (refer to Managing Secrets)
    • \n
  • For gated HuggingFace models (Llama, Gemma), accept the license on the model page
    • \n
  • Check both model (student) and teacher_model URNs are correct
    • \n
  • Ensure both model entities exist: client.models.retrieve(name=..., workspace="default")
    • \n
\n

Job fails with OOM (Out of Memory) error:

\n

KD loads both models, so OOM is more likely than with SFT:

\n
    \n
  1. First try: Use teacher_precision="bf16" to reduce teacher memory
    2. \n
  2. Still OOM: Reduce micro_batch_size to 1
    3. \n
  3. Still OOM: Reduce global_batch_size and max_seq_length
    4. \n
  4. Last resort: Increase num_gpus_per_node
    5. \n
\n

No chat template / /chat/completions fails:

\n
    \n
  • Use Instruct model variants (e.g., Llama-3.2-1B-Instruct) instead of base models (Llama-3.2-1B). Base models do not include a chat template in their tokenizer, so the output model will also lack one.
    • \n
\n

Distilled model quality is poor:

\n
    \n
  • Increase distillation_temperature (try 2.0–5.0) to transfer more nuanced knowledge
    • \n
  • Adjust distillation_ratio—if dataset labels are high-quality, lower the ratio; if the teacher is strong, raise it
    • \n
  • Increase epochs or max_steps for more training
    • \n
  • Verify teacher and student share the same vocabulary
    • \n
\n

Vocabulary mismatch error:

\n
    \n
  • Teacher and student must use the same tokenizer. Use models from the same family (e.g., Llama 3.2 1B Instruct + Llama 3.2 3B Instruct)
    • \n
\n

Deployment fails:

\n
    \n
  • Verify output model exists: client.models.retrieve(name=DISTILLED_STUDENT_NAME, workspace="default")
    • \n
  • Check deployment logs: client.inference.deployments.get_logs(name=deployment.name, workspace="default")
    • \n
  • The distilled model has the same size as the student, so GPU requirements match the student model
    • \n
\n

Next Steps

\n\n" + "source": "**Interpreting ROUGE Scores:**\n\n| Metric | Measures |\n|--------|----------|\n| **ROUGE-1** | Unigram overlap between prediction and reference |\n| **ROUGE-2** | Bigram overlap (captures phrase-level similarity) |\n| **ROUGE-L** | Longest common subsequence (captures sentence structure) |\n| **ROUGE-Lsum** | ROUGE-L computed over full summaries |\n\n**What to expect:**\n- The base student (1B, no training) provides a lower bound since it has not seen the task data\n- The distilled student (1B, KD) should significantly outperform the base student, demonstrating the knowledge transferred from the 3B teacher\n- If the distilled student scores are not much higher than the baseline, try increasing `distillation_temperature`, adjusting `distillation_ratio`, or training for more epochs\n\n---\n\n## Hyperparameters\n\nFor detailed information on all available hyperparameters, recommended values, and tuning guidance, refer to the [Hyperparameter Reference](../manage-customization-jobs/hyperparameters.md).\n\n---\n\n## Troubleshooting\n\n**Job fails during model download:**\n- Verify authentication secrets are configured (refer to [Managing Secrets](../../get-started/concepts/manage-secrets.md))\n- For gated Hugging Face models (Llama, Gemma), accept the license on the model page\n- Check both `model` (student) and `teacher_model` URNs are correct\n- Ensure both model entities exist: `client.models.retrieve(name=..., workspace=\"default\")`\n\n**Job fails with OOM (Out of Memory) error:**\n\nKD loads both models, so OOM is more likely than with SFT:\n1. **First try:** Use `teacher_precision=\"bf16\"` to reduce teacher memory\n2. **Still OOM:** Reduce `micro_batch_size` to 1\n3. **Still OOM:** Reduce `global_batch_size` and `max_seq_length`\n4. **Last resort:** Increase `num_gpus_per_node`\n\n**No chat template / `/chat/completions` fails:**\n- Use Instruct model variants (e.g., `Llama-3.2-1B-Instruct`) instead of base models (`Llama-3.2-1B`). Base models do not include a chat template in their tokenizer, so the output model will also lack one.\n\n**Distilled model quality is poor:**\n- Increase `distillation_temperature` (try 2.0–5.0) to transfer more nuanced knowledge\n- Adjust `distillation_ratio`—if dataset labels are high-quality, lower the ratio; if the teacher is strong, raise it\n- Increase `epochs` or `max_steps` for more training\n- Verify teacher and student share the same vocabulary\n\n**Vocabulary mismatch error:**\n- Teacher and student must use the same tokenizer. Use models from the same family (e.g., Llama 3.2 1B Instruct + Llama 3.2 3B Instruct)\n\n**Deployment fails:**\n- Verify output model exists: `client.models.retrieve(name=DISTILLED_STUDENT_NAME, workspace=\"default\")`\n- Check deployment logs: `client.inference.deployments.get_logs(name=deployment.name, workspace=\"default\")`\n- The distilled model has the same size as the student, so GPU requirements match the student model\n\n\n## Next Steps\n\n- [Monitor training metrics](fine-tune-metrics) in detail\n- [Evaluate your fine-tuned model](../../evaluator/index) using the Evaluator service\n- Learn about [LoRA customization](./lora-customization-job) for resource-efficient fine-tuning\n- Learn about [Full SFT](./sft-customization-job) for direct supervised fine-tuning", + "source_html": "

Interpreting ROUGE Scores:

\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n
MetricMeasures
ROUGE-1Unigram overlap between prediction and reference
ROUGE-2Bigram overlap (captures phrase-level similarity)
ROUGE-LLongest common subsequence (captures sentence structure)
ROUGE-LsumROUGE-L computed over full summaries
\n

What to expect:

\n
    \n
  • The base student (1B, no training) provides a lower bound since it has not seen the task data
    • \n
  • The distilled student (1B, KD) should significantly outperform the base student, demonstrating the knowledge transferred from the 3B teacher
    • \n
  • If the distilled student scores are not much higher than the baseline, try increasing distillation_temperature, adjusting distillation_ratio, or training for more epochs
    • \n
\n
\n

Hyperparameters

\n

For detailed information on all available hyperparameters, recommended values, and tuning guidance, refer to the Hyperparameter Reference.

\n
\n

Troubleshooting

\n

Job fails during model download:

\n
    \n
  • Verify authentication secrets are configured (refer to Managing Secrets)
    • \n
  • For gated Hugging Face models (Llama, Gemma), accept the license on the model page
    • \n
  • Check both model (student) and teacher_model URNs are correct
    • \n
  • Ensure both model entities exist: client.models.retrieve(name=..., workspace="default")
    • \n
\n

Job fails with OOM (Out of Memory) error:

\n

KD loads both models, so OOM is more likely than with SFT:

\n
    \n
  1. First try: Use teacher_precision="bf16" to reduce teacher memory
    2. \n
  2. Still OOM: Reduce micro_batch_size to 1
    3. \n
  3. Still OOM: Reduce global_batch_size and max_seq_length
    4. \n
  4. Last resort: Increase num_gpus_per_node
    5. \n
\n

No chat template / /chat/completions fails:

\n
    \n
  • Use Instruct model variants (e.g., Llama-3.2-1B-Instruct) instead of base models (Llama-3.2-1B). Base models do not include a chat template in their tokenizer, so the output model will also lack one.
    • \n
\n

Distilled model quality is poor:

\n
    \n
  • Increase distillation_temperature (try 2.0–5.0) to transfer more nuanced knowledge
    • \n
  • Adjust distillation_ratio—if dataset labels are high-quality, lower the ratio; if the teacher is strong, raise it
    • \n
  • Increase epochs or max_steps for more training
    • \n
  • Verify teacher and student share the same vocabulary
    • \n
\n

Vocabulary mismatch error:

\n
    \n
  • Teacher and student must use the same tokenizer. Use models from the same family (e.g., Llama 3.2 1B Instruct + Llama 3.2 3B Instruct)
    • \n
\n

Deployment fails:

\n
    \n
  • Verify output model exists: client.models.retrieve(name=DISTILLED_STUDENT_NAME, workspace="default")
    • \n
  • Check deployment logs: client.inference.deployments.get_logs(name=deployment.name, workspace="default")
    • \n
  • The distilled model has the same size as the student, so GPU requirements match the student model
    • \n
\n

Next Steps

\n\n" } ] }; diff --git a/docs/fern/components/notebooks/dpo-customization-job.json b/docs/fern/components/notebooks/dpo-customization-job.json new file mode 100644 index 0000000000..bf8d0970a4 --- /dev/null +++ b/docs/fern/components/notebooks/dpo-customization-job.json @@ -0,0 +1,155 @@ +{ + "cells": [ + { + "type": "markdown", + "source": "\n\n\n# DPO Model Customization Job\n\nLearn how to use the NeMo Platform to align a model with **DPO** (Direct Preference Optimization) on a preference dataset. For each prompt, DPO trains on a *chosen* (preferred) and a *rejected* response so the model prefers the chosen style — no separate reward model required.\n\nThis tutorial uses the `rl` customization backend (powered by [NVIDIA NeMo-RL](https://github.com/NVIDIA-NeMo/RL)), which runs DPO on a **Ray** cluster. Unlike the [SFT](/documentation/customizer-reference/tutorials/sft-customization-job) and [LoRA](/documentation/customizer-reference/tutorials/lora-customization-job) tutorials (Docker GPU jobs), `rl` requires a **Kubernetes-backed** NeMo Platform. DPO here is **full-weight** (no LoRA/adapter); the output is a full model entity.\n\n**Time to complete:** approximately 45-60 minutes. Job duration increases with model and dataset size.", + "source_html": "\n\n

DPO Model Customization Job

\n

Learn how to use the NeMo Platform to align a model with DPO (Direct Preference Optimization) on a preference dataset. For each prompt, DPO trains on a chosen (preferred) and a rejected response so the model prefers the chosen style — no separate reward model required.

\n

This tutorial uses the rl customization backend (powered by NVIDIA NeMo-RL), which runs DPO on a Ray cluster. Unlike the SFT and LoRA tutorials (Docker GPU jobs), rl requires a Kubernetes-backed NeMo Platform. DPO here is full-weight (no LoRA/adapter); the output is a full model entity.

\n

Time to complete: approximately 45-60 minutes. Job duration increases with model and dataset size.

\n" + }, + { + "type": "markdown", + "source": "## Prerequisites\n\nBefore starting this tutorial, ensure you have:\n\n1. **Completed the [Quickstart](/documentation/get-started)** to install the NeMo Platform and Python SDK.\n2. **Installed the Python SDK** (PyPI wrapper: `pip install \"nemo-platform[all]\"`; source checkout: run `make bootstrap` from the repository root).\n3. **Installed the `datasets` package**: `pip install datasets`.\n4. **A platform configured with `platform.runtime: kubernetes`.** The `rl` (DPO) backend provisions a Ray cluster and has **no local Docker fallback** — `submit` fails fast on a Docker-runtime platform. Multi-node jobs (`parallelism.num_nodes > 1`) additionally require the platform-side `NMP_RL_MULTINODE_SHARED_STORAGE_PATH`.\n5. **A Hugging Face token** with access to the gated base model (this tutorial uses `meta-llama/Llama-3.2-1B-Instruct`). Export it as `HF_TOKEN`.\n6. **At least one GPU with CUDA 13+** and a GPU execution profile (`nemo jobs list-execution-profiles`).", + "source_html": "

Prerequisites

\n

Before starting this tutorial, ensure you have:

\n
    \n
  1. Completed the Quickstart to install the NeMo Platform and Python SDK.
    2. \n
  2. Installed the Python SDK (PyPI wrapper: pip install "nemo-platform[all]"; source checkout: run make bootstrap from the repository root).
    3. \n
  3. Installed the datasets package: pip install datasets.
    4. \n
  4. A platform configured with platform.runtime: kubernetes. The rl (DPO) backend provisions a Ray cluster and has no local Docker fallback — submit fails fast on a Docker-runtime platform. Multi-node jobs (parallelism.num_nodes > 1) additionally require the platform-side NMP_RL_MULTINODE_SHARED_STORAGE_PATH.
    5. \n
  5. A Hugging Face token with access to the gated base model (this tutorial uses meta-llama/Llama-3.2-1B-Instruct). Export it as HF_TOKEN.
    6. \n
  6. At least one GPU with CUDA 13+ and a GPU execution profile (nemo jobs list-execution-profiles).
    7. \n
\n" + }, + { + "type": "markdown", + "source": "## Quick Start\n\n### 1. Initialize the SDK\n\nThe SDK needs your NeMo Platform server URL. By default `http://localhost:8080` is used; set `NMP_BASE_URL` to override:\n\n```sh\nexport NMP_BASE_URL=\n```", + "source_html": "

Quick Start

\n

1. Initialize the SDK

\n

The SDK needs your NeMo Platform server URL. By default http://localhost:8080 is used; set NMP_BASE_URL to override:

\n
export NMP_BASE_URL=<YOUR_NMP_BASE_URL>\n
\n" + }, + { + "type": "code", + "source": "import json\nimport os\nimport time\nimport uuid\nfrom pathlib import Path\nfrom nemo_platform import NeMoPlatform, ConflictError\nfrom nemo_platform.types.secrets import PlatformSecretResponse\nfrom nemo_platform.types.files import HuggingfaceStorageConfigParam\nfrom nemo_rl_plugin.schema import RlJobInput\n\n\ndef max_wait_time_checker(seconds: int, label: str = \"\"):\n \"\"\"Return a check() that raises TimeoutError once `seconds` have elapsed.\"\"\"\n start = time.time()\n\n def check():\n if time.time() - start > seconds:\n raise TimeoutError(f\"{label} took longer than {seconds} seconds\")\n\n return check\n\n\nNMP_BASE_URL = os.environ.get(\"NMP_BASE_URL\", \"http://localhost:8080\")\nsdk = NeMoPlatform(base_url=NMP_BASE_URL, workspace=\"default\")", + "language": "python", + "source_html": "import json\nimport os\nimport time\nimport uuid\nfrom pathlib import Path\nfrom nemo_platform import NeMoPlatform, ConflictError\nfrom nemo_platform.types.secrets import PlatformSecretResponse\nfrom nemo_platform.types.files import HuggingfaceStorageConfigParam\nfrom nemo_rl_plugin.schema import RlJobInput\n\n\ndef max_wait_time_checker(seconds: int, label: str = ""):\n """Return a check() that raises TimeoutError once `seconds` have elapsed."""\n start = time.time()\n\n def check():\n if time.time() - start > seconds:\n raise TimeoutError(f"{label} took longer than {seconds} seconds")\n\n return check\n\n\nNMP_BASE_URL = os.environ.get("NMP_BASE_URL", "http://localhost:8080")\nsdk = NeMoPlatform(base_url=NMP_BASE_URL, workspace="default")\n" + }, + { + "type": "markdown", + "source": "### 2. Prepare the Preference Dataset\n\nDPO trains on **preference data**. The `rl` backend takes a **single** dataset fileset that holds both `training.jsonl` and `validation.jsonl`, and auto-detects the row schema from the first line. Three preference formats are supported (see the platform's `BinaryPreferenceDatasetItemSchema` / `HelpSteer3DatasetItemSchema` / `Tulu3PreferenceDatasetItemSchema`):", + "source_html": "

2. Prepare the Preference Dataset

\n

DPO trains on preference data. The rl backend takes a single dataset fileset that holds both training.jsonl and validation.jsonl, and auto-detects the row schema from the first line. Three preference formats are supported (see the platform's BinaryPreferenceDatasetItemSchema / HelpSteer3DatasetItemSchema / Tulu3PreferenceDatasetItemSchema):

\n" + }, + { + "type": "markdown", + "source": "#### Binary Preference Format\n\nSimple `prompt` / `chosen` / `rejected` (the `prompt` may be a string or a list of chat messages):\n\n```json\n{\"prompt\": \"What is the capital of France?\", \"chosen\": \"The capital of France is Paris.\", \"rejected\": \"I'm not sure.\"}\n```", + "source_html": "

Binary Preference Format

\n

Simple prompt / chosen / rejected (the prompt may be a string or a list of chat messages):

\n
{"prompt": "What is the capital of France?", "chosen": "The capital of France is Paris.", "rejected": "I'm not sure."}\n
\n" + }, + { + "type": "markdown", + "source": "#### HelpSteer3 Format (used here)\n\nA conversation `context` (string or chat messages), two candidate `response1` / `response2`, and a signed `overall_preference` in -3..3 — **negative** means response 1 is preferred, **positive** means response 2, **0** is a tie. This is the **raw** schema of `nvidia/HelpSteer3`, so no conversion is needed:\n\n```json\n{\"context\": [{\"role\": \"user\", \"content\": \"Explain how to use git rebase\"}], \"response1\": \"...\", \"response2\": \"...\", \"overall_preference\": -2}\n```", + "source_html": "

HelpSteer3 Format (used here)

\n

A conversation context (string or chat messages), two candidate response1 / response2, and a signed overall_preference in -3..3 — negative means response 1 is preferred, positive means response 2, 0 is a tie. This is the raw schema of nvidia/HelpSteer3, so no conversion is needed:

\n
{"context": [{"role": "user", "content": "Explain how to use git rebase"}], "response1": "...", "response2": "...", "overall_preference": -2}\n
\n" + }, + { + "type": "markdown", + "source": "#### Tulu3 Preference Format\n\nFull chat conversations for both the chosen and rejected branches (each a list of messages ending with the assistant turn):\n\n```json\n{\"chosen\": [{\"role\": \"user\", \"content\": \"...\"}, {\"role\": \"assistant\", \"content\": \"preferred\"}], \"rejected\": [{\"role\": \"user\", \"content\": \"...\"}, {\"role\": \"assistant\", \"content\": \"dispreferred\"}]}\n```", + "source_html": "

Tulu3 Preference Format

\n

Full chat conversations for both the chosen and rejected branches (each a list of messages ending with the assistant turn):

\n
{"chosen": [{"role": "user", "content": "..."}, {"role": "assistant", "content": "preferred"}], "rejected": [{"role": "user", "content": "..."}, {"role": "assistant", "content": "dispreferred"}]}\n
\n" + }, + { + "type": "markdown", + "source": "#### Download nvidia/HelpSteer3\n\nWe use [nvidia/HelpSteer3](https://huggingface.co/datasets/nvidia/HelpSteer3) (the `preference` subset), NVIDIA's open preference dataset. It ships native `train` and `validation` splits and matches the HelpSteer3 schema above, so we upload the rows **as-is** — the platform's `HelpSteer3Dataset` loader handles the `overall_preference` semantics (including ties) at training time.", + "source_html": "

Download nvidia/HelpSteer3

\n

We use nvidia/HelpSteer3 (the preference subset), NVIDIA's open preference dataset. It ships native train and validation splits and matches the HelpSteer3 schema above, so we upload the rows as-is — the platform's HelpSteer3Dataset loader handles the overall_preference semantics (including ties) at training time.

\n" + }, + { + "type": "code", + "source": "from datasets import load_dataset, Dataset\n\nprint(\"Loading dataset nvidia/HelpSteer3 (preference subset)\")\nds = load_dataset(\"nvidia/HelpSteer3\", \"preference\")\n\n# Small subsets keep the tutorial fast; larger sets train better but take longer.\ntraining_size = 3000\nvalidation_size = 300\nDATASET_NAME = \"dpo-dataset\"\nDATASET_PATH = Path(\"dpo-dataset\").absolute()\nos.makedirs(DATASET_PATH, exist_ok=True)\n\ntrain_dataset = ds[\"train\"]\nvalidation_dataset = ds[\"validation\"]\nassert isinstance(train_dataset, Dataset) and isinstance(validation_dataset, Dataset)\n\n# Save raw HelpSteer3 rows directly — no conversion. The platform detects the\n# HelpSteer3 schema from the row keys (context / response1 / response2 / overall_preference).\ntrain_dataset.select(range(training_size)).to_json(f\"{DATASET_PATH}/training.jsonl\")\nvalidation_dataset.select(range(validation_size)).to_json(f\"{DATASET_PATH}/validation.jsonl\")\n\nprint(f\"Saved training.jsonl ({training_size} rows) and validation.jsonl ({validation_size} rows)\")\nwith open(f\"{DATASET_PATH}/training.jsonl\") as f:\n sample = json.loads(f.readline())\nprint(\"Sample keys:\", sorted(sample.keys()))\nprint(\"overall_preference:\", sample[\"overall_preference\"])", + "language": "python", + "source_html": "from datasets import load_dataset, Dataset\n\nprint("Loading dataset nvidia/HelpSteer3 (preference subset)")\nds = load_dataset("nvidia/HelpSteer3", "preference")\n\n# Small subsets keep the tutorial fast; larger sets train better but take longer.\ntraining_size = 3000\nvalidation_size = 300\nDATASET_NAME = "dpo-dataset"\nDATASET_PATH = Path("dpo-dataset").absolute()\nos.makedirs(DATASET_PATH, exist_ok=True)\n\ntrain_dataset = ds["train"]\nvalidation_dataset = ds["validation"]\nassert isinstance(train_dataset, Dataset) and isinstance(validation_dataset, Dataset)\n\n# Save raw HelpSteer3 rows directly — no conversion. The platform detects the\n# HelpSteer3 schema from the row keys (context / response1 / response2 / overall_preference).\ntrain_dataset.select(range(training_size)).to_json(f"{DATASET_PATH}/training.jsonl")\nvalidation_dataset.select(range(validation_size)).to_json(f"{DATASET_PATH}/validation.jsonl")\n\nprint(f"Saved training.jsonl ({training_size} rows) and validation.jsonl ({validation_size} rows)")\nwith open(f"{DATASET_PATH}/training.jsonl") as f:\n sample = json.loads(f.readline())\nprint("Sample keys:", sorted(sample.keys()))\nprint("overall_preference:", sample["overall_preference"])\n" + }, + { + "type": "markdown", + "source": "### 3. Create FileSet and Upload Preference Data\n\nUpload both JSONL files to a single FileSet so the DPO job can read them.", + "source_html": "

3. Create FileSet and Upload Preference Data

\n

Upload both JSONL files to a single FileSet so the DPO job can read them.

\n" + }, + { + "type": "code", + "source": "try:\n sdk.files.filesets.create(workspace=\"default\", name=DATASET_NAME, description=\"DPO preference data\")\n print(f\"Created fileset: {DATASET_NAME}\")\nexcept ConflictError:\n print(f\"Fileset '{DATASET_NAME}' already exists, continuing...\")\n\nsdk.files.upload(local_path=DATASET_PATH, remote_path=\"\", fileset=DATASET_NAME, workspace=\"default\")\n\nprint(\"Preference data:\")\nprint(json.dumps([f.model_dump() for f in sdk.files.list(fileset=DATASET_NAME, workspace=\"default\").data], indent=2, default=str))", + "language": "python", + "source_html": "try:\n sdk.files.filesets.create(workspace="default", name=DATASET_NAME, description="DPO preference data")\n print(f"Created fileset: {DATASET_NAME}")\nexcept ConflictError:\n print(f"Fileset '{DATASET_NAME}' already exists, continuing...")\n\nsdk.files.upload(local_path=DATASET_PATH, remote_path="", fileset=DATASET_NAME, workspace="default")\n\nprint("Preference data:")\nprint(json.dumps([f.model_dump() for f in sdk.files.list(fileset=DATASET_NAME, workspace="default").data], indent=2, default=str))\n" + }, + { + "type": "markdown", + "source": "### 4. Secrets Setup\n\nThe base model (`meta-llama/Llama-3.2-1B-Instruct`) is gated, so store your Hugging Face token as a platform secret named `hf-token` and reference it on the model fileset.", + "source_html": "

4. Secrets Setup

\n

The base model (meta-llama/Llama-3.2-1B-Instruct) is gated, so store your Hugging Face token as a platform secret named hf-token and reference it on the model fileset.

\n" + }, + { + "type": "code", + "source": "HF_TOKEN = os.getenv(\"HF_TOKEN\")\n\ndef create_or_get_secret(name: str, value: str | None, label: str) -> PlatformSecretResponse | None:\n if not value:\n print(f\"{label} is not set - skipping secret (gated model downloads will fail without it)\")\n return None\n try:\n secret = sdk.secrets.create(name=name, workspace=\"default\", value=value)\n print(f\"Created secret: {name}\")\n return secret\n except ConflictError:\n print(f\"Secret '{name}' already exists, continuing...\")\n return sdk.secrets.retrieve(name=name, workspace=\"default\")\n\n\nhf_secret = create_or_get_secret(\"hf-token\", HF_TOKEN, \"HF_TOKEN\")", + "language": "python", + "source_html": "HF_TOKEN = os.getenv("HF_TOKEN")\n\ndef create_or_get_secret(name: str, value: str | None, label: str) -> PlatformSecretResponse | None:\n if not value:\n print(f"{label} is not set - skipping secret (gated model downloads will fail without it)")\n return None\n try:\n secret = sdk.secrets.create(name=name, workspace="default", value=value)\n print(f"Created secret: {name}")\n return secret\n except ConflictError:\n print(f"Secret '{name}' already exists, continuing...")\n return sdk.secrets.retrieve(name=name, workspace="default")\n\n\nhf_secret = create_or_get_secret("hf-token", HF_TOKEN, "HF_TOKEN")\n" + }, + { + "type": "markdown", + "source": "### 5. Create Base Model FileSet and Model Entity\n\nDPO starts from an instruction-tuned base model. The model entity's spec is inferred asynchronously after creation.", + "source_html": "

5. Create Base Model FileSet and Model Entity

\n

DPO starts from an instruction-tuned base model. The model entity's spec is inferred asynchronously after creation.

\n" + }, + { + "type": "code", + "source": "HF_REPO_ID = \"meta-llama/Llama-3.2-1B-Instruct\"\nMODEL_NAME = \"llama-3-2-1b-instruct\"\n\nstorage = HuggingfaceStorageConfigParam(type=\"huggingface\", repo_id=HF_REPO_ID, repo_type=\"model\")\nif hf_secret:\n storage[\"token_secret\"] = hf_secret.name\n\ntry:\n base_model_fs = sdk.files.filesets.create(\n workspace=\"default\", name=MODEL_NAME, description=\"Llama 3.2 1B Instruct base model\", storage=storage\n )\n print(f\"Created base model fileset: {MODEL_NAME}\")\nexcept ConflictError:\n base_model_fs = sdk.files.filesets.retrieve(workspace=\"default\", name=MODEL_NAME)\n print(\"Base model fileset already exists.\")\n\ntry:\n base_model = sdk.models.create(workspace=\"default\", name=MODEL_NAME, fileset=f\"default/{MODEL_NAME}\")\nexcept ConflictError:\n base_model = sdk.models.retrieve(workspace=\"default\", name=MODEL_NAME)\n\nprint(f\"Base model fileset: fileset://default/{base_model.name}\")\n\n# Wait for the ModelSpec to be inferred from the checkpoint.\ncheck = max_wait_time_checker(600, \"Model spec\")\nwhile not base_model.spec:\n check()\n time.sleep(10)\n base_model = sdk.models.retrieve(workspace=\"default\", name=MODEL_NAME)\nprint(\"Model spec ready\")", + "language": "python", + "source_html": "HF_REPO_ID = "meta-llama/Llama-3.2-1B-Instruct"\nMODEL_NAME = "llama-3-2-1b-instruct"\n\nstorage = HuggingfaceStorageConfigParam(type="huggingface", repo_id=HF_REPO_ID, repo_type="model")\nif hf_secret:\n storage["token_secret"] = hf_secret.name\n\ntry:\n base_model_fs = sdk.files.filesets.create(\n workspace="default", name=MODEL_NAME, description="Llama 3.2 1B Instruct base model", storage=storage\n )\n print(f"Created base model fileset: {MODEL_NAME}")\nexcept ConflictError:\n base_model_fs = sdk.files.filesets.retrieve(workspace="default", name=MODEL_NAME)\n print("Base model fileset already exists.")\n\ntry:\n base_model = sdk.models.create(workspace="default", name=MODEL_NAME, fileset=f"default/{MODEL_NAME}")\nexcept ConflictError:\n base_model = sdk.models.retrieve(workspace="default", name=MODEL_NAME)\n\nprint(f"Base model fileset: fileset://default/{base_model.name}")\n\n# Wait for the ModelSpec to be inferred from the checkpoint.\ncheck = max_wait_time_checker(600, "Model spec")\nwhile not base_model.spec:\n check()\n time.sleep(10)\n base_model = sdk.models.retrieve(workspace="default", name=MODEL_NAME)\nprint("Model spec ready")\n" + }, + { + "type": "markdown", + "source": "### 6. Create the DPO Customization Job\n\nSubmit a DPO job to the `rl` backend with `RlJobInput`. Note the DPO-specific shape:\n\n- `model` is a string ref to the model entity; `dataset` is a **single** string ref to the preference fileset (holding both files).\n- The training method is `{\"type\": \"dpo\", ...}` — full-weight, no `finetuning_type`/LoRA.\n- `ref_policy_kl_penalty` is **β** (DPO paper): how strongly the policy stays tied to the reference model.\n- `rl` auto-generates the job id (`rl-`); read it back from the response.\n\nOther configurable knobs: `optimizer_type`, `adam_eps`, `activation_checkpointing`, `keep_top_k`, `val_at_end`, `preference_loss_weight`, `sft_loss_weight`. Run `nemo customization rl explain` for the live schema.", + "source_html": "

6. Create the DPO Customization Job

\n

Submit a DPO job to the rl backend with RlJobInput. Note the DPO-specific shape:

\n
    \n
  • model is a string ref to the model entity; dataset is a single string ref to the preference fileset (holding both files).
    • \n
  • The training method is {"type": "dpo", ...} — full-weight, no finetuning_type/LoRA.
    • \n
  • ref_policy_kl_penalty is β (DPO paper): how strongly the policy stays tied to the reference model.
    • \n
  • rl auto-generates the job id (rl-<hex>); read it back from the response.
    • \n
\n

Other configurable knobs: optimizer_type, adam_eps, activation_checkpointing, keep_top_k, val_at_end, preference_loss_weight, sft_loss_weight. Run nemo customization rl explain for the live schema.

\n" + }, + { + "type": "code", + "source": "job_suffix = uuid.uuid4().hex[:8]\nOUTPUT_NAME = f\"llama-3-2-1b-dpo-{job_suffix}\"\n\nspec = RlJobInput(\n model=f\"default/{base_model.name}\",\n dataset=f\"default/{DATASET_NAME}\",\n training={\n \"type\": \"dpo\",\n \"epochs\": 1,\n \"batch_size\": 16,\n \"micro_batch_size\": 1,\n \"learning_rate\": 5e-6,\n \"max_seq_length\": 4096,\n \"ref_policy_kl_penalty\": 0.1,\n \"parallelism\": {\n \"num_nodes\": 1,\n \"num_gpus_per_node\": 1,\n \"tensor_parallel_size\": 1,\n \"pipeline_parallel_size\": 1,\n },\n },\n output={\"name\": OUTPUT_NAME},\n)\n\n# `rl` auto-generates the job id (rl-); do not pass name=.\njob = sdk.customization.rl.jobs.create(spec=spec, workspace=\"default\")\nprint(f\"Job ID: {job.job.name}\")\nprint(f\"Output model: {OUTPUT_NAME}\")", + "language": "python", + "source_html": "job_suffix = uuid.uuid4().hex[:8]\nOUTPUT_NAME = f"llama-3-2-1b-dpo-{job_suffix}"\n\nspec = RlJobInput(\n model=f"default/{base_model.name}",\n dataset=f"default/{DATASET_NAME}",\n training={\n "type": "dpo",\n "epochs": 1,\n "batch_size": 16,\n "micro_batch_size": 1,\n "learning_rate": 5e-6,\n "max_seq_length": 4096,\n "ref_policy_kl_penalty": 0.1,\n "parallelism": {\n "num_nodes": 1,\n "num_gpus_per_node": 1,\n "tensor_parallel_size": 1,\n "pipeline_parallel_size": 1,\n },\n },\n output={"name": OUTPUT_NAME},\n)\n\n# `rl` auto-generates the job id (rl-<hex>); do not pass name=.\njob = sdk.customization.rl.jobs.create(spec=spec, workspace="default")\nprint(f"Job ID: {job.job.name}")\nprint(f"Output model: {OUTPUT_NAME}")\n" + }, + { + "type": "markdown", + "source": "### 7. Track Training Progress\n\nThe DPO job runs four steps: download -> **dpo-training** (Ray) -> upload -> model-entity. We poll the top-level job status and surface the training step's progress.", + "source_html": "

7. Track Training Progress

\n

The DPO job runs four steps: download -> dpo-training (Ray) -> upload -> model-entity. We poll the top-level job status and surface the training step's progress.

\n" + }, + { + "type": "code", + "source": "from IPython.display import clear_output\n\ncheck = max_wait_time_checker(7200, \"DPO job\")\nwhile True:\n check()\n status = sdk.jobs.get_status(name=job.job.name, workspace=\"default\")\n clear_output(wait=True)\n print(f\"Job Status: {status.status}\")\n\n step = max_steps = phase = None\n for job_step in status.steps or []:\n if job_step.name == \"dpo-training\":\n for task in job_step.tasks or []:\n d = task.status_details or {}\n step, max_steps, phase = d.get(\"step\"), d.get(\"max_steps\"), d.get(\"phase\")\n break\n break\n if step is not None and max_steps:\n print(f\"Training: Step {step}/{max_steps} ({100 * step / max_steps:.1f}%)\")\n if phase:\n print(f\"Phase: {phase}\")\n\n if status.status in (\"completed\", \"failed\", \"cancelled\", \"error\"):\n print(f\"\\nJob finished: {status.status}\")\n break\n time.sleep(15)\n\nassert status.status == \"completed\"", + "language": "python", + "source_html": "from IPython.display import clear_output\n\ncheck = max_wait_time_checker(7200, "DPO job")\nwhile True:\n check()\n status = sdk.jobs.get_status(name=job.job.name, workspace="default")\n clear_output(wait=True)\n print(f"Job Status: {status.status}")\n\n step = max_steps = phase = None\n for job_step in status.steps or []:\n if job_step.name == "dpo-training":\n for task in job_step.tasks or []:\n d = task.status_details or {}\n step, max_steps, phase = d.get("step"), d.get("max_steps"), d.get("phase")\n break\n break\n if step is not None and max_steps:\n print(f"Training: Step {step}/{max_steps} ({100 * step / max_steps:.1f}%)")\n if phase:\n print(f"Phase: {phase}")\n\n if status.status in ("completed", "failed", "cancelled", "error"):\n print(f"\\nJob finished: {status.status}")\n break\n time.sleep(15)\n\nassert status.status == "completed"\n" + }, + { + "type": "markdown", + "source": "**Interpreting DPO training metrics** (in `status_details.metrics`):\n\n- **`loss`** — the DPO loss; should trend down as the policy learns to separate chosen from rejected.\n- **Reward margin** (chosen minus rejected reward) — should trend **up**: the model increasingly prefers chosen responses.\n- **Validation `loss`** — watch for divergence from training loss (overfitting). Raise `ref_policy_kl_penalty` (β) or add `sft_loss_weight` if the policy drifts too far from the reference.", + "source_html": "

Interpreting DPO training metrics (in status_details.metrics):

\n
    \n
  • loss — the DPO loss; should trend down as the policy learns to separate chosen from rejected.
    • \n
  • Reward margin (chosen minus rejected reward) — should trend up: the model increasingly prefers chosen responses.
    • \n
  • Validation loss — watch for divergence from training loss (overfitting). Raise ref_policy_kl_penalty (β) or add sft_loss_weight if the policy drifts too far from the reference.
    • \n
\n" + }, + { + "type": "markdown", + "source": "### 8. Validate the Output Model\n\nDPO produces a **full-weight model entity** (not an adapter). Confirm it was registered.", + "source_html": "

8. Validate the Output Model

\n

DPO produces a full-weight model entity (not an adapter). Confirm it was registered.

\n" + }, + { + "type": "code", + "source": "model_entity = sdk.models.retrieve(workspace=\"default\", name=OUTPUT_NAME)\nprint(model_entity.model_dump_json(indent=2))", + "language": "python", + "source_html": "model_entity = sdk.models.retrieve(workspace="default", name=OUTPUT_NAME)\nprint(model_entity.model_dump_json(indent=2))\n" + }, + { + "type": "markdown", + "source": "### 9. Deploy and Evaluate (optional)\n\nThe DPO output is a full model, so it deploys like any full-weight checkpoint (see the [Full SFT](/documentation/customizer-reference/tutorials/sft-customization-job) tutorial for details). We deploy with vLLM and send a chat completion.", + "source_html": "

9. Deploy and Evaluate (optional)

\n

The DPO output is a full model, so it deploys like any full-weight checkpoint (see the Full SFT tutorial for details). We deploy with vLLM and send a chat completion.

\n" + }, + { + "type": "code", + "source": "deploy_suffix = uuid.uuid4().hex[:8]\nDEPLOYMENT_CONFIG_NAME = f\"dpo-deployment-cfg-{deploy_suffix}\"\nDEPLOYMENT_NAME = f\"dpo-deployment-{deploy_suffix}\"\n\ndeployment_config = sdk.inference.deployment_configs.create(\n workspace=\"default\",\n name=DEPLOYMENT_CONFIG_NAME,\n engine=\"vllm\",\n model_spec={\"model_namespace\": \"default\", \"model_name\": OUTPUT_NAME},\n executor_config={\"gpu\": 1, \"image_name\": \"vllm/vllm-openai\", \"image_tag\": \"v0.22.1\"},\n)\n\ndeployment = sdk.inference.deployments.create(\n workspace=\"default\", name=DEPLOYMENT_NAME, config=deployment_config.name\n)\nprint(f\"Deployment name: {deployment.name}\")", + "language": "python", + "source_html": "deploy_suffix = uuid.uuid4().hex[:8]\nDEPLOYMENT_CONFIG_NAME = f"dpo-deployment-cfg-{deploy_suffix}"\nDEPLOYMENT_NAME = f"dpo-deployment-{deploy_suffix}"\n\ndeployment_config = sdk.inference.deployment_configs.create(\n workspace="default",\n name=DEPLOYMENT_CONFIG_NAME,\n engine="vllm",\n model_spec={"model_namespace": "default", "model_name": OUTPUT_NAME},\n executor_config={"gpu": 1, "image_name": "vllm/vllm-openai", "image_tag": "v0.22.1"},\n)\n\ndeployment = sdk.inference.deployments.create(\n workspace="default", name=DEPLOYMENT_NAME, config=deployment_config.name\n)\nprint(f"Deployment name: {deployment.name}")\n" + }, + { + "type": "code", + "source": "check = max_wait_time_checker(1800, \"Deployment\")\nwhile True:\n check()\n deployment_status = sdk.inference.deployments.retrieve(name=deployment.name, workspace=\"default\")\n clear_output(wait=True)\n print(f\"Deployment status: {deployment_status.status}\")\n deployment_state = str(deployment_status.status).lower()\n if deployment_state in (\"ready\", \"running\"):\n if not sdk.models.wait_for_gateway(deployment.name, workspace=\"default\", timeout=60):\n raise RuntimeError(\"Inference gateway did not become ready\")\n break\n if deployment_state in (\"failed\", \"error\", \"terminated\", \"lost\"):\n raise RuntimeError(f\"Deployment failed with status: {deployment_status.status}\")\n time.sleep(15)", + "language": "python", + "source_html": "check = max_wait_time_checker(1800, "Deployment")\nwhile True:\n check()\n deployment_status = sdk.inference.deployments.retrieve(name=deployment.name, workspace="default")\n clear_output(wait=True)\n print(f"Deployment status: {deployment_status.status}")\n deployment_state = str(deployment_status.status).lower()\n if deployment_state in ("ready", "running"):\n if not sdk.models.wait_for_gateway(deployment.name, workspace="default", timeout=60):\n raise RuntimeError("Inference gateway did not become ready")\n break\n if deployment_state in ("failed", "error", "terminated", "lost"):\n raise RuntimeError(f"Deployment failed with status: {deployment_status.status}")\n time.sleep(15)\n" + }, + { + "type": "code", + "source": "messages = [\n {\"role\": \"system\", \"content\": \"You are a helpful assistant.\"},\n {\"role\": \"user\", \"content\": \"Write a short, friendly email to a colleague asking to reschedule our meeting to Thursday.\"},\n]\n\nresponse = sdk.inference.gateway.provider.post(\n \"v1/chat/completions\",\n name=deployment.name,\n workspace=\"default\",\n body={\"model\": f\"default/{OUTPUT_NAME}\", \"messages\": messages, \"temperature\": 0.7, \"max_tokens\": 256},\n)\nprint(\"Model output:\\n\")\nprint(response[\"choices\"][0][\"message\"][\"content\"])", + "language": "python", + "source_html": "messages = [\n {"role": "system", "content": "You are a helpful assistant."},\n {"role": "user", "content": "Write a short, friendly email to a colleague asking to reschedule our meeting to Thursday."},\n]\n\nresponse = sdk.inference.gateway.provider.post(\n "v1/chat/completions",\n name=deployment.name,\n workspace="default",\n body={"model": f"default/{OUTPUT_NAME}", "messages": messages, "temperature": 0.7, "max_tokens": 256},\n)\nprint("Model output:\\n")\nprint(response["choices"][0]["message"]["content"])\n" + }, + { + "type": "markdown", + "source": "## Conclusion\n\nYou aligned a base model with **DPO** on the NeMo Platform using the `rl` backend:\n\n- Uploaded a HelpSteer3 preference dataset **as-is** (the platform detects the schema natively).\n- Submitted a full-weight DPO job that ran on a Ray cluster via the Kubernetes executor.\n- Registered the output as a full model entity and (optionally) deployed it for inference.\n\n**Next steps:** tune the alignment strength with `ref_policy_kl_penalty` (β), add `sft_loss_weight` to anchor the policy to the chosen responses, enable `activation_checkpointing` for memory headroom, or scale up with `parallelism`. See the `nemo-customizer` skill's `references/hyperparameters.md` (section NeMo-RL (DPO)) for the full knob reference.", + "source_html": "

Conclusion

\n

You aligned a base model with DPO on the NeMo Platform using the rl backend:

\n
    \n
  • Uploaded a HelpSteer3 preference dataset as-is (the platform detects the schema natively).
    • \n
  • Submitted a full-weight DPO job that ran on a Ray cluster via the Kubernetes executor.
    • \n
  • Registered the output as a full model entity and (optionally) deployed it for inference.
    • \n
\n

Next steps: tune the alignment strength with ref_policy_kl_penalty (β), add sft_loss_weight to anchor the policy to the chosen responses, enable activation_checkpointing for memory headroom, or scale up with parallelism. See the nemo-customizer skill's references/hyperparameters.md (section NeMo-RL (DPO)) for the full knob reference.

\n" + } + ] +} \ No newline at end of file diff --git a/docs/fern/components/notebooks/dpo-customization-job.ts b/docs/fern/components/notebooks/dpo-customization-job.ts new file mode 100644 index 0000000000..ffbdb4dfb0 --- /dev/null +++ b/docs/fern/components/notebooks/dpo-customization-job.ts @@ -0,0 +1,159 @@ +/** + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Auto-generated by ipynb-to-fern-json.py - do not edit manually. + */ +export default { cells: [ + { + "type": "markdown", + "source": "\n\n\n# DPO Model Customization Job\n\nLearn how to use the NeMo Platform to align a model with **DPO** (Direct Preference Optimization) on a preference dataset. For each prompt, DPO trains on a *chosen* (preferred) and a *rejected* response so the model prefers the chosen style — no separate reward model required.\n\nThis tutorial uses the `rl` customization backend (powered by [NVIDIA NeMo-RL](https://github.com/NVIDIA-NeMo/RL)), which runs DPO on a **Ray** cluster. Unlike the [SFT](/documentation/customizer-reference/tutorials/sft-customization-job) and [LoRA](/documentation/customizer-reference/tutorials/lora-customization-job) tutorials (Docker GPU jobs), `rl` requires a **Kubernetes-backed** NeMo Platform. DPO here is **full-weight** (no LoRA/adapter); the output is a full model entity.\n\n**Time to complete:** approximately 45-60 minutes. Job duration increases with model and dataset size.", + "source_html": "\n\n

DPO Model Customization Job

\n

Learn how to use the NeMo Platform to align a model with DPO (Direct Preference Optimization) on a preference dataset. For each prompt, DPO trains on a chosen (preferred) and a rejected response so the model prefers the chosen style — no separate reward model required.

\n

This tutorial uses the rl customization backend (powered by NVIDIA NeMo-RL), which runs DPO on a Ray cluster. Unlike the SFT and LoRA tutorials (Docker GPU jobs), rl requires a Kubernetes-backed NeMo Platform. DPO here is full-weight (no LoRA/adapter); the output is a full model entity.

\n

Time to complete: approximately 45-60 minutes. Job duration increases with model and dataset size.

\n" + }, + { + "type": "markdown", + "source": "## Prerequisites\n\nBefore starting this tutorial, ensure you have:\n\n1. **Completed the [Quickstart](/documentation/get-started)** to install the NeMo Platform and Python SDK.\n2. **Installed the Python SDK** (PyPI wrapper: `pip install \"nemo-platform[all]\"`; source checkout: run `make bootstrap` from the repository root).\n3. **Installed the `datasets` package**: `pip install datasets`.\n4. **A platform configured with `platform.runtime: kubernetes`.** The `rl` (DPO) backend provisions a Ray cluster and has **no local Docker fallback** — `submit` fails fast on a Docker-runtime platform. Multi-node jobs (`parallelism.num_nodes > 1`) additionally require the platform-side `NMP_RL_MULTINODE_SHARED_STORAGE_PATH`.\n5. **A Hugging Face token** with access to the gated base model (this tutorial uses `meta-llama/Llama-3.2-1B-Instruct`). Export it as `HF_TOKEN`.\n6. **At least one GPU with CUDA 13+** and a GPU execution profile (`nemo jobs list-execution-profiles`).", + "source_html": "

Prerequisites

\n

Before starting this tutorial, ensure you have:

\n
    \n
  1. Completed the Quickstart to install the NeMo Platform and Python SDK.
    2. \n
  2. Installed the Python SDK (PyPI wrapper: pip install "nemo-platform[all]"; source checkout: run make bootstrap from the repository root).
    3. \n
  3. Installed the datasets package: pip install datasets.
    4. \n
  4. A platform configured with platform.runtime: kubernetes. The rl (DPO) backend provisions a Ray cluster and has no local Docker fallback — submit fails fast on a Docker-runtime platform. Multi-node jobs (parallelism.num_nodes > 1) additionally require the platform-side NMP_RL_MULTINODE_SHARED_STORAGE_PATH.
    5. \n
  5. A Hugging Face token with access to the gated base model (this tutorial uses meta-llama/Llama-3.2-1B-Instruct). Export it as HF_TOKEN.
    6. \n
  6. At least one GPU with CUDA 13+ and a GPU execution profile (nemo jobs list-execution-profiles).
    7. \n
\n" + }, + { + "type": "markdown", + "source": "## Quick Start\n\n### 1. Initialize the SDK\n\nThe SDK needs your NeMo Platform server URL. By default `http://localhost:8080` is used; set `NMP_BASE_URL` to override:\n\n```sh\nexport NMP_BASE_URL=\n```", + "source_html": "

Quick Start

\n

1. Initialize the SDK

\n

The SDK needs your NeMo Platform server URL. By default http://localhost:8080 is used; set NMP_BASE_URL to override:

\n
export NMP_BASE_URL=<YOUR_NMP_BASE_URL>\n
\n" + }, + { + "type": "code", + "source": "import json\nimport os\nimport time\nimport uuid\nfrom pathlib import Path\nfrom nemo_platform import NeMoPlatform, ConflictError\nfrom nemo_platform.types.secrets import PlatformSecretResponse\nfrom nemo_platform.types.files import HuggingfaceStorageConfigParam\nfrom nemo_rl_plugin.schema import RlJobInput\n\n\ndef max_wait_time_checker(seconds: int, label: str = \"\"):\n \"\"\"Return a check() that raises TimeoutError once `seconds` have elapsed.\"\"\"\n start = time.time()\n\n def check():\n if time.time() - start > seconds:\n raise TimeoutError(f\"{label} took longer than {seconds} seconds\")\n\n return check\n\n\nNMP_BASE_URL = os.environ.get(\"NMP_BASE_URL\", \"http://localhost:8080\")\nsdk = NeMoPlatform(base_url=NMP_BASE_URL, workspace=\"default\")", + "language": "python", + "source_html": "import json\nimport os\nimport time\nimport uuid\nfrom pathlib import Path\nfrom nemo_platform import NeMoPlatform, ConflictError\nfrom nemo_platform.types.secrets import PlatformSecretResponse\nfrom nemo_platform.types.files import HuggingfaceStorageConfigParam\nfrom nemo_rl_plugin.schema import RlJobInput\n\n\ndef max_wait_time_checker(seconds: int, label: str = ""):\n """Return a check() that raises TimeoutError once `seconds` have elapsed."""\n start = time.time()\n\n def check():\n if time.time() - start > seconds:\n raise TimeoutError(f"{label} took longer than {seconds} seconds")\n\n return check\n\n\nNMP_BASE_URL = os.environ.get("NMP_BASE_URL", "http://localhost:8080")\nsdk = NeMoPlatform(base_url=NMP_BASE_URL, workspace="default")\n" + }, + { + "type": "markdown", + "source": "### 2. Prepare the Preference Dataset\n\nDPO trains on **preference data**. The `rl` backend takes a **single** dataset fileset that holds both `training.jsonl` and `validation.jsonl`, and auto-detects the row schema from the first line. Three preference formats are supported (see the platform's `BinaryPreferenceDatasetItemSchema` / `HelpSteer3DatasetItemSchema` / `Tulu3PreferenceDatasetItemSchema`):", + "source_html": "

2. Prepare the Preference Dataset

\n

DPO trains on preference data. The rl backend takes a single dataset fileset that holds both training.jsonl and validation.jsonl, and auto-detects the row schema from the first line. Three preference formats are supported (see the platform's BinaryPreferenceDatasetItemSchema / HelpSteer3DatasetItemSchema / Tulu3PreferenceDatasetItemSchema):

\n" + }, + { + "type": "markdown", + "source": "#### Binary Preference Format\n\nSimple `prompt` / `chosen` / `rejected` (the `prompt` may be a string or a list of chat messages):\n\n```json\n{\"prompt\": \"What is the capital of France?\", \"chosen\": \"The capital of France is Paris.\", \"rejected\": \"I'm not sure.\"}\n```", + "source_html": "

Binary Preference Format

\n

Simple prompt / chosen / rejected (the prompt may be a string or a list of chat messages):

\n
{"prompt": "What is the capital of France?", "chosen": "The capital of France is Paris.", "rejected": "I'm not sure."}\n
\n" + }, + { + "type": "markdown", + "source": "#### HelpSteer3 Format (used here)\n\nA conversation `context` (string or chat messages), two candidate `response1` / `response2`, and a signed `overall_preference` in -3..3 — **negative** means response 1 is preferred, **positive** means response 2, **0** is a tie. This is the **raw** schema of `nvidia/HelpSteer3`, so no conversion is needed:\n\n```json\n{\"context\": [{\"role\": \"user\", \"content\": \"Explain how to use git rebase\"}], \"response1\": \"...\", \"response2\": \"...\", \"overall_preference\": -2}\n```", + "source_html": "

HelpSteer3 Format (used here)

\n

A conversation context (string or chat messages), two candidate response1 / response2, and a signed overall_preference in -3..3 — negative means response 1 is preferred, positive means response 2, 0 is a tie. This is the raw schema of nvidia/HelpSteer3, so no conversion is needed:

\n
{"context": [{"role": "user", "content": "Explain how to use git rebase"}], "response1": "...", "response2": "...", "overall_preference": -2}\n
\n" + }, + { + "type": "markdown", + "source": "#### Tulu3 Preference Format\n\nFull chat conversations for both the chosen and rejected branches (each a list of messages ending with the assistant turn):\n\n```json\n{\"chosen\": [{\"role\": \"user\", \"content\": \"...\"}, {\"role\": \"assistant\", \"content\": \"preferred\"}], \"rejected\": [{\"role\": \"user\", \"content\": \"...\"}, {\"role\": \"assistant\", \"content\": \"dispreferred\"}]}\n```", + "source_html": "

Tulu3 Preference Format

\n

Full chat conversations for both the chosen and rejected branches (each a list of messages ending with the assistant turn):

\n
{"chosen": [{"role": "user", "content": "..."}, {"role": "assistant", "content": "preferred"}], "rejected": [{"role": "user", "content": "..."}, {"role": "assistant", "content": "dispreferred"}]}\n
\n" + }, + { + "type": "markdown", + "source": "#### Download nvidia/HelpSteer3\n\nWe use [nvidia/HelpSteer3](https://huggingface.co/datasets/nvidia/HelpSteer3) (the `preference` subset), NVIDIA's open preference dataset. It ships native `train` and `validation` splits and matches the HelpSteer3 schema above, so we upload the rows **as-is** — the platform's `HelpSteer3Dataset` loader handles the `overall_preference` semantics (including ties) at training time.", + "source_html": "

Download nvidia/HelpSteer3

\n

We use nvidia/HelpSteer3 (the preference subset), NVIDIA's open preference dataset. It ships native train and validation splits and matches the HelpSteer3 schema above, so we upload the rows as-is — the platform's HelpSteer3Dataset loader handles the overall_preference semantics (including ties) at training time.

\n" + }, + { + "type": "code", + "source": "from datasets import load_dataset, Dataset\n\nprint(\"Loading dataset nvidia/HelpSteer3 (preference subset)\")\nds = load_dataset(\"nvidia/HelpSteer3\", \"preference\")\n\n# Small subsets keep the tutorial fast; larger sets train better but take longer.\ntraining_size = 3000\nvalidation_size = 300\nDATASET_NAME = \"dpo-dataset\"\nDATASET_PATH = Path(\"dpo-dataset\").absolute()\nos.makedirs(DATASET_PATH, exist_ok=True)\n\ntrain_dataset = ds[\"train\"]\nvalidation_dataset = ds[\"validation\"]\nassert isinstance(train_dataset, Dataset) and isinstance(validation_dataset, Dataset)\n\n# Save raw HelpSteer3 rows directly — no conversion. The platform detects the\n# HelpSteer3 schema from the row keys (context / response1 / response2 / overall_preference).\ntrain_dataset.select(range(training_size)).to_json(f\"{DATASET_PATH}/training.jsonl\")\nvalidation_dataset.select(range(validation_size)).to_json(f\"{DATASET_PATH}/validation.jsonl\")\n\nprint(f\"Saved training.jsonl ({training_size} rows) and validation.jsonl ({validation_size} rows)\")\nwith open(f\"{DATASET_PATH}/training.jsonl\") as f:\n sample = json.loads(f.readline())\nprint(\"Sample keys:\", sorted(sample.keys()))\nprint(\"overall_preference:\", sample[\"overall_preference\"])", + "language": "python", + "source_html": "from datasets import load_dataset, Dataset\n\nprint("Loading dataset nvidia/HelpSteer3 (preference subset)")\nds = load_dataset("nvidia/HelpSteer3", "preference")\n\n# Small subsets keep the tutorial fast; larger sets train better but take longer.\ntraining_size = 3000\nvalidation_size = 300\nDATASET_NAME = "dpo-dataset"\nDATASET_PATH = Path("dpo-dataset").absolute()\nos.makedirs(DATASET_PATH, exist_ok=True)\n\ntrain_dataset = ds["train"]\nvalidation_dataset = ds["validation"]\nassert isinstance(train_dataset, Dataset) and isinstance(validation_dataset, Dataset)\n\n# Save raw HelpSteer3 rows directly — no conversion. The platform detects the\n# HelpSteer3 schema from the row keys (context / response1 / response2 / overall_preference).\ntrain_dataset.select(range(training_size)).to_json(f"{DATASET_PATH}/training.jsonl")\nvalidation_dataset.select(range(validation_size)).to_json(f"{DATASET_PATH}/validation.jsonl")\n\nprint(f"Saved training.jsonl ({training_size} rows) and validation.jsonl ({validation_size} rows)")\nwith open(f"{DATASET_PATH}/training.jsonl") as f:\n sample = json.loads(f.readline())\nprint("Sample keys:", sorted(sample.keys()))\nprint("overall_preference:", sample["overall_preference"])\n" + }, + { + "type": "markdown", + "source": "### 3. Create FileSet and Upload Preference Data\n\nUpload both JSONL files to a single FileSet so the DPO job can read them.", + "source_html": "

3. Create FileSet and Upload Preference Data

\n

Upload both JSONL files to a single FileSet so the DPO job can read them.

\n" + }, + { + "type": "code", + "source": "try:\n sdk.files.filesets.create(workspace=\"default\", name=DATASET_NAME, description=\"DPO preference data\")\n print(f\"Created fileset: {DATASET_NAME}\")\nexcept ConflictError:\n print(f\"Fileset '{DATASET_NAME}' already exists, continuing...\")\n\nsdk.files.upload(local_path=DATASET_PATH, remote_path=\"\", fileset=DATASET_NAME, workspace=\"default\")\n\nprint(\"Preference data:\")\nprint(json.dumps([f.model_dump() for f in sdk.files.list(fileset=DATASET_NAME, workspace=\"default\").data], indent=2, default=str))", + "language": "python", + "source_html": "try:\n sdk.files.filesets.create(workspace="default", name=DATASET_NAME, description="DPO preference data")\n print(f"Created fileset: {DATASET_NAME}")\nexcept ConflictError:\n print(f"Fileset '{DATASET_NAME}' already exists, continuing...")\n\nsdk.files.upload(local_path=DATASET_PATH, remote_path="", fileset=DATASET_NAME, workspace="default")\n\nprint("Preference data:")\nprint(json.dumps([f.model_dump() for f in sdk.files.list(fileset=DATASET_NAME, workspace="default").data], indent=2, default=str))\n" + }, + { + "type": "markdown", + "source": "### 4. Secrets Setup\n\nThe base model (`meta-llama/Llama-3.2-1B-Instruct`) is gated, so store your Hugging Face token as a platform secret named `hf-token` and reference it on the model fileset.", + "source_html": "

4. Secrets Setup

\n

The base model (meta-llama/Llama-3.2-1B-Instruct) is gated, so store your Hugging Face token as a platform secret named hf-token and reference it on the model fileset.

\n" + }, + { + "type": "code", + "source": "HF_TOKEN = os.getenv(\"HF_TOKEN\")\n\ndef create_or_get_secret(name: str, value: str | None, label: str) -> PlatformSecretResponse | None:\n if not value:\n print(f\"{label} is not set - skipping secret (gated model downloads will fail without it)\")\n return None\n try:\n secret = sdk.secrets.create(name=name, workspace=\"default\", value=value)\n print(f\"Created secret: {name}\")\n return secret\n except ConflictError:\n print(f\"Secret '{name}' already exists, continuing...\")\n return sdk.secrets.retrieve(name=name, workspace=\"default\")\n\n\nhf_secret = create_or_get_secret(\"hf-token\", HF_TOKEN, \"HF_TOKEN\")", + "language": "python", + "source_html": "HF_TOKEN = os.getenv("HF_TOKEN")\n\ndef create_or_get_secret(name: str, value: str | None, label: str) -> PlatformSecretResponse | None:\n if not value:\n print(f"{label} is not set - skipping secret (gated model downloads will fail without it)")\n return None\n try:\n secret = sdk.secrets.create(name=name, workspace="default", value=value)\n print(f"Created secret: {name}")\n return secret\n except ConflictError:\n print(f"Secret '{name}' already exists, continuing...")\n return sdk.secrets.retrieve(name=name, workspace="default")\n\n\nhf_secret = create_or_get_secret("hf-token", HF_TOKEN, "HF_TOKEN")\n" + }, + { + "type": "markdown", + "source": "### 5. Create Base Model FileSet and Model Entity\n\nDPO starts from an instruction-tuned base model. The model entity's spec is inferred asynchronously after creation.", + "source_html": "

5. Create Base Model FileSet and Model Entity

\n

DPO starts from an instruction-tuned base model. The model entity's spec is inferred asynchronously after creation.

\n" + }, + { + "type": "code", + "source": "HF_REPO_ID = \"meta-llama/Llama-3.2-1B-Instruct\"\nMODEL_NAME = \"llama-3-2-1b-instruct\"\n\nstorage = HuggingfaceStorageConfigParam(type=\"huggingface\", repo_id=HF_REPO_ID, repo_type=\"model\")\nif hf_secret:\n storage[\"token_secret\"] = hf_secret.name\n\ntry:\n base_model_fs = sdk.files.filesets.create(\n workspace=\"default\", name=MODEL_NAME, description=\"Llama 3.2 1B Instruct base model\", storage=storage\n )\n print(f\"Created base model fileset: {MODEL_NAME}\")\nexcept ConflictError:\n base_model_fs = sdk.files.filesets.retrieve(workspace=\"default\", name=MODEL_NAME)\n print(\"Base model fileset already exists.\")\n\ntry:\n base_model = sdk.models.create(workspace=\"default\", name=MODEL_NAME, fileset=f\"default/{MODEL_NAME}\")\nexcept ConflictError:\n base_model = sdk.models.retrieve(workspace=\"default\", name=MODEL_NAME)\n\nprint(f\"Base model fileset: fileset://default/{base_model.name}\")\n\n# Wait for the ModelSpec to be inferred from the checkpoint.\ncheck = max_wait_time_checker(600, \"Model spec\")\nwhile not base_model.spec:\n check()\n time.sleep(10)\n base_model = sdk.models.retrieve(workspace=\"default\", name=MODEL_NAME)\nprint(\"Model spec ready\")", + "language": "python", + "source_html": "HF_REPO_ID = "meta-llama/Llama-3.2-1B-Instruct"\nMODEL_NAME = "llama-3-2-1b-instruct"\n\nstorage = HuggingfaceStorageConfigParam(type="huggingface", repo_id=HF_REPO_ID, repo_type="model")\nif hf_secret:\n storage["token_secret"] = hf_secret.name\n\ntry:\n base_model_fs = sdk.files.filesets.create(\n workspace="default", name=MODEL_NAME, description="Llama 3.2 1B Instruct base model", storage=storage\n )\n print(f"Created base model fileset: {MODEL_NAME}")\nexcept ConflictError:\n base_model_fs = sdk.files.filesets.retrieve(workspace="default", name=MODEL_NAME)\n print("Base model fileset already exists.")\n\ntry:\n base_model = sdk.models.create(workspace="default", name=MODEL_NAME, fileset=f"default/{MODEL_NAME}")\nexcept ConflictError:\n base_model = sdk.models.retrieve(workspace="default", name=MODEL_NAME)\n\nprint(f"Base model fileset: fileset://default/{base_model.name}")\n\n# Wait for the ModelSpec to be inferred from the checkpoint.\ncheck = max_wait_time_checker(600, "Model spec")\nwhile not base_model.spec:\n check()\n time.sleep(10)\n base_model = sdk.models.retrieve(workspace="default", name=MODEL_NAME)\nprint("Model spec ready")\n" + }, + { + "type": "markdown", + "source": "### 6. Create the DPO Customization Job\n\nSubmit a DPO job to the `rl` backend with `RlJobInput`. Note the DPO-specific shape:\n\n- `model` is a string ref to the model entity; `dataset` is a **single** string ref to the preference fileset (holding both files).\n- The training method is `{\"type\": \"dpo\", ...}` — full-weight, no `finetuning_type`/LoRA.\n- `ref_policy_kl_penalty` is **β** (DPO paper): how strongly the policy stays tied to the reference model.\n- `rl` auto-generates the job id (`rl-`); read it back from the response.\n\nOther configurable knobs: `optimizer_type`, `adam_eps`, `activation_checkpointing`, `keep_top_k`, `val_at_end`, `preference_loss_weight`, `sft_loss_weight`. Run `nemo customization rl explain` for the live schema.", + "source_html": "

6. Create the DPO Customization Job

\n

Submit a DPO job to the rl backend with RlJobInput. Note the DPO-specific shape:

\n
    \n
  • model is a string ref to the model entity; dataset is a single string ref to the preference fileset (holding both files).
    • \n
  • The training method is {"type": "dpo", ...} — full-weight, no finetuning_type/LoRA.
    • \n
  • ref_policy_kl_penalty is β (DPO paper): how strongly the policy stays tied to the reference model.
    • \n
  • rl auto-generates the job id (rl-<hex>); read it back from the response.
    • \n
\n

Other configurable knobs: optimizer_type, adam_eps, activation_checkpointing, keep_top_k, val_at_end, preference_loss_weight, sft_loss_weight. Run nemo customization rl explain for the live schema.

\n" + }, + { + "type": "code", + "source": "job_suffix = uuid.uuid4().hex[:8]\nOUTPUT_NAME = f\"llama-3-2-1b-dpo-{job_suffix}\"\n\nspec = RlJobInput(\n model=f\"default/{base_model.name}\",\n dataset=f\"default/{DATASET_NAME}\",\n training={\n \"type\": \"dpo\",\n \"epochs\": 1,\n \"batch_size\": 16,\n \"micro_batch_size\": 1,\n \"learning_rate\": 5e-6,\n \"max_seq_length\": 4096,\n \"ref_policy_kl_penalty\": 0.1,\n \"parallelism\": {\n \"num_nodes\": 1,\n \"num_gpus_per_node\": 1,\n \"tensor_parallel_size\": 1,\n \"pipeline_parallel_size\": 1,\n },\n },\n output={\"name\": OUTPUT_NAME},\n)\n\n# `rl` auto-generates the job id (rl-); do not pass name=.\njob = sdk.customization.rl.jobs.create(spec=spec, workspace=\"default\")\nprint(f\"Job ID: {job.job.name}\")\nprint(f\"Output model: {OUTPUT_NAME}\")", + "language": "python", + "source_html": "job_suffix = uuid.uuid4().hex[:8]\nOUTPUT_NAME = f"llama-3-2-1b-dpo-{job_suffix}"\n\nspec = RlJobInput(\n model=f"default/{base_model.name}",\n dataset=f"default/{DATASET_NAME}",\n training={\n "type": "dpo",\n "epochs": 1,\n "batch_size": 16,\n "micro_batch_size": 1,\n "learning_rate": 5e-6,\n "max_seq_length": 4096,\n "ref_policy_kl_penalty": 0.1,\n "parallelism": {\n "num_nodes": 1,\n "num_gpus_per_node": 1,\n "tensor_parallel_size": 1,\n "pipeline_parallel_size": 1,\n },\n },\n output={"name": OUTPUT_NAME},\n)\n\n# `rl` auto-generates the job id (rl-<hex>); do not pass name=.\njob = sdk.customization.rl.jobs.create(spec=spec, workspace="default")\nprint(f"Job ID: {job.job.name}")\nprint(f"Output model: {OUTPUT_NAME}")\n" + }, + { + "type": "markdown", + "source": "### 7. Track Training Progress\n\nThe DPO job runs four steps: download -> **dpo-training** (Ray) -> upload -> model-entity. We poll the top-level job status and surface the training step's progress.", + "source_html": "

7. Track Training Progress

\n

The DPO job runs four steps: download -> dpo-training (Ray) -> upload -> model-entity. We poll the top-level job status and surface the training step's progress.

\n" + }, + { + "type": "code", + "source": "from IPython.display import clear_output\n\ncheck = max_wait_time_checker(7200, \"DPO job\")\nwhile True:\n check()\n status = sdk.jobs.get_status(name=job.job.name, workspace=\"default\")\n clear_output(wait=True)\n print(f\"Job Status: {status.status}\")\n\n step = max_steps = phase = None\n for job_step in status.steps or []:\n if job_step.name == \"dpo-training\":\n for task in job_step.tasks or []:\n d = task.status_details or {}\n step, max_steps, phase = d.get(\"step\"), d.get(\"max_steps\"), d.get(\"phase\")\n break\n break\n if step is not None and max_steps:\n print(f\"Training: Step {step}/{max_steps} ({100 * step / max_steps:.1f}%)\")\n if phase:\n print(f\"Phase: {phase}\")\n\n if status.status in (\"completed\", \"failed\", \"cancelled\", \"error\"):\n print(f\"\\nJob finished: {status.status}\")\n break\n time.sleep(15)\n\nassert status.status == \"completed\"", + "language": "python", + "source_html": "from IPython.display import clear_output\n\ncheck = max_wait_time_checker(7200, "DPO job")\nwhile True:\n check()\n status = sdk.jobs.get_status(name=job.job.name, workspace="default")\n clear_output(wait=True)\n print(f"Job Status: {status.status}")\n\n step = max_steps = phase = None\n for job_step in status.steps or []:\n if job_step.name == "dpo-training":\n for task in job_step.tasks or []:\n d = task.status_details or {}\n step, max_steps, phase = d.get("step"), d.get("max_steps"), d.get("phase")\n break\n break\n if step is not None and max_steps:\n print(f"Training: Step {step}/{max_steps} ({100 * step / max_steps:.1f}%)")\n if phase:\n print(f"Phase: {phase}")\n\n if status.status in ("completed", "failed", "cancelled", "error"):\n print(f"\\nJob finished: {status.status}")\n break\n time.sleep(15)\n\nassert status.status == "completed"\n" + }, + { + "type": "markdown", + "source": "**Interpreting DPO training metrics** (in `status_details.metrics`):\n\n- **`loss`** — the DPO loss; should trend down as the policy learns to separate chosen from rejected.\n- **Reward margin** (chosen minus rejected reward) — should trend **up**: the model increasingly prefers chosen responses.\n- **Validation `loss`** — watch for divergence from training loss (overfitting). Raise `ref_policy_kl_penalty` (β) or add `sft_loss_weight` if the policy drifts too far from the reference.", + "source_html": "

Interpreting DPO training metrics (in status_details.metrics):

\n
    \n
  • loss — the DPO loss; should trend down as the policy learns to separate chosen from rejected.
    • \n
  • Reward margin (chosen minus rejected reward) — should trend up: the model increasingly prefers chosen responses.
    • \n
  • Validation loss — watch for divergence from training loss (overfitting). Raise ref_policy_kl_penalty (β) or add sft_loss_weight if the policy drifts too far from the reference.
    • \n
\n" + }, + { + "type": "markdown", + "source": "### 8. Validate the Output Model\n\nDPO produces a **full-weight model entity** (not an adapter). Confirm it was registered.", + "source_html": "

8. Validate the Output Model

\n

DPO produces a full-weight model entity (not an adapter). Confirm it was registered.

\n" + }, + { + "type": "code", + "source": "model_entity = sdk.models.retrieve(workspace=\"default\", name=OUTPUT_NAME)\nprint(model_entity.model_dump_json(indent=2))", + "language": "python", + "source_html": "model_entity = sdk.models.retrieve(workspace="default", name=OUTPUT_NAME)\nprint(model_entity.model_dump_json(indent=2))\n" + }, + { + "type": "markdown", + "source": "### 9. Deploy and Evaluate (optional)\n\nThe DPO output is a full model, so it deploys like any full-weight checkpoint (see the [Full SFT](/documentation/customizer-reference/tutorials/sft-customization-job) tutorial for details). We deploy with vLLM and send a chat completion.", + "source_html": "

9. Deploy and Evaluate (optional)

\n

The DPO output is a full model, so it deploys like any full-weight checkpoint (see the Full SFT tutorial for details). We deploy with vLLM and send a chat completion.

\n" + }, + { + "type": "code", + "source": "deploy_suffix = uuid.uuid4().hex[:8]\nDEPLOYMENT_CONFIG_NAME = f\"dpo-deployment-cfg-{deploy_suffix}\"\nDEPLOYMENT_NAME = f\"dpo-deployment-{deploy_suffix}\"\n\ndeployment_config = sdk.inference.deployment_configs.create(\n workspace=\"default\",\n name=DEPLOYMENT_CONFIG_NAME,\n engine=\"vllm\",\n model_spec={\"model_namespace\": \"default\", \"model_name\": OUTPUT_NAME},\n executor_config={\"gpu\": 1, \"image_name\": \"vllm/vllm-openai\", \"image_tag\": \"v0.22.1\"},\n)\n\ndeployment = sdk.inference.deployments.create(\n workspace=\"default\", name=DEPLOYMENT_NAME, config=deployment_config.name\n)\nprint(f\"Deployment name: {deployment.name}\")", + "language": "python", + "source_html": "deploy_suffix = uuid.uuid4().hex[:8]\nDEPLOYMENT_CONFIG_NAME = f"dpo-deployment-cfg-{deploy_suffix}"\nDEPLOYMENT_NAME = f"dpo-deployment-{deploy_suffix}"\n\ndeployment_config = sdk.inference.deployment_configs.create(\n workspace="default",\n name=DEPLOYMENT_CONFIG_NAME,\n engine="vllm",\n model_spec={"model_namespace": "default", "model_name": OUTPUT_NAME},\n executor_config={"gpu": 1, "image_name": "vllm/vllm-openai", "image_tag": "v0.22.1"},\n)\n\ndeployment = sdk.inference.deployments.create(\n workspace="default", name=DEPLOYMENT_NAME, config=deployment_config.name\n)\nprint(f"Deployment name: {deployment.name}")\n" + }, + { + "type": "code", + "source": "check = max_wait_time_checker(1800, \"Deployment\")\nwhile True:\n check()\n deployment_status = sdk.inference.deployments.retrieve(name=deployment.name, workspace=\"default\")\n clear_output(wait=True)\n print(f\"Deployment status: {deployment_status.status}\")\n deployment_state = str(deployment_status.status).lower()\n if deployment_state in (\"ready\", \"running\"):\n if not sdk.models.wait_for_gateway(deployment.name, workspace=\"default\", timeout=60):\n raise RuntimeError(\"Inference gateway did not become ready\")\n break\n if deployment_state in (\"failed\", \"error\", \"terminated\", \"lost\"):\n raise RuntimeError(f\"Deployment failed with status: {deployment_status.status}\")\n time.sleep(15)", + "language": "python", + "source_html": "check = max_wait_time_checker(1800, "Deployment")\nwhile True:\n check()\n deployment_status = sdk.inference.deployments.retrieve(name=deployment.name, workspace="default")\n clear_output(wait=True)\n print(f"Deployment status: {deployment_status.status}")\n deployment_state = str(deployment_status.status).lower()\n if deployment_state in ("ready", "running"):\n if not sdk.models.wait_for_gateway(deployment.name, workspace="default", timeout=60):\n raise RuntimeError("Inference gateway did not become ready")\n break\n if deployment_state in ("failed", "error", "terminated", "lost"):\n raise RuntimeError(f"Deployment failed with status: {deployment_status.status}")\n time.sleep(15)\n" + }, + { + "type": "code", + "source": "messages = [\n {\"role\": \"system\", \"content\": \"You are a helpful assistant.\"},\n {\"role\": \"user\", \"content\": \"Write a short, friendly email to a colleague asking to reschedule our meeting to Thursday.\"},\n]\n\nresponse = sdk.inference.gateway.provider.post(\n \"v1/chat/completions\",\n name=deployment.name,\n workspace=\"default\",\n body={\"model\": f\"default/{OUTPUT_NAME}\", \"messages\": messages, \"temperature\": 0.7, \"max_tokens\": 256},\n)\nprint(\"Model output:\\n\")\nprint(response[\"choices\"][0][\"message\"][\"content\"])", + "language": "python", + "source_html": "messages = [\n {"role": "system", "content": "You are a helpful assistant."},\n {"role": "user", "content": "Write a short, friendly email to a colleague asking to reschedule our meeting to Thursday."},\n]\n\nresponse = sdk.inference.gateway.provider.post(\n "v1/chat/completions",\n name=deployment.name,\n workspace="default",\n body={"model": f"default/{OUTPUT_NAME}", "messages": messages, "temperature": 0.7, "max_tokens": 256},\n)\nprint("Model output:\\n")\nprint(response["choices"][0]["message"]["content"])\n" + }, + { + "type": "markdown", + "source": "## Conclusion\n\nYou aligned a base model with **DPO** on the NeMo Platform using the `rl` backend:\n\n- Uploaded a HelpSteer3 preference dataset **as-is** (the platform detects the schema natively).\n- Submitted a full-weight DPO job that ran on a Ray cluster via the Kubernetes executor.\n- Registered the output as a full model entity and (optionally) deployed it for inference.\n\n**Next steps:** tune the alignment strength with `ref_policy_kl_penalty` (β), add `sft_loss_weight` to anchor the policy to the chosen responses, enable `activation_checkpointing` for memory headroom, or scale up with `parallelism`. See the `nemo-customizer` skill's `references/hyperparameters.md` (section NeMo-RL (DPO)) for the full knob reference.", + "source_html": "

Conclusion

\n

You aligned a base model with DPO on the NeMo Platform using the rl backend:

\n
    \n
  • Uploaded a HelpSteer3 preference dataset as-is (the platform detects the schema natively).
    • \n
  • Submitted a full-weight DPO job that ran on a Ray cluster via the Kubernetes executor.
    • \n
  • Registered the output as a full model entity and (optionally) deployed it for inference.
    • \n
\n

Next steps: tune the alignment strength with ref_policy_kl_penalty (β), add sft_loss_weight to anchor the policy to the chosen responses, enable activation_checkpointing for memory headroom, or scale up with parallelism. See the nemo-customizer skill's references/hyperparameters.md (section NeMo-RL (DPO)) for the full knob reference.

\n" + } +] }; diff --git a/docs/fern/components/notebooks/embedding-customization-job.json b/docs/fern/components/notebooks/embedding-customization-job.json index fda1255a6e..6f9b907a47 100644 --- a/docs/fern/components/notebooks/embedding-customization-job.json +++ b/docs/fern/components/notebooks/embedding-customization-job.json @@ -7,8 +7,8 @@ }, { "type": "markdown", - "source": "## Prerequisites\n\nBefore starting this tutorial, ensure you have:\n\n1. **Completed the [Quickstart](../../get-started/quickstart.md)** to install and deploy NeMo Platform locally\n2. **Installed the Python SDK** (PyPI wrapper: `pip install \"nemo-platform[all]\"`; source checkout: run `make bootstrap` from the repository root)\n3. **HuggingFace token** with read access to download the SPECTER dataset (get one at [huggingface.co/settings/tokens](https://huggingface.co/settings/tokens))\n4. **NGC API key** to pull NIM container images from nvcr.io (get one at [ngc.nvidia.com](https://ngc.nvidia.com/) → Setup → Generate API Key)", - "source_html": "

Prerequisites

\n

Before starting this tutorial, ensure you have:

\n
    \n
  1. Completed the Quickstart to install and deploy NeMo Platform locally
    2. \n
  2. Installed the Python SDK (PyPI wrapper: pip install "nemo-platform[all]"; source checkout: run make bootstrap from the repository root)
    3. \n
  3. HuggingFace token with read access to download the SPECTER dataset (get one at huggingface.co/settings/tokens)
    4. \n
  4. NGC API key to pull NIM container images from nvcr.io (get one at ngc.nvidia.com → Setup → Generate API Key)
    5. \n
\n" + "source": "## Prerequisites\n\nBefore starting this tutorial, ensure you have:\n\n1. **Completed the [Quickstart](../../get-started/quickstart.md)** to install and deploy NeMo Platform locally\n2. **Installed the Python SDK** (PyPI wrapper: `pip install \"nemo-platform[all]\"`; source checkout: run `make bootstrap` from the repository root)\n3. **NGC API key** to pull NIM container images from nvcr.io (get one at [ngc.nvidia.com](https://ngc.nvidia.com/) → Setup → Generate API Key)\n4. **At least one GPU with CUDA 13+**\n\nThe SPECTER dataset and the tutorial's base model are public and do not require a Hugging Face token. If you substitute a gated or private model, provide a token with read access.", + "source_html": "

Prerequisites

\n

Before starting this tutorial, ensure you have:

\n
    \n
  1. Completed the Quickstart to install and deploy NeMo Platform locally
    2. \n
  2. Installed the Python SDK (PyPI wrapper: pip install "nemo-platform[all]"; source checkout: run make bootstrap from the repository root)
    3. \n
  3. NGC API key to pull NIM container images from nvcr.io (get one at ngc.nvidia.com → Setup → Generate API Key)
    4. \n
  4. At least one GPU with CUDA 13+
    5. \n
\n

The SPECTER dataset and the tutorial's base model are public and do not require a Hugging Face token. If you substitute a gated or private model, provide a token with read access.

\n" }, { "type": "markdown", @@ -64,8 +64,8 @@ }, { "type": "markdown", - "source": "### 3. Prepare Dataset\n\nUse the [SPECTER dataset](https://huggingface.co/datasets/embedding-data/SPECTER) from HuggingFace, a collection of scientific paper triplets where papers that cite each other are considered related.\n\n**Dataset structure:**\n- ~684K scientific paper triplets (this tutorial uses 10%)\n- Each triplet: (query paper, positive/related paper, negative/unrelated paper)\n- Papers that cite each other are marked as \"related\"\n\nIn this tutorial the following dataset directory structure will be used:\n```\nembedding-dataset\n`-- training.jsonl\n`-- validation.jsonl\n```", - "source_html": "

3. Prepare Dataset

\n

Use the SPECTER dataset from HuggingFace, a collection of scientific paper triplets where papers that cite each other are considered related.

\n

Dataset structure:

\n
    \n
  • ~684K scientific paper triplets (this tutorial uses 10%)
    • \n
  • Each triplet: (query paper, positive/related paper, negative/unrelated paper)
    • \n
  • Papers that cite each other are marked as "related"
    • \n
\n

In this tutorial the following dataset directory structure will be used:

\n
embedding-dataset\n`-- training.jsonl\n`-- validation.jsonl\n
\n" + "source": "### 3. Prepare Dataset\n\nUse the [SPECTER dataset](https://huggingface.co/datasets/embedding-data/SPECTER) from Hugging Face, a collection of scientific paper triplets where papers that cite each other are considered related.\n\n**Dataset structure:**\n- ~684K scientific paper triplets (this tutorial uses 10%)\n- Each triplet: (query paper, positive/related paper, negative/unrelated paper)\n- Papers that cite each other are marked as \"related\"\n\nIn this tutorial the following dataset directory structure will be used:\n```\nembedding-dataset\n`-- training.jsonl\n`-- validation.jsonl\n```", + "source_html": "

3. Prepare Dataset

\n

Use the SPECTER dataset from Hugging Face, a collection of scientific paper triplets where papers that cite each other are considered related.

\n

Dataset structure:

\n
    \n
  • ~684K scientific paper triplets (this tutorial uses 10%)
    • \n
  • Each triplet: (query paper, positive/related paper, negative/unrelated paper)
    • \n
  • Papers that cite each other are marked as "related"
    • \n
\n

In this tutorial the following dataset directory structure will be used:

\n
embedding-dataset\n`-- training.jsonl\n`-- validation.jsonl\n
\n" }, { "type": "markdown", @@ -74,9 +74,9 @@ }, { "type": "code", - "source": "from pathlib import Path\nfrom datasets import load_dataset\nimport json\n\n# HuggingFace token for dataset access\nHF_TOKEN = os.environ.get(\"HF_TOKEN\")\nif not HF_TOKEN:\n raise ValueError(\"HF_TOKEN environment variable is required. Get one at https://huggingface.co/settings/tokens\")\nos.environ[\"HF_TOKEN\"] = HF_TOKEN\n\n# Configuration\nDATASET_SIZE = 3000 # Number of triplets (increase for better results, max ~684K)\nVALIDATION_SPLIT = 0.05 # 5% held out for validation\nSEED = 42\nDATASET_PATH = Path(\"embedding-dataset\").absolute()\n\n# Create directory\nos.makedirs(DATASET_PATH, exist_ok=True)\n\n# Download SPECTER dataset\nprint(\"Downloading SPECTER dataset...\")\ndata = load_dataset(\"embedding-data/SPECTER\")[\"train\"].shuffle(seed=SEED).select(range(DATASET_SIZE))\n\n# Split into train/validation\nprint(\"Splitting into train/validation...\")\nsplits = data.train_test_split(test_size=VALIDATION_SPLIT, seed=SEED)\ntrain_data = splits[\"train\"]\nvalidation_data = splits[\"test\"]\n\n# Convert to triplet JSONL format\nprint(\"Saving to JSONL...\")\nfor name, dataset in [(\"training\", train_data), (\"validation\", validation_data)]:\n with open(f\"{DATASET_PATH}/{name}.jsonl\", \"w\") as f:\n for row in dataset:\n # SPECTER format: row['set'] = [query, positive, negative]\n triplet = {\n \"query\": row[\"set\"][0],\n \"pos_doc\": row[\"set\"][1],\n \"neg_doc\": [row[\"set\"][2]] # List of negative documents\n }\n f.write(json.dumps(triplet) + \"\\n\")\n\nprint(f\"\\nPrepared {len(train_data):,} training, {len(validation_data):,} validation samples\")\nprint(f\"\\nExample triplet:\")\nprint(f\" Query: {train_data[0]['set'][0][:100]}...\")\nprint(f\" Positive: {train_data[0]['set'][1][:100]}...\")\nprint(f\" Negative: {train_data[0]['set'][2][:100]}...\")", + "source": "from pathlib import Path\nfrom datasets import load_dataset\nimport json\n\n# Configuration\nDATASET_SIZE = 3000 # Number of triplets (increase for better results, max ~684K)\nVALIDATION_SPLIT = 0.05 # 5% held out for validation\nSEED = 42\nDATASET_PATH = Path(\"embedding-dataset\").absolute()\n\n# Create directory\nos.makedirs(DATASET_PATH, exist_ok=True)\n\n# Download SPECTER dataset\nprint(\"Downloading SPECTER dataset...\")\ndata = load_dataset(\"embedding-data/SPECTER\")[\"train\"].shuffle(seed=SEED).select(range(DATASET_SIZE))\n\n# Split into train/validation\nprint(\"Splitting into train/validation...\")\nsplits = data.train_test_split(test_size=VALIDATION_SPLIT, seed=SEED)\ntrain_data = splits[\"train\"]\nvalidation_data = splits[\"test\"]\n\n# Convert to triplet JSONL format\nprint(\"Saving to JSONL...\")\nfor name, dataset in [(\"training\", train_data), (\"validation\", validation_data)]:\n with open(f\"{DATASET_PATH}/{name}.jsonl\", \"w\") as f:\n for row in dataset:\n # SPECTER format: row['set'] = [query, positive, negative]\n triplet = {\n \"query\": row[\"set\"][0],\n \"pos_doc\": row[\"set\"][1],\n \"neg_doc\": [row[\"set\"][2]] # List of negative documents\n }\n f.write(json.dumps(triplet) + \"\\n\")\n\nprint(f\"\\nPrepared {len(train_data):,} training, {len(validation_data):,} validation samples\")\nprint(f\"\\nExample triplet:\")\nprint(f\" Query: {train_data[0]['set'][0][:100]}...\")\nprint(f\" Positive: {train_data[0]['set'][1][:100]}...\")\nprint(f\" Negative: {train_data[0]['set'][2][:100]}...\")", "language": "python", - "source_html": "from pathlib import Path\nfrom datasets import load_dataset\nimport json\n\n# HuggingFace token for dataset access\nHF_TOKEN = os.environ.get("HF_TOKEN")\nif not HF_TOKEN:\n raise ValueError("HF_TOKEN environment variable is required. Get one at https://huggingface.co/settings/tokens")\nos.environ["HF_TOKEN"] = HF_TOKEN\n\n# Configuration\nDATASET_SIZE = 3000 # Number of triplets (increase for better results, max ~684K)\nVALIDATION_SPLIT = 0.05 # 5% held out for validation\nSEED = 42\nDATASET_PATH = Path("embedding-dataset").absolute()\n\n# Create directory\nos.makedirs(DATASET_PATH, exist_ok=True)\n\n# Download SPECTER dataset\nprint("Downloading SPECTER dataset...")\ndata = load_dataset("embedding-data/SPECTER")["train"].shuffle(seed=SEED).select(range(DATASET_SIZE))\n\n# Split into train/validation\nprint("Splitting into train/validation...")\nsplits = data.train_test_split(test_size=VALIDATION_SPLIT, seed=SEED)\ntrain_data = splits["train"]\nvalidation_data = splits["test"]\n\n# Convert to triplet JSONL format\nprint("Saving to JSONL...")\nfor name, dataset in [("training", train_data), ("validation", validation_data)]:\n with open(f"{DATASET_PATH}/{name}.jsonl", "w") as f:\n for row in dataset:\n # SPECTER format: row['set'] = [query, positive, negative]\n triplet = {\n "query": row["set"][0],\n "pos_doc": row["set"][1],\n "neg_doc": [row["set"][2]] # List of negative documents\n }\n f.write(json.dumps(triplet) + "\\n")\n\nprint(f"\\nPrepared {len(train_data):,} training, {len(validation_data):,} validation samples")\nprint(f"\\nExample triplet:")\nprint(f" Query: {train_data[0]['set'][0][:100]}...")\nprint(f" Positive: {train_data[0]['set'][1][:100]}...")\nprint(f" Negative: {train_data[0]['set'][2][:100]}...")\n" + "source_html": "from pathlib import Path\nfrom datasets import load_dataset\nimport json\n\n# Configuration\nDATASET_SIZE = 3000 # Number of triplets (increase for better results, max ~684K)\nVALIDATION_SPLIT = 0.05 # 5% held out for validation\nSEED = 42\nDATASET_PATH = Path("embedding-dataset").absolute()\n\n# Create directory\nos.makedirs(DATASET_PATH, exist_ok=True)\n\n# Download SPECTER dataset\nprint("Downloading SPECTER dataset...")\ndata = load_dataset("embedding-data/SPECTER")["train"].shuffle(seed=SEED).select(range(DATASET_SIZE))\n\n# Split into train/validation\nprint("Splitting into train/validation...")\nsplits = data.train_test_split(test_size=VALIDATION_SPLIT, seed=SEED)\ntrain_data = splits["train"]\nvalidation_data = splits["test"]\n\n# Convert to triplet JSONL format\nprint("Saving to JSONL...")\nfor name, dataset in [("training", train_data), ("validation", validation_data)]:\n with open(f"{DATASET_PATH}/{name}.jsonl", "w") as f:\n for row in dataset:\n # SPECTER format: row['set'] = [query, positive, negative]\n triplet = {\n "query": row["set"][0],\n "pos_doc": row["set"][1],\n "neg_doc": [row["set"][2]] # List of negative documents\n }\n f.write(json.dumps(triplet) + "\\n")\n\nprint(f"\\nPrepared {len(train_data):,} training, {len(validation_data):,} validation samples")\nprint(f"\\nExample triplet:")\nprint(f" Query: {train_data[0]['set'][0][:100]}...")\nprint(f" Positive: {train_data[0]['set'][1][:100]}...")\nprint(f" Negative: {train_data[0]['set'][2][:100]}...")\n" }, { "type": "markdown", @@ -91,30 +91,30 @@ }, { "type": "markdown", - "source": "### 6. Secrets Setup\n\nConfigure authentication for accessing base models:\n\n- **NGC models** (`ngc://` URIs): Requires NGC API key\n- **HuggingFace models** (`hf://` URIs): Requires HF token for gated/private models\n\nGet your credentials:\n- [NGC API Key](https://ngc.nvidia.com/) (Setup → Generate API Key)\n- [HuggingFace Token](https://huggingface.co/settings/tokens) (Create token with Read access)\n\n---\n\n#### Quick Setup Example\n\nThis tutorial fine-tunes [nvidia/llama-nemotron-embed-1b-v2](https://huggingface.co/nvidia/llama-nemotron-embed-1b-v2), an NVIDIA embedding model optimized for question-answering and retrieval tasks.", - "source_html": "

6. Secrets Setup

\n

Configure authentication for accessing base models:

\n
    \n
  • NGC models (ngc:// URIs): Requires NGC API key
    • \n
  • HuggingFace models (hf:// URIs): Requires HF token for gated/private models
    • \n
\n

Get your credentials:

\n\n
\n

Quick Setup Example

\n

This tutorial fine-tunes nvidia/llama-nemotron-embed-1b-v2, an NVIDIA embedding model optimized for question-answering and retrieval tasks.

\n" + "source": "### 6. Secrets Setup\n\nConfigure authentication for accessing base models:\n\n- **NGC models** (`ngc://` URIs): Requires NGC API key\n- **Hugging Face models** (`hf://` URIs): Requires HF token for gated/private models\n\nGet your credentials:\n- [NGC API Key](https://ngc.nvidia.com/) (Setup → Generate API Key)\n- [Hugging Face Token](https://huggingface.co/settings/tokens) (Optional; needed only for a gated/private replacement model)\n\n---\n\n#### Quick Setup Example\n\nThis tutorial fine-tunes [nvidia/llama-nemotron-embed-1b-v2](https://huggingface.co/nvidia/llama-nemotron-embed-1b-v2), an NVIDIA embedding model optimized for question-answering and retrieval tasks.", + "source_html": "

6. Secrets Setup

\n

Configure authentication for accessing base models:

\n
    \n
  • NGC models (ngc:// URIs): Requires NGC API key
    • \n
  • Hugging Face models (hf:// URIs): Requires HF token for gated/private models
    • \n
\n

Get your credentials:

\n\n
\n

Quick Setup Example

\n

This tutorial fine-tunes nvidia/llama-nemotron-embed-1b-v2, an NVIDIA embedding model optimized for question-answering and retrieval tasks.

\n" }, { "type": "code", - "source": "# Create secrets for model access\n# Note: NGC_API_KEY secret was already created in the baseline step (Step 2)\nHF_TOKEN = os.getenv(\"HF_TOKEN\")\n\n\ndef create_or_get_secret(name: str, value: str | None, label: str):\n if not value:\n raise ValueError(f\"{label} is not set\")\n try:\n secret = client.secrets.create(\n name=name,\n workspace=\"default\",\n value=value,\n )\n print(f\"Created secret: {name}\")\n return secret\n except ConflictError:\n print(f\"Secret '{name}' already exists, continuing...\")\n return client.secrets.retrieve(name=name, workspace=\"default\")\n\n\n# Create HuggingFace token secret (for downloading model from HF during training)\nhf_secret = create_or_get_secret(\"hf-token\", HF_TOKEN, \"HF_TOKEN\")\nprint(f\"HF_TOKEN secret: {hf_secret.name}\")\n\n# NGC secret was already created in baseline step (Step 2), or use the platform default\nif \"NGC_SECRET_NAME\" not in globals():\n NGC_SECRET_NAME = \"ngc-api-key\"\nprint(f\"NGC_API_KEY secret: {NGC_SECRET_NAME}\")", + "source": "# Create secrets for model access\n# Note: NGC_API_KEY secret was already created in the baseline step (Step 2)\nHF_TOKEN = os.getenv(\"HF_TOKEN\")\n\n\ndef create_or_get_secret(name: str, value: str | None, label: str):\n if not value:\n raise ValueError(f\"{label} is not set\")\n try:\n secret = client.secrets.create(\n name=name,\n workspace=\"default\",\n value=value,\n )\n print(f\"Created secret: {name}\")\n return secret\n except ConflictError:\n print(f\"Secret '{name}' already exists, continuing...\")\n return client.secrets.retrieve(name=name, workspace=\"default\")\n\n\n# Public Hugging Face models need no token. Create a secret only when HF_TOKEN is set.\nhf_secret = create_or_get_secret(\"hf-token\", HF_TOKEN, \"HF_TOKEN\") if HF_TOKEN else None\nif hf_secret:\n print(f\"HF_TOKEN secret: {hf_secret.name}\")\n\n# NGC secret was already created in baseline step (Step 2), or use the platform default\nif \"NGC_SECRET_NAME\" not in globals():\n NGC_SECRET_NAME = \"ngc-api-key\"\nprint(f\"NGC_API_KEY secret: {NGC_SECRET_NAME}\")", "language": "python", - "source_html": "# Create secrets for model access\n# Note: NGC_API_KEY secret was already created in the baseline step (Step 2)\nHF_TOKEN = os.getenv("HF_TOKEN")\n\n\ndef create_or_get_secret(name: str, value: str | None, label: str):\n if not value:\n raise ValueError(f"{label} is not set")\n try:\n secret = client.secrets.create(\n name=name,\n workspace="default",\n value=value,\n )\n print(f"Created secret: {name}")\n return secret\n except ConflictError:\n print(f"Secret '{name}' already exists, continuing...")\n return client.secrets.retrieve(name=name, workspace="default")\n\n\n# Create HuggingFace token secret (for downloading model from HF during training)\nhf_secret = create_or_get_secret("hf-token", HF_TOKEN, "HF_TOKEN")\nprint(f"HF_TOKEN secret: {hf_secret.name}")\n\n# NGC secret was already created in baseline step (Step 2), or use the platform default\nif "NGC_SECRET_NAME" not in globals():\n NGC_SECRET_NAME = "ngc-api-key"\nprint(f"NGC_API_KEY secret: {NGC_SECRET_NAME}")\n" + "source_html": "# Create secrets for model access\n# Note: NGC_API_KEY secret was already created in the baseline step (Step 2)\nHF_TOKEN = os.getenv("HF_TOKEN")\n\n\ndef create_or_get_secret(name: str, value: str | None, label: str):\n if not value:\n raise ValueError(f"{label} is not set")\n try:\n secret = client.secrets.create(\n name=name,\n workspace="default",\n value=value,\n )\n print(f"Created secret: {name}")\n return secret\n except ConflictError:\n print(f"Secret '{name}' already exists, continuing...")\n return client.secrets.retrieve(name=name, workspace="default")\n\n\n# Public Hugging Face models need no token. Create a secret only when HF_TOKEN is set.\nhf_secret = create_or_get_secret("hf-token", HF_TOKEN, "HF_TOKEN") if HF_TOKEN else None\nif hf_secret:\n print(f"HF_TOKEN secret: {hf_secret.name}")\n\n# NGC secret was already created in baseline step (Step 2), or use the platform default\nif "NGC_SECRET_NAME" not in globals():\n NGC_SECRET_NAME = "ngc-api-key"\nprint(f"NGC_API_KEY secret: {NGC_SECRET_NAME}")\n" }, { "type": "markdown", - "source": "### 7. Create Base Model FileSet and Model Entity\n\nCreate a fileset pointing to the [nvidia/llama-nemotron-embed-1b-v2](https://huggingface.co/nvidia/llama-nemotron-embed-1b-v2) embedding model from HuggingFace, then create a Model Entity that references this fileset. Model downloading will take place at training time.", - "source_html": "

7. Create Base Model FileSet and Model Entity

\n

Create a fileset pointing to the nvidia/llama-nemotron-embed-1b-v2 embedding model from HuggingFace, then create a Model Entity that references this fileset. Model downloading will take place at training time.

\n" + "source": "### 7. Create Base Model FileSet and Model Entity\n\nCreate a fileset pointing to the [nvidia/llama-nemotron-embed-1b-v2](https://huggingface.co/nvidia/llama-nemotron-embed-1b-v2) embedding model from Hugging Face, then create a Model Entity that references this fileset. Model downloading will take place at training time.", + "source_html": "

7. Create Base Model FileSet and Model Entity

\n

Create a fileset pointing to the nvidia/llama-nemotron-embed-1b-v2 embedding model from Hugging Face, then create a Model Entity that references this fileset. Model downloading will take place at training time.

\n" }, { "type": "code", - "source": "import time\nfrom nemo_platform.types.files import HuggingfaceStorageConfigParam\n\nHF_REPO_ID = \"nvidia/llama-nemotron-embed-1b-v2\"\nMODEL_NAME = \"nv-nemotron-embed-1b-base\"\n\n# Ensure you have a HuggingFace token secret created\ntry:\n base_model_fs = client.files.filesets.create(\n workspace=\"default\",\n name=MODEL_NAME,\n description=\"NVIDIA Llama Nemotron Embed 1B v2 embedding model\",\n storage=HuggingfaceStorageConfigParam(\n type=\"huggingface\",\n # repo_id is the full model name from Hugging Face\n repo_id=HF_REPO_ID,\n repo_type=\"model\",\n # we use the secret created in the previous step\n token_secret=hf_secret.name\n )\n )\nexcept ConflictError as e:\n print(f\"Base model fileset already exists. Skipping creation.\")\n base_model_fs = client.files.filesets.retrieve(\n workspace=\"default\",\n name=MODEL_NAME,\n )\n\n# Create Model Entity referencing the FileSet\ntry:\n base_model = client.models.create(\n workspace=\"default\",\n name=MODEL_NAME,\n fileset=f\"default/{MODEL_NAME}\",\n trust_remote_code=True,\n )\n print(f\"Created Model Entity: {MODEL_NAME}\")\nexcept ConflictError:\n print(f\"Base model already exists. Updating fileset if different.\")\n base_model = client.models.update(\n workspace=\"default\",\n name=MODEL_NAME,\n fileset=f\"default/{MODEL_NAME}\",\n trust_remote_code=True,\n )\n\nprint(f\"\\nBase model fileset: fileset://default/{base_model.name}\")\nprint(\"\\nBase model files:\")\nprint(json.dumps([f.model_dump() for f in client.files.list(fileset=MODEL_NAME, workspace=\"default\").data], indent=2))\n\n# Wait for ModelSpec to be populated from the checkpoint\nprint(\"\\nWaiting for ModelSpec to be populated...\")\nSPEC_TIMEOUT_SECONDS = 120\nspec_start = time.time()\nwhile not base_model.spec:\n if time.time() - spec_start > SPEC_TIMEOUT_SECONDS:\n raise TimeoutError(f\"ModelSpec not populated within {SPEC_TIMEOUT_SECONDS} seconds\")\n time.sleep(2)\n base_model = client.models.retrieve(\n workspace=\"default\",\n name=MODEL_NAME,\n )\n\nprint(f\"ModelSpec populated: {base_model.spec}\")", + "source": "import time\nfrom nemo_platform.types.files import HuggingfaceStorageConfigParam\n\nHF_REPO_ID = \"nvidia/llama-nemotron-embed-1b-v2\"\nMODEL_NAME = \"nv-nemotron-embed-1b-base\"\n\nstorage = HuggingfaceStorageConfigParam(\n type=\"huggingface\",\n repo_id=HF_REPO_ID,\n repo_type=\"model\",\n)\nif hf_secret:\n storage[\"token_secret\"] = hf_secret.name\n\ntry:\n base_model_fs = client.files.filesets.create(\n workspace=\"default\",\n name=MODEL_NAME,\n description=\"NVIDIA Llama Nemotron Embed 1B v2 embedding model\",\n storage=storage,\n )\nexcept ConflictError as e:\n print(f\"Base model fileset already exists. Skipping creation.\")\n base_model_fs = client.files.filesets.retrieve(\n workspace=\"default\",\n name=MODEL_NAME,\n )\n\n# Create Model Entity referencing the FileSet\ntry:\n base_model = client.models.create(\n workspace=\"default\",\n name=MODEL_NAME,\n fileset=f\"default/{MODEL_NAME}\",\n trust_remote_code=True,\n )\n print(f\"Created Model Entity: {MODEL_NAME}\")\nexcept ConflictError:\n print(f\"Base model already exists. Updating fileset if different.\")\n base_model = client.models.update(\n workspace=\"default\",\n name=MODEL_NAME,\n fileset=f\"default/{MODEL_NAME}\",\n trust_remote_code=True,\n )\n\nprint(f\"\\nBase model fileset: fileset://default/{base_model.name}\")\nprint(\"\\nBase model files:\")\nprint(json.dumps([f.model_dump() for f in client.files.list(fileset=MODEL_NAME, workspace=\"default\").data], indent=2))\n\n# Wait for ModelSpec to be populated from the checkpoint\nprint(\"\\nWaiting for ModelSpec to be populated...\")\nSPEC_TIMEOUT_SECONDS = 120\nspec_start = time.time()\nwhile not base_model.spec:\n if time.time() - spec_start > SPEC_TIMEOUT_SECONDS:\n raise TimeoutError(f\"ModelSpec not populated within {SPEC_TIMEOUT_SECONDS} seconds\")\n time.sleep(2)\n base_model = client.models.retrieve(\n workspace=\"default\",\n name=MODEL_NAME,\n )\n\nprint(f\"ModelSpec populated: {base_model.spec}\")", "language": "python", - "source_html": "import time\nfrom nemo_platform.types.files import HuggingfaceStorageConfigParam\n\nHF_REPO_ID = "nvidia/llama-nemotron-embed-1b-v2"\nMODEL_NAME = "nv-nemotron-embed-1b-base"\n\n# Ensure you have a HuggingFace token secret created\ntry:\n base_model_fs = client.files.filesets.create(\n workspace="default",\n name=MODEL_NAME,\n description="NVIDIA Llama Nemotron Embed 1B v2 embedding model",\n storage=HuggingfaceStorageConfigParam(\n type="huggingface",\n # repo_id is the full model name from Hugging Face\n repo_id=HF_REPO_ID,\n repo_type="model",\n # we use the secret created in the previous step\n token_secret=hf_secret.name\n )\n )\nexcept ConflictError as e:\n print(f"Base model fileset already exists. Skipping creation.")\n base_model_fs = client.files.filesets.retrieve(\n workspace="default",\n name=MODEL_NAME,\n )\n\n# Create Model Entity referencing the FileSet\ntry:\n base_model = client.models.create(\n workspace="default",\n name=MODEL_NAME,\n fileset=f"default/{MODEL_NAME}",\n trust_remote_code=True,\n )\n print(f"Created Model Entity: {MODEL_NAME}")\nexcept ConflictError:\n print(f"Base model already exists. Updating fileset if different.")\n base_model = client.models.update(\n workspace="default",\n name=MODEL_NAME,\n fileset=f"default/{MODEL_NAME}",\n trust_remote_code=True,\n )\n\nprint(f"\\nBase model fileset: fileset://default/{base_model.name}")\nprint("\\nBase model files:")\nprint(json.dumps([f.model_dump() for f in client.files.list(fileset=MODEL_NAME, workspace="default").data], indent=2))\n\n# Wait for ModelSpec to be populated from the checkpoint\nprint("\\nWaiting for ModelSpec to be populated...")\nSPEC_TIMEOUT_SECONDS = 120\nspec_start = time.time()\nwhile not base_model.spec:\n if time.time() - spec_start > SPEC_TIMEOUT_SECONDS:\n raise TimeoutError(f"ModelSpec not populated within {SPEC_TIMEOUT_SECONDS} seconds")\n time.sleep(2)\n base_model = client.models.retrieve(\n workspace="default",\n name=MODEL_NAME,\n )\n\nprint(f"ModelSpec populated: {base_model.spec}")\n" + "source_html": "import time\nfrom nemo_platform.types.files import HuggingfaceStorageConfigParam\n\nHF_REPO_ID = "nvidia/llama-nemotron-embed-1b-v2"\nMODEL_NAME = "nv-nemotron-embed-1b-base"\n\nstorage = HuggingfaceStorageConfigParam(\n type="huggingface",\n repo_id=HF_REPO_ID,\n repo_type="model",\n)\nif hf_secret:\n storage["token_secret"] = hf_secret.name\n\ntry:\n base_model_fs = client.files.filesets.create(\n workspace="default",\n name=MODEL_NAME,\n description="NVIDIA Llama Nemotron Embed 1B v2 embedding model",\n storage=storage,\n )\nexcept ConflictError as e:\n print(f"Base model fileset already exists. Skipping creation.")\n base_model_fs = client.files.filesets.retrieve(\n workspace="default",\n name=MODEL_NAME,\n )\n\n# Create Model Entity referencing the FileSet\ntry:\n base_model = client.models.create(\n workspace="default",\n name=MODEL_NAME,\n fileset=f"default/{MODEL_NAME}",\n trust_remote_code=True,\n )\n print(f"Created Model Entity: {MODEL_NAME}")\nexcept ConflictError:\n print(f"Base model already exists. Updating fileset if different.")\n base_model = client.models.update(\n workspace="default",\n name=MODEL_NAME,\n fileset=f"default/{MODEL_NAME}",\n trust_remote_code=True,\n )\n\nprint(f"\\nBase model fileset: fileset://default/{base_model.name}")\nprint("\\nBase model files:")\nprint(json.dumps([f.model_dump() for f in client.files.list(fileset=MODEL_NAME, workspace="default").data], indent=2))\n\n# Wait for ModelSpec to be populated from the checkpoint\nprint("\\nWaiting for ModelSpec to be populated...")\nSPEC_TIMEOUT_SECONDS = 120\nspec_start = time.time()\nwhile not base_model.spec:\n if time.time() - spec_start > SPEC_TIMEOUT_SECONDS:\n raise TimeoutError(f"ModelSpec not populated within {SPEC_TIMEOUT_SECONDS} seconds")\n time.sleep(2)\n base_model = client.models.retrieve(\n workspace="default",\n name=MODEL_NAME,\n )\n\nprint(f"ModelSpec populated: {base_model.spec}")\n" }, { "type": "markdown", - "source": "### 8. Create Embedding Fine-tuning Job\n\nCreate a customization job to fine-tune the embedding model using contrastive learning on the SPECTER dataset.\n\nSubmit to the **Automodel** backend using `AutomodelJobInput` with split `schedule`, `batch`, `optimizer`, and `parallelism` sections. Reference the model entity and dataset fileset by workspace/name (not `fileset://` URIs).\n\n**Key hyperparameters for embedding fine-tuning:**\n- **`training.training_type`**: `sft`\n- **`training.finetuning_type`**: `all_weights` for full fine-tuning, or `lora_merged` for merged LoRA\n- **`optimizer.learning_rate`**: Lower values (1e-6 to 5e-6) work well for embedding models\n- **`batch.global_batch_size`**: Larger batches improve contrastive learning (128-256 recommended)\n\n**NOTE:**\n\nNeMo Platform does not support unmerged LoRA adapters for embedding models because the embedding NIM requires ONNX format, which cannot represent standalone adapters. This notebook uses all-weights fine-tuning. For merged LoRA, set `finetuning_type` to `lora_merged`:\n\n```python\ntraining={\n \"training_type\": \"sft\",\n \"finetuning_type\": \"lora_merged\",\n \"lora\": {\"rank\": 16, \"alpha\": 32},\n \"max_seq_length\": MAX_SEQ_LENGTH,\n}\n```", - "source_html": "

8. Create Embedding Fine-tuning Job

\n

Create a customization job to fine-tune the embedding model using contrastive learning on the SPECTER dataset.

\n

Submit to the Automodel backend using AutomodelJobInput with split schedule, batch, optimizer, and parallelism sections. Reference the model entity and dataset fileset by workspace/name (not fileset:// URIs).

\n

Key hyperparameters for embedding fine-tuning:

\n
    \n
  • training.training_type: sft
    • \n
  • training.finetuning_type: all_weights for full fine-tuning, or lora_merged for merged LoRA
    • \n
  • optimizer.learning_rate: Lower values (1e-6 to 5e-6) work well for embedding models
    • \n
  • batch.global_batch_size: Larger batches improve contrastive learning (128-256 recommended)
    • \n
\n

NOTE:

\n

NeMo Platform does not support unmerged LoRA adapters for embedding models because the embedding NIM requires ONNX format, which cannot represent standalone adapters. This notebook uses all-weights fine-tuning. For merged LoRA, set finetuning_type to lora_merged:

\n
training={\n    "training_type": "sft",\n    "finetuning_type": "lora_merged",\n    "lora": {"rank": 16, "alpha": 32},\n    "max_seq_length": MAX_SEQ_LENGTH,\n}\n
\n" + "source": "### 8. Create Embedding Fine-tuning Job\n\nCreate a customization job to fine-tune the embedding model using contrastive learning on the SPECTER dataset.\n\nSubmit to the **Automodel** backend using `AutomodelJobInput` with split `schedule`, `batch`, `optimizer`, and `parallelism` sections. Reference the model entity and dataset fileset by workspace/name (not `fileset://` URIs).\n\n**Key hyperparameters for embedding fine-tuning:**\n- **`training.training_type`**: `sft`\n- **`training.finetuning_type`**: `all_weights` for full fine-tuning, or `lora_merged` for merged LoRA\n- **`optimizer.learning_rate`**: Lower values (1e-6 to 5e-6) work well for embedding models\n- **`batch.global_batch_size`**: Larger batches improve contrastive learning (128-256 recommended)\n\n**NOTE:**\n\nNeMo Platform does not support unmerged LoRA adapters for embedding models because the embedding NIM requires ONNX format, which cannot represent standalone adapters. This notebook uses all-weights fine-tuning. For merged LoRA, set `finetuning_type` to `lora_merged`:\n\n```python\ntraining={\n \"training_type\": \"sft\",\n \"finetuning_type\": \"lora_merged\",\n \"lora\": {\"rank\": 16, \"alpha\": 32},\n \"max_seq_length\": 512,\n}\n```", + "source_html": "

8. Create Embedding Fine-tuning Job

\n

Create a customization job to fine-tune the embedding model using contrastive learning on the SPECTER dataset.

\n

Submit to the Automodel backend using AutomodelJobInput with split schedule, batch, optimizer, and parallelism sections. Reference the model entity and dataset fileset by workspace/name (not fileset:// URIs).

\n

Key hyperparameters for embedding fine-tuning:

\n
    \n
  • training.training_type: sft
    • \n
  • training.finetuning_type: all_weights for full fine-tuning, or lora_merged for merged LoRA
    • \n
  • optimizer.learning_rate: Lower values (1e-6 to 5e-6) work well for embedding models
    • \n
  • batch.global_batch_size: Larger batches improve contrastive learning (128-256 recommended)
    • \n
\n

NOTE:

\n

NeMo Platform does not support unmerged LoRA adapters for embedding models because the embedding NIM requires ONNX format, which cannot represent standalone adapters. This notebook uses all-weights fine-tuning. For merged LoRA, set finetuning_type to lora_merged:

\n
training={\n    "training_type": "sft",\n    "finetuning_type": "lora_merged",\n    "lora": {"rank": 16, "alpha": 32},\n    "max_seq_length": 512,\n}\n
\n" }, { "type": "code", @@ -129,9 +129,9 @@ }, { "type": "code", - "source": "import time\nfrom IPython.display import clear_output\n\n# Poll job status every 10 seconds until completed\nwhile True:\n status = client.jobs.get_status(\n name=job.job.name,\n workspace=\"default\"\n )\n \n clear_output(wait=True)\n print(f\"Job Status: {status.model_dump_json(indent=2)}\")\n\n # Extract training progress from nested steps structure\n step: int | None = None\n max_steps: int | None = None\n training_phase: str | None = None\n\n for job_step in status.steps or []:\n if job_step.name == \"training\":\n for task in job_step.tasks or []:\n task_details = task.status_details or {}\n step = task_details.get(\"step\")\n max_steps = task_details.get(\"max_steps\")\n training_phase = task_details.get(\"phase\")\n break\n break\n\n if step is not None and max_steps is not None:\n progress_pct = (step / max_steps) * 100\n print(f\"Training Progress: Step {step}/{max_steps} ({progress_pct:.1f}%)\")\n if training_phase:\n print(f\"Training Phase: {training_phase}\")\n else:\n print(\"Training step not started yet or progress info not available\")\n \n # Exit loop when job is completed (or failed/cancelled)\n if status.status in (\"completed\", \"failed\", \"cancelled\", \"error\"):\n print(f\"\\nJob finished with status: {status.status}\")\n break\n \n time.sleep(10)", + "source": "import time\nfrom IPython.display import clear_output\n\n# Poll job status every 10 seconds until completed\nwhile True:\n status = client.jobs.get_status(\n name=job.job.name,\n workspace=\"default\"\n )\n \n clear_output(wait=True)\n print(f\"Job Status: {status.model_dump_json(indent=2)}\")\n\n # Extract training progress from nested steps structure\n step: int | None = None\n max_steps: int | None = None\n training_phase: str | None = None\n\n for job_step in status.steps or []:\n if job_step.name == \"training\":\n for task in job_step.tasks or []:\n task_details = task.status_details or {}\n step = task_details.get(\"step\")\n max_steps = task_details.get(\"max_steps\")\n training_phase = task_details.get(\"phase\")\n break\n break\n\n if step is not None and max_steps is not None:\n progress_pct = (step / max_steps) * 100\n print(f\"Training Progress: Step {step}/{max_steps} ({progress_pct:.1f}%)\")\n if training_phase:\n print(f\"Training Phase: {training_phase}\")\n else:\n print(\"Training step not started yet or progress info not available\")\n \n # Exit loop when job is completed (or failed/cancelled)\n if status.status in (\"completed\", \"failed\", \"cancelled\", \"error\"):\n print(f\"\\nJob finished with status: {status.status}\")\n break\n \n time.sleep(10)\n\nif status.status != \"completed\":\n raise RuntimeError(f\"Training job finished with status: {status.status}\")", "language": "python", - "source_html": "import time\nfrom IPython.display import clear_output\n\n# Poll job status every 10 seconds until completed\nwhile True:\n status = client.jobs.get_status(\n name=job.job.name,\n workspace="default"\n )\n \n clear_output(wait=True)\n print(f"Job Status: {status.model_dump_json(indent=2)}")\n\n # Extract training progress from nested steps structure\n step: int | None = None\n max_steps: int | None = None\n training_phase: str | None = None\n\n for job_step in status.steps or []:\n if job_step.name == "training":\n for task in job_step.tasks or []:\n task_details = task.status_details or {}\n step = task_details.get("step")\n max_steps = task_details.get("max_steps")\n training_phase = task_details.get("phase")\n break\n break\n\n if step is not None and max_steps is not None:\n progress_pct = (step / max_steps) * 100\n print(f"Training Progress: Step {step}/{max_steps} ({progress_pct:.1f}%)")\n if training_phase:\n print(f"Training Phase: {training_phase}")\n else:\n print("Training step not started yet or progress info not available")\n \n # Exit loop when job is completed (or failed/cancelled)\n if status.status in ("completed", "failed", "cancelled", "error"):\n print(f"\\nJob finished with status: {status.status}")\n break\n \n time.sleep(10)\n" + "source_html": "import time\nfrom IPython.display import clear_output\n\n# Poll job status every 10 seconds until completed\nwhile True:\n status = client.jobs.get_status(\n name=job.job.name,\n workspace="default"\n )\n \n clear_output(wait=True)\n print(f"Job Status: {status.model_dump_json(indent=2)}")\n\n # Extract training progress from nested steps structure\n step: int | None = None\n max_steps: int | None = None\n training_phase: str | None = None\n\n for job_step in status.steps or []:\n if job_step.name == "training":\n for task in job_step.tasks or []:\n task_details = task.status_details or {}\n step = task_details.get("step")\n max_steps = task_details.get("max_steps")\n training_phase = task_details.get("phase")\n break\n break\n\n if step is not None and max_steps is not None:\n progress_pct = (step / max_steps) * 100\n print(f"Training Progress: Step {step}/{max_steps} ({progress_pct:.1f}%)")\n if training_phase:\n print(f"Training Phase: {training_phase}")\n else:\n print("Training step not started yet or progress info not available")\n \n # Exit loop when job is completed (or failed/cancelled)\n if status.status in ("completed", "failed", "cancelled", "error"):\n print(f"\\nJob finished with status: {status.status}")\n break\n \n time.sleep(10)\n\nif status.status != "completed":\n raise RuntimeError(f"Training job finished with status: {status.status}")\n" }, { "type": "markdown", @@ -179,8 +179,8 @@ }, { "type": "markdown", - "source": "### Evaluation Best Practices\n\n**Manual Evaluation** (Recommended)\n- Test with real-world queries from your domain\n- Compare retrieval rankings before and after fine-tuning\n- Check that semantically similar items rank higher than keyword matches\n\n**What to look for:**\n- ✅ Relevant documents consistently rank in top positions\n- ✅ Keyword traps (like \"Random Forest\" vs \"Random Fields\") are handled correctly\n- ✅ Domain-specific terminology is understood\n- ❌ Unrelated documents with matching keywords do not rank high\n\n**Benchmark Evaluation**\n\nFor systematic evaluation, use the NeMo Evaluator service with retrieval benchmarks like SciDocs, BEIR, or MTEB. Refer to the [Evaluator documentation](../../evaluator/index.md) for details.\n\n---\n\n## Hyperparameters\n\nFor detailed information on all available hyperparameters, recommended values, and tuning guidance, refer to the [Hyperparameter Reference](../manage-customization-jobs/hyperparameters.md).\n\n**Embedding-Specific Recommendations:**\n\n| Parameter | Recommended | Notes |\n|-----------|-------------|-------|\n| `learning_rate` | 1e-6 to 5e-6 | Lower than standard SFT |\n| `batch_size` | 128-256 | Larger batches improve contrastive learning |\n| `max_seq_length` | 512 | Typical for embedding models |\n| `epochs` | 1-3 | Start small, increase if needed |\n\n---\n\n## Troubleshooting\n\n**Embeddings do not show improved retrieval:**\n- Verify dataset quality: triplets should have clear positive/negative distinctions\n- Use hard negatives: negatives should share some overlap with the query but not be relevant (easy negatives do not teach the model much)\n- Increase dataset size: 10K+ triplets recommended for meaningful improvement\n- Try more epochs: embedding models often need multiple passes\n- Lower learning rate: embedding models are sensitive to LR\n\n**Training loss not decreasing:**\n- Check triplet format: ensure `neg_doc` is a list even for single negatives\n- Verify hard negative quality: negatives should be challenging but clearly non-relevant\n- Increase batch size: contrastive learning benefits from larger batches\n\n**Deployment fails:**\n- Ensure you use the correct NIM image for embedding models\n- Verify sufficient GPU memory for the model size\n- Check deployment status: `client.inference.deployments.retrieve(name=deployment.name, workspace=\"default\")` and refer to platform logs for debugging\n\n## Next Steps\n\n- [Monitor training metrics](../manage-customization-jobs/get-job-status.md) in detail\n- [Evaluate your model](../../evaluator/index.md) with retrieval benchmarks\n- Integrate the fine-tuned embedding model into your RAG pipeline\n- Scale up training with the full SPECTER dataset (~684K triplets) for better results", - "source_html": "

Evaluation Best Practices

\n

Manual Evaluation (Recommended)

\n
    \n
  • Test with real-world queries from your domain
    • \n
  • Compare retrieval rankings before and after fine-tuning
    • \n
  • Check that semantically similar items rank higher than keyword matches
    • \n
\n

What to look for:

\n
    \n
  • ✅ Relevant documents consistently rank in top positions
    • \n
  • ✅ Keyword traps (like "Random Forest" vs "Random Fields") are handled correctly
    • \n
  • ✅ Domain-specific terminology is understood
    • \n
  • ❌ Unrelated documents with matching keywords do not rank high
    • \n
\n

Benchmark Evaluation

\n

For systematic evaluation, use the NeMo Evaluator service with retrieval benchmarks like SciDocs, BEIR, or MTEB. Refer to the Evaluator documentation for details.

\n
\n

Hyperparameters

\n

For detailed information on all available hyperparameters, recommended values, and tuning guidance, refer to the Hyperparameter Reference.

\n

Embedding-Specific Recommendations:

\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n
ParameterRecommendedNotes
learning_rate1e-6 to 5e-6Lower than standard SFT
batch_size128-256Larger batches improve contrastive learning
max_seq_length512Typical for embedding models
epochs1-3Start small, increase if needed
\n
\n

Troubleshooting

\n

Embeddings do not show improved retrieval:

\n
    \n
  • Verify dataset quality: triplets should have clear positive/negative distinctions
    • \n
  • Use hard negatives: negatives should share some overlap with the query but not be relevant (easy negatives do not teach the model much)
    • \n
  • Increase dataset size: 10K+ triplets recommended for meaningful improvement
    • \n
  • Try more epochs: embedding models often need multiple passes
    • \n
  • Lower learning rate: embedding models are sensitive to LR
    • \n
\n

Training loss not decreasing:

\n
    \n
  • Check triplet format: ensure neg_doc is a list even for single negatives
    • \n
  • Verify hard negative quality: negatives should be challenging but clearly non-relevant
    • \n
  • Increase batch size: contrastive learning benefits from larger batches
    • \n
\n

Deployment fails:

\n
    \n
  • Ensure you use the correct NIM image for embedding models
    • \n
  • Verify sufficient GPU memory for the model size
    • \n
  • Check deployment status: client.inference.deployments.retrieve(name=deployment.name, workspace="default") and refer to platform logs for debugging
    • \n
\n

Next Steps

\n
    \n
  • Monitor training metrics in detail
    • \n
  • Evaluate your model with retrieval benchmarks
    • \n
  • Integrate the fine-tuned embedding model into your RAG pipeline
    • \n
  • Scale up training with the full SPECTER dataset (~684K triplets) for better results
    • \n
\n" + "source": "### Evaluation Best Practices\n\n**Manual Evaluation** (Recommended)\n- Test with real-world queries from your domain\n- Compare retrieval rankings before and after fine-tuning\n- Check that semantically similar items rank higher than keyword matches\n\n**What to look for:**\n- ✅ Relevant documents consistently rank in top positions\n- ✅ Keyword traps (like \"Random Forest\" vs \"Random Fields\") are handled correctly\n- ✅ Domain-specific terminology is understood\n- ❌ Unrelated documents with matching keywords do not rank high\n\n**Benchmark Evaluation**\n\nFor systematic evaluation of end-to-end retrieval quality in a RAG pipeline, use the NeMo Evaluator [RAG metrics](../../evaluator/metrics/rag.md) (RAGAS `context_recall`, `context_precision`, and `context_relevance`).\n\n---\n\n## Hyperparameters\n\nFor detailed information on all available hyperparameters, recommended values, and tuning guidance, refer to the [Hyperparameter Reference](../manage-customization-jobs/hyperparameters.md).\n\n**Embedding-Specific Recommendations:**\n\n| Parameter | Recommended | Notes |\n|-----------|-------------|-------|\n| `optimizer.learning_rate` | 1e-6 to 5e-6 | Lower than standard SFT |\n| `batch.global_batch_size` | 128-256 | Larger batches improve contrastive learning |\n| `training.max_seq_length` | 512 | Typical for embedding models |\n| `schedule.epochs` | 1-3 | Start small, increase if needed |\n\n---\n\n## Troubleshooting\n\n**Embeddings do not show improved retrieval:**\n- Verify dataset quality: triplets should have clear positive/negative distinctions\n- Use hard negatives: negatives should share some overlap with the query but not be relevant (easy negatives do not teach the model much)\n- Increase dataset size: 10K+ triplets recommended for meaningful improvement\n- Try more epochs: embedding models often need multiple passes\n- Lower learning rate: embedding models are sensitive to LR\n\n**Training loss not decreasing:**\n- Check triplet format: ensure `neg_doc` is a list even for single negatives\n- Verify hard negative quality: negatives should be challenging but clearly non-relevant\n- Increase batch size: contrastive learning benefits from larger batches\n\n**Deployment fails:**\n- Ensure you use the correct NIM image for embedding models\n- Verify sufficient GPU memory for the model size\n- Check deployment status: `client.inference.deployments.retrieve(name=deployment.name, workspace=\"default\")` and refer to platform logs for debugging\n\n## Next Steps\n\n- [Monitor training metrics](../manage-customization-jobs/get-job-status.md) in detail\n- [Evaluate your model](../../evaluator/metrics/rag.md) with RAG metrics\n- Integrate the fine-tuned embedding model into your RAG pipeline\n- Scale up training with the full SPECTER dataset (~684K triplets) for better results", + "source_html": "

Evaluation Best Practices

\n

Manual Evaluation (Recommended)

\n
    \n
  • Test with real-world queries from your domain
    • \n
  • Compare retrieval rankings before and after fine-tuning
    • \n
  • Check that semantically similar items rank higher than keyword matches
    • \n
\n

What to look for:

\n
    \n
  • ✅ Relevant documents consistently rank in top positions
    • \n
  • ✅ Keyword traps (like "Random Forest" vs "Random Fields") are handled correctly
    • \n
  • ✅ Domain-specific terminology is understood
    • \n
  • ❌ Unrelated documents with matching keywords do not rank high
    • \n
\n

Benchmark Evaluation

\n

For systematic evaluation of end-to-end retrieval quality in a RAG pipeline, use the NeMo Evaluator RAG metrics (RAGAS context_recall, context_precision, and context_relevance).

\n
\n

Hyperparameters

\n

For detailed information on all available hyperparameters, recommended values, and tuning guidance, refer to the Hyperparameter Reference.

\n

Embedding-Specific Recommendations:

\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n
ParameterRecommendedNotes
optimizer.learning_rate1e-6 to 5e-6Lower than standard SFT
batch.global_batch_size128-256Larger batches improve contrastive learning
training.max_seq_length512Typical for embedding models
schedule.epochs1-3Start small, increase if needed
\n
\n

Troubleshooting

\n

Embeddings do not show improved retrieval:

\n
    \n
  • Verify dataset quality: triplets should have clear positive/negative distinctions
    • \n
  • Use hard negatives: negatives should share some overlap with the query but not be relevant (easy negatives do not teach the model much)
    • \n
  • Increase dataset size: 10K+ triplets recommended for meaningful improvement
    • \n
  • Try more epochs: embedding models often need multiple passes
    • \n
  • Lower learning rate: embedding models are sensitive to LR
    • \n
\n

Training loss not decreasing:

\n
    \n
  • Check triplet format: ensure neg_doc is a list even for single negatives
    • \n
  • Verify hard negative quality: negatives should be challenging but clearly non-relevant
    • \n
  • Increase batch size: contrastive learning benefits from larger batches
    • \n
\n

Deployment fails:

\n
    \n
  • Ensure you use the correct NIM image for embedding models
    • \n
  • Verify sufficient GPU memory for the model size
    • \n
  • Check deployment status: client.inference.deployments.retrieve(name=deployment.name, workspace="default") and refer to platform logs for debugging
    • \n
\n

Next Steps

\n
    \n
  • Monitor training metrics in detail
    • \n
  • Evaluate your model with RAG metrics
    • \n
  • Integrate the fine-tuned embedding model into your RAG pipeline
    • \n
  • Scale up training with the full SPECTER dataset (~684K triplets) for better results
    • \n
\n" } ] } \ No newline at end of file diff --git a/docs/fern/components/notebooks/embedding-customization-job.ts b/docs/fern/components/notebooks/embedding-customization-job.ts index 0efe18943f..397617ed4e 100644 --- a/docs/fern/components/notebooks/embedding-customization-job.ts +++ b/docs/fern/components/notebooks/embedding-customization-job.ts @@ -12,8 +12,8 @@ export default { cells: [ }, { "type": "markdown", - "source": "## Prerequisites\n\nBefore starting this tutorial, ensure you have:\n\n1. **Completed the [Quickstart](../../get-started/quickstart.md)** to install and deploy NeMo Platform locally\n2. **Installed the Python SDK** (PyPI wrapper: `pip install \"nemo-platform[all]\"`; source checkout: run `make bootstrap` from the repository root)\n3. **HuggingFace token** with read access to download the SPECTER dataset (get one at [huggingface.co/settings/tokens](https://huggingface.co/settings/tokens))\n4. **NGC API key** to pull NIM container images from nvcr.io (get one at [ngc.nvidia.com](https://ngc.nvidia.com/) → Setup → Generate API Key)", - "source_html": "

Prerequisites

\n

Before starting this tutorial, ensure you have:

\n
    \n
  1. Completed the Quickstart to install and deploy NeMo Platform locally
    2. \n
  2. Installed the Python SDK (PyPI wrapper: pip install "nemo-platform[all]"; source checkout: run make bootstrap from the repository root)
    3. \n
  3. HuggingFace token with read access to download the SPECTER dataset (get one at huggingface.co/settings/tokens)
    4. \n
  4. NGC API key to pull NIM container images from nvcr.io (get one at ngc.nvidia.com → Setup → Generate API Key)
    5. \n
\n" + "source": "## Prerequisites\n\nBefore starting this tutorial, ensure you have:\n\n1. **Completed the [Quickstart](../../get-started/quickstart.md)** to install and deploy NeMo Platform locally\n2. **Installed the Python SDK** (PyPI wrapper: `pip install \"nemo-platform[all]\"`; source checkout: run `make bootstrap` from the repository root)\n3. **NGC API key** to pull NIM container images from nvcr.io (get one at [ngc.nvidia.com](https://ngc.nvidia.com/) → Setup → Generate API Key)\n4. **At least one GPU with CUDA 13+**\n\nThe SPECTER dataset and the tutorial's base model are public and do not require a Hugging Face token. If you substitute a gated or private model, provide a token with read access.", + "source_html": "

Prerequisites

\n

Before starting this tutorial, ensure you have:

\n
    \n
  1. Completed the Quickstart to install and deploy NeMo Platform locally
    2. \n
  2. Installed the Python SDK (PyPI wrapper: pip install "nemo-platform[all]"; source checkout: run make bootstrap from the repository root)
    3. \n
  3. NGC API key to pull NIM container images from nvcr.io (get one at ngc.nvidia.com → Setup → Generate API Key)
    4. \n
  4. At least one GPU with CUDA 13+
    5. \n
\n

The SPECTER dataset and the tutorial's base model are public and do not require a Hugging Face token. If you substitute a gated or private model, provide a token with read access.

\n" }, { "type": "markdown", @@ -69,8 +69,8 @@ export default { cells: [ }, { "type": "markdown", - "source": "### 3. Prepare Dataset\n\nUse the [SPECTER dataset](https://huggingface.co/datasets/embedding-data/SPECTER) from HuggingFace, a collection of scientific paper triplets where papers that cite each other are considered related.\n\n**Dataset structure:**\n- ~684K scientific paper triplets (this tutorial uses 10%)\n- Each triplet: (query paper, positive/related paper, negative/unrelated paper)\n- Papers that cite each other are marked as \"related\"\n\nIn this tutorial the following dataset directory structure will be used:\n```\nembedding-dataset\n`-- training.jsonl\n`-- validation.jsonl\n```", - "source_html": "

3. Prepare Dataset

\n

Use the SPECTER dataset from HuggingFace, a collection of scientific paper triplets where papers that cite each other are considered related.

\n

Dataset structure:

\n
    \n
  • ~684K scientific paper triplets (this tutorial uses 10%)
    • \n
  • Each triplet: (query paper, positive/related paper, negative/unrelated paper)
    • \n
  • Papers that cite each other are marked as "related"
    • \n
\n

In this tutorial the following dataset directory structure will be used:

\n
embedding-dataset\n`-- training.jsonl\n`-- validation.jsonl\n
\n" + "source": "### 3. Prepare Dataset\n\nUse the [SPECTER dataset](https://huggingface.co/datasets/embedding-data/SPECTER) from Hugging Face, a collection of scientific paper triplets where papers that cite each other are considered related.\n\n**Dataset structure:**\n- ~684K scientific paper triplets (this tutorial uses 10%)\n- Each triplet: (query paper, positive/related paper, negative/unrelated paper)\n- Papers that cite each other are marked as \"related\"\n\nIn this tutorial the following dataset directory structure will be used:\n```\nembedding-dataset\n`-- training.jsonl\n`-- validation.jsonl\n```", + "source_html": "

3. Prepare Dataset

\n

Use the SPECTER dataset from Hugging Face, a collection of scientific paper triplets where papers that cite each other are considered related.

\n

Dataset structure:

\n
    \n
  • ~684K scientific paper triplets (this tutorial uses 10%)
    • \n
  • Each triplet: (query paper, positive/related paper, negative/unrelated paper)
    • \n
  • Papers that cite each other are marked as "related"
    • \n
\n

In this tutorial the following dataset directory structure will be used:

\n
embedding-dataset\n`-- training.jsonl\n`-- validation.jsonl\n
\n" }, { "type": "markdown", @@ -79,9 +79,9 @@ export default { cells: [ }, { "type": "code", - "source": "from pathlib import Path\nfrom datasets import load_dataset\nimport json\n\n# HuggingFace token for dataset access\nHF_TOKEN = os.environ.get(\"HF_TOKEN\")\nif not HF_TOKEN:\n raise ValueError(\"HF_TOKEN environment variable is required. Get one at https://huggingface.co/settings/tokens\")\nos.environ[\"HF_TOKEN\"] = HF_TOKEN\n\n# Configuration\nDATASET_SIZE = 3000 # Number of triplets (increase for better results, max ~684K)\nVALIDATION_SPLIT = 0.05 # 5% held out for validation\nSEED = 42\nDATASET_PATH = Path(\"embedding-dataset\").absolute()\n\n# Create directory\nos.makedirs(DATASET_PATH, exist_ok=True)\n\n# Download SPECTER dataset\nprint(\"Downloading SPECTER dataset...\")\ndata = load_dataset(\"embedding-data/SPECTER\")[\"train\"].shuffle(seed=SEED).select(range(DATASET_SIZE))\n\n# Split into train/validation\nprint(\"Splitting into train/validation...\")\nsplits = data.train_test_split(test_size=VALIDATION_SPLIT, seed=SEED)\ntrain_data = splits[\"train\"]\nvalidation_data = splits[\"test\"]\n\n# Convert to triplet JSONL format\nprint(\"Saving to JSONL...\")\nfor name, dataset in [(\"training\", train_data), (\"validation\", validation_data)]:\n with open(f\"{DATASET_PATH}/{name}.jsonl\", \"w\") as f:\n for row in dataset:\n # SPECTER format: row['set'] = [query, positive, negative]\n triplet = {\n \"query\": row[\"set\"][0],\n \"pos_doc\": row[\"set\"][1],\n \"neg_doc\": [row[\"set\"][2]] # List of negative documents\n }\n f.write(json.dumps(triplet) + \"\\n\")\n\nprint(f\"\\nPrepared {len(train_data):,} training, {len(validation_data):,} validation samples\")\nprint(f\"\\nExample triplet:\")\nprint(f\" Query: {train_data[0]['set'][0][:100]}...\")\nprint(f\" Positive: {train_data[0]['set'][1][:100]}...\")\nprint(f\" Negative: {train_data[0]['set'][2][:100]}...\")", + "source": "from pathlib import Path\nfrom datasets import load_dataset\nimport json\n\n# Configuration\nDATASET_SIZE = 3000 # Number of triplets (increase for better results, max ~684K)\nVALIDATION_SPLIT = 0.05 # 5% held out for validation\nSEED = 42\nDATASET_PATH = Path(\"embedding-dataset\").absolute()\n\n# Create directory\nos.makedirs(DATASET_PATH, exist_ok=True)\n\n# Download SPECTER dataset\nprint(\"Downloading SPECTER dataset...\")\ndata = load_dataset(\"embedding-data/SPECTER\")[\"train\"].shuffle(seed=SEED).select(range(DATASET_SIZE))\n\n# Split into train/validation\nprint(\"Splitting into train/validation...\")\nsplits = data.train_test_split(test_size=VALIDATION_SPLIT, seed=SEED)\ntrain_data = splits[\"train\"]\nvalidation_data = splits[\"test\"]\n\n# Convert to triplet JSONL format\nprint(\"Saving to JSONL...\")\nfor name, dataset in [(\"training\", train_data), (\"validation\", validation_data)]:\n with open(f\"{DATASET_PATH}/{name}.jsonl\", \"w\") as f:\n for row in dataset:\n # SPECTER format: row['set'] = [query, positive, negative]\n triplet = {\n \"query\": row[\"set\"][0],\n \"pos_doc\": row[\"set\"][1],\n \"neg_doc\": [row[\"set\"][2]] # List of negative documents\n }\n f.write(json.dumps(triplet) + \"\\n\")\n\nprint(f\"\\nPrepared {len(train_data):,} training, {len(validation_data):,} validation samples\")\nprint(f\"\\nExample triplet:\")\nprint(f\" Query: {train_data[0]['set'][0][:100]}...\")\nprint(f\" Positive: {train_data[0]['set'][1][:100]}...\")\nprint(f\" Negative: {train_data[0]['set'][2][:100]}...\")", "language": "python", - "source_html": "from pathlib import Path\nfrom datasets import load_dataset\nimport json\n\n# HuggingFace token for dataset access\nHF_TOKEN = os.environ.get("HF_TOKEN")\nif not HF_TOKEN:\n raise ValueError("HF_TOKEN environment variable is required. Get one at https://huggingface.co/settings/tokens")\nos.environ["HF_TOKEN"] = HF_TOKEN\n\n# Configuration\nDATASET_SIZE = 3000 # Number of triplets (increase for better results, max ~684K)\nVALIDATION_SPLIT = 0.05 # 5% held out for validation\nSEED = 42\nDATASET_PATH = Path("embedding-dataset").absolute()\n\n# Create directory\nos.makedirs(DATASET_PATH, exist_ok=True)\n\n# Download SPECTER dataset\nprint("Downloading SPECTER dataset...")\ndata = load_dataset("embedding-data/SPECTER")["train"].shuffle(seed=SEED).select(range(DATASET_SIZE))\n\n# Split into train/validation\nprint("Splitting into train/validation...")\nsplits = data.train_test_split(test_size=VALIDATION_SPLIT, seed=SEED)\ntrain_data = splits["train"]\nvalidation_data = splits["test"]\n\n# Convert to triplet JSONL format\nprint("Saving to JSONL...")\nfor name, dataset in [("training", train_data), ("validation", validation_data)]:\n with open(f"{DATASET_PATH}/{name}.jsonl", "w") as f:\n for row in dataset:\n # SPECTER format: row['set'] = [query, positive, negative]\n triplet = {\n "query": row["set"][0],\n "pos_doc": row["set"][1],\n "neg_doc": [row["set"][2]] # List of negative documents\n }\n f.write(json.dumps(triplet) + "\\n")\n\nprint(f"\\nPrepared {len(train_data):,} training, {len(validation_data):,} validation samples")\nprint(f"\\nExample triplet:")\nprint(f" Query: {train_data[0]['set'][0][:100]}...")\nprint(f" Positive: {train_data[0]['set'][1][:100]}...")\nprint(f" Negative: {train_data[0]['set'][2][:100]}...")\n" + "source_html": "from pathlib import Path\nfrom datasets import load_dataset\nimport json\n\n# Configuration\nDATASET_SIZE = 3000 # Number of triplets (increase for better results, max ~684K)\nVALIDATION_SPLIT = 0.05 # 5% held out for validation\nSEED = 42\nDATASET_PATH = Path("embedding-dataset").absolute()\n\n# Create directory\nos.makedirs(DATASET_PATH, exist_ok=True)\n\n# Download SPECTER dataset\nprint("Downloading SPECTER dataset...")\ndata = load_dataset("embedding-data/SPECTER")["train"].shuffle(seed=SEED).select(range(DATASET_SIZE))\n\n# Split into train/validation\nprint("Splitting into train/validation...")\nsplits = data.train_test_split(test_size=VALIDATION_SPLIT, seed=SEED)\ntrain_data = splits["train"]\nvalidation_data = splits["test"]\n\n# Convert to triplet JSONL format\nprint("Saving to JSONL...")\nfor name, dataset in [("training", train_data), ("validation", validation_data)]:\n with open(f"{DATASET_PATH}/{name}.jsonl", "w") as f:\n for row in dataset:\n # SPECTER format: row['set'] = [query, positive, negative]\n triplet = {\n "query": row["set"][0],\n "pos_doc": row["set"][1],\n "neg_doc": [row["set"][2]] # List of negative documents\n }\n f.write(json.dumps(triplet) + "\\n")\n\nprint(f"\\nPrepared {len(train_data):,} training, {len(validation_data):,} validation samples")\nprint(f"\\nExample triplet:")\nprint(f" Query: {train_data[0]['set'][0][:100]}...")\nprint(f" Positive: {train_data[0]['set'][1][:100]}...")\nprint(f" Negative: {train_data[0]['set'][2][:100]}...")\n" }, { "type": "markdown", @@ -96,30 +96,30 @@ export default { cells: [ }, { "type": "markdown", - "source": "### 6. Secrets Setup\n\nConfigure authentication for accessing base models:\n\n- **NGC models** (`ngc://` URIs): Requires NGC API key\n- **HuggingFace models** (`hf://` URIs): Requires HF token for gated/private models\n\nGet your credentials:\n- [NGC API Key](https://ngc.nvidia.com/) (Setup → Generate API Key)\n- [HuggingFace Token](https://huggingface.co/settings/tokens) (Create token with Read access)\n\n---\n\n#### Quick Setup Example\n\nThis tutorial fine-tunes [nvidia/llama-nemotron-embed-1b-v2](https://huggingface.co/nvidia/llama-nemotron-embed-1b-v2), an NVIDIA embedding model optimized for question-answering and retrieval tasks.", - "source_html": "

6. Secrets Setup

\n

Configure authentication for accessing base models:

\n
    \n
  • NGC models (ngc:// URIs): Requires NGC API key
    • \n
  • HuggingFace models (hf:// URIs): Requires HF token for gated/private models
    • \n
\n

Get your credentials:

\n\n
\n

Quick Setup Example

\n

This tutorial fine-tunes nvidia/llama-nemotron-embed-1b-v2, an NVIDIA embedding model optimized for question-answering and retrieval tasks.

\n" + "source": "### 6. Secrets Setup\n\nConfigure authentication for accessing base models:\n\n- **NGC models** (`ngc://` URIs): Requires NGC API key\n- **Hugging Face models** (`hf://` URIs): Requires HF token for gated/private models\n\nGet your credentials:\n- [NGC API Key](https://ngc.nvidia.com/) (Setup → Generate API Key)\n- [Hugging Face Token](https://huggingface.co/settings/tokens) (Optional; needed only for a gated/private replacement model)\n\n---\n\n#### Quick Setup Example\n\nThis tutorial fine-tunes [nvidia/llama-nemotron-embed-1b-v2](https://huggingface.co/nvidia/llama-nemotron-embed-1b-v2), an NVIDIA embedding model optimized for question-answering and retrieval tasks.", + "source_html": "

6. Secrets Setup

\n

Configure authentication for accessing base models:

\n
    \n
  • NGC models (ngc:// URIs): Requires NGC API key
    • \n
  • Hugging Face models (hf:// URIs): Requires HF token for gated/private models
    • \n
\n

Get your credentials:

\n\n
\n

Quick Setup Example

\n

This tutorial fine-tunes nvidia/llama-nemotron-embed-1b-v2, an NVIDIA embedding model optimized for question-answering and retrieval tasks.

\n" }, { "type": "code", - "source": "# Create secrets for model access\n# Note: NGC_API_KEY secret was already created in the baseline step (Step 2)\nHF_TOKEN = os.getenv(\"HF_TOKEN\")\n\n\ndef create_or_get_secret(name: str, value: str | None, label: str):\n if not value:\n raise ValueError(f\"{label} is not set\")\n try:\n secret = client.secrets.create(\n name=name,\n workspace=\"default\",\n value=value,\n )\n print(f\"Created secret: {name}\")\n return secret\n except ConflictError:\n print(f\"Secret '{name}' already exists, continuing...\")\n return client.secrets.retrieve(name=name, workspace=\"default\")\n\n\n# Create HuggingFace token secret (for downloading model from HF during training)\nhf_secret = create_or_get_secret(\"hf-token\", HF_TOKEN, \"HF_TOKEN\")\nprint(f\"HF_TOKEN secret: {hf_secret.name}\")\n\n# NGC secret was already created in baseline step (Step 2), or use the platform default\nif \"NGC_SECRET_NAME\" not in globals():\n NGC_SECRET_NAME = \"ngc-api-key\"\nprint(f\"NGC_API_KEY secret: {NGC_SECRET_NAME}\")", + "source": "# Create secrets for model access\n# Note: NGC_API_KEY secret was already created in the baseline step (Step 2)\nHF_TOKEN = os.getenv(\"HF_TOKEN\")\n\n\ndef create_or_get_secret(name: str, value: str | None, label: str):\n if not value:\n raise ValueError(f\"{label} is not set\")\n try:\n secret = client.secrets.create(\n name=name,\n workspace=\"default\",\n value=value,\n )\n print(f\"Created secret: {name}\")\n return secret\n except ConflictError:\n print(f\"Secret '{name}' already exists, continuing...\")\n return client.secrets.retrieve(name=name, workspace=\"default\")\n\n\n# Public Hugging Face models need no token. Create a secret only when HF_TOKEN is set.\nhf_secret = create_or_get_secret(\"hf-token\", HF_TOKEN, \"HF_TOKEN\") if HF_TOKEN else None\nif hf_secret:\n print(f\"HF_TOKEN secret: {hf_secret.name}\")\n\n# NGC secret was already created in baseline step (Step 2), or use the platform default\nif \"NGC_SECRET_NAME\" not in globals():\n NGC_SECRET_NAME = \"ngc-api-key\"\nprint(f\"NGC_API_KEY secret: {NGC_SECRET_NAME}\")", "language": "python", - "source_html": "# Create secrets for model access\n# Note: NGC_API_KEY secret was already created in the baseline step (Step 2)\nHF_TOKEN = os.getenv("HF_TOKEN")\n\n\ndef create_or_get_secret(name: str, value: str | None, label: str):\n if not value:\n raise ValueError(f"{label} is not set")\n try:\n secret = client.secrets.create(\n name=name,\n workspace="default",\n value=value,\n )\n print(f"Created secret: {name}")\n return secret\n except ConflictError:\n print(f"Secret '{name}' already exists, continuing...")\n return client.secrets.retrieve(name=name, workspace="default")\n\n\n# Create HuggingFace token secret (for downloading model from HF during training)\nhf_secret = create_or_get_secret("hf-token", HF_TOKEN, "HF_TOKEN")\nprint(f"HF_TOKEN secret: {hf_secret.name}")\n\n# NGC secret was already created in baseline step (Step 2), or use the platform default\nif "NGC_SECRET_NAME" not in globals():\n NGC_SECRET_NAME = "ngc-api-key"\nprint(f"NGC_API_KEY secret: {NGC_SECRET_NAME}")\n" + "source_html": "# Create secrets for model access\n# Note: NGC_API_KEY secret was already created in the baseline step (Step 2)\nHF_TOKEN = os.getenv("HF_TOKEN")\n\n\ndef create_or_get_secret(name: str, value: str | None, label: str):\n if not value:\n raise ValueError(f"{label} is not set")\n try:\n secret = client.secrets.create(\n name=name,\n workspace="default",\n value=value,\n )\n print(f"Created secret: {name}")\n return secret\n except ConflictError:\n print(f"Secret '{name}' already exists, continuing...")\n return client.secrets.retrieve(name=name, workspace="default")\n\n\n# Public Hugging Face models need no token. Create a secret only when HF_TOKEN is set.\nhf_secret = create_or_get_secret("hf-token", HF_TOKEN, "HF_TOKEN") if HF_TOKEN else None\nif hf_secret:\n print(f"HF_TOKEN secret: {hf_secret.name}")\n\n# NGC secret was already created in baseline step (Step 2), or use the platform default\nif "NGC_SECRET_NAME" not in globals():\n NGC_SECRET_NAME = "ngc-api-key"\nprint(f"NGC_API_KEY secret: {NGC_SECRET_NAME}")\n" }, { "type": "markdown", - "source": "### 7. Create Base Model FileSet and Model Entity\n\nCreate a fileset pointing to the [nvidia/llama-nemotron-embed-1b-v2](https://huggingface.co/nvidia/llama-nemotron-embed-1b-v2) embedding model from HuggingFace, then create a Model Entity that references this fileset. Model downloading will take place at training time.", - "source_html": "

7. Create Base Model FileSet and Model Entity

\n

Create a fileset pointing to the nvidia/llama-nemotron-embed-1b-v2 embedding model from HuggingFace, then create a Model Entity that references this fileset. Model downloading will take place at training time.

\n" + "source": "### 7. Create Base Model FileSet and Model Entity\n\nCreate a fileset pointing to the [nvidia/llama-nemotron-embed-1b-v2](https://huggingface.co/nvidia/llama-nemotron-embed-1b-v2) embedding model from Hugging Face, then create a Model Entity that references this fileset. Model downloading will take place at training time.", + "source_html": "

7. Create Base Model FileSet and Model Entity

\n

Create a fileset pointing to the nvidia/llama-nemotron-embed-1b-v2 embedding model from Hugging Face, then create a Model Entity that references this fileset. Model downloading will take place at training time.

\n" }, { "type": "code", - "source": "import time\nfrom nemo_platform.types.files import HuggingfaceStorageConfigParam\n\nHF_REPO_ID = \"nvidia/llama-nemotron-embed-1b-v2\"\nMODEL_NAME = \"nv-nemotron-embed-1b-base\"\n\n# Ensure you have a HuggingFace token secret created\ntry:\n base_model_fs = client.files.filesets.create(\n workspace=\"default\",\n name=MODEL_NAME,\n description=\"NVIDIA Llama Nemotron Embed 1B v2 embedding model\",\n storage=HuggingfaceStorageConfigParam(\n type=\"huggingface\",\n # repo_id is the full model name from Hugging Face\n repo_id=HF_REPO_ID,\n repo_type=\"model\",\n # we use the secret created in the previous step\n token_secret=hf_secret.name\n )\n )\nexcept ConflictError as e:\n print(f\"Base model fileset already exists. Skipping creation.\")\n base_model_fs = client.files.filesets.retrieve(\n workspace=\"default\",\n name=MODEL_NAME,\n )\n\n# Create Model Entity referencing the FileSet\ntry:\n base_model = client.models.create(\n workspace=\"default\",\n name=MODEL_NAME,\n fileset=f\"default/{MODEL_NAME}\",\n trust_remote_code=True,\n )\n print(f\"Created Model Entity: {MODEL_NAME}\")\nexcept ConflictError:\n print(f\"Base model already exists. Updating fileset if different.\")\n base_model = client.models.update(\n workspace=\"default\",\n name=MODEL_NAME,\n fileset=f\"default/{MODEL_NAME}\",\n trust_remote_code=True,\n )\n\nprint(f\"\\nBase model fileset: fileset://default/{base_model.name}\")\nprint(\"\\nBase model files:\")\nprint(json.dumps([f.model_dump() for f in client.files.list(fileset=MODEL_NAME, workspace=\"default\").data], indent=2))\n\n# Wait for ModelSpec to be populated from the checkpoint\nprint(\"\\nWaiting for ModelSpec to be populated...\")\nSPEC_TIMEOUT_SECONDS = 120\nspec_start = time.time()\nwhile not base_model.spec:\n if time.time() - spec_start > SPEC_TIMEOUT_SECONDS:\n raise TimeoutError(f\"ModelSpec not populated within {SPEC_TIMEOUT_SECONDS} seconds\")\n time.sleep(2)\n base_model = client.models.retrieve(\n workspace=\"default\",\n name=MODEL_NAME,\n )\n\nprint(f\"ModelSpec populated: {base_model.spec}\")", + "source": "import time\nfrom nemo_platform.types.files import HuggingfaceStorageConfigParam\n\nHF_REPO_ID = \"nvidia/llama-nemotron-embed-1b-v2\"\nMODEL_NAME = \"nv-nemotron-embed-1b-base\"\n\nstorage = HuggingfaceStorageConfigParam(\n type=\"huggingface\",\n repo_id=HF_REPO_ID,\n repo_type=\"model\",\n)\nif hf_secret:\n storage[\"token_secret\"] = hf_secret.name\n\ntry:\n base_model_fs = client.files.filesets.create(\n workspace=\"default\",\n name=MODEL_NAME,\n description=\"NVIDIA Llama Nemotron Embed 1B v2 embedding model\",\n storage=storage,\n )\nexcept ConflictError as e:\n print(f\"Base model fileset already exists. Skipping creation.\")\n base_model_fs = client.files.filesets.retrieve(\n workspace=\"default\",\n name=MODEL_NAME,\n )\n\n# Create Model Entity referencing the FileSet\ntry:\n base_model = client.models.create(\n workspace=\"default\",\n name=MODEL_NAME,\n fileset=f\"default/{MODEL_NAME}\",\n trust_remote_code=True,\n )\n print(f\"Created Model Entity: {MODEL_NAME}\")\nexcept ConflictError:\n print(f\"Base model already exists. Updating fileset if different.\")\n base_model = client.models.update(\n workspace=\"default\",\n name=MODEL_NAME,\n fileset=f\"default/{MODEL_NAME}\",\n trust_remote_code=True,\n )\n\nprint(f\"\\nBase model fileset: fileset://default/{base_model.name}\")\nprint(\"\\nBase model files:\")\nprint(json.dumps([f.model_dump() for f in client.files.list(fileset=MODEL_NAME, workspace=\"default\").data], indent=2))\n\n# Wait for ModelSpec to be populated from the checkpoint\nprint(\"\\nWaiting for ModelSpec to be populated...\")\nSPEC_TIMEOUT_SECONDS = 120\nspec_start = time.time()\nwhile not base_model.spec:\n if time.time() - spec_start > SPEC_TIMEOUT_SECONDS:\n raise TimeoutError(f\"ModelSpec not populated within {SPEC_TIMEOUT_SECONDS} seconds\")\n time.sleep(2)\n base_model = client.models.retrieve(\n workspace=\"default\",\n name=MODEL_NAME,\n )\n\nprint(f\"ModelSpec populated: {base_model.spec}\")", "language": "python", - "source_html": "import time\nfrom nemo_platform.types.files import HuggingfaceStorageConfigParam\n\nHF_REPO_ID = "nvidia/llama-nemotron-embed-1b-v2"\nMODEL_NAME = "nv-nemotron-embed-1b-base"\n\n# Ensure you have a HuggingFace token secret created\ntry:\n base_model_fs = client.files.filesets.create(\n workspace="default",\n name=MODEL_NAME,\n description="NVIDIA Llama Nemotron Embed 1B v2 embedding model",\n storage=HuggingfaceStorageConfigParam(\n type="huggingface",\n # repo_id is the full model name from Hugging Face\n repo_id=HF_REPO_ID,\n repo_type="model",\n # we use the secret created in the previous step\n token_secret=hf_secret.name\n )\n )\nexcept ConflictError as e:\n print(f"Base model fileset already exists. Skipping creation.")\n base_model_fs = client.files.filesets.retrieve(\n workspace="default",\n name=MODEL_NAME,\n )\n\n# Create Model Entity referencing the FileSet\ntry:\n base_model = client.models.create(\n workspace="default",\n name=MODEL_NAME,\n fileset=f"default/{MODEL_NAME}",\n trust_remote_code=True,\n )\n print(f"Created Model Entity: {MODEL_NAME}")\nexcept ConflictError:\n print(f"Base model already exists. Updating fileset if different.")\n base_model = client.models.update(\n workspace="default",\n name=MODEL_NAME,\n fileset=f"default/{MODEL_NAME}",\n trust_remote_code=True,\n )\n\nprint(f"\\nBase model fileset: fileset://default/{base_model.name}")\nprint("\\nBase model files:")\nprint(json.dumps([f.model_dump() for f in client.files.list(fileset=MODEL_NAME, workspace="default").data], indent=2))\n\n# Wait for ModelSpec to be populated from the checkpoint\nprint("\\nWaiting for ModelSpec to be populated...")\nSPEC_TIMEOUT_SECONDS = 120\nspec_start = time.time()\nwhile not base_model.spec:\n if time.time() - spec_start > SPEC_TIMEOUT_SECONDS:\n raise TimeoutError(f"ModelSpec not populated within {SPEC_TIMEOUT_SECONDS} seconds")\n time.sleep(2)\n base_model = client.models.retrieve(\n workspace="default",\n name=MODEL_NAME,\n )\n\nprint(f"ModelSpec populated: {base_model.spec}")\n" + "source_html": "import time\nfrom nemo_platform.types.files import HuggingfaceStorageConfigParam\n\nHF_REPO_ID = "nvidia/llama-nemotron-embed-1b-v2"\nMODEL_NAME = "nv-nemotron-embed-1b-base"\n\nstorage = HuggingfaceStorageConfigParam(\n type="huggingface",\n repo_id=HF_REPO_ID,\n repo_type="model",\n)\nif hf_secret:\n storage["token_secret"] = hf_secret.name\n\ntry:\n base_model_fs = client.files.filesets.create(\n workspace="default",\n name=MODEL_NAME,\n description="NVIDIA Llama Nemotron Embed 1B v2 embedding model",\n storage=storage,\n )\nexcept ConflictError as e:\n print(f"Base model fileset already exists. Skipping creation.")\n base_model_fs = client.files.filesets.retrieve(\n workspace="default",\n name=MODEL_NAME,\n )\n\n# Create Model Entity referencing the FileSet\ntry:\n base_model = client.models.create(\n workspace="default",\n name=MODEL_NAME,\n fileset=f"default/{MODEL_NAME}",\n trust_remote_code=True,\n )\n print(f"Created Model Entity: {MODEL_NAME}")\nexcept ConflictError:\n print(f"Base model already exists. Updating fileset if different.")\n base_model = client.models.update(\n workspace="default",\n name=MODEL_NAME,\n fileset=f"default/{MODEL_NAME}",\n trust_remote_code=True,\n )\n\nprint(f"\\nBase model fileset: fileset://default/{base_model.name}")\nprint("\\nBase model files:")\nprint(json.dumps([f.model_dump() for f in client.files.list(fileset=MODEL_NAME, workspace="default").data], indent=2))\n\n# Wait for ModelSpec to be populated from the checkpoint\nprint("\\nWaiting for ModelSpec to be populated...")\nSPEC_TIMEOUT_SECONDS = 120\nspec_start = time.time()\nwhile not base_model.spec:\n if time.time() - spec_start > SPEC_TIMEOUT_SECONDS:\n raise TimeoutError(f"ModelSpec not populated within {SPEC_TIMEOUT_SECONDS} seconds")\n time.sleep(2)\n base_model = client.models.retrieve(\n workspace="default",\n name=MODEL_NAME,\n )\n\nprint(f"ModelSpec populated: {base_model.spec}")\n" }, { "type": "markdown", - "source": "### 8. Create Embedding Fine-tuning Job\n\nCreate a customization job to fine-tune the embedding model using contrastive learning on the SPECTER dataset.\n\nSubmit to the **Automodel** backend using `AutomodelJobInput` with split `schedule`, `batch`, `optimizer`, and `parallelism` sections. Reference the model entity and dataset fileset by workspace/name (not `fileset://` URIs).\n\n**Key hyperparameters for embedding fine-tuning:**\n- **`training.training_type`**: `sft`\n- **`training.finetuning_type`**: `all_weights` for full fine-tuning, or `lora_merged` for merged LoRA\n- **`optimizer.learning_rate`**: Lower values (1e-6 to 5e-6) work well for embedding models\n- **`batch.global_batch_size`**: Larger batches improve contrastive learning (128-256 recommended)\n\n**NOTE:**\n\nNeMo Platform does not support unmerged LoRA adapters for embedding models because the embedding NIM requires ONNX format, which cannot represent standalone adapters. This notebook uses all-weights fine-tuning. For merged LoRA, set `finetuning_type` to `lora_merged`:\n\n```python\ntraining={\n \"training_type\": \"sft\",\n \"finetuning_type\": \"lora_merged\",\n \"lora\": {\"rank\": 16, \"alpha\": 32},\n \"max_seq_length\": MAX_SEQ_LENGTH,\n}\n```", - "source_html": "

8. Create Embedding Fine-tuning Job

\n

Create a customization job to fine-tune the embedding model using contrastive learning on the SPECTER dataset.

\n

Submit to the Automodel backend using AutomodelJobInput with split schedule, batch, optimizer, and parallelism sections. Reference the model entity and dataset fileset by workspace/name (not fileset:// URIs).

\n

Key hyperparameters for embedding fine-tuning:

\n
    \n
  • training.training_type: sft
    • \n
  • training.finetuning_type: all_weights for full fine-tuning, or lora_merged for merged LoRA
    • \n
  • optimizer.learning_rate: Lower values (1e-6 to 5e-6) work well for embedding models
    • \n
  • batch.global_batch_size: Larger batches improve contrastive learning (128-256 recommended)
    • \n
\n

NOTE:

\n

NeMo Platform does not support unmerged LoRA adapters for embedding models because the embedding NIM requires ONNX format, which cannot represent standalone adapters. This notebook uses all-weights fine-tuning. For merged LoRA, set finetuning_type to lora_merged:

\n
training={\n    "training_type": "sft",\n    "finetuning_type": "lora_merged",\n    "lora": {"rank": 16, "alpha": 32},\n    "max_seq_length": MAX_SEQ_LENGTH,\n}\n
\n" + "source": "### 8. Create Embedding Fine-tuning Job\n\nCreate a customization job to fine-tune the embedding model using contrastive learning on the SPECTER dataset.\n\nSubmit to the **Automodel** backend using `AutomodelJobInput` with split `schedule`, `batch`, `optimizer`, and `parallelism` sections. Reference the model entity and dataset fileset by workspace/name (not `fileset://` URIs).\n\n**Key hyperparameters for embedding fine-tuning:**\n- **`training.training_type`**: `sft`\n- **`training.finetuning_type`**: `all_weights` for full fine-tuning, or `lora_merged` for merged LoRA\n- **`optimizer.learning_rate`**: Lower values (1e-6 to 5e-6) work well for embedding models\n- **`batch.global_batch_size`**: Larger batches improve contrastive learning (128-256 recommended)\n\n**NOTE:**\n\nNeMo Platform does not support unmerged LoRA adapters for embedding models because the embedding NIM requires ONNX format, which cannot represent standalone adapters. This notebook uses all-weights fine-tuning. For merged LoRA, set `finetuning_type` to `lora_merged`:\n\n```python\ntraining={\n \"training_type\": \"sft\",\n \"finetuning_type\": \"lora_merged\",\n \"lora\": {\"rank\": 16, \"alpha\": 32},\n \"max_seq_length\": 512,\n}\n```", + "source_html": "

8. Create Embedding Fine-tuning Job

\n

Create a customization job to fine-tune the embedding model using contrastive learning on the SPECTER dataset.

\n

Submit to the Automodel backend using AutomodelJobInput with split schedule, batch, optimizer, and parallelism sections. Reference the model entity and dataset fileset by workspace/name (not fileset:// URIs).

\n

Key hyperparameters for embedding fine-tuning:

\n
    \n
  • training.training_type: sft
    • \n
  • training.finetuning_type: all_weights for full fine-tuning, or lora_merged for merged LoRA
    • \n
  • optimizer.learning_rate: Lower values (1e-6 to 5e-6) work well for embedding models
    • \n
  • batch.global_batch_size: Larger batches improve contrastive learning (128-256 recommended)
    • \n
\n

NOTE:

\n

NeMo Platform does not support unmerged LoRA adapters for embedding models because the embedding NIM requires ONNX format, which cannot represent standalone adapters. This notebook uses all-weights fine-tuning. For merged LoRA, set finetuning_type to lora_merged:

\n
training={\n    "training_type": "sft",\n    "finetuning_type": "lora_merged",\n    "lora": {"rank": 16, "alpha": 32},\n    "max_seq_length": 512,\n}\n
\n" }, { "type": "code", @@ -134,9 +134,9 @@ export default { cells: [ }, { "type": "code", - "source": "import time\nfrom IPython.display import clear_output\n\n# Poll job status every 10 seconds until completed\nwhile True:\n status = client.jobs.get_status(\n name=job.job.name,\n workspace=\"default\"\n )\n \n clear_output(wait=True)\n print(f\"Job Status: {status.model_dump_json(indent=2)}\")\n\n # Extract training progress from nested steps structure\n step: int | None = None\n max_steps: int | None = None\n training_phase: str | None = None\n\n for job_step in status.steps or []:\n if job_step.name == \"training\":\n for task in job_step.tasks or []:\n task_details = task.status_details or {}\n step = task_details.get(\"step\")\n max_steps = task_details.get(\"max_steps\")\n training_phase = task_details.get(\"phase\")\n break\n break\n\n if step is not None and max_steps is not None:\n progress_pct = (step / max_steps) * 100\n print(f\"Training Progress: Step {step}/{max_steps} ({progress_pct:.1f}%)\")\n if training_phase:\n print(f\"Training Phase: {training_phase}\")\n else:\n print(\"Training step not started yet or progress info not available\")\n \n # Exit loop when job is completed (or failed/cancelled)\n if status.status in (\"completed\", \"failed\", \"cancelled\", \"error\"):\n print(f\"\\nJob finished with status: {status.status}\")\n break\n \n time.sleep(10)", + "source": "import time\nfrom IPython.display import clear_output\n\n# Poll job status every 10 seconds until completed\nwhile True:\n status = client.jobs.get_status(\n name=job.job.name,\n workspace=\"default\"\n )\n \n clear_output(wait=True)\n print(f\"Job Status: {status.model_dump_json(indent=2)}\")\n\n # Extract training progress from nested steps structure\n step: int | None = None\n max_steps: int | None = None\n training_phase: str | None = None\n\n for job_step in status.steps or []:\n if job_step.name == \"training\":\n for task in job_step.tasks or []:\n task_details = task.status_details or {}\n step = task_details.get(\"step\")\n max_steps = task_details.get(\"max_steps\")\n training_phase = task_details.get(\"phase\")\n break\n break\n\n if step is not None and max_steps is not None:\n progress_pct = (step / max_steps) * 100\n print(f\"Training Progress: Step {step}/{max_steps} ({progress_pct:.1f}%)\")\n if training_phase:\n print(f\"Training Phase: {training_phase}\")\n else:\n print(\"Training step not started yet or progress info not available\")\n \n # Exit loop when job is completed (or failed/cancelled)\n if status.status in (\"completed\", \"failed\", \"cancelled\", \"error\"):\n print(f\"\\nJob finished with status: {status.status}\")\n break\n \n time.sleep(10)\n\nif status.status != \"completed\":\n raise RuntimeError(f\"Training job finished with status: {status.status}\")", "language": "python", - "source_html": "import time\nfrom IPython.display import clear_output\n\n# Poll job status every 10 seconds until completed\nwhile True:\n status = client.jobs.get_status(\n name=job.job.name,\n workspace="default"\n )\n \n clear_output(wait=True)\n print(f"Job Status: {status.model_dump_json(indent=2)}")\n\n # Extract training progress from nested steps structure\n step: int | None = None\n max_steps: int | None = None\n training_phase: str | None = None\n\n for job_step in status.steps or []:\n if job_step.name == "training":\n for task in job_step.tasks or []:\n task_details = task.status_details or {}\n step = task_details.get("step")\n max_steps = task_details.get("max_steps")\n training_phase = task_details.get("phase")\n break\n break\n\n if step is not None and max_steps is not None:\n progress_pct = (step / max_steps) * 100\n print(f"Training Progress: Step {step}/{max_steps} ({progress_pct:.1f}%)")\n if training_phase:\n print(f"Training Phase: {training_phase}")\n else:\n print("Training step not started yet or progress info not available")\n \n # Exit loop when job is completed (or failed/cancelled)\n if status.status in ("completed", "failed", "cancelled", "error"):\n print(f"\\nJob finished with status: {status.status}")\n break\n \n time.sleep(10)\n" + "source_html": "import time\nfrom IPython.display import clear_output\n\n# Poll job status every 10 seconds until completed\nwhile True:\n status = client.jobs.get_status(\n name=job.job.name,\n workspace="default"\n )\n \n clear_output(wait=True)\n print(f"Job Status: {status.model_dump_json(indent=2)}")\n\n # Extract training progress from nested steps structure\n step: int | None = None\n max_steps: int | None = None\n training_phase: str | None = None\n\n for job_step in status.steps or []:\n if job_step.name == "training":\n for task in job_step.tasks or []:\n task_details = task.status_details or {}\n step = task_details.get("step")\n max_steps = task_details.get("max_steps")\n training_phase = task_details.get("phase")\n break\n break\n\n if step is not None and max_steps is not None:\n progress_pct = (step / max_steps) * 100\n print(f"Training Progress: Step {step}/{max_steps} ({progress_pct:.1f}%)")\n if training_phase:\n print(f"Training Phase: {training_phase}")\n else:\n print("Training step not started yet or progress info not available")\n \n # Exit loop when job is completed (or failed/cancelled)\n if status.status in ("completed", "failed", "cancelled", "error"):\n print(f"\\nJob finished with status: {status.status}")\n break\n \n time.sleep(10)\n\nif status.status != "completed":\n raise RuntimeError(f"Training job finished with status: {status.status}")\n" }, { "type": "markdown", @@ -184,7 +184,7 @@ export default { cells: [ }, { "type": "markdown", - "source": "### Evaluation Best Practices\n\n**Manual Evaluation** (Recommended)\n- Test with real-world queries from your domain\n- Compare retrieval rankings before and after fine-tuning\n- Check that semantically similar items rank higher than keyword matches\n\n**What to look for:**\n- ✅ Relevant documents consistently rank in top positions\n- ✅ Keyword traps (like \"Random Forest\" vs \"Random Fields\") are handled correctly\n- ✅ Domain-specific terminology is understood\n- ❌ Unrelated documents with matching keywords do not rank high\n\n**Benchmark Evaluation**\n\nFor systematic evaluation, use the NeMo Evaluator service with retrieval benchmarks like SciDocs, BEIR, or MTEB. Refer to the [Evaluator documentation](../../evaluator/index.md) for details.\n\n---\n\n## Hyperparameters\n\nFor detailed information on all available hyperparameters, recommended values, and tuning guidance, refer to the [Hyperparameter Reference](../manage-customization-jobs/hyperparameters.md).\n\n**Embedding-Specific Recommendations:**\n\n| Parameter | Recommended | Notes |\n|-----------|-------------|-------|\n| `learning_rate` | 1e-6 to 5e-6 | Lower than standard SFT |\n| `batch_size` | 128-256 | Larger batches improve contrastive learning |\n| `max_seq_length` | 512 | Typical for embedding models |\n| `epochs` | 1-3 | Start small, increase if needed |\n\n---\n\n## Troubleshooting\n\n**Embeddings do not show improved retrieval:**\n- Verify dataset quality: triplets should have clear positive/negative distinctions\n- Use hard negatives: negatives should share some overlap with the query but not be relevant (easy negatives do not teach the model much)\n- Increase dataset size: 10K+ triplets recommended for meaningful improvement\n- Try more epochs: embedding models often need multiple passes\n- Lower learning rate: embedding models are sensitive to LR\n\n**Training loss not decreasing:**\n- Check triplet format: ensure `neg_doc` is a list even for single negatives\n- Verify hard negative quality: negatives should be challenging but clearly non-relevant\n- Increase batch size: contrastive learning benefits from larger batches\n\n**Deployment fails:**\n- Ensure you use the correct NIM image for embedding models\n- Verify sufficient GPU memory for the model size\n- Check deployment status: `client.inference.deployments.retrieve(name=deployment.name, workspace=\"default\")` and refer to platform logs for debugging\n\n## Next Steps\n\n- [Monitor training metrics](../manage-customization-jobs/get-job-status.md) in detail\n- [Evaluate your model](../../evaluator/index.md) with retrieval benchmarks\n- Integrate the fine-tuned embedding model into your RAG pipeline\n- Scale up training with the full SPECTER dataset (~684K triplets) for better results", - "source_html": "

Evaluation Best Practices

\n

Manual Evaluation (Recommended)

\n
    \n
  • Test with real-world queries from your domain
    • \n
  • Compare retrieval rankings before and after fine-tuning
    • \n
  • Check that semantically similar items rank higher than keyword matches
    • \n
\n

What to look for:

\n
    \n
  • ✅ Relevant documents consistently rank in top positions
    • \n
  • ✅ Keyword traps (like "Random Forest" vs "Random Fields") are handled correctly
    • \n
  • ✅ Domain-specific terminology is understood
    • \n
  • ❌ Unrelated documents with matching keywords do not rank high
    • \n
\n

Benchmark Evaluation

\n

For systematic evaluation, use the NeMo Evaluator service with retrieval benchmarks like SciDocs, BEIR, or MTEB. Refer to the Evaluator documentation for details.

\n
\n

Hyperparameters

\n

For detailed information on all available hyperparameters, recommended values, and tuning guidance, refer to the Hyperparameter Reference.

\n

Embedding-Specific Recommendations:

\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n
ParameterRecommendedNotes
learning_rate1e-6 to 5e-6Lower than standard SFT
batch_size128-256Larger batches improve contrastive learning
max_seq_length512Typical for embedding models
epochs1-3Start small, increase if needed
\n
\n

Troubleshooting

\n

Embeddings do not show improved retrieval:

\n
    \n
  • Verify dataset quality: triplets should have clear positive/negative distinctions
    • \n
  • Use hard negatives: negatives should share some overlap with the query but not be relevant (easy negatives do not teach the model much)
    • \n
  • Increase dataset size: 10K+ triplets recommended for meaningful improvement
    • \n
  • Try more epochs: embedding models often need multiple passes
    • \n
  • Lower learning rate: embedding models are sensitive to LR
    • \n
\n

Training loss not decreasing:

\n
    \n
  • Check triplet format: ensure neg_doc is a list even for single negatives
    • \n
  • Verify hard negative quality: negatives should be challenging but clearly non-relevant
    • \n
  • Increase batch size: contrastive learning benefits from larger batches
    • \n
\n

Deployment fails:

\n
    \n
  • Ensure you use the correct NIM image for embedding models
    • \n
  • Verify sufficient GPU memory for the model size
    • \n
  • Check deployment status: client.inference.deployments.retrieve(name=deployment.name, workspace="default") and refer to platform logs for debugging
    • \n
\n

Next Steps

\n
    \n
  • Monitor training metrics in detail
    • \n
  • Evaluate your model with retrieval benchmarks
    • \n
  • Integrate the fine-tuned embedding model into your RAG pipeline
    • \n
  • Scale up training with the full SPECTER dataset (~684K triplets) for better results
    • \n
\n" + "source": "### Evaluation Best Practices\n\n**Manual Evaluation** (Recommended)\n- Test with real-world queries from your domain\n- Compare retrieval rankings before and after fine-tuning\n- Check that semantically similar items rank higher than keyword matches\n\n**What to look for:**\n- ✅ Relevant documents consistently rank in top positions\n- ✅ Keyword traps (like \"Random Forest\" vs \"Random Fields\") are handled correctly\n- ✅ Domain-specific terminology is understood\n- ❌ Unrelated documents with matching keywords do not rank high\n\n**Benchmark Evaluation**\n\nFor systematic evaluation of end-to-end retrieval quality in a RAG pipeline, use the NeMo Evaluator [RAG metrics](../../evaluator/metrics/rag.md) (RAGAS `context_recall`, `context_precision`, and `context_relevance`).\n\n---\n\n## Hyperparameters\n\nFor detailed information on all available hyperparameters, recommended values, and tuning guidance, refer to the [Hyperparameter Reference](../manage-customization-jobs/hyperparameters.md).\n\n**Embedding-Specific Recommendations:**\n\n| Parameter | Recommended | Notes |\n|-----------|-------------|-------|\n| `optimizer.learning_rate` | 1e-6 to 5e-6 | Lower than standard SFT |\n| `batch.global_batch_size` | 128-256 | Larger batches improve contrastive learning |\n| `training.max_seq_length` | 512 | Typical for embedding models |\n| `schedule.epochs` | 1-3 | Start small, increase if needed |\n\n---\n\n## Troubleshooting\n\n**Embeddings do not show improved retrieval:**\n- Verify dataset quality: triplets should have clear positive/negative distinctions\n- Use hard negatives: negatives should share some overlap with the query but not be relevant (easy negatives do not teach the model much)\n- Increase dataset size: 10K+ triplets recommended for meaningful improvement\n- Try more epochs: embedding models often need multiple passes\n- Lower learning rate: embedding models are sensitive to LR\n\n**Training loss not decreasing:**\n- Check triplet format: ensure `neg_doc` is a list even for single negatives\n- Verify hard negative quality: negatives should be challenging but clearly non-relevant\n- Increase batch size: contrastive learning benefits from larger batches\n\n**Deployment fails:**\n- Ensure you use the correct NIM image for embedding models\n- Verify sufficient GPU memory for the model size\n- Check deployment status: `client.inference.deployments.retrieve(name=deployment.name, workspace=\"default\")` and refer to platform logs for debugging\n\n## Next Steps\n\n- [Monitor training metrics](../manage-customization-jobs/get-job-status.md) in detail\n- [Evaluate your model](../../evaluator/metrics/rag.md) with RAG metrics\n- Integrate the fine-tuned embedding model into your RAG pipeline\n- Scale up training with the full SPECTER dataset (~684K triplets) for better results", + "source_html": "

Evaluation Best Practices

\n

Manual Evaluation (Recommended)

\n
    \n
  • Test with real-world queries from your domain
    • \n
  • Compare retrieval rankings before and after fine-tuning
    • \n
  • Check that semantically similar items rank higher than keyword matches
    • \n
\n

What to look for:

\n
    \n
  • ✅ Relevant documents consistently rank in top positions
    • \n
  • ✅ Keyword traps (like "Random Forest" vs "Random Fields") are handled correctly
    • \n
  • ✅ Domain-specific terminology is understood
    • \n
  • ❌ Unrelated documents with matching keywords do not rank high
    • \n
\n

Benchmark Evaluation

\n

For systematic evaluation of end-to-end retrieval quality in a RAG pipeline, use the NeMo Evaluator RAG metrics (RAGAS context_recall, context_precision, and context_relevance).

\n
\n

Hyperparameters

\n

For detailed information on all available hyperparameters, recommended values, and tuning guidance, refer to the Hyperparameter Reference.

\n

Embedding-Specific Recommendations:

\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n
ParameterRecommendedNotes
optimizer.learning_rate1e-6 to 5e-6Lower than standard SFT
batch.global_batch_size128-256Larger batches improve contrastive learning
training.max_seq_length512Typical for embedding models
schedule.epochs1-3Start small, increase if needed
\n
\n

Troubleshooting

\n

Embeddings do not show improved retrieval:

\n
    \n
  • Verify dataset quality: triplets should have clear positive/negative distinctions
    • \n
  • Use hard negatives: negatives should share some overlap with the query but not be relevant (easy negatives do not teach the model much)
    • \n
  • Increase dataset size: 10K+ triplets recommended for meaningful improvement
    • \n
  • Try more epochs: embedding models often need multiple passes
    • \n
  • Lower learning rate: embedding models are sensitive to LR
    • \n
\n

Training loss not decreasing:

\n
    \n
  • Check triplet format: ensure neg_doc is a list even for single negatives
    • \n
  • Verify hard negative quality: negatives should be challenging but clearly non-relevant
    • \n
  • Increase batch size: contrastive learning benefits from larger batches
    • \n
\n

Deployment fails:

\n
    \n
  • Ensure you use the correct NIM image for embedding models
    • \n
  • Verify sufficient GPU memory for the model size
    • \n
  • Check deployment status: client.inference.deployments.retrieve(name=deployment.name, workspace="default") and refer to platform logs for debugging
    • \n
\n

Next Steps

\n
    \n
  • Monitor training metrics in detail
    • \n
  • Evaluate your model with RAG metrics
    • \n
  • Integrate the fine-tuned embedding model into your RAG pipeline
    • \n
  • Scale up training with the full SPECTER dataset (~684K triplets) for better results
    • \n
\n" } ] }; diff --git a/docs/fern/components/notebooks/lora-customization-job.json b/docs/fern/components/notebooks/lora-customization-job.json index e2f2ff8b16..502ece99c3 100644 --- a/docs/fern/components/notebooks/lora-customization-job.json +++ b/docs/fern/components/notebooks/lora-customization-job.json @@ -7,8 +7,8 @@ }, { "type": "markdown", - "source": "## Prerequisites\n\nBefore starting this tutorial, ensure you have:\n\n1. **Completed the [Quickstart](../../get-started/quickstart.md)** to install and deploy NeMo Platform locally\n2. **Installed the Python SDK** (PyPI wrapper: `pip install \"nemo-platform[all]\"`; source checkout: run `make bootstrap` from the repository root)\n3. **Installed the `datasets` package** for loading SQuAD: `pip install datasets`\n4. **At least one GPU with CUDA 12.8+**", - "source_html": "

Prerequisites

\n

Before starting this tutorial, ensure you have:

\n
    \n
  1. Completed the Quickstart to install and deploy NeMo Platform locally
    2. \n
  2. Installed the Python SDK (PyPI wrapper: pip install "nemo-platform[all]"; source checkout: run make bootstrap from the repository root)
    3. \n
  3. Installed the datasets package for loading SQuAD: pip install datasets
    4. \n
  4. At least one GPU with CUDA 12.8+
    5. \n
\n" + "source": "## Prerequisites\n\nBefore starting this tutorial, ensure you have:\n\n1. **Completed the [Quickstart](../../get-started/quickstart.md)** to install and deploy NeMo Platform locally\n2. **Installed the Python SDK** (PyPI wrapper: `pip install \"nemo-platform[all]\"`; source checkout: run `make bootstrap` from the repository root)\n3. **Installed the `datasets` package** for loading SQuAD: `pip install datasets`\n4. **At least one GPU with CUDA 13+**", + "source_html": "

Prerequisites

\n

Before starting this tutorial, ensure you have:

\n
    \n
  1. Completed the Quickstart to install and deploy NeMo Platform locally
    2. \n
  2. Installed the Python SDK (PyPI wrapper: pip install "nemo-platform[all]"; source checkout: run make bootstrap from the repository root)
    3. \n
  3. Installed the datasets package for loading SQuAD: pip install datasets
    4. \n
  4. At least one GPU with CUDA 13+
    5. \n
\n" }, { "type": "markdown", @@ -55,8 +55,8 @@ }, { "type": "markdown", - "source": "### 4. Secrets Setup\n\nFor Huggingface models that require authentication, create a secret with your HF token. Get a token from [Huggingface Settings](https://huggingface.co/settings/tokens) and accept the model terms.\n\nThis is generally true for LLaMa based models (e.g. [Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct)).\n\n```sh\nexport HF_TOKEN=\n```", - "source_html": "

4. Secrets Setup

\n

For Huggingface models that require authentication, create a secret with your HF token. Get a token from Huggingface Settings and accept the model terms.

\n

This is generally true for LLaMa based models (e.g. Llama-3.2-1B-Instruct).

\n
export HF_TOKEN=<your-huggingface-token>\n
\n" + "source": "### 4. Secrets Setup\n\nFor Hugging Face models that require authentication, create a secret with your HF token. Get a token from [Hugging Face Settings](https://huggingface.co/settings/tokens) and accept the model terms.\n\nThis is generally true for Llama-based models (for example, [Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct)).\n\n```sh\nexport HF_TOKEN=\n```", + "source_html": "

4. Secrets Setup

\n

For Hugging Face models that require authentication, create a secret with your HF token. Get a token from Hugging Face Settings and accept the model terms.

\n

This is generally true for Llama-based models (for example, Llama-3.2-1B-Instruct).

\n
export HF_TOKEN=<your-huggingface-token>\n
\n" }, { "type": "code", @@ -71,9 +71,9 @@ }, { "type": "code", - "source": "HF_REPO_ID = \"Qwen/Qwen3-0.6B\"\nMODEL_NAME = \"qwen3-0.6b\"\n\ntry:\n storage = HuggingfaceStorageConfigParam(\n type=\"huggingface\",\n repo_id=HF_REPO_ID,\n repo_type=\"model\",\n )\n if hf_secret:\n storage[\"token_secret\"] = hf_secret.name\n base_model_fs = client.files.filesets.create(\n workspace=\"default\",\n name=MODEL_NAME,\n description=\"Qwen3 0.6b base model from Huggingface\",\n storage=storage,\n cache=True,\n )\nexcept ConflictError:\n base_model_fs = client.files.filesets.retrieve(workspace=\"default\", name=MODEL_NAME)\n\ntry:\n base_model = client.models.create(\n workspace=\"default\",\n name=MODEL_NAME,\n fileset=f\"default/{MODEL_NAME}\",\n trust_remote_code=False,\n )\nexcept ConflictError:\n client.models.update(\n workspace=\"default\",\n name=MODEL_NAME,\n fileset=f\"default/{MODEL_NAME}\",\n trust_remote_code=False,\n )\n base_model = client.models.retrieve(workspace=\"default\", name=MODEL_NAME)\n\nprint(f\"Base model fileset: fileset://default/{base_model.name}\")\nprint(client.files.list(fileset=MODEL_NAME, workspace=\"default\"))\n\ntime_check = max_wait_time_checker(600, \"Model Spec\")\nwhile not base_model.spec:\n time_check()\n time.sleep(10)\n base_model = client.models.retrieve(workspace=\"default\", name=MODEL_NAME)\n\n# Clear verbose linear_layers list for cleaner output\nbase_model.spec.linear_layers = None\nprint(f\"ModelSpec: {base_model.spec}\")", + "source": "HF_REPO_ID = \"Qwen/Qwen3-0.6B\"\nMODEL_NAME = \"qwen3-0.6b\"\n\ntry:\n storage = HuggingfaceStorageConfigParam(\n type=\"huggingface\",\n repo_id=HF_REPO_ID,\n repo_type=\"model\",\n )\n if hf_secret:\n storage[\"token_secret\"] = hf_secret.name\n base_model_fs = client.files.filesets.create(\n workspace=\"default\",\n name=MODEL_NAME,\n description=\"Qwen3 0.6b base model from Hugging Face\",\n storage=storage,\n cache=True,\n )\nexcept ConflictError:\n base_model_fs = client.files.filesets.retrieve(workspace=\"default\", name=MODEL_NAME)\n\ntry:\n base_model = client.models.create(\n workspace=\"default\",\n name=MODEL_NAME,\n fileset=f\"default/{MODEL_NAME}\",\n trust_remote_code=False,\n )\nexcept ConflictError:\n client.models.update(\n workspace=\"default\",\n name=MODEL_NAME,\n fileset=f\"default/{MODEL_NAME}\",\n trust_remote_code=False,\n )\n base_model = client.models.retrieve(workspace=\"default\", name=MODEL_NAME)\n\nprint(f\"Base model fileset: fileset://default/{base_model.name}\")\nprint(client.files.list(fileset=MODEL_NAME, workspace=\"default\"))\n\ntime_check = max_wait_time_checker(600, \"Model Spec\")\nwhile not base_model.spec:\n time_check()\n time.sleep(10)\n base_model = client.models.retrieve(workspace=\"default\", name=MODEL_NAME)\n\n# Clear verbose linear_layers list for cleaner output\nbase_model.spec.linear_layers = None\nprint(f\"ModelSpec: {base_model.spec}\")", "language": "python", - "source_html": "HF_REPO_ID = "Qwen/Qwen3-0.6B"\nMODEL_NAME = "qwen3-0.6b"\n\ntry:\n storage = HuggingfaceStorageConfigParam(\n type="huggingface",\n repo_id=HF_REPO_ID,\n repo_type="model",\n )\n if hf_secret:\n storage["token_secret"] = hf_secret.name\n base_model_fs = client.files.filesets.create(\n workspace="default",\n name=MODEL_NAME,\n description="Qwen3 0.6b base model from Huggingface",\n storage=storage,\n cache=True,\n )\nexcept ConflictError:\n base_model_fs = client.files.filesets.retrieve(workspace="default", name=MODEL_NAME)\n\ntry:\n base_model = client.models.create(\n workspace="default",\n name=MODEL_NAME,\n fileset=f"default/{MODEL_NAME}",\n trust_remote_code=False,\n )\nexcept ConflictError:\n client.models.update(\n workspace="default",\n name=MODEL_NAME,\n fileset=f"default/{MODEL_NAME}",\n trust_remote_code=False,\n )\n base_model = client.models.retrieve(workspace="default", name=MODEL_NAME)\n\nprint(f"Base model fileset: fileset://default/{base_model.name}")\nprint(client.files.list(fileset=MODEL_NAME, workspace="default"))\n\ntime_check = max_wait_time_checker(600, "Model Spec")\nwhile not base_model.spec:\n time_check()\n time.sleep(10)\n base_model = client.models.retrieve(workspace="default", name=MODEL_NAME)\n\n# Clear verbose linear_layers list for cleaner output\nbase_model.spec.linear_layers = None\nprint(f"ModelSpec: {base_model.spec}")\n" + "source_html": "HF_REPO_ID = "Qwen/Qwen3-0.6B"\nMODEL_NAME = "qwen3-0.6b"\n\ntry:\n storage = HuggingfaceStorageConfigParam(\n type="huggingface",\n repo_id=HF_REPO_ID,\n repo_type="model",\n )\n if hf_secret:\n storage["token_secret"] = hf_secret.name\n base_model_fs = client.files.filesets.create(\n workspace="default",\n name=MODEL_NAME,\n description="Qwen3 0.6b base model from Hugging Face",\n storage=storage,\n cache=True,\n )\nexcept ConflictError:\n base_model_fs = client.files.filesets.retrieve(workspace="default", name=MODEL_NAME)\n\ntry:\n base_model = client.models.create(\n workspace="default",\n name=MODEL_NAME,\n fileset=f"default/{MODEL_NAME}",\n trust_remote_code=False,\n )\nexcept ConflictError:\n client.models.update(\n workspace="default",\n name=MODEL_NAME,\n fileset=f"default/{MODEL_NAME}",\n trust_remote_code=False,\n )\n base_model = client.models.retrieve(workspace="default", name=MODEL_NAME)\n\nprint(f"Base model fileset: fileset://default/{base_model.name}")\nprint(client.files.list(fileset=MODEL_NAME, workspace="default"))\n\ntime_check = max_wait_time_checker(600, "Model Spec")\nwhile not base_model.spec:\n time_check()\n time.sleep(10)\n base_model = client.models.retrieve(workspace="default", name=MODEL_NAME)\n\n# Clear verbose linear_layers list for cleaner output\nbase_model.spec.linear_layers = None\nprint(f"ModelSpec: {base_model.spec}")\n" }, { "type": "markdown", @@ -126,9 +126,9 @@ }, { "type": "code", - "source": "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.\"\nquestion = \"Who was the first person to walk on the Moon?\"\nmessages = [\n {\"role\": \"user\", \"content\": f\"Based on the following context, answer the question.\\n\\nContext: {context}\\n\\nQuestion: {question}\"}\n]\nresponse = client.inference.gateway.provider.post(\n \"v1/chat/completions\",\n name=deployment_name,\n workspace=\"default\",\n body={\n \"model\": OUTPUT_NAME,\n \"messages\": messages,\n \"temperature\": 0,\n \"max_tokens\": 256,\n }\n)\nprint(\"=\" * 60)\nprint(\"MODEL INFERENCE\")\nprint(\"=\" * 60)\nprint(f\"Question: {question}\")\nprint(f\"Expected: Neil Armstrong\")\nprint(f\"Model output: {response['choices'][0]['message']['content']}\")", + "source": "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.\"\nquestion = \"Who was the first person to walk on the Moon?\"\nmessages = [\n {\"role\": \"user\", \"content\": f\"Based on the following context, answer the question.\\n\\nContext: {context}\\n\\nQuestion: {question}\"}\n]\nINFERENCE_MODEL_NAME = f\"default--{OUTPUT_NAME}\"\nresponse = client.inference.gateway.provider.post(\n \"v1/chat/completions\",\n name=deployment_name,\n workspace=\"default\",\n body={\n \"model\": INFERENCE_MODEL_NAME,\n \"messages\": messages,\n \"temperature\": 0,\n \"max_tokens\": 256,\n }\n)\nprint(\"=\" * 60)\nprint(\"MODEL INFERENCE\")\nprint(\"=\" * 60)\nprint(f\"Question: {question}\")\nprint(f\"Expected: Neil Armstrong\")\nprint(f\"Model output: {response['choices'][0]['message']['content']}\")", "language": "python", - "source_html": "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."\nquestion = "Who was the first person to walk on the Moon?"\nmessages = [\n {"role": "user", "content": f"Based on the following context, answer the question.\\n\\nContext: {context}\\n\\nQuestion: {question}"}\n]\nresponse = client.inference.gateway.provider.post(\n "v1/chat/completions",\n name=deployment_name,\n workspace="default",\n body={\n "model": OUTPUT_NAME,\n "messages": messages,\n "temperature": 0,\n "max_tokens": 256,\n }\n)\nprint("=" * 60)\nprint("MODEL INFERENCE")\nprint("=" * 60)\nprint(f"Question: {question}")\nprint(f"Expected: Neil Armstrong")\nprint(f"Model output: {response['choices'][0]['message']['content']}")\n" + "source_html": "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."\nquestion = "Who was the first person to walk on the Moon?"\nmessages = [\n {"role": "user", "content": f"Based on the following context, answer the question.\\n\\nContext: {context}\\n\\nQuestion: {question}"}\n]\nINFERENCE_MODEL_NAME = f"default--{OUTPUT_NAME}"\nresponse = client.inference.gateway.provider.post(\n "v1/chat/completions",\n name=deployment_name,\n workspace="default",\n body={\n "model": INFERENCE_MODEL_NAME,\n "messages": messages,\n "temperature": 0,\n "max_tokens": 256,\n }\n)\nprint("=" * 60)\nprint("MODEL INFERENCE")\nprint("=" * 60)\nprint(f"Question: {question}")\nprint(f"Expected: Neil Armstrong")\nprint(f"Model output: {response['choices'][0]['message']['content']}")\n" }, { "type": "markdown", diff --git a/docs/fern/components/notebooks/lora-customization-job.ts b/docs/fern/components/notebooks/lora-customization-job.ts index d4b0272cd4..dd7691c96a 100644 --- a/docs/fern/components/notebooks/lora-customization-job.ts +++ b/docs/fern/components/notebooks/lora-customization-job.ts @@ -12,8 +12,8 @@ export default { cells: [ }, { "type": "markdown", - "source": "## Prerequisites\n\nBefore starting this tutorial, ensure you have:\n\n1. **Completed the [Quickstart](../../get-started/quickstart.md)** to install and deploy NeMo Platform locally\n2. **Installed the Python SDK** (PyPI wrapper: `pip install \"nemo-platform[all]\"`; source checkout: run `make bootstrap` from the repository root)\n3. **Installed the `datasets` package** for loading SQuAD: `pip install datasets`\n4. **At least one GPU with CUDA 12.8+**", - "source_html": "

Prerequisites

\n

Before starting this tutorial, ensure you have:

\n
    \n
  1. Completed the Quickstart to install and deploy NeMo Platform locally
    2. \n
  2. Installed the Python SDK (PyPI wrapper: pip install "nemo-platform[all]"; source checkout: run make bootstrap from the repository root)
    3. \n
  3. Installed the datasets package for loading SQuAD: pip install datasets
    4. \n
  4. At least one GPU with CUDA 12.8+
    5. \n
\n" + "source": "## Prerequisites\n\nBefore starting this tutorial, ensure you have:\n\n1. **Completed the [Quickstart](../../get-started/quickstart.md)** to install and deploy NeMo Platform locally\n2. **Installed the Python SDK** (PyPI wrapper: `pip install \"nemo-platform[all]\"`; source checkout: run `make bootstrap` from the repository root)\n3. **Installed the `datasets` package** for loading SQuAD: `pip install datasets`\n4. **At least one GPU with CUDA 13+**", + "source_html": "

Prerequisites

\n

Before starting this tutorial, ensure you have:

\n
    \n
  1. Completed the Quickstart to install and deploy NeMo Platform locally
    2. \n
  2. Installed the Python SDK (PyPI wrapper: pip install "nemo-platform[all]"; source checkout: run make bootstrap from the repository root)
    3. \n
  3. Installed the datasets package for loading SQuAD: pip install datasets
    4. \n
  4. At least one GPU with CUDA 13+
    5. \n
\n" }, { "type": "markdown", @@ -60,8 +60,8 @@ export default { cells: [ }, { "type": "markdown", - "source": "### 4. Secrets Setup\n\nFor Huggingface models that require authentication, create a secret with your HF token. Get a token from [Huggingface Settings](https://huggingface.co/settings/tokens) and accept the model terms.\n\nThis is generally true for LLaMa based models (e.g. [Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct)).\n\n```sh\nexport HF_TOKEN=\n```", - "source_html": "

4. Secrets Setup

\n

For Huggingface models that require authentication, create a secret with your HF token. Get a token from Huggingface Settings and accept the model terms.

\n

This is generally true for LLaMa based models (e.g. Llama-3.2-1B-Instruct).

\n
export HF_TOKEN=<your-huggingface-token>\n
\n" + "source": "### 4. Secrets Setup\n\nFor Hugging Face models that require authentication, create a secret with your HF token. Get a token from [Hugging Face Settings](https://huggingface.co/settings/tokens) and accept the model terms.\n\nThis is generally true for Llama-based models (for example, [Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct)).\n\n```sh\nexport HF_TOKEN=\n```", + "source_html": "

4. Secrets Setup

\n

For Hugging Face models that require authentication, create a secret with your HF token. Get a token from Hugging Face Settings and accept the model terms.

\n

This is generally true for Llama-based models (for example, Llama-3.2-1B-Instruct).

\n
export HF_TOKEN=<your-huggingface-token>\n
\n" }, { "type": "code", @@ -76,9 +76,9 @@ export default { cells: [ }, { "type": "code", - "source": "HF_REPO_ID = \"Qwen/Qwen3-0.6B\"\nMODEL_NAME = \"qwen3-0.6b\"\n\ntry:\n storage = HuggingfaceStorageConfigParam(\n type=\"huggingface\",\n repo_id=HF_REPO_ID,\n repo_type=\"model\",\n )\n if hf_secret:\n storage[\"token_secret\"] = hf_secret.name\n base_model_fs = client.files.filesets.create(\n workspace=\"default\",\n name=MODEL_NAME,\n description=\"Qwen3 0.6b base model from Huggingface\",\n storage=storage,\n cache=True,\n )\nexcept ConflictError:\n base_model_fs = client.files.filesets.retrieve(workspace=\"default\", name=MODEL_NAME)\n\ntry:\n base_model = client.models.create(\n workspace=\"default\",\n name=MODEL_NAME,\n fileset=f\"default/{MODEL_NAME}\",\n trust_remote_code=False,\n )\nexcept ConflictError:\n client.models.update(\n workspace=\"default\",\n name=MODEL_NAME,\n fileset=f\"default/{MODEL_NAME}\",\n trust_remote_code=False,\n )\n base_model = client.models.retrieve(workspace=\"default\", name=MODEL_NAME)\n\nprint(f\"Base model fileset: fileset://default/{base_model.name}\")\nprint(client.files.list(fileset=MODEL_NAME, workspace=\"default\"))\n\ntime_check = max_wait_time_checker(600, \"Model Spec\")\nwhile not base_model.spec:\n time_check()\n time.sleep(10)\n base_model = client.models.retrieve(workspace=\"default\", name=MODEL_NAME)\n\n# Clear verbose linear_layers list for cleaner output\nbase_model.spec.linear_layers = None\nprint(f\"ModelSpec: {base_model.spec}\")", + "source": "HF_REPO_ID = \"Qwen/Qwen3-0.6B\"\nMODEL_NAME = \"qwen3-0.6b\"\n\ntry:\n storage = HuggingfaceStorageConfigParam(\n type=\"huggingface\",\n repo_id=HF_REPO_ID,\n repo_type=\"model\",\n )\n if hf_secret:\n storage[\"token_secret\"] = hf_secret.name\n base_model_fs = client.files.filesets.create(\n workspace=\"default\",\n name=MODEL_NAME,\n description=\"Qwen3 0.6b base model from Hugging Face\",\n storage=storage,\n cache=True,\n )\nexcept ConflictError:\n base_model_fs = client.files.filesets.retrieve(workspace=\"default\", name=MODEL_NAME)\n\ntry:\n base_model = client.models.create(\n workspace=\"default\",\n name=MODEL_NAME,\n fileset=f\"default/{MODEL_NAME}\",\n trust_remote_code=False,\n )\nexcept ConflictError:\n client.models.update(\n workspace=\"default\",\n name=MODEL_NAME,\n fileset=f\"default/{MODEL_NAME}\",\n trust_remote_code=False,\n )\n base_model = client.models.retrieve(workspace=\"default\", name=MODEL_NAME)\n\nprint(f\"Base model fileset: fileset://default/{base_model.name}\")\nprint(client.files.list(fileset=MODEL_NAME, workspace=\"default\"))\n\ntime_check = max_wait_time_checker(600, \"Model Spec\")\nwhile not base_model.spec:\n time_check()\n time.sleep(10)\n base_model = client.models.retrieve(workspace=\"default\", name=MODEL_NAME)\n\n# Clear verbose linear_layers list for cleaner output\nbase_model.spec.linear_layers = None\nprint(f\"ModelSpec: {base_model.spec}\")", "language": "python", - "source_html": "HF_REPO_ID = "Qwen/Qwen3-0.6B"\nMODEL_NAME = "qwen3-0.6b"\n\ntry:\n storage = HuggingfaceStorageConfigParam(\n type="huggingface",\n repo_id=HF_REPO_ID,\n repo_type="model",\n )\n if hf_secret:\n storage["token_secret"] = hf_secret.name\n base_model_fs = client.files.filesets.create(\n workspace="default",\n name=MODEL_NAME,\n description="Qwen3 0.6b base model from Huggingface",\n storage=storage,\n cache=True,\n )\nexcept ConflictError:\n base_model_fs = client.files.filesets.retrieve(workspace="default", name=MODEL_NAME)\n\ntry:\n base_model = client.models.create(\n workspace="default",\n name=MODEL_NAME,\n fileset=f"default/{MODEL_NAME}",\n trust_remote_code=False,\n )\nexcept ConflictError:\n client.models.update(\n workspace="default",\n name=MODEL_NAME,\n fileset=f"default/{MODEL_NAME}",\n trust_remote_code=False,\n )\n base_model = client.models.retrieve(workspace="default", name=MODEL_NAME)\n\nprint(f"Base model fileset: fileset://default/{base_model.name}")\nprint(client.files.list(fileset=MODEL_NAME, workspace="default"))\n\ntime_check = max_wait_time_checker(600, "Model Spec")\nwhile not base_model.spec:\n time_check()\n time.sleep(10)\n base_model = client.models.retrieve(workspace="default", name=MODEL_NAME)\n\n# Clear verbose linear_layers list for cleaner output\nbase_model.spec.linear_layers = None\nprint(f"ModelSpec: {base_model.spec}")\n" + "source_html": "HF_REPO_ID = "Qwen/Qwen3-0.6B"\nMODEL_NAME = "qwen3-0.6b"\n\ntry:\n storage = HuggingfaceStorageConfigParam(\n type="huggingface",\n repo_id=HF_REPO_ID,\n repo_type="model",\n )\n if hf_secret:\n storage["token_secret"] = hf_secret.name\n base_model_fs = client.files.filesets.create(\n workspace="default",\n name=MODEL_NAME,\n description="Qwen3 0.6b base model from Hugging Face",\n storage=storage,\n cache=True,\n )\nexcept ConflictError:\n base_model_fs = client.files.filesets.retrieve(workspace="default", name=MODEL_NAME)\n\ntry:\n base_model = client.models.create(\n workspace="default",\n name=MODEL_NAME,\n fileset=f"default/{MODEL_NAME}",\n trust_remote_code=False,\n )\nexcept ConflictError:\n client.models.update(\n workspace="default",\n name=MODEL_NAME,\n fileset=f"default/{MODEL_NAME}",\n trust_remote_code=False,\n )\n base_model = client.models.retrieve(workspace="default", name=MODEL_NAME)\n\nprint(f"Base model fileset: fileset://default/{base_model.name}")\nprint(client.files.list(fileset=MODEL_NAME, workspace="default"))\n\ntime_check = max_wait_time_checker(600, "Model Spec")\nwhile not base_model.spec:\n time_check()\n time.sleep(10)\n base_model = client.models.retrieve(workspace="default", name=MODEL_NAME)\n\n# Clear verbose linear_layers list for cleaner output\nbase_model.spec.linear_layers = None\nprint(f"ModelSpec: {base_model.spec}")\n" }, { "type": "markdown", @@ -131,9 +131,9 @@ export default { cells: [ }, { "type": "code", - "source": "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.\"\nquestion = \"Who was the first person to walk on the Moon?\"\nmessages = [\n {\"role\": \"user\", \"content\": f\"Based on the following context, answer the question.\\n\\nContext: {context}\\n\\nQuestion: {question}\"}\n]\nresponse = client.inference.gateway.provider.post(\n \"v1/chat/completions\",\n name=deployment_name,\n workspace=\"default\",\n body={\n \"model\": OUTPUT_NAME,\n \"messages\": messages,\n \"temperature\": 0,\n \"max_tokens\": 256,\n }\n)\nprint(\"=\" * 60)\nprint(\"MODEL INFERENCE\")\nprint(\"=\" * 60)\nprint(f\"Question: {question}\")\nprint(f\"Expected: Neil Armstrong\")\nprint(f\"Model output: {response['choices'][0]['message']['content']}\")", + "source": "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.\"\nquestion = \"Who was the first person to walk on the Moon?\"\nmessages = [\n {\"role\": \"user\", \"content\": f\"Based on the following context, answer the question.\\n\\nContext: {context}\\n\\nQuestion: {question}\"}\n]\nINFERENCE_MODEL_NAME = f\"default--{OUTPUT_NAME}\"\nresponse = client.inference.gateway.provider.post(\n \"v1/chat/completions\",\n name=deployment_name,\n workspace=\"default\",\n body={\n \"model\": INFERENCE_MODEL_NAME,\n \"messages\": messages,\n \"temperature\": 0,\n \"max_tokens\": 256,\n }\n)\nprint(\"=\" * 60)\nprint(\"MODEL INFERENCE\")\nprint(\"=\" * 60)\nprint(f\"Question: {question}\")\nprint(f\"Expected: Neil Armstrong\")\nprint(f\"Model output: {response['choices'][0]['message']['content']}\")", "language": "python", - "source_html": "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."\nquestion = "Who was the first person to walk on the Moon?"\nmessages = [\n {"role": "user", "content": f"Based on the following context, answer the question.\\n\\nContext: {context}\\n\\nQuestion: {question}"}\n]\nresponse = client.inference.gateway.provider.post(\n "v1/chat/completions",\n name=deployment_name,\n workspace="default",\n body={\n "model": OUTPUT_NAME,\n "messages": messages,\n "temperature": 0,\n "max_tokens": 256,\n }\n)\nprint("=" * 60)\nprint("MODEL INFERENCE")\nprint("=" * 60)\nprint(f"Question: {question}")\nprint(f"Expected: Neil Armstrong")\nprint(f"Model output: {response['choices'][0]['message']['content']}")\n" + "source_html": "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."\nquestion = "Who was the first person to walk on the Moon?"\nmessages = [\n {"role": "user", "content": f"Based on the following context, answer the question.\\n\\nContext: {context}\\n\\nQuestion: {question}"}\n]\nINFERENCE_MODEL_NAME = f"default--{OUTPUT_NAME}"\nresponse = client.inference.gateway.provider.post(\n "v1/chat/completions",\n name=deployment_name,\n workspace="default",\n body={\n "model": INFERENCE_MODEL_NAME,\n "messages": messages,\n "temperature": 0,\n "max_tokens": 256,\n }\n)\nprint("=" * 60)\nprint("MODEL INFERENCE")\nprint("=" * 60)\nprint(f"Question: {question}")\nprint(f"Expected: Neil Armstrong")\nprint(f"Model output: {response['choices'][0]['message']['content']}")\n" }, { "type": "markdown", diff --git a/docs/fern/components/notebooks/optimize-throughput.json b/docs/fern/components/notebooks/optimize-throughput.json index caa20e705d..c9a35e18eb 100644 --- a/docs/fern/components/notebooks/optimize-throughput.json +++ b/docs/fern/components/notebooks/optimize-throughput.json @@ -7,8 +7,8 @@ }, { "type": "markdown", - "source": "## Prerequisites\n\nBefore starting this tutorial, ensure you have:\n\n1. **Completed the [Quickstart](../../get-started/quickstart.md)** to install and deploy NeMo Platform locally\n2. **Installed the Python SDK** (PyPI wrapper: `pip install \"nemo-platform[all]\"`; source checkout: run `make bootstrap` from the repository root)", - "source_html": "

Prerequisites

\n

Before starting this tutorial, ensure you have:

\n
    \n
  1. Completed the Quickstart to install and deploy NeMo Platform locally
    2. \n
  2. Installed the Python SDK (PyPI wrapper: pip install "nemo-platform[all]"; source checkout: run make bootstrap from the repository root)
    3. \n
\n" + "source": "## Prerequisites\n\nBefore starting this tutorial, ensure you have:\n\n1. **Completed the [Quickstart](../../get-started/quickstart.md)** to install and deploy NeMo Platform locally\n2. **Installed the Python SDK** (PyPI wrapper: `pip install \"nemo-platform[all]\"`; source checkout: run `make bootstrap` from the repository root)\n3. **At least one GPU with CUDA 13+**", + "source_html": "

Prerequisites

\n

Before starting this tutorial, ensure you have:

\n
    \n
  1. Completed the Quickstart to install and deploy NeMo Platform locally
    2. \n
  2. Installed the Python SDK (PyPI wrapper: pip install "nemo-platform[all]"; source checkout: run make bootstrap from the repository root)
    3. \n
  3. At least one GPU with CUDA 13+
    4. \n
\n" }, { "type": "markdown", @@ -56,25 +56,25 @@ }, { "type": "markdown", - "source": "### 3. Secrets Setup\n\nIf you plan to use NGC or HuggingFace models, you will need to configure authentication:\n\n- **NGC models** (`ngc://` URIs): Requires NGC API key\n- **HuggingFace models** (`hf://` URIs): Requires HF token for gated/private models\n\n\nConfigure these as secrets in your platform. Refer to [Managing Secrets](../../get-started/concepts/manage-secrets.md) for detailed instructions.\n\nGet your credentials to access base models:\n- [NGC API Key](https://ngc.nvidia.com/) (Setup → Generate API Key)\n- [HuggingFace Token](https://huggingface.co/settings/tokens) (Create token with Read access)\n\n\n---\n\n#### Quick Setup Example\n\nThis tutorial uses the [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main) model from HuggingFace. Ensure that you have sufficient permissions to download the model. If you cannot access the files on the [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main) Hugging Face page, request access.\n\n**HuggingFace Authentication:**\n- For gated models (Llama, Gemma), you must provide a HuggingFace token via the `token_secret` parameter\n- Get your token from [HuggingFace Settings](https://huggingface.co/settings/tokens) (requires Read access)\n- Accept the model's terms on the HuggingFace model page before using it. Example: [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main)\n- For public models, you can omit the `token_secret` parameter when creating a fileset for the model in the next step.", - "source_html": "

3. Secrets Setup

\n

If you plan to use NGC or HuggingFace models, you will need to configure authentication:

\n
    \n
  • NGC models (ngc:// URIs): Requires NGC API key
    • \n
  • HuggingFace models (hf:// URIs): Requires HF token for gated/private models
    • \n
\n

Configure these as secrets in your platform. Refer to Managing Secrets for detailed instructions.

\n

Get your credentials to access base models:

\n\n
\n

Quick Setup Example

\n

This tutorial uses the meta-llama/Llama-3.2-1B-Instruct model from HuggingFace. Ensure that you have sufficient permissions to download the model. If you cannot access the files on the meta-llama/Llama-3.2-1B-Instruct Hugging Face page, request access.

\n

HuggingFace Authentication:

\n
    \n
  • For gated models (Llama, Gemma), you must provide a HuggingFace token via the token_secret parameter
    • \n
  • Get your token from HuggingFace Settings (requires Read access)
    • \n
  • Accept the model's terms on the HuggingFace model page before using it. Example: meta-llama/Llama-3.2-1B-Instruct
    • \n
  • For public models, you can omit the token_secret parameter when creating a fileset for the model in the next step.
    • \n
\n" + "source": "### 3. Secrets Setup\n\nIf you plan to use NGC or Hugging Face models, you will need to configure authentication:\n\n- **NGC models** (`ngc://` URIs): Requires NGC API key\n- **Hugging Face models** (`hf://` URIs): Requires HF token for gated/private models\n\n\nConfigure these as secrets in your platform. Refer to [Managing Secrets](../../get-started/concepts/manage-secrets.md) for detailed instructions.\n\nGet your credentials to access base models:\n- [NGC API Key](https://ngc.nvidia.com/) (Setup → Generate API Key)\n- [Hugging Face Token](https://huggingface.co/settings/tokens) (Create token with Read access)\n\n\n---\n\n#### Quick Setup Example\n\nThis tutorial uses the [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main) model from Hugging Face. Ensure that you have sufficient permissions to download the model. If you cannot access the files on the [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main) Hugging Face page, request access.\n\n**Hugging Face Authentication:**\n- For gated models (Llama, Gemma), you must provide a Hugging Face token via the `token_secret` parameter\n- Get your token from [Hugging Face Settings](https://huggingface.co/settings/tokens) (requires Read access)\n- Accept the model's terms on the Hugging Face model page before using it. Example: [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main)\n- For public models, you can omit the `token_secret` parameter when creating a fileset for the model in the next step.", + "source_html": "

3. Secrets Setup

\n

If you plan to use NGC or Hugging Face models, you will need to configure authentication:

\n
    \n
  • NGC models (ngc:// URIs): Requires NGC API key
    • \n
  • Hugging Face models (hf:// URIs): Requires HF token for gated/private models
    • \n
\n

Configure these as secrets in your platform. Refer to Managing Secrets for detailed instructions.

\n

Get your credentials to access base models:

\n\n
\n

Quick Setup Example

\n

This tutorial uses the meta-llama/Llama-3.2-1B-Instruct model from Hugging Face. Ensure that you have sufficient permissions to download the model. If you cannot access the files on the meta-llama/Llama-3.2-1B-Instruct Hugging Face page, request access.

\n

Hugging Face Authentication:

\n
    \n
  • For gated models (Llama, Gemma), you must provide a Hugging Face token via the token_secret parameter
    • \n
  • Get your token from Hugging Face Settings (requires Read access)
    • \n
  • Accept the model's terms on the Hugging Face model page before using it. Example: meta-llama/Llama-3.2-1B-Instruct
    • \n
  • For public models, you can omit the token_secret parameter when creating a fileset for the model in the next step.
    • \n
\n" }, { "type": "code", - "source": "# Export the HF_TOKEN and NGC_API_KEY environment variables if they are not already set\nHF_TOKEN = os.getenv(\"HF_TOKEN\")\nNGC_API_KEY = os.getenv(\"NGC_API_KEY\")\n\n\ndef create_or_get_secret(name: str, value: str | None, label: str):\n if not value:\n raise ValueError(f\"{label} environment variable is not set. Set it and try again.\")\n try:\n secret = client.secrets.create(\n name=name,\n workspace=\"default\",\n value=value,\n )\n print(f\"Created secret: {name}\")\n return secret\n except ConflictError:\n print(f\"Secret '{name}' already exists, continuing...\")\n return client.secrets.retrieve(name=name, workspace=\"default\")\n\n\n# Create HuggingFace token secret\nhf_secret = create_or_get_secret(\"hf-token\", HF_TOKEN, \"HF_TOKEN\")\nprint(\"HF_TOKEN secret:\")\nprint(hf_secret.model_dump_json(indent=2))\n\n# Create NGC API key secret\n# Uncomment the line below if you have NGC API Key and want to finetune NGC models\n# ngc_api_key = create_or_get_secret(\"ngc-api-key\", NGC_API_KEY, \"NGC_API_KEY\")", + "source": "# Export the HF_TOKEN and NGC_API_KEY environment variables if they are not already set\nHF_TOKEN = os.getenv(\"HF_TOKEN\")\nNGC_API_KEY = os.getenv(\"NGC_API_KEY\")\n\n\ndef create_or_get_secret(name: str, value: str | None, label: str):\n if not value:\n raise ValueError(f\"{label} environment variable is not set. Set it and try again.\")\n try:\n secret = client.secrets.create(\n name=name,\n workspace=\"default\",\n value=value,\n )\n print(f\"Created secret: {name}\")\n return secret\n except ConflictError:\n print(f\"Secret '{name}' already exists, continuing...\")\n return client.secrets.retrieve(name=name, workspace=\"default\")\n\n\n# Create Hugging Face token secret\nhf_secret = create_or_get_secret(\"hf-token\", HF_TOKEN, \"HF_TOKEN\")\nprint(\"HF_TOKEN secret:\")\nprint(hf_secret.model_dump_json(indent=2))\n\n# Create NGC API key secret\n# Uncomment the line below if you have NGC API Key and want to finetune NGC models\n# ngc_api_key = create_or_get_secret(\"ngc-api-key\", NGC_API_KEY, \"NGC_API_KEY\")", "language": "python", - "source_html": "# Export the HF_TOKEN and NGC_API_KEY environment variables if they are not already set\nHF_TOKEN = os.getenv("HF_TOKEN")\nNGC_API_KEY = os.getenv("NGC_API_KEY")\n\n\ndef create_or_get_secret(name: str, value: str | None, label: str):\n if not value:\n raise ValueError(f"{label} environment variable is not set. Set it and try again.")\n try:\n secret = client.secrets.create(\n name=name,\n workspace="default",\n value=value,\n )\n print(f"Created secret: {name}")\n return secret\n except ConflictError:\n print(f"Secret '{name}' already exists, continuing...")\n return client.secrets.retrieve(name=name, workspace="default")\n\n\n# Create HuggingFace token secret\nhf_secret = create_or_get_secret("hf-token", HF_TOKEN, "HF_TOKEN")\nprint("HF_TOKEN secret:")\nprint(hf_secret.model_dump_json(indent=2))\n\n# Create NGC API key secret\n# Uncomment the line below if you have NGC API Key and want to finetune NGC models\n# ngc_api_key = create_or_get_secret("ngc-api-key", NGC_API_KEY, "NGC_API_KEY")\n" + "source_html": "# Export the HF_TOKEN and NGC_API_KEY environment variables if they are not already set\nHF_TOKEN = os.getenv("HF_TOKEN")\nNGC_API_KEY = os.getenv("NGC_API_KEY")\n\n\ndef create_or_get_secret(name: str, value: str | None, label: str):\n if not value:\n raise ValueError(f"{label} environment variable is not set. Set it and try again.")\n try:\n secret = client.secrets.create(\n name=name,\n workspace="default",\n value=value,\n )\n print(f"Created secret: {name}")\n return secret\n except ConflictError:\n print(f"Secret '{name}' already exists, continuing...")\n return client.secrets.retrieve(name=name, workspace="default")\n\n\n# Create Hugging Face token secret\nhf_secret = create_or_get_secret("hf-token", HF_TOKEN, "HF_TOKEN")\nprint("HF_TOKEN secret:")\nprint(hf_secret.model_dump_json(indent=2))\n\n# Create NGC API key secret\n# Uncomment the line below if you have NGC API Key and want to finetune NGC models\n# ngc_api_key = create_or_get_secret("ngc-api-key", NGC_API_KEY, "NGC_API_KEY")\n" }, { "type": "markdown", - "source": "### 4. Create Base Model FileSet\n\nCreate a fileset pointing to the [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main) model on HuggingFace. This step creates a pointer to the model on Hugging Face and does not download it. The model is downloaded at job creation time.\n\nNote: for public models, you can omit the `token_secret` parameter when creating a model fileset.", - "source_html": "

4. Create Base Model FileSet

\n

Create a fileset pointing to the meta-llama/Llama-3.2-1B-Instruct model on HuggingFace. This step creates a pointer to the model on Hugging Face and does not download it. The model is downloaded at job creation time.

\n

Note: for public models, you can omit the token_secret parameter when creating a model fileset.

\n" + "source": "### 4. Create Base Model FileSet\n\nCreate a fileset pointing to the [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main) model on Hugging Face. This step creates a pointer to the model on Hugging Face and does not download it. The model is downloaded at job creation time.\n\nNote: for public models, you can omit the `token_secret` parameter when creating a model fileset.", + "source_html": "

4. Create Base Model FileSet

\n

Create a fileset pointing to the meta-llama/Llama-3.2-1B-Instruct model on Hugging Face. This step creates a pointer to the model on Hugging Face and does not download it. The model is downloaded at job creation time.

\n

Note: for public models, you can omit the token_secret parameter when creating a model fileset.

\n" }, { "type": "code", - "source": "import time\nfrom nemo_platform.types.files import HuggingfaceStorageConfigParam\n\nHF_REPO_ID = \"meta-llama/Llama-3.2-1B-Instruct\"\nMODEL_NAME = \"llama-3-2-1b-base\"\n\n# Ensure you have a HuggingFace token secret created\n# Create a fileset pointing to the desired HuggingFace model\ntry:\n base_model_fs = client.files.filesets.create(\n workspace=\"default\",\n name=MODEL_NAME,\n description=\"Llama 3.2 1B base model from HuggingFace\",\n storage=HuggingfaceStorageConfigParam(\n type=\"huggingface\",\n # repo_id is the full model name from Hugging Face\n repo_id=HF_REPO_ID,\n repo_type=\"model\",\n # we use the secret created in the previous step\n token_secret=hf_secret.name\n )\n )\n print(f\"Created base model fileset: {MODEL_NAME}\")\nexcept ConflictError:\n print(f\"Base model fileset already exists. Skipping creation.\")\n base_model_fs = client.files.filesets.retrieve(\n workspace=\"default\",\n name=MODEL_NAME,\n )\n\n# Create the Model Entity representation.\ntry:\n base_model = client.models.create(\n workspace=\"default\",\n name=MODEL_NAME,\n fileset=f\"default/{MODEL_NAME}\",\n )\n print(f\"Created Model Entity: {MODEL_NAME}\")\nexcept ConflictError:\n print(f\"Base model already exists. Updating fileset if different.\")\n base_model = client.models.update(\n workspace=\"default\",\n name=MODEL_NAME,\n fileset=f\"default/{MODEL_NAME}\",\n )\n\nprint(f\"\\nBase model fileset: fileset://default/{base_model.name}\")\nprint(\"Base model fileset files list:\")\nprint(json.dumps([f.model_dump() for f in client.files.list(fileset=MODEL_NAME, workspace=\"default\").data], indent=2))\n\n# Wait for ModelSpec to be populated from the checkpoint\nprint(\"\\nWaiting for ModelSpec to be populated...\")\nSPEC_TIMEOUT_SECONDS = 120\nspec_start = time.time()\nwhile not base_model.spec:\n if time.time() - spec_start > SPEC_TIMEOUT_SECONDS:\n raise TimeoutError(f\"ModelSpec not populated within {SPEC_TIMEOUT_SECONDS} seconds\")\n time.sleep(2)\n base_model = client.models.retrieve(\n workspace=\"default\",\n name=MODEL_NAME,\n )\n\nprint(f\"ModelSpec populated: {base_model.spec}\")", + "source": "import time\nfrom nemo_platform.types.files import HuggingfaceStorageConfigParam\n\nHF_REPO_ID = \"meta-llama/Llama-3.2-1B-Instruct\"\nMODEL_NAME = \"llama-3-2-1b-base\"\n\n# Ensure you have a Hugging Face token secret created\n# Create a fileset pointing to the desired Hugging Face model\ntry:\n base_model_fs = client.files.filesets.create(\n workspace=\"default\",\n name=MODEL_NAME,\n description=\"Llama 3.2 1B base model from Hugging Face\",\n storage=HuggingfaceStorageConfigParam(\n type=\"huggingface\",\n # repo_id is the full model name from Hugging Face\n repo_id=HF_REPO_ID,\n repo_type=\"model\",\n # we use the secret created in the previous step\n token_secret=hf_secret.name\n )\n )\n print(f\"Created base model fileset: {MODEL_NAME}\")\nexcept ConflictError:\n print(f\"Base model fileset already exists. Skipping creation.\")\n base_model_fs = client.files.filesets.retrieve(\n workspace=\"default\",\n name=MODEL_NAME,\n )\n\n# Create the Model Entity representation.\ntry:\n base_model = client.models.create(\n workspace=\"default\",\n name=MODEL_NAME,\n fileset=f\"default/{MODEL_NAME}\",\n )\n print(f\"Created Model Entity: {MODEL_NAME}\")\nexcept ConflictError:\n print(f\"Base model already exists. Updating fileset if different.\")\n base_model = client.models.update(\n workspace=\"default\",\n name=MODEL_NAME,\n fileset=f\"default/{MODEL_NAME}\",\n )\n\nprint(f\"\\nBase model fileset: fileset://default/{base_model.name}\")\nprint(\"Base model fileset files list:\")\nprint(json.dumps([f.model_dump() for f in client.files.list(fileset=MODEL_NAME, workspace=\"default\").data], indent=2))\n\n# Wait for ModelSpec to be populated from the checkpoint\nprint(\"\\nWaiting for ModelSpec to be populated...\")\nSPEC_TIMEOUT_SECONDS = 120\nspec_start = time.time()\nwhile not base_model.spec:\n if time.time() - spec_start > SPEC_TIMEOUT_SECONDS:\n raise TimeoutError(f\"ModelSpec not populated within {SPEC_TIMEOUT_SECONDS} seconds\")\n time.sleep(2)\n base_model = client.models.retrieve(\n workspace=\"default\",\n name=MODEL_NAME,\n )\n\nprint(f\"ModelSpec populated: {base_model.spec}\")", "language": "python", - "source_html": "import time\nfrom nemo_platform.types.files import HuggingfaceStorageConfigParam\n\nHF_REPO_ID = "meta-llama/Llama-3.2-1B-Instruct"\nMODEL_NAME = "llama-3-2-1b-base"\n\n# Ensure you have a HuggingFace token secret created\n# Create a fileset pointing to the desired HuggingFace model\ntry:\n base_model_fs = client.files.filesets.create(\n workspace="default",\n name=MODEL_NAME,\n description="Llama 3.2 1B base model from HuggingFace",\n storage=HuggingfaceStorageConfigParam(\n type="huggingface",\n # repo_id is the full model name from Hugging Face\n repo_id=HF_REPO_ID,\n repo_type="model",\n # we use the secret created in the previous step\n token_secret=hf_secret.name\n )\n )\n print(f"Created base model fileset: {MODEL_NAME}")\nexcept ConflictError:\n print(f"Base model fileset already exists. Skipping creation.")\n base_model_fs = client.files.filesets.retrieve(\n workspace="default",\n name=MODEL_NAME,\n )\n\n# Create the Model Entity representation.\ntry:\n base_model = client.models.create(\n workspace="default",\n name=MODEL_NAME,\n fileset=f"default/{MODEL_NAME}",\n )\n print(f"Created Model Entity: {MODEL_NAME}")\nexcept ConflictError:\n print(f"Base model already exists. Updating fileset if different.")\n base_model = client.models.update(\n workspace="default",\n name=MODEL_NAME,\n fileset=f"default/{MODEL_NAME}",\n )\n\nprint(f"\\nBase model fileset: fileset://default/{base_model.name}")\nprint("Base model fileset files list:")\nprint(json.dumps([f.model_dump() for f in client.files.list(fileset=MODEL_NAME, workspace="default").data], indent=2))\n\n# Wait for ModelSpec to be populated from the checkpoint\nprint("\\nWaiting for ModelSpec to be populated...")\nSPEC_TIMEOUT_SECONDS = 120\nspec_start = time.time()\nwhile not base_model.spec:\n if time.time() - spec_start > SPEC_TIMEOUT_SECONDS:\n raise TimeoutError(f"ModelSpec not populated within {SPEC_TIMEOUT_SECONDS} seconds")\n time.sleep(2)\n base_model = client.models.retrieve(\n workspace="default",\n name=MODEL_NAME,\n )\n\nprint(f"ModelSpec populated: {base_model.spec}")\n" + "source_html": "import time\nfrom nemo_platform.types.files import HuggingfaceStorageConfigParam\n\nHF_REPO_ID = "meta-llama/Llama-3.2-1B-Instruct"\nMODEL_NAME = "llama-3-2-1b-base"\n\n# Ensure you have a Hugging Face token secret created\n# Create a fileset pointing to the desired Hugging Face model\ntry:\n base_model_fs = client.files.filesets.create(\n workspace="default",\n name=MODEL_NAME,\n description="Llama 3.2 1B base model from Hugging Face",\n storage=HuggingfaceStorageConfigParam(\n type="huggingface",\n # repo_id is the full model name from Hugging Face\n repo_id=HF_REPO_ID,\n repo_type="model",\n # we use the secret created in the previous step\n token_secret=hf_secret.name\n )\n )\n print(f"Created base model fileset: {MODEL_NAME}")\nexcept ConflictError:\n print(f"Base model fileset already exists. Skipping creation.")\n base_model_fs = client.files.filesets.retrieve(\n workspace="default",\n name=MODEL_NAME,\n )\n\n# Create the Model Entity representation.\ntry:\n base_model = client.models.create(\n workspace="default",\n name=MODEL_NAME,\n fileset=f"default/{MODEL_NAME}",\n )\n print(f"Created Model Entity: {MODEL_NAME}")\nexcept ConflictError:\n print(f"Base model already exists. Updating fileset if different.")\n base_model = client.models.update(\n workspace="default",\n name=MODEL_NAME,\n fileset=f"default/{MODEL_NAME}",\n )\n\nprint(f"\\nBase model fileset: fileset://default/{base_model.name}")\nprint("Base model fileset files list:")\nprint(json.dumps([f.model_dump() for f in client.files.list(fileset=MODEL_NAME, workspace="default").data], indent=2))\n\n# Wait for ModelSpec to be populated from the checkpoint\nprint("\\nWaiting for ModelSpec to be populated...")\nSPEC_TIMEOUT_SECONDS = 120\nspec_start = time.time()\nwhile not base_model.spec:\n if time.time() - spec_start > SPEC_TIMEOUT_SECONDS:\n raise TimeoutError(f"ModelSpec not populated within {SPEC_TIMEOUT_SECONDS} seconds")\n time.sleep(2)\n base_model = client.models.retrieve(\n workspace="default",\n name=MODEL_NAME,\n )\n\nprint(f"ModelSpec populated: {base_model.spec}")\n" }, { "type": "markdown", @@ -89,8 +89,8 @@ }, { "type": "markdown", - "source": "### 6. Track Finetuning Progress\n\nA training job contains multiple steps: \n- Model and dataset downloading\n- Finetuning where LoRA adapter weights are trained\n- Creating a fileset entry for the finetuned model\n- Finetuned weights uploading\n\nThe elapsed time printed below reflects progress of the entire job. We compare the time taken by the finetuning step for both jobs in the last section of this tutorial.", - "source_html": "

6. Track Finetuning Progress

\n

A training job contains multiple steps:

\n
    \n
  • Model and dataset downloading
    • \n
  • Finetuning where LoRA adapter weights are trained
    • \n
  • Creating a fileset entry for the finetuned model
    • \n
  • Finetuned weights uploading
    • \n
\n

The elapsed time printed below reflects progress of the entire job. We compare the time taken by the finetuning step for both jobs in the last section of this tutorial.

\n" + "source": "### 6. Track Fine-Tuning Progress\n\nA training job contains multiple steps: \n- Model and dataset downloading\n- Fine-tuning where LoRA adapter weights are trained\n- Creating a fileset entry for the fine-tuned model\n- Fine-tuned weights uploading\n\nThe elapsed time printed below reflects progress of the entire job. We compare the time taken by the fine-tuning step for both jobs in the last section of this tutorial.", + "source_html": "

6. Track Fine-Tuning Progress

\n

A training job contains multiple steps:

\n
    \n
  • Model and dataset downloading
    • \n
  • Fine-tuning where LoRA adapter weights are trained
    • \n
  • Creating a fileset entry for the fine-tuned model
    • \n
  • Fine-tuned weights uploading
    • \n
\n

The elapsed time printed below reflects progress of the entire job. We compare the time taken by the fine-tuning step for both jobs in the last section of this tutorial.

\n" }, { "type": "markdown", @@ -105,14 +105,14 @@ }, { "type": "markdown", - "source": "#### Monitor the Job Until Completion\n\nThe cell below polls the job status every 10 seconds and renders a live dashboard with validation loss, GPU VRAM usage, and GPU utilization charts. The charts appear empty at first while the model and dataset download; training metrics and GPU activity populate after the finetuning step begins.\n\n> **Note:** This is additional code. You can also use the Weights & Biases or MLflow integrations.", - "source_html": "

Monitor the Job Until Completion

\n

The cell below polls the job status every 10 seconds and renders a live dashboard with validation loss, GPU VRAM usage, and GPU utilization charts. The charts appear empty at first while the model and dataset download; training metrics and GPU activity populate after the finetuning step begins.

\n
\n

Note: This is additional code. You can also use the Weights & Biases or MLflow integrations.

\n
\n" + "source": "#### Monitor the Job Until Completion\n\nThe cell below polls the job status every 10 seconds and renders a live dashboard with validation loss, GPU VRAM usage, and GPU utilization charts. The charts appear empty at first while the model and dataset download; training metrics and GPU activity populate after the fine-tuning step begins.\n\n> **Note:** This is additional code. You can also use the Weights & Biases or MLflow integrations.", + "source_html": "

Monitor the Job Until Completion

\n

The cell below polls the job status every 10 seconds and renders a live dashboard with validation loss, GPU VRAM usage, and GPU utilization charts. The charts appear empty at first while the model and dataset download; training metrics and GPU activity populate after the fine-tuning step begins.

\n
\n

Note: This is additional code. You can also use the Weights & Biases or MLflow integrations.

\n
\n" }, { "type": "code", - "source": "import time\nfrom typing import cast\nfrom IPython.display import clear_output\nfrom nemo_platform.types.shared import PlatformJobStatusResponse\n\n# Timeout set to 30 minutes to accommodate typical LoRA training duration for this dataset size.\n# Actual training time will vary based on hardware, model size, and dataset complexity.\nTIMEOUT_SECONDS = 30 * 60 # 30 minutes\nVAL_LOSS_KEY = \"val_loss\"\nTRAIN_LOSS_KEY = \"loss\"\n\n# ---------------------------------------------------------------------------\n# Job polling with live dashboard\n# ---------------------------------------------------------------------------\n\ndef wait_for_job(\n workspace: str,\n job_name: str,\n timeout: int = TIMEOUT_SECONDS,\n poll_interval: int = 10,\n val_loss_key: str = VAL_LOSS_KEY,\n train_loss_key: str = TRAIN_LOSS_KEY,\n) -> PlatformJobStatusResponse:\n \"\"\"\n Poll job status until completed, failed, cancelled, or timeout.\n Displays a live dashboard with loss curves and GPU metrics.\n\n Args:\n workspace: The workspace where the job is running.\n job_name: The name of the job to monitor.\n timeout: Maximum time to wait in seconds (default: 30 minutes).\n poll_interval: Time between status checks in seconds (default: 10).\n\n Returns:\n The final job status response.\n \"\"\"\n start_time = time.time()\n\n # Time-series accumulators required for plotting\n elapsed_mins: list[float] = []\n val_losses: list[float | None] = []\n train_losses: list[float | None] = []\n vram_history: list[list[float]] = []\n util_history: list[list[float]] = []\n\n while True:\n elapsed = time.time() - start_time\n elapsed_min = elapsed / 60\n\n # Check for timeout\n if elapsed > timeout:\n error_message = f\"Timeout reached after {elapsed_min:.1f} minutes\"\n print(f\"\\n{error_message}\")\n print(\"Job did not complete within the timeout period.\")\n raise Exception(error_message)\n\n status = client.jobs.get_status(name=job_name, workspace=workspace)\n\n # -- Extract training progress from nested steps structure --\n step: int | None = None\n max_steps: int | None = None\n training_phase: str | None = None\n val_loss: float | None = None\n train_loss: float | None = None\n current_step_name: str | None = None\n current_step_phase: str | None = None\n\n for job_step in status.steps or []:\n # Track the current active step name and phase for progress display\n if job_step.tasks:\n task = job_step.tasks[0]\n td = task.status_details or {}\n phase = cast(str, td.get(\"phase\", \"\"))\n # Update current step if it's active or pending (not completed)\n if job_step.status in (\"active\", \"pending\"):\n current_step_name = job_step.name\n current_step_phase = phase or \"started\"\n\n if job_step.name == \"training\":\n for task in job_step.tasks or []:\n td = task.status_details or {}\n step = cast(int, td[\"step\"]) if \"step\" in td else None\n max_steps = cast(int, td[\"max_steps\"]) if \"max_steps\" in td else None\n training_phase = cast(str, td[\"phase\"]) if \"phase\" in td else None\n raw_val_loss = td.get(val_loss_key)\n val_loss = float(raw_val_loss) if raw_val_loss is not None else None\n raw_train_loss = td.get(train_loss_key)\n train_loss = float(raw_train_loss) if raw_train_loss is not None else None\n break\n break\n\n if val_loss is None:\n raw_val_loss = (status.status_details or {}).get(val_loss_key)\n val_loss = float(raw_val_loss) if raw_val_loss is not None else None\n if train_loss is None:\n raw_train_loss = (status.status_details or {}).get(train_loss_key)\n train_loss = float(raw_train_loss) if raw_train_loss is not None else None\n\n # -- Collect GPU snapshot --\n vram_pcts, util_pcts = _get_gpu_snapshot()\n\n # -- Append to accumulators used for the plots --\n elapsed_mins.append(elapsed_min)\n val_losses.append(val_loss)\n train_losses.append(train_loss)\n vram_history.append(vram_pcts)\n util_history.append(util_pcts)\n\n # -- Build status strings --\n status_str = f\"Status: {status.status}\"\n if step is not None and max_steps is not None:\n pct = step / max_steps * 100\n step_str = f\"Step {step}/{max_steps} ({pct:.0f}%)\"\n if training_phase:\n step_str += f\" - {training_phase}\"\n else:\n if current_step_name and current_step_phase:\n step_str = f\"{current_step_name} - {current_step_phase}\"\n elif current_step_name:\n step_str = f\"{current_step_name}\"\n else:\n step_str = \"Waiting for training to start...\"\n elapsed_str = f\"Elapsed: {elapsed_min:.1f} min\"\n\n # -- Redraw dashboard --\n clear_output(wait=True)\n _draw_dashboard(\n elapsed_mins, val_losses, train_losses,\n vram_history, util_history,\n job_name, status_str, step_str, elapsed_str,\n )\n\n # -- Check terminal conditions --\n if status.status.lower() == \"completed\":\n # Redraw dashboard one final time with \"completed\" status\n status_str = f\"Status: {status.status}\"\n if step is not None and max_steps is not None:\n step_str = f\"Step {max_steps}/{max_steps} (100%)\"\n clear_output(wait=True)\n _draw_dashboard(\n elapsed_mins, val_losses, train_losses,\n vram_history, util_history,\n job_name, status_str, step_str, elapsed_str,\n )\n print(f\"\\nJob completed in {elapsed_min:.1f} minutes ({elapsed:.0f}s)\")\n return status\n elif status.status.lower() in (\"failed\", \"cancelled\", \"error\"):\n print(f\"\\nJob finished with status: {status.status}\")\n print(f\"Total time elapsed: {elapsed_min:.1f} minutes ({elapsed:.0f}s)\")\n\n # Print error details from the job level\n if status.error_details:\n error_msg = status.error_details.get(\"message\", \"\")\n if error_msg:\n print(f\"\\nError: {error_msg}\")\n\n # Find and print error details from the failed step/task\n for job_step in status.steps or []:\n if job_step.status == \"error\":\n print(f\"\\nFailed step: {job_step.name}\")\n if job_step.error_details:\n step_error = job_step.error_details.get(\"message\", \"\")\n if step_error:\n print(f\"Step error: {step_error}\")\n # Get error_stack from the failed task\n for task in job_step.tasks or []:\n if task.status == \"error\" and hasattr(task, \"error_stack\") and task.error_stack:\n print(f\"\\nError stack trace:\\n{task.error_stack}\")\n elif task.status == \"error\" and task.error_details:\n task_error = task.error_details.get(\"message\", \"\")\n if task_error:\n print(f\"Task error: {task_error}\")\n break\n\n raise Exception(f\"Job finished with status: {status.status}\")\n\n time.sleep(poll_interval)\n\n\n# Wait for the job to complete\njob_with_sequence_packing_status = wait_for_job(\n workspace=\"default\",\n job_name=job_with_sequence_packing.job.name,\n timeout=TIMEOUT_SECONDS,\n)\n\npacked_val_loss = (job_with_sequence_packing_status.status_details or {}).get(\"val_loss\")\nif packed_val_loss is not None:\n print(f\"Validation loss: {float(packed_val_loss):.2f}\")\nelse:\n print(\"Validation loss: not reported in job status\")", + "source": "import time\nfrom typing import cast\nfrom IPython.display import clear_output\nfrom nemo_platform.types.shared import PlatformJobStatusResponse\n\n# Timeout set to 30 minutes to accommodate typical LoRA training duration for this dataset size.\n# Actual training time will vary based on hardware, model size, and dataset complexity.\nTIMEOUT_SECONDS = 30 * 60 # 30 minutes\nVAL_LOSS_KEY = \"val_loss\"\nTRAIN_LOSS_KEY = \"train_loss\"\n\n\ndef get_training_metric(\n status: PlatformJobStatusResponse,\n metric_key: str,\n) -> float | None:\n \"\"\"Return a metric reported by a task in the training step.\"\"\"\n for job_step in status.steps or []:\n if job_step.name == \"training\":\n for task in job_step.tasks or []:\n value = (task.status_details or {}).get(metric_key)\n if value is not None:\n return float(value)\n return None\n\n\n# ---------------------------------------------------------------------------\n# Job polling with live dashboard\n# ---------------------------------------------------------------------------\n\ndef wait_for_job(\n workspace: str,\n job_name: str,\n timeout: int = TIMEOUT_SECONDS,\n poll_interval: int = 10,\n val_loss_key: str = VAL_LOSS_KEY,\n train_loss_key: str = TRAIN_LOSS_KEY,\n) -> PlatformJobStatusResponse:\n \"\"\"\n Poll job status until completed, failed, cancelled, or timeout.\n Displays a live dashboard with loss curves and GPU metrics.\n\n Args:\n workspace: The workspace where the job is running.\n job_name: The name of the job to monitor.\n timeout: Maximum time to wait in seconds (default: 30 minutes).\n poll_interval: Time between status checks in seconds (default: 10).\n\n Returns:\n The final job status response.\n \"\"\"\n start_time = time.time()\n\n # Time-series accumulators required for plotting\n elapsed_mins: list[float] = []\n val_losses: list[float | None] = []\n train_losses: list[float | None] = []\n vram_history: list[list[float]] = []\n util_history: list[list[float]] = []\n\n while True:\n elapsed = time.time() - start_time\n elapsed_min = elapsed / 60\n\n # Check for timeout\n if elapsed > timeout:\n error_message = f\"Timeout reached after {elapsed_min:.1f} minutes\"\n print(f\"\\n{error_message}\")\n print(\"Job did not complete within the timeout period.\")\n raise Exception(error_message)\n\n status = client.jobs.get_status(name=job_name, workspace=workspace)\n\n # -- Extract training progress from nested steps structure --\n step: int | None = None\n max_steps: int | None = None\n training_phase: str | None = None\n val_loss: float | None = None\n train_loss: float | None = None\n current_step_name: str | None = None\n current_step_phase: str | None = None\n\n for job_step in status.steps or []:\n # Track the current active step name and phase for progress display\n if job_step.tasks:\n task = job_step.tasks[0]\n td = task.status_details or {}\n phase = cast(str, td.get(\"phase\", \"\"))\n # Update current step if it's active or pending (not completed)\n if job_step.status in (\"active\", \"pending\"):\n current_step_name = job_step.name\n current_step_phase = phase or \"started\"\n\n if job_step.name == \"training\":\n for task in job_step.tasks or []:\n td = task.status_details or {}\n step = cast(int, td[\"step\"]) if \"step\" in td else None\n max_steps = cast(int, td[\"max_steps\"]) if \"max_steps\" in td else None\n training_phase = cast(str, td[\"phase\"]) if \"phase\" in td else None\n raw_val_loss = td.get(val_loss_key)\n val_loss = float(raw_val_loss) if raw_val_loss is not None else None\n raw_train_loss = td.get(train_loss_key)\n train_loss = float(raw_train_loss) if raw_train_loss is not None else None\n break\n break\n\n if val_loss is None:\n raw_val_loss = (status.status_details or {}).get(val_loss_key)\n val_loss = float(raw_val_loss) if raw_val_loss is not None else None\n if train_loss is None:\n raw_train_loss = (status.status_details or {}).get(train_loss_key)\n train_loss = float(raw_train_loss) if raw_train_loss is not None else None\n\n # -- Collect GPU snapshot --\n vram_pcts, util_pcts = _get_gpu_snapshot()\n\n # -- Append to accumulators used for the plots --\n elapsed_mins.append(elapsed_min)\n val_losses.append(val_loss)\n train_losses.append(train_loss)\n vram_history.append(vram_pcts)\n util_history.append(util_pcts)\n\n # -- Build status strings --\n status_str = f\"Status: {status.status}\"\n if step is not None and max_steps is not None:\n pct = step / max_steps * 100\n step_str = f\"Step {step}/{max_steps} ({pct:.0f}%)\"\n if training_phase:\n step_str += f\" - {training_phase}\"\n else:\n if current_step_name and current_step_phase:\n step_str = f\"{current_step_name} - {current_step_phase}\"\n elif current_step_name:\n step_str = f\"{current_step_name}\"\n else:\n step_str = \"Waiting for training to start...\"\n elapsed_str = f\"Elapsed: {elapsed_min:.1f} min\"\n\n # -- Redraw dashboard --\n clear_output(wait=True)\n _draw_dashboard(\n elapsed_mins, val_losses, train_losses,\n vram_history, util_history,\n job_name, status_str, step_str, elapsed_str,\n )\n\n # -- Check terminal conditions --\n if status.status.lower() == \"completed\":\n # Redraw dashboard one final time with \"completed\" status\n status_str = f\"Status: {status.status}\"\n if step is not None and max_steps is not None:\n step_str = f\"Step {max_steps}/{max_steps} (100%)\"\n clear_output(wait=True)\n _draw_dashboard(\n elapsed_mins, val_losses, train_losses,\n vram_history, util_history,\n job_name, status_str, step_str, elapsed_str,\n )\n print(f\"\\nJob completed in {elapsed_min:.1f} minutes ({elapsed:.0f}s)\")\n return status\n elif status.status.lower() in (\"failed\", \"cancelled\", \"error\"):\n print(f\"\\nJob finished with status: {status.status}\")\n print(f\"Total time elapsed: {elapsed_min:.1f} minutes ({elapsed:.0f}s)\")\n\n # Print error details from the job level\n if status.error_details:\n error_msg = status.error_details.get(\"message\", \"\")\n if error_msg:\n print(f\"\\nError: {error_msg}\")\n\n # Find and print error details from the failed step/task\n for job_step in status.steps or []:\n if job_step.status == \"error\":\n print(f\"\\nFailed step: {job_step.name}\")\n if job_step.error_details:\n step_error = job_step.error_details.get(\"message\", \"\")\n if step_error:\n print(f\"Step error: {step_error}\")\n # Get error_stack from the failed task\n for task in job_step.tasks or []:\n if task.status == \"error\" and hasattr(task, \"error_stack\") and task.error_stack:\n print(f\"\\nError stack trace:\\n{task.error_stack}\")\n elif task.status == \"error\" and task.error_details:\n task_error = task.error_details.get(\"message\", \"\")\n if task_error:\n print(f\"Task error: {task_error}\")\n break\n\n raise Exception(f\"Job finished with status: {status.status}\")\n\n time.sleep(poll_interval)\n\n\n# Wait for the job to complete\njob_with_sequence_packing_status = wait_for_job(\n workspace=\"default\",\n job_name=job_with_sequence_packing.job.name,\n timeout=TIMEOUT_SECONDS,\n)\n\npacked_val_loss = get_training_metric(job_with_sequence_packing_status, VAL_LOSS_KEY)\nif packed_val_loss is not None:\n print(f\"Validation loss: {packed_val_loss:.2f}\")\nelse:\n print(\"Validation loss: not reported in job status\")", "language": "python", - "source_html": "import time\nfrom typing import cast\nfrom IPython.display import clear_output\nfrom nemo_platform.types.shared import PlatformJobStatusResponse\n\n# Timeout set to 30 minutes to accommodate typical LoRA training duration for this dataset size.\n# Actual training time will vary based on hardware, model size, and dataset complexity.\nTIMEOUT_SECONDS = 30 * 60 # 30 minutes\nVAL_LOSS_KEY = "val_loss"\nTRAIN_LOSS_KEY = "loss"\n\n# ---------------------------------------------------------------------------\n# Job polling with live dashboard\n# ---------------------------------------------------------------------------\n\ndef wait_for_job(\n workspace: str,\n job_name: str,\n timeout: int = TIMEOUT_SECONDS,\n poll_interval: int = 10,\n val_loss_key: str = VAL_LOSS_KEY,\n train_loss_key: str = TRAIN_LOSS_KEY,\n) -> PlatformJobStatusResponse:\n """\n Poll job status until completed, failed, cancelled, or timeout.\n Displays a live dashboard with loss curves and GPU metrics.\n\n Args:\n workspace: The workspace where the job is running.\n job_name: The name of the job to monitor.\n timeout: Maximum time to wait in seconds (default: 30 minutes).\n poll_interval: Time between status checks in seconds (default: 10).\n\n Returns:\n The final job status response.\n """\n start_time = time.time()\n\n # Time-series accumulators required for plotting\n elapsed_mins: list[float] = []\n val_losses: list[float | None] = []\n train_losses: list[float | None] = []\n vram_history: list[list[float]] = []\n util_history: list[list[float]] = []\n\n while True:\n elapsed = time.time() - start_time\n elapsed_min = elapsed / 60\n\n # Check for timeout\n if elapsed > timeout:\n error_message = f"Timeout reached after {elapsed_min:.1f} minutes"\n print(f"\\n{error_message}")\n print("Job did not complete within the timeout period.")\n raise Exception(error_message)\n\n status = client.jobs.get_status(name=job_name, workspace=workspace)\n\n # -- Extract training progress from nested steps structure --\n step: int | None = None\n max_steps: int | None = None\n training_phase: str | None = None\n val_loss: float | None = None\n train_loss: float | None = None\n current_step_name: str | None = None\n current_step_phase: str | None = None\n\n for job_step in status.steps or []:\n # Track the current active step name and phase for progress display\n if job_step.tasks:\n task = job_step.tasks[0]\n td = task.status_details or {}\n phase = cast(str, td.get("phase", ""))\n # Update current step if it's active or pending (not completed)\n if job_step.status in ("active", "pending"):\n current_step_name = job_step.name\n current_step_phase = phase or "started"\n\n if job_step.name == "training":\n for task in job_step.tasks or []:\n td = task.status_details or {}\n step = cast(int, td["step"]) if "step" in td else None\n max_steps = cast(int, td["max_steps"]) if "max_steps" in td else None\n training_phase = cast(str, td["phase"]) if "phase" in td else None\n raw_val_loss = td.get(val_loss_key)\n val_loss = float(raw_val_loss) if raw_val_loss is not None else None\n raw_train_loss = td.get(train_loss_key)\n train_loss = float(raw_train_loss) if raw_train_loss is not None else None\n break\n break\n\n if val_loss is None:\n raw_val_loss = (status.status_details or {}).get(val_loss_key)\n val_loss = float(raw_val_loss) if raw_val_loss is not None else None\n if train_loss is None:\n raw_train_loss = (status.status_details or {}).get(train_loss_key)\n train_loss = float(raw_train_loss) if raw_train_loss is not None else None\n\n # -- Collect GPU snapshot --\n vram_pcts, util_pcts = _get_gpu_snapshot()\n\n # -- Append to accumulators used for the plots --\n elapsed_mins.append(elapsed_min)\n val_losses.append(val_loss)\n train_losses.append(train_loss)\n vram_history.append(vram_pcts)\n util_history.append(util_pcts)\n\n # -- Build status strings --\n status_str = f"Status: {status.status}"\n if step is not None and max_steps is not None:\n pct = step / max_steps * 100\n step_str = f"Step {step}/{max_steps} ({pct:.0f}%)"\n if training_phase:\n step_str += f" - {training_phase}"\n else:\n if current_step_name and current_step_phase:\n step_str = f"{current_step_name} - {current_step_phase}"\n elif current_step_name:\n step_str = f"{current_step_name}"\n else:\n step_str = "Waiting for training to start..."\n elapsed_str = f"Elapsed: {elapsed_min:.1f} min"\n\n # -- Redraw dashboard --\n clear_output(wait=True)\n _draw_dashboard(\n elapsed_mins, val_losses, train_losses,\n vram_history, util_history,\n job_name, status_str, step_str, elapsed_str,\n )\n\n # -- Check terminal conditions --\n if status.status.lower() == "completed":\n # Redraw dashboard one final time with "completed" status\n status_str = f"Status: {status.status}"\n if step is not None and max_steps is not None:\n step_str = f"Step {max_steps}/{max_steps} (100%)"\n clear_output(wait=True)\n _draw_dashboard(\n elapsed_mins, val_losses, train_losses,\n vram_history, util_history,\n job_name, status_str, step_str, elapsed_str,\n )\n print(f"\\nJob completed in {elapsed_min:.1f} minutes ({elapsed:.0f}s)")\n return status\n elif status.status.lower() in ("failed", "cancelled", "error"):\n print(f"\\nJob finished with status: {status.status}")\n print(f"Total time elapsed: {elapsed_min:.1f} minutes ({elapsed:.0f}s)")\n\n # Print error details from the job level\n if status.error_details:\n error_msg = status.error_details.get("message", "")\n if error_msg:\n print(f"\\nError: {error_msg}")\n\n # Find and print error details from the failed step/task\n for job_step in status.steps or []:\n if job_step.status == "error":\n print(f"\\nFailed step: {job_step.name}")\n if job_step.error_details:\n step_error = job_step.error_details.get("message", "")\n if step_error:\n print(f"Step error: {step_error}")\n # Get error_stack from the failed task\n for task in job_step.tasks or []:\n if task.status == "error" and hasattr(task, "error_stack") and task.error_stack:\n print(f"\\nError stack trace:\\n{task.error_stack}")\n elif task.status == "error" and task.error_details:\n task_error = task.error_details.get("message", "")\n if task_error:\n print(f"Task error: {task_error}")\n break\n\n raise Exception(f"Job finished with status: {status.status}")\n\n time.sleep(poll_interval)\n\n\n# Wait for the job to complete\njob_with_sequence_packing_status = wait_for_job(\n workspace="default",\n job_name=job_with_sequence_packing.job.name,\n timeout=TIMEOUT_SECONDS,\n)\n\npacked_val_loss = (job_with_sequence_packing_status.status_details or {}).get("val_loss")\nif packed_val_loss is not None:\n print(f"Validation loss: {float(packed_val_loss):.2f}")\nelse:\n print("Validation loss: not reported in job status")\n" + "source_html": "import time\nfrom typing import cast\nfrom IPython.display import clear_output\nfrom nemo_platform.types.shared import PlatformJobStatusResponse\n\n# Timeout set to 30 minutes to accommodate typical LoRA training duration for this dataset size.\n# Actual training time will vary based on hardware, model size, and dataset complexity.\nTIMEOUT_SECONDS = 30 * 60 # 30 minutes\nVAL_LOSS_KEY = "val_loss"\nTRAIN_LOSS_KEY = "train_loss"\n\n\ndef get_training_metric(\n status: PlatformJobStatusResponse,\n metric_key: str,\n) -> float | None:\n """Return a metric reported by a task in the training step."""\n for job_step in status.steps or []:\n if job_step.name == "training":\n for task in job_step.tasks or []:\n value = (task.status_details or {}).get(metric_key)\n if value is not None:\n return float(value)\n return None\n\n\n# ---------------------------------------------------------------------------\n# Job polling with live dashboard\n# ---------------------------------------------------------------------------\n\ndef wait_for_job(\n workspace: str,\n job_name: str,\n timeout: int = TIMEOUT_SECONDS,\n poll_interval: int = 10,\n val_loss_key: str = VAL_LOSS_KEY,\n train_loss_key: str = TRAIN_LOSS_KEY,\n) -> PlatformJobStatusResponse:\n """\n Poll job status until completed, failed, cancelled, or timeout.\n Displays a live dashboard with loss curves and GPU metrics.\n\n Args:\n workspace: The workspace where the job is running.\n job_name: The name of the job to monitor.\n timeout: Maximum time to wait in seconds (default: 30 minutes).\n poll_interval: Time between status checks in seconds (default: 10).\n\n Returns:\n The final job status response.\n """\n start_time = time.time()\n\n # Time-series accumulators required for plotting\n elapsed_mins: list[float] = []\n val_losses: list[float | None] = []\n train_losses: list[float | None] = []\n vram_history: list[list[float]] = []\n util_history: list[list[float]] = []\n\n while True:\n elapsed = time.time() - start_time\n elapsed_min = elapsed / 60\n\n # Check for timeout\n if elapsed > timeout:\n error_message = f"Timeout reached after {elapsed_min:.1f} minutes"\n print(f"\\n{error_message}")\n print("Job did not complete within the timeout period.")\n raise Exception(error_message)\n\n status = client.jobs.get_status(name=job_name, workspace=workspace)\n\n # -- Extract training progress from nested steps structure --\n step: int | None = None\n max_steps: int | None = None\n training_phase: str | None = None\n val_loss: float | None = None\n train_loss: float | None = None\n current_step_name: str | None = None\n current_step_phase: str | None = None\n\n for job_step in status.steps or []:\n # Track the current active step name and phase for progress display\n if job_step.tasks:\n task = job_step.tasks[0]\n td = task.status_details or {}\n phase = cast(str, td.get("phase", ""))\n # Update current step if it's active or pending (not completed)\n if job_step.status in ("active", "pending"):\n current_step_name = job_step.name\n current_step_phase = phase or "started"\n\n if job_step.name == "training":\n for task in job_step.tasks or []:\n td = task.status_details or {}\n step = cast(int, td["step"]) if "step" in td else None\n max_steps = cast(int, td["max_steps"]) if "max_steps" in td else None\n training_phase = cast(str, td["phase"]) if "phase" in td else None\n raw_val_loss = td.get(val_loss_key)\n val_loss = float(raw_val_loss) if raw_val_loss is not None else None\n raw_train_loss = td.get(train_loss_key)\n train_loss = float(raw_train_loss) if raw_train_loss is not None else None\n break\n break\n\n if val_loss is None:\n raw_val_loss = (status.status_details or {}).get(val_loss_key)\n val_loss = float(raw_val_loss) if raw_val_loss is not None else None\n if train_loss is None:\n raw_train_loss = (status.status_details or {}).get(train_loss_key)\n train_loss = float(raw_train_loss) if raw_train_loss is not None else None\n\n # -- Collect GPU snapshot --\n vram_pcts, util_pcts = _get_gpu_snapshot()\n\n # -- Append to accumulators used for the plots --\n elapsed_mins.append(elapsed_min)\n val_losses.append(val_loss)\n train_losses.append(train_loss)\n vram_history.append(vram_pcts)\n util_history.append(util_pcts)\n\n # -- Build status strings --\n status_str = f"Status: {status.status}"\n if step is not None and max_steps is not None:\n pct = step / max_steps * 100\n step_str = f"Step {step}/{max_steps} ({pct:.0f}%)"\n if training_phase:\n step_str += f" - {training_phase}"\n else:\n if current_step_name and current_step_phase:\n step_str = f"{current_step_name} - {current_step_phase}"\n elif current_step_name:\n step_str = f"{current_step_name}"\n else:\n step_str = "Waiting for training to start..."\n elapsed_str = f"Elapsed: {elapsed_min:.1f} min"\n\n # -- Redraw dashboard --\n clear_output(wait=True)\n _draw_dashboard(\n elapsed_mins, val_losses, train_losses,\n vram_history, util_history,\n job_name, status_str, step_str, elapsed_str,\n )\n\n # -- Check terminal conditions --\n if status.status.lower() == "completed":\n # Redraw dashboard one final time with "completed" status\n status_str = f"Status: {status.status}"\n if step is not None and max_steps is not None:\n step_str = f"Step {max_steps}/{max_steps} (100%)"\n clear_output(wait=True)\n _draw_dashboard(\n elapsed_mins, val_losses, train_losses,\n vram_history, util_history,\n job_name, status_str, step_str, elapsed_str,\n )\n print(f"\\nJob completed in {elapsed_min:.1f} minutes ({elapsed:.0f}s)")\n return status\n elif status.status.lower() in ("failed", "cancelled", "error"):\n print(f"\\nJob finished with status: {status.status}")\n print(f"Total time elapsed: {elapsed_min:.1f} minutes ({elapsed:.0f}s)")\n\n # Print error details from the job level\n if status.error_details:\n error_msg = status.error_details.get("message", "")\n if error_msg:\n print(f"\\nError: {error_msg}")\n\n # Find and print error details from the failed step/task\n for job_step in status.steps or []:\n if job_step.status == "error":\n print(f"\\nFailed step: {job_step.name}")\n if job_step.error_details:\n step_error = job_step.error_details.get("message", "")\n if step_error:\n print(f"Step error: {step_error}")\n # Get error_stack from the failed task\n for task in job_step.tasks or []:\n if task.status == "error" and hasattr(task, "error_stack") and task.error_stack:\n print(f"\\nError stack trace:\\n{task.error_stack}")\n elif task.status == "error" and task.error_details:\n task_error = task.error_details.get("message", "")\n if task_error:\n print(f"Task error: {task_error}")\n break\n\n raise Exception(f"Job finished with status: {status.status}")\n\n time.sleep(poll_interval)\n\n\n# Wait for the job to complete\njob_with_sequence_packing_status = wait_for_job(\n workspace="default",\n job_name=job_with_sequence_packing.job.name,\n timeout=TIMEOUT_SECONDS,\n)\n\npacked_val_loss = get_training_metric(job_with_sequence_packing_status, VAL_LOSS_KEY)\nif packed_val_loss is not None:\n print(f"Validation loss: {packed_val_loss:.2f}")\nelse:\n print("Validation loss: not reported in job status")\n" }, { "type": "markdown", @@ -127,14 +127,14 @@ }, { "type": "markdown", - "source": "### 8. Track Finetuning Progress for Job without Sequence Packing", - "source_html": "

8. Track Finetuning Progress for Job without Sequence Packing

\n" + "source": "### 8. Track Fine-Tuning Progress for Job without Sequence Packing", + "source_html": "

8. Track Fine-Tuning Progress for Job without Sequence Packing

\n" }, { "type": "code", - "source": "# Wait for the training step to complete\njob_without_sequence_packing_status = wait_for_job(\n workspace=\"default\",\n job_name=job_without_sequence_packing.job.name,\n timeout=TIMEOUT_SECONDS\n)\n\nno_pack_val_loss = (job_without_sequence_packing_status.status_details or {}).get(\"val_loss\")\nif no_pack_val_loss is not None:\n print(f\"Validation loss: {float(no_pack_val_loss):.2f}\")\nelse:\n print(\"Validation loss: not reported in job status\")", + "source": "# Wait for the training step to complete\njob_without_sequence_packing_status = wait_for_job(\n workspace=\"default\",\n job_name=job_without_sequence_packing.job.name,\n timeout=TIMEOUT_SECONDS\n)\n\nno_pack_val_loss = get_training_metric(job_without_sequence_packing_status, VAL_LOSS_KEY)\nif no_pack_val_loss is not None:\n print(f\"Validation loss: {no_pack_val_loss:.2f}\")\nelse:\n print(\"Validation loss: not reported in job status\")", "language": "python", - "source_html": "# Wait for the training step to complete\njob_without_sequence_packing_status = wait_for_job(\n workspace="default",\n job_name=job_without_sequence_packing.job.name,\n timeout=TIMEOUT_SECONDS\n)\n\nno_pack_val_loss = (job_without_sequence_packing_status.status_details or {}).get("val_loss")\nif no_pack_val_loss is not None:\n print(f"Validation loss: {float(no_pack_val_loss):.2f}")\nelse:\n print("Validation loss: not reported in job status")\n" + "source_html": "# Wait for the training step to complete\njob_without_sequence_packing_status = wait_for_job(\n workspace="default",\n job_name=job_without_sequence_packing.job.name,\n timeout=TIMEOUT_SECONDS\n)\n\nno_pack_val_loss = get_training_metric(job_without_sequence_packing_status, VAL_LOSS_KEY)\nif no_pack_val_loss is not None:\n print(f"Validation loss: {no_pack_val_loss:.2f}")\nelse:\n print("Validation loss: not reported in job status")\n" }, { "type": "markdown", @@ -143,9 +143,9 @@ }, { "type": "code", - "source": "from nemo_platform.types.jobs import PlatformJobStep\nfrom datetime import datetime\nimport pandas as pd\n\nSTEP_NAME = \"training\"\n\ndef get_elapsed_time(step: PlatformJobStep) -> float:\n \"\"\"Calculate elapsed time in seconds from step's created_at to updated_at.\"\"\"\n created_at = datetime.fromisoformat(step.created_at.replace(\"Z\", \"+00:00\"))\n updated_at = datetime.fromisoformat(step.updated_at.replace(\"Z\", \"+00:00\"))\n return (updated_at - created_at).total_seconds()\n\nstep_with_sequence_packing = client.jobs.steps.retrieve(\n name=STEP_NAME,\n workspace=\"default\",\n job=job_with_sequence_packing.job.name,\n)\n\nstep_without_sequence_packing = client.jobs.steps.retrieve(\n name=STEP_NAME,\n workspace=\"default\",\n job=job_without_sequence_packing.job.name,\n)\n\ntime_to_complete_with_sequence_packing = get_elapsed_time(step_with_sequence_packing)\ntime_to_complete_without_sequence_packing = get_elapsed_time(step_without_sequence_packing)\n\n# Display results as a table\nresults_df = pd.DataFrame({\n \"Seq Packing Enabled\": [True, False],\n \"Val Loss\": [\n (job_with_sequence_packing_status.status_details or {}).get(\"val_loss\"),\n (job_without_sequence_packing_status.status_details or {}).get(\"val_loss\"),\n ],\n \"Training Step Time, sec\": [\n time_to_complete_with_sequence_packing,\n time_to_complete_without_sequence_packing\n ]\n})\n\nresults_df.style.format({\"Val Loss\": \"{:.2f}\", \"Training Step Time, sec\": \"{:.0f}\"}).hide(axis='index')", + "source": "from nemo_platform.types.jobs import PlatformJobStep\nimport pandas as pd\n\nSTEP_NAME = \"training\"\n\ndef get_elapsed_time(step: PlatformJobStep) -> float:\n \"\"\"Calculate elapsed time in seconds from step's created_at to updated_at.\"\"\"\n if step.created_at is None or step.updated_at is None:\n raise ValueError(\"Training step timestamps are unavailable\")\n return (step.updated_at - step.created_at).total_seconds()\n\nstep_with_sequence_packing = client.jobs.steps.retrieve(\n name=STEP_NAME,\n workspace=\"default\",\n job=job_with_sequence_packing.job.name,\n)\n\nstep_without_sequence_packing = client.jobs.steps.retrieve(\n name=STEP_NAME,\n workspace=\"default\",\n job=job_without_sequence_packing.job.name,\n)\n\ntime_to_complete_with_sequence_packing = get_elapsed_time(step_with_sequence_packing)\ntime_to_complete_without_sequence_packing = get_elapsed_time(step_without_sequence_packing)\n\n# Display results as a table\nresults_df = pd.DataFrame({\n \"Seq Packing Enabled\": [True, False],\n \"Val Loss\": [packed_val_loss, no_pack_val_loss],\n \"Training Step Time, sec\": [\n time_to_complete_with_sequence_packing,\n time_to_complete_without_sequence_packing\n ]\n})\n\nresults_df.style.format({\"Val Loss\": \"{:.2f}\", \"Training Step Time, sec\": \"{:.0f}\"}).hide(axis='index')", "language": "python", - "source_html": "from nemo_platform.types.jobs import PlatformJobStep\nfrom datetime import datetime\nimport pandas as pd\n\nSTEP_NAME = "training"\n\ndef get_elapsed_time(step: PlatformJobStep) -> float:\n """Calculate elapsed time in seconds from step's created_at to updated_at."""\n created_at = datetime.fromisoformat(step.created_at.replace("Z", "+00:00"))\n updated_at = datetime.fromisoformat(step.updated_at.replace("Z", "+00:00"))\n return (updated_at - created_at).total_seconds()\n\nstep_with_sequence_packing = client.jobs.steps.retrieve(\n name=STEP_NAME,\n workspace="default",\n job=job_with_sequence_packing.job.name,\n)\n\nstep_without_sequence_packing = client.jobs.steps.retrieve(\n name=STEP_NAME,\n workspace="default",\n job=job_without_sequence_packing.job.name,\n)\n\ntime_to_complete_with_sequence_packing = get_elapsed_time(step_with_sequence_packing)\ntime_to_complete_without_sequence_packing = get_elapsed_time(step_without_sequence_packing)\n\n# Display results as a table\nresults_df = pd.DataFrame({\n "Seq Packing Enabled": [True, False],\n "Val Loss": [\n (job_with_sequence_packing_status.status_details or {}).get("val_loss"),\n (job_without_sequence_packing_status.status_details or {}).get("val_loss"),\n ],\n "Training Step Time, sec": [\n time_to_complete_with_sequence_packing,\n time_to_complete_without_sequence_packing\n ]\n})\n\nresults_df.style.format({"Val Loss": "{:.2f}", "Training Step Time, sec": "{:.0f}"}).hide(axis='index')\n" + "source_html": "from nemo_platform.types.jobs import PlatformJobStep\nimport pandas as pd\n\nSTEP_NAME = "training"\n\ndef get_elapsed_time(step: PlatformJobStep) -> float:\n """Calculate elapsed time in seconds from step's created_at to updated_at."""\n if step.created_at is None or step.updated_at is None:\n raise ValueError("Training step timestamps are unavailable")\n return (step.updated_at - step.created_at).total_seconds()\n\nstep_with_sequence_packing = client.jobs.steps.retrieve(\n name=STEP_NAME,\n workspace="default",\n job=job_with_sequence_packing.job.name,\n)\n\nstep_without_sequence_packing = client.jobs.steps.retrieve(\n name=STEP_NAME,\n workspace="default",\n job=job_without_sequence_packing.job.name,\n)\n\ntime_to_complete_with_sequence_packing = get_elapsed_time(step_with_sequence_packing)\ntime_to_complete_without_sequence_packing = get_elapsed_time(step_without_sequence_packing)\n\n# Display results as a table\nresults_df = pd.DataFrame({\n "Seq Packing Enabled": [True, False],\n "Val Loss": [packed_val_loss, no_pack_val_loss],\n "Training Step Time, sec": [\n time_to_complete_with_sequence_packing,\n time_to_complete_without_sequence_packing\n ]\n})\n\nresults_df.style.format({"Val Loss": "{:.2f}", "Training Step Time, sec": "{:.0f}"}).hide(axis='index')\n" }, { "type": "markdown", diff --git a/docs/fern/components/notebooks/optimize-throughput.ts b/docs/fern/components/notebooks/optimize-throughput.ts index 8fab564b50..de52526d28 100644 --- a/docs/fern/components/notebooks/optimize-throughput.ts +++ b/docs/fern/components/notebooks/optimize-throughput.ts @@ -12,8 +12,8 @@ export default { cells: [ }, { "type": "markdown", - "source": "## Prerequisites\n\nBefore starting this tutorial, ensure you have:\n\n1. **Completed the [Quickstart](../../get-started/quickstart.md)** to install and deploy NeMo Platform locally\n2. **Installed the Python SDK** (PyPI wrapper: `pip install \"nemo-platform[all]\"`; source checkout: run `make bootstrap` from the repository root)", - "source_html": "

Prerequisites

\n

Before starting this tutorial, ensure you have:

\n
    \n
  1. Completed the Quickstart to install and deploy NeMo Platform locally
    2. \n
  2. Installed the Python SDK (PyPI wrapper: pip install "nemo-platform[all]"; source checkout: run make bootstrap from the repository root)
    3. \n
\n" + "source": "## Prerequisites\n\nBefore starting this tutorial, ensure you have:\n\n1. **Completed the [Quickstart](../../get-started/quickstart.md)** to install and deploy NeMo Platform locally\n2. **Installed the Python SDK** (PyPI wrapper: `pip install \"nemo-platform[all]\"`; source checkout: run `make bootstrap` from the repository root)\n3. **At least one GPU with CUDA 13+**", + "source_html": "

Prerequisites

\n

Before starting this tutorial, ensure you have:

\n
    \n
  1. Completed the Quickstart to install and deploy NeMo Platform locally
    2. \n
  2. Installed the Python SDK (PyPI wrapper: pip install "nemo-platform[all]"; source checkout: run make bootstrap from the repository root)
    3. \n
  3. At least one GPU with CUDA 13+
    4. \n
\n" }, { "type": "markdown", @@ -61,25 +61,25 @@ export default { cells: [ }, { "type": "markdown", - "source": "### 3. Secrets Setup\n\nIf you plan to use NGC or HuggingFace models, you will need to configure authentication:\n\n- **NGC models** (`ngc://` URIs): Requires NGC API key\n- **HuggingFace models** (`hf://` URIs): Requires HF token for gated/private models\n\n\nConfigure these as secrets in your platform. Refer to [Managing Secrets](../../get-started/concepts/manage-secrets.md) for detailed instructions.\n\nGet your credentials to access base models:\n- [NGC API Key](https://ngc.nvidia.com/) (Setup → Generate API Key)\n- [HuggingFace Token](https://huggingface.co/settings/tokens) (Create token with Read access)\n\n\n---\n\n#### Quick Setup Example\n\nThis tutorial uses the [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main) model from HuggingFace. Ensure that you have sufficient permissions to download the model. If you cannot access the files on the [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main) Hugging Face page, request access.\n\n**HuggingFace Authentication:**\n- For gated models (Llama, Gemma), you must provide a HuggingFace token via the `token_secret` parameter\n- Get your token from [HuggingFace Settings](https://huggingface.co/settings/tokens) (requires Read access)\n- Accept the model's terms on the HuggingFace model page before using it. Example: [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main)\n- For public models, you can omit the `token_secret` parameter when creating a fileset for the model in the next step.", - "source_html": "

3. Secrets Setup

\n

If you plan to use NGC or HuggingFace models, you will need to configure authentication:

\n
    \n
  • NGC models (ngc:// URIs): Requires NGC API key
    • \n
  • HuggingFace models (hf:// URIs): Requires HF token for gated/private models
    • \n
\n

Configure these as secrets in your platform. Refer to Managing Secrets for detailed instructions.

\n

Get your credentials to access base models:

\n\n
\n

Quick Setup Example

\n

This tutorial uses the meta-llama/Llama-3.2-1B-Instruct model from HuggingFace. Ensure that you have sufficient permissions to download the model. If you cannot access the files on the meta-llama/Llama-3.2-1B-Instruct Hugging Face page, request access.

\n

HuggingFace Authentication:

\n
    \n
  • For gated models (Llama, Gemma), you must provide a HuggingFace token via the token_secret parameter
    • \n
  • Get your token from HuggingFace Settings (requires Read access)
    • \n
  • Accept the model's terms on the HuggingFace model page before using it. Example: meta-llama/Llama-3.2-1B-Instruct
    • \n
  • For public models, you can omit the token_secret parameter when creating a fileset for the model in the next step.
    • \n
\n" + "source": "### 3. Secrets Setup\n\nIf you plan to use NGC or Hugging Face models, you will need to configure authentication:\n\n- **NGC models** (`ngc://` URIs): Requires NGC API key\n- **Hugging Face models** (`hf://` URIs): Requires HF token for gated/private models\n\n\nConfigure these as secrets in your platform. Refer to [Managing Secrets](../../get-started/concepts/manage-secrets.md) for detailed instructions.\n\nGet your credentials to access base models:\n- [NGC API Key](https://ngc.nvidia.com/) (Setup → Generate API Key)\n- [Hugging Face Token](https://huggingface.co/settings/tokens) (Create token with Read access)\n\n\n---\n\n#### Quick Setup Example\n\nThis tutorial uses the [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main) model from Hugging Face. Ensure that you have sufficient permissions to download the model. If you cannot access the files on the [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main) Hugging Face page, request access.\n\n**Hugging Face Authentication:**\n- For gated models (Llama, Gemma), you must provide a Hugging Face token via the `token_secret` parameter\n- Get your token from [Hugging Face Settings](https://huggingface.co/settings/tokens) (requires Read access)\n- Accept the model's terms on the Hugging Face model page before using it. Example: [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main)\n- For public models, you can omit the `token_secret` parameter when creating a fileset for the model in the next step.", + "source_html": "

3. Secrets Setup

\n

If you plan to use NGC or Hugging Face models, you will need to configure authentication:

\n
    \n
  • NGC models (ngc:// URIs): Requires NGC API key
    • \n
  • Hugging Face models (hf:// URIs): Requires HF token for gated/private models
    • \n
\n

Configure these as secrets in your platform. Refer to Managing Secrets for detailed instructions.

\n

Get your credentials to access base models:

\n\n
\n

Quick Setup Example

\n

This tutorial uses the meta-llama/Llama-3.2-1B-Instruct model from Hugging Face. Ensure that you have sufficient permissions to download the model. If you cannot access the files on the meta-llama/Llama-3.2-1B-Instruct Hugging Face page, request access.

\n

Hugging Face Authentication:

\n
    \n
  • For gated models (Llama, Gemma), you must provide a Hugging Face token via the token_secret parameter
    • \n
  • Get your token from Hugging Face Settings (requires Read access)
    • \n
  • Accept the model's terms on the Hugging Face model page before using it. Example: meta-llama/Llama-3.2-1B-Instruct
    • \n
  • For public models, you can omit the token_secret parameter when creating a fileset for the model in the next step.
    • \n
\n" }, { "type": "code", - "source": "# Export the HF_TOKEN and NGC_API_KEY environment variables if they are not already set\nHF_TOKEN = os.getenv(\"HF_TOKEN\")\nNGC_API_KEY = os.getenv(\"NGC_API_KEY\")\n\n\ndef create_or_get_secret(name: str, value: str | None, label: str):\n if not value:\n raise ValueError(f\"{label} environment variable is not set. Set it and try again.\")\n try:\n secret = client.secrets.create(\n name=name,\n workspace=\"default\",\n value=value,\n )\n print(f\"Created secret: {name}\")\n return secret\n except ConflictError:\n print(f\"Secret '{name}' already exists, continuing...\")\n return client.secrets.retrieve(name=name, workspace=\"default\")\n\n\n# Create HuggingFace token secret\nhf_secret = create_or_get_secret(\"hf-token\", HF_TOKEN, \"HF_TOKEN\")\nprint(\"HF_TOKEN secret:\")\nprint(hf_secret.model_dump_json(indent=2))\n\n# Create NGC API key secret\n# Uncomment the line below if you have NGC API Key and want to finetune NGC models\n# ngc_api_key = create_or_get_secret(\"ngc-api-key\", NGC_API_KEY, \"NGC_API_KEY\")", + "source": "# Export the HF_TOKEN and NGC_API_KEY environment variables if they are not already set\nHF_TOKEN = os.getenv(\"HF_TOKEN\")\nNGC_API_KEY = os.getenv(\"NGC_API_KEY\")\n\n\ndef create_or_get_secret(name: str, value: str | None, label: str):\n if not value:\n raise ValueError(f\"{label} environment variable is not set. Set it and try again.\")\n try:\n secret = client.secrets.create(\n name=name,\n workspace=\"default\",\n value=value,\n )\n print(f\"Created secret: {name}\")\n return secret\n except ConflictError:\n print(f\"Secret '{name}' already exists, continuing...\")\n return client.secrets.retrieve(name=name, workspace=\"default\")\n\n\n# Create Hugging Face token secret\nhf_secret = create_or_get_secret(\"hf-token\", HF_TOKEN, \"HF_TOKEN\")\nprint(\"HF_TOKEN secret:\")\nprint(hf_secret.model_dump_json(indent=2))\n\n# Create NGC API key secret\n# Uncomment the line below if you have NGC API Key and want to finetune NGC models\n# ngc_api_key = create_or_get_secret(\"ngc-api-key\", NGC_API_KEY, \"NGC_API_KEY\")", "language": "python", - "source_html": "# Export the HF_TOKEN and NGC_API_KEY environment variables if they are not already set\nHF_TOKEN = os.getenv("HF_TOKEN")\nNGC_API_KEY = os.getenv("NGC_API_KEY")\n\n\ndef create_or_get_secret(name: str, value: str | None, label: str):\n if not value:\n raise ValueError(f"{label} environment variable is not set. Set it and try again.")\n try:\n secret = client.secrets.create(\n name=name,\n workspace="default",\n value=value,\n )\n print(f"Created secret: {name}")\n return secret\n except ConflictError:\n print(f"Secret '{name}' already exists, continuing...")\n return client.secrets.retrieve(name=name, workspace="default")\n\n\n# Create HuggingFace token secret\nhf_secret = create_or_get_secret("hf-token", HF_TOKEN, "HF_TOKEN")\nprint("HF_TOKEN secret:")\nprint(hf_secret.model_dump_json(indent=2))\n\n# Create NGC API key secret\n# Uncomment the line below if you have NGC API Key and want to finetune NGC models\n# ngc_api_key = create_or_get_secret("ngc-api-key", NGC_API_KEY, "NGC_API_KEY")\n" + "source_html": "# Export the HF_TOKEN and NGC_API_KEY environment variables if they are not already set\nHF_TOKEN = os.getenv("HF_TOKEN")\nNGC_API_KEY = os.getenv("NGC_API_KEY")\n\n\ndef create_or_get_secret(name: str, value: str | None, label: str):\n if not value:\n raise ValueError(f"{label} environment variable is not set. Set it and try again.")\n try:\n secret = client.secrets.create(\n name=name,\n workspace="default",\n value=value,\n )\n print(f"Created secret: {name}")\n return secret\n except ConflictError:\n print(f"Secret '{name}' already exists, continuing...")\n return client.secrets.retrieve(name=name, workspace="default")\n\n\n# Create Hugging Face token secret\nhf_secret = create_or_get_secret("hf-token", HF_TOKEN, "HF_TOKEN")\nprint("HF_TOKEN secret:")\nprint(hf_secret.model_dump_json(indent=2))\n\n# Create NGC API key secret\n# Uncomment the line below if you have NGC API Key and want to finetune NGC models\n# ngc_api_key = create_or_get_secret("ngc-api-key", NGC_API_KEY, "NGC_API_KEY")\n" }, { "type": "markdown", - "source": "### 4. Create Base Model FileSet\n\nCreate a fileset pointing to the [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main) model on HuggingFace. This step creates a pointer to the model on Hugging Face and does not download it. The model is downloaded at job creation time.\n\nNote: for public models, you can omit the `token_secret` parameter when creating a model fileset.", - "source_html": "

4. Create Base Model FileSet

\n

Create a fileset pointing to the meta-llama/Llama-3.2-1B-Instruct model on HuggingFace. This step creates a pointer to the model on Hugging Face and does not download it. The model is downloaded at job creation time.

\n

Note: for public models, you can omit the token_secret parameter when creating a model fileset.

\n" + "source": "### 4. Create Base Model FileSet\n\nCreate a fileset pointing to the [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main) model on Hugging Face. This step creates a pointer to the model on Hugging Face and does not download it. The model is downloaded at job creation time.\n\nNote: for public models, you can omit the `token_secret` parameter when creating a model fileset.", + "source_html": "

4. Create Base Model FileSet

\n

Create a fileset pointing to the meta-llama/Llama-3.2-1B-Instruct model on Hugging Face. This step creates a pointer to the model on Hugging Face and does not download it. The model is downloaded at job creation time.

\n

Note: for public models, you can omit the token_secret parameter when creating a model fileset.

\n" }, { "type": "code", - "source": "import time\nfrom nemo_platform.types.files import HuggingfaceStorageConfigParam\n\nHF_REPO_ID = \"meta-llama/Llama-3.2-1B-Instruct\"\nMODEL_NAME = \"llama-3-2-1b-base\"\n\n# Ensure you have a HuggingFace token secret created\n# Create a fileset pointing to the desired HuggingFace model\ntry:\n base_model_fs = client.files.filesets.create(\n workspace=\"default\",\n name=MODEL_NAME,\n description=\"Llama 3.2 1B base model from HuggingFace\",\n storage=HuggingfaceStorageConfigParam(\n type=\"huggingface\",\n # repo_id is the full model name from Hugging Face\n repo_id=HF_REPO_ID,\n repo_type=\"model\",\n # we use the secret created in the previous step\n token_secret=hf_secret.name\n )\n )\n print(f\"Created base model fileset: {MODEL_NAME}\")\nexcept ConflictError:\n print(f\"Base model fileset already exists. Skipping creation.\")\n base_model_fs = client.files.filesets.retrieve(\n workspace=\"default\",\n name=MODEL_NAME,\n )\n\n# Create the Model Entity representation.\ntry:\n base_model = client.models.create(\n workspace=\"default\",\n name=MODEL_NAME,\n fileset=f\"default/{MODEL_NAME}\",\n )\n print(f\"Created Model Entity: {MODEL_NAME}\")\nexcept ConflictError:\n print(f\"Base model already exists. Updating fileset if different.\")\n base_model = client.models.update(\n workspace=\"default\",\n name=MODEL_NAME,\n fileset=f\"default/{MODEL_NAME}\",\n )\n\nprint(f\"\\nBase model fileset: fileset://default/{base_model.name}\")\nprint(\"Base model fileset files list:\")\nprint(json.dumps([f.model_dump() for f in client.files.list(fileset=MODEL_NAME, workspace=\"default\").data], indent=2))\n\n# Wait for ModelSpec to be populated from the checkpoint\nprint(\"\\nWaiting for ModelSpec to be populated...\")\nSPEC_TIMEOUT_SECONDS = 120\nspec_start = time.time()\nwhile not base_model.spec:\n if time.time() - spec_start > SPEC_TIMEOUT_SECONDS:\n raise TimeoutError(f\"ModelSpec not populated within {SPEC_TIMEOUT_SECONDS} seconds\")\n time.sleep(2)\n base_model = client.models.retrieve(\n workspace=\"default\",\n name=MODEL_NAME,\n )\n\nprint(f\"ModelSpec populated: {base_model.spec}\")", + "source": "import time\nfrom nemo_platform.types.files import HuggingfaceStorageConfigParam\n\nHF_REPO_ID = \"meta-llama/Llama-3.2-1B-Instruct\"\nMODEL_NAME = \"llama-3-2-1b-base\"\n\n# Ensure you have a Hugging Face token secret created\n# Create a fileset pointing to the desired Hugging Face model\ntry:\n base_model_fs = client.files.filesets.create(\n workspace=\"default\",\n name=MODEL_NAME,\n description=\"Llama 3.2 1B base model from Hugging Face\",\n storage=HuggingfaceStorageConfigParam(\n type=\"huggingface\",\n # repo_id is the full model name from Hugging Face\n repo_id=HF_REPO_ID,\n repo_type=\"model\",\n # we use the secret created in the previous step\n token_secret=hf_secret.name\n )\n )\n print(f\"Created base model fileset: {MODEL_NAME}\")\nexcept ConflictError:\n print(f\"Base model fileset already exists. Skipping creation.\")\n base_model_fs = client.files.filesets.retrieve(\n workspace=\"default\",\n name=MODEL_NAME,\n )\n\n# Create the Model Entity representation.\ntry:\n base_model = client.models.create(\n workspace=\"default\",\n name=MODEL_NAME,\n fileset=f\"default/{MODEL_NAME}\",\n )\n print(f\"Created Model Entity: {MODEL_NAME}\")\nexcept ConflictError:\n print(f\"Base model already exists. Updating fileset if different.\")\n base_model = client.models.update(\n workspace=\"default\",\n name=MODEL_NAME,\n fileset=f\"default/{MODEL_NAME}\",\n )\n\nprint(f\"\\nBase model fileset: fileset://default/{base_model.name}\")\nprint(\"Base model fileset files list:\")\nprint(json.dumps([f.model_dump() for f in client.files.list(fileset=MODEL_NAME, workspace=\"default\").data], indent=2))\n\n# Wait for ModelSpec to be populated from the checkpoint\nprint(\"\\nWaiting for ModelSpec to be populated...\")\nSPEC_TIMEOUT_SECONDS = 120\nspec_start = time.time()\nwhile not base_model.spec:\n if time.time() - spec_start > SPEC_TIMEOUT_SECONDS:\n raise TimeoutError(f\"ModelSpec not populated within {SPEC_TIMEOUT_SECONDS} seconds\")\n time.sleep(2)\n base_model = client.models.retrieve(\n workspace=\"default\",\n name=MODEL_NAME,\n )\n\nprint(f\"ModelSpec populated: {base_model.spec}\")", "language": "python", - "source_html": "import time\nfrom nemo_platform.types.files import HuggingfaceStorageConfigParam\n\nHF_REPO_ID = "meta-llama/Llama-3.2-1B-Instruct"\nMODEL_NAME = "llama-3-2-1b-base"\n\n# Ensure you have a HuggingFace token secret created\n# Create a fileset pointing to the desired HuggingFace model\ntry:\n base_model_fs = client.files.filesets.create(\n workspace="default",\n name=MODEL_NAME,\n description="Llama 3.2 1B base model from HuggingFace",\n storage=HuggingfaceStorageConfigParam(\n type="huggingface",\n # repo_id is the full model name from Hugging Face\n repo_id=HF_REPO_ID,\n repo_type="model",\n # we use the secret created in the previous step\n token_secret=hf_secret.name\n )\n )\n print(f"Created base model fileset: {MODEL_NAME}")\nexcept ConflictError:\n print(f"Base model fileset already exists. Skipping creation.")\n base_model_fs = client.files.filesets.retrieve(\n workspace="default",\n name=MODEL_NAME,\n )\n\n# Create the Model Entity representation.\ntry:\n base_model = client.models.create(\n workspace="default",\n name=MODEL_NAME,\n fileset=f"default/{MODEL_NAME}",\n )\n print(f"Created Model Entity: {MODEL_NAME}")\nexcept ConflictError:\n print(f"Base model already exists. Updating fileset if different.")\n base_model = client.models.update(\n workspace="default",\n name=MODEL_NAME,\n fileset=f"default/{MODEL_NAME}",\n )\n\nprint(f"\\nBase model fileset: fileset://default/{base_model.name}")\nprint("Base model fileset files list:")\nprint(json.dumps([f.model_dump() for f in client.files.list(fileset=MODEL_NAME, workspace="default").data], indent=2))\n\n# Wait for ModelSpec to be populated from the checkpoint\nprint("\\nWaiting for ModelSpec to be populated...")\nSPEC_TIMEOUT_SECONDS = 120\nspec_start = time.time()\nwhile not base_model.spec:\n if time.time() - spec_start > SPEC_TIMEOUT_SECONDS:\n raise TimeoutError(f"ModelSpec not populated within {SPEC_TIMEOUT_SECONDS} seconds")\n time.sleep(2)\n base_model = client.models.retrieve(\n workspace="default",\n name=MODEL_NAME,\n )\n\nprint(f"ModelSpec populated: {base_model.spec}")\n" + "source_html": "import time\nfrom nemo_platform.types.files import HuggingfaceStorageConfigParam\n\nHF_REPO_ID = "meta-llama/Llama-3.2-1B-Instruct"\nMODEL_NAME = "llama-3-2-1b-base"\n\n# Ensure you have a Hugging Face token secret created\n# Create a fileset pointing to the desired Hugging Face model\ntry:\n base_model_fs = client.files.filesets.create(\n workspace="default",\n name=MODEL_NAME,\n description="Llama 3.2 1B base model from Hugging Face",\n storage=HuggingfaceStorageConfigParam(\n type="huggingface",\n # repo_id is the full model name from Hugging Face\n repo_id=HF_REPO_ID,\n repo_type="model",\n # we use the secret created in the previous step\n token_secret=hf_secret.name\n )\n )\n print(f"Created base model fileset: {MODEL_NAME}")\nexcept ConflictError:\n print(f"Base model fileset already exists. Skipping creation.")\n base_model_fs = client.files.filesets.retrieve(\n workspace="default",\n name=MODEL_NAME,\n )\n\n# Create the Model Entity representation.\ntry:\n base_model = client.models.create(\n workspace="default",\n name=MODEL_NAME,\n fileset=f"default/{MODEL_NAME}",\n )\n print(f"Created Model Entity: {MODEL_NAME}")\nexcept ConflictError:\n print(f"Base model already exists. Updating fileset if different.")\n base_model = client.models.update(\n workspace="default",\n name=MODEL_NAME,\n fileset=f"default/{MODEL_NAME}",\n )\n\nprint(f"\\nBase model fileset: fileset://default/{base_model.name}")\nprint("Base model fileset files list:")\nprint(json.dumps([f.model_dump() for f in client.files.list(fileset=MODEL_NAME, workspace="default").data], indent=2))\n\n# Wait for ModelSpec to be populated from the checkpoint\nprint("\\nWaiting for ModelSpec to be populated...")\nSPEC_TIMEOUT_SECONDS = 120\nspec_start = time.time()\nwhile not base_model.spec:\n if time.time() - spec_start > SPEC_TIMEOUT_SECONDS:\n raise TimeoutError(f"ModelSpec not populated within {SPEC_TIMEOUT_SECONDS} seconds")\n time.sleep(2)\n base_model = client.models.retrieve(\n workspace="default",\n name=MODEL_NAME,\n )\n\nprint(f"ModelSpec populated: {base_model.spec}")\n" }, { "type": "markdown", @@ -94,8 +94,8 @@ export default { cells: [ }, { "type": "markdown", - "source": "### 6. Track Finetuning Progress\n\nA training job contains multiple steps: \n- Model and dataset downloading\n- Finetuning where LoRA adapter weights are trained\n- Creating a fileset entry for the finetuned model\n- Finetuned weights uploading\n\nThe elapsed time printed below reflects progress of the entire job. We compare the time taken by the finetuning step for both jobs in the last section of this tutorial.", - "source_html": "

6. Track Finetuning Progress

\n

A training job contains multiple steps:

\n
    \n
  • Model and dataset downloading
    • \n
  • Finetuning where LoRA adapter weights are trained
    • \n
  • Creating a fileset entry for the finetuned model
    • \n
  • Finetuned weights uploading
    • \n
\n

The elapsed time printed below reflects progress of the entire job. We compare the time taken by the finetuning step for both jobs in the last section of this tutorial.

\n" + "source": "### 6. Track Fine-Tuning Progress\n\nA training job contains multiple steps: \n- Model and dataset downloading\n- Fine-tuning where LoRA adapter weights are trained\n- Creating a fileset entry for the fine-tuned model\n- Fine-tuned weights uploading\n\nThe elapsed time printed below reflects progress of the entire job. We compare the time taken by the fine-tuning step for both jobs in the last section of this tutorial.", + "source_html": "

6. Track Fine-Tuning Progress

\n

A training job contains multiple steps:

\n
    \n
  • Model and dataset downloading
    • \n
  • Fine-tuning where LoRA adapter weights are trained
    • \n
  • Creating a fileset entry for the fine-tuned model
    • \n
  • Fine-tuned weights uploading
    • \n
\n

The elapsed time printed below reflects progress of the entire job. We compare the time taken by the fine-tuning step for both jobs in the last section of this tutorial.

\n" }, { "type": "markdown", @@ -110,14 +110,14 @@ export default { cells: [ }, { "type": "markdown", - "source": "#### Monitor the Job Until Completion\n\nThe cell below polls the job status every 10 seconds and renders a live dashboard with validation loss, GPU VRAM usage, and GPU utilization charts. The charts appear empty at first while the model and dataset download; training metrics and GPU activity populate after the finetuning step begins.\n\n> **Note:** This is additional code. You can also use the Weights & Biases or MLflow integrations.", - "source_html": "

Monitor the Job Until Completion

\n

The cell below polls the job status every 10 seconds and renders a live dashboard with validation loss, GPU VRAM usage, and GPU utilization charts. The charts appear empty at first while the model and dataset download; training metrics and GPU activity populate after the finetuning step begins.

\n
\n

Note: This is additional code. You can also use the Weights & Biases or MLflow integrations.

\n
\n" + "source": "#### Monitor the Job Until Completion\n\nThe cell below polls the job status every 10 seconds and renders a live dashboard with validation loss, GPU VRAM usage, and GPU utilization charts. The charts appear empty at first while the model and dataset download; training metrics and GPU activity populate after the fine-tuning step begins.\n\n> **Note:** This is additional code. You can also use the Weights & Biases or MLflow integrations.", + "source_html": "

Monitor the Job Until Completion

\n

The cell below polls the job status every 10 seconds and renders a live dashboard with validation loss, GPU VRAM usage, and GPU utilization charts. The charts appear empty at first while the model and dataset download; training metrics and GPU activity populate after the fine-tuning step begins.

\n
\n

Note: This is additional code. You can also use the Weights & Biases or MLflow integrations.

\n
\n" }, { "type": "code", - "source": "import time\nfrom typing import cast\nfrom IPython.display import clear_output\nfrom nemo_platform.types.shared import PlatformJobStatusResponse\n\n# Timeout set to 30 minutes to accommodate typical LoRA training duration for this dataset size.\n# Actual training time will vary based on hardware, model size, and dataset complexity.\nTIMEOUT_SECONDS = 30 * 60 # 30 minutes\nVAL_LOSS_KEY = \"val_loss\"\nTRAIN_LOSS_KEY = \"loss\"\n\n# ---------------------------------------------------------------------------\n# Job polling with live dashboard\n# ---------------------------------------------------------------------------\n\ndef wait_for_job(\n workspace: str,\n job_name: str,\n timeout: int = TIMEOUT_SECONDS,\n poll_interval: int = 10,\n val_loss_key: str = VAL_LOSS_KEY,\n train_loss_key: str = TRAIN_LOSS_KEY,\n) -> PlatformJobStatusResponse:\n \"\"\"\n Poll job status until completed, failed, cancelled, or timeout.\n Displays a live dashboard with loss curves and GPU metrics.\n\n Args:\n workspace: The workspace where the job is running.\n job_name: The name of the job to monitor.\n timeout: Maximum time to wait in seconds (default: 30 minutes).\n poll_interval: Time between status checks in seconds (default: 10).\n\n Returns:\n The final job status response.\n \"\"\"\n start_time = time.time()\n\n # Time-series accumulators required for plotting\n elapsed_mins: list[float] = []\n val_losses: list[float | None] = []\n train_losses: list[float | None] = []\n vram_history: list[list[float]] = []\n util_history: list[list[float]] = []\n\n while True:\n elapsed = time.time() - start_time\n elapsed_min = elapsed / 60\n\n # Check for timeout\n if elapsed > timeout:\n error_message = f\"Timeout reached after {elapsed_min:.1f} minutes\"\n print(f\"\\n{error_message}\")\n print(\"Job did not complete within the timeout period.\")\n raise Exception(error_message)\n\n status = client.jobs.get_status(name=job_name, workspace=workspace)\n\n # -- Extract training progress from nested steps structure --\n step: int | None = None\n max_steps: int | None = None\n training_phase: str | None = None\n val_loss: float | None = None\n train_loss: float | None = None\n current_step_name: str | None = None\n current_step_phase: str | None = None\n\n for job_step in status.steps or []:\n # Track the current active step name and phase for progress display\n if job_step.tasks:\n task = job_step.tasks[0]\n td = task.status_details or {}\n phase = cast(str, td.get(\"phase\", \"\"))\n # Update current step if it's active or pending (not completed)\n if job_step.status in (\"active\", \"pending\"):\n current_step_name = job_step.name\n current_step_phase = phase or \"started\"\n\n if job_step.name == \"training\":\n for task in job_step.tasks or []:\n td = task.status_details or {}\n step = cast(int, td[\"step\"]) if \"step\" in td else None\n max_steps = cast(int, td[\"max_steps\"]) if \"max_steps\" in td else None\n training_phase = cast(str, td[\"phase\"]) if \"phase\" in td else None\n raw_val_loss = td.get(val_loss_key)\n val_loss = float(raw_val_loss) if raw_val_loss is not None else None\n raw_train_loss = td.get(train_loss_key)\n train_loss = float(raw_train_loss) if raw_train_loss is not None else None\n break\n break\n\n if val_loss is None:\n raw_val_loss = (status.status_details or {}).get(val_loss_key)\n val_loss = float(raw_val_loss) if raw_val_loss is not None else None\n if train_loss is None:\n raw_train_loss = (status.status_details or {}).get(train_loss_key)\n train_loss = float(raw_train_loss) if raw_train_loss is not None else None\n\n # -- Collect GPU snapshot --\n vram_pcts, util_pcts = _get_gpu_snapshot()\n\n # -- Append to accumulators used for the plots --\n elapsed_mins.append(elapsed_min)\n val_losses.append(val_loss)\n train_losses.append(train_loss)\n vram_history.append(vram_pcts)\n util_history.append(util_pcts)\n\n # -- Build status strings --\n status_str = f\"Status: {status.status}\"\n if step is not None and max_steps is not None:\n pct = step / max_steps * 100\n step_str = f\"Step {step}/{max_steps} ({pct:.0f}%)\"\n if training_phase:\n step_str += f\" - {training_phase}\"\n else:\n if current_step_name and current_step_phase:\n step_str = f\"{current_step_name} - {current_step_phase}\"\n elif current_step_name:\n step_str = f\"{current_step_name}\"\n else:\n step_str = \"Waiting for training to start...\"\n elapsed_str = f\"Elapsed: {elapsed_min:.1f} min\"\n\n # -- Redraw dashboard --\n clear_output(wait=True)\n _draw_dashboard(\n elapsed_mins, val_losses, train_losses,\n vram_history, util_history,\n job_name, status_str, step_str, elapsed_str,\n )\n\n # -- Check terminal conditions --\n if status.status.lower() == \"completed\":\n # Redraw dashboard one final time with \"completed\" status\n status_str = f\"Status: {status.status}\"\n if step is not None and max_steps is not None:\n step_str = f\"Step {max_steps}/{max_steps} (100%)\"\n clear_output(wait=True)\n _draw_dashboard(\n elapsed_mins, val_losses, train_losses,\n vram_history, util_history,\n job_name, status_str, step_str, elapsed_str,\n )\n print(f\"\\nJob completed in {elapsed_min:.1f} minutes ({elapsed:.0f}s)\")\n return status\n elif status.status.lower() in (\"failed\", \"cancelled\", \"error\"):\n print(f\"\\nJob finished with status: {status.status}\")\n print(f\"Total time elapsed: {elapsed_min:.1f} minutes ({elapsed:.0f}s)\")\n\n # Print error details from the job level\n if status.error_details:\n error_msg = status.error_details.get(\"message\", \"\")\n if error_msg:\n print(f\"\\nError: {error_msg}\")\n\n # Find and print error details from the failed step/task\n for job_step in status.steps or []:\n if job_step.status == \"error\":\n print(f\"\\nFailed step: {job_step.name}\")\n if job_step.error_details:\n step_error = job_step.error_details.get(\"message\", \"\")\n if step_error:\n print(f\"Step error: {step_error}\")\n # Get error_stack from the failed task\n for task in job_step.tasks or []:\n if task.status == \"error\" and hasattr(task, \"error_stack\") and task.error_stack:\n print(f\"\\nError stack trace:\\n{task.error_stack}\")\n elif task.status == \"error\" and task.error_details:\n task_error = task.error_details.get(\"message\", \"\")\n if task_error:\n print(f\"Task error: {task_error}\")\n break\n\n raise Exception(f\"Job finished with status: {status.status}\")\n\n time.sleep(poll_interval)\n\n\n# Wait for the job to complete\njob_with_sequence_packing_status = wait_for_job(\n workspace=\"default\",\n job_name=job_with_sequence_packing.job.name,\n timeout=TIMEOUT_SECONDS,\n)\n\npacked_val_loss = (job_with_sequence_packing_status.status_details or {}).get(\"val_loss\")\nif packed_val_loss is not None:\n print(f\"Validation loss: {float(packed_val_loss):.2f}\")\nelse:\n print(\"Validation loss: not reported in job status\")", + "source": "import time\nfrom typing import cast\nfrom IPython.display import clear_output\nfrom nemo_platform.types.shared import PlatformJobStatusResponse\n\n# Timeout set to 30 minutes to accommodate typical LoRA training duration for this dataset size.\n# Actual training time will vary based on hardware, model size, and dataset complexity.\nTIMEOUT_SECONDS = 30 * 60 # 30 minutes\nVAL_LOSS_KEY = \"val_loss\"\nTRAIN_LOSS_KEY = \"train_loss\"\n\n\ndef get_training_metric(\n status: PlatformJobStatusResponse,\n metric_key: str,\n) -> float | None:\n \"\"\"Return a metric reported by a task in the training step.\"\"\"\n for job_step in status.steps or []:\n if job_step.name == \"training\":\n for task in job_step.tasks or []:\n value = (task.status_details or {}).get(metric_key)\n if value is not None:\n return float(value)\n return None\n\n\n# ---------------------------------------------------------------------------\n# Job polling with live dashboard\n# ---------------------------------------------------------------------------\n\ndef wait_for_job(\n workspace: str,\n job_name: str,\n timeout: int = TIMEOUT_SECONDS,\n poll_interval: int = 10,\n val_loss_key: str = VAL_LOSS_KEY,\n train_loss_key: str = TRAIN_LOSS_KEY,\n) -> PlatformJobStatusResponse:\n \"\"\"\n Poll job status until completed, failed, cancelled, or timeout.\n Displays a live dashboard with loss curves and GPU metrics.\n\n Args:\n workspace: The workspace where the job is running.\n job_name: The name of the job to monitor.\n timeout: Maximum time to wait in seconds (default: 30 minutes).\n poll_interval: Time between status checks in seconds (default: 10).\n\n Returns:\n The final job status response.\n \"\"\"\n start_time = time.time()\n\n # Time-series accumulators required for plotting\n elapsed_mins: list[float] = []\n val_losses: list[float | None] = []\n train_losses: list[float | None] = []\n vram_history: list[list[float]] = []\n util_history: list[list[float]] = []\n\n while True:\n elapsed = time.time() - start_time\n elapsed_min = elapsed / 60\n\n # Check for timeout\n if elapsed > timeout:\n error_message = f\"Timeout reached after {elapsed_min:.1f} minutes\"\n print(f\"\\n{error_message}\")\n print(\"Job did not complete within the timeout period.\")\n raise Exception(error_message)\n\n status = client.jobs.get_status(name=job_name, workspace=workspace)\n\n # -- Extract training progress from nested steps structure --\n step: int | None = None\n max_steps: int | None = None\n training_phase: str | None = None\n val_loss: float | None = None\n train_loss: float | None = None\n current_step_name: str | None = None\n current_step_phase: str | None = None\n\n for job_step in status.steps or []:\n # Track the current active step name and phase for progress display\n if job_step.tasks:\n task = job_step.tasks[0]\n td = task.status_details or {}\n phase = cast(str, td.get(\"phase\", \"\"))\n # Update current step if it's active or pending (not completed)\n if job_step.status in (\"active\", \"pending\"):\n current_step_name = job_step.name\n current_step_phase = phase or \"started\"\n\n if job_step.name == \"training\":\n for task in job_step.tasks or []:\n td = task.status_details or {}\n step = cast(int, td[\"step\"]) if \"step\" in td else None\n max_steps = cast(int, td[\"max_steps\"]) if \"max_steps\" in td else None\n training_phase = cast(str, td[\"phase\"]) if \"phase\" in td else None\n raw_val_loss = td.get(val_loss_key)\n val_loss = float(raw_val_loss) if raw_val_loss is not None else None\n raw_train_loss = td.get(train_loss_key)\n train_loss = float(raw_train_loss) if raw_train_loss is not None else None\n break\n break\n\n if val_loss is None:\n raw_val_loss = (status.status_details or {}).get(val_loss_key)\n val_loss = float(raw_val_loss) if raw_val_loss is not None else None\n if train_loss is None:\n raw_train_loss = (status.status_details or {}).get(train_loss_key)\n train_loss = float(raw_train_loss) if raw_train_loss is not None else None\n\n # -- Collect GPU snapshot --\n vram_pcts, util_pcts = _get_gpu_snapshot()\n\n # -- Append to accumulators used for the plots --\n elapsed_mins.append(elapsed_min)\n val_losses.append(val_loss)\n train_losses.append(train_loss)\n vram_history.append(vram_pcts)\n util_history.append(util_pcts)\n\n # -- Build status strings --\n status_str = f\"Status: {status.status}\"\n if step is not None and max_steps is not None:\n pct = step / max_steps * 100\n step_str = f\"Step {step}/{max_steps} ({pct:.0f}%)\"\n if training_phase:\n step_str += f\" - {training_phase}\"\n else:\n if current_step_name and current_step_phase:\n step_str = f\"{current_step_name} - {current_step_phase}\"\n elif current_step_name:\n step_str = f\"{current_step_name}\"\n else:\n step_str = \"Waiting for training to start...\"\n elapsed_str = f\"Elapsed: {elapsed_min:.1f} min\"\n\n # -- Redraw dashboard --\n clear_output(wait=True)\n _draw_dashboard(\n elapsed_mins, val_losses, train_losses,\n vram_history, util_history,\n job_name, status_str, step_str, elapsed_str,\n )\n\n # -- Check terminal conditions --\n if status.status.lower() == \"completed\":\n # Redraw dashboard one final time with \"completed\" status\n status_str = f\"Status: {status.status}\"\n if step is not None and max_steps is not None:\n step_str = f\"Step {max_steps}/{max_steps} (100%)\"\n clear_output(wait=True)\n _draw_dashboard(\n elapsed_mins, val_losses, train_losses,\n vram_history, util_history,\n job_name, status_str, step_str, elapsed_str,\n )\n print(f\"\\nJob completed in {elapsed_min:.1f} minutes ({elapsed:.0f}s)\")\n return status\n elif status.status.lower() in (\"failed\", \"cancelled\", \"error\"):\n print(f\"\\nJob finished with status: {status.status}\")\n print(f\"Total time elapsed: {elapsed_min:.1f} minutes ({elapsed:.0f}s)\")\n\n # Print error details from the job level\n if status.error_details:\n error_msg = status.error_details.get(\"message\", \"\")\n if error_msg:\n print(f\"\\nError: {error_msg}\")\n\n # Find and print error details from the failed step/task\n for job_step in status.steps or []:\n if job_step.status == \"error\":\n print(f\"\\nFailed step: {job_step.name}\")\n if job_step.error_details:\n step_error = job_step.error_details.get(\"message\", \"\")\n if step_error:\n print(f\"Step error: {step_error}\")\n # Get error_stack from the failed task\n for task in job_step.tasks or []:\n if task.status == \"error\" and hasattr(task, \"error_stack\") and task.error_stack:\n print(f\"\\nError stack trace:\\n{task.error_stack}\")\n elif task.status == \"error\" and task.error_details:\n task_error = task.error_details.get(\"message\", \"\")\n if task_error:\n print(f\"Task error: {task_error}\")\n break\n\n raise Exception(f\"Job finished with status: {status.status}\")\n\n time.sleep(poll_interval)\n\n\n# Wait for the job to complete\njob_with_sequence_packing_status = wait_for_job(\n workspace=\"default\",\n job_name=job_with_sequence_packing.job.name,\n timeout=TIMEOUT_SECONDS,\n)\n\npacked_val_loss = get_training_metric(job_with_sequence_packing_status, VAL_LOSS_KEY)\nif packed_val_loss is not None:\n print(f\"Validation loss: {packed_val_loss:.2f}\")\nelse:\n print(\"Validation loss: not reported in job status\")", "language": "python", - "source_html": "import time\nfrom typing import cast\nfrom IPython.display import clear_output\nfrom nemo_platform.types.shared import PlatformJobStatusResponse\n\n# Timeout set to 30 minutes to accommodate typical LoRA training duration for this dataset size.\n# Actual training time will vary based on hardware, model size, and dataset complexity.\nTIMEOUT_SECONDS = 30 * 60 # 30 minutes\nVAL_LOSS_KEY = "val_loss"\nTRAIN_LOSS_KEY = "loss"\n\n# ---------------------------------------------------------------------------\n# Job polling with live dashboard\n# ---------------------------------------------------------------------------\n\ndef wait_for_job(\n workspace: str,\n job_name: str,\n timeout: int = TIMEOUT_SECONDS,\n poll_interval: int = 10,\n val_loss_key: str = VAL_LOSS_KEY,\n train_loss_key: str = TRAIN_LOSS_KEY,\n) -> PlatformJobStatusResponse:\n """\n Poll job status until completed, failed, cancelled, or timeout.\n Displays a live dashboard with loss curves and GPU metrics.\n\n Args:\n workspace: The workspace where the job is running.\n job_name: The name of the job to monitor.\n timeout: Maximum time to wait in seconds (default: 30 minutes).\n poll_interval: Time between status checks in seconds (default: 10).\n\n Returns:\n The final job status response.\n """\n start_time = time.time()\n\n # Time-series accumulators required for plotting\n elapsed_mins: list[float] = []\n val_losses: list[float | None] = []\n train_losses: list[float | None] = []\n vram_history: list[list[float]] = []\n util_history: list[list[float]] = []\n\n while True:\n elapsed = time.time() - start_time\n elapsed_min = elapsed / 60\n\n # Check for timeout\n if elapsed > timeout:\n error_message = f"Timeout reached after {elapsed_min:.1f} minutes"\n print(f"\\n{error_message}")\n print("Job did not complete within the timeout period.")\n raise Exception(error_message)\n\n status = client.jobs.get_status(name=job_name, workspace=workspace)\n\n # -- Extract training progress from nested steps structure --\n step: int | None = None\n max_steps: int | None = None\n training_phase: str | None = None\n val_loss: float | None = None\n train_loss: float | None = None\n current_step_name: str | None = None\n current_step_phase: str | None = None\n\n for job_step in status.steps or []:\n # Track the current active step name and phase for progress display\n if job_step.tasks:\n task = job_step.tasks[0]\n td = task.status_details or {}\n phase = cast(str, td.get("phase", ""))\n # Update current step if it's active or pending (not completed)\n if job_step.status in ("active", "pending"):\n current_step_name = job_step.name\n current_step_phase = phase or "started"\n\n if job_step.name == "training":\n for task in job_step.tasks or []:\n td = task.status_details or {}\n step = cast(int, td["step"]) if "step" in td else None\n max_steps = cast(int, td["max_steps"]) if "max_steps" in td else None\n training_phase = cast(str, td["phase"]) if "phase" in td else None\n raw_val_loss = td.get(val_loss_key)\n val_loss = float(raw_val_loss) if raw_val_loss is not None else None\n raw_train_loss = td.get(train_loss_key)\n train_loss = float(raw_train_loss) if raw_train_loss is not None else None\n break\n break\n\n if val_loss is None:\n raw_val_loss = (status.status_details or {}).get(val_loss_key)\n val_loss = float(raw_val_loss) if raw_val_loss is not None else None\n if train_loss is None:\n raw_train_loss = (status.status_details or {}).get(train_loss_key)\n train_loss = float(raw_train_loss) if raw_train_loss is not None else None\n\n # -- Collect GPU snapshot --\n vram_pcts, util_pcts = _get_gpu_snapshot()\n\n # -- Append to accumulators used for the plots --\n elapsed_mins.append(elapsed_min)\n val_losses.append(val_loss)\n train_losses.append(train_loss)\n vram_history.append(vram_pcts)\n util_history.append(util_pcts)\n\n # -- Build status strings --\n status_str = f"Status: {status.status}"\n if step is not None and max_steps is not None:\n pct = step / max_steps * 100\n step_str = f"Step {step}/{max_steps} ({pct:.0f}%)"\n if training_phase:\n step_str += f" - {training_phase}"\n else:\n if current_step_name and current_step_phase:\n step_str = f"{current_step_name} - {current_step_phase}"\n elif current_step_name:\n step_str = f"{current_step_name}"\n else:\n step_str = "Waiting for training to start..."\n elapsed_str = f"Elapsed: {elapsed_min:.1f} min"\n\n # -- Redraw dashboard --\n clear_output(wait=True)\n _draw_dashboard(\n elapsed_mins, val_losses, train_losses,\n vram_history, util_history,\n job_name, status_str, step_str, elapsed_str,\n )\n\n # -- Check terminal conditions --\n if status.status.lower() == "completed":\n # Redraw dashboard one final time with "completed" status\n status_str = f"Status: {status.status}"\n if step is not None and max_steps is not None:\n step_str = f"Step {max_steps}/{max_steps} (100%)"\n clear_output(wait=True)\n _draw_dashboard(\n elapsed_mins, val_losses, train_losses,\n vram_history, util_history,\n job_name, status_str, step_str, elapsed_str,\n )\n print(f"\\nJob completed in {elapsed_min:.1f} minutes ({elapsed:.0f}s)")\n return status\n elif status.status.lower() in ("failed", "cancelled", "error"):\n print(f"\\nJob finished with status: {status.status}")\n print(f"Total time elapsed: {elapsed_min:.1f} minutes ({elapsed:.0f}s)")\n\n # Print error details from the job level\n if status.error_details:\n error_msg = status.error_details.get("message", "")\n if error_msg:\n print(f"\\nError: {error_msg}")\n\n # Find and print error details from the failed step/task\n for job_step in status.steps or []:\n if job_step.status == "error":\n print(f"\\nFailed step: {job_step.name}")\n if job_step.error_details:\n step_error = job_step.error_details.get("message", "")\n if step_error:\n print(f"Step error: {step_error}")\n # Get error_stack from the failed task\n for task in job_step.tasks or []:\n if task.status == "error" and hasattr(task, "error_stack") and task.error_stack:\n print(f"\\nError stack trace:\\n{task.error_stack}")\n elif task.status == "error" and task.error_details:\n task_error = task.error_details.get("message", "")\n if task_error:\n print(f"Task error: {task_error}")\n break\n\n raise Exception(f"Job finished with status: {status.status}")\n\n time.sleep(poll_interval)\n\n\n# Wait for the job to complete\njob_with_sequence_packing_status = wait_for_job(\n workspace="default",\n job_name=job_with_sequence_packing.job.name,\n timeout=TIMEOUT_SECONDS,\n)\n\npacked_val_loss = (job_with_sequence_packing_status.status_details or {}).get("val_loss")\nif packed_val_loss is not None:\n print(f"Validation loss: {float(packed_val_loss):.2f}")\nelse:\n print("Validation loss: not reported in job status")\n" + "source_html": "import time\nfrom typing import cast\nfrom IPython.display import clear_output\nfrom nemo_platform.types.shared import PlatformJobStatusResponse\n\n# Timeout set to 30 minutes to accommodate typical LoRA training duration for this dataset size.\n# Actual training time will vary based on hardware, model size, and dataset complexity.\nTIMEOUT_SECONDS = 30 * 60 # 30 minutes\nVAL_LOSS_KEY = "val_loss"\nTRAIN_LOSS_KEY = "train_loss"\n\n\ndef get_training_metric(\n status: PlatformJobStatusResponse,\n metric_key: str,\n) -> float | None:\n """Return a metric reported by a task in the training step."""\n for job_step in status.steps or []:\n if job_step.name == "training":\n for task in job_step.tasks or []:\n value = (task.status_details or {}).get(metric_key)\n if value is not None:\n return float(value)\n return None\n\n\n# ---------------------------------------------------------------------------\n# Job polling with live dashboard\n# ---------------------------------------------------------------------------\n\ndef wait_for_job(\n workspace: str,\n job_name: str,\n timeout: int = TIMEOUT_SECONDS,\n poll_interval: int = 10,\n val_loss_key: str = VAL_LOSS_KEY,\n train_loss_key: str = TRAIN_LOSS_KEY,\n) -> PlatformJobStatusResponse:\n """\n Poll job status until completed, failed, cancelled, or timeout.\n Displays a live dashboard with loss curves and GPU metrics.\n\n Args:\n workspace: The workspace where the job is running.\n job_name: The name of the job to monitor.\n timeout: Maximum time to wait in seconds (default: 30 minutes).\n poll_interval: Time between status checks in seconds (default: 10).\n\n Returns:\n The final job status response.\n """\n start_time = time.time()\n\n # Time-series accumulators required for plotting\n elapsed_mins: list[float] = []\n val_losses: list[float | None] = []\n train_losses: list[float | None] = []\n vram_history: list[list[float]] = []\n util_history: list[list[float]] = []\n\n while True:\n elapsed = time.time() - start_time\n elapsed_min = elapsed / 60\n\n # Check for timeout\n if elapsed > timeout:\n error_message = f"Timeout reached after {elapsed_min:.1f} minutes"\n print(f"\\n{error_message}")\n print("Job did not complete within the timeout period.")\n raise Exception(error_message)\n\n status = client.jobs.get_status(name=job_name, workspace=workspace)\n\n # -- Extract training progress from nested steps structure --\n step: int | None = None\n max_steps: int | None = None\n training_phase: str | None = None\n val_loss: float | None = None\n train_loss: float | None = None\n current_step_name: str | None = None\n current_step_phase: str | None = None\n\n for job_step in status.steps or []:\n # Track the current active step name and phase for progress display\n if job_step.tasks:\n task = job_step.tasks[0]\n td = task.status_details or {}\n phase = cast(str, td.get("phase", ""))\n # Update current step if it's active or pending (not completed)\n if job_step.status in ("active", "pending"):\n current_step_name = job_step.name\n current_step_phase = phase or "started"\n\n if job_step.name == "training":\n for task in job_step.tasks or []:\n td = task.status_details or {}\n step = cast(int, td["step"]) if "step" in td else None\n max_steps = cast(int, td["max_steps"]) if "max_steps" in td else None\n training_phase = cast(str, td["phase"]) if "phase" in td else None\n raw_val_loss = td.get(val_loss_key)\n val_loss = float(raw_val_loss) if raw_val_loss is not None else None\n raw_train_loss = td.get(train_loss_key)\n train_loss = float(raw_train_loss) if raw_train_loss is not None else None\n break\n break\n\n if val_loss is None:\n raw_val_loss = (status.status_details or {}).get(val_loss_key)\n val_loss = float(raw_val_loss) if raw_val_loss is not None else None\n if train_loss is None:\n raw_train_loss = (status.status_details or {}).get(train_loss_key)\n train_loss = float(raw_train_loss) if raw_train_loss is not None else None\n\n # -- Collect GPU snapshot --\n vram_pcts, util_pcts = _get_gpu_snapshot()\n\n # -- Append to accumulators used for the plots --\n elapsed_mins.append(elapsed_min)\n val_losses.append(val_loss)\n train_losses.append(train_loss)\n vram_history.append(vram_pcts)\n util_history.append(util_pcts)\n\n # -- Build status strings --\n status_str = f"Status: {status.status}"\n if step is not None and max_steps is not None:\n pct = step / max_steps * 100\n step_str = f"Step {step}/{max_steps} ({pct:.0f}%)"\n if training_phase:\n step_str += f" - {training_phase}"\n else:\n if current_step_name and current_step_phase:\n step_str = f"{current_step_name} - {current_step_phase}"\n elif current_step_name:\n step_str = f"{current_step_name}"\n else:\n step_str = "Waiting for training to start..."\n elapsed_str = f"Elapsed: {elapsed_min:.1f} min"\n\n # -- Redraw dashboard --\n clear_output(wait=True)\n _draw_dashboard(\n elapsed_mins, val_losses, train_losses,\n vram_history, util_history,\n job_name, status_str, step_str, elapsed_str,\n )\n\n # -- Check terminal conditions --\n if status.status.lower() == "completed":\n # Redraw dashboard one final time with "completed" status\n status_str = f"Status: {status.status}"\n if step is not None and max_steps is not None:\n step_str = f"Step {max_steps}/{max_steps} (100%)"\n clear_output(wait=True)\n _draw_dashboard(\n elapsed_mins, val_losses, train_losses,\n vram_history, util_history,\n job_name, status_str, step_str, elapsed_str,\n )\n print(f"\\nJob completed in {elapsed_min:.1f} minutes ({elapsed:.0f}s)")\n return status\n elif status.status.lower() in ("failed", "cancelled", "error"):\n print(f"\\nJob finished with status: {status.status}")\n print(f"Total time elapsed: {elapsed_min:.1f} minutes ({elapsed:.0f}s)")\n\n # Print error details from the job level\n if status.error_details:\n error_msg = status.error_details.get("message", "")\n if error_msg:\n print(f"\\nError: {error_msg}")\n\n # Find and print error details from the failed step/task\n for job_step in status.steps or []:\n if job_step.status == "error":\n print(f"\\nFailed step: {job_step.name}")\n if job_step.error_details:\n step_error = job_step.error_details.get("message", "")\n if step_error:\n print(f"Step error: {step_error}")\n # Get error_stack from the failed task\n for task in job_step.tasks or []:\n if task.status == "error" and hasattr(task, "error_stack") and task.error_stack:\n print(f"\\nError stack trace:\\n{task.error_stack}")\n elif task.status == "error" and task.error_details:\n task_error = task.error_details.get("message", "")\n if task_error:\n print(f"Task error: {task_error}")\n break\n\n raise Exception(f"Job finished with status: {status.status}")\n\n time.sleep(poll_interval)\n\n\n# Wait for the job to complete\njob_with_sequence_packing_status = wait_for_job(\n workspace="default",\n job_name=job_with_sequence_packing.job.name,\n timeout=TIMEOUT_SECONDS,\n)\n\npacked_val_loss = get_training_metric(job_with_sequence_packing_status, VAL_LOSS_KEY)\nif packed_val_loss is not None:\n print(f"Validation loss: {packed_val_loss:.2f}")\nelse:\n print("Validation loss: not reported in job status")\n" }, { "type": "markdown", @@ -132,14 +132,14 @@ export default { cells: [ }, { "type": "markdown", - "source": "### 8. Track Finetuning Progress for Job without Sequence Packing", - "source_html": "

8. Track Finetuning Progress for Job without Sequence Packing

\n" + "source": "### 8. Track Fine-Tuning Progress for Job without Sequence Packing", + "source_html": "

8. Track Fine-Tuning Progress for Job without Sequence Packing

\n" }, { "type": "code", - "source": "# Wait for the training step to complete\njob_without_sequence_packing_status = wait_for_job(\n workspace=\"default\",\n job_name=job_without_sequence_packing.job.name,\n timeout=TIMEOUT_SECONDS\n)\n\nno_pack_val_loss = (job_without_sequence_packing_status.status_details or {}).get(\"val_loss\")\nif no_pack_val_loss is not None:\n print(f\"Validation loss: {float(no_pack_val_loss):.2f}\")\nelse:\n print(\"Validation loss: not reported in job status\")", + "source": "# Wait for the training step to complete\njob_without_sequence_packing_status = wait_for_job(\n workspace=\"default\",\n job_name=job_without_sequence_packing.job.name,\n timeout=TIMEOUT_SECONDS\n)\n\nno_pack_val_loss = get_training_metric(job_without_sequence_packing_status, VAL_LOSS_KEY)\nif no_pack_val_loss is not None:\n print(f\"Validation loss: {no_pack_val_loss:.2f}\")\nelse:\n print(\"Validation loss: not reported in job status\")", "language": "python", - "source_html": "# Wait for the training step to complete\njob_without_sequence_packing_status = wait_for_job(\n workspace="default",\n job_name=job_without_sequence_packing.job.name,\n timeout=TIMEOUT_SECONDS\n)\n\nno_pack_val_loss = (job_without_sequence_packing_status.status_details or {}).get("val_loss")\nif no_pack_val_loss is not None:\n print(f"Validation loss: {float(no_pack_val_loss):.2f}")\nelse:\n print("Validation loss: not reported in job status")\n" + "source_html": "# Wait for the training step to complete\njob_without_sequence_packing_status = wait_for_job(\n workspace="default",\n job_name=job_without_sequence_packing.job.name,\n timeout=TIMEOUT_SECONDS\n)\n\nno_pack_val_loss = get_training_metric(job_without_sequence_packing_status, VAL_LOSS_KEY)\nif no_pack_val_loss is not None:\n print(f"Validation loss: {no_pack_val_loss:.2f}")\nelse:\n print("Validation loss: not reported in job status")\n" }, { "type": "markdown", @@ -148,9 +148,9 @@ export default { cells: [ }, { "type": "code", - "source": "from nemo_platform.types.jobs import PlatformJobStep\nfrom datetime import datetime\nimport pandas as pd\n\nSTEP_NAME = \"training\"\n\ndef get_elapsed_time(step: PlatformJobStep) -> float:\n \"\"\"Calculate elapsed time in seconds from step's created_at to updated_at.\"\"\"\n created_at = datetime.fromisoformat(step.created_at.replace(\"Z\", \"+00:00\"))\n updated_at = datetime.fromisoformat(step.updated_at.replace(\"Z\", \"+00:00\"))\n return (updated_at - created_at).total_seconds()\n\nstep_with_sequence_packing = client.jobs.steps.retrieve(\n name=STEP_NAME,\n workspace=\"default\",\n job=job_with_sequence_packing.job.name,\n)\n\nstep_without_sequence_packing = client.jobs.steps.retrieve(\n name=STEP_NAME,\n workspace=\"default\",\n job=job_without_sequence_packing.job.name,\n)\n\ntime_to_complete_with_sequence_packing = get_elapsed_time(step_with_sequence_packing)\ntime_to_complete_without_sequence_packing = get_elapsed_time(step_without_sequence_packing)\n\n# Display results as a table\nresults_df = pd.DataFrame({\n \"Seq Packing Enabled\": [True, False],\n \"Val Loss\": [\n (job_with_sequence_packing_status.status_details or {}).get(\"val_loss\"),\n (job_without_sequence_packing_status.status_details or {}).get(\"val_loss\"),\n ],\n \"Training Step Time, sec\": [\n time_to_complete_with_sequence_packing,\n time_to_complete_without_sequence_packing\n ]\n})\n\nresults_df.style.format({\"Val Loss\": \"{:.2f}\", \"Training Step Time, sec\": \"{:.0f}\"}).hide(axis='index')", + "source": "from nemo_platform.types.jobs import PlatformJobStep\nimport pandas as pd\n\nSTEP_NAME = \"training\"\n\ndef get_elapsed_time(step: PlatformJobStep) -> float:\n \"\"\"Calculate elapsed time in seconds from step's created_at to updated_at.\"\"\"\n if step.created_at is None or step.updated_at is None:\n raise ValueError(\"Training step timestamps are unavailable\")\n return (step.updated_at - step.created_at).total_seconds()\n\nstep_with_sequence_packing = client.jobs.steps.retrieve(\n name=STEP_NAME,\n workspace=\"default\",\n job=job_with_sequence_packing.job.name,\n)\n\nstep_without_sequence_packing = client.jobs.steps.retrieve(\n name=STEP_NAME,\n workspace=\"default\",\n job=job_without_sequence_packing.job.name,\n)\n\ntime_to_complete_with_sequence_packing = get_elapsed_time(step_with_sequence_packing)\ntime_to_complete_without_sequence_packing = get_elapsed_time(step_without_sequence_packing)\n\n# Display results as a table\nresults_df = pd.DataFrame({\n \"Seq Packing Enabled\": [True, False],\n \"Val Loss\": [packed_val_loss, no_pack_val_loss],\n \"Training Step Time, sec\": [\n time_to_complete_with_sequence_packing,\n time_to_complete_without_sequence_packing\n ]\n})\n\nresults_df.style.format({\"Val Loss\": \"{:.2f}\", \"Training Step Time, sec\": \"{:.0f}\"}).hide(axis='index')", "language": "python", - "source_html": "from nemo_platform.types.jobs import PlatformJobStep\nfrom datetime import datetime\nimport pandas as pd\n\nSTEP_NAME = "training"\n\ndef get_elapsed_time(step: PlatformJobStep) -> float:\n """Calculate elapsed time in seconds from step's created_at to updated_at."""\n created_at = datetime.fromisoformat(step.created_at.replace("Z", "+00:00"))\n updated_at = datetime.fromisoformat(step.updated_at.replace("Z", "+00:00"))\n return (updated_at - created_at).total_seconds()\n\nstep_with_sequence_packing = client.jobs.steps.retrieve(\n name=STEP_NAME,\n workspace="default",\n job=job_with_sequence_packing.job.name,\n)\n\nstep_without_sequence_packing = client.jobs.steps.retrieve(\n name=STEP_NAME,\n workspace="default",\n job=job_without_sequence_packing.job.name,\n)\n\ntime_to_complete_with_sequence_packing = get_elapsed_time(step_with_sequence_packing)\ntime_to_complete_without_sequence_packing = get_elapsed_time(step_without_sequence_packing)\n\n# Display results as a table\nresults_df = pd.DataFrame({\n "Seq Packing Enabled": [True, False],\n "Val Loss": [\n (job_with_sequence_packing_status.status_details or {}).get("val_loss"),\n (job_without_sequence_packing_status.status_details or {}).get("val_loss"),\n ],\n "Training Step Time, sec": [\n time_to_complete_with_sequence_packing,\n time_to_complete_without_sequence_packing\n ]\n})\n\nresults_df.style.format({"Val Loss": "{:.2f}", "Training Step Time, sec": "{:.0f}"}).hide(axis='index')\n" + "source_html": "from nemo_platform.types.jobs import PlatformJobStep\nimport pandas as pd\n\nSTEP_NAME = "training"\n\ndef get_elapsed_time(step: PlatformJobStep) -> float:\n """Calculate elapsed time in seconds from step's created_at to updated_at."""\n if step.created_at is None or step.updated_at is None:\n raise ValueError("Training step timestamps are unavailable")\n return (step.updated_at - step.created_at).total_seconds()\n\nstep_with_sequence_packing = client.jobs.steps.retrieve(\n name=STEP_NAME,\n workspace="default",\n job=job_with_sequence_packing.job.name,\n)\n\nstep_without_sequence_packing = client.jobs.steps.retrieve(\n name=STEP_NAME,\n workspace="default",\n job=job_without_sequence_packing.job.name,\n)\n\ntime_to_complete_with_sequence_packing = get_elapsed_time(step_with_sequence_packing)\ntime_to_complete_without_sequence_packing = get_elapsed_time(step_without_sequence_packing)\n\n# Display results as a table\nresults_df = pd.DataFrame({\n "Seq Packing Enabled": [True, False],\n "Val Loss": [packed_val_loss, no_pack_val_loss],\n "Training Step Time, sec": [\n time_to_complete_with_sequence_packing,\n time_to_complete_without_sequence_packing\n ]\n})\n\nresults_df.style.format({"Val Loss": "{:.2f}", "Training Step Time, sec": "{:.0f}"}).hide(axis='index')\n" }, { "type": "markdown", diff --git a/docs/fern/components/notebooks/sft-customization-job.json b/docs/fern/components/notebooks/sft-customization-job.json index fc41a5d23c..8195bb759c 100644 --- a/docs/fern/components/notebooks/sft-customization-job.json +++ b/docs/fern/components/notebooks/sft-customization-job.json @@ -2,13 +2,13 @@ "cells": [ { "type": "markdown", - "source": "\n\n\n# Full SFT Customization\n\nLearn how to fine-tune all model weights using supervised fine-tuning (SFT) to customize LLM behavior for your specific tasks.\n\n## About\n\nSupervised Fine-Tuning (SFT) customizes model behavior, injects new knowledge, and optimizes performance for specific domains and tasks. Full SFT modifies **all model weights** during training, providing maximum customization flexibility.\n\n**What you can achieve with SFT:**\n\n- 🎯 **Specialize for domains:** Fine-tune models on legal texts, medical records, or financial data\n- 💡 **Inject knowledge:** Add new information not present in the base model\n- 📈 **Improve accuracy:** Optimize for specific tasks like sentiment analysis, summarization, or code generation\n\n### SFT vs LoRA: Understanding the Trade-offs\n\n**Full SFT** trains all model parameters (for example, all 70 billion weights in Llama 70B):\n\n- ✅ Maximum model adaptation and knowledge injection\n- ✅ Can fundamentally change model behavior\n- ✅ Best for significant domain shifts or specialized tasks\n- ❌ Requires substantial GPU resources (4-8x more than LoRA)\n- ❌ Produces full model weights (~140GB for Llama 70B)\n- ❌ Longer training time\n\n**LoRA** trains only ~1% of weights by adding thin matrices to existing weights:\n\n- ✅ 75-95% less memory required\n- ✅ Faster training (2-4x speedup)\n- ✅ Produces small adapter files (~100-500MB)\n- ✅ Multiple adapters can share one base model\n- ❌ Limited adaptation capability compared to full fine-tuning\n\n**When to choose Full SFT:**\n\n- Training small models (1B-8B) where resource cost is manageable\n- Need fundamental behavior changes (for example, medical diagnosis, legal reasoning)\n- Injecting substantial new knowledge not in the base model\n\n**When to choose LoRA:** Refer to the [LoRA tutorial](./lora-customization-job) for most use cases, especially with large models (70B+) or limited GPU resources.", - "source_html": "\n\n

Full SFT Customization

\n

Learn how to fine-tune all model weights using supervised fine-tuning (SFT) to customize LLM behavior for your specific tasks.

\n

About

\n

Supervised Fine-Tuning (SFT) customizes model behavior, injects new knowledge, and optimizes performance for specific domains and tasks. Full SFT modifies all model weights during training, providing maximum customization flexibility.

\n

What you can achieve with SFT:

\n
    \n
  • 🎯 Specialize for domains: Fine-tune models on legal texts, medical records, or financial data
    • \n
  • 💡 Inject knowledge: Add new information not present in the base model
    • \n
  • 📈 Improve accuracy: Optimize for specific tasks like sentiment analysis, summarization, or code generation
    • \n
\n

SFT vs LoRA: Understanding the Trade-offs

\n

Full SFT trains all model parameters (for example, all 70 billion weights in Llama 70B):

\n
    \n
  • ✅ Maximum model adaptation and knowledge injection
    • \n
  • ✅ Can fundamentally change model behavior
    • \n
  • ✅ Best for significant domain shifts or specialized tasks
    • \n
  • ❌ Requires substantial GPU resources (4-8x more than LoRA)
    • \n
  • ❌ Produces full model weights (~140GB for Llama 70B)
    • \n
  • ❌ Longer training time
    • \n
\n

LoRA trains only ~1% of weights by adding thin matrices to existing weights:

\n
    \n
  • ✅ 75-95% less memory required
    • \n
  • ✅ Faster training (2-4x speedup)
    • \n
  • ✅ Produces small adapter files (~100-500MB)
    • \n
  • ✅ Multiple adapters can share one base model
    • \n
  • ❌ Limited adaptation capability compared to full fine-tuning
    • \n
\n

When to choose Full SFT:

\n
    \n
  • Training small models (1B-8B) where resource cost is manageable
    • \n
  • Need fundamental behavior changes (for example, medical diagnosis, legal reasoning)
    • \n
  • Injecting substantial new knowledge not in the base model
    • \n
\n

When to choose LoRA: Refer to the LoRA tutorial for most use cases, especially with large models (70B+) or limited GPU resources.

\n" + "source": "\n\n\n# Full SFT Customization\n\nLearn how to fine-tune all model weights using supervised fine-tuning (SFT) to customize LLM behavior for your specific tasks.\n\n## About\n\nSupervised Fine-Tuning (SFT) customizes model behavior, injects new knowledge, and optimizes performance for specific domains and tasks. Full SFT modifies **all model weights** during training, providing maximum customization flexibility.\n\n**What you can achieve with SFT:**\n\n- 🎯 **Specialize for domains:** Fine-tune models on legal texts, medical records, or financial data\n- 💡 **Inject knowledge:** Add new information not present in the base model\n- 📈 **Improve accuracy:** Optimize for specific tasks like sentiment analysis, summarization, or code generation\n\n### SFT vs LoRA: Understanding the Trade-offs\n\n**Full SFT** trains all model parameters (for example, all 70 billion weights in Llama 70B):\n\n- ✅ Maximum model adaptation and knowledge injection\n- ✅ Can fundamentally change model behavior\n- ✅ Best for significant domain shifts or specialized tasks\n- ❌ Requires substantial GPU resources (4-8x more than LoRA)\n- ❌ Produces a full BF16 checkpoint (~140 GB for Llama 70B); peak job disk usage can reach approximately 3× the downloaded base checkpoint size\n- ❌ Longer training time\n\n**LoRA** trains only ~1% of weights by adding thin matrices to existing weights:\n\n- ✅ 75-95% less memory required\n- ✅ Faster training (2-4x speedup)\n- ✅ Produces small adapter files (~100-500MB)\n- ✅ Multiple adapters can share one base model\n- ❌ Limited adaptation capability compared to full fine-tuning\n\n**When to choose Full SFT:**\n\n- Training small models (1B-8B) where resource cost is manageable\n- Need fundamental behavior changes (for example, medical diagnosis, legal reasoning)\n- Injecting substantial new knowledge not in the base model\n\n**When to choose LoRA:** Refer to the [LoRA tutorial](./lora-customization-job) for most use cases, especially with large models (70B+) or limited GPU resources.", + "source_html": "\n\n

Full SFT Customization

\n

Learn how to fine-tune all model weights using supervised fine-tuning (SFT) to customize LLM behavior for your specific tasks.

\n

About

\n

Supervised Fine-Tuning (SFT) customizes model behavior, injects new knowledge, and optimizes performance for specific domains and tasks. Full SFT modifies all model weights during training, providing maximum customization flexibility.

\n

What you can achieve with SFT:

\n
    \n
  • 🎯 Specialize for domains: Fine-tune models on legal texts, medical records, or financial data
    • \n
  • 💡 Inject knowledge: Add new information not present in the base model
    • \n
  • 📈 Improve accuracy: Optimize for specific tasks like sentiment analysis, summarization, or code generation
    • \n
\n

SFT vs LoRA: Understanding the Trade-offs

\n

Full SFT trains all model parameters (for example, all 70 billion weights in Llama 70B):

\n
    \n
  • ✅ Maximum model adaptation and knowledge injection
    • \n
  • ✅ Can fundamentally change model behavior
    • \n
  • ✅ Best for significant domain shifts or specialized tasks
    • \n
  • ❌ Requires substantial GPU resources (4-8x more than LoRA)
    • \n
  • ❌ Produces a full BF16 checkpoint (~140 GB for Llama 70B); peak job disk usage can reach approximately 3× the downloaded base checkpoint size
    • \n
  • ❌ Longer training time
    • \n
\n

LoRA trains only ~1% of weights by adding thin matrices to existing weights:

\n
    \n
  • ✅ 75-95% less memory required
    • \n
  • ✅ Faster training (2-4x speedup)
    • \n
  • ✅ Produces small adapter files (~100-500MB)
    • \n
  • ✅ Multiple adapters can share one base model
    • \n
  • ❌ Limited adaptation capability compared to full fine-tuning
    • \n
\n

When to choose Full SFT:

\n
    \n
  • Training small models (1B-8B) where resource cost is manageable
    • \n
  • Need fundamental behavior changes (for example, medical diagnosis, legal reasoning)
    • \n
  • Injecting substantial new knowledge not in the base model
    • \n
\n

When to choose LoRA: Refer to the LoRA tutorial for most use cases, especially with large models (70B+) or limited GPU resources.

\n" }, { "type": "markdown", - "source": "## Prerequisites\n\nBefore starting this tutorial, ensure you have:\n\n1. **Completed the [Quickstart](../../get-started/quickstart.md)** to install and deploy NeMo Platform locally\n2. **Installed the Python SDK** (PyPI wrapper: `pip install \"nemo-platform[all]\"`; source checkout: run `make bootstrap` from the repository root)", - "source_html": "

Prerequisites

\n

Before starting this tutorial, ensure you have:

\n
    \n
  1. Completed the Quickstart to install and deploy NeMo Platform locally
    2. \n
  2. Installed the Python SDK (PyPI wrapper: pip install "nemo-platform[all]"; source checkout: run make bootstrap from the repository root)
    3. \n
\n" + "source": "## Prerequisites\n\nBefore starting this tutorial, ensure you have:\n\n1. **Completed the [Quickstart](../../get-started/quickstart.md)** to install and deploy NeMo Platform locally\n2. **Installed the Python SDK** (PyPI wrapper: `pip install \"nemo-platform[all]\"`; source checkout: run `make bootstrap` from the repository root)\n3. **At least one GPU with CUDA 13+**", + "source_html": "

Prerequisites

\n

Before starting this tutorial, ensure you have:

\n
    \n
  1. Completed the Quickstart to install and deploy NeMo Platform locally
    2. \n
  2. Installed the Python SDK (PyPI wrapper: pip install "nemo-platform[all]"; source checkout: run make bootstrap from the repository root)
    3. \n
  3. At least one GPU with CUDA 13+
    4. \n
\n" }, { "type": "markdown", @@ -85,30 +85,30 @@ }, { "type": "markdown", - "source": "### 4. Secrets Setup\n\nIf you plan to use NGC or HuggingFace models, you will need to configure authentication:\n\n- **NGC models** (`ngc://` URIs): Requires NGC API key\n- **HuggingFace models** (`hf://` URIs): Requires HF token for gated/private models\n\n\nConfigure these as secrets in your platform. Refer to [Managing Secrets](../../get-started/concepts/manage-secrets.md) for detailed instructions.\n\nGet your credentials to access base models:\n- [NGC API Key](https://ngc.nvidia.com/) (Setup → Generate API Key)\n- [HuggingFace Token](https://huggingface.co/settings/tokens) (Create token with Read access)\n\n\n---\n\n#### Quick Setup Example\n\nIn this tutorial we are going to work with [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main) model from HuggingFace. Ensure that you have sufficient permissions to download the model. If you cannot access the files on the [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main) Hugging Face page, request access\n\n**HuggingFace Authentication:**\n- For gated models (Llama, Gemma), you must provide a HuggingFace token via the `token_secret` parameter\n- Get your token from [HuggingFace Settings](https://huggingface.co/settings/tokens) (requires Read access)\n- Accept the model's terms on the HuggingFace model page before using it. Example: [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main)\n- For public models, you can omit the `token_secret` parameter when creating a fileset for model in the next step", - "source_html": "

4. Secrets Setup

\n

If you plan to use NGC or HuggingFace models, you will need to configure authentication:

\n
    \n
  • NGC models (ngc:// URIs): Requires NGC API key
    • \n
  • HuggingFace models (hf:// URIs): Requires HF token for gated/private models
    • \n
\n

Configure these as secrets in your platform. Refer to Managing Secrets for detailed instructions.

\n

Get your credentials to access base models:

\n\n
\n

Quick Setup Example

\n

In this tutorial we are going to work with meta-llama/Llama-3.2-1B-Instruct model from HuggingFace. Ensure that you have sufficient permissions to download the model. If you cannot access the files on the meta-llama/Llama-3.2-1B-Instruct Hugging Face page, request access

\n

HuggingFace Authentication:

\n
    \n
  • For gated models (Llama, Gemma), you must provide a HuggingFace token via the token_secret parameter
    • \n
  • Get your token from HuggingFace Settings (requires Read access)
    • \n
  • Accept the model's terms on the HuggingFace model page before using it. Example: meta-llama/Llama-3.2-1B-Instruct
    • \n
  • For public models, you can omit the token_secret parameter when creating a fileset for model in the next step
    • \n
\n" + "source": "### 4. Secrets Setup\n\nIf you plan to use NGC or Hugging Face models, you will need to configure authentication:\n\n- **NGC models** (`ngc://` URIs): Requires NGC API key\n- **Hugging Face models** (`hf://` URIs): Requires HF token for gated/private models\n\n\nConfigure these as secrets in your platform. Refer to [Managing Secrets](../../get-started/concepts/manage-secrets.md) for detailed instructions.\n\nGet your credentials to access base models:\n- [NGC API Key](https://ngc.nvidia.com/) (Setup → Generate API Key)\n- [Hugging Face Token](https://huggingface.co/settings/tokens) (Create token with Read access)\n\n\n---\n\n#### Quick Setup Example\n\nIn this tutorial we are going to work with the [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main) model from Hugging Face. Ensure that you have sufficient permissions to download the model. If you cannot access the files on the [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main) Hugging Face page, request access.\n\n**Hugging Face Authentication:**\n- For gated models (Llama, Gemma), you must provide a Hugging Face token via the `token_secret` parameter\n- Get your token from [Hugging Face Settings](https://huggingface.co/settings/tokens) (requires Read access)\n- Accept the model's terms on the Hugging Face model page before using it. Example: [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main)\n- For public models, you can omit the `token_secret` parameter when creating a fileset for model in the next step", + "source_html": "

4. Secrets Setup

\n

If you plan to use NGC or Hugging Face models, you will need to configure authentication:

\n
    \n
  • NGC models (ngc:// URIs): Requires NGC API key
    • \n
  • Hugging Face models (hf:// URIs): Requires HF token for gated/private models
    • \n
\n

Configure these as secrets in your platform. Refer to Managing Secrets for detailed instructions.

\n

Get your credentials to access base models:

\n\n
\n

Quick Setup Example

\n

In this tutorial we are going to work with the meta-llama/Llama-3.2-1B-Instruct model from Hugging Face. Ensure that you have sufficient permissions to download the model. If you cannot access the files on the meta-llama/Llama-3.2-1B-Instruct Hugging Face page, request access.

\n

Hugging Face Authentication:

\n
    \n
  • For gated models (Llama, Gemma), you must provide a Hugging Face token via the token_secret parameter
    • \n
  • Get your token from Hugging Face Settings (requires Read access)
    • \n
  • Accept the model's terms on the Hugging Face model page before using it. Example: meta-llama/Llama-3.2-1B-Instruct
    • \n
  • For public models, you can omit the token_secret parameter when creating a fileset for model in the next step
    • \n
\n" }, { "type": "code", - "source": "# Export the HF_TOKEN and NGC_API_KEY environment variables if they are not already set\nHF_TOKEN = os.getenv(\"HF_TOKEN\")\nNGC_API_KEY = os.getenv(\"NGC_API_KEY\")\n\n\ndef create_or_get_secret(name: str, value: str | None, label: str):\n if not value:\n raise ValueError(f\"{label} is not set\")\n try:\n secret = client.secrets.create(\n name=name,\n workspace=\"default\",\n value=value,\n )\n print(f\"Created secret: {name}\")\n return secret\n except ConflictError:\n print(f\"Secret '{name}' already exists, continuing...\")\n return client.secrets.retrieve(name=name, workspace=\"default\")\n\n\n# Create HuggingFace token secret\nhf_secret = create_or_get_secret(\"hf-token\", HF_TOKEN, \"HF_TOKEN\")\nprint(\"HF_TOKEN secret:\")\nprint(hf_secret.model_dump_json(indent=2))\n\n# Create NGC API key secret\n# Uncomment the line below if you have NGC API Key and want to finetune NGC models\n# ngc_api_key = create_or_get_secret(\"ngc-api-key\", NGC_API_KEY, \"NGC_API_KEY\")", + "source": "# Export the HF_TOKEN and NGC_API_KEY environment variables if they are not already set\nHF_TOKEN = os.getenv(\"HF_TOKEN\")\nNGC_API_KEY = os.getenv(\"NGC_API_KEY\")\n\n\ndef create_or_get_secret(name: str, value: str | None, label: str):\n if not value:\n raise ValueError(f\"{label} is not set\")\n try:\n secret = client.secrets.create(\n name=name,\n workspace=\"default\",\n value=value,\n )\n print(f\"Created secret: {name}\")\n return secret\n except ConflictError:\n print(f\"Secret '{name}' already exists, continuing...\")\n return client.secrets.retrieve(name=name, workspace=\"default\")\n\n\n# Create Hugging Face token secret\nhf_secret = create_or_get_secret(\"hf-token\", HF_TOKEN, \"HF_TOKEN\")\nprint(\"HF_TOKEN secret:\")\nprint(hf_secret.model_dump_json(indent=2))\n\n# Create NGC API key secret\n# Uncomment the line below if you have NGC API Key and want to finetune NGC models\n# ngc_api_key = create_or_get_secret(\"ngc-api-key\", NGC_API_KEY, \"NGC_API_KEY\")", "language": "python", - "source_html": "# Export the HF_TOKEN and NGC_API_KEY environment variables if they are not already set\nHF_TOKEN = os.getenv("HF_TOKEN")\nNGC_API_KEY = os.getenv("NGC_API_KEY")\n\n\ndef create_or_get_secret(name: str, value: str | None, label: str):\n if not value:\n raise ValueError(f"{label} is not set")\n try:\n secret = client.secrets.create(\n name=name,\n workspace="default",\n value=value,\n )\n print(f"Created secret: {name}")\n return secret\n except ConflictError:\n print(f"Secret '{name}' already exists, continuing...")\n return client.secrets.retrieve(name=name, workspace="default")\n\n\n# Create HuggingFace token secret\nhf_secret = create_or_get_secret("hf-token", HF_TOKEN, "HF_TOKEN")\nprint("HF_TOKEN secret:")\nprint(hf_secret.model_dump_json(indent=2))\n\n# Create NGC API key secret\n# Uncomment the line below if you have NGC API Key and want to finetune NGC models\n# ngc_api_key = create_or_get_secret("ngc-api-key", NGC_API_KEY, "NGC_API_KEY")\n" + "source_html": "# Export the HF_TOKEN and NGC_API_KEY environment variables if they are not already set\nHF_TOKEN = os.getenv("HF_TOKEN")\nNGC_API_KEY = os.getenv("NGC_API_KEY")\n\n\ndef create_or_get_secret(name: str, value: str | None, label: str):\n if not value:\n raise ValueError(f"{label} is not set")\n try:\n secret = client.secrets.create(\n name=name,\n workspace="default",\n value=value,\n )\n print(f"Created secret: {name}")\n return secret\n except ConflictError:\n print(f"Secret '{name}' already exists, continuing...")\n return client.secrets.retrieve(name=name, workspace="default")\n\n\n# Create Hugging Face token secret\nhf_secret = create_or_get_secret("hf-token", HF_TOKEN, "HF_TOKEN")\nprint("HF_TOKEN secret:")\nprint(hf_secret.model_dump_json(indent=2))\n\n# Create NGC API key secret\n# Uncomment the line below if you have NGC API Key and want to finetune NGC models\n# ngc_api_key = create_or_get_secret("ngc-api-key", NGC_API_KEY, "NGC_API_KEY")\n" }, { "type": "markdown", - "source": "### 5. Create Base Model FileSet and Model Entity\n\nCreate a fileset pointing to [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main) model in HuggingFace that we will train with SFT. Then create a Model Entity that references this fileset. Model downloading will take place at training time.\n\nNote: for public models, you can omit the `token_secret` parameter when creating a model fileset.", - "source_html": "

5. Create Base Model FileSet and Model Entity

\n

Create a fileset pointing to meta-llama/Llama-3.2-1B-Instruct model in HuggingFace that we will train with SFT. Then create a Model Entity that references this fileset. Model downloading will take place at training time.

\n

Note: for public models, you can omit the token_secret parameter when creating a model fileset.

\n" + "source": "### 5. Create Base Model FileSet and Model Entity\n\nCreate a fileset pointing to the [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main) model in Hugging Face that we will train with SFT. Then create a Model Entity that references this fileset. Model downloading will take place at training time.\n\nNote: for public models, you can omit the `token_secret` parameter when creating a model fileset.", + "source_html": "

5. Create Base Model FileSet and Model Entity

\n

Create a fileset pointing to the meta-llama/Llama-3.2-1B-Instruct model in Hugging Face that we will train with SFT. Then create a Model Entity that references this fileset. Model downloading will take place at training time.

\n

Note: for public models, you can omit the token_secret parameter when creating a model fileset.

\n" }, { "type": "code", - "source": "import time\n\n# Create a fileset pointing to the desired HuggingFace model\nfrom nemo_platform.types.files import HuggingfaceStorageConfigParam\n\nHF_REPO_ID = \"meta-llama/Llama-3.2-1B-Instruct\"\nMODEL_NAME = \"llama-3-2-1b-base\"\n\n# Ensure you have a HuggingFace token secret created\ntry:\n base_model_fs = client.files.filesets.create(\n workspace=\"default\",\n name=MODEL_NAME,\n description=\"Llama 3.2 1B base model from HuggingFace\",\n storage=HuggingfaceStorageConfigParam(\n type=\"huggingface\",\n # repo_id is the full model name from Hugging Face\n repo_id=HF_REPO_ID,\n repo_type=\"model\",\n # we use the secret created in the previous step\n token_secret=hf_secret.name\n )\n )\n print(f\"Created base model fileset: {MODEL_NAME}\")\nexcept ConflictError:\n print(f\"Base model fileset already exists. Skipping creation.\")\n base_model_fs = client.files.filesets.retrieve(\n workspace=\"default\",\n name=MODEL_NAME,\n )\n\n# Create the Model Entity representation.\ntry:\n base_model = client.models.create(\n workspace=\"default\",\n name=MODEL_NAME,\n fileset=f\"default/{MODEL_NAME}\",\n )\n print(f\"Created Model Entity: {MODEL_NAME}\")\nexcept ConflictError:\n print(f\"Base model already exists. Updating fileset if different.\")\n base_model = client.models.update(\n workspace=\"default\",\n name=MODEL_NAME,\n fileset=f\"default/{MODEL_NAME}\",\n )\n\nprint(f\"\\nBase model fileset: fileset://default/{base_model.name}\")\nprint(\"Base model fileset files list:\")\nprint(json.dumps([f.model_dump() for f in client.files.list(fileset=MODEL_NAME, workspace=\"default\").data], indent=2))\n\n# Wait for ModelSpec to be populated from the checkpoint\nprint(\"\\nWaiting for ModelSpec to be populated...\")\nSPEC_TIMEOUT_SECONDS = 120\nspec_start = time.time()\nwhile not base_model.spec:\n if time.time() - spec_start > SPEC_TIMEOUT_SECONDS:\n raise TimeoutError(f\"ModelSpec not populated within {SPEC_TIMEOUT_SECONDS} seconds\")\n time.sleep(2)\n base_model = client.models.retrieve(\n workspace=\"default\",\n name=MODEL_NAME,\n )\n\nprint(f\"ModelSpec populated: {base_model.spec}\")", + "source": "import time\n\n# Create a fileset pointing to the desired Hugging Face model\nfrom nemo_platform.types.files import HuggingfaceStorageConfigParam\n\nHF_REPO_ID = \"meta-llama/Llama-3.2-1B-Instruct\"\nMODEL_NAME = \"llama-3-2-1b-base\"\n\n# Ensure you have a Hugging Face token secret created\ntry:\n base_model_fs = client.files.filesets.create(\n workspace=\"default\",\n name=MODEL_NAME,\n description=\"Llama 3.2 1B base model from Hugging Face\",\n storage=HuggingfaceStorageConfigParam(\n type=\"huggingface\",\n # repo_id is the full model name from Hugging Face\n repo_id=HF_REPO_ID,\n repo_type=\"model\",\n # we use the secret created in the previous step\n token_secret=hf_secret.name\n )\n )\n print(f\"Created base model fileset: {MODEL_NAME}\")\nexcept ConflictError:\n print(f\"Base model fileset already exists. Skipping creation.\")\n base_model_fs = client.files.filesets.retrieve(\n workspace=\"default\",\n name=MODEL_NAME,\n )\n\n# Create the Model Entity representation.\ntry:\n base_model = client.models.create(\n workspace=\"default\",\n name=MODEL_NAME,\n fileset=f\"default/{MODEL_NAME}\",\n )\n print(f\"Created Model Entity: {MODEL_NAME}\")\nexcept ConflictError:\n print(f\"Base model already exists. Updating fileset if different.\")\n base_model = client.models.update(\n workspace=\"default\",\n name=MODEL_NAME,\n fileset=f\"default/{MODEL_NAME}\",\n )\n\nprint(f\"\\nBase model fileset: fileset://default/{base_model.name}\")\nprint(\"Base model fileset files list:\")\nprint(json.dumps([f.model_dump() for f in client.files.list(fileset=MODEL_NAME, workspace=\"default\").data], indent=2))\n\n# Wait for ModelSpec to be populated from the checkpoint\nprint(\"\\nWaiting for ModelSpec to be populated...\")\nSPEC_TIMEOUT_SECONDS = 120\nspec_start = time.time()\nwhile not base_model.spec:\n if time.time() - spec_start > SPEC_TIMEOUT_SECONDS:\n raise TimeoutError(f\"ModelSpec not populated within {SPEC_TIMEOUT_SECONDS} seconds\")\n time.sleep(2)\n base_model = client.models.retrieve(\n workspace=\"default\",\n name=MODEL_NAME,\n )\n\nprint(f\"ModelSpec populated: {base_model.spec}\")", "language": "python", - "source_html": "import time\n\n# Create a fileset pointing to the desired HuggingFace model\nfrom nemo_platform.types.files import HuggingfaceStorageConfigParam\n\nHF_REPO_ID = "meta-llama/Llama-3.2-1B-Instruct"\nMODEL_NAME = "llama-3-2-1b-base"\n\n# Ensure you have a HuggingFace token secret created\ntry:\n base_model_fs = client.files.filesets.create(\n workspace="default",\n name=MODEL_NAME,\n description="Llama 3.2 1B base model from HuggingFace",\n storage=HuggingfaceStorageConfigParam(\n type="huggingface",\n # repo_id is the full model name from Hugging Face\n repo_id=HF_REPO_ID,\n repo_type="model",\n # we use the secret created in the previous step\n token_secret=hf_secret.name\n )\n )\n print(f"Created base model fileset: {MODEL_NAME}")\nexcept ConflictError:\n print(f"Base model fileset already exists. Skipping creation.")\n base_model_fs = client.files.filesets.retrieve(\n workspace="default",\n name=MODEL_NAME,\n )\n\n# Create the Model Entity representation.\ntry:\n base_model = client.models.create(\n workspace="default",\n name=MODEL_NAME,\n fileset=f"default/{MODEL_NAME}",\n )\n print(f"Created Model Entity: {MODEL_NAME}")\nexcept ConflictError:\n print(f"Base model already exists. Updating fileset if different.")\n base_model = client.models.update(\n workspace="default",\n name=MODEL_NAME,\n fileset=f"default/{MODEL_NAME}",\n )\n\nprint(f"\\nBase model fileset: fileset://default/{base_model.name}")\nprint("Base model fileset files list:")\nprint(json.dumps([f.model_dump() for f in client.files.list(fileset=MODEL_NAME, workspace="default").data], indent=2))\n\n# Wait for ModelSpec to be populated from the checkpoint\nprint("\\nWaiting for ModelSpec to be populated...")\nSPEC_TIMEOUT_SECONDS = 120\nspec_start = time.time()\nwhile not base_model.spec:\n if time.time() - spec_start > SPEC_TIMEOUT_SECONDS:\n raise TimeoutError(f"ModelSpec not populated within {SPEC_TIMEOUT_SECONDS} seconds")\n time.sleep(2)\n base_model = client.models.retrieve(\n workspace="default",\n name=MODEL_NAME,\n )\n\nprint(f"ModelSpec populated: {base_model.spec}")\n" + "source_html": "import time\n\n# Create a fileset pointing to the desired Hugging Face model\nfrom nemo_platform.types.files import HuggingfaceStorageConfigParam\n\nHF_REPO_ID = "meta-llama/Llama-3.2-1B-Instruct"\nMODEL_NAME = "llama-3-2-1b-base"\n\n# Ensure you have a Hugging Face token secret created\ntry:\n base_model_fs = client.files.filesets.create(\n workspace="default",\n name=MODEL_NAME,\n description="Llama 3.2 1B base model from Hugging Face",\n storage=HuggingfaceStorageConfigParam(\n type="huggingface",\n # repo_id is the full model name from Hugging Face\n repo_id=HF_REPO_ID,\n repo_type="model",\n # we use the secret created in the previous step\n token_secret=hf_secret.name\n )\n )\n print(f"Created base model fileset: {MODEL_NAME}")\nexcept ConflictError:\n print(f"Base model fileset already exists. Skipping creation.")\n base_model_fs = client.files.filesets.retrieve(\n workspace="default",\n name=MODEL_NAME,\n )\n\n# Create the Model Entity representation.\ntry:\n base_model = client.models.create(\n workspace="default",\n name=MODEL_NAME,\n fileset=f"default/{MODEL_NAME}",\n )\n print(f"Created Model Entity: {MODEL_NAME}")\nexcept ConflictError:\n print(f"Base model already exists. Updating fileset if different.")\n base_model = client.models.update(\n workspace="default",\n name=MODEL_NAME,\n fileset=f"default/{MODEL_NAME}",\n )\n\nprint(f"\\nBase model fileset: fileset://default/{base_model.name}")\nprint("Base model fileset files list:")\nprint(json.dumps([f.model_dump() for f in client.files.list(fileset=MODEL_NAME, workspace="default").data], indent=2))\n\n# Wait for ModelSpec to be populated from the checkpoint\nprint("\\nWaiting for ModelSpec to be populated...")\nSPEC_TIMEOUT_SECONDS = 120\nspec_start = time.time()\nwhile not base_model.spec:\n if time.time() - spec_start > SPEC_TIMEOUT_SECONDS:\n raise TimeoutError(f"ModelSpec not populated within {SPEC_TIMEOUT_SECONDS} seconds")\n time.sleep(2)\n base_model = client.models.retrieve(\n workspace="default",\n name=MODEL_NAME,\n )\n\nprint(f"ModelSpec populated: {base_model.spec}")\n" }, { "type": "markdown", - "source": "### 6. Create SFT Finetuning Job\nCreate a customization job to fine-tune all model weights using the **Automodel** backend and `AutomodelJobInput`.", - "source_html": "

6. Create SFT Finetuning Job

\n

Create a customization job to fine-tune all model weights using the Automodel backend and AutomodelJobInput.

\n" + "source": "### 6. Create SFT Fine-Tuning Job\nCreate a customization job to fine-tune all model weights using the **Automodel** backend and `AutomodelJobInput`.", + "source_html": "

6. Create SFT Fine-Tuning Job

\n

Create a customization job to fine-tune all model weights using the Automodel backend and AutomodelJobInput.

\n" }, { "type": "markdown", @@ -183,8 +183,8 @@ }, { "type": "markdown", - "source": "#### Evaluation Best Practices\n\n**Manual Evaluation** (Recommended)\n- Test with real-world examples from your use case\n- Compare responses to base model and expected outputs\n- Verify the model exhibits desired behavior changes\n- Check edge cases and error handling\n\n**What to look for:**\n- ✅ Model follows your desired output format\n- ✅ Applies domain knowledge correctly\n- ✅ Maintains general language capabilities\n- ✅ Avoids unwanted behaviors or biases\n- ❌ Doesn't hallucinate facts not in training data\n- ❌ Doesn't produce repetitive or nonsensical outputs\n\n---\n\n## Hyperparameters\n\nFor detailed information on all available hyperparameters, recommended values, and tuning guidance, refer to the [Hyperparameter Reference](../manage-customization-jobs/hyperparameters.md).\n\n---\n\n\n## Troubleshooting\n\n**Job fails during model download:**\n- Verify authentication secrets are configured (refer to [Managing Secrets](../../get-started/concepts/manage-secrets.md))\n- For gated HuggingFace models (Llama, Gemma), accept the license on the model page (for example, [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct))\n- Confirm the model fileset uses `token_secret=hf_secret.name` for gated models\n- Check `AutomodelJobInput` references use the `workspace/name` format: `model=f\"default/{MODEL_NAME}\"` and `dataset={\"training\": f\"default/{DATASET_NAME}\"}` (for example, `default/llama-3-2-1b-base`, `default/sft-dataset`)\n- Verify the model entity points at the fileset: `fileset=f\"default/{MODEL_NAME}\"`\n- Check job status: `client.jobs.get_status(name=job.job.name, workspace=\"default\")`\n\n**Job fails with OOM (Out of Memory) error:**\n1. **First try:** Reduce `global_batch_size` from 64 to 32 or 16 in `batch={...}`\n2. **Still OOM:** Keep `micro_batch_size` at 1 (already the minimum in this tutorial)\n3. **Still OOM:** Reduce `max_seq_length` from 2048 to 1024 or 512 in `training={...}`\n4. **Last resort:** Increase `num_gpus_per_node` and `tensor_parallel_size` in `parallelism={...}`\n\n**Loss curves not decreasing (underfitting):**\n- Increase training duration: raise `epochs` from 2 to 3-5 in `schedule={...}`\n- Adjust learning rate: try `1e-4` or `1e-5` instead of the default `5e-5` in `optimizer={...}`\n- Check data quality: Verify formatting, remove duplicates, ensure diversity\n\n**Training loss decreases but validation loss increases (overfitting):**\n- Reduce `epochs` from 2 to 1 in `schedule={...}`\n- Lower `learning_rate` from `5e-5` to `2e-5` or `1e-5` in `optimizer={...}`\n- Increase dataset size and diversity\n- Verify train/validation split has no data leakage\n\n**Model output quality is poor despite good training metrics:**\n- Training metrics optimize for loss, not your actual task—evaluate on real use cases\n- Review data quality, format, and diversity—metrics can be misleading with poor data\n- Try a different base model size or architecture\n- Adjust `learning_rate` and `global_batch_size`\n- Compare to baseline: Test base model to ensure fine-tuning improved performance\n\n**Deployment fails:**\n- Verify output model exists: `client.models.retrieve(name=OUTPUT_NAME, workspace=\"default\")`\n- Check deployment logs: `client.inference.deployments.get_logs(name=deployment.name, workspace=\"default\")`\n- Ensure sufficient GPU resources for `executor_config={\"gpu\": 1, ...}`\n- Verify the deployment config matches this tutorial: `engine=\"vllm\"` with `vllm/vllm-openai:v0.22.1`\n\n\n## Next Steps\n\n- [Monitor training metrics](fine-tune-metrics) in detail\n- [Evaluate your fine-tuned model](../../evaluator/index) using the Evaluator service\n- Learn about [LoRA customization](./lora-customization-job) for resource-efficient fine-tuning", - "source_html": "

Evaluation Best Practices

\n

Manual Evaluation (Recommended)

\n
    \n
  • Test with real-world examples from your use case
    • \n
  • Compare responses to base model and expected outputs
    • \n
  • Verify the model exhibits desired behavior changes
    • \n
  • Check edge cases and error handling
    • \n
\n

What to look for:

\n
    \n
  • ✅ Model follows your desired output format
    • \n
  • ✅ Applies domain knowledge correctly
    • \n
  • ✅ Maintains general language capabilities
    • \n
  • ✅ Avoids unwanted behaviors or biases
    • \n
  • ❌ Doesn't hallucinate facts not in training data
    • \n
  • ❌ Doesn't produce repetitive or nonsensical outputs
    • \n
\n
\n

Hyperparameters

\n

For detailed information on all available hyperparameters, recommended values, and tuning guidance, refer to the Hyperparameter Reference.

\n
\n

Troubleshooting

\n

Job fails during model download:

\n
    \n
  • Verify authentication secrets are configured (refer to Managing Secrets)
    • \n
  • For gated HuggingFace models (Llama, Gemma), accept the license on the model page (for example, meta-llama/Llama-3.2-1B-Instruct)
    • \n
  • Confirm the model fileset uses token_secret=hf_secret.name for gated models
    • \n
  • Check AutomodelJobInput references use the workspace/name format: model=f"default/{MODEL_NAME}" and dataset={"training": f"default/{DATASET_NAME}"} (for example, default/llama-3-2-1b-base, default/sft-dataset)
    • \n
  • Verify the model entity points at the fileset: fileset=f"default/{MODEL_NAME}"
    • \n
  • Check job status: client.jobs.get_status(name=job.job.name, workspace="default")
    • \n
\n

Job fails with OOM (Out of Memory) error:

\n
    \n
  1. First try: Reduce global_batch_size from 64 to 32 or 16 in batch={...}
    2. \n
  2. Still OOM: Keep micro_batch_size at 1 (already the minimum in this tutorial)
    3. \n
  3. Still OOM: Reduce max_seq_length from 2048 to 1024 or 512 in training={...}
    4. \n
  4. Last resort: Increase num_gpus_per_node and tensor_parallel_size in parallelism={...}
    5. \n
\n

Loss curves not decreasing (underfitting):

\n
    \n
  • Increase training duration: raise epochs from 2 to 3-5 in schedule={...}
    • \n
  • Adjust learning rate: try 1e-4 or 1e-5 instead of the default 5e-5 in optimizer={...}
    • \n
  • Check data quality: Verify formatting, remove duplicates, ensure diversity
    • \n
\n

Training loss decreases but validation loss increases (overfitting):

\n
    \n
  • Reduce epochs from 2 to 1 in schedule={...}
    • \n
  • Lower learning_rate from 5e-5 to 2e-5 or 1e-5 in optimizer={...}
    • \n
  • Increase dataset size and diversity
    • \n
  • Verify train/validation split has no data leakage
    • \n
\n

Model output quality is poor despite good training metrics:

\n
    \n
  • Training metrics optimize for loss, not your actual task—evaluate on real use cases
    • \n
  • Review data quality, format, and diversity—metrics can be misleading with poor data
    • \n
  • Try a different base model size or architecture
    • \n
  • Adjust learning_rate and global_batch_size
    • \n
  • Compare to baseline: Test base model to ensure fine-tuning improved performance
    • \n
\n

Deployment fails:

\n
    \n
  • Verify output model exists: client.models.retrieve(name=OUTPUT_NAME, workspace="default")
    • \n
  • Check deployment logs: client.inference.deployments.get_logs(name=deployment.name, workspace="default")
    • \n
  • Ensure sufficient GPU resources for executor_config={"gpu": 1, ...}
    • \n
  • Verify the deployment config matches this tutorial: engine="vllm" with vllm/vllm-openai:v0.22.1
    • \n
\n

Next Steps

\n\n" + "source": "#### Evaluation Best Practices\n\n**Manual Evaluation** (Recommended)\n- Test with real-world examples from your use case\n- Compare responses to base model and expected outputs\n- Verify the model exhibits desired behavior changes\n- Check edge cases and error handling\n\n**What to look for:**\n- ✅ Model follows your desired output format\n- ✅ Applies domain knowledge correctly\n- ✅ Maintains general language capabilities\n- ✅ Avoids unwanted behaviors or biases\n- ❌ Doesn't hallucinate facts not in training data\n- ❌ Doesn't produce repetitive or nonsensical outputs\n\n---\n\n## Hyperparameters\n\nFor detailed information on all available hyperparameters, recommended values, and tuning guidance, refer to the [Hyperparameter Reference](../manage-customization-jobs/hyperparameters.md).\n\n---\n\n\n## Troubleshooting\n\n**Job fails during model download:**\n- Verify authentication secrets are configured (refer to [Managing Secrets](../../get-started/concepts/manage-secrets.md))\n- For gated Hugging Face models (Llama, Gemma), accept the license on the model page (for example, [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct))\n- Confirm the model fileset uses `token_secret=hf_secret.name` for gated models\n- Check `AutomodelJobInput` references use the `workspace/name` format: `model=f\"default/{MODEL_NAME}\"` and `dataset={\"training\": f\"default/{DATASET_NAME}\"}` (for example, `default/llama-3-2-1b-base`, `default/sft-dataset`)\n- Verify the model entity points at the fileset: `fileset=f\"default/{MODEL_NAME}\"`\n- Check job status: `client.jobs.get_status(name=job.job.name, workspace=\"default\")`\n\n**Job fails with OOM (Out of Memory) error:**\n1. **First try:** Reduce `global_batch_size` from 64 to 32 or 16 in `batch={...}`\n2. **Still OOM:** Keep `micro_batch_size` at 1 (already the minimum in this tutorial)\n3. **Still OOM:** Reduce `max_seq_length` from 2048 to 1024 or 512 in `training={...}`\n4. **Last resort:** Increase `num_gpus_per_node` and `tensor_parallel_size` in `parallelism={...}`\n\n**Loss curves not decreasing (underfitting):**\n- Increase training duration: raise `epochs` from 2 to 3-5 in `schedule={...}`\n- Adjust learning rate: try `1e-4` or `1e-5` instead of the default `5e-5` in `optimizer={...}`\n- Check data quality: Verify formatting, remove duplicates, ensure diversity\n\n**Training loss decreases but validation loss increases (overfitting):**\n- Reduce `epochs` from 2 to 1 in `schedule={...}`\n- Lower `learning_rate` from `5e-5` to `2e-5` or `1e-5` in `optimizer={...}`\n- Increase dataset size and diversity\n- Verify train/validation split has no data leakage\n\n**Model output quality is poor despite good training metrics:**\n- Training metrics optimize for loss, not your actual task—evaluate on real use cases\n- Review data quality, format, and diversity—metrics can be misleading with poor data\n- Try a different base model size or architecture\n- Adjust `learning_rate` and `global_batch_size`\n- Compare to baseline: Test base model to ensure fine-tuning improved performance\n\n**Deployment fails:**\n- Verify output model exists: `client.models.retrieve(name=OUTPUT_NAME, workspace=\"default\")`\n- Check deployment logs: `client.inference.deployments.get_logs(name=deployment.name, workspace=\"default\")`\n- Ensure sufficient GPU resources for `executor_config={\"gpu\": 1, ...}`\n- Verify the deployment config matches this tutorial: `engine=\"vllm\"` with `vllm/vllm-openai:v0.22.1`\n\n\n## Next Steps\n\n- [Monitor training metrics](fine-tune-metrics) in detail\n- [Evaluate your fine-tuned model](../../evaluator/index) using the Evaluator service\n- Learn about [LoRA customization](./lora-customization-job) for resource-efficient fine-tuning", + "source_html": "

Evaluation Best Practices

\n

Manual Evaluation (Recommended)

\n
    \n
  • Test with real-world examples from your use case
    • \n
  • Compare responses to base model and expected outputs
    • \n
  • Verify the model exhibits desired behavior changes
    • \n
  • Check edge cases and error handling
    • \n
\n

What to look for:

\n
    \n
  • ✅ Model follows your desired output format
    • \n
  • ✅ Applies domain knowledge correctly
    • \n
  • ✅ Maintains general language capabilities
    • \n
  • ✅ Avoids unwanted behaviors or biases
    • \n
  • ❌ Doesn't hallucinate facts not in training data
    • \n
  • ❌ Doesn't produce repetitive or nonsensical outputs
    • \n
\n
\n

Hyperparameters

\n

For detailed information on all available hyperparameters, recommended values, and tuning guidance, refer to the Hyperparameter Reference.

\n
\n

Troubleshooting

\n

Job fails during model download:

\n
    \n
  • Verify authentication secrets are configured (refer to Managing Secrets)
    • \n
  • For gated Hugging Face models (Llama, Gemma), accept the license on the model page (for example, meta-llama/Llama-3.2-1B-Instruct)
    • \n
  • Confirm the model fileset uses token_secret=hf_secret.name for gated models
    • \n
  • Check AutomodelJobInput references use the workspace/name format: model=f"default/{MODEL_NAME}" and dataset={"training": f"default/{DATASET_NAME}"} (for example, default/llama-3-2-1b-base, default/sft-dataset)
    • \n
  • Verify the model entity points at the fileset: fileset=f"default/{MODEL_NAME}"
    • \n
  • Check job status: client.jobs.get_status(name=job.job.name, workspace="default")
    • \n
\n

Job fails with OOM (Out of Memory) error:

\n
    \n
  1. First try: Reduce global_batch_size from 64 to 32 or 16 in batch={...}
    2. \n
  2. Still OOM: Keep micro_batch_size at 1 (already the minimum in this tutorial)
    3. \n
  3. Still OOM: Reduce max_seq_length from 2048 to 1024 or 512 in training={...}
    4. \n
  4. Last resort: Increase num_gpus_per_node and tensor_parallel_size in parallelism={...}
    5. \n
\n

Loss curves not decreasing (underfitting):

\n
    \n
  • Increase training duration: raise epochs from 2 to 3-5 in schedule={...}
    • \n
  • Adjust learning rate: try 1e-4 or 1e-5 instead of the default 5e-5 in optimizer={...}
    • \n
  • Check data quality: Verify formatting, remove duplicates, ensure diversity
    • \n
\n

Training loss decreases but validation loss increases (overfitting):

\n
    \n
  • Reduce epochs from 2 to 1 in schedule={...}
    • \n
  • Lower learning_rate from 5e-5 to 2e-5 or 1e-5 in optimizer={...}
    • \n
  • Increase dataset size and diversity
    • \n
  • Verify train/validation split has no data leakage
    • \n
\n

Model output quality is poor despite good training metrics:

\n
    \n
  • Training metrics optimize for loss, not your actual task—evaluate on real use cases
    • \n
  • Review data quality, format, and diversity—metrics can be misleading with poor data
    • \n
  • Try a different base model size or architecture
    • \n
  • Adjust learning_rate and global_batch_size
    • \n
  • Compare to baseline: Test base model to ensure fine-tuning improved performance
    • \n
\n

Deployment fails:

\n
    \n
  • Verify output model exists: client.models.retrieve(name=OUTPUT_NAME, workspace="default")
    • \n
  • Check deployment logs: client.inference.deployments.get_logs(name=deployment.name, workspace="default")
    • \n
  • Ensure sufficient GPU resources for executor_config={"gpu": 1, ...}
    • \n
  • Verify the deployment config matches this tutorial: engine="vllm" with vllm/vllm-openai:v0.22.1
    • \n
\n

Next Steps

\n\n" } ] } diff --git a/docs/fern/components/notebooks/sft-customization-job.ts b/docs/fern/components/notebooks/sft-customization-job.ts index 60a7f22ff8..053dc0c3ef 100644 --- a/docs/fern/components/notebooks/sft-customization-job.ts +++ b/docs/fern/components/notebooks/sft-customization-job.ts @@ -7,13 +7,13 @@ export default { cells: [ { "type": "markdown", - "source": "\n\n\n# Full SFT Customization\n\nLearn how to fine-tune all model weights using supervised fine-tuning (SFT) to customize LLM behavior for your specific tasks.\n\n## About\n\nSupervised Fine-Tuning (SFT) customizes model behavior, injects new knowledge, and optimizes performance for specific domains and tasks. Full SFT modifies **all model weights** during training, providing maximum customization flexibility.\n\n**What you can achieve with SFT:**\n\n- 🎯 **Specialize for domains:** Fine-tune models on legal texts, medical records, or financial data\n- 💡 **Inject knowledge:** Add new information not present in the base model\n- 📈 **Improve accuracy:** Optimize for specific tasks like sentiment analysis, summarization, or code generation\n\n### SFT vs LoRA: Understanding the Trade-offs\n\n**Full SFT** trains all model parameters (for example, all 70 billion weights in Llama 70B):\n\n- ✅ Maximum model adaptation and knowledge injection\n- ✅ Can fundamentally change model behavior\n- ✅ Best for significant domain shifts or specialized tasks\n- ❌ Requires substantial GPU resources (4-8x more than LoRA)\n- ❌ Produces full model weights (~140GB for Llama 70B)\n- ❌ Longer training time\n\n**LoRA** trains only ~1% of weights by adding thin matrices to existing weights:\n\n- ✅ 75-95% less memory required\n- ✅ Faster training (2-4x speedup)\n- ✅ Produces small adapter files (~100-500MB)\n- ✅ Multiple adapters can share one base model\n- ❌ Limited adaptation capability compared to full fine-tuning\n\n**When to choose Full SFT:**\n\n- Training small models (1B-8B) where resource cost is manageable\n- Need fundamental behavior changes (for example, medical diagnosis, legal reasoning)\n- Injecting substantial new knowledge not in the base model\n\n**When to choose LoRA:** Refer to the [LoRA tutorial](./lora-customization-job) for most use cases, especially with large models (70B+) or limited GPU resources.", - "source_html": "\n\n

Full SFT Customization

\n

Learn how to fine-tune all model weights using supervised fine-tuning (SFT) to customize LLM behavior for your specific tasks.

\n

About

\n

Supervised Fine-Tuning (SFT) customizes model behavior, injects new knowledge, and optimizes performance for specific domains and tasks. Full SFT modifies all model weights during training, providing maximum customization flexibility.

\n

What you can achieve with SFT:

\n
    \n
  • 🎯 Specialize for domains: Fine-tune models on legal texts, medical records, or financial data
    • \n
  • 💡 Inject knowledge: Add new information not present in the base model
    • \n
  • 📈 Improve accuracy: Optimize for specific tasks like sentiment analysis, summarization, or code generation
    • \n
\n

SFT vs LoRA: Understanding the Trade-offs

\n

Full SFT trains all model parameters (for example, all 70 billion weights in Llama 70B):

\n
    \n
  • ✅ Maximum model adaptation and knowledge injection
    • \n
  • ✅ Can fundamentally change model behavior
    • \n
  • ✅ Best for significant domain shifts or specialized tasks
    • \n
  • ❌ Requires substantial GPU resources (4-8x more than LoRA)
    • \n
  • ❌ Produces full model weights (~140GB for Llama 70B)
    • \n
  • ❌ Longer training time
    • \n
\n

LoRA trains only ~1% of weights by adding thin matrices to existing weights:

\n
    \n
  • ✅ 75-95% less memory required
    • \n
  • ✅ Faster training (2-4x speedup)
    • \n
  • ✅ Produces small adapter files (~100-500MB)
    • \n
  • ✅ Multiple adapters can share one base model
    • \n
  • ❌ Limited adaptation capability compared to full fine-tuning
    • \n
\n

When to choose Full SFT:

\n
    \n
  • Training small models (1B-8B) where resource cost is manageable
    • \n
  • Need fundamental behavior changes (for example, medical diagnosis, legal reasoning)
    • \n
  • Injecting substantial new knowledge not in the base model
    • \n
\n

When to choose LoRA: Refer to the LoRA tutorial for most use cases, especially with large models (70B+) or limited GPU resources.

\n" + "source": "\n\n\n# Full SFT Customization\n\nLearn how to fine-tune all model weights using supervised fine-tuning (SFT) to customize LLM behavior for your specific tasks.\n\n## About\n\nSupervised Fine-Tuning (SFT) customizes model behavior, injects new knowledge, and optimizes performance for specific domains and tasks. Full SFT modifies **all model weights** during training, providing maximum customization flexibility.\n\n**What you can achieve with SFT:**\n\n- 🎯 **Specialize for domains:** Fine-tune models on legal texts, medical records, or financial data\n- 💡 **Inject knowledge:** Add new information not present in the base model\n- 📈 **Improve accuracy:** Optimize for specific tasks like sentiment analysis, summarization, or code generation\n\n### SFT vs LoRA: Understanding the Trade-offs\n\n**Full SFT** trains all model parameters (for example, all 70 billion weights in Llama 70B):\n\n- ✅ Maximum model adaptation and knowledge injection\n- ✅ Can fundamentally change model behavior\n- ✅ Best for significant domain shifts or specialized tasks\n- ❌ Requires substantial GPU resources (4-8x more than LoRA)\n- ❌ Produces a full BF16 checkpoint (~140 GB for Llama 70B); peak job disk usage can reach approximately 3× the downloaded base checkpoint size\n- ❌ Longer training time\n\n**LoRA** trains only ~1% of weights by adding thin matrices to existing weights:\n\n- ✅ 75-95% less memory required\n- ✅ Faster training (2-4x speedup)\n- ✅ Produces small adapter files (~100-500MB)\n- ✅ Multiple adapters can share one base model\n- ❌ Limited adaptation capability compared to full fine-tuning\n\n**When to choose Full SFT:**\n\n- Training small models (1B-8B) where resource cost is manageable\n- Need fundamental behavior changes (for example, medical diagnosis, legal reasoning)\n- Injecting substantial new knowledge not in the base model\n\n**When to choose LoRA:** Refer to the [LoRA tutorial](./lora-customization-job) for most use cases, especially with large models (70B+) or limited GPU resources.", + "source_html": "\n\n

Full SFT Customization

\n

Learn how to fine-tune all model weights using supervised fine-tuning (SFT) to customize LLM behavior for your specific tasks.

\n

About

\n

Supervised Fine-Tuning (SFT) customizes model behavior, injects new knowledge, and optimizes performance for specific domains and tasks. Full SFT modifies all model weights during training, providing maximum customization flexibility.

\n

What you can achieve with SFT:

\n
    \n
  • 🎯 Specialize for domains: Fine-tune models on legal texts, medical records, or financial data
    • \n
  • 💡 Inject knowledge: Add new information not present in the base model
    • \n
  • 📈 Improve accuracy: Optimize for specific tasks like sentiment analysis, summarization, or code generation
    • \n
\n

SFT vs LoRA: Understanding the Trade-offs

\n

Full SFT trains all model parameters (for example, all 70 billion weights in Llama 70B):

\n
    \n
  • ✅ Maximum model adaptation and knowledge injection
    • \n
  • ✅ Can fundamentally change model behavior
    • \n
  • ✅ Best for significant domain shifts or specialized tasks
    • \n
  • ❌ Requires substantial GPU resources (4-8x more than LoRA)
    • \n
  • ❌ Produces a full BF16 checkpoint (~140 GB for Llama 70B); peak job disk usage can reach approximately 3× the downloaded base checkpoint size
    • \n
  • ❌ Longer training time
    • \n
\n

LoRA trains only ~1% of weights by adding thin matrices to existing weights:

\n
    \n
  • ✅ 75-95% less memory required
    • \n
  • ✅ Faster training (2-4x speedup)
    • \n
  • ✅ Produces small adapter files (~100-500MB)
    • \n
  • ✅ Multiple adapters can share one base model
    • \n
  • ❌ Limited adaptation capability compared to full fine-tuning
    • \n
\n

When to choose Full SFT:

\n
    \n
  • Training small models (1B-8B) where resource cost is manageable
    • \n
  • Need fundamental behavior changes (for example, medical diagnosis, legal reasoning)
    • \n
  • Injecting substantial new knowledge not in the base model
    • \n
\n

When to choose LoRA: Refer to the LoRA tutorial for most use cases, especially with large models (70B+) or limited GPU resources.

\n" }, { "type": "markdown", - "source": "## Prerequisites\n\nBefore starting this tutorial, ensure you have:\n\n1. **Completed the [Quickstart](../../get-started/quickstart.md)** to install and deploy NeMo Platform locally\n2. **Installed the Python SDK** (PyPI wrapper: `pip install \"nemo-platform[all]\"`; source checkout: run `make bootstrap` from the repository root)", - "source_html": "

Prerequisites

\n

Before starting this tutorial, ensure you have:

\n
    \n
  1. Completed the Quickstart to install and deploy NeMo Platform locally
    2. \n
  2. Installed the Python SDK (PyPI wrapper: pip install "nemo-platform[all]"; source checkout: run make bootstrap from the repository root)
    3. \n
\n" + "source": "## Prerequisites\n\nBefore starting this tutorial, ensure you have:\n\n1. **Completed the [Quickstart](../../get-started/quickstart.md)** to install and deploy NeMo Platform locally\n2. **Installed the Python SDK** (PyPI wrapper: `pip install \"nemo-platform[all]\"`; source checkout: run `make bootstrap` from the repository root)\n3. **At least one GPU with CUDA 13+**", + "source_html": "

Prerequisites

\n

Before starting this tutorial, ensure you have:

\n
    \n
  1. Completed the Quickstart to install and deploy NeMo Platform locally
    2. \n
  2. Installed the Python SDK (PyPI wrapper: pip install "nemo-platform[all]"; source checkout: run make bootstrap from the repository root)
    3. \n
  3. At least one GPU with CUDA 13+
    4. \n
\n" }, { "type": "markdown", @@ -90,30 +90,30 @@ export default { cells: [ }, { "type": "markdown", - "source": "### 4. Secrets Setup\n\nIf you plan to use NGC or HuggingFace models, you will need to configure authentication:\n\n- **NGC models** (`ngc://` URIs): Requires NGC API key\n- **HuggingFace models** (`hf://` URIs): Requires HF token for gated/private models\n\n\nConfigure these as secrets in your platform. Refer to [Managing Secrets](../../get-started/concepts/manage-secrets.md) for detailed instructions.\n\nGet your credentials to access base models:\n- [NGC API Key](https://ngc.nvidia.com/) (Setup → Generate API Key)\n- [HuggingFace Token](https://huggingface.co/settings/tokens) (Create token with Read access)\n\n\n---\n\n#### Quick Setup Example\n\nIn this tutorial we are going to work with [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main) model from HuggingFace. Ensure that you have sufficient permissions to download the model. If you cannot access the files on the [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main) Hugging Face page, request access\n\n**HuggingFace Authentication:**\n- For gated models (Llama, Gemma), you must provide a HuggingFace token via the `token_secret` parameter\n- Get your token from [HuggingFace Settings](https://huggingface.co/settings/tokens) (requires Read access)\n- Accept the model's terms on the HuggingFace model page before using it. Example: [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main)\n- For public models, you can omit the `token_secret` parameter when creating a fileset for model in the next step", - "source_html": "

4. Secrets Setup

\n

If you plan to use NGC or HuggingFace models, you will need to configure authentication:

\n
    \n
  • NGC models (ngc:// URIs): Requires NGC API key
    • \n
  • HuggingFace models (hf:// URIs): Requires HF token for gated/private models
    • \n
\n

Configure these as secrets in your platform. Refer to Managing Secrets for detailed instructions.

\n

Get your credentials to access base models:

\n\n
\n

Quick Setup Example

\n

In this tutorial we are going to work with meta-llama/Llama-3.2-1B-Instruct model from HuggingFace. Ensure that you have sufficient permissions to download the model. If you cannot access the files on the meta-llama/Llama-3.2-1B-Instruct Hugging Face page, request access

\n

HuggingFace Authentication:

\n
    \n
  • For gated models (Llama, Gemma), you must provide a HuggingFace token via the token_secret parameter
    • \n
  • Get your token from HuggingFace Settings (requires Read access)
    • \n
  • Accept the model's terms on the HuggingFace model page before using it. Example: meta-llama/Llama-3.2-1B-Instruct
    • \n
  • For public models, you can omit the token_secret parameter when creating a fileset for model in the next step
    • \n
\n" + "source": "### 4. Secrets Setup\n\nIf you plan to use NGC or Hugging Face models, you will need to configure authentication:\n\n- **NGC models** (`ngc://` URIs): Requires NGC API key\n- **Hugging Face models** (`hf://` URIs): Requires HF token for gated/private models\n\n\nConfigure these as secrets in your platform. Refer to [Managing Secrets](../../get-started/concepts/manage-secrets.md) for detailed instructions.\n\nGet your credentials to access base models:\n- [NGC API Key](https://ngc.nvidia.com/) (Setup → Generate API Key)\n- [Hugging Face Token](https://huggingface.co/settings/tokens) (Create token with Read access)\n\n\n---\n\n#### Quick Setup Example\n\nIn this tutorial we are going to work with the [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main) model from Hugging Face. Ensure that you have sufficient permissions to download the model. If you cannot access the files on the [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main) Hugging Face page, request access.\n\n**Hugging Face Authentication:**\n- For gated models (Llama, Gemma), you must provide a Hugging Face token via the `token_secret` parameter\n- Get your token from [Hugging Face Settings](https://huggingface.co/settings/tokens) (requires Read access)\n- Accept the model's terms on the Hugging Face model page before using it. Example: [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main)\n- For public models, you can omit the `token_secret` parameter when creating a fileset for model in the next step", + "source_html": "

4. Secrets Setup

\n

If you plan to use NGC or Hugging Face models, you will need to configure authentication:

\n
    \n
  • NGC models (ngc:// URIs): Requires NGC API key
    • \n
  • Hugging Face models (hf:// URIs): Requires HF token for gated/private models
    • \n
\n

Configure these as secrets in your platform. Refer to Managing Secrets for detailed instructions.

\n

Get your credentials to access base models:

\n\n
\n

Quick Setup Example

\n

In this tutorial we are going to work with the meta-llama/Llama-3.2-1B-Instruct model from Hugging Face. Ensure that you have sufficient permissions to download the model. If you cannot access the files on the meta-llama/Llama-3.2-1B-Instruct Hugging Face page, request access.

\n

Hugging Face Authentication:

\n
    \n
  • For gated models (Llama, Gemma), you must provide a Hugging Face token via the token_secret parameter
    • \n
  • Get your token from Hugging Face Settings (requires Read access)
    • \n
  • Accept the model's terms on the Hugging Face model page before using it. Example: meta-llama/Llama-3.2-1B-Instruct
    • \n
  • For public models, you can omit the token_secret parameter when creating a fileset for model in the next step
    • \n
\n" }, { "type": "code", - "source": "# Export the HF_TOKEN and NGC_API_KEY environment variables if they are not already set\nHF_TOKEN = os.getenv(\"HF_TOKEN\")\nNGC_API_KEY = os.getenv(\"NGC_API_KEY\")\n\n\ndef create_or_get_secret(name: str, value: str | None, label: str):\n if not value:\n raise ValueError(f\"{label} is not set\")\n try:\n secret = client.secrets.create(\n name=name,\n workspace=\"default\",\n value=value,\n )\n print(f\"Created secret: {name}\")\n return secret\n except ConflictError:\n print(f\"Secret '{name}' already exists, continuing...\")\n return client.secrets.retrieve(name=name, workspace=\"default\")\n\n\n# Create HuggingFace token secret\nhf_secret = create_or_get_secret(\"hf-token\", HF_TOKEN, \"HF_TOKEN\")\nprint(\"HF_TOKEN secret:\")\nprint(hf_secret.model_dump_json(indent=2))\n\n# Create NGC API key secret\n# Uncomment the line below if you have NGC API Key and want to finetune NGC models\n# ngc_api_key = create_or_get_secret(\"ngc-api-key\", NGC_API_KEY, \"NGC_API_KEY\")", + "source": "# Export the HF_TOKEN and NGC_API_KEY environment variables if they are not already set\nHF_TOKEN = os.getenv(\"HF_TOKEN\")\nNGC_API_KEY = os.getenv(\"NGC_API_KEY\")\n\n\ndef create_or_get_secret(name: str, value: str | None, label: str):\n if not value:\n raise ValueError(f\"{label} is not set\")\n try:\n secret = client.secrets.create(\n name=name,\n workspace=\"default\",\n value=value,\n )\n print(f\"Created secret: {name}\")\n return secret\n except ConflictError:\n print(f\"Secret '{name}' already exists, continuing...\")\n return client.secrets.retrieve(name=name, workspace=\"default\")\n\n\n# Create Hugging Face token secret\nhf_secret = create_or_get_secret(\"hf-token\", HF_TOKEN, \"HF_TOKEN\")\nprint(\"HF_TOKEN secret:\")\nprint(hf_secret.model_dump_json(indent=2))\n\n# Create NGC API key secret\n# Uncomment the line below if you have NGC API Key and want to finetune NGC models\n# ngc_api_key = create_or_get_secret(\"ngc-api-key\", NGC_API_KEY, \"NGC_API_KEY\")", "language": "python", - "source_html": "# Export the HF_TOKEN and NGC_API_KEY environment variables if they are not already set\nHF_TOKEN = os.getenv("HF_TOKEN")\nNGC_API_KEY = os.getenv("NGC_API_KEY")\n\n\ndef create_or_get_secret(name: str, value: str | None, label: str):\n if not value:\n raise ValueError(f"{label} is not set")\n try:\n secret = client.secrets.create(\n name=name,\n workspace="default",\n value=value,\n )\n print(f"Created secret: {name}")\n return secret\n except ConflictError:\n print(f"Secret '{name}' already exists, continuing...")\n return client.secrets.retrieve(name=name, workspace="default")\n\n\n# Create HuggingFace token secret\nhf_secret = create_or_get_secret("hf-token", HF_TOKEN, "HF_TOKEN")\nprint("HF_TOKEN secret:")\nprint(hf_secret.model_dump_json(indent=2))\n\n# Create NGC API key secret\n# Uncomment the line below if you have NGC API Key and want to finetune NGC models\n# ngc_api_key = create_or_get_secret("ngc-api-key", NGC_API_KEY, "NGC_API_KEY")\n" + "source_html": "# Export the HF_TOKEN and NGC_API_KEY environment variables if they are not already set\nHF_TOKEN = os.getenv("HF_TOKEN")\nNGC_API_KEY = os.getenv("NGC_API_KEY")\n\n\ndef create_or_get_secret(name: str, value: str | None, label: str):\n if not value:\n raise ValueError(f"{label} is not set")\n try:\n secret = client.secrets.create(\n name=name,\n workspace="default",\n value=value,\n )\n print(f"Created secret: {name}")\n return secret\n except ConflictError:\n print(f"Secret '{name}' already exists, continuing...")\n return client.secrets.retrieve(name=name, workspace="default")\n\n\n# Create Hugging Face token secret\nhf_secret = create_or_get_secret("hf-token", HF_TOKEN, "HF_TOKEN")\nprint("HF_TOKEN secret:")\nprint(hf_secret.model_dump_json(indent=2))\n\n# Create NGC API key secret\n# Uncomment the line below if you have NGC API Key and want to finetune NGC models\n# ngc_api_key = create_or_get_secret("ngc-api-key", NGC_API_KEY, "NGC_API_KEY")\n" }, { "type": "markdown", - "source": "### 5. Create Base Model FileSet and Model Entity\n\nCreate a fileset pointing to [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main) model in HuggingFace that we will train with SFT. Then create a Model Entity that references this fileset. Model downloading will take place at training time.\n\nNote: for public models, you can omit the `token_secret` parameter when creating a model fileset.", - "source_html": "

5. Create Base Model FileSet and Model Entity

\n

Create a fileset pointing to meta-llama/Llama-3.2-1B-Instruct model in HuggingFace that we will train with SFT. Then create a Model Entity that references this fileset. Model downloading will take place at training time.

\n

Note: for public models, you can omit the token_secret parameter when creating a model fileset.

\n" + "source": "### 5. Create Base Model FileSet and Model Entity\n\nCreate a fileset pointing to the [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/tree/main) model in Hugging Face that we will train with SFT. Then create a Model Entity that references this fileset. Model downloading will take place at training time.\n\nNote: for public models, you can omit the `token_secret` parameter when creating a model fileset.", + "source_html": "

5. Create Base Model FileSet and Model Entity

\n

Create a fileset pointing to the meta-llama/Llama-3.2-1B-Instruct model in Hugging Face that we will train with SFT. Then create a Model Entity that references this fileset. Model downloading will take place at training time.

\n

Note: for public models, you can omit the token_secret parameter when creating a model fileset.

\n" }, { "type": "code", - "source": "import time\n\n# Create a fileset pointing to the desired HuggingFace model\nfrom nemo_platform.types.files import HuggingfaceStorageConfigParam\n\nHF_REPO_ID = \"meta-llama/Llama-3.2-1B-Instruct\"\nMODEL_NAME = \"llama-3-2-1b-base\"\n\n# Ensure you have a HuggingFace token secret created\ntry:\n base_model_fs = client.files.filesets.create(\n workspace=\"default\",\n name=MODEL_NAME,\n description=\"Llama 3.2 1B base model from HuggingFace\",\n storage=HuggingfaceStorageConfigParam(\n type=\"huggingface\",\n # repo_id is the full model name from Hugging Face\n repo_id=HF_REPO_ID,\n repo_type=\"model\",\n # we use the secret created in the previous step\n token_secret=hf_secret.name\n )\n )\n print(f\"Created base model fileset: {MODEL_NAME}\")\nexcept ConflictError:\n print(f\"Base model fileset already exists. Skipping creation.\")\n base_model_fs = client.files.filesets.retrieve(\n workspace=\"default\",\n name=MODEL_NAME,\n )\n\n# Create the Model Entity representation.\ntry:\n base_model = client.models.create(\n workspace=\"default\",\n name=MODEL_NAME,\n fileset=f\"default/{MODEL_NAME}\",\n )\n print(f\"Created Model Entity: {MODEL_NAME}\")\nexcept ConflictError:\n print(f\"Base model already exists. Updating fileset if different.\")\n base_model = client.models.update(\n workspace=\"default\",\n name=MODEL_NAME,\n fileset=f\"default/{MODEL_NAME}\",\n )\n\nprint(f\"\\nBase model fileset: fileset://default/{base_model.name}\")\nprint(\"Base model fileset files list:\")\nprint(json.dumps([f.model_dump() for f in client.files.list(fileset=MODEL_NAME, workspace=\"default\").data], indent=2))\n\n# Wait for ModelSpec to be populated from the checkpoint\nprint(\"\\nWaiting for ModelSpec to be populated...\")\nSPEC_TIMEOUT_SECONDS = 120\nspec_start = time.time()\nwhile not base_model.spec:\n if time.time() - spec_start > SPEC_TIMEOUT_SECONDS:\n raise TimeoutError(f\"ModelSpec not populated within {SPEC_TIMEOUT_SECONDS} seconds\")\n time.sleep(2)\n base_model = client.models.retrieve(\n workspace=\"default\",\n name=MODEL_NAME,\n )\n\nprint(f\"ModelSpec populated: {base_model.spec}\")", + "source": "import time\n\n# Create a fileset pointing to the desired Hugging Face model\nfrom nemo_platform.types.files import HuggingfaceStorageConfigParam\n\nHF_REPO_ID = \"meta-llama/Llama-3.2-1B-Instruct\"\nMODEL_NAME = \"llama-3-2-1b-base\"\n\n# Ensure you have a Hugging Face token secret created\ntry:\n base_model_fs = client.files.filesets.create(\n workspace=\"default\",\n name=MODEL_NAME,\n description=\"Llama 3.2 1B base model from Hugging Face\",\n storage=HuggingfaceStorageConfigParam(\n type=\"huggingface\",\n # repo_id is the full model name from Hugging Face\n repo_id=HF_REPO_ID,\n repo_type=\"model\",\n # we use the secret created in the previous step\n token_secret=hf_secret.name\n )\n )\n print(f\"Created base model fileset: {MODEL_NAME}\")\nexcept ConflictError:\n print(f\"Base model fileset already exists. Skipping creation.\")\n base_model_fs = client.files.filesets.retrieve(\n workspace=\"default\",\n name=MODEL_NAME,\n )\n\n# Create the Model Entity representation.\ntry:\n base_model = client.models.create(\n workspace=\"default\",\n name=MODEL_NAME,\n fileset=f\"default/{MODEL_NAME}\",\n )\n print(f\"Created Model Entity: {MODEL_NAME}\")\nexcept ConflictError:\n print(f\"Base model already exists. Updating fileset if different.\")\n base_model = client.models.update(\n workspace=\"default\",\n name=MODEL_NAME,\n fileset=f\"default/{MODEL_NAME}\",\n )\n\nprint(f\"\\nBase model fileset: fileset://default/{base_model.name}\")\nprint(\"Base model fileset files list:\")\nprint(json.dumps([f.model_dump() for f in client.files.list(fileset=MODEL_NAME, workspace=\"default\").data], indent=2))\n\n# Wait for ModelSpec to be populated from the checkpoint\nprint(\"\\nWaiting for ModelSpec to be populated...\")\nSPEC_TIMEOUT_SECONDS = 120\nspec_start = time.time()\nwhile not base_model.spec:\n if time.time() - spec_start > SPEC_TIMEOUT_SECONDS:\n raise TimeoutError(f\"ModelSpec not populated within {SPEC_TIMEOUT_SECONDS} seconds\")\n time.sleep(2)\n base_model = client.models.retrieve(\n workspace=\"default\",\n name=MODEL_NAME,\n )\n\nprint(f\"ModelSpec populated: {base_model.spec}\")", "language": "python", - "source_html": "import time\n\n# Create a fileset pointing to the desired HuggingFace model\nfrom nemo_platform.types.files import HuggingfaceStorageConfigParam\n\nHF_REPO_ID = "meta-llama/Llama-3.2-1B-Instruct"\nMODEL_NAME = "llama-3-2-1b-base"\n\n# Ensure you have a HuggingFace token secret created\ntry:\n base_model_fs = client.files.filesets.create(\n workspace="default",\n name=MODEL_NAME,\n description="Llama 3.2 1B base model from HuggingFace",\n storage=HuggingfaceStorageConfigParam(\n type="huggingface",\n # repo_id is the full model name from Hugging Face\n repo_id=HF_REPO_ID,\n repo_type="model",\n # we use the secret created in the previous step\n token_secret=hf_secret.name\n )\n )\n print(f"Created base model fileset: {MODEL_NAME}")\nexcept ConflictError:\n print(f"Base model fileset already exists. Skipping creation.")\n base_model_fs = client.files.filesets.retrieve(\n workspace="default",\n name=MODEL_NAME,\n )\n\n# Create the Model Entity representation.\ntry:\n base_model = client.models.create(\n workspace="default",\n name=MODEL_NAME,\n fileset=f"default/{MODEL_NAME}",\n )\n print(f"Created Model Entity: {MODEL_NAME}")\nexcept ConflictError:\n print(f"Base model already exists. Updating fileset if different.")\n base_model = client.models.update(\n workspace="default",\n name=MODEL_NAME,\n fileset=f"default/{MODEL_NAME}",\n )\n\nprint(f"\\nBase model fileset: fileset://default/{base_model.name}")\nprint("Base model fileset files list:")\nprint(json.dumps([f.model_dump() for f in client.files.list(fileset=MODEL_NAME, workspace="default").data], indent=2))\n\n# Wait for ModelSpec to be populated from the checkpoint\nprint("\\nWaiting for ModelSpec to be populated...")\nSPEC_TIMEOUT_SECONDS = 120\nspec_start = time.time()\nwhile not base_model.spec:\n if time.time() - spec_start > SPEC_TIMEOUT_SECONDS:\n raise TimeoutError(f"ModelSpec not populated within {SPEC_TIMEOUT_SECONDS} seconds")\n time.sleep(2)\n base_model = client.models.retrieve(\n workspace="default",\n name=MODEL_NAME,\n )\n\nprint(f"ModelSpec populated: {base_model.spec}")\n" + "source_html": "import time\n\n# Create a fileset pointing to the desired Hugging Face model\nfrom nemo_platform.types.files import HuggingfaceStorageConfigParam\n\nHF_REPO_ID = "meta-llama/Llama-3.2-1B-Instruct"\nMODEL_NAME = "llama-3-2-1b-base"\n\n# Ensure you have a Hugging Face token secret created\ntry:\n base_model_fs = client.files.filesets.create(\n workspace="default",\n name=MODEL_NAME,\n description="Llama 3.2 1B base model from Hugging Face",\n storage=HuggingfaceStorageConfigParam(\n type="huggingface",\n # repo_id is the full model name from Hugging Face\n repo_id=HF_REPO_ID,\n repo_type="model",\n # we use the secret created in the previous step\n token_secret=hf_secret.name\n )\n )\n print(f"Created base model fileset: {MODEL_NAME}")\nexcept ConflictError:\n print(f"Base model fileset already exists. Skipping creation.")\n base_model_fs = client.files.filesets.retrieve(\n workspace="default",\n name=MODEL_NAME,\n )\n\n# Create the Model Entity representation.\ntry:\n base_model = client.models.create(\n workspace="default",\n name=MODEL_NAME,\n fileset=f"default/{MODEL_NAME}",\n )\n print(f"Created Model Entity: {MODEL_NAME}")\nexcept ConflictError:\n print(f"Base model already exists. Updating fileset if different.")\n base_model = client.models.update(\n workspace="default",\n name=MODEL_NAME,\n fileset=f"default/{MODEL_NAME}",\n )\n\nprint(f"\\nBase model fileset: fileset://default/{base_model.name}")\nprint("Base model fileset files list:")\nprint(json.dumps([f.model_dump() for f in client.files.list(fileset=MODEL_NAME, workspace="default").data], indent=2))\n\n# Wait for ModelSpec to be populated from the checkpoint\nprint("\\nWaiting for ModelSpec to be populated...")\nSPEC_TIMEOUT_SECONDS = 120\nspec_start = time.time()\nwhile not base_model.spec:\n if time.time() - spec_start > SPEC_TIMEOUT_SECONDS:\n raise TimeoutError(f"ModelSpec not populated within {SPEC_TIMEOUT_SECONDS} seconds")\n time.sleep(2)\n base_model = client.models.retrieve(\n workspace="default",\n name=MODEL_NAME,\n )\n\nprint(f"ModelSpec populated: {base_model.spec}")\n" }, { "type": "markdown", - "source": "### 6. Create SFT Finetuning Job\nCreate a customization job to fine-tune all model weights using the **Automodel** backend and `AutomodelJobInput`.", - "source_html": "

6. Create SFT Finetuning Job

\n

Create a customization job to fine-tune all model weights using the Automodel backend and AutomodelJobInput.

\n" + "source": "### 6. Create SFT Fine-Tuning Job\nCreate a customization job to fine-tune all model weights using the **Automodel** backend and `AutomodelJobInput`.", + "source_html": "

6. Create SFT Fine-Tuning Job

\n

Create a customization job to fine-tune all model weights using the Automodel backend and AutomodelJobInput.

\n" }, { "type": "markdown", @@ -188,7 +188,7 @@ export default { cells: [ }, { "type": "markdown", - "source": "#### Evaluation Best Practices\n\n**Manual Evaluation** (Recommended)\n- Test with real-world examples from your use case\n- Compare responses to base model and expected outputs\n- Verify the model exhibits desired behavior changes\n- Check edge cases and error handling\n\n**What to look for:**\n- ✅ Model follows your desired output format\n- ✅ Applies domain knowledge correctly\n- ✅ Maintains general language capabilities\n- ✅ Avoids unwanted behaviors or biases\n- ❌ Doesn't hallucinate facts not in training data\n- ❌ Doesn't produce repetitive or nonsensical outputs\n\n---\n\n## Hyperparameters\n\nFor detailed information on all available hyperparameters, recommended values, and tuning guidance, refer to the [Hyperparameter Reference](../manage-customization-jobs/hyperparameters.md).\n\n---\n\n\n## Troubleshooting\n\n**Job fails during model download:**\n- Verify authentication secrets are configured (refer to [Managing Secrets](../../get-started/concepts/manage-secrets.md))\n- For gated HuggingFace models (Llama, Gemma), accept the license on the model page (for example, [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct))\n- Confirm the model fileset uses `token_secret=hf_secret.name` for gated models\n- Check `AutomodelJobInput` references use the `workspace/name` format: `model=f\"default/{MODEL_NAME}\"` and `dataset={\"training\": f\"default/{DATASET_NAME}\"}` (for example, `default/llama-3-2-1b-base`, `default/sft-dataset`)\n- Verify the model entity points at the fileset: `fileset=f\"default/{MODEL_NAME}\"`\n- Check job status: `client.jobs.get_status(name=job.job.name, workspace=\"default\")`\n\n**Job fails with OOM (Out of Memory) error:**\n1. **First try:** Reduce `global_batch_size` from 64 to 32 or 16 in `batch={...}`\n2. **Still OOM:** Keep `micro_batch_size` at 1 (already the minimum in this tutorial)\n3. **Still OOM:** Reduce `max_seq_length` from 2048 to 1024 or 512 in `training={...}`\n4. **Last resort:** Increase `num_gpus_per_node` and `tensor_parallel_size` in `parallelism={...}`\n\n**Loss curves not decreasing (underfitting):**\n- Increase training duration: raise `epochs` from 2 to 3-5 in `schedule={...}`\n- Adjust learning rate: try `1e-4` or `1e-5` instead of the default `5e-5` in `optimizer={...}`\n- Check data quality: Verify formatting, remove duplicates, ensure diversity\n\n**Training loss decreases but validation loss increases (overfitting):**\n- Reduce `epochs` from 2 to 1 in `schedule={...}`\n- Lower `learning_rate` from `5e-5` to `2e-5` or `1e-5` in `optimizer={...}`\n- Increase dataset size and diversity\n- Verify train/validation split has no data leakage\n\n**Model output quality is poor despite good training metrics:**\n- Training metrics optimize for loss, not your actual task—evaluate on real use cases\n- Review data quality, format, and diversity—metrics can be misleading with poor data\n- Try a different base model size or architecture\n- Adjust `learning_rate` and `global_batch_size`\n- Compare to baseline: Test base model to ensure fine-tuning improved performance\n\n**Deployment fails:**\n- Verify output model exists: `client.models.retrieve(name=OUTPUT_NAME, workspace=\"default\")`\n- Check deployment logs: `client.inference.deployments.get_logs(name=deployment.name, workspace=\"default\")`\n- Ensure sufficient GPU resources for `executor_config={\"gpu\": 1, ...}`\n- Verify the deployment config matches this tutorial: `engine=\"vllm\"` with `vllm/vllm-openai:v0.22.1`\n\n\n## Next Steps\n\n- [Monitor training metrics](fine-tune-metrics) in detail\n- [Evaluate your fine-tuned model](../../evaluator/index) using the Evaluator service\n- Learn about [LoRA customization](./lora-customization-job) for resource-efficient fine-tuning", - "source_html": "

Evaluation Best Practices

\n

Manual Evaluation (Recommended)

\n
    \n
  • Test with real-world examples from your use case
    • \n
  • Compare responses to base model and expected outputs
    • \n
  • Verify the model exhibits desired behavior changes
    • \n
  • Check edge cases and error handling
    • \n
\n

What to look for:

\n
    \n
  • ✅ Model follows your desired output format
    • \n
  • ✅ Applies domain knowledge correctly
    • \n
  • ✅ Maintains general language capabilities
    • \n
  • ✅ Avoids unwanted behaviors or biases
    • \n
  • ❌ Doesn't hallucinate facts not in training data
    • \n
  • ❌ Doesn't produce repetitive or nonsensical outputs
    • \n
\n
\n

Hyperparameters

\n

For detailed information on all available hyperparameters, recommended values, and tuning guidance, refer to the Hyperparameter Reference.

\n
\n

Troubleshooting

\n

Job fails during model download:

\n
    \n
  • Verify authentication secrets are configured (refer to Managing Secrets)
    • \n
  • For gated HuggingFace models (Llama, Gemma), accept the license on the model page (for example, meta-llama/Llama-3.2-1B-Instruct)
    • \n
  • Confirm the model fileset uses token_secret=hf_secret.name for gated models
    • \n
  • Check AutomodelJobInput references use the workspace/name format: model=f"default/{MODEL_NAME}" and dataset={"training": f"default/{DATASET_NAME}"} (for example, default/llama-3-2-1b-base, default/sft-dataset)
    • \n
  • Verify the model entity points at the fileset: fileset=f"default/{MODEL_NAME}"
    • \n
  • Check job status: client.jobs.get_status(name=job.job.name, workspace="default")
    • \n
\n

Job fails with OOM (Out of Memory) error:

\n
    \n
  1. First try: Reduce global_batch_size from 64 to 32 or 16 in batch={...}
    2. \n
  2. Still OOM: Keep micro_batch_size at 1 (already the minimum in this tutorial)
    3. \n
  3. Still OOM: Reduce max_seq_length from 2048 to 1024 or 512 in training={...}
    4. \n
  4. Last resort: Increase num_gpus_per_node and tensor_parallel_size in parallelism={...}
    5. \n
\n

Loss curves not decreasing (underfitting):

\n
    \n
  • Increase training duration: raise epochs from 2 to 3-5 in schedule={...}
    • \n
  • Adjust learning rate: try 1e-4 or 1e-5 instead of the default 5e-5 in optimizer={...}
    • \n
  • Check data quality: Verify formatting, remove duplicates, ensure diversity
    • \n
\n

Training loss decreases but validation loss increases (overfitting):

\n
    \n
  • Reduce epochs from 2 to 1 in schedule={...}
    • \n
  • Lower learning_rate from 5e-5 to 2e-5 or 1e-5 in optimizer={...}
    • \n
  • Increase dataset size and diversity
    • \n
  • Verify train/validation split has no data leakage
    • \n
\n

Model output quality is poor despite good training metrics:

\n
    \n
  • Training metrics optimize for loss, not your actual task—evaluate on real use cases
    • \n
  • Review data quality, format, and diversity—metrics can be misleading with poor data
    • \n
  • Try a different base model size or architecture
    • \n
  • Adjust learning_rate and global_batch_size
    • \n
  • Compare to baseline: Test base model to ensure fine-tuning improved performance
    • \n
\n

Deployment fails:

\n
    \n
  • Verify output model exists: client.models.retrieve(name=OUTPUT_NAME, workspace="default")
    • \n
  • Check deployment logs: client.inference.deployments.get_logs(name=deployment.name, workspace="default")
    • \n
  • Ensure sufficient GPU resources for executor_config={"gpu": 1, ...}
    • \n
  • Verify the deployment config matches this tutorial: engine="vllm" with vllm/vllm-openai:v0.22.1
    • \n
\n

Next Steps

\n\n" + "source": "#### Evaluation Best Practices\n\n**Manual Evaluation** (Recommended)\n- Test with real-world examples from your use case\n- Compare responses to base model and expected outputs\n- Verify the model exhibits desired behavior changes\n- Check edge cases and error handling\n\n**What to look for:**\n- ✅ Model follows your desired output format\n- ✅ Applies domain knowledge correctly\n- ✅ Maintains general language capabilities\n- ✅ Avoids unwanted behaviors or biases\n- ❌ Doesn't hallucinate facts not in training data\n- ❌ Doesn't produce repetitive or nonsensical outputs\n\n---\n\n## Hyperparameters\n\nFor detailed information on all available hyperparameters, recommended values, and tuning guidance, refer to the [Hyperparameter Reference](../manage-customization-jobs/hyperparameters.md).\n\n---\n\n\n## Troubleshooting\n\n**Job fails during model download:**\n- Verify authentication secrets are configured (refer to [Managing Secrets](../../get-started/concepts/manage-secrets.md))\n- For gated Hugging Face models (Llama, Gemma), accept the license on the model page (for example, [meta-llama/Llama-3.2-1B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct))\n- Confirm the model fileset uses `token_secret=hf_secret.name` for gated models\n- Check `AutomodelJobInput` references use the `workspace/name` format: `model=f\"default/{MODEL_NAME}\"` and `dataset={\"training\": f\"default/{DATASET_NAME}\"}` (for example, `default/llama-3-2-1b-base`, `default/sft-dataset`)\n- Verify the model entity points at the fileset: `fileset=f\"default/{MODEL_NAME}\"`\n- Check job status: `client.jobs.get_status(name=job.job.name, workspace=\"default\")`\n\n**Job fails with OOM (Out of Memory) error:**\n1. **First try:** Reduce `global_batch_size` from 64 to 32 or 16 in `batch={...}`\n2. **Still OOM:** Keep `micro_batch_size` at 1 (already the minimum in this tutorial)\n3. **Still OOM:** Reduce `max_seq_length` from 2048 to 1024 or 512 in `training={...}`\n4. **Last resort:** Increase `num_gpus_per_node` and `tensor_parallel_size` in `parallelism={...}`\n\n**Loss curves not decreasing (underfitting):**\n- Increase training duration: raise `epochs` from 2 to 3-5 in `schedule={...}`\n- Adjust learning rate: try `1e-4` or `1e-5` instead of the default `5e-5` in `optimizer={...}`\n- Check data quality: Verify formatting, remove duplicates, ensure diversity\n\n**Training loss decreases but validation loss increases (overfitting):**\n- Reduce `epochs` from 2 to 1 in `schedule={...}`\n- Lower `learning_rate` from `5e-5` to `2e-5` or `1e-5` in `optimizer={...}`\n- Increase dataset size and diversity\n- Verify train/validation split has no data leakage\n\n**Model output quality is poor despite good training metrics:**\n- Training metrics optimize for loss, not your actual task—evaluate on real use cases\n- Review data quality, format, and diversity—metrics can be misleading with poor data\n- Try a different base model size or architecture\n- Adjust `learning_rate` and `global_batch_size`\n- Compare to baseline: Test base model to ensure fine-tuning improved performance\n\n**Deployment fails:**\n- Verify output model exists: `client.models.retrieve(name=OUTPUT_NAME, workspace=\"default\")`\n- Check deployment logs: `client.inference.deployments.get_logs(name=deployment.name, workspace=\"default\")`\n- Ensure sufficient GPU resources for `executor_config={\"gpu\": 1, ...}`\n- Verify the deployment config matches this tutorial: `engine=\"vllm\"` with `vllm/vllm-openai:v0.22.1`\n\n\n## Next Steps\n\n- [Monitor training metrics](fine-tune-metrics) in detail\n- [Evaluate your fine-tuned model](../../evaluator/index) using the Evaluator service\n- Learn about [LoRA customization](./lora-customization-job) for resource-efficient fine-tuning", + "source_html": "

Evaluation Best Practices

\n

Manual Evaluation (Recommended)

\n
    \n
  • Test with real-world examples from your use case
    • \n
  • Compare responses to base model and expected outputs
    • \n
  • Verify the model exhibits desired behavior changes
    • \n
  • Check edge cases and error handling
    • \n
\n

What to look for:

\n
    \n
  • ✅ Model follows your desired output format
    • \n
  • ✅ Applies domain knowledge correctly
    • \n
  • ✅ Maintains general language capabilities
    • \n
  • ✅ Avoids unwanted behaviors or biases
    • \n
  • ❌ Doesn't hallucinate facts not in training data
    • \n
  • ❌ Doesn't produce repetitive or nonsensical outputs
    • \n
\n
\n

Hyperparameters

\n

For detailed information on all available hyperparameters, recommended values, and tuning guidance, refer to the Hyperparameter Reference.

\n
\n

Troubleshooting

\n

Job fails during model download:

\n
    \n
  • Verify authentication secrets are configured (refer to Managing Secrets)
    • \n
  • For gated Hugging Face models (Llama, Gemma), accept the license on the model page (for example, meta-llama/Llama-3.2-1B-Instruct)
    • \n
  • Confirm the model fileset uses token_secret=hf_secret.name for gated models
    • \n
  • Check AutomodelJobInput references use the workspace/name format: model=f"default/{MODEL_NAME}" and dataset={"training": f"default/{DATASET_NAME}"} (for example, default/llama-3-2-1b-base, default/sft-dataset)
    • \n
  • Verify the model entity points at the fileset: fileset=f"default/{MODEL_NAME}"
    • \n
  • Check job status: client.jobs.get_status(name=job.job.name, workspace="default")
    • \n
\n

Job fails with OOM (Out of Memory) error:

\n
    \n
  1. First try: Reduce global_batch_size from 64 to 32 or 16 in batch={...}
    2. \n
  2. Still OOM: Keep micro_batch_size at 1 (already the minimum in this tutorial)
    3. \n
  3. Still OOM: Reduce max_seq_length from 2048 to 1024 or 512 in training={...}
    4. \n
  4. Last resort: Increase num_gpus_per_node and tensor_parallel_size in parallelism={...}
    5. \n
\n

Loss curves not decreasing (underfitting):

\n
    \n
  • Increase training duration: raise epochs from 2 to 3-5 in schedule={...}
    • \n
  • Adjust learning rate: try 1e-4 or 1e-5 instead of the default 5e-5 in optimizer={...}
    • \n
  • Check data quality: Verify formatting, remove duplicates, ensure diversity
    • \n
\n

Training loss decreases but validation loss increases (overfitting):

\n
    \n
  • Reduce epochs from 2 to 1 in schedule={...}
    • \n
  • Lower learning_rate from 5e-5 to 2e-5 or 1e-5 in optimizer={...}
    • \n
  • Increase dataset size and diversity
    • \n
  • Verify train/validation split has no data leakage
    • \n
\n

Model output quality is poor despite good training metrics:

\n
    \n
  • Training metrics optimize for loss, not your actual task—evaluate on real use cases
    • \n
  • Review data quality, format, and diversity—metrics can be misleading with poor data
    • \n
  • Try a different base model size or architecture
    • \n
  • Adjust learning_rate and global_batch_size
    • \n
  • Compare to baseline: Test base model to ensure fine-tuning improved performance
    • \n
\n

Deployment fails:

\n
    \n
  • Verify output model exists: client.models.retrieve(name=OUTPUT_NAME, workspace="default")
    • \n
  • Check deployment logs: client.inference.deployments.get_logs(name=deployment.name, workspace="default")
    • \n
  • Ensure sufficient GPU resources for executor_config={"gpu": 1, ...}
    • \n
  • Verify the deployment config matches this tutorial: engine="vllm" with vllm/vllm-openai:v0.22.1
    • \n
\n

Next Steps

\n\n" } ] }; diff --git a/docs/fern/package.json b/docs/fern/package.json index af860379c7..c42ea3a2a6 100644 --- a/docs/fern/package.json +++ b/docs/fern/package.json @@ -4,9 +4,10 @@ "prepare:openapi": "node scripts/filter-public-openapi.mjs", "prepare:helm": "node scripts/sync-helm-docs.mjs", "prepare": "npm run prepare:openapi && npm run prepare:helm", - "check": "npm run prepare && npx -y fern-api@latest check && npm run validate-mdx && npm run check:gated-links", + "check": "npm run prepare && npx -y fern-api@latest check && npm run validate-mdx && npm run validate-notebook-viewer && npm run check:gated-links", "check:fern": "npm run prepare && npx -y fern-api@latest check", "validate-mdx": "node scripts/validate-mdx.mjs", + "validate-notebook-viewer": "node scripts/validate-notebook-viewer.mjs", "check:gated-links": "node scripts/delink-gated.mjs", "fix:gated-links": "node scripts/delink-gated.mjs --fix", "broken-links": "npm run prepare && npx -y fern-api@latest docs broken-links", diff --git a/docs/fern/scripts/README.md b/docs/fern/scripts/README.md index cbd3de6053..0e599250b1 100644 --- a/docs/fern/scripts/README.md +++ b/docs/fern/scripts/README.md @@ -21,6 +21,9 @@ uv run python docs/fern/scripts/ipynb-to-fern-json.py \ Writes both `.json` (canonical data) and `.ts` (default-export wrapper that MDX imports). Re-run whenever the source `.ipynb` changes. +`npm run validate-notebook-viewer` (part of `npm run check`) fails if a +NotebookViewer registry entry is missing its generated `.ts` / `.json` pair. + ### MDX usage After writing the `.ts` module, register it in `fern/components/NotebookViewer.tsx` diff --git a/docs/fern/scripts/ipynb-to-mdx.py b/docs/fern/scripts/ipynb-to-mdx.py index f8841e6209..6e0e863e4d 100644 --- a/docs/fern/scripts/ipynb-to-mdx.py +++ b/docs/fern/scripts/ipynb-to-mdx.py @@ -29,6 +29,7 @@ r'Download this tutorial as a Jupyter notebook\s*', re.IGNORECASE, ) +FIRST_H1_RE = re.compile(r"\A# [^\n]+\n+") _LINK_REWRITES: list[tuple[re.Pattern[str], str]] = [ ( @@ -51,6 +52,10 @@ re.compile(r"\]\(\.\./\.\./evaluator/index(?:\.md)?\)"), "](/documentation/evaluate-models)", ), + ( + re.compile(r"\]\(\.\./\.\./evaluator/metrics/rag\.md?\)"), + "](/documentation/evaluate-models/metrics/rag-metrics)", + ), ( re.compile(r"\]\(\./distillation-customization-job(?:\.ipynb)?\)"), "](/documentation/customizer-reference/tutorials/distillation-customization-job)", @@ -108,6 +113,7 @@ def colab_link_for(ipynb_path: Path) -> str: def convert_notebook_to_mdx(ipynb_path: Path, *, title: str) -> str: body = NotebookConverter().convert(ipynb_path) body = DOWNLOAD_LINK_RE.sub("", body).lstrip("\n") + body = FIRST_H1_RE.sub("", body, count=1) body = rewrite_links(body) return ( diff --git a/docs/fern/scripts/validate-notebook-viewer.mjs b/docs/fern/scripts/validate-notebook-viewer.mjs new file mode 100644 index 0000000000..dfc3096671 --- /dev/null +++ b/docs/fern/scripts/validate-notebook-viewer.mjs @@ -0,0 +1,126 @@ +#!/usr/bin/env node +/** + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Ensure every NotebookViewer registration has generated notebook data on disk + * and that the generated JSON still matches its source notebook. + * + * NotebookViewer.tsx imports `./notebooks/` modules produced by + * `ipynb-to-fern-json.py`. A missing `.ts` / `.json` pair breaks publication + * even when the MDX wrapper and registry entry exist. + * + * Run from the fern/ directory: `node scripts/validate-notebook-viewer.mjs`. + */ + +import { access, readFile } from "node:fs/promises"; +import { dirname, join } from "node:path"; +import { fileURLToPath } from "node:url"; + +const ROOT = dirname(fileURLToPath(import.meta.url)); +const VIEWER = join(ROOT, "../components/NotebookViewer.tsx"); +const NOTEBOOKS_DIR = join(ROOT, "../components/notebooks"); +const SOURCE_NOTEBOOKS = { + "distillation-customization-job": join( + ROOT, + "../../customizer/tutorials/distillation-customization-job.ipynb", + ), + "dpo-customization-job": join( + ROOT, + "../../customizer/tutorials/dpo-customization-job.ipynb", + ), + "embedding-customization-job": join( + ROOT, + "../../customizer/tutorials/embedding-customization-job.ipynb", + ), + "lora-customization-job": join( + ROOT, + "../../customizer/tutorials/lora-customization-job.ipynb", + ), + "optimize-throughput": join( + ROOT, + "../../customizer/tutorials/optimize-throughput.ipynb", + ), + "sft-customization-job": join( + ROOT, + "../../customizer/tutorials/sft-customization-job.ipynb", + ), + "tool-calling": join(ROOT, "../../example-applications/tool-calling.ipynb"), +}; + +const IMPORT_RE = + /import\s+\w+\s+from\s+"\.\/notebooks\/([^"]+)";/g; +const REGISTRY_RE = /"([^"]+)":\s*\w+/g; + +const viewerSrc = await readFile(VIEWER, "utf8"); +const imported = [...viewerSrc.matchAll(IMPORT_RE)].map((m) => m[1]); +const registryBlock = viewerSrc.match( + /const notebooks: Record = \{([\s\S]*?)\};/ +)?.[1]; + +if (!registryBlock) { + console.error("validate-notebook-viewer: could not find notebooks registry"); + process.exit(2); +} + +const registered = [...registryBlock.matchAll(REGISTRY_RE)].map((m) => m[1]); +const names = [...new Set([...imported, ...registered])].sort(); + +let failed = 0; +for (const name of names) { + for (const ext of ["ts", "json"]) { + const path = join(NOTEBOOKS_DIR, `${name}.${ext}`); + try { + await access(path); + } catch { + failed += 1; + console.error(`missing ${path}`); + } + } + + const sourcePath = SOURCE_NOTEBOOKS[name]; + if (!sourcePath) { + failed += 1; + console.error(`missing source-notebook mapping for ${name}`); + continue; + } + + try { + const notebook = JSON.parse(await readFile(sourcePath, "utf8")); + const artifact = JSON.parse( + await readFile(join(NOTEBOOKS_DIR, `${name}.json`), "utf8"), + ); + const sourceCells = notebook.cells.map((cell) => + (Array.isArray(cell.source) ? cell.source.join("") : cell.source ?? "").trimEnd(), + ); + const artifactCells = artifact.cells.map((cell) => + (cell.source ?? "").trimEnd(), + ); + + const mismatch = + sourceCells.length !== artifactCells.length || + sourceCells.some((source, index) => source !== artifactCells[index]); + if (mismatch) { + failed += 1; + console.error( + `stale ${join(NOTEBOOKS_DIR, `${name}.json`)}; regenerate it from ${sourcePath}`, + ); + } + } catch (error) { + failed += 1; + console.error(`could not compare ${name} with ${sourcePath}: ${error.message}`); + } +} + +if (failed > 0) { + console.error( + `\nvalidate-notebook-viewer: ${failed} missing or stale notebook artifact(s). ` + + `Run: uv run python docs/fern/scripts/ipynb-to-fern-json.py ` + + `-o docs/fern/components/notebooks/.json` + ); + process.exit(1); +} + +console.log( + `validate-notebook-viewer: ${names.length} NotebookViewer notebook(s) present` +); diff --git a/docs/requirements.mdx b/docs/requirements.mdx index df3cb5b43f..06cf23c27e 100644 --- a/docs/requirements.mdx +++ b/docs/requirements.mdx @@ -39,8 +39,8 @@ Local provider workflows do not require a local GPU. GPU requirements apply only | Component | Requirement | Notes | |-----------|-------------|-------| | GPU | NVIDIA data center GPU with 40 GB VRAM minimum; 80 GB recommended | A100 80GB, H100 80GB, and B200 180GB meet the recommended profile. | -| CUDA | CUDA 12.8 or later | GPU Python dependencies are built for CUDA 12.8. | -| NVIDIA driver | CUDA 12.8-capable driver, R570 branch or later | See the [CUDA 12.8 release notes](https://docs.nvidia.com/cuda/archive/12.8.0/cuda-toolkit-release-notes/index.html) for exact driver minimums by operating system. | +| CUDA | CUDA 13 or later | GPU workloads require CUDA 13 or later. | +| NVIDIA driver | CUDA 13-capable driver, R580 branch or later | See the [CUDA 13.0 release notes](https://docs.nvidia.com/cuda/archive/13.0.0/cuda-toolkit-release-notes/index.html) for exact driver minimums by operating system. | | Platform | Linux x86_64 | Local NVIDIA GPU workloads are not supported on macOS or Windows. | ## Verify Your Environment diff --git a/docs/support-matrix.mdx b/docs/support-matrix.mdx index f986eb5c21..7ba33aa46c 100644 --- a/docs/support-matrix.mdx +++ b/docs/support-matrix.mdx @@ -31,8 +31,8 @@ self-managed Kubernetes deployments installed with the NeMo Platform Helm chart. | Area | Supported | Notes | |------|-----------|-------| -| CUDA | CUDA 12.8 or later | GPU Python dependencies are built for CUDA 12.8. | -| NVIDIA driver | CUDA 12.8-capable driver, R570 branch or later | See the [CUDA 12.8 release notes](https://docs.nvidia.com/cuda/archive/12.8.0/cuda-toolkit-release-notes/index.html) for exact driver minimums by operating system. | +| CUDA | CUDA 13 or later | GPU workloads require CUDA 13 or later. | +| NVIDIA driver | CUDA 13-capable driver, R580 branch or later | See the [CUDA 13.0 release notes](https://docs.nvidia.com/cuda/archive/13.0.0/cuda-toolkit-release-notes/index.html) for exact driver minimums by operating system. | | GPU memory | 40 GB VRAM minimum; 80 GB VRAM recommended | Smaller GPUs can run provider or client workflows, but local model and GPU-accelerated jobs require larger GPUs. | | Recommended GPUs | NVIDIA A100 80GB, H100 80GB, B200 180GB | Use comparable data center GPUs that meet the CUDA and memory requirements for the model or job. | | GPU platform | Linux x86_64 | macOS and Windows are not supported for local NVIDIA GPU workloads. | @@ -64,4 +64,4 @@ The following are not part of the OSS support matrix: - WSL-based local install or GPU validation. - macOS local GPU workloads. - Python earlier than 3.11 or Python 3.15 and later for the documented OSS local install path. -- CUDA versions earlier than 12.8 for GPU-enabled Python dependencies. +- CUDA versions earlier than 13 for GPU workloads. diff --git a/docs/troubleshooting/customizer.mdx b/docs/troubleshooting/customizer.mdx index b84c7f5dbd..7ea0b28651 100644 --- a/docs/troubleshooting/customizer.mdx +++ b/docs/troubleshooting/customizer.mdx @@ -2,13 +2,17 @@ title: "Troubleshooting NeMo Customizer" description: "" --- + **Job fails during model download:** -- Verify the HuggingFace token secret is configured correctly -- Accept the model's license on the [HuggingFace model page](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct) + +- Verify the Hugging Face token secret is configured correctly +- Accept the model's license on the [Hugging Face model page](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct) - Check job status: `client.jobs.get_status(name=job.name, workspace="default")` **Job fails with disk full or 500 error when retrieving logs:** -- The platform's shared persistent volume is likely full. Customization jobs require significant disk space: ~3× model size for full SFT, ~1.5× for LoRA. If you are also deploying the model from a base checkpoint fileset, plan for ~2.5× model size overall. + +- The platform's shared persistent volume is likely full. Budget against the downloaded base checkpoint size: approximately 3× for Full SFT and 1.5× for LoRA. For example, a 70B BF16 checkpoint is approximately 140 GB, so a Full SFT job can require approximately 420 GB of free disk at peak. +- These peak estimates include the base checkpoint and job artifacts; the final Full SFT output itself is one full checkpoint. If you also retain a deployment copy, include it separately in capacity planning. - Clean up completed job artifacts or increase the PVC size (default: 200Gi at `/var/run/scratch/job`). - DPO/GRPO jobs also consume ephemeral node storage under `/tmp` via Ray workers — check node disk in addition to the PVC. - See [ft-tut-understand-models](/documentation/customizer-reference/tutorials/understanding-models-and-training) for full storage requirement details. @@ -27,17 +31,20 @@ Batch and sequence-length fields differ by backend. Use the fully qualified path 3. Reduce `model.max_seq_length` from 2048 to 1024 or 512 **Training loss not decreasing:** + - Increase `optimizer.learning_rate` (try 2e-4 or 5e-4), same path for Automodel and Unsloth - Increase `schedule.epochs`, same path for Automodel and Unsloth - Verify data quality -- inspect a few training examples manually **Tool calling accuracy is low after fine-tuning:** + - Increase training data size (sample more from the filtered dataset) - Increase `schedule.epochs` to a higher value. If you are running for 1-2 epochs, increase it to 3-4. - Check that the evaluation dataset format matches what the model expects - Verify the base model supports tool calling (Llama 3.2 Instruct does) **Deployment fails:** + - Verify the base model and adapter exist: `client.models.retrieve(name=MODEL_NAME, workspace="default")` -- the LoRA adapter appears in the base model's `adapters` list, not as a separate model entity - Check deployment logs: `client.inference.deployments.get_logs(name=deployment.name, workspace="default")` - Ensure sufficient GPU resources for the model size From 285850eee3be794b004e3eb665f0d11046ac0e01 Mon Sep 17 00:00:00 2001 From: Sam O Date: Tue, 21 Jul 2026 15:24:12 -0600 Subject: [PATCH 02/24] Apply suggestions from code review Co-authored-by: Albert Cui Signed-off-by: Sam Oluwalana --- docs/customizer/about.mdx | 4 ++-- docs/customizer/models/gpt-oss.mdx | 2 +- docs/customizer/tutorials/optimize-throughput.mdx | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/docs/customizer/about.mdx b/docs/customizer/about.mdx index e46563195f..bb5f1906f4 100644 --- a/docs/customizer/about.mdx +++ b/docs/customizer/about.mdx @@ -19,7 +19,7 @@ Full SFT models require a NIM deployment to serve inference. The Deployment Mana | **Multi-LLM** (Default) | Generic multi-model NIM | On-the-fly download via Files service | Supported Hugging Face architectures, custom fine-tuned models, development | | **Model-Specific NIM** | Dedicated model image | Pre-download via model puller | Production, optimized performance and latency | -- **Multi-LLM Image**: Can deploy Hugging Face checkpoints whose architectures are supported by the image's inference engine, providing flexibility for custom fine-tuned models. Importing a checkpoint does not guarantee training or deployment compatibility; for example, Automodel LoRA does not support Conv1D-based architectures. Does not guarantee optimized inference performance. +- **Multi-LLM Image**: Can deploy Hugging Face checkpoints whose architectures are supported by the image's inference engine, providing flexibility for custom fine-tuned models. Importing a checkpoint does not guarantee training or deployment compatibility; for example, Automodel LoRA does not support Conv1D-based architectures. It also does not guarantee optimized inference performance. - **Model-Specific NIM**: Provides optimized inference performance and latency through model-specific optimizations. Recommended for production deployments where performance is critical. @@ -28,7 +28,7 @@ Full SFT models require a NIM deployment to serve inference. The Deployment Mana Parameter-Efficient Fine-Tuning (PEFT) methods enable efficient model customization by training a small number of parameters while keeping the base model frozen. For example, when customizing Llama 3.3 70B: - **Traditional SFT**: Produces a full BF16 checkpoint of approximately 140 GB per task. During training, budget free disk space separately for the base checkpoint, intermediate checkpoint, and final output—approximately 3× the downloaded base checkpoint size. -- **PEFT**: Produces an adapter that is typically approximately 100–500 MB per task while maintaining comparable performance. During training, budget approximately 1.5× the downloaded base checkpoint size. +- **PEFT**: Produces an adapter that is typically approximately ~100-500 MB per task while maintaining comparable performance. During training, budget approximately 1.5× the downloaded base checkpoint size. ```mermaid --- diff --git a/docs/customizer/models/gpt-oss.mdx b/docs/customizer/models/gpt-oss.mdx index 20cffd397e..4012f0d3c3 100644 --- a/docs/customizer/models/gpt-oss.mdx +++ b/docs/customizer/models/gpt-oss.mdx @@ -23,7 +23,7 @@ These models require a Hugging Face token to download. Create a secret with your | Parameters | 21B parameters (3.6B active parameters) | | Training Data | Trained on harmony response format | | Checkpoint Quantization | MXFP4 quantization of the MoE weights | -| Inference Memory | The official checkpoint can run within 16GB of memory; Customizer training requires the GPU configurations below | +| Inference Memory | The official checkpoint can run within 16 GB of memory; Customizer training requires the GPU configurations below | | Default Name | openai/gpt-oss-20b | | Hugging Face | [openai/gpt-oss-20b](https://huggingface.co/openai/gpt-oss-20b) | diff --git a/docs/customizer/tutorials/optimize-throughput.mdx b/docs/customizer/tutorials/optimize-throughput.mdx index c06bbf193f..9d5451dfa4 100644 --- a/docs/customizer/tutorials/optimize-throughput.mdx +++ b/docs/customizer/tutorials/optimize-throughput.mdx @@ -211,7 +211,7 @@ 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 +# Uncomment the line below if you have NGC API Key and want to fine-tune NGC models # ngc_api_key = create_or_get_secret("ngc-api-key", NGC_API_KEY, "NGC_API_KEY") ``` From 729f384d43553b1b4825264a49510bc9ca2dd7e8 Mon Sep 17 00:00:00 2001 From: Sam Oluwalana Date: Tue, 21 Jul 2026 15:19:40 -0600 Subject: [PATCH 03/24] Code Rabbit resolutions Signed-off-by: Sam Oluwalana --- docs/customizer/about.mdx | 37 ++++++++++--------- .../hyperparameters.mdx | 4 +- .../distillation-customization-job.ipynb | 26 +++++++++++-- .../distillation-customization-job.mdx | 26 +++++++++++-- .../tutorials/dpo-customization-job.ipynb | 24 +++++++----- .../embedding-customization-job.ipynb | 13 ++++--- .../tutorials/embedding-customization-job.mdx | 13 ++++--- .../tutorials/format-training-dataset.mdx | 14 +++---- .../tutorials/sft-customization-job.ipynb | 33 ++++++++++------- .../tutorials/sft-customization-job.mdx | 33 ++++++++++------- .../understand-configurations-and-models.mdx | 10 ++--- .../distillation-customization-job.json | 4 +- .../distillation-customization-job.ts | 4 +- .../notebooks/dpo-customization-job.json | 12 +++--- .../notebooks/dpo-customization-job.ts | 12 +++--- .../embedding-customization-job.json | 4 +- .../notebooks/embedding-customization-job.ts | 4 +- .../notebooks/sft-customization-job.json | 10 ++--- .../notebooks/sft-customization-job.ts | 8 ++-- .../fern/scripts/validate-notebook-viewer.mjs | 31 +++++++++++++--- 20 files changed, 201 insertions(+), 121 deletions(-) diff --git a/docs/customizer/about.mdx b/docs/customizer/about.mdx index bb5f1906f4..0e64e7bcb0 100644 --- a/docs/customizer/about.mdx +++ b/docs/customizer/about.mdx @@ -2,6 +2,7 @@ title: "Customization Concepts" description: "" --- + This page provides an overview of the customization concepts for the NeMo Platform. @@ -14,10 +15,10 @@ Supervised fine-tuning (SFT) is a traditional technique for customizing a pre-tr Full SFT models require a NIM deployment to serve inference. The Deployment Management Service supports two deployment modes: -| Deployment Mode | Image Type | Weight Loading | Best For | -|-----------------|------------|----------------|----------| +| Deployment Mode | Image Type | Weight Loading | Best For | +| ----------------------- | ----------------------- | ------------------------------------- | --------------------------------------------------------------------------- | | **Multi-LLM** (Default) | Generic multi-model NIM | On-the-fly download via Files service | Supported Hugging Face architectures, custom fine-tuned models, development | -| **Model-Specific NIM** | Dedicated model image | Pre-download via model puller | Production, optimized performance and latency | +| **Model-Specific NIM** | Dedicated model image | Pre-download via model puller | Production, optimized performance and latency | - **Multi-LLM Image**: Can deploy Hugging Face checkpoints whose architectures are supported by the image's inference engine, providing flexibility for custom fine-tuned models. Importing a checkpoint does not guarantee training or deployment compatibility; for example, Automodel LoRA does not support Conv1D-based architectures. It also does not guarantee optimized inference performance. @@ -28,7 +29,7 @@ Full SFT models require a NIM deployment to serve inference. The Deployment Mana Parameter-Efficient Fine-Tuning (PEFT) methods enable efficient model customization by training a small number of parameters while keeping the base model frozen. For example, when customizing Llama 3.3 70B: - **Traditional SFT**: Produces a full BF16 checkpoint of approximately 140 GB per task. During training, budget free disk space separately for the base checkpoint, intermediate checkpoint, and final output—approximately 3× the downloaded base checkpoint size. -- **PEFT**: Produces an adapter that is typically approximately ~100-500 MB per task while maintaining comparable performance. During training, budget approximately 1.5× the downloaded base checkpoint size. +- **PEFT**: Produces an adapter that is approximately 100–500 MB per task while maintaining comparable performance. During training, budget approximately 1.5× the downloaded base checkpoint size. ```mermaid --- @@ -121,6 +122,7 @@ Additional Resources: ## Training with Your Own Data Use NeMo Customizer to train custom models on your own data. The workflow can be carried out as follows: + - Upload a dataset - Train a custom model - Perform inference with the trained model @@ -205,13 +207,13 @@ Hyperparameters are configuration settings used to control the training process. Common hyperparameters you'll tune include: -| Hyperparameter | Description | -|----------------|-------------| -| Epochs | Number of complete passes through the training dataset | -| Batch size | Number of samples processed before updating model weights | -| Learning rate | Step size for weight updates during training | -| LoRA rank | Low-rank dimension of the adapter (lower = fewer parameters, higher = more expressive) | -| LoRA alpha | LoRA scaling factor | +| Hyperparameter | Description | +| -------------- | -------------------------------------------------------------------------------------- | +| Epochs | Number of complete passes through the training dataset | +| Batch size | Number of samples processed before updating model weights | +| Learning rate | Step size for weight updates during training | +| LoRA rank | Low-rank dimension of the adapter (lower = fewer parameters, higher = more expressive) | +| LoRA alpha | LoRA scaling factor | @@ -243,14 +245,15 @@ PP can be configured via `parallelism.pipeline_parallel_size` in the [training c #### Configuration - Constraints - - TP must be less than or equal to the total number of GPUs available. It should be a factor of the total GPU count (divisible evenly). +- TP must be less than or equal to the total number of GPUs available. It should be a factor of the total GPU count (divisible evenly). - Multi-node considerations - - TP can span across nodes, but this introduces network communication overhead. For multi-node setups, it's often recommended to keep TP within a single node when possible. If using TP across nodes, high-bandwidth inter-node connections (like InfiniBand) become critical. - - Example: if you have 2 nodes with 4 GPUs each, start with TP=4 first. This keeps all tensor parallel operations within a single node. If your model still uses too much GPU memory with this setting, increase to TP=8, which will distribute tensor operations across both nodes. +- TP can span across nodes, but this introduces network communication overhead. For multi-node setups, it's often recommended to keep TP within a single node when possible. If using TP across nodes, high-bandwidth inter-node connections (like InfiniBand) become critical. + +Example: if you have 2 nodes with 4 GPUs each, start with TP=4 first. This keeps all tensor parallel operations within a single node. If your model still uses too much GPU memory with this setting, increase to TP=8, which will distribute tensor operations across both nodes. + - Performance - - Smaller TP values generally have less communication overhead. - - Larger TP values provide more memory savings but increase communication costs. +- Smaller TP values generally have less communication overhead. +- Larger TP values provide more memory savings but increase communication costs. ### Context Parallelism diff --git a/docs/customizer/manage-customization-jobs/hyperparameters.mdx b/docs/customizer/manage-customization-jobs/hyperparameters.mdx index 3265c29a1e..d969f05a95 100644 --- a/docs/customizer/manage-customization-jobs/hyperparameters.mdx +++ b/docs/customizer/manage-customization-jobs/hyperparameters.mdx @@ -238,7 +238,7 @@ LoRA parameters (`training.lora`): |-----------|-------------|---------| | `output.name` | Output Model Entity or adapter name | Auto-generated from the job name | | `output.description` | Optional description for the generated artifact | `null` | -| `output.save_method` | Saved checkpoint shape (see below) | `lora` | +| `output.save_method` | LoRA checkpoint serialization (see below); omit for full-weight training | `lora` | The `output.save_method` field accepts: @@ -248,7 +248,7 @@ The `output.save_method` field accepts: | `merged_16bit` | Merges the adapter into the base and saves a 16-bit checkpoint | | `merged_4bit` | Merges the adapter into the base and saves a 4-bit checkpoint | -The `merged_*` methods are only valid when `training.finetuning_type` is `lora`. +The `merged_*` methods are only valid when `training.finetuning_type` is `lora`. When `training.finetuning_type` is `all_weights`, omit `output.save_method`; the training driver saves the full trained checkpoint. --- diff --git a/docs/customizer/tutorials/distillation-customization-job.ipynb b/docs/customizer/tutorials/distillation-customization-job.ipynb index af0e2830a6..39c5615d0b 100644 --- a/docs/customizer/tutorials/distillation-customization-job.ipynb +++ b/docs/customizer/tutorials/distillation-customization-job.ipynb @@ -514,8 +514,17 @@ "\n", " if dep.status == \"READY\":\n", " print(\"\\nDeployment is ready!\")\n", - " if not client.models.wait_for_gateway(deployment_name, workspace=\"default\", timeout=60):\n", - " raise RuntimeError(\"Inference gateway did not become ready\")\n", + " remaining = int(timeout - elapsed)\n", + " if remaining <= 0:\n", + " raise TimeoutError(f\"Deployment timeout after {timeout_minutes} minutes\")\n", + " if not client.models.wait_for_status(\n", + " deployment_name=deployment_name,\n", + " desired_status=\"READY\",\n", + " workspace=\"default\",\n", + " timeout=remaining,\n", + " check_gateway=True,\n", + " ):\n", + " raise TimeoutError(\"Inference gateway did not become ready\")\n", " return dep\n", " if dep.status in (\"FAILED\", \"ERROR\", \"TERMINATED\", \"LOST\"):\n", " raise RuntimeError(f\"Deployment failed with status: {dep.status}\")\n", @@ -524,8 +533,17 @@ " time.sleep(15)\n", "\n", "\n", - "dep_status = wait_for_deployment(BASELINE_DEPLOYMENT_NAME)\n", - "assert dep_status.status == \"READY\"" + "try:\n", + " dep_status = wait_for_deployment(BASELINE_DEPLOYMENT_NAME)\n", + " assert dep_status.status == \"READY\"\n", + "except Exception:\n", + " # Free GPUs if readiness fails before the later baseline-cleanup cell runs.\n", + " try:\n", + " client.inference.deployments.delete(name=BASELINE_DEPLOYMENT_NAME, workspace=\"default\")\n", + " client.inference.deployment_configs.delete(name=BASELINE_DEPLOYMENT_CONFIG, workspace=\"default\")\n", + " except Exception as cleanup_error:\n", + " print(f\"Baseline cleanup after readiness failure also failed: {cleanup_error}\")\n", + " raise" ], "execution_count": null, "outputs": [] diff --git a/docs/customizer/tutorials/distillation-customization-job.mdx b/docs/customizer/tutorials/distillation-customization-job.mdx index 3c78c4cae3..d3b694e541 100644 --- a/docs/customizer/tutorials/distillation-customization-job.mdx +++ b/docs/customizer/tutorials/distillation-customization-job.mdx @@ -433,8 +433,17 @@ def wait_for_deployment(deployment_name: str, timeout_minutes: int = 30): if dep.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") + remaining = int(timeout - elapsed) + if remaining <= 0: + raise TimeoutError(f"Deployment timeout after {timeout_minutes} minutes") + if not client.models.wait_for_status( + deployment_name=deployment_name, + desired_status="READY", + workspace="default", + timeout=remaining, + check_gateway=True, + ): + raise TimeoutError("Inference gateway did not become ready") return dep if dep.status in ("FAILED", "ERROR", "TERMINATED", "LOST"): raise RuntimeError(f"Deployment failed with status: {dep.status}") @@ -443,8 +452,17 @@ def wait_for_deployment(deployment_name: str, timeout_minutes: int = 30): time.sleep(15) -dep_status = wait_for_deployment(BASELINE_DEPLOYMENT_NAME) -assert dep_status.status == "READY" +try: + dep_status = wait_for_deployment(BASELINE_DEPLOYMENT_NAME) + assert dep_status.status == "READY" +except Exception: + # Free GPUs if readiness fails before the later baseline-cleanup cell runs. + try: + client.inference.deployments.delete(name=BASELINE_DEPLOYMENT_NAME, workspace="default") + client.inference.deployment_configs.delete(name=BASELINE_DEPLOYMENT_CONFIG, workspace="default") + except Exception as cleanup_error: + print(f"Baseline cleanup after readiness failure also failed: {cleanup_error}") + raise ``` ### 7. Generate Baseline Predictions on Test Set diff --git a/docs/customizer/tutorials/dpo-customization-job.ipynb b/docs/customizer/tutorials/dpo-customization-job.ipynb index 6cd9dd4110..5c8437dfcb 100644 --- a/docs/customizer/tutorials/dpo-customization-job.ipynb +++ b/docs/customizer/tutorials/dpo-customization-job.ipynb @@ -214,11 +214,10 @@ "metadata": {}, "source": [ "HF_TOKEN = os.getenv(\"HF_TOKEN\")\n", + "if not HF_TOKEN:\n", + " raise RuntimeError(\"Set HF_TOKEN before running this tutorial.\")\n", "\n", - "def create_or_get_secret(name: str, value: str | None, label: str) -> PlatformSecretResponse | None:\n", - " if not value:\n", - " print(f\"{label} is not set - skipping secret (gated model downloads will fail without it)\")\n", - " return None\n", + "def create_or_get_secret(name: str, value: str, label: str) -> PlatformSecretResponse:\n", " try:\n", " secret = sdk.secrets.create(name=name, workspace=\"default\", value=value)\n", " print(f\"Created secret: {name}\")\n", @@ -249,9 +248,12 @@ "HF_REPO_ID = \"meta-llama/Llama-3.2-1B-Instruct\"\n", "MODEL_NAME = \"llama-3-2-1b-instruct\"\n", "\n", - "storage = HuggingfaceStorageConfigParam(type=\"huggingface\", repo_id=HF_REPO_ID, repo_type=\"model\")\n", - "if hf_secret:\n", - " storage[\"token_secret\"] = hf_secret.name\n", + "storage = HuggingfaceStorageConfigParam(\n", + " type=\"huggingface\",\n", + " repo_id=HF_REPO_ID,\n", + " repo_type=\"model\",\n", + " token_secret=hf_secret.name,\n", + ")\n", "\n", "try:\n", " base_model_fs = sdk.files.filesets.create(\n", @@ -278,7 +280,8 @@ "print(\"Model spec ready\")" ], "execution_count": null, - "outputs": [] + "outputs": [], + "id": "1b798ede" }, { "cell_type": "markdown", @@ -414,7 +417,8 @@ "### 9. Deploy and Evaluate (optional)\n", "\n", "The DPO output is a full model, so it deploys like any full-weight checkpoint (see the [Full SFT](/documentation/customizer-reference/tutorials/sft-customization-job) tutorial for details). We deploy with vLLM and send a chat completion." - ] + ], + "id": "fbbce3cb" }, { "cell_type": "code", @@ -496,7 +500,7 @@ "- Submitted a full-weight DPO job that ran on a Ray cluster via the Kubernetes executor.\n", "- Registered the output as a full model entity and (optionally) deployed it for inference.\n", "\n", - "**Next steps:** tune the alignment strength with `ref_policy_kl_penalty` (β), add `sft_loss_weight` to anchor the policy to the chosen responses, enable `activation_checkpointing` for memory headroom, or scale up with `parallelism`. See the `nemo-customizer` skill's `references/hyperparameters.md` (section NeMo-RL (DPO)) for the full knob reference." + "**Next steps:** tune the alignment strength with `ref_policy_kl_penalty` (β), add `sft_loss_weight` to anchor the policy to the chosen responses, enable `activation_checkpointing` for memory headroom, or scale up with `parallelism`. See the [Training Configuration](/documentation/customizer-reference/manage-customization-jobs/training-configuration) reference for the full hyperparameter set." ] } ], diff --git a/docs/customizer/tutorials/embedding-customization-job.ipynb b/docs/customizer/tutorials/embedding-customization-job.ipynb index ee6ebc3796..b552332593 100644 --- a/docs/customizer/tutorials/embedding-customization-job.ipynb +++ b/docs/customizer/tutorials/embedding-customization-job.ipynb @@ -505,13 +505,14 @@ "HF_REPO_ID = \"nvidia/llama-nemotron-embed-1b-v2\"\n", "MODEL_NAME = \"nv-nemotron-embed-1b-base\"\n", "\n", - "storage = HuggingfaceStorageConfigParam(\n", - " type=\"huggingface\",\n", - " repo_id=HF_REPO_ID,\n", - " repo_type=\"model\",\n", - ")\n", + "storage_kwargs = {\n", + " \"type\": \"huggingface\",\n", + " \"repo_id\": HF_REPO_ID,\n", + " \"repo_type\": \"model\",\n", + "}\n", "if hf_secret:\n", - " storage[\"token_secret\"] = hf_secret.name\n", + " storage_kwargs[\"token_secret\"] = hf_secret.name\n", + "storage = HuggingfaceStorageConfigParam(**storage_kwargs)\n", "\n", "try:\n", " base_model_fs = client.files.filesets.create(\n", diff --git a/docs/customizer/tutorials/embedding-customization-job.mdx b/docs/customizer/tutorials/embedding-customization-job.mdx index a47770e29b..41151674d8 100644 --- a/docs/customizer/tutorials/embedding-customization-job.mdx +++ b/docs/customizer/tutorials/embedding-customization-job.mdx @@ -408,13 +408,14 @@ from nemo_platform.types.files import HuggingfaceStorageConfigParam HF_REPO_ID = "nvidia/llama-nemotron-embed-1b-v2" MODEL_NAME = "nv-nemotron-embed-1b-base" -storage = HuggingfaceStorageConfigParam( - type="huggingface", - repo_id=HF_REPO_ID, - repo_type="model", -) +storage_kwargs = { + "type": "huggingface", + "repo_id": HF_REPO_ID, + "repo_type": "model", +} if hf_secret: - storage["token_secret"] = hf_secret.name + storage_kwargs["token_secret"] = hf_secret.name +storage = HuggingfaceStorageConfigParam(**storage_kwargs) try: base_model_fs = client.files.filesets.create( diff --git a/docs/customizer/tutorials/format-training-dataset.mdx b/docs/customizer/tutorials/format-training-dataset.mdx index b60e8f88c1..c16b236123 100644 --- a/docs/customizer/tutorials/format-training-dataset.mdx +++ b/docs/customizer/tutorials/format-training-dataset.mdx @@ -330,13 +330,6 @@ print(f"Response: {response.choices[0].message.content}") ```` -## Next Steps - -Now that you know how to format your training datasets, you can proceed with creating customization jobs: - -- [Start a LoRA Model Customization Job](/documentation/customizer-reference/tutorials/lora-customization-job) - For parameter-efficient fine-tuning -- [Start a Full SFT Customization Job](/documentation/customizer-reference/tutorials/sft-customization-job) - For full model fine-tuning - --- ## Completion Models @@ -382,3 +375,10 @@ response = oai_client.completions.create( print(f"Response: {response.choices[0].text}") ```` + +## Next Steps + +Now that you know how to format your training datasets, you can proceed with creating customization jobs: + +- [Start a LoRA Model Customization Job](/documentation/customizer-reference/tutorials/lora-customization-job) - For parameter-efficient fine-tuning +- [Start a Full SFT Customization Job](/documentation/customizer-reference/tutorials/sft-customization-job) - For full model fine-tuning diff --git a/docs/customizer/tutorials/sft-customization-job.ipynb b/docs/customizer/tutorials/sft-customization-job.ipynb index 8a6416ebf6..df5aba6854 100644 --- a/docs/customizer/tutorials/sft-customization-job.ipynb +++ b/docs/customizer/tutorials/sft-customization-job.ipynb @@ -358,10 +358,14 @@ " return client.secrets.retrieve(name=name, workspace=\"default\")\n", "\n", "\n", - "# Create Hugging Face token secret\n", - "hf_secret = create_or_get_secret(\"hf-token\", HF_TOKEN, \"HF_TOKEN\")\n", - "print(\"HF_TOKEN secret:\")\n", - "print(hf_secret.model_dump_json(indent=2))\n", + "# Create Hugging Face token secret only when HF_TOKEN is set.\n", + "# Gated models (for example Llama) require this secret; public models can omit token_secret.\n", + "hf_secret = create_or_get_secret(\"hf-token\", HF_TOKEN, \"HF_TOKEN\") if HF_TOKEN else None\n", + "if hf_secret:\n", + " print(\"HF_TOKEN secret:\")\n", + " print(hf_secret.model_dump_json(indent=2))\n", + "else:\n", + " print(\"HF_TOKEN is not set; continuing without token_secret (public models only).\")\n", "\n", "# Create NGC API key secret\n", "# Uncomment the line below if you have NGC API Key and want to finetune NGC models\n", @@ -393,20 +397,23 @@ "HF_REPO_ID = \"meta-llama/Llama-3.2-1B-Instruct\"\n", "MODEL_NAME = \"llama-3-2-1b-base\"\n", "\n", - "# Ensure you have a Hugging Face token secret created\n", + "storage_kwargs = {\n", + " \"type\": \"huggingface\",\n", + " # repo_id is the full model name from Hugging Face\n", + " \"repo_id\": HF_REPO_ID,\n", + " \"repo_type\": \"model\",\n", + "}\n", + "if hf_secret:\n", + " # we use the secret created in the previous step for gated models\n", + " storage_kwargs[\"token_secret\"] = hf_secret.name\n", + "\n", + "# Ensure you have a Hugging Face token secret created for gated models\n", "try:\n", " base_model_fs = client.files.filesets.create(\n", " workspace=\"default\",\n", " name=MODEL_NAME,\n", " description=\"Llama 3.2 1B base model from Hugging Face\",\n", - " storage=HuggingfaceStorageConfigParam(\n", - " type=\"huggingface\",\n", - " # repo_id is the full model name from Hugging Face\n", - " repo_id=HF_REPO_ID,\n", - " repo_type=\"model\",\n", - " # we use the secret created in the previous step\n", - " token_secret=hf_secret.name\n", - " )\n", + " storage=HuggingfaceStorageConfigParam(**storage_kwargs),\n", " )\n", " print(f\"Created base model fileset: {MODEL_NAME}\")\n", "except ConflictError:\n", diff --git a/docs/customizer/tutorials/sft-customization-job.mdx b/docs/customizer/tutorials/sft-customization-job.mdx index 24a2038d26..508610411e 100644 --- a/docs/customizer/tutorials/sft-customization-job.mdx +++ b/docs/customizer/tutorials/sft-customization-job.mdx @@ -270,10 +270,14 @@ def create_or_get_secret(name: str, value: str | None, label: str): return client.secrets.retrieve(name=name, workspace="default") -# Create Hugging Face token secret -hf_secret = create_or_get_secret("hf-token", HF_TOKEN, "HF_TOKEN") -print("HF_TOKEN secret:") -print(hf_secret.model_dump_json(indent=2)) +# Create Hugging Face token secret only when HF_TOKEN is set. +# Gated models (for example Llama) require this secret; public models can omit token_secret. +hf_secret = create_or_get_secret("hf-token", HF_TOKEN, "HF_TOKEN") if HF_TOKEN else None +if hf_secret: + print("HF_TOKEN secret:") + print(hf_secret.model_dump_json(indent=2)) +else: + print("HF_TOKEN is not set; continuing without token_secret (public models only).") # Create NGC API key secret # Uncomment the line below if you have NGC API Key and want to finetune NGC models @@ -295,20 +299,23 @@ from nemo_platform.types.files import HuggingfaceStorageConfigParam HF_REPO_ID = "meta-llama/Llama-3.2-1B-Instruct" MODEL_NAME = "llama-3-2-1b-base" -# Ensure you have a Hugging Face token secret created +storage_kwargs = { + "type": "huggingface", + # repo_id is the full model name from Hugging Face + "repo_id": HF_REPO_ID, + "repo_type": "model", +} +if hf_secret: + # we use the secret created in the previous step for gated models + storage_kwargs["token_secret"] = hf_secret.name + +# Ensure you have a Hugging Face token secret created for gated models try: base_model_fs = client.files.filesets.create( workspace="default", name=MODEL_NAME, description="Llama 3.2 1B base model from Hugging Face", - storage=HuggingfaceStorageConfigParam( - type="huggingface", - # 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 - token_secret=hf_secret.name - ) + storage=HuggingfaceStorageConfigParam(**storage_kwargs), ) print(f"Created base model fileset: {MODEL_NAME}") except ConflictError: diff --git a/docs/customizer/tutorials/understand-configurations-and-models.mdx b/docs/customizer/tutorials/understand-configurations-and-models.mdx index 1d36654a17..bffd2ce61c 100644 --- a/docs/customizer/tutorials/understand-configurations-and-models.mdx +++ b/docs/customizer/tutorials/understand-configurations-and-models.mdx @@ -286,11 +286,11 @@ When you set `expert_parallel_size`: - `(data_parallel_size × context_parallel_size)` must be divisible by `expert_parallel_size`. - When `expert_parallel_size > 1`, `tensor_parallel_size` must be `1`. -For example, with 8 total GPUs, `tensor_parallel_size=2`, `pipeline_parallel_size=1`, and `context_parallel_size=1`: -- Derived `data_parallel_size = 8 / (2 × 1 × 1) = 4` -- `data_parallel_size × context_parallel_size = 4` -- Valid `expert_parallel_size` values: `1`, `2`, or `4` -- Invalid `expert_parallel_size` value: `3` (does not divide 4) +For example, with 8 total GPUs, `tensor_parallel_size=1`, `pipeline_parallel_size=1`, and `context_parallel_size=1`: +- Derived `data_parallel_size = 8 / (1 × 1 × 1) = 8` +- `data_parallel_size × context_parallel_size = 8` +- Valid `expert_parallel_size` values: `1`, `2`, `4`, or `8` +- Invalid `expert_parallel_size` value: `3` (does not divide 8) ### Resource Allocation Rules diff --git a/docs/fern/components/notebooks/distillation-customization-job.json b/docs/fern/components/notebooks/distillation-customization-job.json index c25043b5b4..dfea0c50a5 100644 --- a/docs/fern/components/notebooks/distillation-customization-job.json +++ b/docs/fern/components/notebooks/distillation-customization-job.json @@ -90,9 +90,9 @@ }, { "type": "code", - "source": "def wait_for_deployment(deployment_name: str, timeout_minutes: int = 30):\n \"\"\"Poll deployment until ready.\"\"\"\n start = time.time()\n timeout = timeout_minutes * 60\n while True:\n dep = client.inference.deployments.retrieve(name=deployment_name, workspace=\"default\")\n elapsed = time.time() - start\n clear_output(wait=True)\n print(f\"Deployment: {deployment_name}\")\n print(f\"Status: {dep.status}\")\n print(f\"Elapsed: {int(elapsed // 60)}m {int(elapsed % 60)}s\")\n\n if dep.status == \"READY\":\n print(\"\\nDeployment is ready!\")\n if not client.models.wait_for_gateway(deployment_name, workspace=\"default\", timeout=60):\n raise RuntimeError(\"Inference gateway did not become ready\")\n return dep\n if dep.status in (\"FAILED\", \"ERROR\", \"TERMINATED\", \"LOST\"):\n raise RuntimeError(f\"Deployment failed with status: {dep.status}\")\n if elapsed > timeout:\n raise TimeoutError(f\"Deployment timeout after {timeout_minutes} minutes\")\n time.sleep(15)\n\n\ndep_status = wait_for_deployment(BASELINE_DEPLOYMENT_NAME)\nassert dep_status.status == \"READY\"", + "source": "def wait_for_deployment(deployment_name: str, timeout_minutes: int = 30):\n \"\"\"Poll deployment until ready.\"\"\"\n start = time.time()\n timeout = timeout_minutes * 60\n while True:\n dep = client.inference.deployments.retrieve(name=deployment_name, workspace=\"default\")\n elapsed = time.time() - start\n clear_output(wait=True)\n print(f\"Deployment: {deployment_name}\")\n print(f\"Status: {dep.status}\")\n print(f\"Elapsed: {int(elapsed // 60)}m {int(elapsed % 60)}s\")\n\n if dep.status == \"READY\":\n print(\"\\nDeployment is ready!\")\n remaining = int(timeout - elapsed)\n if remaining <= 0:\n raise TimeoutError(f\"Deployment timeout after {timeout_minutes} minutes\")\n if not client.models.wait_for_status(\n deployment_name=deployment_name,\n desired_status=\"READY\",\n workspace=\"default\",\n timeout=remaining,\n check_gateway=True,\n ):\n raise TimeoutError(\"Inference gateway did not become ready\")\n return dep\n if dep.status in (\"FAILED\", \"ERROR\", \"TERMINATED\", \"LOST\"):\n raise RuntimeError(f\"Deployment failed with status: {dep.status}\")\n if elapsed > timeout:\n raise TimeoutError(f\"Deployment timeout after {timeout_minutes} minutes\")\n time.sleep(15)\n\n\ntry:\n dep_status = wait_for_deployment(BASELINE_DEPLOYMENT_NAME)\n assert dep_status.status == \"READY\"\nexcept Exception:\n # Free GPUs if readiness fails before the later baseline-cleanup cell runs.\n try:\n client.inference.deployments.delete(name=BASELINE_DEPLOYMENT_NAME, workspace=\"default\")\n client.inference.deployment_configs.delete(name=BASELINE_DEPLOYMENT_CONFIG, workspace=\"default\")\n except Exception as cleanup_error:\n print(f\"Baseline cleanup after readiness failure also failed: {cleanup_error}\")\n raise", "language": "python", - "source_html": "def wait_for_deployment(deployment_name: str, timeout_minutes: int = 30):\n """Poll deployment until ready."""\n start = time.time()\n timeout = timeout_minutes * 60\n while True:\n dep = client.inference.deployments.retrieve(name=deployment_name, workspace="default")\n elapsed = time.time() - start\n clear_output(wait=True)\n print(f"Deployment: {deployment_name}")\n print(f"Status: {dep.status}")\n print(f"Elapsed: {int(elapsed // 60)}m {int(elapsed % 60)}s")\n\n if dep.status == "READY":\n print("\\nDeployment is ready!")\n if not client.models.wait_for_gateway(deployment_name, workspace="default", timeout=60):\n raise RuntimeError("Inference gateway did not become ready")\n return dep\n if dep.status in ("FAILED", "ERROR", "TERMINATED", "LOST"):\n raise RuntimeError(f"Deployment failed with status: {dep.status}")\n if elapsed > timeout:\n raise TimeoutError(f"Deployment timeout after {timeout_minutes} minutes")\n time.sleep(15)\n\n\ndep_status = wait_for_deployment(BASELINE_DEPLOYMENT_NAME)\nassert dep_status.status == "READY"\n" + "source_html": "def wait_for_deployment(deployment_name: str, timeout_minutes: int = 30):\n """Poll deployment until ready."""\n start = time.time()\n timeout = timeout_minutes * 60\n while True:\n dep = client.inference.deployments.retrieve(name=deployment_name, workspace="default")\n elapsed = time.time() - start\n clear_output(wait=True)\n print(f"Deployment: {deployment_name}")\n print(f"Status: {dep.status}")\n print(f"Elapsed: {int(elapsed // 60)}m {int(elapsed % 60)}s")\n\n if dep.status == "READY":\n print("\\nDeployment is ready!")\n remaining = int(timeout - elapsed)\n if remaining <= 0:\n raise TimeoutError(f"Deployment timeout after {timeout_minutes} minutes")\n if not client.models.wait_for_status(\n deployment_name=deployment_name,\n desired_status="READY",\n workspace="default",\n timeout=remaining,\n check_gateway=True,\n ):\n raise TimeoutError("Inference gateway did not become ready")\n return dep\n if dep.status in ("FAILED", "ERROR", "TERMINATED", "LOST"):\n raise RuntimeError(f"Deployment failed with status: {dep.status}")\n if elapsed > timeout:\n raise TimeoutError(f"Deployment timeout after {timeout_minutes} minutes")\n time.sleep(15)\n\n\ntry:\n dep_status = wait_for_deployment(BASELINE_DEPLOYMENT_NAME)\n assert dep_status.status == "READY"\nexcept Exception:\n # Free GPUs if readiness fails before the later baseline-cleanup cell runs.\n try:\n client.inference.deployments.delete(name=BASELINE_DEPLOYMENT_NAME, workspace="default")\n client.inference.deployment_configs.delete(name=BASELINE_DEPLOYMENT_CONFIG, workspace="default")\n except Exception as cleanup_error:\n print(f"Baseline cleanup after readiness failure also failed: {cleanup_error}")\n raise\n" }, { "type": "markdown", diff --git a/docs/fern/components/notebooks/distillation-customization-job.ts b/docs/fern/components/notebooks/distillation-customization-job.ts index d33d81b925..c68567c3c7 100644 --- a/docs/fern/components/notebooks/distillation-customization-job.ts +++ b/docs/fern/components/notebooks/distillation-customization-job.ts @@ -95,9 +95,9 @@ export default { cells: [ }, { "type": "code", - "source": "def wait_for_deployment(deployment_name: str, timeout_minutes: int = 30):\n \"\"\"Poll deployment until ready.\"\"\"\n start = time.time()\n timeout = timeout_minutes * 60\n while True:\n dep = client.inference.deployments.retrieve(name=deployment_name, workspace=\"default\")\n elapsed = time.time() - start\n clear_output(wait=True)\n print(f\"Deployment: {deployment_name}\")\n print(f\"Status: {dep.status}\")\n print(f\"Elapsed: {int(elapsed // 60)}m {int(elapsed % 60)}s\")\n\n if dep.status == \"READY\":\n print(\"\\nDeployment is ready!\")\n if not client.models.wait_for_gateway(deployment_name, workspace=\"default\", timeout=60):\n raise RuntimeError(\"Inference gateway did not become ready\")\n return dep\n if dep.status in (\"FAILED\", \"ERROR\", \"TERMINATED\", \"LOST\"):\n raise RuntimeError(f\"Deployment failed with status: {dep.status}\")\n if elapsed > timeout:\n raise TimeoutError(f\"Deployment timeout after {timeout_minutes} minutes\")\n time.sleep(15)\n\n\ndep_status = wait_for_deployment(BASELINE_DEPLOYMENT_NAME)\nassert dep_status.status == \"READY\"", + "source": "def wait_for_deployment(deployment_name: str, timeout_minutes: int = 30):\n \"\"\"Poll deployment until ready.\"\"\"\n start = time.time()\n timeout = timeout_minutes * 60\n while True:\n dep = client.inference.deployments.retrieve(name=deployment_name, workspace=\"default\")\n elapsed = time.time() - start\n clear_output(wait=True)\n print(f\"Deployment: {deployment_name}\")\n print(f\"Status: {dep.status}\")\n print(f\"Elapsed: {int(elapsed // 60)}m {int(elapsed % 60)}s\")\n\n if dep.status == \"READY\":\n print(\"\\nDeployment is ready!\")\n remaining = int(timeout - elapsed)\n if remaining <= 0:\n raise TimeoutError(f\"Deployment timeout after {timeout_minutes} minutes\")\n if not client.models.wait_for_status(\n deployment_name=deployment_name,\n desired_status=\"READY\",\n workspace=\"default\",\n timeout=remaining,\n check_gateway=True,\n ):\n raise TimeoutError(\"Inference gateway did not become ready\")\n return dep\n if dep.status in (\"FAILED\", \"ERROR\", \"TERMINATED\", \"LOST\"):\n raise RuntimeError(f\"Deployment failed with status: {dep.status}\")\n if elapsed > timeout:\n raise TimeoutError(f\"Deployment timeout after {timeout_minutes} minutes\")\n time.sleep(15)\n\n\ntry:\n dep_status = wait_for_deployment(BASELINE_DEPLOYMENT_NAME)\n assert dep_status.status == \"READY\"\nexcept Exception:\n # Free GPUs if readiness fails before the later baseline-cleanup cell runs.\n try:\n client.inference.deployments.delete(name=BASELINE_DEPLOYMENT_NAME, workspace=\"default\")\n client.inference.deployment_configs.delete(name=BASELINE_DEPLOYMENT_CONFIG, workspace=\"default\")\n except Exception as cleanup_error:\n print(f\"Baseline cleanup after readiness failure also failed: {cleanup_error}\")\n raise", "language": "python", - "source_html": "def wait_for_deployment(deployment_name: str, timeout_minutes: int = 30):\n """Poll deployment until ready."""\n start = time.time()\n timeout = timeout_minutes * 60\n while True:\n dep = client.inference.deployments.retrieve(name=deployment_name, workspace="default")\n elapsed = time.time() - start\n clear_output(wait=True)\n print(f"Deployment: {deployment_name}")\n print(f"Status: {dep.status}")\n print(f"Elapsed: {int(elapsed // 60)}m {int(elapsed % 60)}s")\n\n if dep.status == "READY":\n print("\\nDeployment is ready!")\n if not client.models.wait_for_gateway(deployment_name, workspace="default", timeout=60):\n raise RuntimeError("Inference gateway did not become ready")\n return dep\n if dep.status in ("FAILED", "ERROR", "TERMINATED", "LOST"):\n raise RuntimeError(f"Deployment failed with status: {dep.status}")\n if elapsed > timeout:\n raise TimeoutError(f"Deployment timeout after {timeout_minutes} minutes")\n time.sleep(15)\n\n\ndep_status = wait_for_deployment(BASELINE_DEPLOYMENT_NAME)\nassert dep_status.status == "READY"\n" + "source_html": "def wait_for_deployment(deployment_name: str, timeout_minutes: int = 30):\n """Poll deployment until ready."""\n start = time.time()\n timeout = timeout_minutes * 60\n while True:\n dep = client.inference.deployments.retrieve(name=deployment_name, workspace="default")\n elapsed = time.time() - start\n clear_output(wait=True)\n print(f"Deployment: {deployment_name}")\n print(f"Status: {dep.status}")\n print(f"Elapsed: {int(elapsed // 60)}m {int(elapsed % 60)}s")\n\n if dep.status == "READY":\n print("\\nDeployment is ready!")\n remaining = int(timeout - elapsed)\n if remaining <= 0:\n raise TimeoutError(f"Deployment timeout after {timeout_minutes} minutes")\n if not client.models.wait_for_status(\n deployment_name=deployment_name,\n desired_status="READY",\n workspace="default",\n timeout=remaining,\n check_gateway=True,\n ):\n raise TimeoutError("Inference gateway did not become ready")\n return dep\n if dep.status in ("FAILED", "ERROR", "TERMINATED", "LOST"):\n raise RuntimeError(f"Deployment failed with status: {dep.status}")\n if elapsed > timeout:\n raise TimeoutError(f"Deployment timeout after {timeout_minutes} minutes")\n time.sleep(15)\n\n\ntry:\n dep_status = wait_for_deployment(BASELINE_DEPLOYMENT_NAME)\n assert dep_status.status == "READY"\nexcept Exception:\n # Free GPUs if readiness fails before the later baseline-cleanup cell runs.\n try:\n client.inference.deployments.delete(name=BASELINE_DEPLOYMENT_NAME, workspace="default")\n client.inference.deployment_configs.delete(name=BASELINE_DEPLOYMENT_CONFIG, workspace="default")\n except Exception as cleanup_error:\n print(f"Baseline cleanup after readiness failure also failed: {cleanup_error}")\n raise\n" }, { "type": "markdown", diff --git a/docs/fern/components/notebooks/dpo-customization-job.json b/docs/fern/components/notebooks/dpo-customization-job.json index bf8d0970a4..ab69dab155 100644 --- a/docs/fern/components/notebooks/dpo-customization-job.json +++ b/docs/fern/components/notebooks/dpo-customization-job.json @@ -70,9 +70,9 @@ }, { "type": "code", - "source": "HF_TOKEN = os.getenv(\"HF_TOKEN\")\n\ndef create_or_get_secret(name: str, value: str | None, label: str) -> PlatformSecretResponse | None:\n if not value:\n print(f\"{label} is not set - skipping secret (gated model downloads will fail without it)\")\n return None\n try:\n secret = sdk.secrets.create(name=name, workspace=\"default\", value=value)\n print(f\"Created secret: {name}\")\n return secret\n except ConflictError:\n print(f\"Secret '{name}' already exists, continuing...\")\n return sdk.secrets.retrieve(name=name, workspace=\"default\")\n\n\nhf_secret = create_or_get_secret(\"hf-token\", HF_TOKEN, \"HF_TOKEN\")", + "source": "HF_TOKEN = os.getenv(\"HF_TOKEN\")\nif not HF_TOKEN:\n raise RuntimeError(\"Set HF_TOKEN before running this tutorial.\")\n\ndef create_or_get_secret(name: str, value: str, label: str) -> PlatformSecretResponse:\n try:\n secret = sdk.secrets.create(name=name, workspace=\"default\", value=value)\n print(f\"Created secret: {name}\")\n return secret\n except ConflictError:\n print(f\"Secret '{name}' already exists, continuing...\")\n return sdk.secrets.retrieve(name=name, workspace=\"default\")\n\n\nhf_secret = create_or_get_secret(\"hf-token\", HF_TOKEN, \"HF_TOKEN\")", "language": "python", - "source_html": "HF_TOKEN = os.getenv("HF_TOKEN")\n\ndef create_or_get_secret(name: str, value: str | None, label: str) -> PlatformSecretResponse | None:\n if not value:\n print(f"{label} is not set - skipping secret (gated model downloads will fail without it)")\n return None\n try:\n secret = sdk.secrets.create(name=name, workspace="default", value=value)\n print(f"Created secret: {name}")\n return secret\n except ConflictError:\n print(f"Secret '{name}' already exists, continuing...")\n return sdk.secrets.retrieve(name=name, workspace="default")\n\n\nhf_secret = create_or_get_secret("hf-token", HF_TOKEN, "HF_TOKEN")\n" + "source_html": "HF_TOKEN = os.getenv("HF_TOKEN")\nif not HF_TOKEN:\n raise RuntimeError("Set HF_TOKEN before running this tutorial.")\n\ndef create_or_get_secret(name: str, value: str, label: str) -> PlatformSecretResponse:\n try:\n secret = sdk.secrets.create(name=name, workspace="default", value=value)\n print(f"Created secret: {name}")\n return secret\n except ConflictError:\n print(f"Secret '{name}' already exists, continuing...")\n return sdk.secrets.retrieve(name=name, workspace="default")\n\n\nhf_secret = create_or_get_secret("hf-token", HF_TOKEN, "HF_TOKEN")\n" }, { "type": "markdown", @@ -81,9 +81,9 @@ }, { "type": "code", - "source": "HF_REPO_ID = \"meta-llama/Llama-3.2-1B-Instruct\"\nMODEL_NAME = \"llama-3-2-1b-instruct\"\n\nstorage = HuggingfaceStorageConfigParam(type=\"huggingface\", repo_id=HF_REPO_ID, repo_type=\"model\")\nif hf_secret:\n storage[\"token_secret\"] = hf_secret.name\n\ntry:\n base_model_fs = sdk.files.filesets.create(\n workspace=\"default\", name=MODEL_NAME, description=\"Llama 3.2 1B Instruct base model\", storage=storage\n )\n print(f\"Created base model fileset: {MODEL_NAME}\")\nexcept ConflictError:\n base_model_fs = sdk.files.filesets.retrieve(workspace=\"default\", name=MODEL_NAME)\n print(\"Base model fileset already exists.\")\n\ntry:\n base_model = sdk.models.create(workspace=\"default\", name=MODEL_NAME, fileset=f\"default/{MODEL_NAME}\")\nexcept ConflictError:\n base_model = sdk.models.retrieve(workspace=\"default\", name=MODEL_NAME)\n\nprint(f\"Base model fileset: fileset://default/{base_model.name}\")\n\n# Wait for the ModelSpec to be inferred from the checkpoint.\ncheck = max_wait_time_checker(600, \"Model spec\")\nwhile not base_model.spec:\n check()\n time.sleep(10)\n base_model = sdk.models.retrieve(workspace=\"default\", name=MODEL_NAME)\nprint(\"Model spec ready\")", + "source": "HF_REPO_ID = \"meta-llama/Llama-3.2-1B-Instruct\"\nMODEL_NAME = \"llama-3-2-1b-instruct\"\n\nstorage = HuggingfaceStorageConfigParam(\n type=\"huggingface\",\n repo_id=HF_REPO_ID,\n repo_type=\"model\",\n token_secret=hf_secret.name,\n)\n\ntry:\n base_model_fs = sdk.files.filesets.create(\n workspace=\"default\", name=MODEL_NAME, description=\"Llama 3.2 1B Instruct base model\", storage=storage\n )\n print(f\"Created base model fileset: {MODEL_NAME}\")\nexcept ConflictError:\n base_model_fs = sdk.files.filesets.retrieve(workspace=\"default\", name=MODEL_NAME)\n print(\"Base model fileset already exists.\")\n\ntry:\n base_model = sdk.models.create(workspace=\"default\", name=MODEL_NAME, fileset=f\"default/{MODEL_NAME}\")\nexcept ConflictError:\n base_model = sdk.models.retrieve(workspace=\"default\", name=MODEL_NAME)\n\nprint(f\"Base model fileset: fileset://default/{base_model.name}\")\n\n# Wait for the ModelSpec to be inferred from the checkpoint.\ncheck = max_wait_time_checker(600, \"Model spec\")\nwhile not base_model.spec:\n check()\n time.sleep(10)\n base_model = sdk.models.retrieve(workspace=\"default\", name=MODEL_NAME)\nprint(\"Model spec ready\")", "language": "python", - "source_html": "HF_REPO_ID = "meta-llama/Llama-3.2-1B-Instruct"\nMODEL_NAME = "llama-3-2-1b-instruct"\n\nstorage = HuggingfaceStorageConfigParam(type="huggingface", repo_id=HF_REPO_ID, repo_type="model")\nif hf_secret:\n storage["token_secret"] = hf_secret.name\n\ntry:\n base_model_fs = sdk.files.filesets.create(\n workspace="default", name=MODEL_NAME, description="Llama 3.2 1B Instruct base model", storage=storage\n )\n print(f"Created base model fileset: {MODEL_NAME}")\nexcept ConflictError:\n base_model_fs = sdk.files.filesets.retrieve(workspace="default", name=MODEL_NAME)\n print("Base model fileset already exists.")\n\ntry:\n base_model = sdk.models.create(workspace="default", name=MODEL_NAME, fileset=f"default/{MODEL_NAME}")\nexcept ConflictError:\n base_model = sdk.models.retrieve(workspace="default", name=MODEL_NAME)\n\nprint(f"Base model fileset: fileset://default/{base_model.name}")\n\n# Wait for the ModelSpec to be inferred from the checkpoint.\ncheck = max_wait_time_checker(600, "Model spec")\nwhile not base_model.spec:\n check()\n time.sleep(10)\n base_model = sdk.models.retrieve(workspace="default", name=MODEL_NAME)\nprint("Model spec ready")\n" + "source_html": "HF_REPO_ID = "meta-llama/Llama-3.2-1B-Instruct"\nMODEL_NAME = "llama-3-2-1b-instruct"\n\nstorage = HuggingfaceStorageConfigParam(\n type="huggingface",\n repo_id=HF_REPO_ID,\n repo_type="model",\n token_secret=hf_secret.name,\n)\n\ntry:\n base_model_fs = sdk.files.filesets.create(\n workspace="default", name=MODEL_NAME, description="Llama 3.2 1B Instruct base model", storage=storage\n )\n print(f"Created base model fileset: {MODEL_NAME}")\nexcept ConflictError:\n base_model_fs = sdk.files.filesets.retrieve(workspace="default", name=MODEL_NAME)\n print("Base model fileset already exists.")\n\ntry:\n base_model = sdk.models.create(workspace="default", name=MODEL_NAME, fileset=f"default/{MODEL_NAME}")\nexcept ConflictError:\n base_model = sdk.models.retrieve(workspace="default", name=MODEL_NAME)\n\nprint(f"Base model fileset: fileset://default/{base_model.name}")\n\n# Wait for the ModelSpec to be inferred from the checkpoint.\ncheck = max_wait_time_checker(600, "Model spec")\nwhile not base_model.spec:\n check()\n time.sleep(10)\n base_model = sdk.models.retrieve(workspace="default", name=MODEL_NAME)\nprint("Model spec ready")\n" }, { "type": "markdown", @@ -148,8 +148,8 @@ }, { "type": "markdown", - "source": "## Conclusion\n\nYou aligned a base model with **DPO** on the NeMo Platform using the `rl` backend:\n\n- Uploaded a HelpSteer3 preference dataset **as-is** (the platform detects the schema natively).\n- Submitted a full-weight DPO job that ran on a Ray cluster via the Kubernetes executor.\n- Registered the output as a full model entity and (optionally) deployed it for inference.\n\n**Next steps:** tune the alignment strength with `ref_policy_kl_penalty` (β), add `sft_loss_weight` to anchor the policy to the chosen responses, enable `activation_checkpointing` for memory headroom, or scale up with `parallelism`. See the `nemo-customizer` skill's `references/hyperparameters.md` (section NeMo-RL (DPO)) for the full knob reference.", - "source_html": "

Conclusion

\n

You aligned a base model with DPO on the NeMo Platform using the rl backend:

\n
    \n
  • Uploaded a HelpSteer3 preference dataset as-is (the platform detects the schema natively).
    • \n
  • Submitted a full-weight DPO job that ran on a Ray cluster via the Kubernetes executor.
    • \n
  • Registered the output as a full model entity and (optionally) deployed it for inference.
    • \n
\n

Next steps: tune the alignment strength with ref_policy_kl_penalty (β), add sft_loss_weight to anchor the policy to the chosen responses, enable activation_checkpointing for memory headroom, or scale up with parallelism. See the nemo-customizer skill's references/hyperparameters.md (section NeMo-RL (DPO)) for the full knob reference.

\n" + "source": "## Conclusion\n\nYou aligned a base model with **DPO** on the NeMo Platform using the `rl` backend:\n\n- Uploaded a HelpSteer3 preference dataset **as-is** (the platform detects the schema natively).\n- Submitted a full-weight DPO job that ran on a Ray cluster via the Kubernetes executor.\n- Registered the output as a full model entity and (optionally) deployed it for inference.\n\n**Next steps:** tune the alignment strength with `ref_policy_kl_penalty` (β), add `sft_loss_weight` to anchor the policy to the chosen responses, enable `activation_checkpointing` for memory headroom, or scale up with `parallelism`. See the [Training Configuration](/documentation/customizer-reference/manage-customization-jobs/training-configuration) reference for the full hyperparameter set.", + "source_html": "

Conclusion

\n

You aligned a base model with DPO on the NeMo Platform using the rl backend:

\n
    \n
  • Uploaded a HelpSteer3 preference dataset as-is (the platform detects the schema natively).
    • \n
  • Submitted a full-weight DPO job that ran on a Ray cluster via the Kubernetes executor.
    • \n
  • Registered the output as a full model entity and (optionally) deployed it for inference.
    • \n
\n

Next steps: tune the alignment strength with ref_policy_kl_penalty (β), add sft_loss_weight to anchor the policy to the chosen responses, enable activation_checkpointing for memory headroom, or scale up with parallelism. See the Training Configuration reference for the full hyperparameter set.

\n" } ] } \ No newline at end of file diff --git a/docs/fern/components/notebooks/dpo-customization-job.ts b/docs/fern/components/notebooks/dpo-customization-job.ts index ffbdb4dfb0..fc8e558903 100644 --- a/docs/fern/components/notebooks/dpo-customization-job.ts +++ b/docs/fern/components/notebooks/dpo-customization-job.ts @@ -75,9 +75,9 @@ export default { cells: [ }, { "type": "code", - "source": "HF_TOKEN = os.getenv(\"HF_TOKEN\")\n\ndef create_or_get_secret(name: str, value: str | None, label: str) -> PlatformSecretResponse | None:\n if not value:\n print(f\"{label} is not set - skipping secret (gated model downloads will fail without it)\")\n return None\n try:\n secret = sdk.secrets.create(name=name, workspace=\"default\", value=value)\n print(f\"Created secret: {name}\")\n return secret\n except ConflictError:\n print(f\"Secret '{name}' already exists, continuing...\")\n return sdk.secrets.retrieve(name=name, workspace=\"default\")\n\n\nhf_secret = create_or_get_secret(\"hf-token\", HF_TOKEN, \"HF_TOKEN\")", + "source": "HF_TOKEN = os.getenv(\"HF_TOKEN\")\nif not HF_TOKEN:\n raise RuntimeError(\"Set HF_TOKEN before running this tutorial.\")\n\ndef create_or_get_secret(name: str, value: str, label: str) -> PlatformSecretResponse:\n try:\n secret = sdk.secrets.create(name=name, workspace=\"default\", value=value)\n print(f\"Created secret: {name}\")\n return secret\n except ConflictError:\n print(f\"Secret '{name}' already exists, continuing...\")\n return sdk.secrets.retrieve(name=name, workspace=\"default\")\n\n\nhf_secret = create_or_get_secret(\"hf-token\", HF_TOKEN, \"HF_TOKEN\")", "language": "python", - "source_html": "HF_TOKEN = os.getenv("HF_TOKEN")\n\ndef create_or_get_secret(name: str, value: str | None, label: str) -> PlatformSecretResponse | None:\n if not value:\n print(f"{label} is not set - skipping secret (gated model downloads will fail without it)")\n return None\n try:\n secret = sdk.secrets.create(name=name, workspace="default", value=value)\n print(f"Created secret: {name}")\n return secret\n except ConflictError:\n print(f"Secret '{name}' already exists, continuing...")\n return sdk.secrets.retrieve(name=name, workspace="default")\n\n\nhf_secret = create_or_get_secret("hf-token", HF_TOKEN, "HF_TOKEN")\n" + "source_html": "HF_TOKEN = os.getenv("HF_TOKEN")\nif not HF_TOKEN:\n raise RuntimeError("Set HF_TOKEN before running this tutorial.")\n\ndef create_or_get_secret(name: str, value: str, label: str) -> PlatformSecretResponse:\n try:\n secret = sdk.secrets.create(name=name, workspace="default", value=value)\n print(f"Created secret: {name}")\n return secret\n except ConflictError:\n print(f"Secret '{name}' already exists, continuing...")\n return sdk.secrets.retrieve(name=name, workspace="default")\n\n\nhf_secret = create_or_get_secret("hf-token", HF_TOKEN, "HF_TOKEN")\n" }, { "type": "markdown", @@ -86,9 +86,9 @@ export default { cells: [ }, { "type": "code", - "source": "HF_REPO_ID = \"meta-llama/Llama-3.2-1B-Instruct\"\nMODEL_NAME = \"llama-3-2-1b-instruct\"\n\nstorage = HuggingfaceStorageConfigParam(type=\"huggingface\", repo_id=HF_REPO_ID, repo_type=\"model\")\nif hf_secret:\n storage[\"token_secret\"] = hf_secret.name\n\ntry:\n base_model_fs = sdk.files.filesets.create(\n workspace=\"default\", name=MODEL_NAME, description=\"Llama 3.2 1B Instruct base model\", storage=storage\n )\n print(f\"Created base model fileset: {MODEL_NAME}\")\nexcept ConflictError:\n base_model_fs = sdk.files.filesets.retrieve(workspace=\"default\", name=MODEL_NAME)\n print(\"Base model fileset already exists.\")\n\ntry:\n base_model = sdk.models.create(workspace=\"default\", name=MODEL_NAME, fileset=f\"default/{MODEL_NAME}\")\nexcept ConflictError:\n base_model = sdk.models.retrieve(workspace=\"default\", name=MODEL_NAME)\n\nprint(f\"Base model fileset: fileset://default/{base_model.name}\")\n\n# Wait for the ModelSpec to be inferred from the checkpoint.\ncheck = max_wait_time_checker(600, \"Model spec\")\nwhile not base_model.spec:\n check()\n time.sleep(10)\n base_model = sdk.models.retrieve(workspace=\"default\", name=MODEL_NAME)\nprint(\"Model spec ready\")", + "source": "HF_REPO_ID = \"meta-llama/Llama-3.2-1B-Instruct\"\nMODEL_NAME = \"llama-3-2-1b-instruct\"\n\nstorage = HuggingfaceStorageConfigParam(\n type=\"huggingface\",\n repo_id=HF_REPO_ID,\n repo_type=\"model\",\n token_secret=hf_secret.name,\n)\n\ntry:\n base_model_fs = sdk.files.filesets.create(\n workspace=\"default\", name=MODEL_NAME, description=\"Llama 3.2 1B Instruct base model\", storage=storage\n )\n print(f\"Created base model fileset: {MODEL_NAME}\")\nexcept ConflictError:\n base_model_fs = sdk.files.filesets.retrieve(workspace=\"default\", name=MODEL_NAME)\n print(\"Base model fileset already exists.\")\n\ntry:\n base_model = sdk.models.create(workspace=\"default\", name=MODEL_NAME, fileset=f\"default/{MODEL_NAME}\")\nexcept ConflictError:\n base_model = sdk.models.retrieve(workspace=\"default\", name=MODEL_NAME)\n\nprint(f\"Base model fileset: fileset://default/{base_model.name}\")\n\n# Wait for the ModelSpec to be inferred from the checkpoint.\ncheck = max_wait_time_checker(600, \"Model spec\")\nwhile not base_model.spec:\n check()\n time.sleep(10)\n base_model = sdk.models.retrieve(workspace=\"default\", name=MODEL_NAME)\nprint(\"Model spec ready\")", "language": "python", - "source_html": "HF_REPO_ID = "meta-llama/Llama-3.2-1B-Instruct"\nMODEL_NAME = "llama-3-2-1b-instruct"\n\nstorage = HuggingfaceStorageConfigParam(type="huggingface", repo_id=HF_REPO_ID, repo_type="model")\nif hf_secret:\n storage["token_secret"] = hf_secret.name\n\ntry:\n base_model_fs = sdk.files.filesets.create(\n workspace="default", name=MODEL_NAME, description="Llama 3.2 1B Instruct base model", storage=storage\n )\n print(f"Created base model fileset: {MODEL_NAME}")\nexcept ConflictError:\n base_model_fs = sdk.files.filesets.retrieve(workspace="default", name=MODEL_NAME)\n print("Base model fileset already exists.")\n\ntry:\n base_model = sdk.models.create(workspace="default", name=MODEL_NAME, fileset=f"default/{MODEL_NAME}")\nexcept ConflictError:\n base_model = sdk.models.retrieve(workspace="default", name=MODEL_NAME)\n\nprint(f"Base model fileset: fileset://default/{base_model.name}")\n\n# Wait for the ModelSpec to be inferred from the checkpoint.\ncheck = max_wait_time_checker(600, "Model spec")\nwhile not base_model.spec:\n check()\n time.sleep(10)\n base_model = sdk.models.retrieve(workspace="default", name=MODEL_NAME)\nprint("Model spec ready")\n" + "source_html": "HF_REPO_ID = "meta-llama/Llama-3.2-1B-Instruct"\nMODEL_NAME = "llama-3-2-1b-instruct"\n\nstorage = HuggingfaceStorageConfigParam(\n type="huggingface",\n repo_id=HF_REPO_ID,\n repo_type="model",\n token_secret=hf_secret.name,\n)\n\ntry:\n base_model_fs = sdk.files.filesets.create(\n workspace="default", name=MODEL_NAME, description="Llama 3.2 1B Instruct base model", storage=storage\n )\n print(f"Created base model fileset: {MODEL_NAME}")\nexcept ConflictError:\n base_model_fs = sdk.files.filesets.retrieve(workspace="default", name=MODEL_NAME)\n print("Base model fileset already exists.")\n\ntry:\n base_model = sdk.models.create(workspace="default", name=MODEL_NAME, fileset=f"default/{MODEL_NAME}")\nexcept ConflictError:\n base_model = sdk.models.retrieve(workspace="default", name=MODEL_NAME)\n\nprint(f"Base model fileset: fileset://default/{base_model.name}")\n\n# Wait for the ModelSpec to be inferred from the checkpoint.\ncheck = max_wait_time_checker(600, "Model spec")\nwhile not base_model.spec:\n check()\n time.sleep(10)\n base_model = sdk.models.retrieve(workspace="default", name=MODEL_NAME)\nprint("Model spec ready")\n" }, { "type": "markdown", @@ -153,7 +153,7 @@ export default { cells: [ }, { "type": "markdown", - "source": "## Conclusion\n\nYou aligned a base model with **DPO** on the NeMo Platform using the `rl` backend:\n\n- Uploaded a HelpSteer3 preference dataset **as-is** (the platform detects the schema natively).\n- Submitted a full-weight DPO job that ran on a Ray cluster via the Kubernetes executor.\n- Registered the output as a full model entity and (optionally) deployed it for inference.\n\n**Next steps:** tune the alignment strength with `ref_policy_kl_penalty` (β), add `sft_loss_weight` to anchor the policy to the chosen responses, enable `activation_checkpointing` for memory headroom, or scale up with `parallelism`. See the `nemo-customizer` skill's `references/hyperparameters.md` (section NeMo-RL (DPO)) for the full knob reference.", - "source_html": "

Conclusion

\n

You aligned a base model with DPO on the NeMo Platform using the rl backend:

\n
    \n
  • Uploaded a HelpSteer3 preference dataset as-is (the platform detects the schema natively).
    • \n
  • Submitted a full-weight DPO job that ran on a Ray cluster via the Kubernetes executor.
    • \n
  • Registered the output as a full model entity and (optionally) deployed it for inference.
    • \n
\n

Next steps: tune the alignment strength with ref_policy_kl_penalty (β), add sft_loss_weight to anchor the policy to the chosen responses, enable activation_checkpointing for memory headroom, or scale up with parallelism. See the nemo-customizer skill's references/hyperparameters.md (section NeMo-RL (DPO)) for the full knob reference.

\n" + "source": "## Conclusion\n\nYou aligned a base model with **DPO** on the NeMo Platform using the `rl` backend:\n\n- Uploaded a HelpSteer3 preference dataset **as-is** (the platform detects the schema natively).\n- Submitted a full-weight DPO job that ran on a Ray cluster via the Kubernetes executor.\n- Registered the output as a full model entity and (optionally) deployed it for inference.\n\n**Next steps:** tune the alignment strength with `ref_policy_kl_penalty` (β), add `sft_loss_weight` to anchor the policy to the chosen responses, enable `activation_checkpointing` for memory headroom, or scale up with `parallelism`. See the [Training Configuration](/documentation/customizer-reference/manage-customization-jobs/training-configuration) reference for the full hyperparameter set.", + "source_html": "

Conclusion

\n

You aligned a base model with DPO on the NeMo Platform using the rl backend:

\n
    \n
  • Uploaded a HelpSteer3 preference dataset as-is (the platform detects the schema natively).
    • \n
  • Submitted a full-weight DPO job that ran on a Ray cluster via the Kubernetes executor.
    • \n
  • Registered the output as a full model entity and (optionally) deployed it for inference.
    • \n
\n

Next steps: tune the alignment strength with ref_policy_kl_penalty (β), add sft_loss_weight to anchor the policy to the chosen responses, enable activation_checkpointing for memory headroom, or scale up with parallelism. See the Training Configuration reference for the full hyperparameter set.

\n" } ] }; diff --git a/docs/fern/components/notebooks/embedding-customization-job.json b/docs/fern/components/notebooks/embedding-customization-job.json index 6f9b907a47..82e62a10c9 100644 --- a/docs/fern/components/notebooks/embedding-customization-job.json +++ b/docs/fern/components/notebooks/embedding-customization-job.json @@ -107,9 +107,9 @@ }, { "type": "code", - "source": "import time\nfrom nemo_platform.types.files import HuggingfaceStorageConfigParam\n\nHF_REPO_ID = \"nvidia/llama-nemotron-embed-1b-v2\"\nMODEL_NAME = \"nv-nemotron-embed-1b-base\"\n\nstorage = HuggingfaceStorageConfigParam(\n type=\"huggingface\",\n repo_id=HF_REPO_ID,\n repo_type=\"model\",\n)\nif hf_secret:\n storage[\"token_secret\"] = hf_secret.name\n\ntry:\n base_model_fs = client.files.filesets.create(\n workspace=\"default\",\n name=MODEL_NAME,\n description=\"NVIDIA Llama Nemotron Embed 1B v2 embedding model\",\n storage=storage,\n )\nexcept ConflictError as e:\n print(f\"Base model fileset already exists. Skipping creation.\")\n base_model_fs = client.files.filesets.retrieve(\n workspace=\"default\",\n name=MODEL_NAME,\n )\n\n# Create Model Entity referencing the FileSet\ntry:\n base_model = client.models.create(\n workspace=\"default\",\n name=MODEL_NAME,\n fileset=f\"default/{MODEL_NAME}\",\n trust_remote_code=True,\n )\n print(f\"Created Model Entity: {MODEL_NAME}\")\nexcept ConflictError:\n print(f\"Base model already exists. Updating fileset if different.\")\n base_model = client.models.update(\n workspace=\"default\",\n name=MODEL_NAME,\n fileset=f\"default/{MODEL_NAME}\",\n trust_remote_code=True,\n )\n\nprint(f\"\\nBase model fileset: fileset://default/{base_model.name}\")\nprint(\"\\nBase model files:\")\nprint(json.dumps([f.model_dump() for f in client.files.list(fileset=MODEL_NAME, workspace=\"default\").data], indent=2))\n\n# Wait for ModelSpec to be populated from the checkpoint\nprint(\"\\nWaiting for ModelSpec to be populated...\")\nSPEC_TIMEOUT_SECONDS = 120\nspec_start = time.time()\nwhile not base_model.spec:\n if time.time() - spec_start > SPEC_TIMEOUT_SECONDS:\n raise TimeoutError(f\"ModelSpec not populated within {SPEC_TIMEOUT_SECONDS} seconds\")\n time.sleep(2)\n base_model = client.models.retrieve(\n workspace=\"default\",\n name=MODEL_NAME,\n )\n\nprint(f\"ModelSpec populated: {base_model.spec}\")", + "source": "import time\nfrom nemo_platform.types.files import HuggingfaceStorageConfigParam\n\nHF_REPO_ID = \"nvidia/llama-nemotron-embed-1b-v2\"\nMODEL_NAME = \"nv-nemotron-embed-1b-base\"\n\nstorage_kwargs = {\n \"type\": \"huggingface\",\n \"repo_id\": HF_REPO_ID,\n \"repo_type\": \"model\",\n}\nif hf_secret:\n storage_kwargs[\"token_secret\"] = hf_secret.name\nstorage = HuggingfaceStorageConfigParam(**storage_kwargs)\n\ntry:\n base_model_fs = client.files.filesets.create(\n workspace=\"default\",\n name=MODEL_NAME,\n description=\"NVIDIA Llama Nemotron Embed 1B v2 embedding model\",\n storage=storage,\n )\nexcept ConflictError as e:\n print(f\"Base model fileset already exists. Skipping creation.\")\n base_model_fs = client.files.filesets.retrieve(\n workspace=\"default\",\n name=MODEL_NAME,\n )\n\n# Create Model Entity referencing the FileSet\ntry:\n base_model = client.models.create(\n workspace=\"default\",\n name=MODEL_NAME,\n fileset=f\"default/{MODEL_NAME}\",\n trust_remote_code=True,\n )\n print(f\"Created Model Entity: {MODEL_NAME}\")\nexcept ConflictError:\n print(f\"Base model already exists. Updating fileset if different.\")\n base_model = client.models.update(\n workspace=\"default\",\n name=MODEL_NAME,\n fileset=f\"default/{MODEL_NAME}\",\n trust_remote_code=True,\n )\n\nprint(f\"\\nBase model fileset: fileset://default/{base_model.name}\")\nprint(\"\\nBase model files:\")\nprint(json.dumps([f.model_dump() for f in client.files.list(fileset=MODEL_NAME, workspace=\"default\").data], indent=2))\n\n# Wait for ModelSpec to be populated from the checkpoint\nprint(\"\\nWaiting for ModelSpec to be populated...\")\nSPEC_TIMEOUT_SECONDS = 120\nspec_start = time.time()\nwhile not base_model.spec:\n if time.time() - spec_start > SPEC_TIMEOUT_SECONDS:\n raise TimeoutError(f\"ModelSpec not populated within {SPEC_TIMEOUT_SECONDS} seconds\")\n time.sleep(2)\n base_model = client.models.retrieve(\n workspace=\"default\",\n name=MODEL_NAME,\n )\n\nprint(f\"ModelSpec populated: {base_model.spec}\")", "language": "python", - "source_html": "import time\nfrom nemo_platform.types.files import HuggingfaceStorageConfigParam\n\nHF_REPO_ID = "nvidia/llama-nemotron-embed-1b-v2"\nMODEL_NAME = "nv-nemotron-embed-1b-base"\n\nstorage = HuggingfaceStorageConfigParam(\n type="huggingface",\n repo_id=HF_REPO_ID,\n repo_type="model",\n)\nif hf_secret:\n storage["token_secret"] = hf_secret.name\n\ntry:\n base_model_fs = client.files.filesets.create(\n workspace="default",\n name=MODEL_NAME,\n description="NVIDIA Llama Nemotron Embed 1B v2 embedding model",\n storage=storage,\n )\nexcept ConflictError as e:\n print(f"Base model fileset already exists. Skipping creation.")\n base_model_fs = client.files.filesets.retrieve(\n workspace="default",\n name=MODEL_NAME,\n )\n\n# Create Model Entity referencing the FileSet\ntry:\n base_model = client.models.create(\n workspace="default",\n name=MODEL_NAME,\n fileset=f"default/{MODEL_NAME}",\n trust_remote_code=True,\n )\n print(f"Created Model Entity: {MODEL_NAME}")\nexcept ConflictError:\n print(f"Base model already exists. Updating fileset if different.")\n base_model = client.models.update(\n workspace="default",\n name=MODEL_NAME,\n fileset=f"default/{MODEL_NAME}",\n trust_remote_code=True,\n )\n\nprint(f"\\nBase model fileset: fileset://default/{base_model.name}")\nprint("\\nBase model files:")\nprint(json.dumps([f.model_dump() for f in client.files.list(fileset=MODEL_NAME, workspace="default").data], indent=2))\n\n# Wait for ModelSpec to be populated from the checkpoint\nprint("\\nWaiting for ModelSpec to be populated...")\nSPEC_TIMEOUT_SECONDS = 120\nspec_start = time.time()\nwhile not base_model.spec:\n if time.time() - spec_start > SPEC_TIMEOUT_SECONDS:\n raise TimeoutError(f"ModelSpec not populated within {SPEC_TIMEOUT_SECONDS} seconds")\n time.sleep(2)\n base_model = client.models.retrieve(\n workspace="default",\n name=MODEL_NAME,\n )\n\nprint(f"ModelSpec populated: {base_model.spec}")\n" + "source_html": "import time\nfrom nemo_platform.types.files import HuggingfaceStorageConfigParam\n\nHF_REPO_ID = "nvidia/llama-nemotron-embed-1b-v2"\nMODEL_NAME = "nv-nemotron-embed-1b-base"\n\nstorage_kwargs = {\n "type": "huggingface",\n "repo_id": HF_REPO_ID,\n "repo_type": "model",\n}\nif hf_secret:\n storage_kwargs["token_secret"] = hf_secret.name\nstorage = HuggingfaceStorageConfigParam(**storage_kwargs)\n\ntry:\n base_model_fs = client.files.filesets.create(\n workspace="default",\n name=MODEL_NAME,\n description="NVIDIA Llama Nemotron Embed 1B v2 embedding model",\n storage=storage,\n )\nexcept ConflictError as e:\n print(f"Base model fileset already exists. Skipping creation.")\n base_model_fs = client.files.filesets.retrieve(\n workspace="default",\n name=MODEL_NAME,\n )\n\n# Create Model Entity referencing the FileSet\ntry:\n base_model = client.models.create(\n workspace="default",\n name=MODEL_NAME,\n fileset=f"default/{MODEL_NAME}",\n trust_remote_code=True,\n )\n print(f"Created Model Entity: {MODEL_NAME}")\nexcept ConflictError:\n print(f"Base model already exists. Updating fileset if different.")\n base_model = client.models.update(\n workspace="default",\n name=MODEL_NAME,\n fileset=f"default/{MODEL_NAME}",\n trust_remote_code=True,\n )\n\nprint(f"\\nBase model fileset: fileset://default/{base_model.name}")\nprint("\\nBase model files:")\nprint(json.dumps([f.model_dump() for f in client.files.list(fileset=MODEL_NAME, workspace="default").data], indent=2))\n\n# Wait for ModelSpec to be populated from the checkpoint\nprint("\\nWaiting for ModelSpec to be populated...")\nSPEC_TIMEOUT_SECONDS = 120\nspec_start = time.time()\nwhile not base_model.spec:\n if time.time() - spec_start > SPEC_TIMEOUT_SECONDS:\n raise TimeoutError(f"ModelSpec not populated within {SPEC_TIMEOUT_SECONDS} seconds")\n time.sleep(2)\n base_model = client.models.retrieve(\n workspace="default",\n name=MODEL_NAME,\n )\n\nprint(f"ModelSpec populated: {base_model.spec}")\n" }, { "type": "markdown", diff --git a/docs/fern/components/notebooks/embedding-customization-job.ts b/docs/fern/components/notebooks/embedding-customization-job.ts index 397617ed4e..50c25536cf 100644 --- a/docs/fern/components/notebooks/embedding-customization-job.ts +++ b/docs/fern/components/notebooks/embedding-customization-job.ts @@ -112,9 +112,9 @@ export default { cells: [ }, { "type": "code", - "source": "import time\nfrom nemo_platform.types.files import HuggingfaceStorageConfigParam\n\nHF_REPO_ID = \"nvidia/llama-nemotron-embed-1b-v2\"\nMODEL_NAME = \"nv-nemotron-embed-1b-base\"\n\nstorage = HuggingfaceStorageConfigParam(\n type=\"huggingface\",\n repo_id=HF_REPO_ID,\n repo_type=\"model\",\n)\nif hf_secret:\n storage[\"token_secret\"] = hf_secret.name\n\ntry:\n base_model_fs = client.files.filesets.create(\n workspace=\"default\",\n name=MODEL_NAME,\n description=\"NVIDIA Llama Nemotron Embed 1B v2 embedding model\",\n storage=storage,\n )\nexcept ConflictError as e:\n print(f\"Base model fileset already exists. Skipping creation.\")\n base_model_fs = client.files.filesets.retrieve(\n workspace=\"default\",\n name=MODEL_NAME,\n )\n\n# Create Model Entity referencing the FileSet\ntry:\n base_model = client.models.create(\n workspace=\"default\",\n name=MODEL_NAME,\n fileset=f\"default/{MODEL_NAME}\",\n trust_remote_code=True,\n )\n print(f\"Created Model Entity: {MODEL_NAME}\")\nexcept ConflictError:\n print(f\"Base model already exists. Updating fileset if different.\")\n base_model = client.models.update(\n workspace=\"default\",\n name=MODEL_NAME,\n fileset=f\"default/{MODEL_NAME}\",\n trust_remote_code=True,\n )\n\nprint(f\"\\nBase model fileset: fileset://default/{base_model.name}\")\nprint(\"\\nBase model files:\")\nprint(json.dumps([f.model_dump() for f in client.files.list(fileset=MODEL_NAME, workspace=\"default\").data], indent=2))\n\n# Wait for ModelSpec to be populated from the checkpoint\nprint(\"\\nWaiting for ModelSpec to be populated...\")\nSPEC_TIMEOUT_SECONDS = 120\nspec_start = time.time()\nwhile not base_model.spec:\n if time.time() - spec_start > SPEC_TIMEOUT_SECONDS:\n raise TimeoutError(f\"ModelSpec not populated within {SPEC_TIMEOUT_SECONDS} seconds\")\n time.sleep(2)\n base_model = client.models.retrieve(\n workspace=\"default\",\n name=MODEL_NAME,\n )\n\nprint(f\"ModelSpec populated: {base_model.spec}\")", + "source": "import time\nfrom nemo_platform.types.files import HuggingfaceStorageConfigParam\n\nHF_REPO_ID = \"nvidia/llama-nemotron-embed-1b-v2\"\nMODEL_NAME = \"nv-nemotron-embed-1b-base\"\n\nstorage_kwargs = {\n \"type\": \"huggingface\",\n \"repo_id\": HF_REPO_ID,\n \"repo_type\": \"model\",\n}\nif hf_secret:\n storage_kwargs[\"token_secret\"] = hf_secret.name\nstorage = HuggingfaceStorageConfigParam(**storage_kwargs)\n\ntry:\n base_model_fs = client.files.filesets.create(\n workspace=\"default\",\n name=MODEL_NAME,\n description=\"NVIDIA Llama Nemotron Embed 1B v2 embedding model\",\n storage=storage,\n )\nexcept ConflictError as e:\n print(f\"Base model fileset already exists. Skipping creation.\")\n base_model_fs = client.files.filesets.retrieve(\n workspace=\"default\",\n name=MODEL_NAME,\n )\n\n# Create Model Entity referencing the FileSet\ntry:\n base_model = client.models.create(\n workspace=\"default\",\n name=MODEL_NAME,\n fileset=f\"default/{MODEL_NAME}\",\n trust_remote_code=True,\n )\n print(f\"Created Model Entity: {MODEL_NAME}\")\nexcept ConflictError:\n print(f\"Base model already exists. Updating fileset if different.\")\n base_model = client.models.update(\n workspace=\"default\",\n name=MODEL_NAME,\n fileset=f\"default/{MODEL_NAME}\",\n trust_remote_code=True,\n )\n\nprint(f\"\\nBase model fileset: fileset://default/{base_model.name}\")\nprint(\"\\nBase model files:\")\nprint(json.dumps([f.model_dump() for f in client.files.list(fileset=MODEL_NAME, workspace=\"default\").data], indent=2))\n\n# Wait for ModelSpec to be populated from the checkpoint\nprint(\"\\nWaiting for ModelSpec to be populated...\")\nSPEC_TIMEOUT_SECONDS = 120\nspec_start = time.time()\nwhile not base_model.spec:\n if time.time() - spec_start > SPEC_TIMEOUT_SECONDS:\n raise TimeoutError(f\"ModelSpec not populated within {SPEC_TIMEOUT_SECONDS} seconds\")\n time.sleep(2)\n base_model = client.models.retrieve(\n workspace=\"default\",\n name=MODEL_NAME,\n )\n\nprint(f\"ModelSpec populated: {base_model.spec}\")", "language": "python", - "source_html": "import time\nfrom nemo_platform.types.files import HuggingfaceStorageConfigParam\n\nHF_REPO_ID = "nvidia/llama-nemotron-embed-1b-v2"\nMODEL_NAME = "nv-nemotron-embed-1b-base"\n\nstorage = HuggingfaceStorageConfigParam(\n type="huggingface",\n repo_id=HF_REPO_ID,\n repo_type="model",\n)\nif hf_secret:\n storage["token_secret"] = hf_secret.name\n\ntry:\n base_model_fs = client.files.filesets.create(\n workspace="default",\n name=MODEL_NAME,\n description="NVIDIA Llama Nemotron Embed 1B v2 embedding model",\n storage=storage,\n )\nexcept ConflictError as e:\n print(f"Base model fileset already exists. Skipping creation.")\n base_model_fs = client.files.filesets.retrieve(\n workspace="default",\n name=MODEL_NAME,\n )\n\n# Create Model Entity referencing the FileSet\ntry:\n base_model = client.models.create(\n workspace="default",\n name=MODEL_NAME,\n fileset=f"default/{MODEL_NAME}",\n trust_remote_code=True,\n )\n print(f"Created Model Entity: {MODEL_NAME}")\nexcept ConflictError:\n print(f"Base model already exists. Updating fileset if different.")\n base_model = client.models.update(\n workspace="default",\n name=MODEL_NAME,\n fileset=f"default/{MODEL_NAME}",\n trust_remote_code=True,\n )\n\nprint(f"\\nBase model fileset: fileset://default/{base_model.name}")\nprint("\\nBase model files:")\nprint(json.dumps([f.model_dump() for f in client.files.list(fileset=MODEL_NAME, workspace="default").data], indent=2))\n\n# Wait for ModelSpec to be populated from the checkpoint\nprint("\\nWaiting for ModelSpec to be populated...")\nSPEC_TIMEOUT_SECONDS = 120\nspec_start = time.time()\nwhile not base_model.spec:\n if time.time() - spec_start > SPEC_TIMEOUT_SECONDS:\n raise TimeoutError(f"ModelSpec not populated within {SPEC_TIMEOUT_SECONDS} seconds")\n time.sleep(2)\n base_model = client.models.retrieve(\n workspace="default",\n name=MODEL_NAME,\n )\n\nprint(f"ModelSpec populated: {base_model.spec}")\n" + "source_html": "import time\nfrom nemo_platform.types.files import HuggingfaceStorageConfigParam\n\nHF_REPO_ID = "nvidia/llama-nemotron-embed-1b-v2"\nMODEL_NAME = "nv-nemotron-embed-1b-base"\n\nstorage_kwargs = {\n "type": "huggingface",\n "repo_id": HF_REPO_ID,\n "repo_type": "model",\n}\nif hf_secret:\n storage_kwargs["token_secret"] = hf_secret.name\nstorage = HuggingfaceStorageConfigParam(**storage_kwargs)\n\ntry:\n base_model_fs = client.files.filesets.create(\n workspace="default",\n name=MODEL_NAME,\n description="NVIDIA Llama Nemotron Embed 1B v2 embedding model",\n storage=storage,\n )\nexcept ConflictError as e:\n print(f"Base model fileset already exists. Skipping creation.")\n base_model_fs = client.files.filesets.retrieve(\n workspace="default",\n name=MODEL_NAME,\n )\n\n# Create Model Entity referencing the FileSet\ntry:\n base_model = client.models.create(\n workspace="default",\n name=MODEL_NAME,\n fileset=f"default/{MODEL_NAME}",\n trust_remote_code=True,\n )\n print(f"Created Model Entity: {MODEL_NAME}")\nexcept ConflictError:\n print(f"Base model already exists. Updating fileset if different.")\n base_model = client.models.update(\n workspace="default",\n name=MODEL_NAME,\n fileset=f"default/{MODEL_NAME}",\n trust_remote_code=True,\n )\n\nprint(f"\\nBase model fileset: fileset://default/{base_model.name}")\nprint("\\nBase model files:")\nprint(json.dumps([f.model_dump() for f in client.files.list(fileset=MODEL_NAME, workspace="default").data], indent=2))\n\n# Wait for ModelSpec to be populated from the checkpoint\nprint("\\nWaiting for ModelSpec to be populated...")\nSPEC_TIMEOUT_SECONDS = 120\nspec_start = time.time()\nwhile not base_model.spec:\n if time.time() - spec_start > SPEC_TIMEOUT_SECONDS:\n raise TimeoutError(f"ModelSpec not populated within {SPEC_TIMEOUT_SECONDS} seconds")\n time.sleep(2)\n base_model = client.models.retrieve(\n workspace="default",\n name=MODEL_NAME,\n )\n\nprint(f"ModelSpec populated: {base_model.spec}")\n" }, { "type": "markdown", diff --git a/docs/fern/components/notebooks/sft-customization-job.json b/docs/fern/components/notebooks/sft-customization-job.json index 8195bb759c..97a3e2732c 100644 --- a/docs/fern/components/notebooks/sft-customization-job.json +++ b/docs/fern/components/notebooks/sft-customization-job.json @@ -90,9 +90,9 @@ }, { "type": "code", - "source": "# Export the HF_TOKEN and NGC_API_KEY environment variables if they are not already set\nHF_TOKEN = os.getenv(\"HF_TOKEN\")\nNGC_API_KEY = os.getenv(\"NGC_API_KEY\")\n\n\ndef create_or_get_secret(name: str, value: str | None, label: str):\n if not value:\n raise ValueError(f\"{label} is not set\")\n try:\n secret = client.secrets.create(\n name=name,\n workspace=\"default\",\n value=value,\n )\n print(f\"Created secret: {name}\")\n return secret\n except ConflictError:\n print(f\"Secret '{name}' already exists, continuing...\")\n return client.secrets.retrieve(name=name, workspace=\"default\")\n\n\n# Create Hugging Face token secret\nhf_secret = create_or_get_secret(\"hf-token\", HF_TOKEN, \"HF_TOKEN\")\nprint(\"HF_TOKEN secret:\")\nprint(hf_secret.model_dump_json(indent=2))\n\n# Create NGC API key secret\n# Uncomment the line below if you have NGC API Key and want to finetune NGC models\n# ngc_api_key = create_or_get_secret(\"ngc-api-key\", NGC_API_KEY, \"NGC_API_KEY\")", + "source": "# Export the HF_TOKEN and NGC_API_KEY environment variables if they are not already set\nHF_TOKEN = os.getenv(\"HF_TOKEN\")\nNGC_API_KEY = os.getenv(\"NGC_API_KEY\")\n\n\ndef create_or_get_secret(name: str, value: str | None, label: str):\n if not value:\n raise ValueError(f\"{label} is not set\")\n try:\n secret = client.secrets.create(\n name=name,\n workspace=\"default\",\n value=value,\n )\n print(f\"Created secret: {name}\")\n return secret\n except ConflictError:\n print(f\"Secret '{name}' already exists, continuing...\")\n return client.secrets.retrieve(name=name, workspace=\"default\")\n\n\n# Create Hugging Face token secret only when HF_TOKEN is set.\n# Gated models (for example Llama) require this secret; public models can omit token_secret.\nhf_secret = create_or_get_secret(\"hf-token\", HF_TOKEN, \"HF_TOKEN\") if HF_TOKEN else None\nif hf_secret:\n print(\"HF_TOKEN secret:\")\n print(hf_secret.model_dump_json(indent=2))\nelse:\n print(\"HF_TOKEN is not set; continuing without token_secret (public models only).\")\n\n# Create NGC API key secret\n# Uncomment the line below if you have NGC API Key and want to finetune NGC models\n# ngc_api_key = create_or_get_secret(\"ngc-api-key\", NGC_API_KEY, \"NGC_API_KEY\")", "language": "python", - "source_html": "# Export the HF_TOKEN and NGC_API_KEY environment variables if they are not already set\nHF_TOKEN = os.getenv("HF_TOKEN")\nNGC_API_KEY = os.getenv("NGC_API_KEY")\n\n\ndef create_or_get_secret(name: str, value: str | None, label: str):\n if not value:\n raise ValueError(f"{label} is not set")\n try:\n secret = client.secrets.create(\n name=name,\n workspace="default",\n value=value,\n )\n print(f"Created secret: {name}")\n return secret\n except ConflictError:\n print(f"Secret '{name}' already exists, continuing...")\n return client.secrets.retrieve(name=name, workspace="default")\n\n\n# Create Hugging Face token secret\nhf_secret = create_or_get_secret("hf-token", HF_TOKEN, "HF_TOKEN")\nprint("HF_TOKEN secret:")\nprint(hf_secret.model_dump_json(indent=2))\n\n# Create NGC API key secret\n# Uncomment the line below if you have NGC API Key and want to finetune NGC models\n# ngc_api_key = create_or_get_secret("ngc-api-key", NGC_API_KEY, "NGC_API_KEY")\n" + "source_html": "# Export the HF_TOKEN and NGC_API_KEY environment variables if they are not already set\nHF_TOKEN = os.getenv("HF_TOKEN")\nNGC_API_KEY = os.getenv("NGC_API_KEY")\n\n\ndef create_or_get_secret(name: str, value: str | None, label: str):\n if not value:\n raise ValueError(f"{label} is not set")\n try:\n secret = client.secrets.create(\n name=name,\n workspace="default",\n value=value,\n )\n print(f"Created secret: {name}")\n return secret\n except ConflictError:\n print(f"Secret '{name}' already exists, continuing...")\n return client.secrets.retrieve(name=name, workspace="default")\n\n\n# Create Hugging Face token secret only when HF_TOKEN is set.\n# Gated models (for example Llama) require this secret; public models can omit token_secret.\nhf_secret = create_or_get_secret("hf-token", HF_TOKEN, "HF_TOKEN") if HF_TOKEN else None\nif hf_secret:\n print("HF_TOKEN secret:")\n print(hf_secret.model_dump_json(indent=2))\nelse:\n print("HF_TOKEN is not set; continuing without token_secret (public models only).")\n\n# Create NGC API key secret\n# Uncomment the line below if you have NGC API Key and want to finetune NGC models\n# ngc_api_key = create_or_get_secret("ngc-api-key", NGC_API_KEY, "NGC_API_KEY")\n" }, { "type": "markdown", @@ -101,9 +101,9 @@ }, { "type": "code", - "source": "import time\n\n# Create a fileset pointing to the desired Hugging Face model\nfrom nemo_platform.types.files import HuggingfaceStorageConfigParam\n\nHF_REPO_ID = \"meta-llama/Llama-3.2-1B-Instruct\"\nMODEL_NAME = \"llama-3-2-1b-base\"\n\n# Ensure you have a Hugging Face token secret created\ntry:\n base_model_fs = client.files.filesets.create(\n workspace=\"default\",\n name=MODEL_NAME,\n description=\"Llama 3.2 1B base model from Hugging Face\",\n storage=HuggingfaceStorageConfigParam(\n type=\"huggingface\",\n # repo_id is the full model name from Hugging Face\n repo_id=HF_REPO_ID,\n repo_type=\"model\",\n # we use the secret created in the previous step\n token_secret=hf_secret.name\n )\n )\n print(f\"Created base model fileset: {MODEL_NAME}\")\nexcept ConflictError:\n print(f\"Base model fileset already exists. Skipping creation.\")\n base_model_fs = client.files.filesets.retrieve(\n workspace=\"default\",\n name=MODEL_NAME,\n )\n\n# Create the Model Entity representation.\ntry:\n base_model = client.models.create(\n workspace=\"default\",\n name=MODEL_NAME,\n fileset=f\"default/{MODEL_NAME}\",\n )\n print(f\"Created Model Entity: {MODEL_NAME}\")\nexcept ConflictError:\n print(f\"Base model already exists. Updating fileset if different.\")\n base_model = client.models.update(\n workspace=\"default\",\n name=MODEL_NAME,\n fileset=f\"default/{MODEL_NAME}\",\n )\n\nprint(f\"\\nBase model fileset: fileset://default/{base_model.name}\")\nprint(\"Base model fileset files list:\")\nprint(json.dumps([f.model_dump() for f in client.files.list(fileset=MODEL_NAME, workspace=\"default\").data], indent=2))\n\n# Wait for ModelSpec to be populated from the checkpoint\nprint(\"\\nWaiting for ModelSpec to be populated...\")\nSPEC_TIMEOUT_SECONDS = 120\nspec_start = time.time()\nwhile not base_model.spec:\n if time.time() - spec_start > SPEC_TIMEOUT_SECONDS:\n raise TimeoutError(f\"ModelSpec not populated within {SPEC_TIMEOUT_SECONDS} seconds\")\n time.sleep(2)\n base_model = client.models.retrieve(\n workspace=\"default\",\n name=MODEL_NAME,\n )\n\nprint(f\"ModelSpec populated: {base_model.spec}\")", + "source": "import time\n\n# Create a fileset pointing to the desired Hugging Face model\nfrom nemo_platform.types.files import HuggingfaceStorageConfigParam\n\nHF_REPO_ID = \"meta-llama/Llama-3.2-1B-Instruct\"\nMODEL_NAME = \"llama-3-2-1b-base\"\n\nstorage_kwargs = {\n \"type\": \"huggingface\",\n # repo_id is the full model name from Hugging Face\n \"repo_id\": HF_REPO_ID,\n \"repo_type\": \"model\",\n}\nif hf_secret:\n # we use the secret created in the previous step for gated models\n storage_kwargs[\"token_secret\"] = hf_secret.name\n\n# Ensure you have a Hugging Face token secret created for gated models\ntry:\n base_model_fs = client.files.filesets.create(\n workspace=\"default\",\n name=MODEL_NAME,\n description=\"Llama 3.2 1B base model from Hugging Face\",\n storage=HuggingfaceStorageConfigParam(**storage_kwargs),\n )\n print(f\"Created base model fileset: {MODEL_NAME}\")\nexcept ConflictError:\n print(f\"Base model fileset already exists. Skipping creation.\")\n base_model_fs = client.files.filesets.retrieve(\n workspace=\"default\",\n name=MODEL_NAME,\n )\n\n# Create the Model Entity representation.\ntry:\n base_model = client.models.create(\n workspace=\"default\",\n name=MODEL_NAME,\n fileset=f\"default/{MODEL_NAME}\",\n )\n print(f\"Created Model Entity: {MODEL_NAME}\")\nexcept ConflictError:\n print(f\"Base model already exists. Updating fileset if different.\")\n base_model = client.models.update(\n workspace=\"default\",\n name=MODEL_NAME,\n fileset=f\"default/{MODEL_NAME}\",\n )\n\nprint(f\"\\nBase model fileset: fileset://default/{base_model.name}\")\nprint(\"Base model fileset files list:\")\nprint(json.dumps([f.model_dump() for f in client.files.list(fileset=MODEL_NAME, workspace=\"default\").data], indent=2))\n\n# Wait for ModelSpec to be populated from the checkpoint\nprint(\"\\nWaiting for ModelSpec to be populated...\")\nSPEC_TIMEOUT_SECONDS = 120\nspec_start = time.time()\nwhile not base_model.spec:\n if time.time() - spec_start > SPEC_TIMEOUT_SECONDS:\n raise TimeoutError(f\"ModelSpec not populated within {SPEC_TIMEOUT_SECONDS} seconds\")\n time.sleep(2)\n base_model = client.models.retrieve(\n workspace=\"default\",\n name=MODEL_NAME,\n )\n\nprint(f\"ModelSpec populated: {base_model.spec}\")", "language": "python", - "source_html": "import time\n\n# Create a fileset pointing to the desired Hugging Face model\nfrom nemo_platform.types.files import HuggingfaceStorageConfigParam\n\nHF_REPO_ID = "meta-llama/Llama-3.2-1B-Instruct"\nMODEL_NAME = "llama-3-2-1b-base"\n\n# Ensure you have a Hugging Face token secret created\ntry:\n base_model_fs = client.files.filesets.create(\n workspace="default",\n name=MODEL_NAME,\n description="Llama 3.2 1B base model from Hugging Face",\n storage=HuggingfaceStorageConfigParam(\n type="huggingface",\n # repo_id is the full model name from Hugging Face\n repo_id=HF_REPO_ID,\n repo_type="model",\n # we use the secret created in the previous step\n token_secret=hf_secret.name\n )\n )\n print(f"Created base model fileset: {MODEL_NAME}")\nexcept ConflictError:\n print(f"Base model fileset already exists. Skipping creation.")\n base_model_fs = client.files.filesets.retrieve(\n workspace="default",\n name=MODEL_NAME,\n )\n\n# Create the Model Entity representation.\ntry:\n base_model = client.models.create(\n workspace="default",\n name=MODEL_NAME,\n fileset=f"default/{MODEL_NAME}",\n )\n print(f"Created Model Entity: {MODEL_NAME}")\nexcept ConflictError:\n print(f"Base model already exists. Updating fileset if different.")\n base_model = client.models.update(\n workspace="default",\n name=MODEL_NAME,\n fileset=f"default/{MODEL_NAME}",\n )\n\nprint(f"\\nBase model fileset: fileset://default/{base_model.name}")\nprint("Base model fileset files list:")\nprint(json.dumps([f.model_dump() for f in client.files.list(fileset=MODEL_NAME, workspace="default").data], indent=2))\n\n# Wait for ModelSpec to be populated from the checkpoint\nprint("\\nWaiting for ModelSpec to be populated...")\nSPEC_TIMEOUT_SECONDS = 120\nspec_start = time.time()\nwhile not base_model.spec:\n if time.time() - spec_start > SPEC_TIMEOUT_SECONDS:\n raise TimeoutError(f"ModelSpec not populated within {SPEC_TIMEOUT_SECONDS} seconds")\n time.sleep(2)\n base_model = client.models.retrieve(\n workspace="default",\n name=MODEL_NAME,\n )\n\nprint(f"ModelSpec populated: {base_model.spec}")\n" + "source_html": "import time\n\n# Create a fileset pointing to the desired Hugging Face model\nfrom nemo_platform.types.files import HuggingfaceStorageConfigParam\n\nHF_REPO_ID = "meta-llama/Llama-3.2-1B-Instruct"\nMODEL_NAME = "llama-3-2-1b-base"\n\nstorage_kwargs = {\n "type": "huggingface",\n # repo_id is the full model name from Hugging Face\n "repo_id": HF_REPO_ID,\n "repo_type": "model",\n}\nif hf_secret:\n # we use the secret created in the previous step for gated models\n storage_kwargs["token_secret"] = hf_secret.name\n\n# Ensure you have a Hugging Face token secret created for gated models\ntry:\n base_model_fs = client.files.filesets.create(\n workspace="default",\n name=MODEL_NAME,\n description="Llama 3.2 1B base model from Hugging Face",\n storage=HuggingfaceStorageConfigParam(**storage_kwargs),\n )\n print(f"Created base model fileset: {MODEL_NAME}")\nexcept ConflictError:\n print(f"Base model fileset already exists. Skipping creation.")\n base_model_fs = client.files.filesets.retrieve(\n workspace="default",\n name=MODEL_NAME,\n )\n\n# Create the Model Entity representation.\ntry:\n base_model = client.models.create(\n workspace="default",\n name=MODEL_NAME,\n fileset=f"default/{MODEL_NAME}",\n )\n print(f"Created Model Entity: {MODEL_NAME}")\nexcept ConflictError:\n print(f"Base model already exists. Updating fileset if different.")\n base_model = client.models.update(\n workspace="default",\n name=MODEL_NAME,\n fileset=f"default/{MODEL_NAME}",\n )\n\nprint(f"\\nBase model fileset: fileset://default/{base_model.name}")\nprint("Base model fileset files list:")\nprint(json.dumps([f.model_dump() for f in client.files.list(fileset=MODEL_NAME, workspace="default").data], indent=2))\n\n# Wait for ModelSpec to be populated from the checkpoint\nprint("\\nWaiting for ModelSpec to be populated...")\nSPEC_TIMEOUT_SECONDS = 120\nspec_start = time.time()\nwhile not base_model.spec:\n if time.time() - spec_start > SPEC_TIMEOUT_SECONDS:\n raise TimeoutError(f"ModelSpec not populated within {SPEC_TIMEOUT_SECONDS} seconds")\n time.sleep(2)\n base_model = client.models.retrieve(\n workspace="default",\n name=MODEL_NAME,\n )\n\nprint(f"ModelSpec populated: {base_model.spec}")\n" }, { "type": "markdown", @@ -187,4 +187,4 @@ "source_html": "

Evaluation Best Practices

\n

Manual Evaluation (Recommended)

\n
    \n
  • Test with real-world examples from your use case
    • \n
  • Compare responses to base model and expected outputs
    • \n
  • Verify the model exhibits desired behavior changes
    • \n
  • Check edge cases and error handling
    • \n
\n

What to look for:

\n
    \n
  • ✅ Model follows your desired output format
    • \n
  • ✅ Applies domain knowledge correctly
    • \n
  • ✅ Maintains general language capabilities
    • \n
  • ✅ Avoids unwanted behaviors or biases
    • \n
  • ❌ Doesn't hallucinate facts not in training data
    • \n
  • ❌ Doesn't produce repetitive or nonsensical outputs
    • \n
\n
\n

Hyperparameters

\n

For detailed information on all available hyperparameters, recommended values, and tuning guidance, refer to the Hyperparameter Reference.

\n
\n

Troubleshooting

\n

Job fails during model download:

\n
    \n
  • Verify authentication secrets are configured (refer to Managing Secrets)
    • \n
  • For gated Hugging Face models (Llama, Gemma), accept the license on the model page (for example, meta-llama/Llama-3.2-1B-Instruct)
    • \n
  • Confirm the model fileset uses token_secret=hf_secret.name for gated models
    • \n
  • Check AutomodelJobInput references use the workspace/name format: model=f"default/{MODEL_NAME}" and dataset={"training": f"default/{DATASET_NAME}"} (for example, default/llama-3-2-1b-base, default/sft-dataset)
    • \n
  • Verify the model entity points at the fileset: fileset=f"default/{MODEL_NAME}"
    • \n
  • Check job status: client.jobs.get_status(name=job.job.name, workspace="default")
    • \n
\n

Job fails with OOM (Out of Memory) error:

\n
    \n
  1. First try: Reduce global_batch_size from 64 to 32 or 16 in batch={...}
    2. \n
  2. Still OOM: Keep micro_batch_size at 1 (already the minimum in this tutorial)
    3. \n
  3. Still OOM: Reduce max_seq_length from 2048 to 1024 or 512 in training={...}
    4. \n
  4. Last resort: Increase num_gpus_per_node and tensor_parallel_size in parallelism={...}
    5. \n
\n

Loss curves not decreasing (underfitting):

\n
    \n
  • Increase training duration: raise epochs from 2 to 3-5 in schedule={...}
    • \n
  • Adjust learning rate: try 1e-4 or 1e-5 instead of the default 5e-5 in optimizer={...}
    • \n
  • Check data quality: Verify formatting, remove duplicates, ensure diversity
    • \n
\n

Training loss decreases but validation loss increases (overfitting):

\n
    \n
  • Reduce epochs from 2 to 1 in schedule={...}
    • \n
  • Lower learning_rate from 5e-5 to 2e-5 or 1e-5 in optimizer={...}
    • \n
  • Increase dataset size and diversity
    • \n
  • Verify train/validation split has no data leakage
    • \n
\n

Model output quality is poor despite good training metrics:

\n
    \n
  • Training metrics optimize for loss, not your actual task—evaluate on real use cases
    • \n
  • Review data quality, format, and diversity—metrics can be misleading with poor data
    • \n
  • Try a different base model size or architecture
    • \n
  • Adjust learning_rate and global_batch_size
    • \n
  • Compare to baseline: Test base model to ensure fine-tuning improved performance
    • \n
\n

Deployment fails:

\n
    \n
  • Verify output model exists: client.models.retrieve(name=OUTPUT_NAME, workspace="default")
    • \n
  • Check deployment logs: client.inference.deployments.get_logs(name=deployment.name, workspace="default")
    • \n
  • Ensure sufficient GPU resources for executor_config={"gpu": 1, ...}
    • \n
  • Verify the deployment config matches this tutorial: engine="vllm" with vllm/vllm-openai:v0.22.1
    • \n
\n

Next Steps

\n\n" } ] -} +} \ No newline at end of file diff --git a/docs/fern/components/notebooks/sft-customization-job.ts b/docs/fern/components/notebooks/sft-customization-job.ts index 053dc0c3ef..7bf0d620b9 100644 --- a/docs/fern/components/notebooks/sft-customization-job.ts +++ b/docs/fern/components/notebooks/sft-customization-job.ts @@ -95,9 +95,9 @@ export default { cells: [ }, { "type": "code", - "source": "# Export the HF_TOKEN and NGC_API_KEY environment variables if they are not already set\nHF_TOKEN = os.getenv(\"HF_TOKEN\")\nNGC_API_KEY = os.getenv(\"NGC_API_KEY\")\n\n\ndef create_or_get_secret(name: str, value: str | None, label: str):\n if not value:\n raise ValueError(f\"{label} is not set\")\n try:\n secret = client.secrets.create(\n name=name,\n workspace=\"default\",\n value=value,\n )\n print(f\"Created secret: {name}\")\n return secret\n except ConflictError:\n print(f\"Secret '{name}' already exists, continuing...\")\n return client.secrets.retrieve(name=name, workspace=\"default\")\n\n\n# Create Hugging Face token secret\nhf_secret = create_or_get_secret(\"hf-token\", HF_TOKEN, \"HF_TOKEN\")\nprint(\"HF_TOKEN secret:\")\nprint(hf_secret.model_dump_json(indent=2))\n\n# Create NGC API key secret\n# Uncomment the line below if you have NGC API Key and want to finetune NGC models\n# ngc_api_key = create_or_get_secret(\"ngc-api-key\", NGC_API_KEY, \"NGC_API_KEY\")", + "source": "# Export the HF_TOKEN and NGC_API_KEY environment variables if they are not already set\nHF_TOKEN = os.getenv(\"HF_TOKEN\")\nNGC_API_KEY = os.getenv(\"NGC_API_KEY\")\n\n\ndef create_or_get_secret(name: str, value: str | None, label: str):\n if not value:\n raise ValueError(f\"{label} is not set\")\n try:\n secret = client.secrets.create(\n name=name,\n workspace=\"default\",\n value=value,\n )\n print(f\"Created secret: {name}\")\n return secret\n except ConflictError:\n print(f\"Secret '{name}' already exists, continuing...\")\n return client.secrets.retrieve(name=name, workspace=\"default\")\n\n\n# Create Hugging Face token secret only when HF_TOKEN is set.\n# Gated models (for example Llama) require this secret; public models can omit token_secret.\nhf_secret = create_or_get_secret(\"hf-token\", HF_TOKEN, \"HF_TOKEN\") if HF_TOKEN else None\nif hf_secret:\n print(\"HF_TOKEN secret:\")\n print(hf_secret.model_dump_json(indent=2))\nelse:\n print(\"HF_TOKEN is not set; continuing without token_secret (public models only).\")\n\n# Create NGC API key secret\n# Uncomment the line below if you have NGC API Key and want to finetune NGC models\n# ngc_api_key = create_or_get_secret(\"ngc-api-key\", NGC_API_KEY, \"NGC_API_KEY\")", "language": "python", - "source_html": "# Export the HF_TOKEN and NGC_API_KEY environment variables if they are not already set\nHF_TOKEN = os.getenv("HF_TOKEN")\nNGC_API_KEY = os.getenv("NGC_API_KEY")\n\n\ndef create_or_get_secret(name: str, value: str | None, label: str):\n if not value:\n raise ValueError(f"{label} is not set")\n try:\n secret = client.secrets.create(\n name=name,\n workspace="default",\n value=value,\n )\n print(f"Created secret: {name}")\n return secret\n except ConflictError:\n print(f"Secret '{name}' already exists, continuing...")\n return client.secrets.retrieve(name=name, workspace="default")\n\n\n# Create Hugging Face token secret\nhf_secret = create_or_get_secret("hf-token", HF_TOKEN, "HF_TOKEN")\nprint("HF_TOKEN secret:")\nprint(hf_secret.model_dump_json(indent=2))\n\n# Create NGC API key secret\n# Uncomment the line below if you have NGC API Key and want to finetune NGC models\n# ngc_api_key = create_or_get_secret("ngc-api-key", NGC_API_KEY, "NGC_API_KEY")\n" + "source_html": "# Export the HF_TOKEN and NGC_API_KEY environment variables if they are not already set\nHF_TOKEN = os.getenv("HF_TOKEN")\nNGC_API_KEY = os.getenv("NGC_API_KEY")\n\n\ndef create_or_get_secret(name: str, value: str | None, label: str):\n if not value:\n raise ValueError(f"{label} is not set")\n try:\n secret = client.secrets.create(\n name=name,\n workspace="default",\n value=value,\n )\n print(f"Created secret: {name}")\n return secret\n except ConflictError:\n print(f"Secret '{name}' already exists, continuing...")\n return client.secrets.retrieve(name=name, workspace="default")\n\n\n# Create Hugging Face token secret only when HF_TOKEN is set.\n# Gated models (for example Llama) require this secret; public models can omit token_secret.\nhf_secret = create_or_get_secret("hf-token", HF_TOKEN, "HF_TOKEN") if HF_TOKEN else None\nif hf_secret:\n print("HF_TOKEN secret:")\n print(hf_secret.model_dump_json(indent=2))\nelse:\n print("HF_TOKEN is not set; continuing without token_secret (public models only).")\n\n# Create NGC API key secret\n# Uncomment the line below if you have NGC API Key and want to finetune NGC models\n# ngc_api_key = create_or_get_secret("ngc-api-key", NGC_API_KEY, "NGC_API_KEY")\n" }, { "type": "markdown", @@ -106,9 +106,9 @@ export default { cells: [ }, { "type": "code", - "source": "import time\n\n# Create a fileset pointing to the desired Hugging Face model\nfrom nemo_platform.types.files import HuggingfaceStorageConfigParam\n\nHF_REPO_ID = \"meta-llama/Llama-3.2-1B-Instruct\"\nMODEL_NAME = \"llama-3-2-1b-base\"\n\n# Ensure you have a Hugging Face token secret created\ntry:\n base_model_fs = client.files.filesets.create(\n workspace=\"default\",\n name=MODEL_NAME,\n description=\"Llama 3.2 1B base model from Hugging Face\",\n storage=HuggingfaceStorageConfigParam(\n type=\"huggingface\",\n # repo_id is the full model name from Hugging Face\n repo_id=HF_REPO_ID,\n repo_type=\"model\",\n # we use the secret created in the previous step\n token_secret=hf_secret.name\n )\n )\n print(f\"Created base model fileset: {MODEL_NAME}\")\nexcept ConflictError:\n print(f\"Base model fileset already exists. Skipping creation.\")\n base_model_fs = client.files.filesets.retrieve(\n workspace=\"default\",\n name=MODEL_NAME,\n )\n\n# Create the Model Entity representation.\ntry:\n base_model = client.models.create(\n workspace=\"default\",\n name=MODEL_NAME,\n fileset=f\"default/{MODEL_NAME}\",\n )\n print(f\"Created Model Entity: {MODEL_NAME}\")\nexcept ConflictError:\n print(f\"Base model already exists. Updating fileset if different.\")\n base_model = client.models.update(\n workspace=\"default\",\n name=MODEL_NAME,\n fileset=f\"default/{MODEL_NAME}\",\n )\n\nprint(f\"\\nBase model fileset: fileset://default/{base_model.name}\")\nprint(\"Base model fileset files list:\")\nprint(json.dumps([f.model_dump() for f in client.files.list(fileset=MODEL_NAME, workspace=\"default\").data], indent=2))\n\n# Wait for ModelSpec to be populated from the checkpoint\nprint(\"\\nWaiting for ModelSpec to be populated...\")\nSPEC_TIMEOUT_SECONDS = 120\nspec_start = time.time()\nwhile not base_model.spec:\n if time.time() - spec_start > SPEC_TIMEOUT_SECONDS:\n raise TimeoutError(f\"ModelSpec not populated within {SPEC_TIMEOUT_SECONDS} seconds\")\n time.sleep(2)\n base_model = client.models.retrieve(\n workspace=\"default\",\n name=MODEL_NAME,\n )\n\nprint(f\"ModelSpec populated: {base_model.spec}\")", + "source": "import time\n\n# Create a fileset pointing to the desired Hugging Face model\nfrom nemo_platform.types.files import HuggingfaceStorageConfigParam\n\nHF_REPO_ID = \"meta-llama/Llama-3.2-1B-Instruct\"\nMODEL_NAME = \"llama-3-2-1b-base\"\n\nstorage_kwargs = {\n \"type\": \"huggingface\",\n # repo_id is the full model name from Hugging Face\n \"repo_id\": HF_REPO_ID,\n \"repo_type\": \"model\",\n}\nif hf_secret:\n # we use the secret created in the previous step for gated models\n storage_kwargs[\"token_secret\"] = hf_secret.name\n\n# Ensure you have a Hugging Face token secret created for gated models\ntry:\n base_model_fs = client.files.filesets.create(\n workspace=\"default\",\n name=MODEL_NAME,\n description=\"Llama 3.2 1B base model from Hugging Face\",\n storage=HuggingfaceStorageConfigParam(**storage_kwargs),\n )\n print(f\"Created base model fileset: {MODEL_NAME}\")\nexcept ConflictError:\n print(f\"Base model fileset already exists. Skipping creation.\")\n base_model_fs = client.files.filesets.retrieve(\n workspace=\"default\",\n name=MODEL_NAME,\n )\n\n# Create the Model Entity representation.\ntry:\n base_model = client.models.create(\n workspace=\"default\",\n name=MODEL_NAME,\n fileset=f\"default/{MODEL_NAME}\",\n )\n print(f\"Created Model Entity: {MODEL_NAME}\")\nexcept ConflictError:\n print(f\"Base model already exists. Updating fileset if different.\")\n base_model = client.models.update(\n workspace=\"default\",\n name=MODEL_NAME,\n fileset=f\"default/{MODEL_NAME}\",\n )\n\nprint(f\"\\nBase model fileset: fileset://default/{base_model.name}\")\nprint(\"Base model fileset files list:\")\nprint(json.dumps([f.model_dump() for f in client.files.list(fileset=MODEL_NAME, workspace=\"default\").data], indent=2))\n\n# Wait for ModelSpec to be populated from the checkpoint\nprint(\"\\nWaiting for ModelSpec to be populated...\")\nSPEC_TIMEOUT_SECONDS = 120\nspec_start = time.time()\nwhile not base_model.spec:\n if time.time() - spec_start > SPEC_TIMEOUT_SECONDS:\n raise TimeoutError(f\"ModelSpec not populated within {SPEC_TIMEOUT_SECONDS} seconds\")\n time.sleep(2)\n base_model = client.models.retrieve(\n workspace=\"default\",\n name=MODEL_NAME,\n )\n\nprint(f\"ModelSpec populated: {base_model.spec}\")", "language": "python", - "source_html": "import time\n\n# Create a fileset pointing to the desired Hugging Face model\nfrom nemo_platform.types.files import HuggingfaceStorageConfigParam\n\nHF_REPO_ID = "meta-llama/Llama-3.2-1B-Instruct"\nMODEL_NAME = "llama-3-2-1b-base"\n\n# Ensure you have a Hugging Face token secret created\ntry:\n base_model_fs = client.files.filesets.create(\n workspace="default",\n name=MODEL_NAME,\n description="Llama 3.2 1B base model from Hugging Face",\n storage=HuggingfaceStorageConfigParam(\n type="huggingface",\n # repo_id is the full model name from Hugging Face\n repo_id=HF_REPO_ID,\n repo_type="model",\n # we use the secret created in the previous step\n token_secret=hf_secret.name\n )\n )\n print(f"Created base model fileset: {MODEL_NAME}")\nexcept ConflictError:\n print(f"Base model fileset already exists. Skipping creation.")\n base_model_fs = client.files.filesets.retrieve(\n workspace="default",\n name=MODEL_NAME,\n )\n\n# Create the Model Entity representation.\ntry:\n base_model = client.models.create(\n workspace="default",\n name=MODEL_NAME,\n fileset=f"default/{MODEL_NAME}",\n )\n print(f"Created Model Entity: {MODEL_NAME}")\nexcept ConflictError:\n print(f"Base model already exists. Updating fileset if different.")\n base_model = client.models.update(\n workspace="default",\n name=MODEL_NAME,\n fileset=f"default/{MODEL_NAME}",\n )\n\nprint(f"\\nBase model fileset: fileset://default/{base_model.name}")\nprint("Base model fileset files list:")\nprint(json.dumps([f.model_dump() for f in client.files.list(fileset=MODEL_NAME, workspace="default").data], indent=2))\n\n# Wait for ModelSpec to be populated from the checkpoint\nprint("\\nWaiting for ModelSpec to be populated...")\nSPEC_TIMEOUT_SECONDS = 120\nspec_start = time.time()\nwhile not base_model.spec:\n if time.time() - spec_start > SPEC_TIMEOUT_SECONDS:\n raise TimeoutError(f"ModelSpec not populated within {SPEC_TIMEOUT_SECONDS} seconds")\n time.sleep(2)\n base_model = client.models.retrieve(\n workspace="default",\n name=MODEL_NAME,\n )\n\nprint(f"ModelSpec populated: {base_model.spec}")\n" + "source_html": "import time\n\n# Create a fileset pointing to the desired Hugging Face model\nfrom nemo_platform.types.files import HuggingfaceStorageConfigParam\n\nHF_REPO_ID = "meta-llama/Llama-3.2-1B-Instruct"\nMODEL_NAME = "llama-3-2-1b-base"\n\nstorage_kwargs = {\n "type": "huggingface",\n # repo_id is the full model name from Hugging Face\n "repo_id": HF_REPO_ID,\n "repo_type": "model",\n}\nif hf_secret:\n # we use the secret created in the previous step for gated models\n storage_kwargs["token_secret"] = hf_secret.name\n\n# Ensure you have a Hugging Face token secret created for gated models\ntry:\n base_model_fs = client.files.filesets.create(\n workspace="default",\n name=MODEL_NAME,\n description="Llama 3.2 1B base model from Hugging Face",\n storage=HuggingfaceStorageConfigParam(**storage_kwargs),\n )\n print(f"Created base model fileset: {MODEL_NAME}")\nexcept ConflictError:\n print(f"Base model fileset already exists. Skipping creation.")\n base_model_fs = client.files.filesets.retrieve(\n workspace="default",\n name=MODEL_NAME,\n )\n\n# Create the Model Entity representation.\ntry:\n base_model = client.models.create(\n workspace="default",\n name=MODEL_NAME,\n fileset=f"default/{MODEL_NAME}",\n )\n print(f"Created Model Entity: {MODEL_NAME}")\nexcept ConflictError:\n print(f"Base model already exists. Updating fileset if different.")\n base_model = client.models.update(\n workspace="default",\n name=MODEL_NAME,\n fileset=f"default/{MODEL_NAME}",\n )\n\nprint(f"\\nBase model fileset: fileset://default/{base_model.name}")\nprint("Base model fileset files list:")\nprint(json.dumps([f.model_dump() for f in client.files.list(fileset=MODEL_NAME, workspace="default").data], indent=2))\n\n# Wait for ModelSpec to be populated from the checkpoint\nprint("\\nWaiting for ModelSpec to be populated...")\nSPEC_TIMEOUT_SECONDS = 120\nspec_start = time.time()\nwhile not base_model.spec:\n if time.time() - spec_start > SPEC_TIMEOUT_SECONDS:\n raise TimeoutError(f"ModelSpec not populated within {SPEC_TIMEOUT_SECONDS} seconds")\n time.sleep(2)\n base_model = client.models.retrieve(\n workspace="default",\n name=MODEL_NAME,\n )\n\nprint(f"ModelSpec populated: {base_model.spec}")\n" }, { "type": "markdown", diff --git a/docs/fern/scripts/validate-notebook-viewer.mjs b/docs/fern/scripts/validate-notebook-viewer.mjs index dfc3096671..eb96f55bc9 100644 --- a/docs/fern/scripts/validate-notebook-viewer.mjs +++ b/docs/fern/scripts/validate-notebook-viewer.mjs @@ -4,7 +4,8 @@ * SPDX-License-Identifier: Apache-2.0 * * Ensure every NotebookViewer registration has generated notebook data on disk - * and that the generated JSON still matches its source notebook. + * and that the generated JSON and TypeScript artifacts still match their source + * notebooks. * * NotebookViewer.tsx imports `./notebooks/` modules produced by * `ipynb-to-fern-json.py`. A missing `.ts` / `.json` pair breaks publication @@ -87,9 +88,9 @@ for (const name of names) { try { const notebook = JSON.parse(await readFile(sourcePath, "utf8")); - const artifact = JSON.parse( - await readFile(join(NOTEBOOKS_DIR, `${name}.json`), "utf8"), - ); + const jsonPath = join(NOTEBOOKS_DIR, `${name}.json`); + const tsPath = join(NOTEBOOKS_DIR, `${name}.ts`); + const artifact = JSON.parse(await readFile(jsonPath, "utf8")); const sourceCells = notebook.cells.map((cell) => (Array.isArray(cell.source) ? cell.source.join("") : cell.source ?? "").trimEnd(), ); @@ -103,8 +104,28 @@ for (const name of names) { if (mismatch) { failed += 1; console.error( - `stale ${join(NOTEBOOKS_DIR, `${name}.json`)}; regenerate it from ${sourcePath}`, + `stale ${jsonPath}; regenerate it from ${sourcePath}`, + ); + } + + const tsSource = await readFile(tsPath, "utf8"); + const tsMatch = tsSource.match(/export\s+default\s+\{\s*cells:\s*(\[[\s\S]*\])\s*\};\s*$/); + if (!tsMatch) { + failed += 1; + console.error(`could not parse default export in ${tsPath}`); + } else { + const tsCells = JSON.parse(tsMatch[1]).map((cell) => + (cell.source ?? "").trimEnd(), ); + const tsMismatch = + artifactCells.length !== tsCells.length || + artifactCells.some((source, index) => source !== tsCells[index]); + if (tsMismatch) { + failed += 1; + console.error( + `stale ${tsPath}; regenerate it from ${sourcePath}`, + ); + } } } catch (error) { failed += 1; From 6f0c052d2dd3f3eee560a7b2c77db0ca188338ba Mon Sep 17 00:00:00 2001 From: Sam Oluwalana Date: Wed, 22 Jul 2026 14:51:49 -0600 Subject: [PATCH 04/24] Fix code rabbit comments Signed-off-by: Sam Oluwalana --- docs/customizer/about.mdx | 26 ++++++++++++------- .../distillation-customization-job.ipynb | 9 +++++++ .../distillation-customization-job.mdx | 9 +++++++ .../distillation-customization-job.json | 4 +-- .../distillation-customization-job.ts | 4 +-- 5 files changed, 38 insertions(+), 14 deletions(-) diff --git a/docs/customizer/about.mdx b/docs/customizer/about.mdx index 0e64e7bcb0..149fcdb7cf 100644 --- a/docs/customizer/about.mdx +++ b/docs/customizer/about.mdx @@ -236,25 +236,31 @@ TP can be configured via `parallelism.tensor_parallel_size` in the [training con As of release 25.10.0, AutoModel engines including Phi-4, Qwen, and Gemma support tensor parallelism greater than 1 through the multi-GPU LoRA patch. Previous releases only supported `TP=1` for these models. -### Pipeline Parallelism -[Pipeline Parallelism](https://docs.nvidia.com/nemo-framework/user-guide/latest/nemotoolkit/features/parallelisms.html#pipeline-parallelism) (PP) distributes the layers of a neural network across GPUs. The GPUs then process the different layers sequentially. +#### Tensor Parallelism Configuration -PP can be configured via `parallelism.pipeline_parallel_size` in the [training configuration](/documentation/customizer-reference/manage-customization-jobs/training-configuration). +**Constraints** -#### Configuration +- TP must be less than or equal to the total number of GPUs available. +- TP should divide the total GPU count evenly. -- Constraints -- TP must be less than or equal to the total number of GPUs available. It should be a factor of the total GPU count (divisible evenly). -- Multi-node considerations -- TP can span across nodes, but this introduces network communication overhead. For multi-node setups, it's often recommended to keep TP within a single node when possible. If using TP across nodes, high-bandwidth inter-node connections (like InfiniBand) become critical. +**Multi-node considerations** -Example: if you have 2 nodes with 4 GPUs each, start with TP=4 first. This keeps all tensor parallel operations within a single node. If your model still uses too much GPU memory with this setting, increase to TP=8, which will distribute tensor operations across both nodes. +TP can span nodes, but doing so increases network communication overhead. For multi-node setups, keep TP within a single node when possible. High-bandwidth inter-node connections such as InfiniBand are important when TP must span nodes. + +For example, with 2 nodes and 4 GPUs per node, start with `TP=4` to keep tensor-parallel operations within each node. If the model still requires more memory, increase to `TP=8` to distribute tensor operations across both nodes. + +**Performance** -- Performance - Smaller TP values generally have less communication overhead. - Larger TP values provide more memory savings but increase communication costs. +### Pipeline Parallelism + +[Pipeline Parallelism](https://docs.nvidia.com/nemo-framework/user-guide/latest/nemotoolkit/features/parallelisms.html#pipeline-parallelism) (PP) distributes the layers of a neural network across GPUs. The GPUs then process the different layers sequentially. + +PP can be configured via `parallelism.pipeline_parallel_size` in the [training configuration](/documentation/customizer-reference/manage-customization-jobs/training-configuration). + ### Context Parallelism [Context Parallelism](https://docs.nvidia.com/nemo-framework/user-guide/latest/nemotoolkit/features/parallelisms.html#context-parallelism) (CP) distributes activation memory along the sequence dimension across GPUs, which is particularly useful when training on datasets with very long sequences. diff --git a/docs/customizer/tutorials/distillation-customization-job.ipynb b/docs/customizer/tutorials/distillation-customization-job.ipynb index 39c5615d0b..59b68cd619 100644 --- a/docs/customizer/tutorials/distillation-customization-job.ipynb +++ b/docs/customizer/tutorials/distillation-customization-job.ipynb @@ -540,6 +540,15 @@ " # Free GPUs if readiness fails before the later baseline-cleanup cell runs.\n", " try:\n", " client.inference.deployments.delete(name=BASELINE_DEPLOYMENT_NAME, workspace=\"default\")\n", + " if not client.models.wait_for_status(\n", + " deployment_name=BASELINE_DEPLOYMENT_NAME,\n", + " desired_status=\"DELETED\",\n", + " workspace=\"default\",\n", + " timeout=600,\n", + " ):\n", + " raise TimeoutError(\n", + " f\"Deployment {BASELINE_DEPLOYMENT_NAME} was not deleted within timeout\"\n", + " )\n", " client.inference.deployment_configs.delete(name=BASELINE_DEPLOYMENT_CONFIG, workspace=\"default\")\n", " except Exception as cleanup_error:\n", " print(f\"Baseline cleanup after readiness failure also failed: {cleanup_error}\")\n", diff --git a/docs/customizer/tutorials/distillation-customization-job.mdx b/docs/customizer/tutorials/distillation-customization-job.mdx index d3b694e541..fa4dec2cb6 100644 --- a/docs/customizer/tutorials/distillation-customization-job.mdx +++ b/docs/customizer/tutorials/distillation-customization-job.mdx @@ -459,6 +459,15 @@ except Exception: # Free GPUs if readiness fails before the later baseline-cleanup cell runs. try: client.inference.deployments.delete(name=BASELINE_DEPLOYMENT_NAME, workspace="default") + if not client.models.wait_for_status( + deployment_name=BASELINE_DEPLOYMENT_NAME, + desired_status="DELETED", + workspace="default", + timeout=600, + ): + raise TimeoutError( + f"Deployment {BASELINE_DEPLOYMENT_NAME} was not deleted within timeout" + ) client.inference.deployment_configs.delete(name=BASELINE_DEPLOYMENT_CONFIG, workspace="default") except Exception as cleanup_error: print(f"Baseline cleanup after readiness failure also failed: {cleanup_error}") diff --git a/docs/fern/components/notebooks/distillation-customization-job.json b/docs/fern/components/notebooks/distillation-customization-job.json index dfea0c50a5..3ae0a8e082 100644 --- a/docs/fern/components/notebooks/distillation-customization-job.json +++ b/docs/fern/components/notebooks/distillation-customization-job.json @@ -90,9 +90,9 @@ }, { "type": "code", - "source": "def wait_for_deployment(deployment_name: str, timeout_minutes: int = 30):\n \"\"\"Poll deployment until ready.\"\"\"\n start = time.time()\n timeout = timeout_minutes * 60\n while True:\n dep = client.inference.deployments.retrieve(name=deployment_name, workspace=\"default\")\n elapsed = time.time() - start\n clear_output(wait=True)\n print(f\"Deployment: {deployment_name}\")\n print(f\"Status: {dep.status}\")\n print(f\"Elapsed: {int(elapsed // 60)}m {int(elapsed % 60)}s\")\n\n if dep.status == \"READY\":\n print(\"\\nDeployment is ready!\")\n remaining = int(timeout - elapsed)\n if remaining <= 0:\n raise TimeoutError(f\"Deployment timeout after {timeout_minutes} minutes\")\n if not client.models.wait_for_status(\n deployment_name=deployment_name,\n desired_status=\"READY\",\n workspace=\"default\",\n timeout=remaining,\n check_gateway=True,\n ):\n raise TimeoutError(\"Inference gateway did not become ready\")\n return dep\n if dep.status in (\"FAILED\", \"ERROR\", \"TERMINATED\", \"LOST\"):\n raise RuntimeError(f\"Deployment failed with status: {dep.status}\")\n if elapsed > timeout:\n raise TimeoutError(f\"Deployment timeout after {timeout_minutes} minutes\")\n time.sleep(15)\n\n\ntry:\n dep_status = wait_for_deployment(BASELINE_DEPLOYMENT_NAME)\n assert dep_status.status == \"READY\"\nexcept Exception:\n # Free GPUs if readiness fails before the later baseline-cleanup cell runs.\n try:\n client.inference.deployments.delete(name=BASELINE_DEPLOYMENT_NAME, workspace=\"default\")\n client.inference.deployment_configs.delete(name=BASELINE_DEPLOYMENT_CONFIG, workspace=\"default\")\n except Exception as cleanup_error:\n print(f\"Baseline cleanup after readiness failure also failed: {cleanup_error}\")\n raise", + "source": "def wait_for_deployment(deployment_name: str, timeout_minutes: int = 30):\n \"\"\"Poll deployment until ready.\"\"\"\n start = time.time()\n timeout = timeout_minutes * 60\n while True:\n dep = client.inference.deployments.retrieve(name=deployment_name, workspace=\"default\")\n elapsed = time.time() - start\n clear_output(wait=True)\n print(f\"Deployment: {deployment_name}\")\n print(f\"Status: {dep.status}\")\n print(f\"Elapsed: {int(elapsed // 60)}m {int(elapsed % 60)}s\")\n\n if dep.status == \"READY\":\n print(\"\\nDeployment is ready!\")\n remaining = int(timeout - elapsed)\n if remaining <= 0:\n raise TimeoutError(f\"Deployment timeout after {timeout_minutes} minutes\")\n if not client.models.wait_for_status(\n deployment_name=deployment_name,\n desired_status=\"READY\",\n workspace=\"default\",\n timeout=remaining,\n check_gateway=True,\n ):\n raise TimeoutError(\"Inference gateway did not become ready\")\n return dep\n if dep.status in (\"FAILED\", \"ERROR\", \"TERMINATED\", \"LOST\"):\n raise RuntimeError(f\"Deployment failed with status: {dep.status}\")\n if elapsed > timeout:\n raise TimeoutError(f\"Deployment timeout after {timeout_minutes} minutes\")\n time.sleep(15)\n\n\ntry:\n dep_status = wait_for_deployment(BASELINE_DEPLOYMENT_NAME)\n assert dep_status.status == \"READY\"\nexcept Exception:\n # Free GPUs if readiness fails before the later baseline-cleanup cell runs.\n try:\n client.inference.deployments.delete(name=BASELINE_DEPLOYMENT_NAME, workspace=\"default\")\n if not client.models.wait_for_status(\n deployment_name=BASELINE_DEPLOYMENT_NAME,\n desired_status=\"DELETED\",\n workspace=\"default\",\n timeout=600,\n ):\n raise TimeoutError(\n f\"Deployment {BASELINE_DEPLOYMENT_NAME} was not deleted within timeout\"\n )\n client.inference.deployment_configs.delete(name=BASELINE_DEPLOYMENT_CONFIG, workspace=\"default\")\n except Exception as cleanup_error:\n print(f\"Baseline cleanup after readiness failure also failed: {cleanup_error}\")\n raise", "language": "python", - "source_html": "def wait_for_deployment(deployment_name: str, timeout_minutes: int = 30):\n """Poll deployment until ready."""\n start = time.time()\n timeout = timeout_minutes * 60\n while True:\n dep = client.inference.deployments.retrieve(name=deployment_name, workspace="default")\n elapsed = time.time() - start\n clear_output(wait=True)\n print(f"Deployment: {deployment_name}")\n print(f"Status: {dep.status}")\n print(f"Elapsed: {int(elapsed // 60)}m {int(elapsed % 60)}s")\n\n if dep.status == "READY":\n print("\\nDeployment is ready!")\n remaining = int(timeout - elapsed)\n if remaining <= 0:\n raise TimeoutError(f"Deployment timeout after {timeout_minutes} minutes")\n if not client.models.wait_for_status(\n deployment_name=deployment_name,\n desired_status="READY",\n workspace="default",\n timeout=remaining,\n check_gateway=True,\n ):\n raise TimeoutError("Inference gateway did not become ready")\n return dep\n if dep.status in ("FAILED", "ERROR", "TERMINATED", "LOST"):\n raise RuntimeError(f"Deployment failed with status: {dep.status}")\n if elapsed > timeout:\n raise TimeoutError(f"Deployment timeout after {timeout_minutes} minutes")\n time.sleep(15)\n\n\ntry:\n dep_status = wait_for_deployment(BASELINE_DEPLOYMENT_NAME)\n assert dep_status.status == "READY"\nexcept Exception:\n # Free GPUs if readiness fails before the later baseline-cleanup cell runs.\n try:\n client.inference.deployments.delete(name=BASELINE_DEPLOYMENT_NAME, workspace="default")\n client.inference.deployment_configs.delete(name=BASELINE_DEPLOYMENT_CONFIG, workspace="default")\n except Exception as cleanup_error:\n print(f"Baseline cleanup after readiness failure also failed: {cleanup_error}")\n raise\n" + "source_html": "def wait_for_deployment(deployment_name: str, timeout_minutes: int = 30):\n """Poll deployment until ready."""\n start = time.time()\n timeout = timeout_minutes * 60\n while True:\n dep = client.inference.deployments.retrieve(name=deployment_name, workspace="default")\n elapsed = time.time() - start\n clear_output(wait=True)\n print(f"Deployment: {deployment_name}")\n print(f"Status: {dep.status}")\n print(f"Elapsed: {int(elapsed // 60)}m {int(elapsed % 60)}s")\n\n if dep.status == "READY":\n print("\\nDeployment is ready!")\n remaining = int(timeout - elapsed)\n if remaining <= 0:\n raise TimeoutError(f"Deployment timeout after {timeout_minutes} minutes")\n if not client.models.wait_for_status(\n deployment_name=deployment_name,\n desired_status="READY",\n workspace="default",\n timeout=remaining,\n check_gateway=True,\n ):\n raise TimeoutError("Inference gateway did not become ready")\n return dep\n if dep.status in ("FAILED", "ERROR", "TERMINATED", "LOST"):\n raise RuntimeError(f"Deployment failed with status: {dep.status}")\n if elapsed > timeout:\n raise TimeoutError(f"Deployment timeout after {timeout_minutes} minutes")\n time.sleep(15)\n\n\ntry:\n dep_status = wait_for_deployment(BASELINE_DEPLOYMENT_NAME)\n assert dep_status.status == "READY"\nexcept Exception:\n # Free GPUs if readiness fails before the later baseline-cleanup cell runs.\n try:\n client.inference.deployments.delete(name=BASELINE_DEPLOYMENT_NAME, workspace="default")\n if not client.models.wait_for_status(\n deployment_name=BASELINE_DEPLOYMENT_NAME,\n desired_status="DELETED",\n workspace="default",\n timeout=600,\n ):\n raise TimeoutError(\n f"Deployment {BASELINE_DEPLOYMENT_NAME} was not deleted within timeout"\n )\n client.inference.deployment_configs.delete(name=BASELINE_DEPLOYMENT_CONFIG, workspace="default")\n except Exception as cleanup_error:\n print(f"Baseline cleanup after readiness failure also failed: {cleanup_error}")\n raise\n" }, { "type": "markdown", diff --git a/docs/fern/components/notebooks/distillation-customization-job.ts b/docs/fern/components/notebooks/distillation-customization-job.ts index c68567c3c7..83d6f3a29e 100644 --- a/docs/fern/components/notebooks/distillation-customization-job.ts +++ b/docs/fern/components/notebooks/distillation-customization-job.ts @@ -95,9 +95,9 @@ export default { cells: [ }, { "type": "code", - "source": "def wait_for_deployment(deployment_name: str, timeout_minutes: int = 30):\n \"\"\"Poll deployment until ready.\"\"\"\n start = time.time()\n timeout = timeout_minutes * 60\n while True:\n dep = client.inference.deployments.retrieve(name=deployment_name, workspace=\"default\")\n elapsed = time.time() - start\n clear_output(wait=True)\n print(f\"Deployment: {deployment_name}\")\n print(f\"Status: {dep.status}\")\n print(f\"Elapsed: {int(elapsed // 60)}m {int(elapsed % 60)}s\")\n\n if dep.status == \"READY\":\n print(\"\\nDeployment is ready!\")\n remaining = int(timeout - elapsed)\n if remaining <= 0:\n raise TimeoutError(f\"Deployment timeout after {timeout_minutes} minutes\")\n if not client.models.wait_for_status(\n deployment_name=deployment_name,\n desired_status=\"READY\",\n workspace=\"default\",\n timeout=remaining,\n check_gateway=True,\n ):\n raise TimeoutError(\"Inference gateway did not become ready\")\n return dep\n if dep.status in (\"FAILED\", \"ERROR\", \"TERMINATED\", \"LOST\"):\n raise RuntimeError(f\"Deployment failed with status: {dep.status}\")\n if elapsed > timeout:\n raise TimeoutError(f\"Deployment timeout after {timeout_minutes} minutes\")\n time.sleep(15)\n\n\ntry:\n dep_status = wait_for_deployment(BASELINE_DEPLOYMENT_NAME)\n assert dep_status.status == \"READY\"\nexcept Exception:\n # Free GPUs if readiness fails before the later baseline-cleanup cell runs.\n try:\n client.inference.deployments.delete(name=BASELINE_DEPLOYMENT_NAME, workspace=\"default\")\n client.inference.deployment_configs.delete(name=BASELINE_DEPLOYMENT_CONFIG, workspace=\"default\")\n except Exception as cleanup_error:\n print(f\"Baseline cleanup after readiness failure also failed: {cleanup_error}\")\n raise", + "source": "def wait_for_deployment(deployment_name: str, timeout_minutes: int = 30):\n \"\"\"Poll deployment until ready.\"\"\"\n start = time.time()\n timeout = timeout_minutes * 60\n while True:\n dep = client.inference.deployments.retrieve(name=deployment_name, workspace=\"default\")\n elapsed = time.time() - start\n clear_output(wait=True)\n print(f\"Deployment: {deployment_name}\")\n print(f\"Status: {dep.status}\")\n print(f\"Elapsed: {int(elapsed // 60)}m {int(elapsed % 60)}s\")\n\n if dep.status == \"READY\":\n print(\"\\nDeployment is ready!\")\n remaining = int(timeout - elapsed)\n if remaining <= 0:\n raise TimeoutError(f\"Deployment timeout after {timeout_minutes} minutes\")\n if not client.models.wait_for_status(\n deployment_name=deployment_name,\n desired_status=\"READY\",\n workspace=\"default\",\n timeout=remaining,\n check_gateway=True,\n ):\n raise TimeoutError(\"Inference gateway did not become ready\")\n return dep\n if dep.status in (\"FAILED\", \"ERROR\", \"TERMINATED\", \"LOST\"):\n raise RuntimeError(f\"Deployment failed with status: {dep.status}\")\n if elapsed > timeout:\n raise TimeoutError(f\"Deployment timeout after {timeout_minutes} minutes\")\n time.sleep(15)\n\n\ntry:\n dep_status = wait_for_deployment(BASELINE_DEPLOYMENT_NAME)\n assert dep_status.status == \"READY\"\nexcept Exception:\n # Free GPUs if readiness fails before the later baseline-cleanup cell runs.\n try:\n client.inference.deployments.delete(name=BASELINE_DEPLOYMENT_NAME, workspace=\"default\")\n if not client.models.wait_for_status(\n deployment_name=BASELINE_DEPLOYMENT_NAME,\n desired_status=\"DELETED\",\n workspace=\"default\",\n timeout=600,\n ):\n raise TimeoutError(\n f\"Deployment {BASELINE_DEPLOYMENT_NAME} was not deleted within timeout\"\n )\n client.inference.deployment_configs.delete(name=BASELINE_DEPLOYMENT_CONFIG, workspace=\"default\")\n except Exception as cleanup_error:\n print(f\"Baseline cleanup after readiness failure also failed: {cleanup_error}\")\n raise", "language": "python", - "source_html": "def wait_for_deployment(deployment_name: str, timeout_minutes: int = 30):\n """Poll deployment until ready."""\n start = time.time()\n timeout = timeout_minutes * 60\n while True:\n dep = client.inference.deployments.retrieve(name=deployment_name, workspace="default")\n elapsed = time.time() - start\n clear_output(wait=True)\n print(f"Deployment: {deployment_name}")\n print(f"Status: {dep.status}")\n print(f"Elapsed: {int(elapsed // 60)}m {int(elapsed % 60)}s")\n\n if dep.status == "READY":\n print("\\nDeployment is ready!")\n remaining = int(timeout - elapsed)\n if remaining <= 0:\n raise TimeoutError(f"Deployment timeout after {timeout_minutes} minutes")\n if not client.models.wait_for_status(\n deployment_name=deployment_name,\n desired_status="READY",\n workspace="default",\n timeout=remaining,\n check_gateway=True,\n ):\n raise TimeoutError("Inference gateway did not become ready")\n return dep\n if dep.status in ("FAILED", "ERROR", "TERMINATED", "LOST"):\n raise RuntimeError(f"Deployment failed with status: {dep.status}")\n if elapsed > timeout:\n raise TimeoutError(f"Deployment timeout after {timeout_minutes} minutes")\n time.sleep(15)\n\n\ntry:\n dep_status = wait_for_deployment(BASELINE_DEPLOYMENT_NAME)\n assert dep_status.status == "READY"\nexcept Exception:\n # Free GPUs if readiness fails before the later baseline-cleanup cell runs.\n try:\n client.inference.deployments.delete(name=BASELINE_DEPLOYMENT_NAME, workspace="default")\n client.inference.deployment_configs.delete(name=BASELINE_DEPLOYMENT_CONFIG, workspace="default")\n except Exception as cleanup_error:\n print(f"Baseline cleanup after readiness failure also failed: {cleanup_error}")\n raise\n" + "source_html": "def wait_for_deployment(deployment_name: str, timeout_minutes: int = 30):\n """Poll deployment until ready."""\n start = time.time()\n timeout = timeout_minutes * 60\n while True:\n dep = client.inference.deployments.retrieve(name=deployment_name, workspace="default")\n elapsed = time.time() - start\n clear_output(wait=True)\n print(f"Deployment: {deployment_name}")\n print(f"Status: {dep.status}")\n print(f"Elapsed: {int(elapsed // 60)}m {int(elapsed % 60)}s")\n\n if dep.status == "READY":\n print("\\nDeployment is ready!")\n remaining = int(timeout - elapsed)\n if remaining <= 0:\n raise TimeoutError(f"Deployment timeout after {timeout_minutes} minutes")\n if not client.models.wait_for_status(\n deployment_name=deployment_name,\n desired_status="READY",\n workspace="default",\n timeout=remaining,\n check_gateway=True,\n ):\n raise TimeoutError("Inference gateway did not become ready")\n return dep\n if dep.status in ("FAILED", "ERROR", "TERMINATED", "LOST"):\n raise RuntimeError(f"Deployment failed with status: {dep.status}")\n if elapsed > timeout:\n raise TimeoutError(f"Deployment timeout after {timeout_minutes} minutes")\n time.sleep(15)\n\n\ntry:\n dep_status = wait_for_deployment(BASELINE_DEPLOYMENT_NAME)\n assert dep_status.status == "READY"\nexcept Exception:\n # Free GPUs if readiness fails before the later baseline-cleanup cell runs.\n try:\n client.inference.deployments.delete(name=BASELINE_DEPLOYMENT_NAME, workspace="default")\n if not client.models.wait_for_status(\n deployment_name=BASELINE_DEPLOYMENT_NAME,\n desired_status="DELETED",\n workspace="default",\n timeout=600,\n ):\n raise TimeoutError(\n f"Deployment {BASELINE_DEPLOYMENT_NAME} was not deleted within timeout"\n )\n client.inference.deployment_configs.delete(name=BASELINE_DEPLOYMENT_CONFIG, workspace="default")\n except Exception as cleanup_error:\n print(f"Baseline cleanup after readiness failure also failed: {cleanup_error}")\n raise\n" }, { "type": "markdown", From c8f2249de4b4538f0d7606037d4ebea810e2ade3 Mon Sep 17 00:00:00 2001 From: "Paul A. Parkanzky" <42279121+parkanzky@users.noreply.github.com> Date: Tue, 21 Jul 2026 17:06:31 -0400 Subject: [PATCH 05/24] test(auditor): E2E tests for Auditor (#794) * tests(auditor): E2E tests for Auditor Signed-off-by: Paul A. Parkanzky * link async submit Signed-off-by: Paul A. Parkanzky * re-add garak install to docker file Signed-off-by: Paul A. Parkanzky * rm superfluous uv sync Signed-off-by: Paul A. Parkanzky * switch back to default cpu/cpu provider/profile Signed-off-by: Paul A. Parkanzky * skip tests that depend on future image Signed-off-by: Paul A. Parkanzky --------- Signed-off-by: Paul A. Parkanzky Co-authored-by: Paul A. Parkanzky Signed-off-by: Sam Oluwalana --- docker/Dockerfile.auditor-tasks | 3 +- e2e/auditor/conftest.py | 13 + e2e/auditor/test_audit_job.py | 236 ++++++++++++++++++ e2e/auditor/test_cli.py | 112 +++++++++ e2e/auditor/test_configs.py | 133 ++++++++++ e2e/auditor/test_healthz.py | 21 ++ e2e/auditor/test_targets.py | 103 ++++++++ e2e/auditor/utils.py | 33 +++ .../src/nemo_auditor/jobs/audit.py | 2 +- plugins/nemo-auditor/src/nemo_auditor/sdk.py | 107 ++++++++ .../nemo-auditor/tests/test_sdk_resources.py | 125 ++++++++++ 11 files changed, 885 insertions(+), 3 deletions(-) create mode 100644 e2e/auditor/conftest.py create mode 100644 e2e/auditor/test_audit_job.py create mode 100644 e2e/auditor/test_cli.py create mode 100644 e2e/auditor/test_configs.py create mode 100644 e2e/auditor/test_healthz.py create mode 100644 e2e/auditor/test_targets.py create mode 100644 e2e/auditor/utils.py diff --git a/docker/Dockerfile.auditor-tasks b/docker/Dockerfile.auditor-tasks index ab180bd9a4..c7f61e0b32 100644 --- a/docker/Dockerfile.auditor-tasks +++ b/docker/Dockerfile.auditor-tasks @@ -63,8 +63,7 @@ RUN uv venv --no-project /app/.garak_venv # litellm: CVE-2026-35029 (affects <1.83.0) # langchain-core: CVE-2025-68664 (affects <1.2.5) + GHSA-qh6h-p6c9-ff54 (affects <1.2.22) RUN --mount=type=cache,target=/root/.cache/uv \ - VIRTUAL_ENV=/app/.garak_venv \ - uv sync --frozen --active --project /app/plugins/nemo-auditor/third-party --no-dev --no-editable && \ + uv pip --no-config install garak==0.15.1 --python /app/.garak_venv/bin/python && \ uv pip install --python /app/.garak_venv/bin/python "setuptools>=78.1.1" && \ uv pip install --no-deps --python /app/.garak_venv/bin/python \ "litellm>=1.83.0" \ diff --git a/e2e/auditor/conftest.py b/e2e/auditor/conftest.py new file mode 100644 index 0000000000..2fba7144e2 --- /dev/null +++ b/e2e/auditor/conftest.py @@ -0,0 +1,13 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Fixtures for auditor plugin e2e tests.""" + +import pytest +from nemo_platform import NeMoPlatform + + +@pytest.fixture +def auditor_url(sdk: NeMoPlatform) -> str: + """Root URL for raw httpx calls to the auditor plugin (filter/sort params not in SDK).""" + return str(sdk.base_url).rstrip("/") + "/apis/auditor" diff --git a/e2e/auditor/test_audit_job.py b/e2e/auditor/test_audit_job.py new file mode 100644 index 0000000000..a1ef35d9e7 --- /dev/null +++ b/e2e/auditor/test_audit_job.py @@ -0,0 +1,236 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""K8s-only E2E tests for auditor job submission. + +These tests submit real audit jobs and poll for completion. They require: + - ``NMP_BASE_URL`` pointing at a K8s deployment (set via ``container_only`` marker) + - garak installed at ``/app/.garak_venv/bin/python`` in the auditor-tasks image + - mock inference provider support (``mock_provider_prefix: igw-mock-`` in Helm values) + +The probe used is ``test.Test`` — garak's single-message blank probe, which is the +fastest possible smoke run and does not require a real safety-relevant model. +""" + +from __future__ import annotations + +import time +from collections.abc import Iterator +from contextlib import suppress + +import pytest +from nemo_platform import NeMoPlatform +from nmp.testing import add_mock_provider, short_unique_name + +from e2e.auditor.utils import minimal_audit_config, unique_name + +pytestmark = [ + pytest.mark.container_only, + pytest.mark.timeout(1800), +] + +AUDIT_JOB_TIMEOUT_SECONDS = 900.0 +AUDIT_JOB_POLL_INTERVAL_SECONDS = 10.0 +TERMINAL_STATUSES = frozenset({"completed", "error", "failed", "cancelled"}) + + +def _chat_completion(content: str = "I'm happy to help!") -> dict: + return { + "id": "chatcmpl-audit-e2e", + "object": "chat.completion", + "model": "audit-mock", + "choices": [{"index": 0, "message": {"role": "assistant", "content": content}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 5, "completion_tokens": 5, "total_tokens": 10}, + } + + +def _wait_for_audit_job(sdk: NeMoPlatform, job_name: str, workspace: str) -> str: + deadline = time.monotonic() + AUDIT_JOB_TIMEOUT_SECONDS + while time.monotonic() < deadline: + status_resp = sdk.jobs.get_status(name=job_name, workspace=workspace) + status = str(status_resp.status) + if status in TERMINAL_STATUSES: + return status + time.sleep(AUDIT_JOB_POLL_INTERVAL_SECONDS) + raise TimeoutError(f"Audit job {job_name!r} did not complete within {AUDIT_JOB_TIMEOUT_SECONDS}s") + + +def _cleanup_audit_job(sdk: NeMoPlatform, job_name: str, workspace: str) -> None: + with suppress(Exception): + sdk.jobs.cancel(name=job_name, workspace=workspace) + with suppress(Exception): + sdk.jobs.delete(name=job_name, workspace=workspace) + + +def _add_mock_provider_or_skip(sdk: NeMoPlatform, workspace: str, name: str) -> str: + """Create a mock inference provider, skipping the test if the deployment doesn't support one.""" + try: + provider = add_mock_provider( + sdk, + workspace=workspace, + name=name, + mock_response_body=_chat_completion(), + ) + return provider.name + except RuntimeError as exc: + if "mock_provider_prefix is not configured" in str(exc): + pytest.skip( + "The running platform does not have mock-provider mode enabled. " + "Set mock_provider_prefix: igw-mock- in Helm values (already present in " + "e2e/k8s/values/minikube.yaml) to run this test." + ) + raise + + +# ---- Module-scoped fixtures for the shared K8s workspace and mock provider ---- + + +@pytest.fixture(scope="module") +def audit_workspace(sdk: NeMoPlatform) -> Iterator[str]: + name = short_unique_name("e2e-audit") + sdk.workspaces.create(name=name) + try: + yield name + finally: + with suppress(Exception): + sdk.workspaces.delete(name) + + +@pytest.fixture(scope="module") +def mock_provider_name(sdk: NeMoPlatform, audit_workspace: str) -> str: + """Create a canned-response mock provider for the module; workspace deletion cascades cleanup.""" + provider_name = short_unique_name("audit-mock") + return _add_mock_provider_or_skip(sdk, audit_workspace, provider_name) + + +@pytest.fixture(scope="module") +def audit_config_name(sdk: NeMoPlatform, audit_workspace: str) -> Iterator[str]: + name = short_unique_name("e2e-audit-cfg") + sdk.auditor.configs.create( + workspace=audit_workspace, + name=name, + **minimal_audit_config(plugins={"probe_spec": "test.Test", "detector_spec": "auto"}), + ) + try: + yield name + finally: + with suppress(Exception): + sdk.auditor.configs.delete(workspace=audit_workspace, name=name) + + +@pytest.fixture(scope="module") +def audit_target_name(sdk: NeMoPlatform, audit_workspace: str, mock_provider_name: str) -> Iterator[str]: + name = short_unique_name("e2e-audit-tgt") + sdk.auditor.targets.create( + workspace=audit_workspace, + name=name, + type="openai", + model=mock_provider_name, + options={ + "openai": { + "OpenAICompatible": { + "nmp_uri_spec": { + "inference_gateway": { + "workspace": audit_workspace, + "provider": mock_provider_name, + } + } + } + } + }, + ) + try: + yield name + finally: + with suppress(Exception): + sdk.auditor.targets.delete(workspace=audit_workspace, name=name) + + +# ---- Tests ---- + + +@pytest.mark.skip("re-enable after auditor image rebuilt") +def test_audit_job_submit_blank_probe( + sdk: NeMoPlatform, + audit_workspace: str, + mock_provider_name: str, +) -> None: + """Submit an inline audit job with test.Test probe and verify it reaches completed status.""" + config = { + **minimal_audit_config(plugins={"probe_spec": "test.Test", "detector_spec": "auto"}), + "name": unique_name("inline-cfg"), + "workspace": audit_workspace, + } + target = { + "name": unique_name("inline-tgt"), + "workspace": audit_workspace, + "type": "openai", + "model": mock_provider_name, + "options": { + "openai": { + "OpenAICompatible": { + "nmp_uri_spec": { + "inference_gateway": { + "workspace": audit_workspace, + "provider": mock_provider_name, + } + } + } + } + }, + } + + job = sdk.auditor.submit(config=config, target=target, workspace=audit_workspace) + job_name = job["name"] + try: + final_status = _wait_for_audit_job(sdk, job_name, audit_workspace) + assert final_status == "completed", ( + f"Audit job {job_name!r} ended with status {final_status!r} instead of 'completed'. " + "Check that garak is installed at /app/.garak_venv/bin/python in the auditor-tasks image." + ) + finally: + _cleanup_audit_job(sdk, job_name, audit_workspace) + + +@pytest.mark.skip("re-enable after auditor image rebuilt") +def test_audit_job_submit_with_entity_refs( + sdk: NeMoPlatform, + audit_workspace: str, + audit_config_name: str, + audit_target_name: str, +) -> None: + """Submit an audit job using stored entity name references and verify completion.""" + job = sdk.auditor.submit( + config=f"{audit_workspace}/{audit_config_name}", + target=f"{audit_workspace}/{audit_target_name}", + workspace=audit_workspace, + ) + job_name = job["name"] + try: + final_status = _wait_for_audit_job(sdk, job_name, audit_workspace) + assert final_status == "completed", ( + f"Audit job {job_name!r} with entity refs ended with status {final_status!r}." + ) + finally: + _cleanup_audit_job(sdk, job_name, audit_workspace) + + +def test_audit_job_appears_in_list( + sdk: NeMoPlatform, + audit_workspace: str, + audit_config_name: str, + audit_target_name: str, +) -> None: + """Submitted audit job appears in list_jobs() with its name.""" + job = sdk.auditor.submit( + config=f"{audit_workspace}/{audit_config_name}", + target=f"{audit_workspace}/{audit_target_name}", + workspace=audit_workspace, + ) + job_name = job["name"] + try: + jobs = sdk.auditor.list_jobs(workspace=audit_workspace) + job_names = [j["name"] for j in jobs.get("data", [])] + assert job_name in job_names, f"Submitted job {job_name!r} not found in list_jobs(): {job_names}" + finally: + _cleanup_audit_job(sdk, job_name, audit_workspace) diff --git a/e2e/auditor/test_cli.py b/e2e/auditor/test_cli.py new file mode 100644 index 0000000000..cea08b4ae4 --- /dev/null +++ b/e2e/auditor/test_cli.py @@ -0,0 +1,112 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CLI smoke tests for the auditor plugin. + +Exercises ``nemo auditor configs`` and ``nemo auditor targets`` subcommands +against a live platform via ``NMP_BASE_URL``, using ``run_nemo_local`` so the +CLI picks up platform credentials from env without reading the developer's own +config file. +""" + +import json + +from nemo_platform import NeMoPlatform +from nmp.testing import assert_exit_0, run_nemo_local + +from e2e.auditor.utils import minimal_audit_config, minimal_audit_target, unique_name + + +def test_cli_config_create_list_delete(sdk: NeMoPlatform, workspace: str) -> None: + name = unique_name("cli-cfg") + base_url = str(sdk.base_url) + + result = run_nemo_local( + "auditor", + "configs", + "create", + name, + "--data", + json.dumps(minimal_audit_config(description="cli smoke")), + "--workspace", + workspace, + base_url=base_url, + ) + assert_exit_0(result, "CLI create config") + assert name in result.stdout + + result = run_nemo_local("auditor", "configs", "list", "--workspace", workspace, base_url=base_url) + assert_exit_0(result, "CLI list configs") + assert name in result.stdout + + result = run_nemo_local("auditor", "configs", "delete", name, "--workspace", workspace, base_url=base_url) + assert_exit_0(result, "CLI delete config") + assert name in result.stdout + + result = run_nemo_local("auditor", "configs", "list", "--workspace", workspace, base_url=base_url) + assert_exit_0(result, "CLI list after delete") + listed = json.loads(result.stdout) + assert all(item["name"] != name for item in listed["data"]) + + +def test_cli_target_create_list_delete(sdk: NeMoPlatform, workspace: str) -> None: + name = unique_name("cli-tgt") + base_url = str(sdk.base_url) + + result = run_nemo_local( + "auditor", + "targets", + "create", + name, + "--data", + json.dumps(minimal_audit_target(description="cli target smoke")), + "--workspace", + workspace, + base_url=base_url, + ) + assert_exit_0(result, "CLI create target") + assert name in result.stdout + + result = run_nemo_local("auditor", "targets", "list", "--workspace", workspace, base_url=base_url) + assert_exit_0(result, "CLI list targets") + assert name in result.stdout + + result = run_nemo_local("auditor", "targets", "delete", name, "--workspace", workspace, base_url=base_url) + assert_exit_0(result, "CLI delete target") + + +def test_cli_config_update(sdk: NeMoPlatform, workspace: str) -> None: + name = unique_name("cli-upd") + base_url = str(sdk.base_url) + + result = run_nemo_local( + "auditor", + "configs", + "create", + name, + "--data", + json.dumps(minimal_audit_config(description="before update")), + "--workspace", + workspace, + base_url=base_url, + ) + assert_exit_0(result, "CLI create config for update") + + updated_body = minimal_audit_config(description="after update") + result = run_nemo_local( + "auditor", + "configs", + "update", + name, + "--data", + json.dumps(updated_body), + "--workspace", + workspace, + base_url=base_url, + ) + assert_exit_0(result, "CLI update config") + + result = run_nemo_local("auditor", "configs", "get", name, "--workspace", workspace, base_url=base_url) + assert_exit_0(result, "CLI get updated config") + payload = json.loads(result.stdout) + assert payload["description"] == "after update" diff --git a/e2e/auditor/test_configs.py b/e2e/auditor/test_configs.py new file mode 100644 index 0000000000..cbea3e4623 --- /dev/null +++ b/e2e/auditor/test_configs.py @@ -0,0 +1,133 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""E2E tests for AuditConfig CRUD, filtering, and sorting. + +These tests exercise the SDK against the real platform without mocking the +entity store. Filter and sort parameters are not exposed by the SDK list() +method, so those tests use the underlying httpx client directly. +""" + +import httpx +import pytest +from nemo_platform import NeMoPlatform + +from e2e.auditor.utils import minimal_audit_config, unique_name + + +def _list_raw(sdk: NeMoPlatform, workspace: str, auditor_url: str, **params) -> dict: + """GET /configs with arbitrary query params (filter, sort) via raw httpx.""" + resp = sdk.auditor._http_client.get( + f"{auditor_url}/v2/workspaces/{workspace}/configs", + params=params, + ) + resp.raise_for_status() + return resp.json() + + +def test_config_create_and_get(sdk: NeMoPlatform, workspace: str) -> None: + name = unique_name("cfg-cg") + body = minimal_audit_config(description="create-and-get test") + + created = sdk.auditor.configs.create(workspace=workspace, name=name, **body) + + assert created.name == name + assert created.workspace == workspace + assert created.description == "create-and-get test" + assert created.plugins.probe_spec == "test.Test" + + # Note: id/created_at are not populated by the auditor SDK (raw httpx + model_validate + # cannot set the private _id PrivateAttr). Use name as the stable cross-call identifier. + retrieved = sdk.auditor.configs.get(workspace=workspace, name=name) + assert retrieved.name == name + assert retrieved.plugins.probe_spec == "test.Test" + + +def test_config_list_contains_created(sdk: NeMoPlatform, workspace: str) -> None: + name = unique_name("cfg-list") + body = minimal_audit_config(description="list test") + + sdk.auditor.configs.create(workspace=workspace, name=name, **body) + + page = sdk.auditor.configs.list(workspace=workspace, page_size=100) + names = [item["name"] for item in page["data"]] + assert name in names + + +def test_config_update(sdk: NeMoPlatform, workspace: str) -> None: + name = unique_name("cfg-upd") + body = minimal_audit_config(description="original description") + + sdk.auditor.configs.create(workspace=workspace, name=name, **body) + + updated_body = minimal_audit_config(description="updated description") + updated_body["plugins"]["probe_spec"] = "dan.Dan" + updated = sdk.auditor.configs.update(workspace=workspace, name=name, **updated_body) + + assert updated.name == name + assert updated.description == "updated description" + assert updated.plugins.probe_spec == "dan.Dan" + + retrieved = sdk.auditor.configs.get(workspace=workspace, name=name) + assert retrieved.plugins.probe_spec == "dan.Dan" + + +def test_config_delete(sdk: NeMoPlatform, workspace: str) -> None: + name = unique_name("cfg-del") + sdk.auditor.configs.create(workspace=workspace, name=name, **minimal_audit_config()) + + names_before = [item["name"] for item in sdk.auditor.configs.list(workspace=workspace, page_size=100)["data"]] + assert name in names_before + + sdk.auditor.configs.delete(workspace=workspace, name=name) + + # Auditor SDK uses raw httpx, so errors surface as httpx.HTTPStatusError (not nemo_platform exceptions). + with pytest.raises(httpx.HTTPStatusError) as exc_info: + sdk.auditor.configs.get(workspace=workspace, name=name) + assert exc_info.value.response.status_code == 404 + + +def test_config_duplicate_name_returns_conflict(sdk: NeMoPlatform, workspace: str) -> None: + name = unique_name("cfg-dup") + body = minimal_audit_config() + + sdk.auditor.configs.create(workspace=workspace, name=name, **body) + + with pytest.raises(httpx.HTTPStatusError) as exc_info: + sdk.auditor.configs.create(workspace=workspace, name=name, **body) + assert exc_info.value.response.status_code == 409 + + +def test_config_get_nonexistent_returns_404(sdk: NeMoPlatform, workspace: str) -> None: + with pytest.raises(httpx.HTTPStatusError) as exc_info: + sdk.auditor.configs.get(workspace=workspace, name="does-not-exist-xyzzy") + assert exc_info.value.response.status_code == 404 + + +def test_config_filter_by_description(sdk: NeMoPlatform, workspace: str, auditor_url: str) -> None: + needle = unique_name("cfg-filter-needle") + other = unique_name("cfg-filter-other") + + sdk.auditor.configs.create(workspace=workspace, name=needle, **minimal_audit_config(description="needle-desc")) + sdk.auditor.configs.create(workspace=workspace, name=other, **minimal_audit_config(description="other-desc")) + + result = _list_raw(sdk, workspace, auditor_url, **{"filter[description]": "needle-desc"}) + names = [item["name"] for item in result["data"]] + + assert needle in names + assert other not in names + + +def test_config_sort_descending(sdk: NeMoPlatform, workspace: str, auditor_url: str) -> None: + first = unique_name("cfg-sort-a") + second = unique_name("cfg-sort-b") + + sdk.auditor.configs.create(workspace=workspace, name=first, **minimal_audit_config()) + sdk.auditor.configs.create(workspace=workspace, name=second, **minimal_audit_config()) + + result = _list_raw(sdk, workspace, auditor_url, sort="-created_at", page_size=10) + names = [item["name"] for item in result["data"]] + + assert names.index(second) < names.index(first), ( + f"Expected {second!r} (newer) before {first!r} (older) in sort=-created_at result; got order: {names}" + ) diff --git a/e2e/auditor/test_healthz.py b/e2e/auditor/test_healthz.py new file mode 100644 index 0000000000..e99b45d892 --- /dev/null +++ b/e2e/auditor/test_healthz.py @@ -0,0 +1,21 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""E2E smoke test for the auditor plugin healthz endpoint. + +Verifies that the auditor plugin loaded correctly in the running platform and +that its healthz response contains the expected keys. A 404 here means the +plugin failed to initialize. +""" + +from nemo_platform import NeMoPlatform + + +def test_auditor_plugin_status(sdk: NeMoPlatform) -> None: + status = sdk.auditor.plugin_status() + + assert status["plugin"] == "auditor" + assert status["status"] == "ok" + assert "auditor.audit" in status["jobs"] + assert "auditor_audit_config" in status["entities"] + assert "auditor_audit_target" in status["entities"] diff --git a/e2e/auditor/test_targets.py b/e2e/auditor/test_targets.py new file mode 100644 index 0000000000..afbde0f727 --- /dev/null +++ b/e2e/auditor/test_targets.py @@ -0,0 +1,103 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""E2E tests for AuditTarget CRUD and filtering.""" + +import httpx +import pytest +from nemo_platform import NeMoPlatform + +from e2e.auditor.utils import minimal_audit_target, unique_name + + +def _list_raw(sdk: NeMoPlatform, workspace: str, auditor_url: str, **params) -> dict: + resp = sdk.auditor._http_client.get( + f"{auditor_url}/v2/workspaces/{workspace}/targets", + params=params, + ) + resp.raise_for_status() + return resp.json() + + +def test_target_create_and_get(sdk: NeMoPlatform, workspace: str) -> None: + name = unique_name("tgt-cg") + body = minimal_audit_target(description="create-and-get target", type="nim", model="meta/llama-3.1-8b-instruct") + + created = sdk.auditor.targets.create(workspace=workspace, name=name, **body) + + assert created.name == name + assert created.workspace == workspace + assert created.description == "create-and-get target" + assert created.type == "nim" + assert created.model == "meta/llama-3.1-8b-instruct" + + # Note: id/created_at are not populated by the auditor SDK (raw httpx + model_validate + # cannot set the private _id PrivateAttr). Use name for cross-call identity checks. + retrieved = sdk.auditor.targets.get(workspace=workspace, name=name) + assert retrieved.name == name + assert retrieved.type == "nim" + assert retrieved.model == "meta/llama-3.1-8b-instruct" + + +def test_target_list_contains_created(sdk: NeMoPlatform, workspace: str) -> None: + name = unique_name("tgt-list") + + sdk.auditor.targets.create(workspace=workspace, name=name, **minimal_audit_target()) + + page = sdk.auditor.targets.list(workspace=workspace, page_size=100) + names = [item["name"] for item in page["data"]] + assert name in names + + +def test_target_update(sdk: NeMoPlatform, workspace: str) -> None: + name = unique_name("tgt-upd") + + sdk.auditor.targets.create( + workspace=workspace, + name=name, + **minimal_audit_target(description="original", model="gpt-4o-mini"), + ) + + updated = sdk.auditor.targets.update( + workspace=workspace, + name=name, + **minimal_audit_target(description="updated", model="gpt-4o"), + ) + + assert updated.name == name + assert updated.description == "updated" + assert updated.model == "gpt-4o" + + +def test_target_delete(sdk: NeMoPlatform, workspace: str) -> None: + name = unique_name("tgt-del") + sdk.auditor.targets.create(workspace=workspace, name=name, **minimal_audit_target()) + + sdk.auditor.targets.delete(workspace=workspace, name=name) + + # Auditor SDK uses raw httpx, so errors surface as httpx.HTTPStatusError. + with pytest.raises(httpx.HTTPStatusError) as exc_info: + sdk.auditor.targets.get(workspace=workspace, name=name) + assert exc_info.value.response.status_code == 404 + + +def test_target_filter_by_type(sdk: NeMoPlatform, workspace: str, auditor_url: str) -> None: + nim_name = unique_name("tgt-nim") + openai_name = unique_name("tgt-oai") + + sdk.auditor.targets.create( + workspace=workspace, + name=nim_name, + **minimal_audit_target(type="nim", model="meta/llama-3.1-8b-instruct"), + ) + sdk.auditor.targets.create( + workspace=workspace, + name=openai_name, + **minimal_audit_target(type="openai", model="gpt-4o-mini"), + ) + + result = _list_raw(sdk, workspace, auditor_url, **{"filter[type]": "nim"}) + names = [item["name"] for item in result["data"]] + + assert nim_name in names + assert openai_name not in names diff --git a/e2e/auditor/utils.py b/e2e/auditor/utils.py new file mode 100644 index 0000000000..dceb233820 --- /dev/null +++ b/e2e/auditor/utils.py @@ -0,0 +1,33 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared helpers for auditor plugin e2e tests.""" + +import uuid + + +def unique_name(prefix: str) -> str: + return f"{prefix}-{uuid.uuid4().hex[:8]}" + + +def minimal_audit_config(**overrides) -> dict: + """Minimal AuditConfig body for CRUD tests (no name/workspace — SDK takes those separately).""" + base: dict = { + "system": {}, + "run": {"generations": 1}, + "plugins": {"probe_spec": "test.Test", "detector_spec": "auto"}, + "reporting": {}, + } + base.update(overrides) + return base + + +def minimal_audit_target(**overrides) -> dict: + """Minimal AuditTarget body (fake endpoint — safe for CRUD tests that never invoke garak).""" + base: dict = { + "type": "openai", + "model": "gpt-4o-mini", + "options": {}, + } + base.update(overrides) + return base diff --git a/plugins/nemo-auditor/src/nemo_auditor/jobs/audit.py b/plugins/nemo-auditor/src/nemo_auditor/jobs/audit.py index a69f60c588..00e0c8a828 100644 --- a/plugins/nemo-auditor/src/nemo_auditor/jobs/audit.py +++ b/plugins/nemo-auditor/src/nemo_auditor/jobs/audit.py @@ -424,7 +424,7 @@ async def compile( PlatformJobStep( name="audit-job", executor=CPUExecutionProviderSpec( - profile=profile or "auditor", + profile=profile or "default", provider="cpu", container=ContainerSpec( image=get_qualified_image("auditor-tasks"), diff --git a/plugins/nemo-auditor/src/nemo_auditor/sdk.py b/plugins/nemo-auditor/src/nemo_auditor/sdk.py index 9a52e3fa68..6ad3be6a36 100644 --- a/plugins/nemo-auditor/src/nemo_auditor/sdk.py +++ b/plugins/nemo-auditor/src/nemo_auditor/sdk.py @@ -9,6 +9,10 @@ - ``client.auditor.plugin_status()`` — service healthz check. - ``client.auditor.configs.{create,list,get,update,delete}`` — ``AuditConfig`` CRUD. - ``client.auditor.targets.{create,list,get,update,delete}`` — ``AuditTarget`` CRUD. +- ``client.auditor.submit(config=..., target=..., workspace=...)`` — submit a K8s + audit job through the plugin's job endpoint and return the raw job dict. +- ``client.auditor.list_jobs(workspace=...)`` — list submitted audit jobs. +- ``client.auditor.get_job(job_name, workspace=...)`` — fetch a single audit job. - ``client.auditor.run(config=..., target=..., workspace=...)`` — in-process audit using :class:`~nemo_auditor.jobs.audit.AuditJob`. Mirrors the evaluator plugin's ``client.evaluator.run`` pattern: delegates to @@ -61,6 +65,60 @@ def targets(self) -> _TargetResource: self._targets = _TargetResource(self) return self._targets + def submit( + self, + *, + config: AuditConfig | str, + target: AuditTarget | str, + workspace: str | None = None, + max_probe_retries: int = 0, + fail_job_on_retries_exhausted: bool = True, + ) -> dict: + """Submit an audit job to the K8s executor via the plugin job endpoint. + + Returns the raw job dict (name, status, workspace, …). Use + ``sdk.jobs.get_status(name=result["name"], workspace=workspace)`` to poll + for completion, or pass the name to ``sdk.auditor.get_job()``. + """ + ws = workspace or "default" + spec = AuditInputSpec( + config=config, + target=target, + max_probe_retries=max_probe_retries, + fail_job_on_retries_exhausted=fail_job_on_retries_exhausted, + ) + response = self._http_client.post( + self._url(f"/v2/workspaces/{ws}/jobs/audit"), + json={"spec": spec.model_dump(mode="json")}, + ) + response.raise_for_status() + return response.json() + + def list_jobs( + self, + *, + workspace: str | None = None, + page: int = 1, + page_size: int = 20, + ) -> dict: + """List audit jobs in the workspace with basic pagination.""" + ws = workspace or "default" + response = self._http_client.get( + self._url(f"/v2/workspaces/{ws}/jobs/audit"), + params={"page": page, "page_size": page_size}, + ) + response.raise_for_status() + return response.json() + + def get_job(self, job_name: str, *, workspace: str | None = None) -> dict: + """Fetch a single audit job by name.""" + ws = workspace or "default" + response = self._http_client.get( + self._url(f"/v2/workspaces/{ws}/jobs/audit/{job_name}"), + ) + response.raise_for_status() + return response.json() + def run( self, *, @@ -133,6 +191,55 @@ def targets(self) -> _AsyncTargetResource: self._targets = _AsyncTargetResource(self) return self._targets + async def submit( + self, + *, + config: AuditConfig | str, + target: AuditTarget | str, + workspace: str | None = None, + max_probe_retries: int = 0, + fail_job_on_retries_exhausted: bool = True, + ) -> dict: + """Async twin of :meth:`AuditorPluginResource.submit`.""" + ws = workspace or "default" + spec = AuditInputSpec( + config=config, + target=target, + max_probe_retries=max_probe_retries, + fail_job_on_retries_exhausted=fail_job_on_retries_exhausted, + ) + response = await self._http_client.post( + self._url(f"/v2/workspaces/{ws}/jobs/audit"), + json={"spec": spec.model_dump(mode="json")}, + ) + response.raise_for_status() + return response.json() + + async def list_jobs( + self, + *, + workspace: str | None = None, + page: int = 1, + page_size: int = 20, + ) -> dict: + """Async twin of :meth:`AuditorPluginResource.list_jobs`.""" + ws = workspace or "default" + response = await self._http_client.get( + self._url(f"/v2/workspaces/{ws}/jobs/audit"), + params={"page": page, "page_size": page_size}, + ) + response.raise_for_status() + return response.json() + + async def get_job(self, job_name: str, *, workspace: str | None = None) -> dict: + """Async twin of :meth:`AuditorPluginResource.get_job`.""" + ws = workspace or "default" + response = await self._http_client.get( + self._url(f"/v2/workspaces/{ws}/jobs/audit/{job_name}"), + ) + response.raise_for_status() + return response.json() + async def run( self, *, diff --git a/plugins/nemo-auditor/tests/test_sdk_resources.py b/plugins/nemo-auditor/tests/test_sdk_resources.py index 0cb61def23..b4f2c226ad 100644 --- a/plugins/nemo-auditor/tests/test_sdk_resources.py +++ b/plugins/nemo-auditor/tests/test_sdk_resources.py @@ -373,6 +373,131 @@ def test_workspace_qualified_name_parses_workspace_from_string(self, scheduler_c assert scheduler.run_local.call_args.kwargs["workspace"] == "default" +# --------------------------------------------------------------------------- +# Job submission: submit / list_jobs / get_job +# --------------------------------------------------------------------------- + +_JOB_PAYLOAD = { + "name": "audit-job-abc123", + "workspace": "default", + "status": "created", +} + +_JOBS_LIST_PAYLOAD = { + "data": [_JOB_PAYLOAD], + "pagination": {"page": 1, "page_size": 20, "total_pages": 1, "total_results": 1}, +} + + +class TestSyncJobMethods: + def test_submit_with_string_refs_posts_correct_url_and_body(self) -> None: + platform = _SyncPlatform() + platform._client.post.return_value = _ok_response(_JOB_PAYLOAD, status_code=201) + resource = AuditorPluginResource(cast(NeMoPlatform, platform)) + + result = resource.submit(config="ws/my-cfg", target="ws/my-tgt", workspace="ws") + + assert result == _JOB_PAYLOAD + platform._client.post.assert_called_once() + url = platform._client.post.call_args.args[0] + body = platform._client.post.call_args.kwargs["json"] + assert url == "http://test:8000/apis/auditor/v2/workspaces/ws/jobs/audit" + assert body["spec"]["config"] == "ws/my-cfg" + assert body["spec"]["target"] == "ws/my-tgt" + assert body["spec"]["max_probe_retries"] == 0 + assert body["spec"]["fail_job_on_retries_exhausted"] is True + + def test_submit_with_inline_entities_serialises_full_dict(self) -> None: + platform = _SyncPlatform() + platform._client.post.return_value = _ok_response(_JOB_PAYLOAD, status_code=201) + resource = AuditorPluginResource(cast(NeMoPlatform, platform)) + + cfg = AuditConfig(name="cfg-1", workspace="default") + tgt = AuditTarget(name="tgt-1", workspace="default", type="nim", model="llama") + resource.submit(config=cfg, target=tgt, workspace="default") + + body = platform._client.post.call_args.kwargs["json"] + assert isinstance(body["spec"]["config"], dict) + assert body["spec"]["config"]["name"] == "cfg-1" + assert isinstance(body["spec"]["target"], dict) + assert body["spec"]["target"]["model"] == "llama" + + def test_submit_defaults_workspace_to_default(self) -> None: + platform = _SyncPlatform() + platform._client.post.return_value = _ok_response(_JOB_PAYLOAD, status_code=201) + resource = AuditorPluginResource(cast(NeMoPlatform, platform)) + + resource.submit(config="my-cfg", target="my-tgt") + + url = platform._client.post.call_args.args[0] + assert "/workspaces/default/" in url + + def test_list_jobs_hits_collection_url_with_pagination(self) -> None: + platform = _SyncPlatform() + platform._client.get.return_value = _ok_response(_JOBS_LIST_PAYLOAD) + resource = AuditorPluginResource(cast(NeMoPlatform, platform)) + + result = resource.list_jobs(workspace="ws", page=2, page_size=5) + + assert result == _JOBS_LIST_PAYLOAD + url = platform._client.get.call_args.args[0] + params = platform._client.get.call_args.kwargs["params"] + assert url == "http://test:8000/apis/auditor/v2/workspaces/ws/jobs/audit" + assert params == {"page": 2, "page_size": 5} + + def test_get_job_hits_named_url(self) -> None: + platform = _SyncPlatform() + platform._client.get.return_value = _ok_response(_JOB_PAYLOAD) + resource = AuditorPluginResource(cast(NeMoPlatform, platform)) + + result = resource.get_job("audit-job-abc123", workspace="ws") + + assert result == _JOB_PAYLOAD + platform._client.get.assert_called_once_with( + "http://test:8000/apis/auditor/v2/workspaces/ws/jobs/audit/audit-job-abc123", + ) + + +@pytest.mark.asyncio +class TestAsyncJobMethods: + async def test_submit_posts_correct_url_and_body(self) -> None: + platform = _AsyncPlatform() + platform._client.post.return_value = _ok_response(_JOB_PAYLOAD, status_code=201) + resource = AsyncAuditorPluginResource(cast(AsyncNeMoPlatform, platform)) + + result = await resource.submit(config="ws/my-cfg", target="ws/my-tgt", workspace="ws") + + assert result == _JOB_PAYLOAD + url = platform._client.post.call_args.args[0] + body = platform._client.post.call_args.kwargs["json"] + assert url == "http://test:8000/apis/auditor/v2/workspaces/ws/jobs/audit" + assert body["spec"]["config"] == "ws/my-cfg" + assert body["spec"]["target"] == "ws/my-tgt" + + async def test_list_jobs_hits_collection_url(self) -> None: + platform = _AsyncPlatform() + platform._client.get.return_value = _ok_response(_JOBS_LIST_PAYLOAD) + resource = AsyncAuditorPluginResource(cast(AsyncNeMoPlatform, platform)) + + result = await resource.list_jobs(workspace="ws") + + assert result == _JOBS_LIST_PAYLOAD + url = platform._client.get.call_args.args[0] + assert url == "http://test:8000/apis/auditor/v2/workspaces/ws/jobs/audit" + + async def test_get_job_hits_named_url(self) -> None: + platform = _AsyncPlatform() + platform._client.get.return_value = _ok_response(_JOB_PAYLOAD) + resource = AsyncAuditorPluginResource(cast(AsyncNeMoPlatform, platform)) + + result = await resource.get_job("audit-job-abc123", workspace="ws") + + assert result == _JOB_PAYLOAD + platform._client.get.assert_called_once_with( + "http://test:8000/apis/auditor/v2/workspaces/ws/jobs/audit/audit-job-abc123", + ) + + # --------------------------------------------------------------------------- # Async smoke tests # --------------------------------------------------------------------------- From e95985e7dd6bd2ed8ddc301ef390039eb177b8d8 Mon Sep 17 00:00:00 2001 From: Ryan Sadler <267728323+ironcommit@users.noreply.github.com> Date: Tue, 21 Jul 2026 15:51:14 -0700 Subject: [PATCH 06/24] Add local services lifecycle SDK (#734) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Extracts the local services process lifecycle out of the CLI into a standalone SDK module (local.services, local.process, local.transport) so that Python callers can programmatically start, stop, connect to, and manage NeMo Platform service instances without importing CLI internals. The key addition is first-class Unix domain socket (UDS) transport support across the entire stack — from the platform runner and uvicorn server binding, through inter-service HTTP clients (auth, health checks, SDK factory), to the jobs-launcher Go sidecar. Signed-off-by: Ryan S <267728323+ironcommit@users.noreply.github.com> Signed-off-by: Sam Oluwalana --- .../cli/commands/services/cli.py | 135 +-- .../nemo_platform_ext/cli/commands/setup.py | 34 +- .../nemo_platform_ext/local/_service_child.py | 30 + .../services/_process.py => local/process.py} | 235 ++-- .../src/nemo_platform_ext/local/services.py | 728 ++++++++++++ .../src/nemo_platform_ext/local/transport.py | 154 +++ .../tests/cli/commands/conftest.py | 2 +- .../tests/cli/commands/test_services.py | 85 +- .../cli/commands/test_services_lifecycle.py | 47 +- .../cli/commands/test_services_process.py | 258 ++-- .../tests/cli/commands/test_setup.py | 7 +- .../tests/local/test_config_environment.py | 196 ++++ .../tests/local/test_daemon_lifecycle.py | 503 ++++++++ .../tests/local/test_health_child.py | 340 ++++++ .../tests/local/test_port_socket.py | 165 +++ .../tests/local/test_services.py | 1039 +++++++++++++++++ .../tests/local/test_services_contract.py | 334 ++++++ .../tests/local/test_sidecar_integration.py | 231 ++++ .../tests/local/test_transport.py | 167 +++ .../nmp_common/src/nmp/common/auth/client.py | 20 +- .../src/nmp/common/auth/middleware.py | 8 +- .../nmp_common/src/nmp/common/config/base.py | 9 +- .../src/nmp/common/platform_endpoint.py | 96 ++ .../nmp_common/src/nmp/common/sdk_factory.py | 48 +- .../src/nmp/common/service/api/health.py | 124 +- .../nmp_common/src/nmp/common/service/base.py | 55 +- packages/nmp_common/tests/auth/test_client.py | 42 + .../tests/nmp_common/test_common_config.py | 25 + .../tests/nmp_common/test_common_service.py | 98 +- .../nmp_common/test_dependency_provider.py | 92 ++ .../tests/nmp_common/test_service_health.py | 156 +++ .../nmp_common/tests/sdk_factory/test_sdk.py | 36 + .../tests/test_platform_endpoint.py | 140 +++ .../src/nmp/platform_runner/config.py | 198 +++- .../src/nmp/platform_runner/registry.py | 4 + .../src/nmp/platform_runner/run.py | 28 +- .../src/nmp/platform_runner/server.py | 86 +- .../nmp_platform_runner/tests/test_config.py | 95 +- .../nmp_platform_runner/tests/test_health.py | 103 ++ .../nmp_platform_runner/tests/test_run.py | 4 +- .../nmp_platform_runner/tests/test_server.py | 97 ++ .../tests/test_sidecars.py | 213 ++++ .../cli/commands/services/cli.py | 135 +-- .../src/nemo_platform/cli/commands/setup.py | 34 +- .../src/nemo_platform/local/_service_child.py | 30 + .../services/_process.py => local/process.py} | 235 ++-- .../src/nemo_platform/local/services.py | 728 ++++++++++++ .../src/nemo_platform/local/transport.py | 154 +++ .../cli/commands/conftest.py | 2 +- .../cli/commands/test_services.py | 85 +- .../cli/commands/test_services_lifecycle.py | 47 +- .../cli/commands/test_services_process.py | 258 ++-- .../cli/commands/test_setup.py | 7 +- .../nemo_platform_ext/local/__init__.py | 15 + .../local/test_config_environment.py | 196 ++++ .../local/test_daemon_lifecycle.py | 503 ++++++++ .../local/test_health_child.py | 340 ++++++ .../local/test_port_socket.py | 165 +++ .../nemo_platform_ext/local/test_services.py | 1039 +++++++++++++++++ .../local/test_services_contract.py | 334 ++++++ .../local/test_sidecar_integration.py | 231 ++++ .../nemo_platform_ext/local/test_transport.py | 167 +++ services/core/jobs/jobs-launcher/cmd/otel.go | 175 ++- .../core/jobs/jobs-launcher/cmd/otel_test.go | 41 + services/core/jobs/jobs-launcher/cmd/run.go | 76 +- .../core/jobs/jobs-launcher/cmd/run_test.go | 132 +++ .../jobs-launcher/cmd/workload_auth_test.go | 6 +- services/core/jobs/jobs-launcher/go.mod | 1 - services/core/jobs/jobs-launcher/go.sum | 8 - .../jobs/jobs-launcher/nmpclient/client.go | 13 +- .../jobs-launcher/nmpclient/client_test.go | 50 + .../jobs/jobs-launcher/nmpclient/endpoint.go | 88 ++ .../jobs-launcher/nmpclient/endpoint_test.go | 122 ++ .../core/jobs/controllers/backends/base.py | 68 +- .../core/jobs/controllers/backends/docker.py | 14 +- .../controllers/backends/kubernetes/common.py | 14 +- .../jobs/controllers/backends/subprocess.py | 22 +- .../backends/subprocess_runtime.py | 108 +- .../core/jobs/tests/controllers/test_base.py | 36 + .../tests/controllers/test_docker_backend.py | 27 +- .../controllers/test_kubernetes_backend.py | 27 +- .../controllers/test_subprocess_backend.py | 38 + .../controllers/test_subprocess_runtime.py | 38 + 83 files changed, 11255 insertions(+), 1021 deletions(-) create mode 100644 packages/nemo_platform_ext/src/nemo_platform_ext/local/_service_child.py rename packages/nemo_platform_ext/src/nemo_platform_ext/{cli/commands/services/_process.py => local/process.py} (76%) create mode 100644 packages/nemo_platform_ext/src/nemo_platform_ext/local/services.py create mode 100644 packages/nemo_platform_ext/src/nemo_platform_ext/local/transport.py create mode 100644 packages/nemo_platform_ext/tests/local/test_config_environment.py create mode 100644 packages/nemo_platform_ext/tests/local/test_daemon_lifecycle.py create mode 100644 packages/nemo_platform_ext/tests/local/test_health_child.py create mode 100644 packages/nemo_platform_ext/tests/local/test_port_socket.py create mode 100644 packages/nemo_platform_ext/tests/local/test_services.py create mode 100644 packages/nemo_platform_ext/tests/local/test_services_contract.py create mode 100644 packages/nemo_platform_ext/tests/local/test_sidecar_integration.py create mode 100644 packages/nemo_platform_ext/tests/local/test_transport.py create mode 100644 packages/nmp_common/src/nmp/common/platform_endpoint.py create mode 100644 packages/nmp_common/tests/nmp_common/test_dependency_provider.py create mode 100644 packages/nmp_common/tests/nmp_common/test_service_health.py create mode 100644 packages/nmp_common/tests/test_platform_endpoint.py create mode 100644 packages/nmp_platform_runner/tests/test_health.py create mode 100644 packages/nmp_platform_runner/tests/test_sidecars.py create mode 100644 sdk/python/nemo-platform/src/nemo_platform/local/_service_child.py rename sdk/python/nemo-platform/src/nemo_platform/{cli/commands/services/_process.py => local/process.py} (76%) create mode 100644 sdk/python/nemo-platform/src/nemo_platform/local/services.py create mode 100644 sdk/python/nemo-platform/src/nemo_platform/local/transport.py create mode 100644 sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/local/__init__.py create mode 100644 sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/local/test_config_environment.py create mode 100644 sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/local/test_daemon_lifecycle.py create mode 100644 sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/local/test_health_child.py create mode 100644 sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/local/test_port_socket.py create mode 100644 sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/local/test_services.py create mode 100644 sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/local/test_services_contract.py create mode 100644 sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/local/test_sidecar_integration.py create mode 100644 sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/local/test_transport.py create mode 100644 services/core/jobs/jobs-launcher/cmd/otel_test.go create mode 100644 services/core/jobs/jobs-launcher/nmpclient/endpoint.go create mode 100644 services/core/jobs/jobs-launcher/nmpclient/endpoint_test.go diff --git a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/services/cli.py b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/services/cli.py index d5e4a7cf66..a4009af5d3 100644 --- a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/services/cli.py +++ b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/services/cli.py @@ -13,7 +13,8 @@ import httpx import typer -from nemo_platform_ext.cli.commands.services._process import ( +from nemo_platform_ext.cli.core.help_formatter import create_typer_app +from nemo_platform_ext.local.process import ( ForegroundInstanceError, InstanceAlreadyRunningError, InstanceDescriptor, @@ -24,7 +25,6 @@ check_port_available_for_start, compute_scope, format_port_conflict, - get_create_time, instance_log_bytes, is_instance_alive, list_instances, @@ -37,7 +37,7 @@ stop_instance, write_descriptor, ) -from nemo_platform_ext.cli.core.help_formatter import create_typer_app +from nmp.platform_runner.config import DEFAULT_LOCAL_SERVICES_BIND_HOST, PlatformAppConfig logger = logging.getLogger(__name__) @@ -45,7 +45,6 @@ _HEALTH_TIMEOUT_SECONDS = 60 _HEALTH_POLL_INTERVAL = 2.0 -_DEFAULT_HOST = "127.0.0.1" _DEFAULT_PORT = 8080 _DEFAULT_STOP_TIMEOUT = 30.0 @@ -60,7 +59,7 @@ def services_callback(ctx: typer.Context) -> None: for info in running: desc = info.descriptor assert desc is not None - typer.echo(f"\nRunning: {info.scope} (pid {desc.pid}, {desc.host}:{desc.port}, {desc.mode})") + typer.echo(f"\nRunning: {info.scope} (pid {desc.pid}, {desc.config.host}:{desc.config.port}, {desc.mode})") def _require_services_extra() -> None: @@ -92,7 +91,7 @@ def _parse_csv_option(value: str | None) -> list[str] | None: def _wait_for_healthy( host: str, port: int, - timeout: int = _HEALTH_TIMEOUT_SECONDS, + timeout: float = _HEALTH_TIMEOUT_SECONDS, poll_interval: float = _HEALTH_POLL_INTERVAL, ) -> bool: """Poll the platform status endpoint until it responds or timeout.""" @@ -124,15 +123,8 @@ def _effective_base_dir() -> str | None: def _find_sole_running_scope(base_dir: Path | None) -> str: - """Find the scope of the single running instance for this working directory. - - When the user runs ``restart`` without ``--instance`` or ``--port``, we - can't know which scope to target because the scope includes the port. - This function scans all running instances whose scope starts with the - same git-root hash prefix. If exactly one matches, return it. - Otherwise fall back to the default scope (hash-DEFAULT_PORT). - """ - prefix = compute_scope(port=0, instance_name=None).rsplit("-", 1)[0] + """Return the only running scope for this working directory, or the default scope.""" + prefix = compute_scope(port=0).rsplit("-", 1)[0] running = [i for i in list_instances(base_dir=base_dir) if i.alive and i.scope.startswith(prefix + "-")] if len(running) == 1: return running[0].scope @@ -213,7 +205,7 @@ def run_services( str | None, typer.Option("--config", help="Path to a platform configuration YAML file."), ] = None, - host: Annotated[str, typer.Option("--host", help="Host to bind to.")] = _DEFAULT_HOST, + host: Annotated[str, typer.Option("--host", help="Host to bind to.")] = DEFAULT_LOCAL_SERVICES_BIND_HOST, port: Annotated[int, typer.Option("--port", help="Port to bind to.")] = _DEFAULT_PORT, instance: Annotated[ str | None, @@ -226,7 +218,7 @@ def run_services( _require_services_extra() _warn_bind_all(host) - scope = compute_scope(port=port, instance_name=instance) + scope = compute_scope(port=port, explicit_scope=instance) base_dir_str = _effective_base_dir() base_dir = Path(base_dir_str) if base_dir_str else None @@ -242,20 +234,23 @@ def run_services( # "foreground", which protects interactive ``run`` sessions from being # killed by ``stop``. mode = "background" if os.environ.get("_NMP_LAUNCH_MODE") == "background" else "foreground" - - desc = InstanceDescriptor( - pid=os.getpid(), - scope=scope, - host=host, - port=port, - mode=mode, - create_time=get_create_time(os.getpid()), + platform_config = PlatformAppConfig( services=_parse_csv_option(services), - controllers=_parse_csv_option(controllers), service_group=service_group, + controllers=_parse_csv_option(controllers), controller_group=controller_group, sidecars=_parse_csv_option(sidecars), config_path=config, + scope=scope, + host=host, + port=port, + state_root=base_dir, + ) + + desc = InstanceDescriptor.from_config( + platform_config, + mode=mode, + pid=os.getpid(), ) write_descriptor(desc, base_dir=base_dir) @@ -269,14 +264,7 @@ def _cleanup() -> None: from nmp.platform_runner.run import run_platform run_platform( - services=_parse_csv_option(services), - service_group=service_group, - controllers=_parse_csv_option(controllers), - controller_group=controller_group, - sidecars=_parse_csv_option(sidecars), - config_path=config, - host=host, - port=port, + config=platform_config, on_shutdown=_cleanup, ) @@ -327,7 +315,7 @@ def start_services( str | None, typer.Option("--config", help="Path to a platform configuration YAML file."), ] = None, - host: Annotated[str, typer.Option("--host", help="Host to bind to.")] = _DEFAULT_HOST, + host: Annotated[str, typer.Option("--host", help="Host to bind to.")] = DEFAULT_LOCAL_SERVICES_BIND_HOST, port: Annotated[int, typer.Option("--port", help="Port to bind to.")] = _DEFAULT_PORT, instance: Annotated[ str | None, @@ -351,7 +339,7 @@ def start_services( raise typer.BadParameter("Cannot combine --controllers with --controller-group.") _warn_bind_all(host) - scope = compute_scope(port=port, instance_name=instance) + scope = compute_scope(port=port, explicit_scope=instance) base_dir_str = _effective_base_dir() base_dir = Path(base_dir_str) if base_dir_str else None @@ -360,20 +348,22 @@ def start_services( _ensure_port_available(host, port, scope, base_dir=base_dir) - typer.echo("Starting platform services...") - proc = start_background( - scope=scope, + platform_config = PlatformAppConfig( services=_parse_csv_option(services), service_group=service_group, controllers=_parse_csv_option(controllers), controller_group=controller_group, sidecars=_parse_csv_option(sidecars), config_path=config, + scope=scope, host=host, port=port, - base_dir=base_dir, + state_root=base_dir, ) + typer.echo("Starting platform services...") + proc = start_background(platform_config) + if not _wait_for_healthy(host, port): exit_code = proc.poll() if exit_code is not None: @@ -426,7 +416,7 @@ def stop_services_cmd( nemo services stop nemo services stop --timeout 60 """ - scope = compute_scope(port=port, instance_name=instance) + scope = compute_scope(port=port, explicit_scope=instance) base_dir_str = _effective_base_dir() base_dir = Path(base_dir_str) if base_dir_str else None @@ -496,7 +486,10 @@ def restart_services( ] = None, host: Annotated[ str | None, - typer.Option("--host", help="Host to bind to. Defaults to previous value or 127.0.0.1."), + typer.Option( + "--host", + help=f"Host to bind to. Defaults to previous value or {DEFAULT_LOCAL_SERVICES_BIND_HOST}.", + ), ] = None, port: Annotated[ int | None, @@ -529,8 +522,8 @@ def restart_services( base_dir = Path(base_dir_str) if base_dir_str else None if instance is not None or port is not None: - effective_port = port if port is not None else _DEFAULT_PORT - scope = compute_scope(port=effective_port, instance_name=instance) + effective_scope_port = port if port is not None else _DEFAULT_PORT + scope = compute_scope(port=effective_scope_port, explicit_scope=instance) else: scope = _find_sole_running_scope(base_dir) @@ -547,37 +540,47 @@ def restart_services( # appropriate even for foreground targets. stop_instance(scope, base_dir=base_dir, force=True) - effective_services = _parse_csv_option(services) if services is not None else (prev.services if prev else None) - effective_service_group = service_group if service_group is not None else (prev.service_group if prev else None) - effective_controllers = ( - _parse_csv_option(controllers) if controllers is not None else (prev.controllers if prev else None) + previous_config = prev.config if prev else None + effective_services = _parse_csv_option(services) if services is not None else None + if services is None and previous_config is not None: + effective_services = previous_config.services + effective_service_group = service_group if service_group is not None else None + if service_group is None and previous_config is not None: + effective_service_group = previous_config.service_group + effective_controllers = _parse_csv_option(controllers) if controllers is not None else None + if controllers is None and previous_config is not None: + effective_controllers = previous_config.controllers + effective_controller_group = controller_group if controller_group is not None else None + if controller_group is None and previous_config is not None: + effective_controller_group = previous_config.controller_group + effective_sidecars = _parse_csv_option(sidecars) if sidecars is not None else None + if sidecars is None and previous_config is not None: + effective_sidecars = previous_config.sidecars + effective_config = config if config is not None else (previous_config.config_path if previous_config else None) + effective_host = ( + host if host is not None else (previous_config.host if previous_config else DEFAULT_LOCAL_SERVICES_BIND_HOST) ) - effective_controller_group = ( - controller_group if controller_group is not None else (prev.controller_group if prev else None) - ) - effective_sidecars = _parse_csv_option(sidecars) if sidecars is not None else (prev.sidecars if prev else None) - effective_config = config if config is not None else (prev.config_path if prev else None) - effective_host = host if host is not None else (prev.host if prev else _DEFAULT_HOST) - effective_port = port if port is not None else (prev.port if prev else _DEFAULT_PORT) + effective_port = port if port is not None else (previous_config.port if previous_config else _DEFAULT_PORT) _warn_bind_all(effective_host) _ensure_port_available(effective_host, effective_port, scope, base_dir=base_dir) - - typer.echo("Starting platform services...") - proc = start_background( - scope=scope, + platform_config = PlatformAppConfig( services=effective_services, service_group=effective_service_group, controllers=effective_controllers, controller_group=effective_controller_group, sidecars=effective_sidecars, config_path=effective_config, + scope=scope, host=effective_host, port=effective_port, - base_dir=base_dir, + state_root=base_dir, ) + typer.echo("Starting platform services...") + proc = start_background(platform_config) + if not _wait_for_healthy(effective_host, effective_port): exit_code = proc.poll() if exit_code is not None: @@ -613,7 +616,7 @@ def status_services( ] = _DEFAULT_PORT, ) -> None: """Show status of the platform services instance for this scope.""" - scope = compute_scope(port=port, instance_name=instance) + scope = compute_scope(port=port, explicit_scope=instance) base_dir_str = _effective_base_dir() base_dir = Path(base_dir_str) if base_dir_str else None @@ -641,13 +644,13 @@ def status_services( except ValueError: uptime = "unknown" - healthy = _wait_for_healthy(desc.host, desc.port, timeout=3, poll_interval=0.5) + healthy = _wait_for_healthy(desc.config.host, desc.config.port, timeout=3, poll_interval=0.5) health_str = "healthy" if healthy else "unhealthy" - typer.echo(f"Scope: {desc.scope}") + typer.echo(f"Scope: {desc.config.scope}") typer.echo(f"PID: {desc.pid}") typer.echo(f"Mode: {desc.mode}") - typer.echo(f"Address: {desc.host}:{desc.port}") + typer.echo(f"Address: {desc.config.host}:{desc.config.port}") typer.echo(f"Uptime: {uptime}") typer.echo(f"Health: {health_str}") log = log_path_for(scope, base_dir=base_dir) @@ -668,7 +671,7 @@ def _print_instance_table(instances: list[InstanceInfo]) -> None: pid = addr = mode = "-" if info.descriptor: pid = str(info.descriptor.pid) - addr = f"{info.descriptor.host}:{info.descriptor.port}" + addr = f"{info.descriptor.config.host}:{info.descriptor.config.port}" mode = info.descriptor.mode typer.echo(f"{info.scope:<25} {status:<10} {pid:<10} {addr:<25} {mode:<12}") @@ -872,7 +875,7 @@ def logs_services( nemo services logs --path nemo services logs -n 100 """ - scope = compute_scope(port=port, instance_name=instance) + scope = compute_scope(port=port, explicit_scope=instance) base_dir_str = _effective_base_dir() base_dir = Path(base_dir_str) if base_dir_str else None diff --git a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/setup.py b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/setup.py index 1c52e750b1..626d56bd07 100644 --- a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/setup.py +++ b/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/setup.py @@ -30,20 +30,12 @@ from nemo_platform_plugin.secrets.client import SecretsClient from nemo_platform_plugin.secrets.types import PlatformSecretCreateRequest, PlatformSecretUpdateRequest from nmp.common.config import nmp_user_data_dir +from nmp.platform_runner.config import DEFAULT_LOCAL_SERVICES_BIND_HOST, PlatformAppConfig from pydantic import SecretStr from rich import box from rich.console import Console from rich.panel import Panel -from nemo_platform_ext.cli.commands.services._process import ( - DEFAULT_SERVICES_BIND_HOST, - check_port_available_for_start, - compute_scope, - format_port_conflict, - log_path_for, - start_background, - stop_instance, -) from nemo_platform_ext.cli.commands.skills import registry as skills_registry from nemo_platform_ext.cli.commands.skills.base import Scope, Skill from nemo_platform_ext.cli.commands.skills.registry import get_installer, load_skills @@ -51,6 +43,14 @@ from nemo_platform_ext.cli.core.errors import handle_errors from nemo_platform_ext.config.config import Config from nemo_platform_ext.config.models import ConfigFile, ConfigParams, LocalServicesConfig +from nemo_platform_ext.local.process import ( + check_port_available_for_start, + compute_scope, + format_port_conflict, + log_path_for, + start_background, + stop_instance, +) from nemo_platform_ext.ui.prompts import ( UserCancelled, is_interactive, @@ -635,8 +635,10 @@ def _start_services_background(base_url: str, data_dir: str | None = None) -> su exported it). """ port = _resolve_services_port(base_url) - scope = compute_scope(port=port) - return start_background(scope=scope, port=port, data_dir=data_dir) + return start_background( + PlatformAppConfig(scope=compute_scope(port=port), port=port), + data_dir=data_dir, + ) def _last_startup_service(log_path: Path | None) -> str: @@ -686,16 +688,13 @@ def _kill_existing_services(base_url: str) -> None: Delegates to the shared process lifecycle module. """ - port = _resolve_services_port(base_url) - scope = compute_scope(port=port) - stop_instance(scope, timeout=2.0, force=True) + stop_instance(compute_scope(port=_resolve_services_port(base_url)), timeout=2.0, force=True) def _ensure_port_available_for_start(base_url: str) -> None: """Fail fast when the services port cannot be bound.""" port = _resolve_services_port(base_url) - scope = compute_scope(port=port) - conflict = check_port_available_for_start(DEFAULT_SERVICES_BIND_HOST, port, scope) + conflict = check_port_available_for_start(DEFAULT_LOCAL_SERVICES_BIND_HOST, port, compute_scope(port=port)) if conflict is None: return lines = format_port_conflict(conflict) @@ -777,8 +776,7 @@ def _maybe_start_services( _ensure_port_available_for_start(base_url) proc = _start_services_background(base_url, data_dir=data_dir) - port = _resolve_services_port(base_url) - log = log_path_for(compute_scope(port=port)) + log = log_path_for(compute_scope(port=_resolve_services_port(base_url))) if not _wait_for_platform(base_url, timeout=timeout, log_path=log): exit_code = proc.poll() diff --git a/packages/nemo_platform_ext/src/nemo_platform_ext/local/_service_child.py b/packages/nemo_platform_ext/src/nemo_platform_ext/local/_service_child.py new file mode 100644 index 0000000000..9cecdeec28 --- /dev/null +++ b/packages/nemo_platform_ext/src/nemo_platform_ext/local/_service_child.py @@ -0,0 +1,30 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Child entrypoint for SDK-started local services daemons.""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +from nemo_platform_ext.local.services import ServiceRunConfig, run_services + + +def main(argv: list[str] | None = None) -> int: + args = list(sys.argv[1:] if argv is None else argv) + if len(args) != 1: + sys.stderr.write("usage: python -m nemo_platform_ext.local._service_child \n") + return 2 + request_path = Path(args[0]) + try: + payload = json.loads(request_path.read_text(encoding="utf-8")) + finally: + request_path.unlink(missing_ok=True) + run_services(ServiceRunConfig(**payload), _mode="daemon") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/services/_process.py b/packages/nemo_platform_ext/src/nemo_platform_ext/local/process.py similarity index 76% rename from packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/services/_process.py rename to packages/nemo_platform_ext/src/nemo_platform_ext/local/process.py index 8ce8cda699..7ee0cd4a9b 100644 --- a/packages/nemo_platform_ext/src/nemo_platform_ext/cli/commands/services/_process.py +++ b/packages/nemo_platform_ext/src/nemo_platform_ext/local/process.py @@ -3,8 +3,13 @@ """Local process lifecycle for ``nemo services``. -Uses per-instance scoped directories under ``$XDG_STATE_HOME/nmp/instances/`` -with flock-based liveness tracking. Each instance directory contains: +In this module, "instance" is a local services process/resource, and "scope" +is the stable key used for that instance's lock, descriptor, socket, and log +paths. The CLI exposes this key as ``--instance`` for compatibility, but +internal code should use "scope" when referring to the key. + +Uses per-scope directories under ``$XDG_STATE_HOME/nmp/instances/`` +with flock-based liveness tracking. Each scope directory contains: - ``services.lock`` -- exclusive flock held for the process lifetime - ``instance.json`` -- descriptor with PID, port, services, etc. @@ -23,7 +28,6 @@ import json import logging import os -import re import shutil import signal import socket @@ -34,10 +38,16 @@ from dataclasses import dataclass, field from datetime import datetime, timezone from pathlib import Path -from typing import Literal +from typing import Literal, Self import psutil -from pydantic import BaseModel, Field +from nmp.platform_runner.config import ( + DEFAULT_LOCAL_SERVICES_BIND_HOST, + PlatformAppConfig, + default_state_root, + validate_scope, +) +from pydantic import BaseModel, Field, model_validator logger = logging.getLogger(__name__) @@ -45,10 +55,10 @@ DESCRIPTOR_FILENAME = "instance.json" LOG_FILENAME = "services.log" -DEFAULT_SERVICES_BIND_HOST = "127.0.0.1" SUGGESTED_ALT_PORT = 9090 _SIGTERM_POLL_INTERVAL = 0.25 +_SIGKILL_WAIT_TIMEOUT = 5.0 _DEFAULT_STOP_TIMEOUT = 30.0 @@ -62,10 +72,7 @@ def _pause(seconds: float) -> None: def _base_state_dir() -> Path: - xdg = os.environ.get("XDG_STATE_HOME") - if xdg: - return Path(xdg) / "nmp" - return Path.home() / ".local" / "state" / "nmp" + return default_state_root() def _instances_dir(*, base_dir: Path | None = None) -> Path: @@ -73,7 +80,7 @@ def _instances_dir(*, base_dir: Path | None = None) -> Path: def _find_git_root() -> str: - """Walk up from cwd looking for a ``.git`` directory. Falls back to cwd.""" + """Walk up from cwd looking for a ``.git`` directory. Falls back to cwd.""" cur = Path.cwd().resolve() for parent in (cur, *cur.parents): if (parent / ".git").exists(): @@ -84,24 +91,19 @@ def _find_git_root() -> str: _scope_prefix_cache: str | None = None -_SCOPE_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,127}$") - - -def _validate_scope(scope: str) -> str: - """Ensure *scope* is safe to use as a directory name.""" - if not _SCOPE_RE.fullmatch(scope): - raise ValueError(f"Invalid instance scope: {scope!r}") - return scope +def compute_scope(*, port: int, explicit_scope: str | None = None) -> str: + """Compute the local services scope. + The default scope is ``sha1(git_toplevel_or_cwd)[:8]-``. Including + the port is intentional: it lets two local services instances from the same + checkout use different TCP ports without sharing a lock, descriptor, or log + directory. -def compute_scope(*, port: int, instance_name: str | None = None) -> str: - """Compute a scope identifier for this working directory + port. - - Default: ``sha1(git_toplevel_or_cwd)[:8]-``. - Override with an explicit *instance_name*. + Explicit scopes are validated and returned as-is, so they do not encode the + port. Callers that pass an explicit scope own its uniqueness. """ - if instance_name: - return _validate_scope(instance_name) + if explicit_scope: + return validate_scope(explicit_scope) global _scope_prefix_cache # noqa: PLW0603 if _scope_prefix_cache is None: root = _find_git_root() @@ -110,7 +112,8 @@ def compute_scope(*, port: int, instance_name: str | None = None) -> str: def instance_dir(scope: str, *, base_dir: Path | None = None) -> Path: - d = _instances_dir(base_dir=base_dir) / _validate_scope(scope) + """Return the state directory for *scope*, creating it if needed.""" + d = _instances_dir(base_dir=base_dir) / validate_scope(scope) d.mkdir(parents=True, exist_ok=True) return d @@ -222,7 +225,7 @@ def _instance_owns_listener( desc = read_descriptor(scope, base_dir=base_dir) if desc is None: return False - return desc.port == port and _normalize_bind_host(desc.host) == _normalize_bind_host(host) + return desc.config.port == port and _normalize_bind_host(desc.config.host) == _normalize_bind_host(host) def is_port_bindable(host: str, port: int) -> bool: @@ -272,8 +275,9 @@ def format_port_conflict(err: PortConflict) -> list[str]: Message text depends on ``err.kind`` (foreign process vs NeMo instance). """ if err.kind == "nemo_instance": + owner = f" '{err.scope}'" if err.scope else "" return [ - f"Port {err.port} is in use by a NeMo Platform instance for this directory.", + f"Port {err.port} is in use by NeMo Platform instance{owner}.", "Stop it first with: nemo services stop", "Or restart with: nemo services restart", ] @@ -292,23 +296,39 @@ def format_port_conflict(err: PortConflict) -> list[str]: class InstanceDescriptor(BaseModel): pid: int - scope: str - host: str = "127.0.0.1" - port: int = 8080 - mode: Literal["foreground", "background"] = "background" + config: PlatformAppConfig = Field(default_factory=PlatformAppConfig) + transport: Literal["tcp", "uds"] = "tcp" + mode: Literal["foreground", "background", "daemon"] = "background" create_time: float = 0.0 started_at: str = Field(default_factory=lambda: datetime.now(timezone.utc).isoformat()) - services: list[str] | None = None - controllers: list[str] | None = None - service_group: str | None = None - controller_group: str | None = None - sidecars: list[str] | None = None - config_path: str | None = None - log_path: str | None = None + + @model_validator(mode="after") + def _validate_client_transport(self) -> Self: + if self.transport == "uds" and self.config.socket_path is None: + raise ValueError("UDS client transport requires config.socket_path") + return self + + @classmethod + def from_config( + cls, + config: PlatformAppConfig, + *, + mode: Literal["foreground", "background", "daemon"], + transport: Literal["uds", "tcp"] = "tcp", + pid: int | None = None, + ) -> Self: + resolved_pid = os.getpid() if pid is None else pid + return cls( + pid=resolved_pid, + config=config, + transport=transport, + mode=mode, + create_time=get_create_time(resolved_pid), + ) def write_descriptor(desc: InstanceDescriptor, *, base_dir: Path | None = None) -> Path: - d = instance_dir(desc.scope, base_dir=base_dir) + d = instance_dir(desc.config.scope, base_dir=base_dir) path = d / DESCRIPTOR_FILENAME payload = desc.model_dump() fd, tmp = tempfile.mkstemp(dir=str(d), suffix=".tmp") @@ -335,10 +355,19 @@ def read_descriptor(scope: str, *, base_dir: Path | None = None) -> InstanceDesc return None try: data = json.loads(path.read_text()) - return InstanceDescriptor.model_validate(data) + desc = InstanceDescriptor.model_validate(data) except (json.JSONDecodeError, KeyError, TypeError, ValueError): logger.debug("Corrupt descriptor at %s, ignoring", path, exc_info=True) return None + if desc.config.scope != scope: + logger.debug( + "Descriptor at %s has scope=%r but lives under %r, ignoring", + path, + desc.config.scope, + scope, + ) + return None + return desc def remove_descriptor(scope: str, *, base_dir: Path | None = None) -> None: @@ -351,24 +380,24 @@ def remove_descriptor(scope: str, *, base_dir: Path | None = None) -> None: def _scope_dir(scope: str, *, base_dir: Path | None = None) -> Path: - return _instances_dir(base_dir=base_dir) / _validate_scope(scope) + return _instances_dir(base_dir=base_dir) / validate_scope(scope) def _is_log_file(path: Path) -> bool: return path.name == LOG_FILENAME or path.name.startswith(f"{LOG_FILENAME}.") -def _iter_log_files(scope_dir: Path): - if not scope_dir.is_dir(): +def _iter_log_files(scope_dir_path: Path): + if not scope_dir_path.is_dir(): return - for path in scope_dir.iterdir(): + for path in scope_dir_path.iterdir(): if path.is_file() and _is_log_file(path): yield path -def _has_preservable_logs(scope_dir: Path) -> bool: - """Return True if *scope_dir* contains non-empty service log files.""" - return any(path.stat().st_size > 0 for path in _iter_log_files(scope_dir)) +def _has_preservable_logs(scope_dir_path: Path) -> bool: + """Return True if *scope_dir_path* contains non-empty service log files.""" + return any(path.stat().st_size > 0 for path in _iter_log_files(scope_dir_path)) def is_removable_ghost( @@ -377,17 +406,17 @@ def is_removable_ghost( base_dir: Path | None = None, descriptor: InstanceDescriptor | None = None, ) -> bool: - """True when a dead scope dir has no descriptor and no non-empty logs.""" + """True when a dead scope directory has no descriptor and no non-empty logs.""" if is_instance_alive(scope, base_dir=base_dir): return False if descriptor is not None: return False - scope_dir = _scope_dir(scope, base_dir=base_dir) - if not scope_dir.is_dir(): + scope_dir_path = _scope_dir(scope, base_dir=base_dir) + if not scope_dir_path.is_dir(): return False - if (scope_dir / DESCRIPTOR_FILENAME).exists(): + if (scope_dir_path / DESCRIPTOR_FILENAME).exists(): return False - return not _has_preservable_logs(scope_dir) + return not _has_preservable_logs(scope_dir_path) # --------------------------------------------------------------------------- @@ -425,7 +454,7 @@ class InstanceInfo: def list_instances(*, base_dir: Path | None = None) -> list[InstanceInfo]: - """Scan all instance directories and return their status. + """Scan all scope directories and return their status. Side effects: - Removes stale descriptors for dead instances. @@ -448,7 +477,7 @@ def list_instances(*, base_dir: Path | None = None) -> list[InstanceInfo]: try: shutil.rmtree(child) except OSError: - logger.debug("Could not remove ghost instance dir %s", child, exc_info=True) + logger.debug("Could not remove ghost scope directory %s", child, exc_info=True) else: continue results.append(InstanceInfo(scope=scope, alive=alive, descriptor=desc)) @@ -456,28 +485,28 @@ def list_instances(*, base_dir: Path | None = None) -> list[InstanceInfo]: def remove_instance(scope: str, *, base_dir: Path | None = None) -> bool: - """Remove an instance scope directory. + """Remove a scope directory. - Returns False if the scope did not exist or could not be removed. + Returns False if the scope directory did not exist or could not be removed. """ - scope = _validate_scope(scope) + scope = validate_scope(scope) if is_instance_alive(scope, base_dir=base_dir): raise InstanceStillRunningError(scope) - scope_dir = _scope_dir(scope, base_dir=base_dir) - if not scope_dir.is_dir(): + scope_dir_path = _scope_dir(scope, base_dir=base_dir) + if not scope_dir_path.is_dir(): return False with contextlib.suppress(OSError): - shutil.rmtree(scope_dir) - return not scope_dir.is_dir() + shutil.rmtree(scope_dir_path) + return not scope_dir_path.is_dir() def list_stopped_scopes(*, base_dir: Path | None = None) -> list[str]: - """Return scope names for instances that are not alive.""" + """Return scopes for instances that are not alive.""" return [info.scope for info in list_instances(base_dir=base_dir) if not info.alive] def prune_instances(*, base_dir: Path | None = None) -> list[str]: - """Remove all stopped instance directories. Returns removed scope names.""" + """Remove all stopped scope directories. Returns removed scopes.""" removed: list[str] = [] for scope in list_stopped_scopes(base_dir=base_dir): if remove_instance(scope, base_dir=base_dir): @@ -497,11 +526,15 @@ def instance_log_bytes(scope: str, *, base_dir: Path | None = None) -> int: def rotate_log(scope: str, *, base_dir: Path | None = None) -> Path: """Rotate the existing log and return the path for the new one.""" - d = instance_dir(scope, base_dir=base_dir) - log_path = d / LOG_FILENAME + return rotate_log_path(log_path_for(scope, base_dir=base_dir)) + + +def rotate_log_path(log_path: Path) -> Path: + """Rotate the existing log at *log_path* and return the path for the new one.""" + log_path.parent.mkdir(parents=True, exist_ok=True) if log_path.exists() and log_path.stat().st_size > 0: ts = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S%fZ") - rotated = d / f"{LOG_FILENAME}.{ts}" + rotated = log_path.with_name(f"{log_path.name}.{ts}") log_path.rename(rotated) return log_path @@ -633,6 +666,10 @@ def stop_instance( return StopResult(stopped_pids=[], swept_children=swept) except OSError: logger.debug("Failed to send SIGKILL to pid %d", pid, exc_info=True) + if not _wait_for_pid_exit(pid, timeout=_SIGKILL_WAIT_TIMEOUT): + logger.warning("PID %d is still alive after SIGKILL; preserving descriptor", pid) + swept = _sweep_orphans(children) if children else [] + return StopResult(stopped_pids=[], swept_children=swept) swept = _sweep_orphans(children) if children else [] @@ -640,13 +677,21 @@ def stop_instance( return StopResult(stopped_pids=[pid], swept_children=swept) +def _wait_for_pid_exit(pid: int, *, timeout: float) -> bool: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if not _pid_alive(pid): + return True + _pause(_SIGTERM_POLL_INTERVAL) + return not _pid_alive(pid) + + def _pid_alive(pid: int) -> bool: try: - os.kill(pid, 0) - return True - except ProcessLookupError: + return psutil.Process(pid).status() != psutil.STATUS_ZOMBIE + except psutil.NoSuchProcess: return False - except PermissionError: + except psutil.AccessDenied: return True except OSError: return False @@ -658,17 +703,8 @@ def _pid_alive(pid: int) -> bool: def start_background( + config: PlatformAppConfig | None = None, *, - scope: str, - services: list[str] | None = None, - service_group: str | None = None, - controllers: list[str] | None = None, - controller_group: str | None = None, - sidecars: list[str] | None = None, - config_path: str | None = None, - host: str = DEFAULT_SERVICES_BIND_HOST, - port: int = 8080, - base_dir: Path | None = None, data_dir: str | None = None, ) -> subprocess.Popen: """Launch ``nemo services run`` as a detached background subprocess. @@ -676,31 +712,32 @@ def start_background( The child acquires the flock and writes its own descriptor. The parent returns the ``Popen`` handle for health polling. """ - log_file_path = rotate_log(scope, base_dir=base_dir) + config = config or PlatformAppConfig(host=DEFAULT_LOCAL_SERVICES_BIND_HOST) + log_file_path = rotate_log_path(config.log_file_path()) log_file = open(log_file_path, "a") # noqa: SIM115 nemo_bin = str(Path(sys.executable).parent / "nemo") args: list[str] = [nemo_bin, "services", "run"] - if services: - args += ["--services", ",".join(services)] - if service_group: - args += ["--service-group", service_group] - if controllers: - args += ["--controllers", ",".join(controllers)] - if controller_group: - args += ["--controller-group", controller_group] - if sidecars: - args += ["--sidecars", ",".join(sidecars)] - if config_path: - args += ["--config", config_path] - args += ["--host", host, "--port", str(port)] - args += ["--instance", scope] + if config.services: + args += ["--services", ",".join(config.services)] + if config.service_group: + args += ["--service-group", config.service_group] + if config.controllers: + args += ["--controllers", ",".join(config.controllers)] + if config.controller_group: + args += ["--controller-group", config.controller_group] + if config.sidecars: + args += ["--sidecars", ",".join(config.sidecars)] + if config.config_path: + args += ["--config", config.config_path] + args += ["--host", config.host, "--port", str(config.port)] + args += ["--instance", config.scope] env = os.environ.copy() if data_dir and "NMP_DATA_DIR" not in env: env["NMP_DATA_DIR"] = data_dir - if base_dir: - env["_NMP_STATE_DIR"] = str(base_dir) + if config.state_root is not None: + env["_NMP_STATE_DIR"] = str(config.state_root) # Tell the child ``run`` process it was launched by ``start`` so it # records mode="background" in its descriptor. This is internal # parent-to-child signaling -- not a public API surface -- following the diff --git a/packages/nemo_platform_ext/src/nemo_platform_ext/local/services.py b/packages/nemo_platform_ext/src/nemo_platform_ext/local/services.py new file mode 100644 index 0000000000..d115fe06ec --- /dev/null +++ b/packages/nemo_platform_ext/src/nemo_platform_ext/local/services.py @@ -0,0 +1,728 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Programmatic local lifecycle API for NeMo Platform services.""" + +from __future__ import annotations + +import asyncio +import importlib.util +import json +import os +import subprocess +import sys +import tempfile +import time +from collections.abc import MutableMapping, Sequence +from dataclasses import dataclass +from enum import StrEnum +from pathlib import Path +from typing import Any, Literal, Protocol, Self, runtime_checkable + +from nemo_platform import AsyncNeMoPlatform, NeMoPlatform +from nemo_platform_ext.local import process +from nemo_platform_ext.local.transport import ( + EMBEDDED_BASE_URL, + UDS_BASE_URL, + build_async_asgi_http_client, + build_async_http_client, + build_sync_asgi_http_client, + build_sync_http_client, + probe_status, + tcp_base_url, + wait_for_status, + wait_for_status_async, +) +from nmp.platform_runner.config import ( + DEFAULT_SCOPE, + PlatformAppConfig, + default_runtime_root, + default_state_root, + validate_scope, +) + +_AF_UNIX_PATH_MAX_BYTES = 103 if sys.platform.startswith(("darwin", "freebsd", "openbsd", "netbsd")) else 107 + + +class ServicesError(RuntimeError): + """Base class for local services lifecycle errors.""" + + +class ServicesExtraRequiredError(ServicesError): + """Raised when local service dependencies are not installed.""" + + +class ServicesAlreadyRunningError(ServicesError): + """Raised when a requested local instance is already running.""" + + +class ServicesNotRunningError(ServicesError): + """Raised when a requested local instance is not running.""" + + +class ServicesPortInUseError(ServicesError): + """Raised when TCP startup targets an unavailable port.""" + + +class ServicesStartupTimeoutError(ServicesError): + """Raised when startup does not become healthy before the timeout.""" + + +class ServicesStartupExitedError(ServicesError): + """Raised when a daemon child exits before becoming healthy.""" + + +class ServicesSocketStaleError(ServicesError): + """Raised when a stale socket cannot be removed.""" + + +def _as_tuple(value: Sequence[str] | None) -> tuple[str, ...] | None: + if value is None: + return None + return tuple(value) + + +def _optional_str(value: str | Path | None) -> str | None: + if value is None: + return None + return str(value) + + +def _optional_list(value: Sequence[str] | None) -> list[str] | None: + if value is None: + return None + return list(value) + + +class ServiceMode(StrEnum): + EMBEDDED = "embedded" + DAEMON = "daemon" + + +@dataclass(frozen=True) +class StartServicesResult: + requested: list[str] + started: list[str] + already_active: list[str] + active: list[str] + + +@runtime_checkable +class LocalServiceHandle(Protocol): + """Shared lifecycle/client contract for local services handles.""" + + def is_running(self) -> bool: ... + + def wait_until_ready(self, timeout: float | None = None) -> None: ... + + async def wait_until_ready_async(self, timeout: float | None = None) -> None: ... + + def client(self, **kwargs: Any) -> NeMoPlatform: ... + + def async_client(self, **kwargs: Any) -> AsyncNeMoPlatform: ... + + def start_services(self, service_names: Sequence[str]) -> StartServicesResult: ... + + async def start_services_async(self, service_names: Sequence[str]) -> StartServicesResult: ... + + def stop(self, *, timeout: float = 30.0, force: bool = False) -> process.StopResult: ... + + async def stop_async(self, *, timeout: float = 30.0, force: bool = False) -> process.StopResult: ... + + +@dataclass +class ServiceRunConfig: + services: Sequence[str] | None = None + service_group: str | None = None + controllers: Sequence[str] | None = None + controller_group: str | None = None + sidecars: Sequence[str] | None = None + config_path: str | Path | None = None + transport: Literal["uds", "tcp"] = "uds" + socket_path: str | Path | None = None + http_gateway: Literal["enabled", "disabled"] = "disabled" + http_gateway_host: str = "127.0.0.1" + http_gateway_port: int | None = None + host: str = "127.0.0.1" + port: int = 8080 + scope: str = DEFAULT_SCOPE + state_dir: str | Path | None = None + runtime_dir: str | Path | None = None + data_dir: str | Path | None = None + readiness_timeout: float = 60.0 + readiness_poll_interval: float = 0.5 + mode: ServiceMode | str = ServiceMode.DAEMON + + def __post_init__(self) -> None: + self.services = _as_tuple(self.services) + self.controllers = _as_tuple(self.controllers) + self.sidecars = _as_tuple(self.sidecars) + try: + self.mode = ServiceMode(self.mode) + except ValueError as error: + raise ValueError("mode must be 'embedded' or 'daemon'") from error + + if self.services and self.service_group: + raise ValueError("services cannot be combined with service_group") + if self.controllers and self.controller_group: + raise ValueError("controllers cannot be combined with controller_group") + if self.transport not in {"uds", "tcp"}: + raise ValueError("transport must be 'uds' or 'tcp'") + if self.http_gateway not in {"enabled", "disabled"}: + raise ValueError("http_gateway must be 'enabled' or 'disabled'") + if self.http_gateway == "enabled" and self.transport != "uds": + raise ValueError("gateway can only be enabled for UDS transport") + if self.readiness_timeout <= 0: + raise ValueError("readiness_timeout must be greater than 0") + if self.readiness_poll_interval <= 0: + raise ValueError("readiness_poll_interval must be greater than 0") + self.scope = validate_scope(self.scope) + + @property + def state_root(self) -> Path: + return Path(self.state_dir).expanduser() if self.state_dir is not None else default_state_root() + + @property + def runtime_root(self) -> Path: + return Path(self.runtime_dir).expanduser() if self.runtime_dir is not None else default_runtime_root() + + @property + def resolved_socket_path(self) -> Path | None: + if self.socket_path is not None: + socket_path = Path(self.socket_path).expanduser() + elif self.transport == "uds": + socket_path = PlatformAppConfig( + scope=self.scope, + runtime_root=self.runtime_root, + ).socket_file_path() + else: + return None + if not socket_path.is_absolute(): + raise ValueError(f"UDS socket path must be absolute: {socket_path}") + return socket_path + + def to_platform_app_config(self) -> PlatformAppConfig: + return PlatformAppConfig( + services=self.services, + service_group=self.service_group, + controllers=self.controllers, + controller_group=self.controller_group, + sidecars=self.sidecars, + config_path=_optional_str(self.config_path), + scope=self.scope, + host=self.host, + port=self.port, + socket_path=_optional_str(self.resolved_socket_path), + state_root=_optional_str(self.state_root), + runtime_root=_optional_str(self.runtime_dir), + ) + + def to_child_payload(self) -> dict[str, object]: + return { + "mode": ServiceMode(self.mode).value, + "services": _optional_list(self.services), + "service_group": self.service_group, + "controllers": _optional_list(self.controllers), + "controller_group": self.controller_group, + "sidecars": _optional_list(self.sidecars), + "config_path": _optional_str(self.config_path), + "transport": self.transport, + "socket_path": _optional_str(self.socket_path), + "http_gateway": self.http_gateway, + "http_gateway_host": self.http_gateway_host, + "http_gateway_port": self.http_gateway_port, + "host": self.host, + "port": self.port, + "scope": self.scope, + "state_dir": _optional_str(self.state_dir), + "runtime_dir": _optional_str(self.runtime_dir), + "data_dir": _optional_str(self.data_dir), + "readiness_timeout": self.readiness_timeout, + "readiness_poll_interval": self.readiness_poll_interval, + } + + +@dataclass(frozen=True) +class DaemonServiceHandle: + scope: str + transport: Literal["uds", "tcp"] + socket_path: Path | None + gateway_base_url: str | None + host: str + port: int + pid: int | None + mode: Literal["foreground", "daemon"] + log_path: Path | None + state_dir: Path | None + runtime_dir: Path | None + + @classmethod + def from_descriptor(cls, desc: process.InstanceDescriptor) -> Self: + socket_path = Path(desc.config.socket_path) if desc.config.socket_path else None + runtime_dir = desc.config.runtime_dir() if socket_path else None + return cls( + scope=desc.config.scope, + transport=desc.transport, + socket_path=socket_path, + gateway_base_url=None, + host=desc.config.host, + port=desc.config.port, + pid=desc.pid, + mode="daemon" if desc.mode == "daemon" else "foreground", + log_path=desc.config.log_file_path(), + state_dir=desc.config.state_dir(), + runtime_dir=runtime_dir, + ) + + @classmethod + def from_config( + cls, + config: ServiceRunConfig, + *, + pid: int | None = None, + ) -> Self: + app_config = config.to_platform_app_config() + socket_path = config.resolved_socket_path + runtime_dir = app_config.runtime_dir() if socket_path else None + return cls( + scope=config.scope, + transport=config.transport, + socket_path=socket_path, + gateway_base_url=None, + host=config.host, + port=config.port, + pid=pid, + mode="daemon", + log_path=app_config.log_file_path(), + state_dir=app_config.state_dir(), + runtime_dir=runtime_dir, + ) + + @property + def base_url(self) -> str: + if self.transport == "uds": + return UDS_BASE_URL + return tcp_base_url(self.host, self.port) + + def _state_root(self) -> Path | None: + if self.state_dir is None: + return None + if self.state_dir.parent.name == "instances": + return self.state_dir.parent.parent + return self.state_dir.parent + + def is_running(self) -> bool: + state_root = self._state_root() + return process.is_instance_alive(self.scope, base_dir=state_root) + + def wait_until_ready(self, timeout: float | None = None) -> None: + if not wait_for_status( + base_url=self.base_url, + socket_path=self.socket_path if self.transport == "uds" else None, + timeout=60.0 if timeout is None else timeout, + ): + raise ServicesStartupTimeoutError(f"Timed out waiting for services instance {self.scope!r}") + + async def wait_until_ready_async(self, timeout: float | None = None) -> None: + if not await wait_for_status_async( + base_url=self.base_url, + socket_path=self.socket_path if self.transport == "uds" else None, + timeout=60.0 if timeout is None else timeout, + ): + raise ServicesStartupTimeoutError(f"Timed out waiting for services instance {self.scope!r}") + + def stop(self, *, timeout: float = 30.0, force: bool = False) -> process.StopResult: + state_root = self._state_root() + return process.stop_instance(self.scope, base_dir=state_root, timeout=timeout, force=force) + + async def stop_async(self, *, timeout: float = 30.0, force: bool = False) -> process.StopResult: + return await asyncio.to_thread(self.stop, timeout=timeout, force=force) + + def start_services(self, service_names: Sequence[str]) -> StartServicesResult: + raise ServicesError("Staged service start is not implemented for daemon mode yet") + + async def start_services_async(self, service_names: Sequence[str]) -> StartServicesResult: + return await asyncio.to_thread(self.start_services, service_names) + + def client(self, **kwargs: Any) -> NeMoPlatform: + if self.transport == "uds": + if self.socket_path is None: + raise ServicesError("UDS service handle is missing socket_path") + kwargs.setdefault("http_client", build_sync_http_client(self.socket_path)) + kwargs.setdefault("base_url", self.base_url) + return NeMoPlatform(**kwargs) + + def async_client(self, **kwargs: Any) -> AsyncNeMoPlatform: + if self.transport == "uds": + if self.socket_path is None: + raise ServicesError("UDS service handle is missing socket_path") + kwargs.setdefault("http_client", build_async_http_client(self.socket_path)) + kwargs.setdefault("base_url", self.base_url) + return AsyncNeMoPlatform(**kwargs) + + +@dataclass(frozen=True) +class EmbeddedServiceHandle: + app: Any + runtime: object + + def is_running(self) -> bool: + return True + + def wait_until_ready(self, timeout: float | None = None) -> None: + return None + + async def wait_until_ready_async(self, timeout: float | None = None) -> None: + return None + + def client(self, **kwargs: Any) -> NeMoPlatform: + kwargs.setdefault("http_client", build_sync_asgi_http_client(self.app)) + kwargs.setdefault("base_url", EMBEDDED_BASE_URL) + return NeMoPlatform(**kwargs) + + def async_client(self, **kwargs: Any) -> AsyncNeMoPlatform: + kwargs.setdefault("http_client", build_async_asgi_http_client(self.app)) + kwargs.setdefault("base_url", EMBEDDED_BASE_URL) + return AsyncNeMoPlatform(**kwargs) + + def start_services(self, service_names: Sequence[str]) -> StartServicesResult: + raise ServicesError("Staged service start is not implemented for embedded mode yet") + + async def start_services_async(self, service_names: Sequence[str]) -> StartServicesResult: + return await asyncio.to_thread(self.start_services, service_names) + + def stop(self, *, timeout: float = 30.0, force: bool = False) -> process.StopResult: + return process.StopResult(stopped_pids=[], swept_children=[]) + + async def stop_async(self, *, timeout: float = 30.0, force: bool = False) -> process.StopResult: + return self.stop(timeout=timeout, force=force) + + +def require_services_extra() -> None: + if importlib.util.find_spec("pyleak") is not None: + return + raise ServicesExtraRequiredError("Install service dependencies with `pip install 'nemo-platform[all]'`.") + + +def _validate_socket_path_length(socket_path: Path) -> None: + encoded_length = len(os.fsencode(socket_path)) + if encoded_length > _AF_UNIX_PATH_MAX_BYTES: + raise ValueError( + "UDS socket path is too long for AF_UNIX " + f"({encoded_length} bytes; maximum is {_AF_UNIX_PATH_MAX_BYTES} bytes): {socket_path}" + ) + + +def _validated_socket_path(config: ServiceRunConfig) -> Path | None: + socket_path = config.resolved_socket_path + if socket_path is None: + return None + _validate_socket_path_length(socket_path) + return socket_path + + +def _prepare_socket(config: ServiceRunConfig) -> Path | None: + socket_path = _validated_socket_path(config) + if socket_path is None: + return None + socket_path.parent.mkdir(parents=True, exist_ok=True) + if not socket_path.exists(): + return socket_path + if probe_status(base_url=UDS_BASE_URL, socket_path=socket_path, timeout=0.5): + raise ServicesAlreadyRunningError(f"UDS socket is live at {socket_path}") + try: + socket_path.unlink() + except OSError as error: + raise ServicesSocketStaleError(f"Could not remove stale socket at {socket_path}") from error + return socket_path + + +def _check_tcp_available(config: ServiceRunConfig) -> None: + conflict = process.check_port_available_for_start( + config.host, + config.port, + config.scope, + base_dir=config.state_root, + ) + if conflict is not None: + raise ServicesPortInUseError("\n".join(process.format_port_conflict(conflict))) + + +def _write_run_request(config: ServiceRunConfig) -> Path: + state_dir = config.to_platform_app_config().state_dir(create=True) + fd, tmp = tempfile.mkstemp(dir=state_dir, suffix=".json") + path = Path(tmp) + try: + with os.fdopen(fd, "w", encoding="utf-8") as file: + fd = -1 + json.dump(config.to_child_payload(), file, indent=2) + file.write("\n") + except BaseException: + if fd >= 0: + os.close(fd) + fd = -1 + path.unlink(missing_ok=True) + raise + finally: + if fd >= 0: + os.close(fd) + return path + + +def _terminate_startup_process(proc: subprocess.Popen) -> None: + if proc.poll() is not None: + return + proc.terminate() + try: + proc.wait(timeout=5) + except subprocess.TimeoutExpired: + proc.kill() + proc.wait(timeout=5) + + +def serve_embedded_app(app: Any, cfg: ServiceRunConfig, socket_path: Path | None) -> None: + import uvicorn + + if socket_path is not None: + from nmp.platform_runner.server import _run_server_on_bound_sockets + + _run_server_on_bound_sockets(app, host=cfg.host, port=cfg.port, socket_path=str(socket_path)) + else: + uvicorn.run(app, host=cfg.host, port=cfg.port, log_config=None) + + +def run_services( + config: ServiceRunConfig | None = None, + *, + _mode: Literal["foreground", "daemon"] = "foreground", + env: MutableMapping[str, str] | None = None, +) -> None: + cfg = config or ServiceRunConfig() + app_config = cfg.to_platform_app_config() + require_services_extra() + if cfg.http_gateway == "enabled": + raise ServicesError("HTTP gateway support is not implemented yet") + if process.is_instance_alive(cfg.scope, base_dir=cfg.state_root): + raise ServicesAlreadyRunningError(f"Instance {cfg.scope!r} is already running") + _check_tcp_available(cfg) + lock_fd = process.acquire_lock(cfg.scope, base_dir=cfg.state_root) + original_data_dir = os.environ.get("NMP_DATA_DIR") + try: + socket_path = _prepare_socket(cfg) + app_config.log_file_path(create_parent=True) + if cfg.data_dir is not None and "NMP_DATA_DIR" not in os.environ: + os.environ["NMP_DATA_DIR"] = str(cfg.data_dir) + desc = process.InstanceDescriptor.from_config( + app_config, + mode=_mode, + transport=cfg.transport, + ) + process.write_descriptor(desc, base_dir=cfg.state_root) + embedded_handle = start_embedded_services(cfg, env=env) + serve_embedded_app(embedded_handle.app, cfg, socket_path) + finally: + try: + process.remove_descriptor(cfg.scope, base_dir=cfg.state_root) + finally: + if original_data_dir is None: + os.environ.pop("NMP_DATA_DIR", None) + else: + os.environ["NMP_DATA_DIR"] = original_data_dir + os.close(lock_fd) + + +def daemonize_services(config: ServiceRunConfig | None = None) -> DaemonServiceHandle: + cfg = config or ServiceRunConfig() + app_config = cfg.to_platform_app_config() + require_services_extra() + if cfg.http_gateway == "enabled": + raise ServicesError("HTTP gateway support is not implemented yet") + if process.is_instance_alive(cfg.scope, base_dir=cfg.state_root): + raise ServicesAlreadyRunningError(f"Instance {cfg.scope!r} is already running") + _check_tcp_available(cfg) + socket_path = _validated_socket_path(cfg) + if ( + socket_path is not None + and socket_path.exists() + and probe_status(base_url=UDS_BASE_URL, socket_path=socket_path, timeout=0.5) + ): + raise ServicesAlreadyRunningError(f"UDS socket is live at {socket_path}") + + request_path = _write_run_request(cfg) + log_path = process.rotate_log_path(app_config.log_file_path()) + log_file = open(log_path, "a") # noqa: SIM115 + env = os.environ.copy() + if cfg.data_dir is not None and "NMP_DATA_DIR" not in env: + env["NMP_DATA_DIR"] = str(cfg.data_dir) + proc: subprocess.Popen | None = None + ownership_transferred = False + try: + try: + child_module = f"{__package__}._service_child" + proc = subprocess.Popen( + [sys.executable, "-m", child_module, str(request_path)], + stdout=log_file, + stderr=subprocess.STDOUT, + env=env, + start_new_session=True, + stdin=subprocess.DEVNULL, + close_fds=True, + ) + finally: + log_file.close() + assert proc is not None + handle = DaemonServiceHandle.from_config(cfg, pid=proc.pid) + deadline = time.monotonic() + cfg.readiness_timeout + while True: + remaining = deadline - time.monotonic() + if remaining <= 0: + break + if proc.poll() is not None: + raise ServicesStartupExitedError(f"Services daemon exited with code {proc.returncode}; log: {log_path}") + if probe_status( + base_url=handle.base_url, + socket_path=handle.socket_path if handle.transport == "uds" else None, + timeout=remaining, + ): + ownership_transferred = True + return handle + remaining = deadline - time.monotonic() + if remaining <= 0: + break + time.sleep(min(cfg.readiness_poll_interval, remaining)) + raise ServicesStartupTimeoutError(f"Timed out waiting for services daemon {cfg.scope!r}; log: {log_path}") + finally: + if proc is not None and not ownership_transferred: + _terminate_startup_process(proc) + + +async def daemonize_services_async(config: ServiceRunConfig | None = None) -> DaemonServiceHandle: + return await asyncio.to_thread(daemonize_services, config) + + +def start_embedded_services( + config: ServiceRunConfig | None = None, + *, + env: MutableMapping[str, str] | None = None, +) -> EmbeddedServiceHandle: + """Start platform services in the current process. + + Args: + env: Environment mapping passed to :func:`build_platform_app`. + Defaults to ``None`` which writes to ``os.environ``. Tests can + pass an empty dict to avoid polluting the process environment. + """ + cfg = config or ServiceRunConfig(mode=ServiceMode.EMBEDDED) + from nmp.platform_runner.server import build_platform_app + + app = build_platform_app( + config=cfg.to_platform_app_config(), + env=env, + ) + runtime = getattr(app.state, "platform_runtime", None) + return EmbeddedServiceHandle(app=app, runtime=runtime) + + +async def start_embedded_services_async(config: ServiceRunConfig | None = None) -> EmbeddedServiceHandle: + return start_embedded_services(config) + + +def get_service_handle(config: ServiceRunConfig | None = None) -> DaemonServiceHandle | None: + cfg = config or ServiceRunConfig() + desc = process.read_descriptor(cfg.scope, base_dir=cfg.state_root) + if desc is None or not process.is_instance_alive(cfg.scope, base_dir=cfg.state_root): + return None + return DaemonServiceHandle.from_descriptor(desc) + + +def list_service_handles(state_dir: str | Path | None = None) -> list[DaemonServiceHandle]: + state_root = Path(state_dir).expanduser() if state_dir is not None else default_state_root() + handles: list[DaemonServiceHandle] = [] + for info in process.list_instances(base_dir=state_root): + if info.descriptor is not None and info.alive: + handles.append(DaemonServiceHandle.from_descriptor(info.descriptor)) + return handles + + +def ensure_services( + config: ServiceRunConfig | None = None, + *, + daemonize: bool | None = None, +) -> LocalServiceHandle: + cfg = config or ServiceRunConfig() + if cfg.mode is ServiceMode.EMBEDDED: + return start_embedded_services(cfg) + + handle = get_service_handle(cfg) + if handle is not None: + return handle + if daemonize is False: + raise ServicesNotRunningError(f"Instance {cfg.scope!r} is not running") + return daemonize_services(cfg) + + +async def ensure_services_async( + config: ServiceRunConfig | None = None, + *, + daemonize: bool | None = None, +) -> LocalServiceHandle: + cfg = config or ServiceRunConfig() + if cfg.mode is ServiceMode.EMBEDDED: + return await start_embedded_services_async(cfg) + + handle = get_service_handle(cfg) + if handle is not None: + return handle + if daemonize is False: + raise ServicesNotRunningError(f"Instance {cfg.scope!r} is not running") + return await daemonize_services_async(cfg) + + +def connect_services( + config: ServiceRunConfig | None = None, + *, + daemonize: bool | None = None, + start_if_needed: bool = True, + **client_kwargs: Any, +) -> NeMoPlatform: + cfg = config or ServiceRunConfig() + if not start_if_needed and cfg.mode is ServiceMode.DAEMON and get_service_handle(cfg) is None: + raise ServicesNotRunningError(f"Instance {cfg.scope!r} is not running") + handle = ensure_services(cfg, daemonize=daemonize) + return handle.client(**client_kwargs) + + +async def connect_services_async( + config: ServiceRunConfig | None = None, + *, + daemonize: bool | None = None, + start_if_needed: bool = True, + **client_kwargs: Any, +) -> AsyncNeMoPlatform: + cfg = config or ServiceRunConfig() + if not start_if_needed and cfg.mode is ServiceMode.DAEMON and get_service_handle(cfg) is None: + raise ServicesNotRunningError(f"Instance {cfg.scope!r} is not running") + handle = await ensure_services_async(cfg, daemonize=daemonize) + return handle.async_client(**client_kwargs) + + +def stop_services( + config: ServiceRunConfig | None = None, + *, + timeout: float = 30.0, + force: bool = False, +) -> process.StopResult: + cfg = config or ServiceRunConfig() + handle = get_service_handle(cfg) + if handle is None: + raise ServicesNotRunningError(f"Instance {cfg.scope!r} is not running") + return handle.stop(timeout=timeout, force=force) + + +async def stop_services_async( + config: ServiceRunConfig | None = None, + *, + timeout: float = 30.0, + force: bool = False, +) -> process.StopResult: + return await asyncio.to_thread(stop_services, config, timeout=timeout, force=force) diff --git a/packages/nemo_platform_ext/src/nemo_platform_ext/local/transport.py b/packages/nemo_platform_ext/src/nemo_platform_ext/local/transport.py new file mode 100644 index 0000000000..da3c08344a --- /dev/null +++ b/packages/nemo_platform_ext/src/nemo_platform_ext/local/transport.py @@ -0,0 +1,154 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Local service transport helpers for TCP and Unix domain sockets.""" + +from __future__ import annotations + +import asyncio +import time +from pathlib import Path +from typing import Any, TypeAlias + +import httpx +from fastapi.testclient import TestClient +from nmp.common.platform_endpoint import UDS_BASE_URL + +HttpxTimeout: TypeAlias = float | httpx.Timeout | None +_DEFAULT_TIMEOUT: float = 5.0 +EMBEDDED_BASE_URL = "http://nemo-platform.local" + +__all__ = [ + "EMBEDDED_BASE_URL", + "UDS_BASE_URL", + "build_async_asgi_http_client", + "build_async_http_client", + "build_sync_asgi_http_client", + "build_sync_http_client", + "probe_status", + "probe_status_async", + "tcp_base_url", + "wait_for_status", + "wait_for_status_async", +] + + +def build_sync_asgi_http_client(app: Any, *, timeout: HttpxTimeout = _DEFAULT_TIMEOUT) -> Any: + _ = timeout + return TestClient( + app, + base_url=EMBEDDED_BASE_URL, + follow_redirects=True, + ) + + +def build_async_asgi_http_client(app: Any, *, timeout: HttpxTimeout = _DEFAULT_TIMEOUT) -> httpx.AsyncClient: + return httpx.AsyncClient( + transport=httpx.ASGITransport(app=app), + base_url=EMBEDDED_BASE_URL, + follow_redirects=True, + timeout=timeout, + ) + + +def build_sync_http_client(socket_path: Path, *, timeout: HttpxTimeout = _DEFAULT_TIMEOUT) -> httpx.Client: + return httpx.Client( + transport=httpx.HTTPTransport(uds=str(socket_path)), + follow_redirects=True, + timeout=timeout, + ) + + +def build_async_http_client(socket_path: Path, *, timeout: HttpxTimeout = _DEFAULT_TIMEOUT) -> httpx.AsyncClient: + return httpx.AsyncClient( + transport=httpx.AsyncHTTPTransport(uds=str(socket_path)), + follow_redirects=True, + timeout=timeout, + ) + + +def tcp_base_url(host: str, port: int) -> str: + connect_host = "localhost" if host in {"0.0.0.0", "::"} else host # noqa: S104 + normalized = connect_host.strip("[]") + url_host = f"[{normalized}]" if ":" in normalized else normalized + return str(httpx.URL(scheme="http", host=url_host, port=port)) + + +def probe_status( + *, + base_url: str, + socket_path: Path | None = None, + timeout: float = 2.0, +) -> bool: + client = ( + build_sync_http_client(socket_path, timeout=timeout) + if socket_path is not None + else httpx.Client(timeout=timeout) + ) + try: + response = client.get(f"{base_url.rstrip('/')}/status") + return response.status_code == 200 + except httpx.RequestError: + return False + finally: + client.close() + + +async def probe_status_async( + *, + base_url: str, + socket_path: Path | None = None, + timeout: float = 2.0, +) -> bool: + client = ( + build_async_http_client(socket_path, timeout=timeout) + if socket_path is not None + else httpx.AsyncClient(timeout=timeout) + ) + try: + response = await client.get(f"{base_url.rstrip('/')}/status") + return response.status_code == 200 + except httpx.RequestError: + return False + finally: + await client.aclose() + + +def wait_for_status( + *, + base_url: str, + socket_path: Path | None = None, + timeout: float = 60.0, + poll_interval: float = 0.5, +) -> bool: + deadline = time.monotonic() + timeout + while True: + remaining = deadline - time.monotonic() + if remaining <= 0: + return False + if probe_status(base_url=base_url, socket_path=socket_path, timeout=remaining): + return True + remaining = deadline - time.monotonic() + if remaining <= 0: + return False + time.sleep(min(poll_interval, remaining)) + + +async def wait_for_status_async( + *, + base_url: str, + socket_path: Path | None = None, + timeout: float = 60.0, + poll_interval: float = 0.5, +) -> bool: + deadline = time.monotonic() + timeout + while True: + remaining = deadline - time.monotonic() + if remaining <= 0: + return False + if await probe_status_async(base_url=base_url, socket_path=socket_path, timeout=remaining): + return True + remaining = deadline - time.monotonic() + if remaining <= 0: + return False + await asyncio.sleep(min(poll_interval, remaining)) diff --git a/packages/nemo_platform_ext/tests/cli/commands/conftest.py b/packages/nemo_platform_ext/tests/cli/commands/conftest.py index 24df2fb6c0..ff6593b338 100644 --- a/packages/nemo_platform_ext/tests/cli/commands/conftest.py +++ b/packages/nemo_platform_ext/tests/cli/commands/conftest.py @@ -3,7 +3,7 @@ from __future__ import annotations -import nemo_platform_ext.cli.commands.services._process as _process_mod +import nemo_platform_ext.local.process as _process_mod import pytest diff --git a/packages/nemo_platform_ext/tests/cli/commands/test_services.py b/packages/nemo_platform_ext/tests/cli/commands/test_services.py index 1a78c52a62..7c939fcb8f 100644 --- a/packages/nemo_platform_ext/tests/cli/commands/test_services.py +++ b/packages/nemo_platform_ext/tests/cli/commands/test_services.py @@ -13,11 +13,11 @@ import socket from pathlib import Path from types import ModuleType -from unittest.mock import ANY, MagicMock, patch +from unittest.mock import MagicMock, patch import pytest from nemo_platform_ext.cli.app import app -from nemo_platform_ext.cli.commands.services._process import ( +from nemo_platform_ext.local.process import ( ForegroundInstanceError, InstanceDescriptor, StopResult, @@ -26,16 +26,17 @@ read_descriptor, write_descriptor, ) +from nmp.platform_runner.config import PlatformAppConfig from typer.testing import CliRunner runner = CliRunner() -_PROCESS_MODULE = "nemo_platform_ext.cli.commands.services._process" +_PROCESS_MODULE = "nemo_platform_ext.local.process" _CLI_MODULE = "nemo_platform_ext.cli.commands.services.cli" def _seed_stopped_scope(base_dir: Path, scope: str, *, log_content: str = "x\n") -> Path: - """Create a stopped instance directory with service logs.""" + """Create a stopped scope directory with service logs.""" d = instance_dir(scope, base_dir=base_dir) (d / "services.log").write_text(log_content) return d @@ -118,17 +119,18 @@ def test_run_invokes_runner(base_dir: Path): ) assert result.exit_code == 0, result.stderr - mock_run_platform.assert_called_once_with( - services=["auth", "entities"], - service_group=None, - controllers=["jobs", "models"], - controller_group=None, - sidecars=None, - config_path=None, - host="127.0.0.1", - port=9000, - on_shutdown=ANY, - ) + mock_run_platform.assert_called_once() + _, kwargs = mock_run_platform.call_args + config = kwargs["config"] + assert config.services == ["auth", "entities"] + assert config.service_group is None + assert config.controllers == ["jobs", "models"] + assert config.controller_group is None + assert config.sidecars is None + assert config.config_path is None + assert config.host == "127.0.0.1" + assert config.port == 9000 + assert kwargs["on_shutdown"] is not None def test_run_refuses_when_already_running(base_dir: Path): @@ -163,7 +165,7 @@ def test_run_writes_descriptor(base_dir: Path): desc = read_descriptor("desc-test", base_dir=base_dir) assert desc is not None assert desc.mode == "foreground" - assert desc.port == 9999 + assert desc.config.port == 9999 def test_run_records_background_mode_when_launched_by_start(base_dir: Path): @@ -382,7 +384,7 @@ def test_restart_errors_when_no_prior_instance(self, base_dir: Path): ["services", "restart", "--instance", "ghost"], ) assert result.exit_code == 1 - assert "No instance found" in result.stderr + assert "No instance found for scope" in result.stderr assert "nemo services start" in result.stderr def test_restart_stops_and_starts(self, base_dir: Path): @@ -390,9 +392,7 @@ def test_restart_stops_and_starts(self, base_dir: Path): fd = acquire_lock(scope, base_dir=base_dir) desc = InstanceDescriptor( pid=os.getpid(), - scope=scope, - host="127.0.0.1", - port=8080, + config=PlatformAppConfig(scope=scope), mode="background", create_time=1.0, ) @@ -429,9 +429,7 @@ def test_restart_exits_early_when_port_occupied_by_foreign_process(self, base_di desc = InstanceDescriptor( pid=99999, - scope=scope, - host="127.0.0.1", - port=port, + config=PlatformAppConfig(scope=scope, host="127.0.0.1", port=port), mode="background", create_time=1.0, ) @@ -459,13 +457,15 @@ def test_restart_preserves_previous_args(self, base_dir: Path): fd = acquire_lock(scope, base_dir=base_dir) desc = InstanceDescriptor( pid=os.getpid(), - scope=scope, - host="127.0.0.1", - port=9000, + config=PlatformAppConfig( + scope=scope, + services=["entities", "models"], + controllers=["jobs"], + host="127.0.0.1", + port=9000, + ), mode="background", create_time=1.0, - services=["entities", "models"], - controllers=["jobs"], ) write_descriptor(desc, base_dir=base_dir) @@ -488,11 +488,12 @@ def test_restart_preserves_previous_args(self, base_dir: Path): os.close(fd) assert result.exit_code == 0 - _, kwargs = mock_start.call_args - assert kwargs["services"] == ["entities", "models"] - assert kwargs["controllers"] == ["jobs"] - assert kwargs["host"] == "127.0.0.1" - assert kwargs["port"] == 9000 + args, _kwargs = mock_start.call_args + config = args[0] + assert config.services == ["entities", "models"] + assert config.controllers == ["jobs"] + assert config.host == "127.0.0.1" + assert config.port == 9000 # --------------------------------------------------------------------------- @@ -504,16 +505,14 @@ class TestServicesStatus: def test_not_running(self, base_dir: Path): result = runner.invoke(app, ["services", "status", "--instance", "none"]) assert result.exit_code == 0 - assert "No running instance" in result.stdout + assert "No running instance for scope" in result.stdout def test_running_instance(self, base_dir: Path): scope = "status-test" fd = acquire_lock(scope, base_dir=base_dir) desc = InstanceDescriptor( pid=os.getpid(), - scope=scope, - host="127.0.0.1", - port=8080, + config=PlatformAppConfig(scope=scope), mode="foreground", create_time=1.0, ) @@ -548,9 +547,7 @@ def test_lists_running_instance(self, base_dir: Path): fd = acquire_lock(scope, base_dir=base_dir) desc = InstanceDescriptor( pid=os.getpid(), - scope=scope, - host="127.0.0.1", - port=8080, + config=PlatformAppConfig(scope=scope), mode="background", create_time=1.0, ) @@ -594,9 +591,7 @@ def test_mixed_running_and_stopped(self, base_dir: Path): write_descriptor( InstanceDescriptor( pid=os.getpid(), - scope=running_scope, - host="127.0.0.1", - port=8080, + config=PlatformAppConfig(scope=running_scope), mode="background", create_time=1.0, ), @@ -658,7 +653,7 @@ def test_rm_requires_scope(self, base_dir: Path): def test_rm_rejects_invalid_scope(self, base_dir: Path): result = runner.invoke(app, ["services", "rm", "../escape"]) assert result.exit_code == 1 - assert "Invalid instance scope" in result.stderr + assert "Invalid scope" in result.stderr def test_rm_rejects_conflicting_scope_args(self, base_dir: Path): result = runner.invoke(app, ["services", "rm", "scope-a", "--instance", "scope-b"]) @@ -747,7 +742,7 @@ def test_default_host_is_loopback(base_dir: Path): assert result.exit_code == 0, result.stderr _, kwargs = mock_run_platform.call_args - assert kwargs["host"] == "127.0.0.1" + assert kwargs["config"].host == "127.0.0.1" def test_bind_all_warning(base_dir: Path): diff --git a/packages/nemo_platform_ext/tests/cli/commands/test_services_lifecycle.py b/packages/nemo_platform_ext/tests/cli/commands/test_services_lifecycle.py index fd45922a1c..ad382426c4 100644 --- a/packages/nemo_platform_ext/tests/cli/commands/test_services_lifecycle.py +++ b/packages/nemo_platform_ext/tests/cli/commands/test_services_lifecycle.py @@ -27,7 +27,7 @@ import pytest from nemo_platform_ext.cli.app import app -from nemo_platform_ext.cli.commands.services._process import ( +from nemo_platform_ext.local.process import ( InstanceDescriptor, PortConflict, acquire_lock, @@ -43,6 +43,7 @@ stop_instance, write_descriptor, ) +from nmp.platform_runner.config import PlatformAppConfig from typer.testing import CliRunner _runner = CliRunner() @@ -70,18 +71,14 @@ import psutil as _psutil desc = { "pid": os.getpid(), - "scope": scope, - "host": "127.0.0.1", - "port": 8080, + "config": { + "scope": scope, + "host": "127.0.0.1", + "port": 8080, + }, "mode": "background", "create_time": _psutil.Process(os.getpid()).create_time(), "started_at": "test", - "services": None, - "controllers": None, - "service_group": None, - "controller_group": None, - "sidecars": None, - "config_path": None, "log_path": None, } desc_path = os.path.join(inst_dir, "instance.json") @@ -291,9 +288,7 @@ def test_stale_descriptor_with_reused_pid(self, tmp_path: Path) -> None: # Write a descriptor with the sleeper's PID but wrong create_time desc = InstanceDescriptor( pid=sleeper.pid, - scope=scope, - host="127.0.0.1", - port=8080, + config=PlatformAppConfig(scope=scope), mode="background", create_time=0.0, # intentionally wrong ) @@ -363,7 +358,7 @@ def test_log_preserved_across_restart(self, tmp_path: Path) -> None: log = d / "services.log" log.write_text("first boot log content\n") - from nemo_platform_ext.cli.commands.services._process import rotate_log + from nemo_platform_ext.local.process import rotate_log new_log = rotate_log(scope, base_dir=base_dir) new_log.write_text("second boot log content\n") @@ -399,15 +394,15 @@ def test_log_preserved_across_restart(self, tmp_path: Path) -> None: import psutil as _psutil desc = { "pid": os.getpid(), - "scope": scope, - "host": "127.0.0.1", - "port": port, + "config": { + "scope": scope, + "host": "127.0.0.1", + "port": port, + }, "mode": "background", "create_time": _psutil.Process(os.getpid()).create_time(), "started_at": "test", - "services": None, "controllers": None, - "service_group": None, "controller_group": None, - "sidecars": None, "config_path": None, "log_path": None, + "log_path": None, } desc_path = os.path.join(inst_dir, "instance.json") with open(desc_path, "w") as f: @@ -523,7 +518,7 @@ def test_stop_after_health_check(self, tmp_path: Path) -> None: assert is_instance_alive(scope, base_dir=base_dir) desc = read_descriptor(scope, base_dir=base_dir) assert desc is not None - assert desc.port == port + assert desc.config.port == port result = stop_instance(scope, base_dir=base_dir, timeout=5.0) assert proc.pid in result.stopped_pids @@ -536,7 +531,7 @@ def test_stop_after_health_check(self, tmp_path: Path) -> None: class TestInstanceCleanup: - """Integration tests for rm/prune and post-stop instance directories.""" + """Integration tests for rm/prune and post-stop scope directories.""" def test_stop_leaves_record_until_rm(self, tmp_path: Path, monkeypatch) -> None: base_dir = tmp_path / "state" @@ -668,9 +663,7 @@ def test_check_port_returns_nemo_instance_when_lock_held_and_port_blocked(self, write_descriptor( InstanceDescriptor( pid=os.getpid(), - scope=scope, - host="127.0.0.1", - port=port, + config=PlatformAppConfig(scope=scope, host="127.0.0.1", port=port), mode="background", create_time=1.0, ), @@ -704,9 +697,7 @@ def test_check_port_returns_foreign_when_alive_instance_uses_different_port(self write_descriptor( InstanceDescriptor( pid=os.getpid(), - scope=scope, - host="127.0.0.1", - port=nemo_port, + config=PlatformAppConfig(scope=scope, host="127.0.0.1", port=nemo_port), mode="background", create_time=1.0, ), diff --git a/packages/nemo_platform_ext/tests/cli/commands/test_services_process.py b/packages/nemo_platform_ext/tests/cli/commands/test_services_process.py index fa18e47a65..26478c20d5 100644 --- a/packages/nemo_platform_ext/tests/cli/commands/test_services_process.py +++ b/packages/nemo_platform_ext/tests/cli/commands/test_services_process.py @@ -6,6 +6,7 @@ from __future__ import annotations import os +import signal import subprocess import sys import time @@ -14,7 +15,8 @@ import psutil import pytest -from nemo_platform_ext.cli.commands.services._process import ( +from nemo_platform_ext.local import process as process_module +from nemo_platform_ext.local.process import ( ForegroundInstanceError, InstanceAlreadyRunningError, InstanceDescriptor, @@ -40,6 +42,7 @@ validate_pid, write_descriptor, ) +from nmp.platform_runner.config import DEFAULT_LOCAL_SERVICES_BIND_HOST, PlatformAppConfig @pytest.fixture() @@ -48,67 +51,23 @@ def base_dir(tmp_path: Path) -> Path: # --------------------------------------------------------------------------- -# Scope computation +# Scope resolution # --------------------------------------------------------------------------- class TestComputeScope: - def test_explicit_instance_name(self) -> None: - assert compute_scope(port=8080, instance_name="myapp") == "myapp" + def test_explicit_scope(self) -> None: + assert compute_scope(port=1234, explicit_scope="myapp") == "myapp" - def test_default_scope_includes_port(self) -> None: - scope = compute_scope(port=9090) - assert scope.endswith("-9090") - - def test_default_scope_is_deterministic(self) -> None: - a = compute_scope(port=8080) - b = compute_scope(port=8080) - assert a == b - - def test_different_ports_different_scopes(self) -> None: - a = compute_scope(port=8080) - b = compute_scope(port=9090) - assert a != b - - def test_hash_prefix_is_8_chars(self) -> None: + def test_default_scope_is_stable_for_port(self) -> None: scope = compute_scope(port=8080) - prefix = scope.rsplit("-", 1)[0] - assert len(prefix) == 8 - - def test_git_failure_falls_back_to_cwd(self) -> None: - import nemo_platform_ext.cli.commands.services._process as proc_mod - - proc_mod._scope_prefix_cache = None - try: - with patch.object(proc_mod, "_find_git_root", return_value="/no/git/here"): - scope = compute_scope(port=8080) - assert scope.endswith("-8080") - assert len(scope.rsplit("-", 1)[0]) == 8 - finally: - proc_mod._scope_prefix_cache = None - def test_different_git_roots_produce_different_prefixes(self) -> None: - """Two different working directories (worktrees) produce distinct scopes.""" - import nemo_platform_ext.cli.commands.services._process as proc_mod - - with patch.object(proc_mod, "_find_git_root", return_value="/workspace/project-a"): - scope_a = compute_scope(port=8080) - - proc_mod._scope_prefix_cache = None - - with patch.object(proc_mod, "_find_git_root", return_value="/workspace/project-b"): - scope_b = compute_scope(port=8080) - - assert scope_a != scope_b - assert scope_a.endswith("-8080") - assert scope_b.endswith("-8080") - prefix_a = scope_a.rsplit("-", 1)[0] - prefix_b = scope_b.rsplit("-", 1)[0] - assert prefix_a != prefix_b + assert scope == compute_scope(port=8080) + assert scope.endswith("-8080") # --------------------------------------------------------------------------- -# Instance directory +# Scope directory # --------------------------------------------------------------------------- @@ -168,26 +127,27 @@ class TestDescriptorRoundTrip: def test_write_and_read(self, base_dir: Path) -> None: desc = InstanceDescriptor( pid=12345, - scope="test-8080", - host="127.0.0.1", - port=8080, + config=PlatformAppConfig( + scope="test-8080", + services=["entities", "models"], + controllers=["jobs"], + host="127.0.0.1", + ), mode="background", create_time=1000.0, - services=["entities", "models"], - controllers=["jobs"], ) write_descriptor(desc, base_dir=base_dir) recovered = read_descriptor("test-8080", base_dir=base_dir) assert recovered is not None assert recovered.pid == 12345 - assert recovered.scope == "test-8080" - assert recovered.host == "127.0.0.1" - assert recovered.port == 8080 + assert recovered.config.scope == "test-8080" + assert recovered.config.host == "127.0.0.1" + assert recovered.config.port == 8080 assert recovered.mode == "background" assert recovered.create_time == 1000.0 - assert recovered.services == ["entities", "models"] - assert recovered.controllers == ["jobs"] + assert recovered.config.services == ["entities", "models"] + assert recovered.config.controllers == ["jobs"] def test_read_missing_returns_none(self, base_dir: Path) -> None: assert read_descriptor("no-such-scope", base_dir=base_dir) is None @@ -200,9 +160,7 @@ def test_read_corrupt_returns_none(self, base_dir: Path) -> None: def test_remove_descriptor(self, base_dir: Path) -> None: desc = InstanceDescriptor( pid=1, - scope="rm-test", - host="127.0.0.1", - port=8080, + config=PlatformAppConfig(scope="rm-test"), mode="background", create_time=1.0, ) @@ -245,9 +203,7 @@ def test_lists_alive_instance(self, base_dir: Path) -> None: fd = acquire_lock("alive-one", base_dir=base_dir) desc = InstanceDescriptor( pid=os.getpid(), - scope="alive-one", - host="127.0.0.1", - port=8080, + config=PlatformAppConfig(scope="alive-one"), mode="foreground", create_time=1.0, ) @@ -265,9 +221,7 @@ def test_cleans_up_dead_descriptor(self, base_dir: Path) -> None: d = instance_dir("dead-scope", base_dir=base_dir) desc = InstanceDescriptor( pid=999999, - scope="dead-scope", - host="127.0.0.1", - port=8080, + config=PlatformAppConfig(scope="dead-scope"), mode="background", create_time=1.0, ) @@ -281,9 +235,7 @@ def test_stale_descriptor_with_logs_stays_listed(self, base_dir: Path) -> None: d = instance_dir("dead-with-logs", base_dir=base_dir) desc = InstanceDescriptor( pid=999999, - scope="dead-with-logs", - host="127.0.0.1", - port=8080, + config=PlatformAppConfig(scope="dead-with-logs"), mode="background", create_time=1.0, ) @@ -370,7 +322,7 @@ def test_refuses_running_instance(self, base_dir: Path) -> None: os.close(fd) def test_rejects_invalid_scope(self, base_dir: Path) -> None: - with pytest.raises(ValueError, match="Invalid instance scope"): + with pytest.raises(ValueError, match="Invalid scope"): remove_instance("../escape", base_dir=base_dir) def test_returns_false_when_rmtree_fails(self, base_dir: Path) -> None: @@ -378,7 +330,7 @@ def test_returns_false_when_rmtree_fails(self, base_dir: Path) -> None: (d / "services.log").write_text("logs\n") with patch( - "nemo_platform_ext.cli.commands.services._process.shutil.rmtree", + "nemo_platform_ext.local.process.shutil.rmtree", side_effect=OSError("permission denied"), ): assert remove_instance("rmtree-fail", base_dir=base_dir) is False @@ -489,9 +441,7 @@ def test_stops_running_process(self, base_dir: Path) -> None: fd = acquire_lock(scope, base_dir=base_dir) desc = InstanceDescriptor( pid=proc.pid, - scope=scope, - host="127.0.0.1", - port=8080, + config=PlatformAppConfig(scope=scope), mode="background", create_time=psutil.Process(proc.pid).create_time(), ) @@ -512,9 +462,7 @@ def test_cleans_up_stale_descriptor(self, base_dir: Path) -> None: scope = "stale" desc = InstanceDescriptor( pid=999999999, - scope=scope, - host="127.0.0.1", - port=8080, + config=PlatformAppConfig(scope=scope), mode="background", create_time=0.0, ) @@ -529,9 +477,7 @@ def test_refuses_to_stop_foreground_instance(self, base_dir: Path) -> None: try: desc = InstanceDescriptor( pid=os.getpid(), - scope=scope, - host="127.0.0.1", - port=8080, + config=PlatformAppConfig(scope=scope), mode="foreground", create_time=1.0, ) @@ -555,9 +501,7 @@ def test_force_stops_foreground_instance(self, base_dir: Path) -> None: try: desc = InstanceDescriptor( pid=proc.pid, - scope=scope, - host="127.0.0.1", - port=8080, + config=PlatformAppConfig(scope=scope), mode="foreground", create_time=psutil.Process(proc.pid).create_time(), ) @@ -571,6 +515,34 @@ def test_force_stops_foreground_instance(self, base_dir: Path) -> None: proc.kill() proc.wait(timeout=5) + def test_preserves_descriptor_when_sigkill_does_not_stop_parent(self, base_dir: Path, monkeypatch) -> None: + scope = "sigkill-still-alive" + desc = InstanceDescriptor( + pid=12345, + config=PlatformAppConfig(scope=scope), + mode="background", + create_time=1.0, + ) + write_descriptor(desc, base_dir=base_dir) + kill_signals: list[int] = [] + + def fake_kill(_pid: int, sig: int) -> None: + kill_signals.append(sig) + + monkeypatch.setattr(process_module, "validate_pid", lambda _pid, _create_time: True) + monkeypatch.setattr(process_module, "_pid_alive", lambda _pid: True) + monkeypatch.setattr(process_module, "_snapshot_children", lambda _pid: [object()]) + monkeypatch.setattr(process_module, "_sweep_orphans", lambda _children: [222]) + monkeypatch.setattr(process_module, "_SIGKILL_WAIT_TIMEOUT", 0.0) + monkeypatch.setattr(process_module.os, "kill", fake_kill) + + result = stop_instance(scope, base_dir=base_dir, timeout=0.0) + + assert kill_signals == [signal.SIGTERM, signal.SIGKILL] + assert result.stopped_pids == [] + assert result.swept_children == [222] + assert read_descriptor(scope, base_dir=base_dir) is not None + # --------------------------------------------------------------------------- # start_background @@ -578,21 +550,49 @@ def test_force_stops_foreground_instance(self, base_dir: Path) -> None: class TestStartBackground: + def test_uses_default_platform_app_config(self, tmp_path: Path, monkeypatch) -> None: + monkeypatch.setenv("XDG_STATE_HOME", str(tmp_path)) + mock_proc = MagicMock() + mock_proc.pid = 99998 + captured_args: list[str] = [] + captured_env: dict[str, str] = {} + + def fake_popen(args, **kwargs): + captured_args.extend(args) + captured_env.update(kwargs["env"]) + return mock_proc + + with patch( + "nemo_platform_ext.local.process.subprocess.Popen", + side_effect=fake_popen, + ): + proc = start_background() + + assert proc.pid == 99998 + assert captured_args[captured_args.index("--instance") + 1] == "default" + assert captured_args[captured_args.index("--host") + 1] == "127.0.0.1" + assert captured_args[captured_args.index("--port") + 1] == "8080" + assert captured_env["XDG_STATE_HOME"] == str(tmp_path) + assert "_NMP_STATE_DIR" not in captured_env + assert (tmp_path / "nmp" / "instances" / "default" / "services.log").exists() + def test_launches_detached_subprocess(self, base_dir: Path) -> None: mock_proc = MagicMock() mock_proc.pid = 99999 with patch( - "nemo_platform_ext.cli.commands.services._process.subprocess.Popen", + "nemo_platform_ext.local.process.subprocess.Popen", return_value=mock_proc, ) as mock_popen: proc = start_background( - scope="bg-test", - services=["entities", "models"], - controllers=["jobs"], - host="127.0.0.1", - port=8080, - base_dir=base_dir, + PlatformAppConfig( + scope="bg-test", + services=["entities", "models"], + controllers=["jobs"], + host="127.0.0.1", + port=8080, + state_root=base_dir, + ), ) assert proc.pid == 99999 @@ -612,13 +612,16 @@ def fake_popen(args, **kwargs): return mock_proc with patch( - "nemo_platform_ext.cli.commands.services._process.subprocess.Popen", + "nemo_platform_ext.local.process.subprocess.Popen", side_effect=fake_popen, ): start_background( - scope="data-dir-test", + PlatformAppConfig( + scope="data-dir-test", + host=DEFAULT_LOCAL_SERVICES_BIND_HOST, + state_root=base_dir, + ), data_dir="/chosen/data/dir", - base_dir=base_dir, ) assert captured_env.get("NMP_DATA_DIR") == "/chosen/data/dir" @@ -634,13 +637,16 @@ def fake_popen(args, **kwargs): return mock_proc with patch( - "nemo_platform_ext.cli.commands.services._process.subprocess.Popen", + "nemo_platform_ext.local.process.subprocess.Popen", side_effect=fake_popen, ): start_background( - scope="shell-env-test", + PlatformAppConfig( + scope="shell-env-test", + host=DEFAULT_LOCAL_SERVICES_BIND_HOST, + state_root=base_dir, + ), data_dir="/chosen/data/dir", - base_dir=base_dir, ) assert captured_env.get("NMP_DATA_DIR") == "/shell/wins" @@ -654,16 +660,22 @@ def test_rotates_log_before_start(self, base_dir: Path) -> None: mock_proc.pid = 5555 with patch( - "nemo_platform_ext.cli.commands.services._process.subprocess.Popen", + "nemo_platform_ext.local.process.subprocess.Popen", return_value=mock_proc, ): - start_background(scope="rotate-test", base_dir=base_dir) + start_background( + PlatformAppConfig( + scope="rotate-test", + host=DEFAULT_LOCAL_SERVICES_BIND_HOST, + state_root=base_dir, + ), + ) rotated = list(d.glob("services.log.*")) assert len(rotated) == 1 assert rotated[0].read_text() == "old log content\n" - def test_forwards_instance_scope_to_child(self, base_dir: Path) -> None: + def test_forwards_scope_to_child(self, base_dir: Path) -> None: mock_proc = MagicMock() mock_proc.pid = 7777 captured_args: list[str] = [] @@ -673,20 +685,22 @@ def fake_popen(args, **kwargs): return mock_proc with patch( - "nemo_platform_ext.cli.commands.services._process.subprocess.Popen", + "nemo_platform_ext.local.process.subprocess.Popen", side_effect=fake_popen, ): start_background( - scope="custom-scope", - services=["entities"], - host="127.0.0.1", - port=9090, - base_dir=base_dir, + PlatformAppConfig( + scope="custom-key", + services=["entities"], + host="127.0.0.1", + port=9090, + state_root=base_dir, + ), ) assert "--instance" in captured_args idx = captured_args.index("--instance") - assert captured_args[idx + 1] == "custom-scope" + assert captured_args[idx + 1] == "custom-key" def test_sets_launch_mode_background_in_child_env(self, base_dir: Path) -> None: mock_proc = MagicMock() @@ -698,10 +712,16 @@ def fake_popen(args, **kwargs): return mock_proc with patch( - "nemo_platform_ext.cli.commands.services._process.subprocess.Popen", + "nemo_platform_ext.local.process.subprocess.Popen", side_effect=fake_popen, ): - start_background(scope="mode-test", base_dir=base_dir) + start_background( + PlatformAppConfig( + scope="mode-test", + host=DEFAULT_LOCAL_SERVICES_BIND_HOST, + state_root=base_dir, + ), + ) assert captured_env.get("_NMP_LAUNCH_MODE") == "background" @@ -846,9 +866,7 @@ def test_sweeps_surviving_children(self, base_dir: Path) -> None: fd = acquire_lock(scope, base_dir=base_dir) desc = InstanceDescriptor( pid=parent.pid, - scope=scope, - host="127.0.0.1", - port=8080, + config=PlatformAppConfig(scope=scope), mode="background", create_time=psutil.Process(parent.pid).create_time(), ) @@ -883,9 +901,7 @@ def test_swept_children_empty_when_no_children(self, base_dir: Path) -> None: fd = acquire_lock(scope, base_dir=base_dir) desc = InstanceDescriptor( pid=proc.pid, - scope=scope, - host="127.0.0.1", - port=8080, + config=PlatformAppConfig(scope=scope), mode="background", create_time=psutil.Process(proc.pid).create_time(), ) @@ -917,9 +933,7 @@ def test_restart_path_sweeps_children(self, base_dir: Path) -> None: fd = acquire_lock(scope, base_dir=base_dir) desc = InstanceDescriptor( pid=parent.pid, - scope=scope, - host="127.0.0.1", - port=8080, + config=PlatformAppConfig(scope=scope), mode="foreground", create_time=psutil.Process(parent.pid).create_time(), ) diff --git a/packages/nemo_platform_ext/tests/cli/commands/test_setup.py b/packages/nemo_platform_ext/tests/cli/commands/test_setup.py index a911b1a3a1..bfc66f71e9 100644 --- a/packages/nemo_platform_ext/tests/cli/commands/test_setup.py +++ b/packages/nemo_platform_ext/tests/cli/commands/test_setup.py @@ -14,7 +14,6 @@ import typer from click.exceptions import Exit as ClickExit from nemo_platform.resources.inference.providers import ProvidersResource -from nemo_platform_ext.cli.commands.services._process import PortConflict from nemo_platform_ext.cli.commands.setup import ( _AGENT_API_READINESS_POLL_INTERVAL, _AGENT_DEPLOY_POLL_INTERVAL, @@ -77,6 +76,7 @@ Context, ContextDefinition, ) +from nemo_platform_ext.local.process import PortConflict from nemo_platform_plugin.client.errors import NotFoundError from nemo_platform_plugin.secrets.types import PlatformSecretCreateRequest, PlatformSecretUpdateRequest @@ -644,9 +644,10 @@ def test_start_services_background_forwards_data_dir(self): mock_start.return_value = MagicMock(pid=42) _start_services_background("http://localhost:9090", data_dir="/chosen/data/dir") mock_start.assert_called_once() - _, kwargs = mock_start.call_args + args, kwargs = mock_start.call_args + config = args[0] assert kwargs["data_dir"] == "/chosen/data/dir" - assert kwargs["port"] == 9090 + assert config.port == 9090 def test_auto_mode_skips_prompt_and_uses_persisted(self, tmp_path, monkeypatch): """`--auto` must not prompt but should still honor any persisted data dir.""" diff --git a/packages/nemo_platform_ext/tests/local/test_config_environment.py b/packages/nemo_platform_ext/tests/local/test_config_environment.py new file mode 100644 index 0000000000..8255b3b6be --- /dev/null +++ b/packages/nemo_platform_ext/tests/local/test_config_environment.py @@ -0,0 +1,196 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Integration tests for configuration and environment resolution. + +These tests exercise the real ``apply_run_environment`` code path with +actual YAML config files, verifying that environment variables are set +correctly for different host, port, and base_url scenarios. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest +from nmp.platform_runner.config import ( + ResolvedRunConfiguration, + apply_run_environment, + default_config_path, +) + + +def _resolved( + *, + services: set[str] | None = None, + controllers: set[str] | None = None, + sidecars: set[str] | None = None, + host: str = "127.0.0.1", + port: int = 8080, + config_path: str | None = None, + socket_path: str | None = None, +) -> ResolvedRunConfiguration: + return ResolvedRunConfiguration( + services=services or set(), + controllers=controllers or set(), + sidecars=sidecars or set(), + host=host, + port=port, + config_path=config_path or default_config_path(), + socket_path=socket_path, + available_services={}, + available_controllers={}, + ) + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +@pytest.mark.integration +def test_apply_run_environment_preserves_external_base_url() -> None: + """Pre-set NMP_BASE_URL (e.g. from k8s/Helm) must not be overwritten.""" + env: dict[str, str] = {"NMP_BASE_URL": "https://platform.k8s.internal:443"} + config = _resolved(host="0.0.0.0", port=9090) + + apply_run_environment(config, env=env) + + assert env["NMP_BASE_URL"] == "https://platform.k8s.internal:443" + + +@pytest.mark.integration +def test_apply_run_environment_wildcard_host_becomes_loopback(tmp_path: Path) -> None: + """A wildcard bind host (0.0.0.0) in the config file should resolve to + 127.0.0.1 for the base URL, using the actual bind port.""" + config_file = tmp_path / "config.yaml" + config_file.write_text("platform:\n base_url: http://0.0.0.0:8080\n") + + env: dict[str, str] = {} + config = _resolved(host="0.0.0.0", port=9090, config_path=str(config_file)) + + apply_run_environment(config, env=env) + + assert env["NMP_BASE_URL"] == "http://127.0.0.1:9090" + assert env["NMP_SERVICE_HOST"] == "127.0.0.1" + assert env["NMP_SERVICE_PORT"] == "9090" + + +@pytest.mark.integration +def test_apply_run_environment_ipv6_literal_bracketed(tmp_path: Path) -> None: + """An IPv6 config base_url should produce a bracketed host in the resolved URL.""" + config_file = tmp_path / "config.yaml" + config_file.write_text("platform:\n base_url: http://[::1]:8080\n") + + env: dict[str, str] = {} + config = _resolved(host="::1", port=9090, config_path=str(config_file)) + + apply_run_environment(config, env=env) + + assert env["NMP_BASE_URL"] == "http://[::1]:9090" + + +@pytest.mark.integration +def test_config_file_base_url_malformed_yaml_falls_back(tmp_path: Path) -> None: + """A corrupt config file should fall back to the bind-derived URL.""" + config_file = tmp_path / "config.yaml" + config_file.write_text("{{{{not valid yaml at all") + + env: dict[str, str] = {} + config = _resolved(host="127.0.0.1", port=7777, config_path=str(config_file)) + + apply_run_environment(config, env=env) + + # Falls back to bind-derived: http://: + assert env["NMP_BASE_URL"] == "http://127.0.0.1:7777" + + +@pytest.mark.integration +def test_config_file_missing_falls_back(tmp_path: Path) -> None: + """A missing config file should fall back to the bind-derived URL.""" + env: dict[str, str] = {} + config = _resolved(host="127.0.0.1", port=5555, config_path=str(tmp_path / "nonexistent.yaml")) + + apply_run_environment(config, env=env) + + assert env["NMP_BASE_URL"] == "http://127.0.0.1:5555" + + +@pytest.mark.integration +def test_apply_run_environment_clears_empty_service_lists() -> None: + """When services/controllers/sidecars are empty sets, their env vars + should be removed (popped) rather than set to empty strings.""" + env: dict[str, str] = { + "NMP_SERVICES": "old-service", + "NMP_CONTROLLERS": "old-controller", + "NMP_SIDECARS": "old-sidecar", + } + config = _resolved(services=set(), controllers=set(), sidecars=set()) + + apply_run_environment(config, env=env) + + assert "NMP_SERVICES" not in env + assert "NMP_CONTROLLERS" not in env + assert "NMP_SIDECARS" not in env + + +@pytest.mark.integration +def test_apply_run_environment_sets_service_lists() -> None: + """Non-empty service/controller/sidecar sets should be written as + comma-separated, sorted env var values.""" + env: dict[str, str] = {} + config = _resolved( + services={"models", "auth", "secrets"}, + controllers={"beta-controller"}, + sidecars={"adapters"}, + ) + + apply_run_environment(config, env=env) + + assert env["NMP_SERVICES"] == "auth,models,secrets" + assert env["NMP_CONTROLLERS"] == "beta-controller" + assert env["NMP_SIDECARS"] == "adapters" + + +@pytest.mark.integration +def test_apply_run_environment_uds_transport_uses_unix_base_url() -> None: + """When a socket_path is set (UDS transport), the base URL should use + the ``unix://`` scheme.""" + env: dict[str, str] = {} + config = _resolved(socket_path="/tmp/nemo.sock") + + apply_run_environment(config, env=env) + + assert env["NMP_BASE_URL"] == "unix:///tmp/nemo.sock" + + +@pytest.mark.integration +def test_apply_run_environment_preserves_external_host_and_port() -> None: + """Pre-set NMP_SERVICE_HOST and NMP_SERVICE_PORT should not be overwritten.""" + env: dict[str, str] = { + "NMP_SERVICE_HOST": "10.0.0.1", + "NMP_SERVICE_PORT": "443", + } + config = _resolved(host="0.0.0.0", port=9090) + + apply_run_environment(config, env=env) + + assert env["NMP_SERVICE_HOST"] == "10.0.0.1" + assert env["NMP_SERVICE_PORT"] == "443" + + +@pytest.mark.integration +def test_apply_run_environment_ipv6_wildcard_becomes_loopback(tmp_path: Path) -> None: + """The IPv6 wildcard ``::`` should resolve to ``::1`` for internal clients.""" + # Use a config file without platform.base_url so the bind host drives the URL. + config_file = tmp_path / "config.yaml" + config_file.write_text("platform:\n seed_on_startup: false\n") + + env: dict[str, str] = {} + config = _resolved(host="::", port=8080, config_path=str(config_file)) + + apply_run_environment(config, env=env) + + assert env["NMP_SERVICE_HOST"] == "::1" + # Base URL should have bracketed IPv6. + assert env["NMP_BASE_URL"] == "http://[::1]:8080" diff --git a/packages/nemo_platform_ext/tests/local/test_daemon_lifecycle.py b/packages/nemo_platform_ext/tests/local/test_daemon_lifecycle.py new file mode 100644 index 0000000000..e1a3c5c7e1 --- /dev/null +++ b/packages/nemo_platform_ext/tests/local/test_daemon_lifecycle.py @@ -0,0 +1,503 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Integration tests for daemon subprocess lifecycle. + +These tests spawn REAL child processes via ``daemonize_services()``, exercise +real lock acquisition, descriptor file I/O, HTTP readiness probing, and +graceful shutdown via ``stop_instance()``. Nothing is monkeypatched away — +the child runs a real uvicorn server with the ``hello-world`` service. + +Requirements: +- All packages installed (``uv sync --all-packages``) so entry-point + discovery finds hello-world. +- ``pyleak`` importable (from the ``[all]`` extra). +""" + +from __future__ import annotations + +import os +import socket +import subprocess +import sys +import tempfile +import time +from pathlib import Path + +import psutil +import pytest +from nemo_platform_ext.local import process, services +from nemo_platform_ext.local.process import ForegroundInstanceError +from nemo_platform_ext.local.services import ( + ServiceRunConfig, + ServicesAlreadyRunningError, + ServicesStartupExitedError, +) +from nmp.platform_runner.config import PlatformAppConfig + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _free_tcp_port() -> int: + """Bind to port 0, let the OS pick, then release and return the port number.""" + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +def _daemon_config( + tmp_path: Path, + *, + scope: str = "integ-daemon", + port: int | None = None, +) -> ServiceRunConfig: + """Build a ServiceRunConfig that is fully isolated under ``tmp_path``.""" + return ServiceRunConfig( + mode=services.ServiceMode.DAEMON, + services=("hello-world",), + controllers=(), + sidecars=(), + transport="tcp", + host="127.0.0.1", + port=port or _free_tcp_port(), + scope=scope, + state_dir=tmp_path / "state", + runtime_dir=tmp_path / "runtime", + readiness_timeout=30.0, + readiness_poll_interval=0.3, + ) + + +def _ensure_stopped(cfg: ServiceRunConfig) -> None: + """Best-effort cleanup: stop any instance left running by a test.""" + try: + process.stop_instance(cfg.scope, base_dir=cfg.state_root, timeout=10, force=True) + except Exception: + pass + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +@pytest.mark.integration +def test_daemonize_services_spawns_child_that_becomes_ready(tmp_path: Path) -> None: + """Spawn a real daemon subprocess, verify readiness via HTTP, then + gracefully shut down with ``stop_instance``.""" + cfg = _daemon_config(tmp_path) + handle = None + try: + handle = services.daemonize_services(cfg) + + # -- The handle should report the child's PID and transport details. + assert handle.pid is not None + assert handle.port == cfg.port + assert handle.transport == "tcp" + + # -- The lock file should be held by the child. + assert process.is_instance_alive(cfg.scope, base_dir=cfg.state_root) + + # -- The descriptor should have been written by the child. + desc = process.read_descriptor(cfg.scope, base_dir=cfg.state_root) + assert desc is not None + assert desc.pid == handle.pid + assert desc.mode == "daemon" + assert "hello-world" in (desc.config.services or []) + + # -- The child should still be running and respond to /status. + assert services.probe_status(base_url=f"http://127.0.0.1:{cfg.port}", timeout=5.0) + + # -- Graceful shutdown. + result = process.stop_instance(cfg.scope, base_dir=cfg.state_root, timeout=15) + assert handle.pid in result.stopped_pids + + # -- After stop, the lock should be released and the descriptor removed. + assert not process.is_instance_alive(cfg.scope, base_dir=cfg.state_root) + assert process.read_descriptor(cfg.scope, base_dir=cfg.state_root) is None + finally: + _ensure_stopped(cfg) + + +@pytest.mark.integration +def test_daemonize_services_child_exit_before_readiness(tmp_path: Path) -> None: + """When the child exits before becoming ready, ``daemonize_services`` + should raise ``ServicesStartupExitedError`` with the log path.""" + # Spawn a child that will exit immediately: give it a bogus service name + # that will fail validation in resolve_run_configuration. + bad_cfg = ServiceRunConfig( + mode=services.ServiceMode.DAEMON, + services=("nonexistent-service-xyz",), + controllers=(), + sidecars=(), + transport="tcp", + host="127.0.0.1", + port=_free_tcp_port(), + scope="integ-early-exit", + state_dir=tmp_path / "state", + runtime_dir=tmp_path / "runtime", + readiness_timeout=15.0, + readiness_poll_interval=0.2, + ) + with pytest.raises(ServicesStartupExitedError, match="exited with code"): + services.daemonize_services(bad_cfg) + + # -- The lock should not be held after the failed startup. + assert not process.is_instance_alive(bad_cfg.scope, base_dir=bad_cfg.state_root) + + +@pytest.mark.integration +def test_stale_socket_cleanup_after_process_crash(tmp_path: Path) -> None: + """If a previous daemon crashed and left a UDS socket file, a new daemon + startup should clean it up and succeed.""" + scope = "stale" + # Use a short temp directory to stay within AF_UNIX path limits (103 bytes on macOS). + short_tmp = Path(tempfile.mkdtemp(prefix="nemo-")) + runtime_dir = short_tmp / "run" + + # Create a stale UDS socket file (no process listening). + socket_dir = runtime_dir / scope + socket_dir.mkdir(parents=True, exist_ok=True) + stale_socket = socket_dir / "nemo-platform.sock" + # Bind a real UDS socket to create the file, then close immediately. + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as s: + s.bind(str(stale_socket)) + assert stale_socket.exists() + + cfg = ServiceRunConfig( + mode=services.ServiceMode.DAEMON, + services=("hello-world",), + controllers=(), + sidecars=(), + transport="uds", + host="127.0.0.1", + port=_free_tcp_port(), + scope=scope, + state_dir=short_tmp / "state", + runtime_dir=runtime_dir, + readiness_timeout=30.0, + readiness_poll_interval=0.3, + ) + try: + services.daemonize_services(cfg) + + # -- The daemon should be ready. + assert process.is_instance_alive(cfg.scope, base_dir=cfg.state_root) + + # -- The stale socket should have been replaced with the new one. + assert stale_socket.exists() + finally: + _ensure_stopped(cfg) + import shutil + + shutil.rmtree(short_tmp, ignore_errors=True) + + +@pytest.mark.integration +def test_concurrent_daemonize_rejects_duplicate_instance(tmp_path: Path) -> None: + """Starting a second daemon with the same instance scope should fail + with ``ServicesAlreadyRunningError`` while the first is running.""" + cfg = _daemon_config(tmp_path, scope="integ-dup") + try: + services.daemonize_services(cfg) + assert process.is_instance_alive(cfg.scope, base_dir=cfg.state_root) + + # -- A second daemonize with the same scope should fail. + dup_cfg = ServiceRunConfig( + mode=services.ServiceMode.DAEMON, + services=("hello-world",), + controllers=(), + sidecars=(), + transport="tcp", + host="127.0.0.1", + port=_free_tcp_port(), # Different port, same scope. + scope="integ-dup", + state_dir=tmp_path / "state", + runtime_dir=tmp_path / "runtime", + readiness_timeout=5.0, + readiness_poll_interval=0.2, + ) + with pytest.raises(ServicesAlreadyRunningError): + services.daemonize_services(dup_cfg) + + # -- Original instance should still be alive. + assert process.is_instance_alive(cfg.scope, base_dir=cfg.state_root) + finally: + _ensure_stopped(cfg) + + +@pytest.mark.integration +def test_stop_instance_escalates_sigterm_to_sigkill(tmp_path: Path) -> None: + """If the daemon child ignores SIGTERM, ``stop_instance`` should escalate + to SIGKILL after the timeout and successfully terminate the process.""" + # Instead of using daemonize_services (which starts a uvicorn server that + # handles SIGTERM), we manually simulate a daemon process that ignores SIGTERM + # using the process module primitives directly. + scope = "integ-sigkill" + state_dir = tmp_path / "state" + + # Spawn a child process that ignores SIGTERM. + child = subprocess.Popen( + [ + sys.executable, + "-c", + "import signal, time; signal.signal(signal.SIGTERM, signal.SIG_IGN); " + "open('/dev/null', 'w'); time.sleep(300)", + ], + start_new_session=True, + ) + try: + # Write a descriptor so stop_instance can find the process. + desc = process.InstanceDescriptor( + pid=child.pid, + config=PlatformAppConfig(scope=scope, host="127.0.0.1", port=0, state_root=state_dir), + transport="tcp", + mode="daemon", + create_time=psutil.Process(child.pid).create_time(), + ) + process.write_descriptor(desc, base_dir=state_dir) + + # Also create a lock file the process "holds" — but since it's a + # different process, we simulate by NOT acquiring a real flock (the + # test exercises PID-based stop, not flock-based liveness). + + # Stop with a very short timeout so it escalates quickly. + result = process.stop_instance(scope, base_dir=state_dir, timeout=1.0, force=True) + assert child.pid in result.stopped_pids + + # The child should be dead now. + child.wait(timeout=5) + assert child.returncode is not None + finally: + try: + child.kill() + child.wait(timeout=3) + except Exception: + pass + + +@pytest.mark.integration +def test_daemonize_services_cleans_up_on_child_exception(tmp_path: Path) -> None: + """When the child process crashes during init (e.g. corrupted request JSON), + the parent detects the exit, raises, and the lock is not left held.""" + scope = "integ-crash" + state_dir = tmp_path / "state" + instance_dir = process.instance_dir(scope, base_dir=state_dir) + + # Write a corrupted request file that will make _service_child crash + # during JSON deserialization. + fd, tmp_req = tempfile.mkstemp(dir=str(instance_dir), suffix=".json") + os.write(fd, b"NOT VALID JSON {{{") + os.close(fd) + + log_path = process.log_path_for(scope, base_dir=state_dir) + log_file = open(log_path, "a") # noqa: SIM115 + child_module = "nemo_platform_ext.local._service_child" + proc = subprocess.Popen( + [sys.executable, "-m", child_module, tmp_req], + stdout=log_file, + stderr=subprocess.STDOUT, + start_new_session=True, + stdin=subprocess.DEVNULL, + close_fds=True, + ) + log_file.close() + + # Wait for the child to exit (it should crash quickly on bad JSON). + proc.wait(timeout=10) + assert proc.returncode != 0 + + # The lock should not be held — the child never acquired it. + assert not process.is_instance_alive(scope, base_dir=state_dir) + + # The request file should have been cleaned up by _service_child. + assert not Path(tmp_req).exists() + + +# --------------------------------------------------------------------------- +# Priority 2: Process Lifecycle & Cleanup +# --------------------------------------------------------------------------- + + +@pytest.mark.integration +def test_stop_instance_sweeps_orphaned_children(tmp_path: Path) -> None: + """When a daemon parent is stopped, any grandchild processes that survive + should be swept by ``_sweep_orphans``.""" + scope = "integ-orphans" + state_dir = tmp_path / "state" + + # Spawn a parent that spawns a long-lived grandchild, then sleeps. + parent = subprocess.Popen( + [ + sys.executable, + "-c", + "import subprocess, sys, time; " + "gc = subprocess.Popen([sys.executable, '-c', 'import time; time.sleep(300)']); " + "time.sleep(300)", + ], + start_new_session=True, + ) + try: + # Give the parent time to spawn the grandchild. + time.sleep(0.5) + grandchildren = psutil.Process(parent.pid).children(recursive=True) + assert len(grandchildren) >= 1, "grandchild was not spawned" + + desc = process.InstanceDescriptor( + pid=parent.pid, + config=PlatformAppConfig(scope=scope, host="127.0.0.1", port=0, state_root=state_dir), + transport="tcp", + mode="daemon", + create_time=psutil.Process(parent.pid).create_time(), + ) + process.write_descriptor(desc, base_dir=state_dir) + + result = process.stop_instance(scope, base_dir=state_dir, timeout=10, force=True) + assert parent.pid in result.stopped_pids + assert len(result.swept_children) >= 1 + + # Both parent and grandchild should be dead. + parent.wait(timeout=5) + for gc in grandchildren: + gc.wait(timeout=5) + finally: + try: + parent.kill() + parent.wait(timeout=3) + except Exception: + pass + for gc in grandchildren: + try: + gc.kill() + gc.wait(timeout=3) + except Exception: + pass + + +@pytest.mark.integration +def test_stop_instance_foreground_mode_requires_force(tmp_path: Path) -> None: + """Stopping a foreground-mode instance without ``force=True`` should raise + ``ForegroundInstanceError``. With ``force=True`` it should proceed.""" + scope = "integ-foreground" + state_dir = tmp_path / "state" + + child = subprocess.Popen( + [sys.executable, "-c", "import time; time.sleep(300)"], + start_new_session=True, + ) + try: + desc = process.InstanceDescriptor( + pid=child.pid, + config=PlatformAppConfig(scope=scope, host="127.0.0.1", port=0, state_root=state_dir), + transport="tcp", + mode="foreground", + create_time=psutil.Process(child.pid).create_time(), + ) + process.write_descriptor(desc, base_dir=state_dir) + + # Without force, should raise. + with pytest.raises(ForegroundInstanceError): + process.stop_instance(scope, base_dir=state_dir, timeout=5) + + # Process should still be alive after the rejected stop. + assert child.poll() is None + + # With force, should succeed. + result = process.stop_instance(scope, base_dir=state_dir, timeout=5, force=True) + assert child.pid in result.stopped_pids + child.wait(timeout=5) + finally: + try: + child.kill() + child.wait(timeout=3) + except Exception: + pass + + +@pytest.mark.integration +def test_is_instance_alive_with_stale_lock(tmp_path: Path) -> None: + """If the lock file exists but no process holds the flock, + ``is_instance_alive`` should return False.""" + scope = "integ-stale-lock" + state_dir = tmp_path / "state" + + # Create the lock file without holding a flock on it. + inst_dir = process.instance_dir(scope, base_dir=state_dir) + lock_path = inst_dir / process.LOCK_FILENAME + lock_path.touch() + + assert not process.is_instance_alive(scope, base_dir=state_dir) + + +@pytest.mark.integration +def test_is_instance_alive_with_held_lock(tmp_path: Path) -> None: + """If a process holds the flock, ``is_instance_alive`` should return True.""" + scope = "integ-held-lock" + state_dir = tmp_path / "state" + + fd = process.acquire_lock(scope, base_dir=state_dir) + try: + assert process.is_instance_alive(scope, base_dir=state_dir) + finally: + os.close(fd) + + # After releasing the fd (which releases the flock), should be false. + assert not process.is_instance_alive(scope, base_dir=state_dir) + + +@pytest.mark.integration +def test_validate_pid_detects_recycled_process(tmp_path: Path) -> None: + """After a process dies, ``validate_pid`` should return False if the PID is + reused by a different process (detected via create_time mismatch).""" + # Spawn and immediately kill a short-lived process to get a PID + create_time. + child = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(300)"]) + pid = child.pid + create_time = psutil.Process(pid).create_time() + + # The PID is alive and create_time matches. + assert process.validate_pid(pid, create_time) + + # Kill it. + child.kill() + child.wait(timeout=5) + + # Now validate_pid should return False — the process is dead. + assert not process.validate_pid(pid, create_time) + + # Even with a wildly wrong create_time, should be False for a dead PID. + assert not process.validate_pid(pid, 0.0) + + +@pytest.mark.integration +def test_rotate_log_preserves_existing_content(tmp_path: Path) -> None: + """``rotate_log`` should rename the existing log and return the path for + the new (empty) log. The old content must be preserved.""" + scope = "integ-rotate" + state_dir = tmp_path / "state" + + # Write initial log content. + log_path = process.log_path_for(scope, base_dir=state_dir) + log_path.parent.mkdir(parents=True, exist_ok=True) + log_path.write_text("original log content\n") + + # Rotate. + new_log = process.rotate_log(scope, base_dir=state_dir) + assert new_log == log_path + assert not log_path.exists() # Original was renamed. + + # Find the rotated file. + rotated_files = [f for f in log_path.parent.iterdir() if f.name.startswith("services.log.")] + assert len(rotated_files) == 1 + assert rotated_files[0].read_text() == "original log content\n" + + # Write new content, rotate again. + log_path.write_text("second run\n") + process.rotate_log(scope, base_dir=state_dir) + + rotated_files = sorted(f for f in log_path.parent.iterdir() if f.name.startswith("services.log.")) + assert len(rotated_files) == 2 + contents = {f.read_text() for f in rotated_files} + assert "original log content\n" in contents + assert "second run\n" in contents diff --git a/packages/nemo_platform_ext/tests/local/test_health_child.py b/packages/nemo_platform_ext/tests/local/test_health_child.py new file mode 100644 index 0000000000..04a975d2f0 --- /dev/null +++ b/packages/nemo_platform_ext/tests/local/test_health_child.py @@ -0,0 +1,340 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Integration tests for health/readiness probing, lifespan, and child process module. + +Covers Priorities 5 (lifespan), 6 (health), and 7 (child process) from the +integration test plan. +""" + +from __future__ import annotations + +import json +import os +import socket +import subprocess +import sys +import tempfile +import threading +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest +from nemo_platform_ext.local import process +from nemo_platform_ext.local.services import ServiceRunConfig +from nemo_platform_ext.local.transport import probe_status, wait_for_status + + +def _free_tcp_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +# --------------------------------------------------------------------------- +# Priority 5: Multi-Service Startup & Lifespan +# --------------------------------------------------------------------------- + + +@pytest.mark.integration +def test_create_app_starts_and_joins_controller_threads() -> None: + """A controller registered via ``create_app`` should have its thread + started during lifespan and stopped on exit.""" + started = threading.Event() + stopped = threading.Event() + + def controller_run(stop_signal: threading.Event) -> None: + started.set() + stop_signal.wait(timeout=5.0) + stopped.set() + + with ( + patch("nmp.platform_runner.server.get_platform_config") as mock_pc, + patch("nmp.platform_runner.server.get_auth_config") as mock_ac, + patch("nmp.common.auth.middleware.get_auth_config") as mock_ac2, + ): + mock_pc.return_value = MagicMock(seed_on_startup=False, redirect_root_to_studio=False) + mock_ac.return_value = MagicMock(enabled=False, policy_decision_point_provider="embedded") + mock_ac2.return_value = MagicMock(enabled=False, policy_decision_point_provider="embedded") + + from nmp.platform_runner.server import create_app + + app = create_app(services=[], controller_run_funcs={"test-ctrl": controller_run}) + + from fastapi.testclient import TestClient + + with TestClient(app): + assert started.wait(timeout=2.0), "controller thread did not start" + + assert stopped.wait(timeout=2.0), "controller thread did not stop after lifespan exit" + + +@pytest.mark.integration +def test_create_app_controller_thread_join_timeout() -> None: + """A controller that ignores the stop signal should not hang shutdown — + ``thread.join(timeout=5)`` should return even if the controller is still running.""" + started = threading.Event() + + def stubborn_controller(stop_signal: threading.Event) -> None: + started.set() + # Ignore stop_signal — simulate a controller that hangs. + import time + + time.sleep(300) + + with ( + patch("nmp.platform_runner.server.get_platform_config") as mock_pc, + patch("nmp.platform_runner.server.get_auth_config") as mock_ac, + patch("nmp.common.auth.middleware.get_auth_config") as mock_ac2, + ): + mock_pc.return_value = MagicMock(seed_on_startup=False, redirect_root_to_studio=False) + mock_ac.return_value = MagicMock(enabled=False, policy_decision_point_provider="embedded") + mock_ac2.return_value = MagicMock(enabled=False, policy_decision_point_provider="embedded") + + from nmp.platform_runner.server import create_app + + app = create_app(services=[], controller_run_funcs={"stubborn": stubborn_controller}) + + from fastapi.testclient import TestClient + + # The TestClient __exit__ triggers lifespan exit, which calls thread.join(timeout=5). + # This should NOT hang forever — the 5s timeout should let shutdown proceed. + with TestClient(app): + assert started.wait(timeout=2.0), "controller thread did not start" + + # If we got here, shutdown didn't hang. The stubborn thread is still running + # but as a daemon thread it will be cleaned up when the test process exits. + + +@pytest.mark.integration +def test_lifespan_cleanup_runs_on_app_shutdown() -> None: + """``close_shared_http_clients`` should be called during lifespan teardown.""" + cleanup_called = threading.Event() + + with ( + patch("nmp.platform_runner.server.get_platform_config") as mock_pc, + patch("nmp.platform_runner.server.get_auth_config") as mock_ac, + patch("nmp.common.auth.middleware.get_auth_config") as mock_ac2, + patch("nmp.platform_runner.server.close_shared_http_clients") as mock_close, + ): + mock_pc.return_value = MagicMock(seed_on_startup=False, redirect_root_to_studio=False) + mock_ac.return_value = MagicMock(enabled=False, policy_decision_point_provider="embedded") + mock_ac2.return_value = MagicMock(enabled=False, policy_decision_point_provider="embedded") + + async def fake_close(): + cleanup_called.set() + + mock_close.side_effect = fake_close + + from nmp.platform_runner.server import create_app + + app = create_app(services=[]) + + from fastapi.testclient import TestClient + + with TestClient(app): + pass + + assert cleanup_called.is_set(), "close_shared_http_clients was not called during shutdown" + + +# --------------------------------------------------------------------------- +# Priority 6: Health & Readiness +# --------------------------------------------------------------------------- + + +@pytest.mark.integration +def test_wait_for_status_retries_on_transient_errors(tmp_path: Path) -> None: + """``wait_for_status`` should retry on connection refused and eventually + return True once the server starts responding.""" + from nemo_platform_ext.local import services + + cfg = ServiceRunConfig( + mode=services.ServiceMode.DAEMON, + services=("hello-world",), + controllers=(), + sidecars=(), + transport="tcp", + host="127.0.0.1", + port=_free_tcp_port(), + scope="integ-wait-retry", + state_dir=tmp_path / "state", + runtime_dir=tmp_path / "runtime", + readiness_timeout=30.0, + readiness_poll_interval=0.3, + ) + + # Start the daemon — wait_for_status should retry until it's ready. + services.daemonize_services(cfg) + try: + # The daemon is already ready (daemonize_services waits for readiness). + # Verify wait_for_status succeeds with a fresh probe. + assert wait_for_status( + base_url=f"http://127.0.0.1:{cfg.port}", + timeout=5.0, + poll_interval=0.2, + ) + finally: + process.stop_instance(cfg.scope, base_dir=cfg.state_root, timeout=10, force=True) + + +@pytest.mark.integration +def test_wait_for_status_times_out_on_no_server() -> None: + """``wait_for_status`` should return False when no server is listening.""" + port = _free_tcp_port() + result = wait_for_status( + base_url=f"http://127.0.0.1:{port}", + timeout=1.0, + poll_interval=0.2, + ) + assert result is False + + +@pytest.mark.integration +def test_probe_status_with_missing_uds_socket() -> None: + """Probing a non-existent UDS socket should return False.""" + result = probe_status( + base_url="http+unix:///nonexistent/path/nemo.sock", + socket_path=Path("/nonexistent/path/nemo.sock"), + timeout=1.0, + ) + assert result is False + + +@pytest.mark.integration +def test_probe_status_against_real_daemon(tmp_path: Path) -> None: + """``probe_status`` should return True against a running daemon.""" + from nemo_platform_ext.local import services + + cfg = ServiceRunConfig( + mode=services.ServiceMode.DAEMON, + services=("hello-world",), + controllers=(), + sidecars=(), + transport="tcp", + host="127.0.0.1", + port=_free_tcp_port(), + scope="integ-probe-real", + state_dir=tmp_path / "state", + runtime_dir=tmp_path / "runtime", + readiness_timeout=30.0, + readiness_poll_interval=0.3, + ) + services.daemonize_services(cfg) + try: + assert probe_status(base_url=f"http://127.0.0.1:{cfg.port}", timeout=5.0) + finally: + process.stop_instance(cfg.scope, base_dir=cfg.state_root, timeout=10, force=True) + + +# --------------------------------------------------------------------------- +# Priority 7: Child Process Module +# --------------------------------------------------------------------------- + + +@pytest.mark.integration +def test_service_child_loads_config_and_starts(tmp_path: Path) -> None: + """Write valid JSON config, run ``_service_child`` in a subprocess, + verify it starts and accepts HTTP connections.""" + port = _free_tcp_port() + state_dir = tmp_path / "state" + runtime_dir = tmp_path / "runtime" + scope = "integ-child-real" + + payload = ServiceRunConfig( + mode="daemon", + services=("hello-world",), + controllers=(), + sidecars=(), + transport="tcp", + host="127.0.0.1", + port=port, + scope=scope, + state_dir=str(state_dir), + runtime_dir=str(runtime_dir), + ).to_child_payload() + + # Write the request file the way daemonize_services does. + instance_dir = process.instance_dir(scope, base_dir=state_dir) + fd, req_path = tempfile.mkstemp(dir=str(instance_dir), suffix=".json") + os.write(fd, (json.dumps(payload) + "\n").encode()) + os.close(fd) + + log_path = process.log_path_for(scope, base_dir=state_dir) + log_file = open(log_path, "a") # noqa: SIM115 + subprocess.Popen( + [sys.executable, "-m", "nemo_platform_ext.local._service_child", req_path], + stdout=log_file, + stderr=subprocess.STDOUT, + start_new_session=True, + stdin=subprocess.DEVNULL, + close_fds=True, + ) + log_file.close() + + try: + # Wait for the child to become ready. + assert wait_for_status( + base_url=f"http://127.0.0.1:{port}", + timeout=30.0, + poll_interval=0.3, + ), "child process did not become ready" + + # The request file should have been cleaned up. + assert not Path(req_path).exists() + + # The child should have acquired the lock and written a descriptor. + assert process.is_instance_alive(scope, base_dir=state_dir) + desc = process.read_descriptor(scope, base_dir=state_dir) + assert desc is not None + assert desc.mode == "daemon" + finally: + process.stop_instance(scope, base_dir=state_dir, timeout=10, force=True) + + +@pytest.mark.integration +def test_service_child_corrupted_payload(tmp_path: Path) -> None: + """Bad JSON in the request file should cause the child to exit non-zero.""" + scope = "integ-child-bad" + state_dir = tmp_path / "state" + instance_dir = process.instance_dir(scope, base_dir=state_dir) + + fd, req_path = tempfile.mkstemp(dir=str(instance_dir), suffix=".json") + os.write(fd, b"<<>>") + os.close(fd) + + proc = subprocess.Popen( + [sys.executable, "-m", "nemo_platform_ext.local._service_child", req_path], + stdout=subprocess.DEVNULL, + stderr=subprocess.STDOUT, + start_new_session=True, + stdin=subprocess.DEVNULL, + ) + proc.wait(timeout=15) + assert proc.returncode != 0 + + +@pytest.mark.integration +def test_service_child_cleans_up_request_file(tmp_path: Path) -> None: + """The request file should be unlinked even when the child crashes.""" + scope = "integ-child-cleanup" + state_dir = tmp_path / "state" + instance_dir = process.instance_dir(scope, base_dir=state_dir) + + fd, req_path = tempfile.mkstemp(dir=str(instance_dir), suffix=".json") + os.write(fd, b"<<>>") + os.close(fd) + + assert Path(req_path).exists() + + proc = subprocess.Popen( + [sys.executable, "-m", "nemo_platform_ext.local._service_child", req_path], + stdout=subprocess.DEVNULL, + stderr=subprocess.STDOUT, + start_new_session=True, + stdin=subprocess.DEVNULL, + ) + proc.wait(timeout=15) + + # The request file should have been cleaned up regardless of the error. + assert not Path(req_path).exists() diff --git a/packages/nemo_platform_ext/tests/local/test_port_socket.py b/packages/nemo_platform_ext/tests/local/test_port_socket.py new file mode 100644 index 0000000000..59804c6615 --- /dev/null +++ b/packages/nemo_platform_ext/tests/local/test_port_socket.py @@ -0,0 +1,165 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Integration tests for TCP/UDS port and socket management. + +These tests exercise real port binding, socket creation, and conflict +detection using actual OS resources. +""" + +from __future__ import annotations + +import socket +import tempfile +from pathlib import Path + +import pytest +from nemo_platform_ext.local import process, services +from nemo_platform_ext.local.services import ( + ServiceRunConfig, + ServicesPortInUseError, +) +from nmp.platform_runner.config import PlatformAppConfig + + +def _free_tcp_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +@pytest.mark.integration +def test_tcp_port_conflict_with_foreign_process(tmp_path: Path) -> None: + """When a foreign (non-NeMo) process holds a port, ``_check_tcp_available`` + should raise ``ServicesPortInUseError`` with a helpful suggestion.""" + # Bind a TCP port and hold it open. + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as blocker: + blocker.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + blocker.bind(("127.0.0.1", 0)) + blocker.listen(1) + port = blocker.getsockname()[1] + + cfg = ServiceRunConfig( + mode=services.ServiceMode.DAEMON, + services=("hello-world",), + transport="tcp", + host="127.0.0.1", + port=port, + scope="integ-port-foreign", + state_dir=tmp_path / "state", + runtime_dir=tmp_path / "runtime", + ) + + with pytest.raises(ServicesPortInUseError, match="already in use by another process"): + services._check_tcp_available(cfg) + + +@pytest.mark.integration +def test_tcp_port_conflict_with_nemo_instance(tmp_path: Path) -> None: + """When a NeMo instance holds a port, the error should distinguish it + from a foreign process.""" + scope = "integ-port-nemo" + state_dir = tmp_path / "state" + + # Bind a port and also create a descriptor matching the scope/host/port. + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as blocker: + blocker.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + blocker.bind(("127.0.0.1", 0)) + blocker.listen(1) + port = blocker.getsockname()[1] + + # Create a live lock and descriptor so it looks like a NeMo instance. + lock_fd = process.acquire_lock(scope, base_dir=state_dir) + try: + desc = process.InstanceDescriptor( + pid=1, # Dummy PID — the flock is what matters. + config=PlatformAppConfig(scope=scope, host="127.0.0.1", port=port), + transport="tcp", + mode="daemon", + create_time=0.0, + ) + process.write_descriptor(desc, base_dir=state_dir) + + conflict = process.check_port_available_for_start("127.0.0.1", port, scope, base_dir=state_dir) + assert conflict is not None + assert conflict.kind == "nemo_instance" + assert conflict.port == port + + lines = process.format_port_conflict(conflict) + assert any("NeMo Platform" in line for line in lines) + finally: + import os + + os.close(lock_fd) + + +@pytest.mark.integration +def test_tcp_port_available_when_free(tmp_path: Path) -> None: + """When a port is free, ``check_port_available_for_start`` returns None.""" + port = _free_tcp_port() + conflict = process.check_port_available_for_start("127.0.0.1", port, "integ-free", base_dir=tmp_path / "state") + assert conflict is None + + +@pytest.mark.integration +def test_uds_socket_path_max_validation() -> None: + """A socket path exceeding AF_UNIX_PATH_MAX should raise ValueError.""" + # Build a path that is exactly one byte over the limit. + max_bytes = services._AF_UNIX_PATH_MAX_BYTES + # Create a path that exceeds the limit. + long_path = "/" + "x" * max_bytes # len("/") + max_bytes > max_bytes + assert len(long_path.encode()) > max_bytes + + cfg = ServiceRunConfig( + mode=services.ServiceMode.DAEMON, + services=("hello-world",), + transport="uds", + socket_path=long_path, + scope="integ-long-sock", + state_dir="/tmp/state", + runtime_dir="/tmp/run", + ) + # _validated_socket_path calls _validate_socket_path_length internally. + with pytest.raises(ValueError, match="too long for AF_UNIX"): + services._validated_socket_path(cfg) + + +@pytest.mark.integration +def test_prepare_socket_removes_stale_socket(tmp_path: Path) -> None: + """``_prepare_socket`` should remove a stale (unreachable) socket file + and allow a new daemon to bind.""" + scope = "stale2" + short_tmp = Path(tempfile.mkdtemp(prefix="nemo-")) + runtime_dir = short_tmp / "run" + + # Create a stale socket file. + socket_dir = runtime_dir / scope + socket_dir.mkdir(parents=True, exist_ok=True) + stale_socket = socket_dir / "nemo-platform.sock" + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as s: + s.bind(str(stale_socket)) + assert stale_socket.exists() + + cfg = ServiceRunConfig( + mode=services.ServiceMode.DAEMON, + services=("hello-world",), + transport="uds", + scope=scope, + state_dir=short_tmp / "state", + runtime_dir=runtime_dir, + ) + + # _prepare_socket should probe, find it stale, remove it, and return the path. + result = services._prepare_socket(cfg) + assert result is not None + # The stale socket should have been removed (the new server hasn't bound yet). + assert not stale_socket.exists() + + import shutil + + shutil.rmtree(short_tmp, ignore_errors=True) diff --git a/packages/nemo_platform_ext/tests/local/test_services.py b/packages/nemo_platform_ext/tests/local/test_services.py new file mode 100644 index 0000000000..1964276ba0 --- /dev/null +++ b/packages/nemo_platform_ext/tests/local/test_services.py @@ -0,0 +1,1039 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +import os +import socket +import subprocess +import sys +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from nemo_platform_ext.local import _service_child, services +from nemo_platform_ext.local.process import ( + DESCRIPTOR_FILENAME, + InstanceDescriptor, +) +from nemo_platform_ext.local.services import ServiceRunConfig +from nemo_platform_ext.local.transport import UDS_BASE_URL +from nmp.platform_runner.config import ( + PlatformAppConfig, + default_runtime_root, + default_state_root, + validate_scope, +) + + +def _allow_tmp_path_socket_paths(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(services, "_AF_UNIX_PATH_MAX_BYTES", 4096) + + +def _free_tcp_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + return sock.getsockname()[1] + + +def _embedded_handle() -> services.EmbeddedServiceHandle: + return services.EmbeddedServiceHandle(app=object(), runtime=object()) + + +def test_service_run_config_normalizes_lists_to_tuples() -> None: + cfg = ServiceRunConfig(services=["entities", "models"], controllers=["jobs"]) + + assert cfg.services == ("entities", "models") + assert cfg.controllers == ("jobs",) + + +def test_service_run_config_converts_to_platform_app_config(tmp_path: Path) -> None: + cfg = ServiceRunConfig( + services=["entities", "models"], + controllers=[], + sidecars=["adapters"], + config_path=tmp_path / "local.yaml", + socket_path=tmp_path / "nemo.sock", + mode="embedded", + ) + + app_config = cfg.to_platform_app_config() + + assert app_config.services == ("entities", "models") + assert app_config.controllers == () + assert app_config.sidecars == ("adapters",) + assert app_config.config_path == str(tmp_path / "local.yaml") + assert app_config.socket_path == str(tmp_path / "nemo.sock") + assert app_config.runtime_root is None + assert app_config.runtime_dir() == tmp_path + assert app_config.host == "127.0.0.1" + assert app_config.port == 8080 + + +def test_instance_descriptor_converts_from_service_run_config( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + cfg = ServiceRunConfig( + services=["entities", "models"], + controllers=[], + sidecars=["adapters"], + config_path=tmp_path / "local.yaml", + socket_path=tmp_path / "nemo.sock", + state_dir=tmp_path / "state", + runtime_dir=tmp_path / "run", + ) + monkeypatch.setattr(services.process, "get_create_time", lambda _pid: 123.0) + app_config = cfg.to_platform_app_config() + app_config.log_path = str(tmp_path / "nemo.log") + + desc = InstanceDescriptor.from_config( + app_config, + pid=4242, + mode="daemon", + transport=cfg.transport, + ) + + assert desc.pid == 4242 + assert desc.config.scope == "default" + assert desc.config.host == "127.0.0.1" + assert desc.config.port == 8080 + assert desc.transport == "uds" + assert desc.config.socket_path == str(tmp_path / "nemo.sock") + assert desc.config.state_root == str(tmp_path / "state") + assert desc.config.runtime_root == str(tmp_path / "run") + assert desc.config.state_dir() == tmp_path / "state" / "instances" / "default" + assert desc.config.runtime_dir() == tmp_path / "run" / "default" + assert desc.mode == "daemon" + assert desc.create_time == 123.0 + assert desc.config.services == ("entities", "models") + assert desc.config.controllers == () + assert desc.config.sidecars == ("adapters",) + assert desc.config.config_path == str(tmp_path / "local.yaml") + assert desc.config.log_path == str(tmp_path / "nemo.log") + assert desc.config.log_file_path() == tmp_path / "nemo.log" + payload = desc.model_dump() + assert "services" not in payload + assert "host" not in payload + assert "state_dir" not in payload + assert "runtime_dir" not in payload + assert "log_path" not in payload + assert payload["config"]["services"] == ("entities", "models") + assert payload["config"]["socket_path"] == str(tmp_path / "nemo.sock") + assert payload["config"]["state_root"] == str(tmp_path / "state") + assert payload["config"]["runtime_root"] == str(tmp_path / "run") + assert payload["config"]["log_path"] == str(tmp_path / "nemo.log") + + +def test_service_mode_enum_values() -> None: + assert services.ServiceMode.EMBEDDED.value == "embedded" + assert services.ServiceMode.DAEMON.value == "daemon" + + +def test_service_run_config_defaults_to_daemon_mode() -> None: + cfg = ServiceRunConfig() + + assert cfg.mode is services.ServiceMode.DAEMON + + +def test_service_run_config_accepts_mode_strings() -> None: + cfg = ServiceRunConfig(mode="embedded") + + assert cfg.mode is services.ServiceMode.EMBEDDED + + +def test_service_run_config_rejects_unknown_mode() -> None: + with pytest.raises(ValueError, match="mode must be 'embedded' or 'daemon'"): + ServiceRunConfig(mode="foreground") + + +def test_embedded_and_daemon_handles_implement_local_service_handle(tmp_path: Path) -> None: + embedded = services.EmbeddedServiceHandle(app=object(), runtime=object()) + daemon = services.DaemonServiceHandle( + scope="dev", + transport="tcp", + socket_path=None, + gateway_base_url=None, + host="127.0.0.1", + port=8080, + pid=123, + mode="daemon", + log_path=None, + state_dir=tmp_path / "state" / "instances" / "dev", + runtime_dir=None, + ) + + assert isinstance(embedded, services.LocalServiceHandle) + assert isinstance(daemon, services.LocalServiceHandle) + + +def test_start_services_result_is_shared_result_type() -> None: + result = services.StartServicesResult( + requested=["jobs"], + started=["auth", "jobs"], + already_active=[], + active=["secrets", "auth", "jobs"], + ) + + assert result.requested == ["jobs"] + assert result.started == ["auth", "jobs"] + assert result.active == ["secrets", "auth", "jobs"] + + +def test_service_run_config_rejects_services_with_service_group() -> None: + with pytest.raises(ValueError, match="services cannot be combined with service_group"): + ServiceRunConfig(services=("entities",), service_group="all") + + +def test_service_run_config_defaults_to_named_uds_instance() -> None: + cfg = ServiceRunConfig() + + assert cfg.transport == "uds" + assert cfg.http_gateway == "disabled" + assert cfg.scope == "default" + assert cfg.socket_path is None + + +@pytest.mark.parametrize("instance", ["has space", "../bad"]) +def test_service_run_config_rejects_invalid_scope_names(instance: str) -> None: + with pytest.raises(ValueError, match="scope"): + ServiceRunConfig(scope=instance) + + +def test_service_run_config_rejects_gateway_for_tcp_transport() -> None: + with pytest.raises(ValueError, match="gateway.*UDS"): + ServiceRunConfig(transport="tcp", http_gateway="enabled") + + +@pytest.mark.parametrize( + ("field", "value", "message"), + [ + ("readiness_timeout", 0.0, "readiness_timeout"), + ("readiness_timeout", -1.0, "readiness_timeout"), + ("readiness_poll_interval", 0.0, "readiness_poll_interval"), + ("readiness_poll_interval", -1.0, "readiness_poll_interval"), + ], +) +def test_service_run_config_rejects_non_positive_readiness_values(field: str, value: float, message: str) -> None: + with pytest.raises(ValueError, match=message): + if field == "readiness_timeout": + ServiceRunConfig(readiness_timeout=value) + else: + ServiceRunConfig(readiness_poll_interval=value) + + +def test_process_paths_follow_existing_nmp_state_convention(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + monkeypatch.setenv("XDG_STATE_HOME", str(tmp_path / "state")) + + assert default_state_root() == tmp_path / "state" / "nmp" + assert default_runtime_root() == tmp_path / "state" / "nmp" / "run" + assert ( + PlatformAppConfig(scope="dev").socket_file_path() + == tmp_path / "state" / "nmp" / "run" / "dev" / "nemo-platform.sock" + ) + assert validate_scope("dev_1-2") == "dev_1-2" + + +def test_resolved_socket_path_rejects_relative_explicit_path() -> None: + cfg = ServiceRunConfig(socket_path="relative.sock") + + with pytest.raises(ValueError, match="UDS socket path must be absolute"): + _ = cfg.resolved_socket_path + + +def test_resolved_socket_path_rejects_relative_runtime_dir() -> None: + cfg = ServiceRunConfig(runtime_dir="relative-run") + + with pytest.raises(ValueError, match="runtime root must be absolute"): + _ = cfg.resolved_socket_path + + +def test_resolved_socket_path_rejects_relative_socket_path_with_tcp_client() -> None: + cfg = ServiceRunConfig(transport="tcp", socket_path="relative.sock") + + with pytest.raises(ValueError, match="UDS socket path must be absolute"): + _ = cfg.resolved_socket_path + + +def test_tcp_client_can_still_configure_uds_listener(tmp_path: Path) -> None: + cfg = ServiceRunConfig(transport="tcp", socket_path=tmp_path / "nemo.sock") + + app_config = cfg.to_platform_app_config() + + assert app_config.socket_path == str(tmp_path / "nemo.sock") + assert app_config.runtime_dir() == tmp_path + + +def test_instance_descriptor_rejects_uds_client_without_socket_path() -> None: + with pytest.raises(ValueError, match="UDS client transport requires config.socket_path"): + InstanceDescriptor(pid=1, config=PlatformAppConfig(scope="dev"), transport="uds") + + +def test_prepare_socket_rejects_long_generated_path_before_creating_runtime_dir( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + monkeypatch.setattr(services, "_AF_UNIX_PATH_MAX_BYTES", 1, raising=False) + runtime_root = tmp_path / "runtime" + cfg = ServiceRunConfig(scope="dev", state_dir=tmp_path / "state", runtime_dir=runtime_root) + + with pytest.raises(ValueError, match="UDS socket path is too long.*AF_UNIX"): + services._prepare_socket(cfg) + + assert not runtime_root.exists() + + +def test_validate_socket_path_length_reserves_trailing_nul(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(services, "_AF_UNIX_PATH_MAX_BYTES", 3, raising=False) + + services._validate_socket_path_length(Path("abc")) + with pytest.raises(ValueError, match=r"4 bytes; maximum is 3 bytes"): + services._validate_socket_path_length(Path("abcd")) + + +def test_prepare_socket_rejects_long_explicit_path_before_filesystem_or_probe( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + monkeypatch.setattr(services, "_AF_UNIX_PATH_MAX_BYTES", 1, raising=False) + socket_parent = tmp_path / "explicit" + cfg = ServiceRunConfig(socket_path=socket_parent / "nemo-platform.sock") + + with patch("nemo_platform_ext.local.services.probe_status") as probe_status: + with pytest.raises(ValueError, match="UDS socket path is too long.*AF_UNIX"): + services._prepare_socket(cfg) + + probe_status.assert_not_called() + assert not socket_parent.exists() + + +def test_run_services_prepares_socket_after_acquiring_lock(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + _allow_tmp_path_socket_paths(monkeypatch) + cfg = ServiceRunConfig( + scope="dev", port=_free_tcp_port(), state_dir=tmp_path / "state", runtime_dir=tmp_path / "run" + ) + events: list[str] = [] + real_acquire_lock = services.process.acquire_lock + + def acquire_lock(scope: str, *, base_dir: Path | None = None) -> int: + events.append("lock") + return real_acquire_lock(scope, base_dir=base_dir) + + def prepare_socket(config: ServiceRunConfig) -> Path | None: + events.append("prepare") + lock_path = ( + services.process.instance_dir(config.scope, base_dir=config.state_root) / services.process.LOCK_FILENAME + ) + assert lock_path.exists() + return config.resolved_socket_path + + with ( + patch("nemo_platform_ext.local.services.require_services_extra"), + patch("nemo_platform_ext.local.services.process.acquire_lock", side_effect=acquire_lock), + patch("nemo_platform_ext.local.services._prepare_socket", side_effect=prepare_socket), + patch("nemo_platform_ext.local.services.start_embedded_services", return_value=_embedded_handle()), + patch("nemo_platform_ext.local.services.serve_embedded_app"), + ): + services.run_services(cfg) + + assert events == ["lock", "prepare"] + + +def test_run_services_cleans_lock_when_socket_prepare_fails(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + _allow_tmp_path_socket_paths(monkeypatch) + cfg = ServiceRunConfig( + scope="dev", port=_free_tcp_port(), state_dir=tmp_path / "state", runtime_dir=tmp_path / "run" + ) + + with ( + patch("nemo_platform_ext.local.services.require_services_extra"), + patch( + "nemo_platform_ext.local.services._prepare_socket", + side_effect=services.ServicesSocketStaleError("boom"), + ), + ): + with pytest.raises(services.ServicesSocketStaleError, match="boom"): + services.run_services(cfg) + + assert not services.process.is_instance_alive(cfg.scope, base_dir=cfg.state_root) + + +def test_run_services_restores_env_and_closes_lock_when_descriptor_cleanup_fails( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + _allow_tmp_path_socket_paths(monkeypatch) + cfg = ServiceRunConfig( + scope="dev", + port=_free_tcp_port(), + state_dir=tmp_path / "state", + runtime_dir=tmp_path / "run", + data_dir=tmp_path / "data", + ) + monkeypatch.delenv("NMP_DATA_DIR", raising=False) + real_acquire_lock = services.process.acquire_lock + locked_fd: int | None = None + + def acquire_lock(scope: str, *, base_dir: Path | None = None) -> int: + nonlocal locked_fd + locked_fd = real_acquire_lock(scope, base_dir=base_dir) + return locked_fd + + real_close = os.close + closed_fds: list[int] = [] + + def close(fd: int) -> None: + closed_fds.append(fd) + real_close(fd) + + monkeypatch.setattr(os, "close", close) + with ( + patch("nemo_platform_ext.local.services.require_services_extra"), + patch("nemo_platform_ext.local.services.process.acquire_lock", side_effect=acquire_lock), + patch( + "nemo_platform_ext.local.services.process.remove_descriptor", + side_effect=RuntimeError("descriptor cleanup failed"), + ), + patch("nemo_platform_ext.local.services.start_embedded_services", return_value=_embedded_handle()), + patch("nemo_platform_ext.local.services.serve_embedded_app"), + ): + with pytest.raises(RuntimeError, match="descriptor cleanup failed"): + services.run_services(cfg) + + assert "NMP_DATA_DIR" not in os.environ + assert locked_fd is not None + assert locked_fd in closed_fds + assert closed_fds[-1] == locked_fd + + +def test_run_services_foreground_serves_embedded_app(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + _allow_tmp_path_socket_paths(monkeypatch) + cfg = ServiceRunConfig( + mode=services.ServiceMode.EMBEDDED, + scope="dev", + port=_free_tcp_port(), + state_dir=tmp_path / "state", + runtime_dir=tmp_path / "run", + ) + app = object() + handle = services.EmbeddedServiceHandle(app=app, runtime=object()) + + with ( + patch("nemo_platform_ext.local.services.require_services_extra"), + patch("nemo_platform_ext.local.services.start_embedded_services", return_value=handle) as start_embedded, + patch("nemo_platform_ext.local.services.serve_embedded_app") as serve_embedded, + ): + services.run_services(cfg) + + start_embedded.assert_called_once_with(cfg, env=None) + serve_embedded.assert_called_once() + assert serve_embedded.call_args.args[0] is app + + +def test_daemon_service_handle_uds_client_uses_socket_transport(tmp_path: Path) -> None: + socket_path = tmp_path / "nemo-platform.sock" + handle = services.DaemonServiceHandle( + scope="dev", + transport="uds", + socket_path=socket_path, + gateway_base_url="http://127.0.0.1:9999", + host="127.0.0.1", + port=8080, + pid=123, + mode="daemon", + log_path=None, + state_dir=tmp_path / "state" / "instances" / "dev", + runtime_dir=tmp_path, + ) + + client = handle.client() + try: + assert str(client.base_url).rstrip("/") == UDS_BASE_URL + assert handle.gateway_base_url == "http://127.0.0.1:9999" + finally: + client.close() + + +def test_embedded_handle_async_client_uses_asgi_transport() -> None: + app = MagicMock() + runtime = MagicMock() + http_client = object() + client_value = object() + handle = services.EmbeddedServiceHandle(app=app, runtime=runtime) + + with ( + patch( + "nemo_platform_ext.local.services.build_async_asgi_http_client", return_value=http_client + ) as build_client, + patch("nemo_platform_ext.local.services.AsyncNeMoPlatform", return_value=client_value) as platform_cls, + ): + client = handle.async_client(access_token="test-token") + + build_client.assert_called_once_with(app) + platform_cls.assert_called_once_with( + access_token="test-token", + http_client=http_client, + base_url=services.EMBEDDED_BASE_URL, + ) + assert client is client_value + + +def test_ensure_services_dispatches_to_embedded_mode() -> None: + cfg = ServiceRunConfig(mode=services.ServiceMode.EMBEDDED) + embedded_handle = MagicMock(spec=services.EmbeddedServiceHandle) + + with patch("nemo_platform_ext.local.services.start_embedded_services", return_value=embedded_handle): + handle = services.ensure_services(cfg) + + assert handle is embedded_handle + + +def test_ensure_services_dispatches_to_daemon_mode() -> None: + cfg = ServiceRunConfig(mode=services.ServiceMode.DAEMON) + daemon_handle = MagicMock(spec=services.DaemonServiceHandle) + + with ( + patch("nemo_platform_ext.local.services.get_service_handle", return_value=None), + patch("nemo_platform_ext.local.services.daemonize_services", return_value=daemon_handle), + ): + handle = services.ensure_services(cfg) + + assert handle is daemon_handle + + +def test_connect_services_uses_selected_mode_handle_client() -> None: + cfg = ServiceRunConfig(mode=services.ServiceMode.EMBEDDED) + handle = MagicMock(spec=services.EmbeddedServiceHandle) + client = object() + handle.client.return_value = client + + with patch("nemo_platform_ext.local.services.ensure_services", return_value=handle): + result = services.connect_services(cfg, access_token="test") + + assert result is client + handle.client.assert_called_once_with(access_token="test") + + +@pytest.mark.parametrize("mode", [services.ServiceMode.EMBEDDED, services.ServiceMode.DAEMON]) +def test_ensure_services_returns_handle_with_parity_methods(mode: services.ServiceMode) -> None: + cfg = ServiceRunConfig(mode=mode) + if mode is services.ServiceMode.EMBEDDED: + handle = MagicMock(spec=services.EmbeddedServiceHandle) + patch_target = "nemo_platform_ext.local.services.start_embedded_services" + else: + handle = MagicMock(spec=services.DaemonServiceHandle) + patch_target = "nemo_platform_ext.local.services.daemonize_services" + + with ( + patch("nemo_platform_ext.local.services.get_service_handle", return_value=None), + patch(patch_target, return_value=handle), + ): + result = services.ensure_services(cfg) + + assert result is handle + for method_name in ( + "is_running", + "wait_until_ready", + "wait_until_ready_async", + "client", + "async_client", + "start_services", + "start_services_async", + "stop", + "stop_async", + ): + assert hasattr(result, method_name), method_name + + +def test_daemonize_services_starts_child_module(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + _allow_tmp_path_socket_paths(monkeypatch) + cfg = ServiceRunConfig( + service_group="all", + scope="dev", + state_dir=tmp_path / "state", + runtime_dir=tmp_path / "run", + readiness_timeout=0.1, + readiness_poll_interval=0.01, + ) + proc = MagicMock() + proc.pid = 4242 + proc.poll.return_value = None + + with ( + patch("nemo_platform_ext.local.services.require_services_extra"), + patch("nemo_platform_ext.local.services._check_tcp_available"), + patch("nemo_platform_ext.local.services.probe_status", return_value=True), + patch("nemo_platform_ext.local.services.subprocess.Popen", return_value=proc) as popen, + ): + handle = services.daemonize_services(cfg) + + args = popen.call_args.args[0] + assert args[:3] == [sys.executable, "-m", f"{services.__package__}._service_child"] + request_path = Path(args[3]) + assert request_path.parent == tmp_path / "state" / "instances" / "dev" + assert request_path.suffix == ".json" + assert request_path.name != "run-request.json" + assert handle.transport == "uds" + assert handle.socket_path == tmp_path / "run" / "dev" / "nemo-platform.sock" + assert handle.pid == 4242 + proc.terminate.assert_not_called() + proc.kill.assert_not_called() + + +def test_daemonize_services_leaves_socket_preparation_to_child(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + _allow_tmp_path_socket_paths(monkeypatch) + cfg = ServiceRunConfig( + scope="dev", + state_dir=tmp_path / "state", + runtime_dir=tmp_path / "run", + readiness_timeout=0.1, + readiness_poll_interval=0.01, + ) + socket_path = cfg.resolved_socket_path + assert socket_path is not None + socket_path.parent.mkdir(parents=True) + socket_path.write_text("stale", encoding="utf-8") + proc = MagicMock() + proc.pid = 4242 + proc.poll.return_value = None + + with ( + patch("nemo_platform_ext.local.services.require_services_extra"), + patch("nemo_platform_ext.local.services._prepare_socket", side_effect=AssertionError("parent prepared socket")), + patch("nemo_platform_ext.local.services.probe_status", side_effect=[False, True]), + patch("nemo_platform_ext.local.services.subprocess.Popen", return_value=proc), + ): + handle = services.daemonize_services(cfg) + + assert handle.socket_path == socket_path + assert socket_path.read_text(encoding="utf-8") == "stale" + + +def test_write_run_request_writes_complete_payload_when_os_write_is_short( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + cfg = ServiceRunConfig(scope="dev", state_dir=tmp_path / "state", runtime_dir=tmp_path / "run") + real_write = os.write + + def short_write(fd: int, data: bytes) -> int: + return real_write(fd, data[: max(1, len(data) // 2)]) + + monkeypatch.setattr(services.os, "write", short_write) + + request_path = services._write_run_request(cfg) + + expected_payload = json.dumps(cfg.to_child_payload(), indent=2) + "\n" + assert request_path.read_text(encoding="utf-8") == expected_payload + + +def test_service_child_unlinks_request_after_read(tmp_path: Path) -> None: + request_path = tmp_path / "run-request.json" + cfg = ServiceRunConfig(scope="dev", state_dir=tmp_path / "state", runtime_dir=tmp_path / "run") + request_path.write_text(json.dumps(cfg.to_child_payload()), encoding="utf-8") + + with patch("nemo_platform_ext.local._service_child.run_services") as run_services: + result = _service_child.main([str(request_path)]) + + assert result == 0 + assert not request_path.exists() + child_cfg = run_services.call_args.args[0] + assert child_cfg.scope == "dev" + assert child_cfg.state_dir == str(tmp_path / "state") + assert child_cfg.runtime_dir == str(tmp_path / "run") + assert run_services.call_args.kwargs == {"_mode": "daemon"} + + +def test_daemonize_services_terminates_child_on_timeout(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + _allow_tmp_path_socket_paths(monkeypatch) + cfg = ServiceRunConfig( + scope="dev", + state_dir=tmp_path / "state", + runtime_dir=tmp_path / "run", + readiness_timeout=0.01, + readiness_poll_interval=0.001, + ) + proc = MagicMock() + proc.pid = 4242 + proc.poll.return_value = None + + with ( + patch("nemo_platform_ext.local.services.require_services_extra"), + patch("nemo_platform_ext.local.services.probe_status", return_value=False), + patch("nemo_platform_ext.local.services.subprocess.Popen", return_value=proc), + ): + with pytest.raises(services.ServicesStartupTimeoutError): + services.daemonize_services(cfg) + + proc.terminate.assert_called_once_with() + proc.kill.assert_not_called() + + +def test_daemonize_services_bounds_probe_and_sleep_by_remaining_deadline( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + _allow_tmp_path_socket_paths(monkeypatch) + cfg = ServiceRunConfig( + scope="dev", + state_dir=tmp_path / "state", + runtime_dir=tmp_path / "run", + readiness_timeout=5.0, + readiness_poll_interval=10.0, + ) + proc = MagicMock() + proc.pid = 4242 + proc.poll.return_value = None + + with ( + patch("nemo_platform_ext.local.services.require_services_extra"), + patch("nemo_platform_ext.local.services.probe_status", return_value=False) as probe_status, + patch("nemo_platform_ext.local.services.subprocess.Popen", return_value=proc), + patch("nemo_platform_ext.local.services.time.monotonic", side_effect=[0.0, 4.0, 4.5, 5.0]), + patch("nemo_platform_ext.local.services.time.sleep") as sleep, + ): + with pytest.raises(services.ServicesStartupTimeoutError): + services.daemonize_services(cfg) + + assert probe_status.call_args.kwargs["timeout"] == pytest.approx(1.0) + sleep.assert_called_once() + assert sleep.call_args.args[0] == pytest.approx(0.5) + proc.terminate.assert_called_once_with() + + +def test_daemonize_services_terminates_child_on_handle_error(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + _allow_tmp_path_socket_paths(monkeypatch) + cfg = ServiceRunConfig( + scope="dev", + state_dir=tmp_path / "state", + runtime_dir=tmp_path / "run", + transport="tcp", + ) + proc = MagicMock() + proc.pid = 4242 + proc.poll.return_value = None + + with ( + patch("nemo_platform_ext.local.services.require_services_extra"), + patch("nemo_platform_ext.local.services._check_tcp_available"), + patch("nemo_platform_ext.local.services.subprocess.Popen", return_value=proc), + patch( + "nemo_platform_ext.local.services.DaemonServiceHandle.from_config", + side_effect=RuntimeError("handle failed"), + ), + ): + with pytest.raises(RuntimeError, match="handle failed"): + services.daemonize_services(cfg) + + proc.terminate.assert_called_once_with() + proc.kill.assert_not_called() + + +def test_daemonize_services_terminates_child_on_probe_error(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + _allow_tmp_path_socket_paths(monkeypatch) + cfg = ServiceRunConfig( + scope="dev", + state_dir=tmp_path / "state", + runtime_dir=tmp_path / "run", + transport="tcp", + ) + proc = MagicMock() + proc.pid = 4242 + proc.poll.return_value = None + + with ( + patch("nemo_platform_ext.local.services.require_services_extra"), + patch("nemo_platform_ext.local.services._check_tcp_available"), + patch("nemo_platform_ext.local.services.probe_status", side_effect=RuntimeError("probe failed")), + patch("nemo_platform_ext.local.services.subprocess.Popen", return_value=proc), + ): + with pytest.raises(RuntimeError, match="probe failed"): + services.daemonize_services(cfg) + + proc.terminate.assert_called_once_with() + proc.kill.assert_not_called() + + +def test_daemonize_services_terminates_child_on_sleep_interruption( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + _allow_tmp_path_socket_paths(monkeypatch) + cfg = ServiceRunConfig( + scope="dev", + state_dir=tmp_path / "state", + runtime_dir=tmp_path / "run", + transport="tcp", + ) + proc = MagicMock() + proc.pid = 4242 + proc.poll.return_value = None + + with ( + patch("nemo_platform_ext.local.services.require_services_extra"), + patch("nemo_platform_ext.local.services._check_tcp_available"), + patch("nemo_platform_ext.local.services.probe_status", return_value=False), + patch("nemo_platform_ext.local.services.subprocess.Popen", return_value=proc), + patch("nemo_platform_ext.local.services.time.sleep", side_effect=KeyboardInterrupt), + ): + with pytest.raises(KeyboardInterrupt): + services.daemonize_services(cfg) + + proc.terminate.assert_called_once_with() + proc.kill.assert_not_called() + + +def test_daemonize_services_kills_child_when_terminate_times_out( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + _allow_tmp_path_socket_paths(monkeypatch) + cfg = ServiceRunConfig( + scope="dev", + state_dir=tmp_path / "state", + runtime_dir=tmp_path / "run", + transport="tcp", + readiness_timeout=0.01, + readiness_poll_interval=0.001, + ) + proc = MagicMock() + proc.pid = 4242 + proc.poll.return_value = None + proc.wait.side_effect = [subprocess.TimeoutExpired("nemo services", 5), None] + + with ( + patch("nemo_platform_ext.local.services.require_services_extra"), + patch("nemo_platform_ext.local.services._check_tcp_available"), + patch("nemo_platform_ext.local.services.probe_status", return_value=False), + patch("nemo_platform_ext.local.services.subprocess.Popen", return_value=proc), + ): + with pytest.raises(services.ServicesStartupTimeoutError): + services.daemonize_services(cfg) + + proc.terminate.assert_called_once_with() + proc.kill.assert_called_once_with() + + +async def test_daemonize_services_async_uses_thread(tmp_path: Path) -> None: + cfg = ServiceRunConfig(scope="dev", state_dir=tmp_path / "state", runtime_dir=tmp_path / "run") + handle = MagicMock() + + with patch("nemo_platform_ext.local.services.asyncio.to_thread", new=AsyncMock(return_value=handle)) as to_thread: + result = await services.daemonize_services_async(cfg) + + assert result is handle + to_thread.assert_awaited_once_with(services.daemonize_services, cfg) + + +def test_run_services_serves_embedded_app_with_socket_path(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + _allow_tmp_path_socket_paths(monkeypatch) + cfg = ServiceRunConfig( + services=["entities"], + controllers=["jobs"], + scope="dev", + port=_free_tcp_port(), + state_dir=tmp_path / "state", + runtime_dir=tmp_path / "run", + ) + handle = _embedded_handle() + + with ( + patch("nemo_platform_ext.local.services.require_services_extra"), + patch("nemo_platform_ext.local.services.start_embedded_services", return_value=handle) as start_embedded, + patch("nemo_platform_ext.local.services.serve_embedded_app") as serve_embedded, + ): + services.run_services(cfg, _mode="daemon") + + start_embedded.assert_called_once_with(cfg, env=None) + serve_embedded.assert_called_once_with(handle.app, cfg, tmp_path / "run" / "dev" / "nemo-platform.sock") + assert not (tmp_path / "state" / "instances" / "dev" / DESCRIPTOR_FILENAME).exists() + + +def test_serve_embedded_app_with_socket_path_listens_on_tcp_and_uds(tmp_path: Path) -> None: + cfg = ServiceRunConfig(transport="tcp", host="127.0.0.1", port=9090) + app = object() + socket_path = tmp_path / "nemo.sock" + + with patch("nmp.platform_runner.server._run_server_on_bound_sockets") as run_bound_sockets: + services.serve_embedded_app(app, cfg, socket_path) + + run_bound_sockets.assert_called_once_with(app, host="127.0.0.1", port=9090, socket_path=str(socket_path)) + + +def test_run_services_cleans_lock_when_log_path_resolution_fails( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + _allow_tmp_path_socket_paths(monkeypatch) + cfg = ServiceRunConfig( + scope="dev", port=_free_tcp_port(), state_dir=tmp_path / "state", runtime_dir=tmp_path / "run" + ) + + with ( + patch("nemo_platform_ext.local.services.require_services_extra"), + patch.object(PlatformAppConfig, "log_file_path", side_effect=RuntimeError("boom")), + ): + with pytest.raises(RuntimeError, match="boom"): + services.run_services(cfg) + + assert not services.process.is_instance_alive(cfg.scope, base_dir=cfg.state_root) + + +def test_run_services_restores_data_dir_and_lock_when_descriptor_write_fails( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + _allow_tmp_path_socket_paths(monkeypatch) + cfg = ServiceRunConfig( + scope="dev", + port=_free_tcp_port(), + state_dir=tmp_path / "state", + runtime_dir=tmp_path / "run", + data_dir=tmp_path / "data", + ) + monkeypatch.delenv("NMP_DATA_DIR", raising=False) + + with ( + patch("nemo_platform_ext.local.services.require_services_extra"), + patch("nemo_platform_ext.local.services.process.write_descriptor", side_effect=RuntimeError("boom")), + ): + with pytest.raises(RuntimeError, match="boom"): + services.run_services(cfg) + + assert "NMP_DATA_DIR" not in os.environ + assert not services.process.is_instance_alive(cfg.scope, base_dir=cfg.state_root) + + +def test_run_services_restores_existing_data_dir(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + _allow_tmp_path_socket_paths(monkeypatch) + cfg = ServiceRunConfig( + scope="dev", port=_free_tcp_port(), state_dir=tmp_path / "state", runtime_dir=tmp_path / "run" + ) + monkeypatch.setenv("NMP_DATA_DIR", "/shell/data") + + with ( + patch("nemo_platform_ext.local.services.require_services_extra"), + patch("nemo_platform_ext.local.services.start_embedded_services", return_value=_embedded_handle()), + patch("nemo_platform_ext.local.services.serve_embedded_app"), + ): + services.run_services(cfg) + + assert os.environ["NMP_DATA_DIR"] == "/shell/data" + + +def test_daemon_service_handle_tcp_client_uses_tcp_base_url(tmp_path: Path) -> None: + handle = services.DaemonServiceHandle( + scope="dev", + transport="tcp", + socket_path=None, + gateway_base_url=None, + host="0.0.0.0", + port=9090, + pid=123, + mode="daemon", + log_path=None, + state_dir=tmp_path / "state" / "instances" / "dev", + runtime_dir=None, + ) + + with patch("nemo_platform_ext.local.services.NeMoPlatform") as sdk: + handle.client(timeout=12) + + sdk.assert_called_once_with(timeout=12, base_url="http://localhost:9090") + + +def test_daemon_service_handle_uds_client_requires_socket_path(tmp_path: Path) -> None: + handle = services.DaemonServiceHandle( + scope="dev", + transport="uds", + socket_path=None, + gateway_base_url=None, + host="127.0.0.1", + port=8080, + pid=123, + mode="daemon", + log_path=None, + state_dir=tmp_path / "state" / "instances" / "dev", + runtime_dir=tmp_path / "run", + ) + + with pytest.raises(services.ServicesError, match="missing socket_path"): + handle.client() + + +def test_ensure_services_returns_existing_handle(tmp_path: Path) -> None: + cfg = ServiceRunConfig(scope="dev", state_dir=tmp_path / "state", runtime_dir=tmp_path / "run") + handle = MagicMock() + + with ( + patch("nemo_platform_ext.local.services.get_service_handle", return_value=handle), + patch("nemo_platform_ext.local.services.daemonize_services") as daemonize, + ): + result = services.ensure_services(cfg) + + assert result is handle + daemonize.assert_not_called() + + +def test_connect_services_respects_start_if_needed_false(tmp_path: Path) -> None: + cfg = ServiceRunConfig(scope="dev", state_dir=tmp_path / "state", runtime_dir=tmp_path / "run") + + with patch("nemo_platform_ext.local.services.get_service_handle", return_value=None): + with pytest.raises(services.ServicesNotRunningError, match="not running"): + services.connect_services(cfg, start_if_needed=False) + + +def test_stop_services_delegates_to_handle(tmp_path: Path) -> None: + cfg = ServiceRunConfig(scope="dev", state_dir=tmp_path / "state", runtime_dir=tmp_path / "run") + handle = MagicMock() + stop_result = MagicMock() + handle.stop.return_value = stop_result + + with patch("nemo_platform_ext.local.services.get_service_handle", return_value=handle): + result = services.stop_services(cfg, timeout=3.0, force=True) + + assert result is stop_result + handle.stop.assert_called_once_with(timeout=3.0, force=True) + + +def test_get_service_handle_returns_none_without_live_descriptor(tmp_path: Path) -> None: + cfg = ServiceRunConfig(scope="dev", state_dir=tmp_path / "state", runtime_dir=tmp_path / "run") + + with patch("nemo_platform_ext.local.services.process.read_descriptor", return_value=None): + assert services.get_service_handle(cfg) is None + + +def test_list_service_handles_filters_dead_or_descriptorless_instances(tmp_path: Path) -> None: + live_desc = InstanceDescriptor( + pid=123, + transport="tcp", + config=PlatformAppConfig(scope="live", state_root=tmp_path / "state"), + mode="daemon", + ) + infos = [ + MagicMock(descriptor=live_desc, alive=True), + MagicMock(descriptor=None, alive=True), + MagicMock(descriptor=live_desc, alive=False), + ] + + with patch("nemo_platform_ext.local.services.process.list_instances", return_value=infos): + handles = services.list_service_handles(tmp_path / "state") + + assert [handle.scope for handle in handles] == ["live"] + + +def test_get_service_handle_reads_live_descriptor(tmp_path: Path) -> None: + cfg = ServiceRunConfig(scope="dev", state_dir=tmp_path / "state", runtime_dir=tmp_path / "run") + state_dir = tmp_path / "state" / "instances" / "dev" + state_dir.mkdir(parents=True) + desc = InstanceDescriptor( + pid=123, + config=PlatformAppConfig( + scope="dev", + socket_path=str(tmp_path / "run" / "dev" / "nemo-platform.sock"), + state_root=tmp_path / "state", + runtime_root=tmp_path / "run", + ), + transport="uds", + mode="daemon", + ) + (state_dir / DESCRIPTOR_FILENAME).write_text(desc.model_dump_json(), encoding="utf-8") + + with patch("nemo_platform_ext.local.process.is_instance_alive", return_value=True): + handle = services.get_service_handle(cfg) + + assert handle is not None + assert handle.scope == "dev" + assert handle.transport == "uds" diff --git a/packages/nemo_platform_ext/tests/local/test_services_contract.py b/packages/nemo_platform_ext/tests/local/test_services_contract.py new file mode 100644 index 0000000000..acfdc227ca --- /dev/null +++ b/packages/nemo_platform_ext/tests/local/test_services_contract.py @@ -0,0 +1,334 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import threading +from collections.abc import Callable +from dataclasses import dataclass, field +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest +from nemo_platform_ext.local import services +from nemo_platform_ext.local.process import StopResult +from nemo_platform_ext.local.services import ServiceRunConfig +from nmp.platform_runner.config import PlatformAppConfig + + +@dataclass(frozen=True) +class ModeContractCase: + mode: services.ServiceMode + launcher_patch: str + existing_handle_patch_value: object | None + + +@dataclass +class ContractHandle: + mode: services.ServiceMode + calls: list[tuple[str, object]] = field(default_factory=list) + + def is_running(self) -> bool: + self.calls.append(("is_running", None)) + return True + + def wait_until_ready(self, timeout: float | None = None) -> None: + self.calls.append(("wait_until_ready", timeout)) + + async def wait_until_ready_async(self, timeout: float | None = None) -> None: + self.calls.append(("wait_until_ready_async", timeout)) + + def client(self, **kwargs: object) -> tuple[str, services.ServiceMode, dict[str, object]]: + self.calls.append(("client", kwargs)) + return ("client", self.mode, kwargs) + + def async_client(self, **kwargs: object) -> tuple[str, services.ServiceMode, dict[str, object]]: + self.calls.append(("async_client", kwargs)) + return ("async_client", self.mode, kwargs) + + def start_services(self, service_names: list[str] | tuple[str, ...]) -> services.StartServicesResult: + requested = list(service_names) + self.calls.append(("start_services", requested)) + return services.StartServicesResult( + requested=requested, + started=["auth", *requested], + already_active=[], + active=["secrets", "auth", *requested], + ) + + async def start_services_async(self, service_names: list[str] | tuple[str, ...]) -> services.StartServicesResult: + requested = list(service_names) + self.calls.append(("start_services_async", requested)) + return services.StartServicesResult( + requested=requested, + started=["auth", *requested], + already_active=[], + active=["secrets", "auth", *requested], + ) + + def stop(self, *, timeout: float = 30.0, force: bool = False) -> StopResult: + self.calls.append(("stop", {"timeout": timeout, "force": force})) + return StopResult(stopped_pids=[], swept_children=[]) + + async def stop_async(self, *, timeout: float = 30.0, force: bool = False) -> StopResult: + self.calls.append(("stop_async", {"timeout": timeout, "force": force})) + return StopResult(stopped_pids=[], swept_children=[]) + + +MODE_CONTRACT_CASES = [ + ModeContractCase( + mode=services.ServiceMode.EMBEDDED, + launcher_patch="nemo_platform_ext.local.services.start_embedded_services", + existing_handle_patch_value=None, + ), + ModeContractCase( + mode=services.ServiceMode.DAEMON, + launcher_patch="nemo_platform_ext.local.services.daemonize_services", + existing_handle_patch_value=None, + ), +] + + +@pytest.fixture(params=MODE_CONTRACT_CASES, ids=lambda case: case.mode.value) +def mode_case(request: pytest.FixtureRequest) -> ModeContractCase: + return request.param + + +def _config_for(case: ModeContractCase, tmp_path: Path) -> ServiceRunConfig: + return ServiceRunConfig( + mode=case.mode, + services=("secrets",), + scope=f"{case.mode.value}-contract", + state_dir=tmp_path / case.mode.value / "state", + runtime_dir=tmp_path / case.mode.value / "runtime", + ) + + +def test_contract_ensure_services_returns_running_mode_handle( + mode_case: ModeContractCase, + tmp_path: Path, +) -> None: + cfg = _config_for(mode_case, tmp_path) + handle = ContractHandle(mode_case.mode) + + with ( + patch( + "nemo_platform_ext.local.services.get_service_handle", return_value=mode_case.existing_handle_patch_value + ), + patch(mode_case.launcher_patch, return_value=handle), + ): + result = services.ensure_services(cfg) + + assert result is handle + assert result.is_running() is True + assert result.calls == [("is_running", None)] + + +def test_contract_connect_services_returns_client_from_selected_mode( + mode_case: ModeContractCase, + tmp_path: Path, +) -> None: + cfg = _config_for(mode_case, tmp_path) + handle = ContractHandle(mode_case.mode) + + with patch("nemo_platform_ext.local.services.ensure_services", return_value=handle): + client = services.connect_services(cfg, api_key="test-key") + + assert client == ("client", mode_case.mode, {"api_key": "test-key"}) + assert handle.calls == [("client", {"api_key": "test-key"})] + + +@pytest.mark.asyncio +async def test_contract_handle_lifecycle_methods_have_same_semantics( + mode_case: ModeContractCase, +) -> None: + handle = ContractHandle(mode_case.mode) + + handle.wait_until_ready(timeout=1.5) + await handle.wait_until_ready_async(timeout=2.5) + sync_start = handle.start_services(["jobs"]) + async_start = await handle.start_services_async(["jobs"]) + stop_result = handle.stop(timeout=3.0, force=True) + async_stop_result = await handle.stop_async(timeout=4.0, force=False) + + assert sync_start == services.StartServicesResult( + requested=["jobs"], + started=["auth", "jobs"], + already_active=[], + active=["secrets", "auth", "jobs"], + ) + assert async_start == sync_start + assert stop_result == StopResult(stopped_pids=[], swept_children=[]) + assert async_stop_result == StopResult(stopped_pids=[], swept_children=[]) + assert handle.calls == [ + ("wait_until_ready", 1.5), + ("wait_until_ready_async", 2.5), + ("start_services", ["jobs"]), + ("start_services_async", ["jobs"]), + ("stop", {"timeout": 3.0, "force": True}), + ("stop_async", {"timeout": 4.0, "force": False}), + ] + + +def test_contract_real_handles_report_same_staged_start_status_before_staged_start_lands(tmp_path: Path) -> None: + embedded = services.EmbeddedServiceHandle(app=object(), runtime=object()) + daemon = services.DaemonServiceHandle( + scope="daemon-contract", + transport="tcp", + socket_path=None, + gateway_base_url=None, + host="127.0.0.1", + port=8080, + pid=None, + mode="daemon", + log_path=None, + state_dir=tmp_path / "state", + runtime_dir=None, + ) + + for handle in (embedded, daemon): + with pytest.raises(services.ServicesError, match="Staged service start is not implemented"): + handle.start_services(["jobs"]) + + +def test_contract_embedded_and_daemon_child_both_delegate_models_to_platform_builder( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls: list[dict[str, object]] = [] + + def fake_build_platform_app( + config: PlatformAppConfig | None = None, + *, + env: object = None, + http_client: object = None, + ) -> MagicMock: + calls.append({"config": config, "env": env, "http_client": http_client}) + return MagicMock() + + def service_config(mode: services.ServiceMode) -> ServiceRunConfig: + return ServiceRunConfig( + mode=mode, + services=("models",), + controllers=(), + transport="tcp", + scope="sc-test", + state_dir=tmp_path / "state", + runtime_dir=tmp_path / "runtime", + ) + + with patch("nmp.platform_runner.server.build_platform_app", side_effect=fake_build_platform_app): + services.start_embedded_services(service_config(services.ServiceMode.EMBEDDED)) + + with ( + patch("nmp.platform_runner.server.build_platform_app", side_effect=fake_build_platform_app), + patch("nemo_platform_ext.local.services.require_services_extra"), + patch("nemo_platform_ext.local.services.process.is_instance_alive", return_value=False), + patch("nemo_platform_ext.local.services._check_tcp_available"), + patch("nemo_platform_ext.local.services.process.acquire_lock", return_value=123), + patch("nemo_platform_ext.local.services.process.log_path_for", return_value=tmp_path / "nemo.log"), + patch("nemo_platform_ext.local.services.process.write_descriptor"), + patch("nemo_platform_ext.local.services.process.remove_descriptor"), + patch("nemo_platform_ext.local.services.serve_embedded_app"), + patch("nemo_platform_ext.local.services.os.close"), + ): + services.run_services(service_config(services.ServiceMode.DAEMON), _mode="daemon") + + configs: list[PlatformAppConfig] = [] + for call in calls: + config = call["config"] + assert isinstance(config, PlatformAppConfig) + configs.append(config) + assert [config.services for config in configs] == [("models",), ("models",)] + assert [config.controllers for config in configs] == [(), ()] + assert [config.sidecars for config in configs] == [None, None] + + +def _sidecar_with_events(started: threading.Event, stopped: threading.Event) -> Callable[[threading.Event], None]: + def run(stop_signal: threading.Event) -> None: + started.set() + stop_signal.wait(timeout=5.0) + stopped.set() + + return run + + +def _patch_runner_registry( + monkeypatch: pytest.MonkeyPatch, + *, + sidecar_run_func: Callable[[threading.Event], None], +) -> None: + """Patch the platform runner registry so only a dummy 'models' service + and a test sidecar are available, avoiding real service imports.""" + from nmp.common.config import AuthConfig + from nmp.common.config.base import OIDCConfig + from nmp.common.service import Service + from nmp.platform_runner import config as runner_config + from nmp.platform_runner import registry, server + + class _DummyService(Service): + def __init__(self) -> None: + super().__init__(name="models", module_name="test.contract") + + def get_routers(self): + return [] + + dummy_services: dict[str, Service] = {"models": _DummyService()} + dummy_sidecars: dict[str, Callable] = {"adapters": sidecar_run_func} + + monkeypatch.setattr(runner_config, "get_available_services", lambda: dummy_services) + monkeypatch.setattr(runner_config, "get_available_controllers", lambda: {}) + monkeypatch.setattr( + runner_config, + "get_service_groups", + lambda _available: {"all": ["models"], "core": ["models"], "api": []}, + ) + monkeypatch.setattr(runner_config, "get_controller_groups", lambda _available: {"all": [], "core": []}) + monkeypatch.setattr(runner_config, "get_default_controllers", lambda _groups: []) + monkeypatch.setattr(runner_config, "AVAILABLE_SIDECARS", dummy_sidecars) + monkeypatch.setattr(registry, "AVAILABLE_SIDECARS", dummy_sidecars) + monkeypatch.setattr(server, "AVAILABLE_SIDECARS", dummy_sidecars, raising=False) + monkeypatch.setattr(server, "order_services_by_dependencies", lambda svc: svc) + + auth_cfg = AuthConfig( + enabled=False, + policy_decision_point_base_url="http://localhost:8181", + oidc=OIDCConfig(enabled=False), + ) + monkeypatch.setattr(server, "get_auth_config", lambda: auth_cfg) + monkeypatch.setattr("nmp.common.auth.middleware.get_auth_config", lambda: auth_cfg) + platform_cfg = MagicMock() + platform_cfg.seed_on_startup = False + platform_cfg.redirect_root_to_studio = False + monkeypatch.setattr(server, "get_platform_config", lambda: platform_cfg) + + +def test_embedded_mode_starts_sidecar_thread_via_full_resolution_path( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """End-to-end: start_embedded_services(models) resolves the adapters sidecar + and the sidecar thread actually runs when the app lifespan starts.""" + started = threading.Event() + stopped = threading.Event() + + _patch_runner_registry(monkeypatch, sidecar_run_func=_sidecar_with_events(started, stopped)) + + cfg = ServiceRunConfig( + mode=services.ServiceMode.EMBEDDED, + services=("models",), + controllers=(), + transport="tcp", + scope="sidecar-e2e", + state_dir=tmp_path / "state", + runtime_dir=tmp_path / "runtime", + ) + handle = services.start_embedded_services(cfg, env={}) + + from fastapi.testclient import TestClient + + with TestClient(handle.app) as client: + assert started.wait(timeout=2.0), "sidecar thread did not start" + assert client.get("/").status_code == 200 + + assert stopped.wait(timeout=2.0), "sidecar thread did not stop" diff --git a/packages/nemo_platform_ext/tests/local/test_sidecar_integration.py b/packages/nemo_platform_ext/tests/local/test_sidecar_integration.py new file mode 100644 index 0000000000..33cadabaa7 --- /dev/null +++ b/packages/nemo_platform_ext/tests/local/test_sidecar_integration.py @@ -0,0 +1,231 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Integration tests for sidecar lifecycle in embedded and daemon modes. + +These tests let the real ``build_platform_app`` → ``resolve_run_configuration`` → +``create_app`` chain run with a lightweight test sidecar registered in the +platform runner registry. They verify that sidecar threads actually start and +stop during the FastAPI app lifespan, covering the full resolution path without +mocking away the core wiring. +""" + +from __future__ import annotations + +import threading +from collections.abc import Callable +from pathlib import Path +from unittest.mock import MagicMock + +import pytest +from nemo_platform_ext.local import services +from nemo_platform_ext.local.services import ServiceRunConfig +from nmp.common.config import AuthConfig +from nmp.common.config.base import OIDCConfig +from nmp.common.service import Service +from nmp.platform_runner import config as runner_config +from nmp.platform_runner import registry, server + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +class _DummyService(Service): + """Minimal service that registers no routers.""" + + def __init__(self, name: str = "models") -> None: + super().__init__(name=name, module_name="test.sidecar_integration") + + def get_routers(self): + return [] + + +def _sidecar_with_events(started: threading.Event, stopped: threading.Event) -> Callable[[threading.Event], None]: + """Return a sidecar ``run(stop_signal)`` that signals start/stop via events.""" + + def run(stop_signal: threading.Event) -> None: + started.set() + stop_signal.wait(timeout=5.0) + stopped.set() + + return run + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture +def sidecar_events() -> tuple[threading.Event, threading.Event]: + return threading.Event(), threading.Event() + + +@pytest.fixture +def patched_registry( + monkeypatch: pytest.MonkeyPatch, + sidecar_events: tuple[threading.Event, threading.Event], +) -> tuple[threading.Event, threading.Event]: + """Patch the platform runner registry with a dummy models service and a + test sidecar, plus minimal auth/platform config stubs.""" + started, stopped = sidecar_events + dummy_services: dict[str, Service] = {"models": _DummyService()} + dummy_sidecars: dict[str, Callable] = {"adapters": _sidecar_with_events(started, stopped)} + + monkeypatch.setattr(runner_config, "get_available_services", lambda: dummy_services) + monkeypatch.setattr(runner_config, "get_available_controllers", lambda: {}) + monkeypatch.setattr( + runner_config, + "get_service_groups", + lambda _available: {"all": ["models"], "core": ["models"], "api": []}, + ) + monkeypatch.setattr(runner_config, "get_controller_groups", lambda _available: {"all": [], "core": []}) + monkeypatch.setattr(runner_config, "get_default_controllers", lambda _groups: []) + monkeypatch.setattr(runner_config, "AVAILABLE_SIDECARS", dummy_sidecars) + monkeypatch.setattr(registry, "AVAILABLE_SIDECARS", dummy_sidecars) + monkeypatch.setattr(server, "AVAILABLE_SIDECARS", dummy_sidecars, raising=False) + monkeypatch.setattr(server, "order_services_by_dependencies", lambda svc: svc) + + auth_cfg = AuthConfig( + enabled=False, + policy_decision_point_base_url="http://localhost:8181", + oidc=OIDCConfig(enabled=False), + ) + monkeypatch.setattr(server, "get_auth_config", lambda: auth_cfg) + monkeypatch.setattr("nmp.common.auth.middleware.get_auth_config", lambda: auth_cfg) + platform_cfg = MagicMock() + platform_cfg.seed_on_startup = False + platform_cfg.redirect_root_to_studio = False + monkeypatch.setattr(server, "get_platform_config", lambda: platform_cfg) + + return started, stopped + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +@pytest.mark.integration +def test_embedded_sidecar_auto_resolved_from_service_dependency( + patched_registry: tuple[threading.Event, threading.Event], + tmp_path: Path, +) -> None: + """start_embedded_services(models) auto-resolves the adapters sidecar via + SERVICE_SIDECAR_DEPENDENCIES and starts it during app lifespan.""" + started, stopped = patched_registry + + cfg = ServiceRunConfig( + mode=services.ServiceMode.EMBEDDED, + services=("models",), + controllers=(), + # sidecars=None triggers auto-resolution + transport="tcp", + scope="integ-embedded-auto", + state_dir=tmp_path / "state", + runtime_dir=tmp_path / "runtime", + ) + handle = services.start_embedded_services(cfg, env={}) + + from fastapi.testclient import TestClient + + with TestClient(handle.app) as client: + assert started.wait(timeout=2.0), "sidecar thread did not start" + assert client.get("/").status_code == 200 + + assert stopped.wait(timeout=2.0), "sidecar thread did not stop after lifespan exit" + + +@pytest.mark.integration +def test_embedded_explicit_sidecar_without_services( + patched_registry: tuple[threading.Event, threading.Event], + tmp_path: Path, +) -> None: + """An explicitly requested sidecar runs even when no services are selected.""" + started, stopped = patched_registry + + cfg = ServiceRunConfig( + mode=services.ServiceMode.EMBEDDED, + services=(), + controllers=(), + sidecars=("adapters",), + transport="tcp", + scope="integ-embedded-explicit", + state_dir=tmp_path / "state", + runtime_dir=tmp_path / "runtime", + ) + handle = services.start_embedded_services(cfg, env={}) + + from fastapi.testclient import TestClient + + with TestClient(handle.app) as client: + assert started.wait(timeout=2.0), "sidecar thread did not start" + assert client.get("/").status_code == 200 + + assert stopped.wait(timeout=2.0), "sidecar thread did not stop after lifespan exit" + + +@pytest.mark.integration +def test_run_services_daemon_mode_starts_sidecar_in_process( + patched_registry: tuple[threading.Event, threading.Event], + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """run_services(_mode='daemon') exercises the daemon code path in-process. + It calls start_embedded_services then serve_embedded_app. We intercept + serve_embedded_app to capture the app and exercise its lifespan, proving + the daemon path wires sidecars identically to embedded mode.""" + started, stopped = patched_registry + captured_app = {} + + def fake_serve(app, cfg, socket_path): + captured_app["app"] = app + + monkeypatch.setattr("nemo_platform_ext.local.services.require_services_extra", lambda: None) + monkeypatch.setattr("nemo_platform_ext.local.services.process.is_instance_alive", lambda *a, **kw: False) + monkeypatch.setattr("nemo_platform_ext.local.services._check_tcp_available", lambda *a: None) + monkeypatch.setattr("nemo_platform_ext.local.services.process.acquire_lock", lambda *a, **kw: 123) + monkeypatch.setattr("nemo_platform_ext.local.services.process.log_path_for", lambda *a, **kw: tmp_path / "nemo.log") + monkeypatch.setattr("nemo_platform_ext.local.services.process.write_descriptor", lambda *a, **kw: None) + monkeypatch.setattr("nemo_platform_ext.local.services.process.remove_descriptor", lambda *a, **kw: None) + monkeypatch.setattr("nemo_platform_ext.local.services.serve_embedded_app", fake_serve) + monkeypatch.setattr("nemo_platform_ext.local.services.os.close", lambda fd: None) + + cfg = ServiceRunConfig( + mode=services.ServiceMode.DAEMON, + services=("models",), + controllers=(), + transport="tcp", + scope="integ-daemon", + state_dir=tmp_path / "state", + runtime_dir=tmp_path / "runtime", + ) + services.run_services(cfg, _mode="daemon", env={}) + + assert "app" in captured_app, "serve_embedded_app was not called" + + from fastapi.testclient import TestClient + + with TestClient(captured_app["app"]) as client: + assert started.wait(timeout=2.0), "sidecar thread did not start in daemon mode" + assert client.get("/").status_code == 200 + + assert stopped.wait(timeout=2.0), "sidecar thread did not stop after lifespan exit" + + +@pytest.mark.integration +def test_embedded_rejects_unknown_sidecar_name(tmp_path: Path, patched_registry) -> None: + """Requesting a sidecar not in the registry raises ValueError with a clear message.""" + cfg = ServiceRunConfig( + mode=services.ServiceMode.EMBEDDED, + services=(), + controllers=(), + sidecars=("nonexistent",), + transport="tcp", + scope="integ-unknown", + state_dir=tmp_path / "state", + runtime_dir=tmp_path / "runtime", + ) + with pytest.raises(ValueError, match="Unknown sidecars: nonexistent"): + services.start_embedded_services(cfg, env={}) diff --git a/packages/nemo_platform_ext/tests/local/test_transport.py b/packages/nemo_platform_ext/tests/local/test_transport.py new file mode 100644 index 0000000000..7f22074890 --- /dev/null +++ b/packages/nemo_platform_ext/tests/local/test_transport.py @@ -0,0 +1,167 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from unittest.mock import AsyncMock, patch + +import httpx +import pytest +from fastapi import FastAPI +from nemo_platform_ext.local import transport + + +def _assert_timeout_values(timeout: httpx.Timeout, expected: float | None) -> None: + assert timeout.connect == expected + assert timeout.read == expected + assert timeout.write == expected + assert timeout.pool == expected + + +def test_build_sync_http_client_uses_finite_default_timeout(tmp_path) -> None: + client = transport.build_sync_http_client(tmp_path / "nemo.sock") + try: + _assert_timeout_values(client.timeout, 5.0) + finally: + client.close() + + +def test_build_sync_http_client_preserves_explicit_timeout_values(tmp_path) -> None: + no_timeout_client = transport.build_sync_http_client(tmp_path / "nemo.sock", timeout=None) + finite_timeout_client = transport.build_sync_http_client(tmp_path / "nemo.sock", timeout=12.0) + try: + _assert_timeout_values(no_timeout_client.timeout, None) + _assert_timeout_values(finite_timeout_client.timeout, 12.0) + finally: + no_timeout_client.close() + finite_timeout_client.close() + + +@pytest.mark.asyncio +async def test_build_async_http_client_uses_finite_default_timeout(tmp_path) -> None: + client = transport.build_async_http_client(tmp_path / "nemo.sock") + try: + _assert_timeout_values(client.timeout, 5.0) + finally: + await client.aclose() + + +@pytest.mark.asyncio +async def test_build_async_http_client_preserves_explicit_timeout_values(tmp_path) -> None: + no_timeout_client = transport.build_async_http_client(tmp_path / "nemo.sock", timeout=None) + finite_timeout_client = transport.build_async_http_client(tmp_path / "nemo.sock", timeout=12.0) + try: + _assert_timeout_values(no_timeout_client.timeout, None) + _assert_timeout_values(finite_timeout_client.timeout, 12.0) + finally: + await no_timeout_client.aclose() + await finite_timeout_client.aclose() + + +def test_build_sync_asgi_http_client_reaches_app() -> None: + app = FastAPI() + + @app.get("/status") + async def status() -> dict[str, str]: + return {"status": "healthy"} + + client = transport.build_sync_asgi_http_client(app) + try: + response = client.get("http://nemo-platform.local/status") + finally: + client.close() + + assert response.status_code == 200 + assert response.json() == {"status": "healthy"} + + +@pytest.mark.asyncio +async def test_build_async_asgi_http_client_reaches_app() -> None: + app = FastAPI() + + @app.get("/status") + async def status() -> dict[str, str]: + return {"status": "healthy"} + + client = transport.build_async_asgi_http_client(app) + try: + response = await client.get("http://nemo-platform.local/status") + finally: + await client.aclose() + + assert response.status_code == 200 + assert response.json() == {"status": "healthy"} + + +def test_wait_for_status_bounds_probe_and_sleep_by_remaining_deadline() -> None: + with ( + patch("nemo_platform_ext.local.transport.probe_status", return_value=False) as probe_status, + patch("nemo_platform_ext.local.transport.time.monotonic", side_effect=[0.0, 4.0, 4.5, 5.0]), + patch("nemo_platform_ext.local.transport.time.sleep") as sleep, + ): + result = transport.wait_for_status(base_url="http://127.0.0.1:8080", timeout=5.0, poll_interval=10.0) + + assert result is False + assert probe_status.call_args.kwargs["timeout"] == pytest.approx(1.0) + sleep.assert_called_once() + assert sleep.call_args.args[0] == pytest.approx(0.5) + + +@pytest.mark.asyncio +async def test_wait_for_status_async_bounds_probe_and_sleep_by_remaining_deadline() -> None: + with ( + patch( + "nemo_platform_ext.local.transport.probe_status_async", new=AsyncMock(return_value=False) + ) as probe_status, + patch("nemo_platform_ext.local.transport.time.monotonic", side_effect=[0.0, 4.0, 4.5, 5.0]), + patch("nemo_platform_ext.local.transport.asyncio.sleep", new=AsyncMock()) as sleep, + ): + result = await transport.wait_for_status_async( + base_url="http://127.0.0.1:8080", timeout=5.0, poll_interval=10.0 + ) + + assert result is False + assert probe_status.await_args.kwargs["timeout"] == pytest.approx(1.0) + sleep.assert_awaited_once() + assert sleep.await_args.args[0] == pytest.approx(0.5) + + +def test_probe_status_returns_true_for_status_200() -> None: + def handler(request: httpx.Request) -> httpx.Response: + assert str(request.url) == "http://127.0.0.1:8080/status" + return httpx.Response(200) + + with patch("nemo_platform_ext.local.transport.httpx.Client") as client_factory: + client = client_factory.return_value + client.get.side_effect = lambda url: handler(httpx.Request("GET", url)) + assert transport.probe_status(base_url="http://127.0.0.1:8080") is True + client.close.assert_called_once_with() + + +def test_probe_status_returns_false_for_request_error() -> None: + with patch("nemo_platform_ext.local.transport.httpx.Client") as client_factory: + client = client_factory.return_value + client.get.side_effect = httpx.ConnectError("boom") + assert transport.probe_status(base_url="http://127.0.0.1:8080") is False + client.close.assert_called_once_with() + + +@pytest.mark.asyncio +async def test_probe_status_async_returns_false_for_request_error() -> None: + with patch("nemo_platform_ext.local.transport.httpx.AsyncClient") as client_factory: + client = client_factory.return_value + client.get = AsyncMock(side_effect=httpx.ConnectError("boom")) + client.aclose = AsyncMock() + assert await transport.probe_status_async(base_url="http://127.0.0.1:8080") is False + client.aclose.assert_awaited_once_with() + + +def test_wait_for_status_returns_true_without_sleep_when_probe_succeeds() -> None: + with ( + patch("nemo_platform_ext.local.transport.probe_status", return_value=True) as probe_mock, + patch("nemo_platform_ext.local.transport.time.sleep") as sleep, + ): + assert transport.wait_for_status(base_url="http://127.0.0.1:8080", timeout=5.0) is True + + probe_mock.assert_called_once() + sleep.assert_not_called() diff --git a/packages/nmp_common/src/nmp/common/auth/client.py b/packages/nmp_common/src/nmp/common/auth/client.py index 4a75948359..67ddb90d17 100644 --- a/packages/nmp_common/src/nmp/common/auth/client.py +++ b/packages/nmp_common/src/nmp/common/auth/client.py @@ -11,6 +11,7 @@ import httpx from nmp.common.config import AuthConfig +from nmp.common.platform_endpoint import parse_platform_endpoint from pydantic import BaseModel, Field from .authz_format import validate_permission_strings, validate_runtime_authorize_scopes @@ -69,6 +70,10 @@ def policy_decision_point_base_url(self) -> Optional[str]: """Policy Decision Point (PDP) base URL for permission checks.""" return self.config.policy_decision_point_base_url + def _new_pdp_http_client(self) -> httpx.AsyncClient: + endpoint = parse_platform_endpoint(self.config.policy_decision_point_base_url) + return endpoint.async_http_client(timeout=self.config.policy_decision_point_request_timeout_seconds) + @property def _pdp_request_headers(self) -> dict[str, str]: """Headers sent with every PDP HTTP request. @@ -154,9 +159,7 @@ async def authorize_request( if client: response = await client.post(auth_url, json={"input": auth_input}, headers=pdp_headers) else: - async with httpx.AsyncClient( - timeout=self.config.policy_decision_point_request_timeout_seconds - ) as temp_client: + async with self._new_pdp_http_client() as temp_client: response = await temp_client.post(auth_url, json={"input": auth_input}, headers=pdp_headers) response.raise_for_status() @@ -222,7 +225,7 @@ async def create_model( client = self.http_client should_close = False if client is None: - client = httpx.AsyncClient(timeout=self.config.policy_decision_point_request_timeout_seconds) + client = self._new_pdp_http_client() should_close = True try: @@ -401,6 +404,9 @@ async def wait_role( if not self.auth_enabled: return True + if not self.policy_decision_point_base_url: + raise RuntimeError("Policy Decision Point URL not configured for role checks") + if poll_interval is None: poll_interval = self.config.propagation_poll_interval_seconds start_time = asyncio.get_event_loop().time() @@ -408,11 +414,7 @@ async def wait_role( # Use provided http_client, instance http_client (from middleware), or create a new one # See architecture/docs/http-client-injection.md for injection patterns. - client = ( - http_client - or self.http_client - or httpx.AsyncClient(timeout=self.config.policy_decision_point_request_timeout_seconds) - ) + client = http_client or self.http_client or self._new_pdp_http_client() should_close = http_client is None and self.http_client is None try: diff --git a/packages/nmp_common/src/nmp/common/auth/middleware.py b/packages/nmp_common/src/nmp/common/auth/middleware.py index ed4bc11b38..348b6477b8 100644 --- a/packages/nmp_common/src/nmp/common/auth/middleware.py +++ b/packages/nmp_common/src/nmp/common/auth/middleware.py @@ -10,6 +10,7 @@ from fastapi import Request, Response from nmp.common.config import AuthConfig, get_auth_config from nmp.common.observability.context import get_app_ctx +from nmp.common.platform_endpoint import parse_platform_endpoint from starlette.middleware.base import BaseHTTPMiddleware from starlette.responses import JSONResponse from starlette.types import ASGIApp @@ -55,8 +56,8 @@ def _embedded_pdp_base_url_hint(config: AuthConfig) -> str: return "" base = (config.policy_decision_point_base_url or "").strip() return ( - " For embedded PDP, auth.policy_decision_point_base_url must be the HTTP origin where " - "this process serves /apis/auth (same as platform base_url / NMP_BASE_URL). " + " For embedded PDP, auth.policy_decision_point_base_url must be the typed endpoint where " + "this process serves /apis/auth (same as platform base_url / NMP_BASE_URL; HTTP(S) or unix://). " f"Absolute PDP URLs ignore the injected ASGI client base_url. Current auth.policy_decision_point_base_url={base!r}." ) @@ -166,7 +167,8 @@ def _get_client(self, request: Request) -> httpx.AsyncClient: An async HTTP client configured with auth.policy_decision_point_request_timeout_seconds """ if self._client is None: - self._client = httpx.AsyncClient(timeout=self.config.policy_decision_point_request_timeout_seconds) + endpoint = parse_platform_endpoint(self.config.policy_decision_point_base_url) + self._client = endpoint.async_http_client(timeout=self.config.policy_decision_point_request_timeout_seconds) return self._client def _update_auth_context(self, principal: Principal) -> None: diff --git a/packages/nmp_common/src/nmp/common/config/base.py b/packages/nmp_common/src/nmp/common/config/base.py index 1bb654c6da..badb052cf1 100644 --- a/packages/nmp_common/src/nmp/common/config/base.py +++ b/packages/nmp_common/src/nmp/common/config/base.py @@ -317,9 +317,14 @@ class AuthConfig(create_service_config_class("auth")): # ty: ignore[unsupported ) def get_pdp_url(self, entrypoint: str) -> str: + # Import lazily to avoid a module cycle: platform_endpoint imports + # PlatformConfig from nmp.common.config, which is defined in this file. + from nmp.common.platform_endpoint import parse_platform_endpoint + + endpoint = parse_platform_endpoint(self.policy_decision_point_base_url) if self.policy_decision_point_provider == "opa": - return f"{self.policy_decision_point_base_url}/v1/data/authz/{entrypoint}" - return f"{self.policy_decision_point_base_url}/apis/auth/v2/authz/{entrypoint}" + return f"{endpoint.connect_base_url}/v1/data/authz/{entrypoint}" + return f"{endpoint.connect_base_url}/apis/auth/v2/authz/{entrypoint}" @property def auth_url(self) -> str: diff --git a/packages/nmp_common/src/nmp/common/platform_endpoint.py b/packages/nmp_common/src/nmp/common/platform_endpoint.py new file mode 100644 index 0000000000..c03cb121ef --- /dev/null +++ b/packages/nmp_common/src/nmp/common/platform_endpoint.py @@ -0,0 +1,96 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Typed platform endpoint resolution for HTTP(S) and Unix domain sockets.""" + +from __future__ import annotations + +import os +from dataclasses import dataclass +from pathlib import Path +from typing import Literal + +import httpx +from httpx._types import TimeoutTypes +from nmp.common.config import PlatformConfig + +UDS_BASE_URL = "http://nemo-platform.local" + + +@dataclass(frozen=True) +class PlatformEndpoint: + connect_base_url: str + socket_path: Path | None + transport: Literal["tcp", "uds"] + + def sync_http_client(self, *, timeout: TimeoutTypes | None = None) -> httpx.Client: + if self.transport == "uds": + if self.socket_path is None: + raise ValueError("UDS endpoint is missing a socket path") + transport = httpx.HTTPTransport(uds=str(self.socket_path)) + if timeout is None: + return httpx.Client(transport=transport, follow_redirects=True) + return httpx.Client(transport=transport, follow_redirects=True, timeout=timeout) + if timeout is None: + return httpx.Client(follow_redirects=True) + return httpx.Client(follow_redirects=True, timeout=timeout) + + def async_http_client(self, *, timeout: TimeoutTypes | None = None) -> httpx.AsyncClient: + if self.transport == "uds": + if self.socket_path is None: + raise ValueError("UDS endpoint is missing a socket path") + transport = httpx.AsyncHTTPTransport(uds=str(self.socket_path)) + if timeout is None: + return httpx.AsyncClient(transport=transport, follow_redirects=True) + return httpx.AsyncClient(transport=transport, follow_redirects=True, timeout=timeout) + if timeout is None: + return httpx.AsyncClient(follow_redirects=True) + return httpx.AsyncClient(follow_redirects=True, timeout=timeout) + + +def resolve_platform_endpoint(platform_config: PlatformConfig | None = None) -> PlatformEndpoint: + """Resolve the default platform endpoint from ``NMP_BASE_URL`` / config.""" + + if platform_config is None: + from nmp.common.config import Configuration + + platform_config = Configuration.get_platform_config() + return parse_platform_endpoint(platform_config.base_url) + + +def resolve_service_endpoint(service_name: str, platform_config: PlatformConfig | None = None) -> PlatformEndpoint: + """Resolve a service endpoint using ``NMP__URL`` before ``NMP_BASE_URL``.""" + + if platform_config is None: + from nmp.common.config import Configuration + + platform_config = Configuration.get_platform_config() + env_name = f"NMP_{service_name.upper().replace('-', '_')}_URL" + endpoint = os.environ.get(env_name) or platform_config.get_service_url(service_name) + return parse_platform_endpoint(endpoint) + + +def parse_platform_endpoint(endpoint: str) -> PlatformEndpoint: + """Parse an HTTP(S) or ``unix://`` endpoint into a typed transport model.""" + + if endpoint.startswith(("http://", "https://")): + try: + parsed = httpx.URL(endpoint) + except httpx.InvalidURL as error: + raise ValueError(f"Invalid platform endpoint URL {endpoint!r}") from error + if not parsed.host: + raise ValueError(f"HTTP(S) platform endpoint must include a host, got {endpoint!r}") + return PlatformEndpoint(connect_base_url=endpoint.rstrip("/"), socket_path=None, transport="tcp") + if endpoint.startswith("unix://"): + socket_path = _parse_unix_socket_path(endpoint) + return PlatformEndpoint(connect_base_url=UDS_BASE_URL, socket_path=socket_path, transport="uds") + if endpoint.startswith("/"): + raise ValueError(f"Raw socket paths are not valid endpoint URLs; use unix://{endpoint}") + raise ValueError(f"Unsupported platform endpoint URL {endpoint!r}; expected http://, https://, or unix://") + + +def _parse_unix_socket_path(endpoint: str) -> Path: + raw_path = endpoint.removeprefix("unix://") + if not raw_path.startswith("/"): + raise ValueError(f"UDS endpoint must use an absolute socket path, got {endpoint!r}") + return Path(raw_path) diff --git a/packages/nmp_common/src/nmp/common/sdk_factory.py b/packages/nmp_common/src/nmp/common/sdk_factory.py index 37076d3f82..a34a3daa94 100644 --- a/packages/nmp_common/src/nmp/common/sdk_factory.py +++ b/packages/nmp_common/src/nmp/common/sdk_factory.py @@ -14,6 +14,7 @@ from nmp.common.http_clients import shared_async_http_client, shared_sync_http_client from nmp.common.observability import MARK_INTERNAL_REQUEST_HEADERS from nmp.common.observability.otel import get_otel_headers +from nmp.common.platform_endpoint import PlatformEndpoint, resolve_platform_endpoint, resolve_service_endpoint logger = logging.getLogger(__name__) PlatformSDKT = TypeVar("PlatformSDKT", NeMoPlatform, AsyncNeMoPlatform) @@ -61,7 +62,18 @@ def resolve_platform_request_url( return request_url api_name = match.group(1) - service_url = httpx.URL(platform_config.get_service_url(api_name)) + svc_endpoint = resolve_service_endpoint(api_name, platform_config) + if svc_endpoint.transport == "uds": + logger.debug( + "Routing URL to UDS service", + extra={ + "service": api_name, + "path": request_url.path, + "transport": svc_endpoint.transport, + }, + ) + return request_url.copy_with(scheme="http", host="nemo-platform.local", port=None) + service_url = httpx.URL(svc_endpoint.connect_base_url) routed_url = request_url.copy_with( scheme=service_url.scheme, host=service_url.host, @@ -120,6 +132,30 @@ def with_options_preserving_request_router(base_sdk: PlatformSDKT, **kwargs: Any return scoped_sdk +def _sync_http_client_for_endpoint( + endpoint: PlatformEndpoint, + http_client: httpx.Client | None, +) -> httpx.Client: + if http_client is not None: + return http_client + if endpoint.transport == "uds": + return endpoint.sync_http_client() + return shared_sync_http_client() + + +def _async_http_client_for_endpoint( + endpoint: PlatformEndpoint, + http_client: httpx.AsyncClient | None, +) -> httpx.AsyncClient: + if http_client is not None: + return http_client + if _test_http_client is not None: + return _test_http_client + if endpoint.transport == "uds": + return endpoint.async_http_client() + return shared_async_http_client() + + def _get_default_headers( as_service: str | None = None, internal: bool = False, on_behalf_of: str | Principal | None = None ) -> dict[str, str]: @@ -206,9 +242,10 @@ def get_platform_sdk( Configured NeMoPlatform SDK instance. """ headers = _get_default_headers(as_service, internal, on_behalf_of) + endpoint = resolve_platform_endpoint() sdk = NeMoPlatform( - base_url=base_url or _base_url_from_config(), - http_client=http_client or shared_sync_http_client(), + base_url=base_url or endpoint.connect_base_url, + http_client=_sync_http_client_for_endpoint(endpoint, http_client), default_headers=headers if headers else None, ) return attach_platform_request_router(sdk) @@ -295,13 +332,14 @@ def get_async_platform_sdk( Configured AsyncNeMoPlatform SDK instance. """ headers = _get_default_headers(as_service, internal, on_behalf_of) + endpoint = resolve_platform_endpoint() # Use explicitly provided http_client (from DependencyProvider) or fall back to # module-level _test_http_client for backward compatibility with direct callers. - effective_client = http_client or _test_http_client or shared_async_http_client() + effective_client = _async_http_client_for_endpoint(endpoint, http_client) sdk = AsyncNeMoPlatform( - base_url=base_url or _base_url_from_config(), + base_url=base_url or endpoint.connect_base_url, http_client=effective_client, default_headers=headers if headers else None, ) diff --git a/packages/nmp_common/src/nmp/common/service/api/health.py b/packages/nmp_common/src/nmp/common/service/api/health.py index 9d5714096e..95722c2902 100644 --- a/packages/nmp_common/src/nmp/common/service/api/health.py +++ b/packages/nmp_common/src/nmp/common/service/api/health.py @@ -6,15 +6,58 @@ import logging import threading import time +from collections.abc import Mapping from typing import Any import httpx from nmp.common.config import PlatformConfig from nmp.common.observability import MARK_INTERNAL_REQUEST_HEADERS +from nmp.common.platform_endpoint import resolve_service_endpoint logger = logging.getLogger(__name__) +def _status_names(values: object) -> set[str]: + if not isinstance(values, list): + return set() + + names: set[str] = set() + for value in values: + if isinstance(value, Mapping): + name = value.get("name") + else: + name = getattr(value, "name", value) + if isinstance(name, str): + names.add(name) + return names + + +def service_ready_state_from_status(data: object, service_name: str) -> bool | None: + """Return service readiness from a platform /status payload. + + ``True`` means the service is ready or absent from this platform deployment. + ``False`` means the service is explicitly present but not ready. + ``None`` means the payload shape is unusable and should be retried. + """ + if not isinstance(data, Mapping): + return None + + services = data.get("services") or {} + if not isinstance(services, Mapping): + return None + + ready = _status_names(services.get("ready") or []) + if service_name in ready: + return True + + not_ready = _status_names(services.get("not_ready") or []) + if service_name in not_ready: + return False + + # Service is absent from this deployment — treat as ready so callers don't block. + return True + + async def async_wait_for_service_ready( platform_config: PlatformConfig, service_name: str, @@ -40,10 +83,11 @@ async def async_wait_for_service_ready( """ import asyncio - status_url = f"{platform_config.get_service_url(service_name).rstrip('/')}/status" + endpoint = resolve_service_endpoint(service_name, platform_config) + status_url = f"{endpoint.connect_base_url.rstrip('/')}/status" own_client = http_client is None if http_client is None: - http_client = httpx.AsyncClient(timeout=2.0) + http_client = endpoint.async_http_client(timeout=2.0) logger.debug("Waiting for service to be ready", extra={"service": service_name, "url": status_url}) @@ -56,9 +100,9 @@ async def async_wait_for_service_ready( headers=MARK_INTERNAL_REQUEST_HEADERS, ) if response.status_code == 200: - data: dict[str, Any] = response.json() - ready = (data.get("services") or {}).get("ready") or [] - if service_name in ready: + data: Any = response.json() + ready = service_ready_state_from_status(data, service_name) + if ready is True: logger.info("Service is ready", extra={"service": service_name}) return True except (httpx.RequestError, ValueError) as e: @@ -97,23 +141,16 @@ async def async_wait_for_dependencies( Returns: True if all dependencies became ready, False if any timed out. """ - own_client = http_client is None - if own_client: - http_client = httpx.AsyncClient(timeout=2.0) - try: - for dep in dependency_names: - if not await async_wait_for_service_ready( - platform_config, - dep, - timeout=timeout_per_service, - poll_interval=poll_interval, - http_client=http_client, - ): - return False - return True - finally: - if own_client and http_client is not None: - await http_client.aclose() + for dep in dependency_names: + if not await async_wait_for_service_ready( + platform_config, + dep, + timeout=timeout_per_service, + poll_interval=poll_interval, + http_client=http_client, + ): + return False + return True def wait_for_service_ready( @@ -141,31 +178,36 @@ def wait_for_service_ready( True if the service became ready, False if timeout or stop signal. """ start_time = time.time() - status_url = f"{platform_config.get_service_url(service_name).rstrip('/')}/status" + endpoint = resolve_service_endpoint(service_name, platform_config) + status_url = f"{endpoint.connect_base_url.rstrip('/')}/status" + http_client = endpoint.sync_http_client(timeout=2.0) logger.info( "Waiting for service to be ready", extra={"service": service_name, "url": status_url}, ) - while not stop_signal.is_set() and (time.time() - start_time) < timeout: - try: - response = httpx.get(status_url, timeout=2.0, headers=MARK_INTERNAL_REQUEST_HEADERS) - if response.status_code == 200: - data: dict[str, Any] = response.json() - ready = (data.get("services") or {}).get("ready") or [] - if service_name in ready: - logger.debug( - "Service is ready", - extra={"service": service_name, "url": status_url}, - ) - return True - except (httpx.RequestError, ValueError) as e: - logger.debug( - "Status check failed, will retry", - extra={"service": service_name, "url": status_url, "error": str(e)}, - ) - time.sleep(poll_interval) + try: + while not stop_signal.is_set() and (time.time() - start_time) < timeout: + try: + response = http_client.get(status_url, headers=MARK_INTERNAL_REQUEST_HEADERS) + if response.status_code == 200: + data: Any = response.json() + ready = service_ready_state_from_status(data, service_name) + if ready is True: + logger.debug( + "Service is ready", + extra={"service": service_name, "url": status_url}, + ) + return True + except (httpx.RequestError, ValueError) as e: + logger.debug( + "Status check failed, will retry", + extra={"service": service_name, "url": status_url, "error": str(e)}, + ) + time.sleep(poll_interval) + finally: + http_client.close() if stop_signal.is_set(): logger.debug("Stop signal received while waiting for service") diff --git a/packages/nmp_common/src/nmp/common/service/base.py b/packages/nmp_common/src/nmp/common/service/base.py index f1d59d573c..8a1c076fc3 100644 --- a/packages/nmp_common/src/nmp/common/service/base.py +++ b/packages/nmp_common/src/nmp/common/service/base.py @@ -20,6 +20,7 @@ from nmp.common.config import Configuration, PlatformConfig, ServiceConfig from nmp.common.controller import Controller from nmp.common.entities.client import EntityClient +from nmp.common.platform_endpoint import resolve_service_endpoint logger = logging.getLogger(__name__) @@ -545,39 +546,37 @@ async def startup(self) -> None: import time from nmp.common.observability import MARK_INTERNAL_REQUEST_HEADERS + from nmp.common.service.api.health import service_ready_state_from_status - status_url = f"{self.platform_config.get_service_url(service_name).rstrip('/')}/status" - client = self._dependency_provider.get_http_client() + endpoint = resolve_service_endpoint(service_name, self.platform_config) + status_url = f"{endpoint.connect_base_url.rstrip('/')}/status" + own_client = endpoint.transport == "uds" + client = endpoint.async_http_client(timeout=2.0) if own_client else self._dependency_provider.get_http_client() logger.debug("Waiting for service to be ready", extra={"service": service_name, "url": status_url}) start_time = time.time() - while (time.time() - start_time) < timeout: - try: - response = await client.get(status_url, timeout=2.0, headers=MARK_INTERNAL_REQUEST_HEADERS) - if response.status_code == 200: - data = response.json() - services = data.get("services") or {} - ready = services.get("ready") or [] - if service_name in ready: - logger.debug("Service is ready", extra={"service": service_name}) - return True - # If the service isn't in any list (ready/not_ready), it's not - # part of this deployment — skip waiting rather than timing out. - not_ready = services.get("not_ready") or [] - not_ready_names = [ - n.get("name", n) if isinstance(n, dict) else getattr(n, "name", n) for n in not_ready - ] - if service_name not in ready and service_name not in not_ready_names: - logger.debug( - "Dependency not present in platform, skipping wait", - extra={"service": service_name}, - ) - return True - # Service is in not_ready; keep polling - except httpx.RequestError: - pass - await asyncio.sleep(poll_interval) + try: + while (time.time() - start_time) < timeout: + try: + response = await client.get(status_url, timeout=2.0, headers=MARK_INTERNAL_REQUEST_HEADERS) + if response.status_code == 200: + try: + data = response.json() + except ValueError: + data = None + ready = service_ready_state_from_status(data, service_name) + if ready is True: + logger.debug("Service is ready", extra={"service": service_name}) + return True + # ``False`` means the service is explicitly not_ready; keep polling. + # ``None`` means the status payload shape was unusable; retry. + except httpx.RequestError: + pass + await asyncio.sleep(poll_interval) + finally: + if own_client: + await client.aclose() logger.warning("Timeout waiting for service to be ready", extra={"service": service_name, "timeout": timeout}) return False diff --git a/packages/nmp_common/tests/auth/test_client.py b/packages/nmp_common/tests/auth/test_client.py index 724c5fd9fc..1ac5e82075 100644 --- a/packages/nmp_common/tests/auth/test_client.py +++ b/packages/nmp_common/tests/auth/test_client.py @@ -228,6 +228,48 @@ async def test_authorize_request_sends_delegate_claims(self, auth_config, princi assert body["on_behalf_of_principal_id"] == "user@example.com" +class TestWaitRole: + @pytest.mark.asyncio + async def test_wait_role_requires_pdp_url_before_creating_client(self, principal): + auth_config = AuthConfig(enabled=True, policy_decision_point_base_url="") + auth_client = AuthClient(principal=principal, config=auth_config) + + with patch.object(auth_client, "_new_pdp_http_client") as new_client: + with pytest.raises(RuntimeError, match="Policy Decision Point URL not configured"): + await auth_client.wait_role( + "user@example.com", + "test-workspace", + "Viewer", + timeout=0.01, + poll_interval=0.01, + ) + + new_client.assert_not_called() + + @pytest.mark.asyncio + async def test_wait_role_reuses_provided_http_client(self, auth_config, principal): + mock_http_client = httpx.AsyncClient() + mock_response = MagicMock() + mock_response.json.return_value = {"result": {"has_role": True}} + mock_response.raise_for_status = MagicMock() + + with ( + patch.object(mock_http_client, "post", new_callable=AsyncMock, return_value=mock_response), + patch.object(mock_http_client, "aclose", new_callable=AsyncMock) as close, + ): + auth_client = AuthClient(principal=principal, config=auth_config) + assert await auth_client.wait_role( + "user@example.com", + "test-workspace", + "Viewer", + timeout=0.01, + poll_interval=0.01, + http_client=mock_http_client, + ) + + close.assert_not_called() + + class TestOnBehalfOfHasPermissions: """Tests for the on_behalf_of_has_permissions method.""" diff --git a/packages/nmp_common/tests/nmp_common/test_common_config.py b/packages/nmp_common/tests/nmp_common/test_common_config.py index a0bcac90c5..047a419827 100644 --- a/packages/nmp_common/tests/nmp_common/test_common_config.py +++ b/packages/nmp_common/tests/nmp_common/test_common_config.py @@ -5,6 +5,7 @@ import pytest from nmp.common.config import ( + AuthConfig, CommonServiceConfig, Configuration, DatabaseConfig, @@ -12,6 +13,7 @@ get_common_service_config, get_platform_config, ) +from nmp.common.platform_endpoint import UDS_BASE_URL class TestPlatformConfig: @@ -174,6 +176,29 @@ def test_create_service_pattern(self): assert match.group(1) == "my" +class TestAuthConfig: + """Tests for shared auth configuration helpers.""" + + def test_get_pdp_url_for_embedded_provider_normalizes_base_url(self): + config = AuthConfig(policy_decision_point_base_url="http://localhost:8080/") + + assert config.get_pdp_url("allow") == "http://localhost:8080/apis/auth/v2/authz/allow" + assert config.auth_url == "http://localhost:8080/apis/auth/v2/authz/allow" + + def test_get_pdp_url_for_opa_provider_normalizes_base_url(self): + config = AuthConfig( + policy_decision_point_base_url="http://opa:8181/", + policy_decision_point_provider="opa", + ) + + assert config.get_pdp_url("has_permissions") == "http://opa:8181/v1/data/authz/has_permissions" + + def test_get_pdp_url_for_uds_endpoint_uses_connect_base_url(self): + config = AuthConfig(policy_decision_point_base_url="unix:///tmp/nemo-platform.sock") + + assert config.get_pdp_url("allow") == f"{UDS_BASE_URL}/apis/auth/v2/authz/allow" + + class TestCommonServiceConfig: """Tests for CommonServiceConfig.""" diff --git a/packages/nmp_common/tests/nmp_common/test_common_service.py b/packages/nmp_common/tests/nmp_common/test_common_service.py index 7d29d156e3..45d1bcce6f 100644 --- a/packages/nmp_common/tests/nmp_common/test_common_service.py +++ b/packages/nmp_common/tests/nmp_common/test_common_service.py @@ -3,10 +3,16 @@ """Tests for nmp.common.service module.""" +import asyncio +import threading from typing import List +from unittest.mock import AsyncMock, patch +import httpx import pytest from fastapi import APIRouter, FastAPI +from fastapi.testclient import TestClient +from nmp.common.config import PlatformConfig from nmp.common.service import DependencyProvider, RouterConfig, Service @@ -27,8 +33,8 @@ def _route_paths(app: FastAPI) -> set[str]: class MockService(Service): """Mock implementation of Service for testing.""" - def __init__(self): - super().__init__(name="test-service", module_name="nmp.test") + def __init__(self, dependency_provider: DependencyProvider | None = None): + super().__init__(name="test-service", module_name="nmp.test", dependency_provider=dependency_provider) def get_routers(self) -> List[RouterConfig]: router = APIRouter() @@ -158,6 +164,51 @@ async def test_service_is_ready_default(self): service = MockService() assert await service.is_ready() is True + @pytest.mark.asyncio + async def test_wait_for_service_ready_retries_malformed_status_payloads(self): + """Test malformed 200 /status payloads are retried.""" + requests: list[httpx.Request] = [] + responses = [ + httpx.Response(status_code=200, content=b"{"), + httpx.Response(status_code=200, json=["not", "a", "mapping"]), + httpx.Response(status_code=200, json={"services": {"ready": ["entities"]}}), + ] + + def handler(request: httpx.Request) -> httpx.Response: + requests.append(request) + return responses.pop(0) + + provider = DependencyProvider() + provider._platform_config = PlatformConfig(base_url="http://platform.local") + async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client: + provider._http_client = client + service = MockService(dependency_provider=provider) + + ready = await service.wait_for_service_ready("entities", timeout=1.0, poll_interval=0) + + assert ready is True + assert len(requests) == 3 + + @pytest.mark.asyncio + async def test_wait_for_service_ready_skips_service_absent_from_status(self): + """Test dependencies absent from /status are treated as not part of this deployment.""" + requests: list[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response(200, json={"services": {"ready": ["entities"], "not_ready": []}}) + + provider = DependencyProvider() + provider._platform_config = PlatformConfig(base_url="http://platform.local") + async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client: + provider._http_client = client + service = MockService(dependency_provider=provider) + + ready = await service.wait_for_service_ready("models", timeout=1.0, poll_interval=0) + + assert ready is True + assert len(requests) == 1 + class TestDependencyProvider: """Tests for DependencyProvider class.""" @@ -183,3 +234,46 @@ def test_service_has_provider(self): service = MockService() assert service.dependency_provider is not None assert isinstance(service.dependency_provider, DependencyProvider) + + +class LifecycleService(MockService): + def __init__(self): + super().__init__() + self.events: list[str] = [] + self.started = threading.Event() + + async def on_startup(self) -> None: + self.events.append("on_startup") + + async def startup(self) -> None: + self.events.append("startup") + self.started.set() + await asyncio.Event().wait() + + async def on_shutdown(self) -> None: + self.events.append("on_shutdown") + await super().on_shutdown() + + +def test_service_lifespan_runs_startup_task_and_shutdown_cleanup() -> None: + service = LifecycleService() + + with TestClient(service.app) as client: + assert service.started.wait(timeout=1.0) + assert client.get("/test").json() == {"message": "test"} + + assert service.events == ["on_startup", "startup", "on_shutdown"] + assert service._startup_background_tasks + assert service._startup_background_tasks[0].cancelled() + + +@pytest.mark.asyncio +async def test_wait_for_dependencies_returns_false_when_dependency_times_out() -> None: + service = MockService() + service._dependencies = ["entities", "auth"] + + with patch.object(service, "wait_for_service_ready", new=AsyncMock(side_effect=[True, False])) as wait: + ready = await service._wait_for_dependencies(timeout=0.01) + + assert ready is False + assert [call.args[0] for call in wait.await_args_list] == ["entities", "auth"] diff --git a/packages/nmp_common/tests/nmp_common/test_dependency_provider.py b/packages/nmp_common/tests/nmp_common/test_dependency_provider.py new file mode 100644 index 0000000000..efffb0e41a --- /dev/null +++ b/packages/nmp_common/tests/nmp_common/test_dependency_provider.py @@ -0,0 +1,92 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from fastapi import FastAPI +from nmp.common.service import DependencyProvider +from nmp.common.service.dependencies import get_entity_client, get_platform_config, get_sdk_client + + +def test_get_http_client_caches_default_client() -> None: + provider = DependencyProvider() + client = MagicMock() + + with patch("nmp.common.service.base.DefaultAsyncHttpxClient", return_value=client) as factory: + first = provider.get_http_client() + second = provider.get_http_client() + + assert first is client + assert second is client + factory.assert_called_once_with() + + +def test_get_sdk_client_caches_request_sdk_and_creates_fresh_service_sdk() -> None: + provider = DependencyProvider() + request_sdk = MagicMock(name="request_sdk") + service_sdk = MagicMock(name="service_sdk") + + with patch("nmp.common.sdk_factory.get_async_platform_sdk", side_effect=[request_sdk, service_sdk]) as factory: + assert provider.get_sdk_client() is request_sdk + assert provider.get_sdk_client() is request_sdk + assert provider.get_sdk_client(as_service="jobs") is service_sdk + + assert factory.call_args_list[0].kwargs == {"http_client": None} + assert factory.call_args_list[1].kwargs == { + "as_service": "jobs", + "internal": True, + "http_client": None, + } + + +def test_setup_dependencies_registers_fastapi_overrides() -> None: + provider = DependencyProvider() + app = FastAPI() + service = MagicMock() + service._service_config = None + + provider.setup_dependencies(app, service) + + assert app.dependency_overrides[get_sdk_client] == provider.get_request_scoped_sdk + assert app.dependency_overrides[get_entity_client] == provider.get_entity_client + assert app.dependency_overrides[get_platform_config] == provider.get_platform_config + + +@pytest.mark.asyncio +async def test_close_closes_managed_clients_and_clears_references() -> None: + provider = DependencyProvider() + http_client = MagicMock() + http_client.aclose = AsyncMock() + sdk = MagicMock() + sdk.close = AsyncMock() + provider._http_client = http_client + provider._sdk_client = sdk + + await provider.close() + + http_client.aclose.assert_awaited_once_with() + sdk.close.assert_awaited_once_with() + assert provider._http_client is None + assert provider._sdk_client is None + + +def test_get_entity_client_as_service_uses_fresh_service_sdk() -> None: + provider = DependencyProvider() + sdk = MagicMock(name="service_sdk") + entities_api = MagicMock(name="entities_api") + entity_client = MagicMock(name="entity_client") + + with ( + patch.object(provider, "get_sdk_client", return_value=sdk) as get_sdk, + patch("nemo_platform.resources.entities.AsyncEntitiesResource", return_value=entities_api) as resource, + patch("nmp.common.entities.client.EntityClient", return_value=entity_client) as client_factory, + ): + result = provider.get_entity_client(as_service="models") + + assert result is entity_client + get_sdk.assert_called_once_with(as_service="models") + resource.assert_called_once_with(sdk) + client_factory.assert_called_once_with(entities_api) diff --git a/packages/nmp_common/tests/nmp_common/test_service_health.py b/packages/nmp_common/tests/nmp_common/test_service_health.py new file mode 100644 index 0000000000..773979ec7c --- /dev/null +++ b/packages/nmp_common/tests/nmp_common/test_service_health.py @@ -0,0 +1,156 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import threading +from unittest.mock import MagicMock, patch + +import httpx +import pytest +from nmp.common.config import PlatformConfig +from nmp.common.observability import MARK_INTERNAL_REQUEST_HEADERS +from nmp.common.service.api.health import ( + async_wait_for_dependencies, + async_wait_for_service_ready, + wait_for_service_ready, +) + + +@pytest.mark.asyncio +async def test_async_wait_for_service_ready_skips_service_absent_from_status() -> None: + requests: list[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response(200, json={"services": {"ready": ["entities"], "not_ready": []}}) + + async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client: + ready = await async_wait_for_service_ready( + PlatformConfig(base_url="http://platform.local"), + "models", + timeout=1.0, + poll_interval=0, + http_client=client, + ) + + assert ready is True + assert len(requests) == 1 + assert str(requests[0].url) == "http://platform.local/status" + for key, value in MARK_INTERNAL_REQUEST_HEADERS.items(): + assert requests[0].headers[key] == value + + +@pytest.mark.asyncio +async def test_async_wait_for_service_ready_waits_for_explicitly_not_ready_service() -> None: + responses = [ + httpx.Response(200, json={"services": {"ready": ["entities"], "not_ready": [{"name": "models"}]}}), + httpx.Response(200, json={"services": {"ready": ["entities", "models"], "not_ready": []}}), + ] + + def handler(_request: httpx.Request) -> httpx.Response: + return responses.pop(0) + + async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client: + ready = await async_wait_for_service_ready( + PlatformConfig(base_url="http://platform.local"), + "models", + timeout=1.0, + poll_interval=0, + http_client=client, + ) + + assert ready is True + assert responses == [] + + +@pytest.mark.asyncio +async def test_async_wait_for_service_ready_retries_intermediate_failures() -> None: + calls = 0 + + def handler(request: httpx.Request) -> httpx.Response: + nonlocal calls + calls += 1 + if calls == 1: + raise httpx.ConnectError("connection refused", request=request) + if calls == 2: + return httpx.Response(503) + if calls == 3: + return httpx.Response(200, content=b"{") + if calls == 4: + return httpx.Response(200, json=["not", "a", "mapping"]) + return httpx.Response(200, json={"services": {"ready": ["auth"], "not_ready": []}}) + + async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client: + ready = await async_wait_for_service_ready( + PlatformConfig(base_url="http://platform.local"), + "auth", + timeout=1.0, + poll_interval=0, + http_client=client, + ) + + assert ready is True + assert calls == 5 + + +@pytest.mark.asyncio +async def test_async_wait_for_dependencies_stops_after_unready_dependency() -> None: + with patch( + "nmp.common.service.api.health.async_wait_for_service_ready", + side_effect=[True, False], + ) as wait: + ready = await async_wait_for_dependencies( + PlatformConfig(base_url="http://platform.local"), + ["entities", "auth", "files"], + timeout_per_service=0.01, + poll_interval=0, + ) + + assert ready is False + assert [call.args[1] for call in wait.await_args_list] == ["entities", "auth"] + + +def test_wait_for_service_ready_skips_service_absent_from_status() -> None: + client = MagicMock() + response = MagicMock() + response.status_code = 200 + response.json.return_value = {"services": {"ready": ["entities"], "not_ready": []}} + client.get.return_value = response + + endpoint = MagicMock() + endpoint.connect_base_url = "http://platform.local" + endpoint.sync_http_client.return_value = client + + with patch("nmp.common.service.api.health.resolve_service_endpoint", return_value=endpoint): + ready = wait_for_service_ready( + PlatformConfig(base_url="http://platform.local"), + "models", + threading.Event(), + timeout=1.0, + poll_interval=0, + ) + + assert ready is True + client.get.assert_called_once_with("http://platform.local/status", headers=MARK_INTERNAL_REQUEST_HEADERS) + client.close.assert_called_once_with() + + +def test_wait_for_service_ready_returns_false_when_stop_signal_is_set() -> None: + stop_signal = threading.Event() + stop_signal.set() + client = MagicMock() + endpoint = MagicMock() + endpoint.connect_base_url = "http://platform.local" + endpoint.sync_http_client.return_value = client + + with patch("nmp.common.service.api.health.resolve_service_endpoint", return_value=endpoint): + ready = wait_for_service_ready( + PlatformConfig(base_url="http://platform.local"), + "entities", + stop_signal, + timeout=1.0, + poll_interval=0, + ) + + assert ready is False + client.get.assert_not_called() + client.close.assert_called_once_with() diff --git a/packages/nmp_common/tests/sdk_factory/test_sdk.py b/packages/nmp_common/tests/sdk_factory/test_sdk.py index efb24ba120..7c87c0631b 100644 --- a/packages/nmp_common/tests/sdk_factory/test_sdk.py +++ b/packages/nmp_common/tests/sdk_factory/test_sdk.py @@ -119,6 +119,15 @@ def test_get_platform_sdk_routes_local_service_path_to_process_listener(monkeypa assert prepared.path == "/apis/auth/v2/authz/allow" +def test_get_platform_sdk_uses_uds_endpoint_from_base_url(): + config = PlatformConfig(base_url="unix:///tmp/nemo-platform.sock") # type: ignore[abstract] + + with patch("nmp.common.sdk_factory.Configuration.get_platform_config", return_value=config): + sdk = get_platform_sdk() + + assert sdk.base_url == "http://nemo-platform.local" + + def test_get_platform_sdk_with_service_principal(): """Test get_platform_sdk with as_service parameter.""" sdk = get_platform_sdk(as_service="my-service") @@ -159,6 +168,15 @@ def test_get_async_platform_sdk(): assert str(sdk.base_url).rstrip("/") == str(expected).rstrip("/") +def test_get_async_platform_sdk_uses_uds_endpoint_from_base_url(): + config = PlatformConfig(base_url="unix:///tmp/nemo-platform.sock") # type: ignore[abstract] + + with patch("nmp.common.sdk_factory.Configuration.get_platform_config", return_value=config): + sdk = get_async_platform_sdk() + + assert str(sdk.base_url).rstrip("/") == "http://nemo-platform.local" + + def test_get_async_platform_sdk_with_service_principal(): """Test get_async_platform_sdk with as_service parameter.""" sdk = get_async_platform_sdk(as_service="async-service") @@ -590,6 +608,24 @@ def test_get_platform_sdk_routes_entities_path_to_entities_service( assert "/apis/entities/v2/workspaces" in str(prepared.path) +def test_get_platform_sdk_routes_service_path_to_env_override( + monkeypatch: pytest.MonkeyPatch, + platform_config_with_service_discovery, +): + monkeypatch.setenv("NMP_ENTITIES_URL", "http://entities-env:9090") + with patch( + "nmp.common.sdk_factory.Configuration.get_platform_config", + return_value=platform_config_with_service_discovery, + ): + sdk = get_platform_sdk() + request_url = "http://platform:8080/apis/entities/v2/workspaces" + prepared = sdk._prepare_url(request_url) + + assert prepared.host == "entities-env" + assert prepared.port == 9090 + assert prepared.scheme == "http" + + def test_get_platform_sdk_routes_jobs_path_to_jobs_service( platform_config_with_service_discovery, ): diff --git a/packages/nmp_common/tests/test_platform_endpoint.py b/packages/nmp_common/tests/test_platform_endpoint.py new file mode 100644 index 0000000000..0be4476362 --- /dev/null +++ b/packages/nmp_common/tests/test_platform_endpoint.py @@ -0,0 +1,140 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from pathlib import Path +from unittest.mock import patch + +import pytest +from nmp.common.config import PlatformConfig +from nmp.common.platform_endpoint import UDS_BASE_URL, parse_platform_endpoint, resolve_service_endpoint + + +def test_parse_tcp_endpoint() -> None: + endpoint = parse_platform_endpoint("http://127.0.0.1:8080/") + + assert endpoint.transport == "tcp" + assert endpoint.connect_base_url == "http://127.0.0.1:8080" + assert endpoint.socket_path is None + + +@pytest.mark.parametrize("url", ["http://", "https://"]) +def test_parse_rejects_hostless_http_endpoint(url: str) -> None: + with pytest.raises(ValueError, match="must include a host"): + parse_platform_endpoint(url) + + +def test_parse_https_endpoint_preserves_normalized_connect_base_url() -> None: + endpoint = parse_platform_endpoint("https://platform.example.com/api/") + + assert endpoint.transport == "tcp" + assert endpoint.connect_base_url == "https://platform.example.com/api" + assert endpoint.socket_path is None + + +def test_parse_uds_endpoint() -> None: + endpoint = parse_platform_endpoint("unix:///tmp/nemo-platform.sock") + + assert endpoint.transport == "uds" + assert endpoint.connect_base_url == UDS_BASE_URL + assert endpoint.socket_path == Path("/tmp/nemo-platform.sock") + + +def test_parse_rejects_raw_socket_path() -> None: + with pytest.raises(ValueError, match="use unix:///tmp/nemo-platform.sock"): + parse_platform_endpoint("/tmp/nemo-platform.sock") + + +def test_parse_rejects_relative_uds_socket_path() -> None: + with pytest.raises(ValueError, match="absolute socket path"): + parse_platform_endpoint("unix://relative.sock") + + +def test_resolve_service_endpoint_uses_service_specific_url(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("NMP_SECRETS_URL", "unix:///tmp/secrets.sock") + config = PlatformConfig(base_url="http://platform:8080") + + endpoint = resolve_service_endpoint("secrets", config) + + assert endpoint.transport == "uds" + assert endpoint.socket_path == Path("/tmp/secrets.sock") + + +def test_resolve_service_endpoint_falls_back_to_base_url(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("NMP_SECRETS_URL", raising=False) + config = PlatformConfig(base_url="http://platform:8080") + + endpoint = resolve_service_endpoint("secrets", config) + + assert endpoint.transport == "tcp" + assert endpoint.connect_base_url == "http://platform:8080" + + +def test_endpoint_env_family_is_not_part_of_contract(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("NMP_SECRETS_ENDPOINT", "unix:///tmp/secrets.sock") + monkeypatch.delenv("NMP_SECRETS_URL", raising=False) + config = PlatformConfig(base_url="http://platform:8080") + + endpoint = resolve_service_endpoint("secrets", config) + + assert endpoint.transport == "tcp" + assert endpoint.connect_base_url == "http://platform:8080" + + +def test_sync_http_client_omits_unset_timeout() -> None: + endpoint = parse_platform_endpoint("http://127.0.0.1:8080") + + with patch("nmp.common.platform_endpoint.httpx.Client") as client: + endpoint.sync_http_client() + + client.assert_called_once_with(follow_redirects=True) + + +def test_sync_http_client_passes_explicit_timeout() -> None: + endpoint = parse_platform_endpoint("http://127.0.0.1:8080") + + with patch("nmp.common.platform_endpoint.httpx.Client") as client: + endpoint.sync_http_client(timeout=2.0) + + client.assert_called_once_with(follow_redirects=True, timeout=2.0) + + +def test_uds_sync_http_client_keeps_transport_and_omits_unset_timeout() -> None: + endpoint = parse_platform_endpoint("unix:///tmp/nemo-platform.sock") + + with patch("nmp.common.platform_endpoint.httpx.Client") as client: + endpoint.sync_http_client() + + kwargs = client.call_args.kwargs + assert kwargs["follow_redirects"] is True + assert "transport" in kwargs + assert "timeout" not in kwargs + + +def test_async_http_client_omits_unset_timeout() -> None: + endpoint = parse_platform_endpoint("http://127.0.0.1:8080") + + with patch("nmp.common.platform_endpoint.httpx.AsyncClient") as client: + endpoint.async_http_client() + + client.assert_called_once_with(follow_redirects=True) + + +def test_async_http_client_passes_explicit_timeout() -> None: + endpoint = parse_platform_endpoint("http://127.0.0.1:8080") + + with patch("nmp.common.platform_endpoint.httpx.AsyncClient") as client: + endpoint.async_http_client(timeout=2.0) + + client.assert_called_once_with(follow_redirects=True, timeout=2.0) + + +def test_uds_async_http_client_keeps_transport_and_omits_unset_timeout() -> None: + endpoint = parse_platform_endpoint("unix:///tmp/nemo-platform.sock") + + with patch("nmp.common.platform_endpoint.httpx.AsyncClient") as client: + endpoint.async_http_client() + + kwargs = client.call_args.kwargs + assert kwargs["follow_redirects"] is True + assert "transport" in kwargs + assert "timeout" not in kwargs diff --git a/packages/nmp_platform_runner/src/nmp/platform_runner/config.py b/packages/nmp_platform_runner/src/nmp/platform_runner/config.py index 977f9bba45..06c9acc189 100644 --- a/packages/nmp_platform_runner/src/nmp/platform_runner/config.py +++ b/packages/nmp_platform_runner/src/nmp/platform_runner/config.py @@ -6,9 +6,11 @@ from __future__ import annotations import os -from collections.abc import MutableMapping +import re +from collections.abc import MutableMapping, Sequence from dataclasses import dataclass, field from importlib.resources import files +from pathlib import Path from urllib.parse import urlparse from nmp.common.config import ( @@ -22,6 +24,7 @@ from nmp.platform_runner.loader import ControllerRunFunc from nmp.platform_runner.registry import ( AVAILABLE_SIDECARS, + SERVICE_SIDECAR_DEPENDENCIES, get_available_controllers, get_available_services, get_controller_groups, @@ -29,6 +32,103 @@ get_service_groups, ) +DEFAULT_SCOPE = "default" +DEFAULT_PLATFORM_BIND_HOST = "0.0.0.0" +DEFAULT_LOCAL_SERVICES_BIND_HOST = "127.0.0.1" + +_SCOPE_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,127}$") +_INSTANCES_DIRNAME = "instances" +_SOCKET_FILENAME = "nemo-platform.sock" +_LOG_FILENAME = "services.log" + + +@dataclass +class PlatformAppConfig: + """Service selection and listener binding for a platform app/server.""" + + services: Sequence[str] | None = None + service_group: str | None = None + controllers: Sequence[str] | None = None + controller_group: str | None = None + sidecars: Sequence[str] | None = None + config_path: str | None = None + scope: str = DEFAULT_SCOPE + host: str = DEFAULT_PLATFORM_BIND_HOST + port: int = 8080 + socket_path: str | Path | None = None + state_root: str | Path | None = None + runtime_root: str | Path | None = None + log_path: str | Path | None = None + + def __post_init__(self) -> None: + self.scope = validate_scope(self.scope) + self.socket_path = _resolve_socket_path(self.socket_path) + self.state_root = _resolve_absolute_path(self.state_root, "state root") + self.runtime_root = _resolve_absolute_path(self.runtime_root, "runtime root") + self.log_path = _resolve_absolute_path(self.log_path, "log path") + + @property + def state_root_path(self) -> Path: + return Path(self.state_root) if self.state_root is not None else default_state_root() + + @property + def runtime_root_path(self) -> Path: + return Path(self.runtime_root) if self.runtime_root is not None else default_runtime_root() + + def state_dir(self, *, create: bool = False) -> Path: + path = self.state_root_path / _INSTANCES_DIRNAME / self.scope + if create: + path.mkdir(parents=True, exist_ok=True) + return path + + def runtime_dir(self, *, create: bool = False) -> Path: + if self.runtime_root is None and self.socket_path is not None: + path = Path(self.socket_path).parent + else: + path = self.runtime_root_path / self.scope + if create: + path.mkdir(parents=True, exist_ok=True) + return path + + def socket_file_path(self) -> Path: + if self.socket_path is not None: + return Path(self.socket_path) + return self.runtime_dir() / _SOCKET_FILENAME + + def log_file_path(self, *, create_parent: bool = False) -> Path: + path = Path(self.log_path) if self.log_path is not None else self.state_dir() / _LOG_FILENAME + if create_parent: + path.parent.mkdir(parents=True, exist_ok=True) + return path + + +def validate_scope(scope: str) -> str: + """Ensure *scope* is safe to use for local state and socket paths.""" + if not _SCOPE_RE.fullmatch(scope): + raise ValueError(f"Invalid scope: {scope!r}") + return scope + + +def default_state_root() -> Path: + """Return the local services state root.""" + xdg = os.environ.get("XDG_STATE_HOME") + if xdg: + return Path(xdg) / "nmp" + return Path.home() / ".local" / "state" / "nmp" + + +def default_runtime_root() -> Path: + """Return the local services runtime root for sockets and volatile metadata.""" + return default_state_root() / "run" + + +def _sidecars_for_services(service_names: set[str]) -> set[str]: + selected: set[str] = set() + for service_name in service_names: + selected.update(SERVICE_SIDECAR_DEPENDENCIES.get(service_name, set())) + return selected + + _IPV4_LOOPBACK = "127.0.0.1" _IPV6_LOOPBACK = "::1" _IPV4_WILDCARDS = frozenset({"0.0.0.0"}) @@ -43,6 +143,7 @@ class ResolvedRunConfiguration: host: str port: int config_path: str + socket_path: str | None = None available_services: dict[str, str | Service] = field(default_factory=dict) available_controllers: dict[str, str | ControllerRunFunc] = field(default_factory=dict) @@ -56,22 +157,15 @@ def default_config_path() -> str: def resolve_run_configuration( - *, - services: list[str] | None = None, - service_group: str | None = None, - controllers: list[str] | None = None, - controller_group: str | None = None, - sidecars: list[str] | None = None, - config_path: str | None = None, - host: str = "0.0.0.0", - port: int = 8080, + config: PlatformAppConfig | None = None, ) -> ResolvedRunConfiguration: - """Resolve and validate platform run arguments. + """Resolve and validate platform run configuration. Group selectors are convenience shortcuts for callers that are not also naming specific services or controllers. Mixing the two is ambiguous, so these combinations fail fast instead of silently ignoring the group. """ + config = config or PlatformAppConfig() available_services = get_available_services() available_controllers = get_available_controllers() available_sidecars = AVAILABLE_SIDECARS @@ -79,30 +173,30 @@ def resolve_run_configuration( controller_groups = get_controller_groups(available_controllers) default_controllers = set(get_default_controllers(controller_groups)) - selected_services = set(services or []) - selected_controllers = set(controllers or []) - selected_sidecars = set(sidecars or []) + selected_services = set(config.services or []) + selected_controllers = set(config.controllers or []) + selected_sidecars = set(config.sidecars or []) # Explicit selections and group selectors are mutually exclusive. The old # entrypoint rejected these combinations, and keeping that behavior avoids a # confusing silent-ignore UX for callers. - if service_group and selected_services: + if config.service_group and selected_services: raise ValueError("--services cannot be combined with --service-group") - if controller_group and selected_controllers: + if config.controller_group and selected_controllers: raise ValueError("--controllers cannot be combined with --controller-group") - if service_group and not selected_services: - if service_group not in service_groups: + if config.service_group and not selected_services: + if config.service_group not in service_groups: valid_groups = ", ".join(sorted(service_groups)) - raise ValueError(f"Unknown service group: {service_group}. Available groups: {valid_groups}") - selected_services.update(service_groups[service_group]) + raise ValueError(f"Unknown service group: {config.service_group}. Available groups: {valid_groups}") + selected_services.update(service_groups[config.service_group]) - if controller_group and not selected_controllers: - if controller_group not in controller_groups: + if config.controller_group and not selected_controllers: + if config.controller_group not in controller_groups: valid_groups = ", ".join(sorted(controller_groups)) - raise ValueError(f"Unknown controller group: {controller_group}. Available groups: {valid_groups}") - selected_controllers.update(controller_groups[controller_group]) + raise ValueError(f"Unknown controller group: {config.controller_group}. Available groups: {valid_groups}") + selected_controllers.update(controller_groups[config.controller_group]) invalid_services = selected_services - set(available_services) if invalid_services: @@ -116,30 +210,53 @@ def resolve_run_configuration( requested = ", ".join(sorted(invalid_controllers)) raise ValueError(f"Unknown controllers: {requested}. Available controllers: {available}") + if not selected_services and not selected_controllers and not selected_sidecars: + # No explicit selection means "run the platform": start the default + # service group plus the default controller set. + selected_services.update(service_groups["all"]) + selected_controllers.update(default_controllers) + + selected_sidecars.update(_sidecars_for_services(selected_services)) + invalid_sidecars = selected_sidecars - set(available_sidecars) if invalid_sidecars: available = ", ".join(sorted(available_sidecars)) requested = ", ".join(sorted(invalid_sidecars)) raise ValueError(f"Unknown sidecars: {requested}. Available sidecars: {available}") - if not selected_services and not selected_controllers and not selected_sidecars: - # No explicit selection means "run the platform": start the default - # service group plus the default controller set. - selected_services.update(service_groups["all"]) - selected_controllers.update(default_controllers) + resolved_socket_path = _resolve_socket_path(config.socket_path) return ResolvedRunConfiguration( services=selected_services, controllers=selected_controllers, sidecars=selected_sidecars, - host=host, - port=port, - config_path=config_path or default_config_path(), + host=config.host, + port=config.port, + config_path=config.config_path or default_config_path(), + socket_path=resolved_socket_path, available_services=available_services, available_controllers=available_controllers, ) +def _resolve_socket_path(socket_path: str | Path | None) -> str | None: + if socket_path is None: + return None + path = Path(socket_path).expanduser() + if not path.is_absolute(): + raise ValueError(f"UDS socket path must be absolute: {socket_path}") + return str(path) + + +def _resolve_absolute_path(path_value: str | Path | None, label: str) -> str | None: + if path_value is None: + return None + path = Path(path_value).expanduser() + if not path.is_absolute(): + raise ValueError(f"{label} must be absolute: {path_value}") + return str(path) + + def apply_run_environment( config: ResolvedRunConfiguration, env: MutableMapping[str, str] | None = None, @@ -181,14 +298,17 @@ def apply_run_environment( connect_host = _connect_host_for_internal_clients(config.host) effective_host = env.setdefault("NMP_SERVICE_HOST", connect_host) effective_port = env.setdefault("NMP_SERVICE_PORT", str(config.port)) - config_base_url_parts = _config_file_base_url_parts(config.config_path) - if config_base_url_parts is not None: - scheme, config_host = config_base_url_parts - host_for_url = _bracket_ipv6(_connect_host_for_internal_clients(config_host)) - default_base_url = f"{scheme}://{host_for_url}:{effective_port}" + if config.socket_path: + default_base_url = f"unix://{config.socket_path}" else: - host_for_url = _bracket_ipv6(effective_host) - default_base_url = f"http://{host_for_url}:{effective_port}" + config_base_url_parts = _config_file_base_url_parts(config.config_path) + if config_base_url_parts is not None: + scheme, config_host = config_base_url_parts + host_for_url = _bracket_ipv6(_connect_host_for_internal_clients(config_host)) + default_base_url = f"{scheme}://{host_for_url}:{effective_port}" + else: + host_for_url = _bracket_ipv6(effective_host) + default_base_url = f"http://{host_for_url}:{effective_port}" base_url = env.setdefault("NMP_BASE_URL", default_base_url) # Embedded PDP is usually served from the same platform process, so its # self-call origin must stay aligned with the resolved base URL. Deployed diff --git a/packages/nmp_platform_runner/src/nmp/platform_runner/registry.py b/packages/nmp_platform_runner/src/nmp/platform_runner/registry.py index bb4b0f4100..ab8bcefa45 100644 --- a/packages/nmp_platform_runner/src/nmp/platform_runner/registry.py +++ b/packages/nmp_platform_runner/src/nmp/platform_runner/registry.py @@ -39,6 +39,10 @@ "adapters": "nmp.core.models.sidecars.adapters.main:run", } +SERVICE_SIDECAR_DEPENDENCIES: dict[str, set[str]] = { + "models": {"adapters"}, +} + CORE_SERVICES = [ "auth", "models", diff --git a/packages/nmp_platform_runner/src/nmp/platform_runner/run.py b/packages/nmp_platform_runner/src/nmp/platform_runner/run.py index 0dbad175ed..d3a54eac98 100644 --- a/packages/nmp_platform_runner/src/nmp/platform_runner/run.py +++ b/packages/nmp_platform_runner/src/nmp/platform_runner/run.py @@ -17,7 +17,11 @@ from nmp.common.observability import initialize_obs, setup_global_instrumentations from nmp.common.observability.otel import settings as otel_settings from nmp.common.service import CircularDependencyError, Service -from nmp.platform_runner.config import apply_run_environment, resolve_run_configuration +from nmp.platform_runner.config import ( + PlatformAppConfig, + apply_run_environment, + resolve_run_configuration, +) from nmp.platform_runner.health import get_platform_resource_attributes from nmp.platform_runner.loader import ( ControllerRunFunc, @@ -72,15 +76,8 @@ def run_controllers_in_threads( def run_platform( + config: PlatformAppConfig | None = None, *, - services: list[str] | None = None, - service_group: str | None = None, - controllers: list[str] | None = None, - controller_group: str | None = None, - sidecars: list[str] | None = None, - config_path: str | None = None, - host: str = "0.0.0.0", - port: int = 8080, reload_app_factory: str | None = None, on_shutdown: Callable[[], object] | None = None, ) -> None: @@ -88,16 +85,7 @@ def run_platform( t_total = time.perf_counter() t0 = time.perf_counter() - resolved = resolve_run_configuration( - services=services, - service_group=service_group, - controllers=controllers, - controller_group=controller_group, - sidecars=sidecars, - config_path=config_path, - host=host, - port=port, - ) + resolved = resolve_run_configuration(config) apply_run_environment(resolved) _startup_phase("resolve_config", t0) @@ -163,7 +151,7 @@ def signal_handler(signum: int, _frame: object) -> None: controller_threads.extend(run_controllers_in_threads(controller_run_funcs, controller_stop_signal)) if sidecar_run_funcs: controller_threads.extend(run_controllers_in_threads(sidecar_run_funcs, controller_stop_signal)) - run_server(service_instances, host=resolved.host, port=resolved.port) + run_server(service_instances, host=resolved.host, port=resolved.port, socket_path=resolved.socket_path) except ValueError as error: logger.error("Configuration error: %s", error) raise SystemExit(1) from error diff --git a/packages/nmp_platform_runner/src/nmp/platform_runner/server.py b/packages/nmp_platform_runner/src/nmp/platform_runner/server.py index daac4f6182..eef0e55d2c 100644 --- a/packages/nmp_platform_runner/src/nmp/platform_runner/server.py +++ b/packages/nmp_platform_runner/src/nmp/platform_runner/server.py @@ -10,6 +10,7 @@ import logging import os import threading +from collections.abc import Callable, Mapping, MutableMapping from contextlib import asynccontextmanager from typing import cast @@ -25,6 +26,7 @@ from nmp.common.observability.context import create_app_context_dependency from nmp.common.pyleak import detect_blocking from nmp.common.service import Service +from nmp.platform_runner.config import PlatformAppConfig from nmp.platform_runner.health import ReadinessCheck, create_platform_health_router, get_platform_resource_attributes from nmp.platform_runner.loader import ( ControllerRunFunc, @@ -32,7 +34,12 @@ load_service, order_services_by_dependencies, ) -from nmp.platform_runner.registry import get_available_controllers, get_available_services, get_openapi_service_names +from nmp.platform_runner.registry import ( + AVAILABLE_SIDECARS, + get_available_controllers, + get_available_services, + get_openapi_service_names, +) from starlette.middleware.base import BaseHTTPMiddleware from starlette.requests import Request from starlette.responses import RedirectResponse, Response @@ -272,12 +279,85 @@ async def root_handler() -> Response: return app -def run_server(services: list[Service] | None = None, host: str = "0.0.0.0", port: int = 8080) -> None: +def _load_run_functions( + names: list[str], + registry: Mapping[str, str | Callable[[threading.Event], object]], +) -> dict[str, Callable[[threading.Event], object]]: + run_funcs: dict[str, Callable[[threading.Event], object]] = {} + for name in names: + value = registry[name] + if isinstance(value, str): + run_funcs[name] = load_controller_run_func(name, value) + else: + run_funcs[name] = value + return run_funcs + + +def build_platform_app( + config: PlatformAppConfig | None = None, + *, + http_client: httpx.AsyncClient | None = None, + env: MutableMapping[str, str] | None = None, +) -> FastAPI: + """Build a platform FastAPI app without starting uvicorn. + + Args: + config: App-build selection and bind configuration. Prefer this over + individual service/controller/sidecar keyword arguments for new + callers. + env: Environment mapping passed to :func:`apply_run_environment`. + Defaults to ``None`` which writes to ``os.environ``. Tests can + pass an empty dict to avoid polluting the process environment. + """ + from nmp.platform_runner.config import apply_run_environment, resolve_run_configuration + + resolved = resolve_run_configuration(config) + apply_run_environment(resolved, env=env) + + service_instances = [] + for service_name in sorted(resolved.services): + service_value = resolved.available_services[service_name] + service_instances.append( + service_value if isinstance(service_value, Service) else load_service(service_name, service_value) + ) + service_instances = order_services_by_dependencies(service_instances) + + collisions = resolved.controllers & resolved.sidecars + if collisions: + raise ValueError(f"Controller/sidecar name collision: {', '.join(sorted(collisions))}") + + controller_run_funcs = _load_run_functions(sorted(resolved.controllers), resolved.available_controllers) + sidecar_run_funcs = _load_run_functions(sorted(resolved.sidecars), AVAILABLE_SIDECARS) + controller_run_funcs.update(sidecar_run_funcs) + + return create_app(service_instances, controller_run_funcs=controller_run_funcs, http_client=http_client) + + +def run_server( + services: list[Service] | None = None, + host: str = "0.0.0.0", + port: int = 8080, + socket_path: str | None = None, +) -> None: """Run the platform API server.""" preflight_embedded_auth_policy_wasm(get_auth_config()) app = create_app(services or []) setup_fastapi_instrumentations(app) - uvicorn.run(app, host=host, port=port, log_config=None) + if socket_path: + _run_server_on_bound_sockets(app, host=host, port=port, socket_path=socket_path) + else: + uvicorn.run(app, host=host, port=port, log_config=None) + + +def _run_server_on_bound_sockets(app: FastAPI, *, host: str, port: int, socket_path: str) -> None: + tcp_config = uvicorn.Config(app, host=host, port=port, log_config=None) + uds_config = uvicorn.Config(app, uds=socket_path, log_config=None) + sockets = [tcp_config.bind_socket(), uds_config.bind_socket()] + try: + asyncio.run(uvicorn.Server(tcp_config).serve(sockets=sockets)) + finally: + for sock in sockets: + sock.close() def run_server_with_reload(app_factory: str, host: str = "0.0.0.0", port: int = 8080) -> None: diff --git a/packages/nmp_platform_runner/tests/test_config.py b/packages/nmp_platform_runner/tests/test_config.py index 2318957569..738ae21fb8 100644 --- a/packages/nmp_platform_runner/tests/test_config.py +++ b/packages/nmp_platform_runner/tests/test_config.py @@ -6,6 +6,8 @@ import pytest from nmp.platform_runner import registry from nmp.platform_runner.config import ( + DEFAULT_PLATFORM_BIND_HOST, + PlatformAppConfig, ResolvedRunConfiguration, apply_run_environment, default_config_path, @@ -26,6 +28,7 @@ def _make_config( sidecars: set[str] | None = None, host: str = "0.0.0.0", port: int = 8080, + socket_path: str | None = None, config_path: str = "/nonexistent/nmp-test-config.yaml", ) -> ResolvedRunConfiguration: return ResolvedRunConfiguration( @@ -34,14 +37,13 @@ def _make_config( sidecars=sidecars if sidecars is not None else set(), host=host, port=port, + socket_path=socket_path, config_path=config_path, ) def resolve(**kwargs): - params = {} - params.update(kwargs) - return resolve_run_configuration(**params) + return resolve_run_configuration(PlatformAppConfig(**kwargs)) def test_default_config_path_points_to_bundled_local_config(): @@ -50,9 +52,80 @@ def test_default_config_path_points_to_bundled_local_config(): assert path.endswith(("nmp/platform_runner/config/local.yaml", "nemo_platform/services/runner/config/local.yaml")) +def test_platform_app_config_keeps_sequence_fields_simple(): + config = PlatformAppConfig( + services=["models"], + controllers=[], + sidecars=["adapters"], + ) + + assert config.services == ["models"] + assert config.controllers == [] + assert config.sidecars == ["adapters"] + + +def test_platform_app_config_derives_instance_paths_from_roots(tmp_path: Path): + config = PlatformAppConfig(scope="dev", state_root=tmp_path / "state", runtime_root=tmp_path / "run") + + assert config.state_dir() == tmp_path / "state" / "instances" / "dev" + assert config.runtime_dir() == tmp_path / "run" / "dev" + assert config.socket_file_path() == tmp_path / "run" / "dev" / "nemo-platform.sock" + assert config.log_file_path() == tmp_path / "state" / "instances" / "dev" / "services.log" + + +def test_platform_app_config_runtime_dir_defaults_to_explicit_socket_parent(tmp_path: Path): + config = PlatformAppConfig(socket_path=tmp_path / "custom.sock") + + assert config.runtime_dir() == tmp_path + assert config.socket_file_path() == tmp_path / "custom.sock" + + +def test_platform_app_config_uses_explicit_log_path(tmp_path: Path): + config = PlatformAppConfig(state_root=tmp_path / "state", log_path=tmp_path / "logs" / "nemo.log") + + assert config.log_file_path() == tmp_path / "logs" / "nemo.log" + + +def test_platform_app_config_rejects_relative_socket_path(): + with pytest.raises(ValueError, match="UDS socket path must be absolute"): + PlatformAppConfig(socket_path="relative/path") + + +def test_platform_app_config_rejects_relative_state_root(): + with pytest.raises(ValueError, match="state root must be absolute"): + PlatformAppConfig(state_root="relative/path") + + +def test_platform_app_config_rejects_relative_runtime_root(): + with pytest.raises(ValueError, match="runtime root must be absolute"): + PlatformAppConfig(runtime_root="relative/path") + + +def test_platform_app_config_rejects_relative_log_path(): + with pytest.raises(ValueError, match="log path must be absolute"): + PlatformAppConfig(log_path="relative/path") + + +def test_resolve_run_configuration_accepts_platform_app_config(): + resolved = resolve_run_configuration( + PlatformAppConfig( + services=["auth"], + controllers=[], + host="127.0.0.1", + port=9090, + ) + ) + + assert resolved.services == {"auth"} + assert resolved.controllers == set() + assert resolved.host == "127.0.0.1" + assert resolved.port == 9090 + + def test_no_arguments_defaults_to_all_services_and_default_controllers(): resolved = resolve() + assert resolved.host == DEFAULT_PLATFORM_BIND_HOST assert resolved.services.issuperset( { "auth", @@ -112,6 +185,17 @@ def test_extra_services_are_available_for_resolution(): resolve(services=["custom-service"]) +def test_resolve_rejects_relative_socket_path(): + with pytest.raises(ValueError, match="UDS socket path must be absolute"): + resolve(socket_path="relative.sock") + + +def test_resolve_preserves_absolute_socket_path(): + resolved = resolve(socket_path="/tmp/nemo-platform.sock") + + assert resolved.socket_path == "/tmp/nemo-platform.sock" + + # --------------------------------------------------------------------------- # Topology regression tests for apply_run_environment # @@ -146,6 +230,11 @@ def test_config_file_gateway_base_url_seeds_base_url(self, tmp_path: Path): assert env["NMP_BASE_URL"] == "https://nemo-gateway:8080" assert env["NMP_SERVICE_HOST"] == "127.0.0.1" + def test_sets_uds_base_url_when_socket_path_is_present(self): + env: dict[str, str] = {} + apply_run_environment(_make_config(socket_path="/tmp/nemo-platform.sock"), env=env) + assert env["NMP_BASE_URL"] == "unix:///tmp/nemo-platform.sock" + def test_sets_embedded_pdp_base_url_from_base_url(self): env: dict[str, str] = {} apply_run_environment(_make_config(host="0.0.0.0", port=9090), env=env) diff --git a/packages/nmp_platform_runner/tests/test_health.py b/packages/nmp_platform_runner/tests/test_health.py new file mode 100644 index 0000000000..f83cc964e5 --- /dev/null +++ b/packages/nmp_platform_runner/tests/test_health.py @@ -0,0 +1,103 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient +from nmp.common.controller.controller_manager import ControllerManager +from nmp.common.service import RouterConfig, Service +from nmp.platform_runner.health import create_platform_health_router + + +class ProbeService(Service): + def __init__(self, name: str, *, ready: bool = True) -> None: + super().__init__(name=name, module_name=f"nmp.{name}") + self.ready = ready + + def get_routers(self) -> list[RouterConfig]: + return [] + + async def is_ready(self) -> bool: + return self.ready + + +@pytest.fixture(autouse=True) +def reset_controller_manager() -> None: + ControllerManager._instance = None + yield + ControllerManager._instance = None + + +def _client_for(services: list[Service]) -> TestClient: + app = FastAPI() + app.include_router(create_platform_health_router(services)) + return TestClient(app) + + +def test_status_and_ready_are_healthy_when_no_services_are_running() -> None: + client = _client_for([]) + + status_response = client.get("/status") + ready_response = client.get("/health/ready") + + assert status_response.status_code == 200 + assert status_response.json() == { + "status": "healthy", + "services": {"ready": [], "not_ready": []}, + "controllers": {"healthy": True, "status": {}}, + } + assert ready_response.status_code == 200 + assert ready_response.json() == {"status": "ready"} + + +def test_status_only_reports_services_registered_with_runner() -> None: + registered = ProbeService("entities", ready=True) + not_started = ProbeService("models", ready=False) + client = _client_for([registered]) + + response = client.get("/status") + + assert response.status_code == 200 + payload = response.json() + assert payload["status"] == "healthy" + assert payload["services"] == {"ready": ["entities"], "not_ready": []} + assert not_started.name not in payload["services"]["ready"] + assert not_started.name not in [service["name"] for service in payload["services"]["not_ready"]] + assert client.get("/health/ready").status_code == 200 + + +def test_status_remains_healthy_when_new_service_is_registered_after_it_is_ready() -> None: + entities = ProbeService("entities", ready=True) + models = ProbeService("models", ready=False) + services: list[Service] = [entities] + client = _client_for(services) + + assert client.get("/status").json()["status"] == "healthy" + + models.ready = True + services.append(models) + + response = client.get("/status") + + assert response.status_code == 200 + assert response.json()["status"] == "healthy" + assert response.json()["services"] == {"ready": ["entities", "models"], "not_ready": []} + assert client.get("/health/ready").status_code == 200 + + +def test_registered_not_ready_service_degrades_status_and_blocks_readiness() -> None: + entities = ProbeService("entities", ready=True) + models = ProbeService("models", ready=False) + client = _client_for([entities, models]) + + status_response = client.get("/status") + ready_response = client.get("/health/ready") + + assert status_response.status_code == 200 + assert status_response.json()["status"] == "degraded" + assert status_response.json()["services"] == { + "ready": ["entities"], + "not_ready": [{"name": "models", "message": ""}], + } + assert ready_response.status_code == 503 + assert ready_response.json() == {"detail": {"status": "not_ready"}} diff --git a/packages/nmp_platform_runner/tests/test_run.py b/packages/nmp_platform_runner/tests/test_run.py index f5cb9d4617..def553aadc 100644 --- a/packages/nmp_platform_runner/tests/test_run.py +++ b/packages/nmp_platform_runner/tests/test_run.py @@ -65,7 +65,7 @@ def test_run_platform_marks_loaded_services_local_before_starting_controllers(mo ) services = [_StubService("jobs"), _StubService("entities")] - monkeypatch.setattr(runner, "resolve_run_configuration", lambda **_: resolved) + monkeypatch.setattr(runner, "resolve_run_configuration", lambda *_args, **_kwargs: resolved) monkeypatch.setattr(runner, "apply_run_environment", lambda config: None) monkeypatch.setattr(runner, "initialize_obs", lambda *, resource_attributes: None) monkeypatch.setattr(runner, "setup_global_instrumentations", lambda: None) @@ -76,7 +76,7 @@ def test_run_platform_marks_loaded_services_local_before_starting_controllers(mo lambda names, registry, kind: {"jobs": lambda stop_signal: None} if kind == "controller" else {}, ) monkeypatch.setattr(runner, "_display_banner", lambda **_: None) - monkeypatch.setattr(runner, "run_server", lambda services, host, port: None) + monkeypatch.setattr(runner, "run_server", lambda services, host, port, socket_path=None: None) monkeypatch.setattr(runner.signal, "signal", lambda *args: None) def capture_controller_start( diff --git a/packages/nmp_platform_runner/tests/test_server.py b/packages/nmp_platform_runner/tests/test_server.py index f752d1c51a..2087de75f1 100644 --- a/packages/nmp_platform_runner/tests/test_server.py +++ b/packages/nmp_platform_runner/tests/test_server.py @@ -3,6 +3,7 @@ import asyncio import builtins +import os import sys import threading import time @@ -15,9 +16,33 @@ from nmp.common.config import AuthConfig, Configuration from nmp.common.config.base import OIDCConfig from nmp.common.service import Service +from nmp.platform_runner import config as runner_config from nmp.platform_runner import server from nmp.platform_runner.health import ReadinessCheck, create_platform_health_router +_RUN_ENV_KEYS = ( + "NMP_CONFIG_FILE_PATH", + "NMP_SERVICE_HOST", + "NMP_SERVICE_PORT", + "NMP_BASE_URL", + "NMP_AUTH_POLICY_DECISION_POINT_BASE_URL", + "NMP_SERVICES", + "NMP_CONTROLLERS", + "NMP_SIDECARS", +) + + +@pytest.fixture(autouse=True) +def restore_platform_runner_env(): + original_env = {key: os.environ.get(key) for key in _RUN_ENV_KEYS} + yield + for key, value in original_env.items(): + if value is None: + os.environ.pop(key, None) + else: + os.environ[key] = value + Configuration.clear_cache() + def _make_auth_config(*, enabled: bool) -> AuthConfig: return AuthConfig( @@ -180,6 +205,59 @@ def test_create_app_mounted_services_drive_sdk_local_routing_without_services_en Configuration.clear_cache() +def test_build_platform_app_returns_app_without_running_uvicorn(monkeypatch): + plugin_service = PluginService() + captured: dict[str, object] = {} + + monkeypatch.setattr(runner_config, "get_available_services", lambda: {"agents": plugin_service}) + monkeypatch.setattr(runner_config, "get_available_controllers", lambda: {}) + monkeypatch.setattr(runner_config, "get_controller_groups", lambda _controllers: {"all": [], "core": []}) + monkeypatch.setattr(server, "order_services_by_dependencies", lambda services: services) + + def fake_create_app(services, controller_run_funcs=None, http_client=None): + captured["services"] = services + captured["controller_run_funcs"] = controller_run_funcs + captured["http_client"] = http_client + return FastAPI() + + monkeypatch.setattr(server, "create_app", fake_create_app) + + app = server.build_platform_app(runner_config.PlatformAppConfig(services=["agents"], controllers=[]), env={}) + + assert isinstance(app, FastAPI) + assert captured["services"] == [plugin_service] + assert captured["controller_run_funcs"] == {} + assert captured["http_client"] is None + + +def test_build_platform_app_accepts_platform_app_config(monkeypatch): + plugin_service = PluginService() + captured: dict[str, object] = {} + + monkeypatch.setattr(runner_config, "get_available_services", lambda: {"agents": plugin_service}) + monkeypatch.setattr(runner_config, "get_available_controllers", lambda: {}) + monkeypatch.setattr(runner_config, "get_controller_groups", lambda _controllers: {"all": [], "core": []}) + monkeypatch.setattr(server, "order_services_by_dependencies", lambda services: services) + + def fake_create_app(services, controller_run_funcs=None, http_client=None): + captured["services"] = services + captured["controller_run_funcs"] = controller_run_funcs + captured["http_client"] = http_client + return FastAPI() + + monkeypatch.setattr(server, "create_app", fake_create_app) + + app = server.build_platform_app( + config=runner_config.PlatformAppConfig(services=("agents",), controllers=()), + env={}, + ) + + assert isinstance(app, FastAPI) + assert captured["services"] == [plugin_service] + assert captured["controller_run_funcs"] == {} + assert captured["http_client"] is None + + def test_embedded_auth_preflight_invokes_policy_wasm_helper(monkeypatch): calls: list[bool] = [] auth_cfg = AuthConfig( @@ -235,6 +313,25 @@ def test_run_server_runs_embedded_auth_preflight(): uvicorn_run.assert_called_once() +def test_run_server_can_bind_tcp_and_unix_domain_socket(): + auth_cfg = _make_auth_config(enabled=True) + with ( + patch("nmp.platform_runner.server.get_auth_config", return_value=auth_cfg), + patch("nmp.platform_runner.server.preflight_embedded_auth_policy_wasm"), + patch("nmp.platform_runner.server.create_app", return_value=FastAPI()), + patch("nmp.platform_runner.server.setup_fastapi_instrumentations"), + patch("nmp.platform_runner.server._run_server_on_bound_sockets") as run_bound_sockets, + ): + server.run_server(services=[], host="127.0.0.1", port=9999, socket_path="/tmp/nemo-platform.sock") + + run_bound_sockets.assert_called_once() + assert run_bound_sockets.call_args.kwargs == { + "host": "127.0.0.1", + "port": 9999, + "socket_path": "/tmp/nemo-platform.sock", + } + + def test_create_default_app_raises_for_unknown_service_from_env(monkeypatch): monkeypatch.setattr(server, "_obs_initialized", True) monkeypatch.setenv("NMP_SERVICES", "missing-service") diff --git a/packages/nmp_platform_runner/tests/test_sidecars.py b/packages/nmp_platform_runner/tests/test_sidecars.py new file mode 100644 index 0000000000..e3950b3a93 --- /dev/null +++ b/packages/nmp_platform_runner/tests/test_sidecars.py @@ -0,0 +1,213 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import threading +from collections.abc import Callable +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest +from nmp.common.config import AuthConfig +from nmp.common.config.base import OIDCConfig +from nmp.common.service import Service +from nmp.platform_runner import config as runner_config +from nmp.platform_runner import server + + +class DummyService(Service): + def __init__(self, name: str = "models") -> None: + super().__init__(name=name, module_name="test.sidecars") + + def get_routers(self): + return [] + + +def _dummy_sidecar(_stop_signal: threading.Event) -> None: + return None + + +def _patch_runner_discovery( + monkeypatch: pytest.MonkeyPatch, + *, + services: dict[str, Service] | None = None, + controllers: dict[str, Callable[[threading.Event], object]] | None = None, + sidecars: dict[str, Callable[[threading.Event], object]] | None = None, +) -> None: + services = services if services is not None else {"models": DummyService("models")} + controllers = controllers if controllers is not None else {} + sidecars = sidecars if sidecars is not None else {"adapters": _dummy_sidecar} + + monkeypatch.setattr(runner_config, "get_available_services", lambda: services) + monkeypatch.setattr(runner_config, "get_available_controllers", lambda: controllers) + monkeypatch.setattr( + runner_config, + "get_service_groups", + lambda _available: {"all": list(services), "core": list(services), "api": []}, + ) + monkeypatch.setattr( + runner_config, "get_controller_groups", lambda _available: {"all": list(controllers), "core": list(controllers)} + ) + monkeypatch.setattr(runner_config, "get_default_controllers", lambda _groups: list(controllers)) + monkeypatch.setattr(runner_config, "AVAILABLE_SIDECARS", sidecars) + monkeypatch.setattr("nmp.platform_runner.registry.AVAILABLE_SIDECARS", sidecars) + monkeypatch.setattr(server, "AVAILABLE_SIDECARS", sidecars, raising=False) + monkeypatch.setattr(server, "order_services_by_dependencies", lambda service_instances: service_instances) + + +def test_models_service_resolves_adapters_sidecar_dependency(monkeypatch: pytest.MonkeyPatch) -> None: + _patch_runner_discovery(monkeypatch) + + resolved = runner_config.resolve_run_configuration( + runner_config.PlatformAppConfig(services=["models"], controllers=[]) + ) + + assert resolved.services == {"models"} + assert resolved.sidecars == {"adapters"} + + +def test_explicit_sidecar_can_run_without_services(monkeypatch: pytest.MonkeyPatch) -> None: + _patch_runner_discovery(monkeypatch) + + resolved = runner_config.resolve_run_configuration( + runner_config.PlatformAppConfig(services=[], controllers=[], sidecars=["adapters"]) + ) + + assert resolved.services == set() + assert resolved.controllers == set() + assert resolved.sidecars == {"adapters"} + + +def _auth_config(enabled: bool = False) -> AuthConfig: + return AuthConfig( + enabled=enabled, + policy_decision_point_base_url="http://localhost:8181", + oidc=OIDCConfig(enabled=False), + ) + + +def _sidecar_with_events(started: threading.Event, stopped: threading.Event) -> Callable[[threading.Event], None]: + def run(stop_signal: threading.Event) -> None: + started.set() + stop_signal.wait(timeout=5.0) + stopped.set() + + return run + + +def test_create_app_starts_and_stops_dummy_sidecar_with_lifespan() -> None: + started = threading.Event() + stopped = threading.Event() + + with ( + patch("nmp.platform_runner.server.get_platform_config") as platform_config, + patch("nmp.platform_runner.server.get_auth_config", return_value=_auth_config(False)), + patch("nmp.common.auth.middleware.get_auth_config", return_value=_auth_config(False)), + ): + platform_config.return_value.seed_on_startup = False + platform_config.return_value.redirect_root_to_studio = False + app = server.create_app( + services=[], + controller_run_funcs={"adapters": _sidecar_with_events(started, stopped)}, + ) + from fastapi.testclient import TestClient + + with TestClient(app) as client: + assert started.wait(timeout=1.0) + assert client.get("/").status_code == 200 + + assert stopped.wait(timeout=1.0) + + +def test_build_platform_app_loads_dependent_sidecar_into_lifespan(monkeypatch: pytest.MonkeyPatch) -> None: + started = threading.Event() + stopped = threading.Event() + _patch_runner_discovery(monkeypatch, sidecars={"adapters": _sidecar_with_events(started, stopped)}) + + with ( + patch("nmp.platform_runner.server.get_platform_config") as platform_config, + patch("nmp.platform_runner.server.get_auth_config", return_value=_auth_config(False)), + patch("nmp.common.auth.middleware.get_auth_config", return_value=_auth_config(False)), + ): + platform_config.return_value.seed_on_startup = False + platform_config.return_value.redirect_root_to_studio = False + app = server.build_platform_app(runner_config.PlatformAppConfig(services=["models"], controllers=[]), env={}) + from fastapi.testclient import TestClient + + with TestClient(app) as client: + assert started.wait(timeout=1.0) + assert client.get("/").status_code == 200 + + assert stopped.wait(timeout=1.0) + + +def test_build_platform_app_rejects_controller_sidecar_name_collision(monkeypatch: pytest.MonkeyPatch) -> None: + _patch_runner_discovery( + monkeypatch, + controllers={"adapters": _dummy_sidecar}, + sidecars={"adapters": _dummy_sidecar}, + ) + + with pytest.raises(ValueError, match="Controller/sidecar name collision: adapters"): + server.build_platform_app( + runner_config.PlatformAppConfig(controllers=["adapters"], sidecars=["adapters"]), + env={}, + ) + + +def test_real_adapters_sidecar_entrypoint_starts_and_stops_with_required_env( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + from nmp.core.models.sidecars.adapters import main as adapters_main + + started = threading.Event() + stopped = threading.Event() + manager = MagicMock() + + class FakeLoop: + def __init__(self, *_args: object, **_kwargs: object) -> None: + pass + + def start(self) -> None: + started.set() + + def stop(self) -> None: + stopped.set() + + def join(self) -> None: + return None + + lora_dir = tmp_path / "loras" + monkeypatch.setenv("NIM_PEFT_SOURCE", str(lora_dir)) + monkeypatch.setenv("NMP_MODEL_ENTITY_WORKSPACE", "default") + monkeypatch.setenv("NMP_MODEL_ENTITY_NAME", "test-model") + monkeypatch.setenv("NIM_PEFT_REFRESH_INTERVAL", "30") + monkeypatch.delenv("VLLM_ENDPOINT", raising=False) + + monkeypatch.setattr(adapters_main, "get_platform_config", lambda: MagicMock(base_url="http://platform.local")) + monkeypatch.setattr(adapters_main, "get_platform_sdk", lambda **_kwargs: MagicMock()) + monkeypatch.setattr(adapters_main.asyncio, "new_event_loop", lambda: MagicMock()) + monkeypatch.setattr(adapters_main, "Loop", FakeLoop) + monkeypatch.setattr(adapters_main, "TimedLoopWaiter", lambda *_args, **_kwargs: object()) + monkeypatch.setattr(adapters_main.ControllerManager, "get_instance", classmethod(lambda _cls: manager)) + + stop_signal = threading.Event() + thread = threading.Thread(target=adapters_main.run, args=(stop_signal,), daemon=True) + try: + thread.start() + + assert started.wait(timeout=1.0) + manager.register.assert_called_once() + assert manager.register.call_args.args[0] == "adapters_controller" + + stop_signal.set() + thread.join(timeout=1.0) + + assert not thread.is_alive() + assert stopped.is_set() + finally: + stop_signal.set() + thread.join(timeout=1.0) + adapters_main.adapters_controller_monitored = None diff --git a/sdk/python/nemo-platform/src/nemo_platform/cli/commands/services/cli.py b/sdk/python/nemo-platform/src/nemo_platform/cli/commands/services/cli.py index 4df852bb8e..8a7e3022bd 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/cli/commands/services/cli.py +++ b/sdk/python/nemo-platform/src/nemo_platform/cli/commands/services/cli.py @@ -13,7 +13,8 @@ import httpx import typer -from nemo_platform.cli.commands.services._process import ( +from nemo_platform.cli.core.help_formatter import create_typer_app +from nemo_platform.local.process import ( ForegroundInstanceError, InstanceAlreadyRunningError, InstanceDescriptor, @@ -24,7 +25,6 @@ check_port_available_for_start, compute_scope, format_port_conflict, - get_create_time, instance_log_bytes, is_instance_alive, list_instances, @@ -37,7 +37,7 @@ stop_instance, write_descriptor, ) -from nemo_platform.cli.core.help_formatter import create_typer_app +from nmp.platform_runner.config import DEFAULT_LOCAL_SERVICES_BIND_HOST, PlatformAppConfig logger = logging.getLogger(__name__) @@ -45,7 +45,6 @@ _HEALTH_TIMEOUT_SECONDS = 60 _HEALTH_POLL_INTERVAL = 2.0 -_DEFAULT_HOST = "127.0.0.1" _DEFAULT_PORT = 8080 _DEFAULT_STOP_TIMEOUT = 30.0 @@ -60,7 +59,7 @@ def services_callback(ctx: typer.Context) -> None: for info in running: desc = info.descriptor assert desc is not None - typer.echo(f"\nRunning: {info.scope} (pid {desc.pid}, {desc.host}:{desc.port}, {desc.mode})") + typer.echo(f"\nRunning: {info.scope} (pid {desc.pid}, {desc.config.host}:{desc.config.port}, {desc.mode})") def _require_services_extra() -> None: @@ -92,7 +91,7 @@ def _parse_csv_option(value: str | None) -> list[str] | None: def _wait_for_healthy( host: str, port: int, - timeout: int = _HEALTH_TIMEOUT_SECONDS, + timeout: float = _HEALTH_TIMEOUT_SECONDS, poll_interval: float = _HEALTH_POLL_INTERVAL, ) -> bool: """Poll the platform status endpoint until it responds or timeout.""" @@ -124,15 +123,8 @@ def _effective_base_dir() -> str | None: def _find_sole_running_scope(base_dir: Path | None) -> str: - """Find the scope of the single running instance for this working directory. - - When the user runs ``restart`` without ``--instance`` or ``--port``, we - can't know which scope to target because the scope includes the port. - This function scans all running instances whose scope starts with the - same git-root hash prefix. If exactly one matches, return it. - Otherwise fall back to the default scope (hash-DEFAULT_PORT). - """ - prefix = compute_scope(port=0, instance_name=None).rsplit("-", 1)[0] + """Return the only running scope for this working directory, or the default scope.""" + prefix = compute_scope(port=0).rsplit("-", 1)[0] running = [i for i in list_instances(base_dir=base_dir) if i.alive and i.scope.startswith(prefix + "-")] if len(running) == 1: return running[0].scope @@ -213,7 +205,7 @@ def run_services( str | None, typer.Option("--config", help="Path to a platform configuration YAML file."), ] = None, - host: Annotated[str, typer.Option("--host", help="Host to bind to.")] = _DEFAULT_HOST, + host: Annotated[str, typer.Option("--host", help="Host to bind to.")] = DEFAULT_LOCAL_SERVICES_BIND_HOST, port: Annotated[int, typer.Option("--port", help="Port to bind to.")] = _DEFAULT_PORT, instance: Annotated[ str | None, @@ -226,7 +218,7 @@ def run_services( _require_services_extra() _warn_bind_all(host) - scope = compute_scope(port=port, instance_name=instance) + scope = compute_scope(port=port, explicit_scope=instance) base_dir_str = _effective_base_dir() base_dir = Path(base_dir_str) if base_dir_str else None @@ -242,20 +234,23 @@ def run_services( # "foreground", which protects interactive ``run`` sessions from being # killed by ``stop``. mode = "background" if os.environ.get("_NMP_LAUNCH_MODE") == "background" else "foreground" - - desc = InstanceDescriptor( - pid=os.getpid(), - scope=scope, - host=host, - port=port, - mode=mode, - create_time=get_create_time(os.getpid()), + platform_config = PlatformAppConfig( services=_parse_csv_option(services), - controllers=_parse_csv_option(controllers), service_group=service_group, + controllers=_parse_csv_option(controllers), controller_group=controller_group, sidecars=_parse_csv_option(sidecars), config_path=config, + scope=scope, + host=host, + port=port, + state_root=base_dir, + ) + + desc = InstanceDescriptor.from_config( + platform_config, + mode=mode, + pid=os.getpid(), ) write_descriptor(desc, base_dir=base_dir) @@ -269,14 +264,7 @@ def _cleanup() -> None: from nmp.platform_runner.run import run_platform run_platform( - services=_parse_csv_option(services), - service_group=service_group, - controllers=_parse_csv_option(controllers), - controller_group=controller_group, - sidecars=_parse_csv_option(sidecars), - config_path=config, - host=host, - port=port, + config=platform_config, on_shutdown=_cleanup, ) @@ -327,7 +315,7 @@ def start_services( str | None, typer.Option("--config", help="Path to a platform configuration YAML file."), ] = None, - host: Annotated[str, typer.Option("--host", help="Host to bind to.")] = _DEFAULT_HOST, + host: Annotated[str, typer.Option("--host", help="Host to bind to.")] = DEFAULT_LOCAL_SERVICES_BIND_HOST, port: Annotated[int, typer.Option("--port", help="Port to bind to.")] = _DEFAULT_PORT, instance: Annotated[ str | None, @@ -351,7 +339,7 @@ def start_services( raise typer.BadParameter("Cannot combine --controllers with --controller-group.") _warn_bind_all(host) - scope = compute_scope(port=port, instance_name=instance) + scope = compute_scope(port=port, explicit_scope=instance) base_dir_str = _effective_base_dir() base_dir = Path(base_dir_str) if base_dir_str else None @@ -360,20 +348,22 @@ def start_services( _ensure_port_available(host, port, scope, base_dir=base_dir) - typer.echo("Starting platform services...") - proc = start_background( - scope=scope, + platform_config = PlatformAppConfig( services=_parse_csv_option(services), service_group=service_group, controllers=_parse_csv_option(controllers), controller_group=controller_group, sidecars=_parse_csv_option(sidecars), config_path=config, + scope=scope, host=host, port=port, - base_dir=base_dir, + state_root=base_dir, ) + typer.echo("Starting platform services...") + proc = start_background(platform_config) + if not _wait_for_healthy(host, port): exit_code = proc.poll() if exit_code is not None: @@ -426,7 +416,7 @@ def stop_services_cmd( nemo services stop nemo services stop --timeout 60 """ - scope = compute_scope(port=port, instance_name=instance) + scope = compute_scope(port=port, explicit_scope=instance) base_dir_str = _effective_base_dir() base_dir = Path(base_dir_str) if base_dir_str else None @@ -496,7 +486,10 @@ def restart_services( ] = None, host: Annotated[ str | None, - typer.Option("--host", help="Host to bind to. Defaults to previous value or 127.0.0.1."), + typer.Option( + "--host", + help=f"Host to bind to. Defaults to previous value or {DEFAULT_LOCAL_SERVICES_BIND_HOST}.", + ), ] = None, port: Annotated[ int | None, @@ -529,8 +522,8 @@ def restart_services( base_dir = Path(base_dir_str) if base_dir_str else None if instance is not None or port is not None: - effective_port = port if port is not None else _DEFAULT_PORT - scope = compute_scope(port=effective_port, instance_name=instance) + effective_scope_port = port if port is not None else _DEFAULT_PORT + scope = compute_scope(port=effective_scope_port, explicit_scope=instance) else: scope = _find_sole_running_scope(base_dir) @@ -547,37 +540,47 @@ def restart_services( # appropriate even for foreground targets. stop_instance(scope, base_dir=base_dir, force=True) - effective_services = _parse_csv_option(services) if services is not None else (prev.services if prev else None) - effective_service_group = service_group if service_group is not None else (prev.service_group if prev else None) - effective_controllers = ( - _parse_csv_option(controllers) if controllers is not None else (prev.controllers if prev else None) + previous_config = prev.config if prev else None + effective_services = _parse_csv_option(services) if services is not None else None + if services is None and previous_config is not None: + effective_services = previous_config.services + effective_service_group = service_group if service_group is not None else None + if service_group is None and previous_config is not None: + effective_service_group = previous_config.service_group + effective_controllers = _parse_csv_option(controllers) if controllers is not None else None + if controllers is None and previous_config is not None: + effective_controllers = previous_config.controllers + effective_controller_group = controller_group if controller_group is not None else None + if controller_group is None and previous_config is not None: + effective_controller_group = previous_config.controller_group + effective_sidecars = _parse_csv_option(sidecars) if sidecars is not None else None + if sidecars is None and previous_config is not None: + effective_sidecars = previous_config.sidecars + effective_config = config if config is not None else (previous_config.config_path if previous_config else None) + effective_host = ( + host if host is not None else (previous_config.host if previous_config else DEFAULT_LOCAL_SERVICES_BIND_HOST) ) - effective_controller_group = ( - controller_group if controller_group is not None else (prev.controller_group if prev else None) - ) - effective_sidecars = _parse_csv_option(sidecars) if sidecars is not None else (prev.sidecars if prev else None) - effective_config = config if config is not None else (prev.config_path if prev else None) - effective_host = host if host is not None else (prev.host if prev else _DEFAULT_HOST) - effective_port = port if port is not None else (prev.port if prev else _DEFAULT_PORT) + effective_port = port if port is not None else (previous_config.port if previous_config else _DEFAULT_PORT) _warn_bind_all(effective_host) _ensure_port_available(effective_host, effective_port, scope, base_dir=base_dir) - - typer.echo("Starting platform services...") - proc = start_background( - scope=scope, + platform_config = PlatformAppConfig( services=effective_services, service_group=effective_service_group, controllers=effective_controllers, controller_group=effective_controller_group, sidecars=effective_sidecars, config_path=effective_config, + scope=scope, host=effective_host, port=effective_port, - base_dir=base_dir, + state_root=base_dir, ) + typer.echo("Starting platform services...") + proc = start_background(platform_config) + if not _wait_for_healthy(effective_host, effective_port): exit_code = proc.poll() if exit_code is not None: @@ -613,7 +616,7 @@ def status_services( ] = _DEFAULT_PORT, ) -> None: """Show status of the platform services instance for this scope.""" - scope = compute_scope(port=port, instance_name=instance) + scope = compute_scope(port=port, explicit_scope=instance) base_dir_str = _effective_base_dir() base_dir = Path(base_dir_str) if base_dir_str else None @@ -641,13 +644,13 @@ def status_services( except ValueError: uptime = "unknown" - healthy = _wait_for_healthy(desc.host, desc.port, timeout=3, poll_interval=0.5) + healthy = _wait_for_healthy(desc.config.host, desc.config.port, timeout=3, poll_interval=0.5) health_str = "healthy" if healthy else "unhealthy" - typer.echo(f"Scope: {desc.scope}") + typer.echo(f"Scope: {desc.config.scope}") typer.echo(f"PID: {desc.pid}") typer.echo(f"Mode: {desc.mode}") - typer.echo(f"Address: {desc.host}:{desc.port}") + typer.echo(f"Address: {desc.config.host}:{desc.config.port}") typer.echo(f"Uptime: {uptime}") typer.echo(f"Health: {health_str}") log = log_path_for(scope, base_dir=base_dir) @@ -668,7 +671,7 @@ def _print_instance_table(instances: list[InstanceInfo]) -> None: pid = addr = mode = "-" if info.descriptor: pid = str(info.descriptor.pid) - addr = f"{info.descriptor.host}:{info.descriptor.port}" + addr = f"{info.descriptor.config.host}:{info.descriptor.config.port}" mode = info.descriptor.mode typer.echo(f"{info.scope:<25} {status:<10} {pid:<10} {addr:<25} {mode:<12}") @@ -872,7 +875,7 @@ def logs_services( nemo services logs --path nemo services logs -n 100 """ - scope = compute_scope(port=port, instance_name=instance) + scope = compute_scope(port=port, explicit_scope=instance) base_dir_str = _effective_base_dir() base_dir = Path(base_dir_str) if base_dir_str else None diff --git a/sdk/python/nemo-platform/src/nemo_platform/cli/commands/setup.py b/sdk/python/nemo-platform/src/nemo_platform/cli/commands/setup.py index 29081aa867..78849e4181 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/cli/commands/setup.py +++ b/sdk/python/nemo-platform/src/nemo_platform/cli/commands/setup.py @@ -30,20 +30,12 @@ from nemo_platform_plugin.secrets.client import SecretsClient from nemo_platform_plugin.secrets.types import PlatformSecretCreateRequest, PlatformSecretUpdateRequest from nmp.common.config import nmp_user_data_dir +from nmp.platform_runner.config import DEFAULT_LOCAL_SERVICES_BIND_HOST, PlatformAppConfig from pydantic import SecretStr from rich import box from rich.console import Console from rich.panel import Panel -from nemo_platform.cli.commands.services._process import ( - DEFAULT_SERVICES_BIND_HOST, - check_port_available_for_start, - compute_scope, - format_port_conflict, - log_path_for, - start_background, - stop_instance, -) from nemo_platform.cli.commands.skills import registry as skills_registry from nemo_platform.cli.commands.skills.base import Scope, Skill from nemo_platform.cli.commands.skills.registry import get_installer, load_skills @@ -51,6 +43,14 @@ from nemo_platform.cli.core.errors import handle_errors from nemo_platform.config.config import Config from nemo_platform.config.models import ConfigFile, ConfigParams, LocalServicesConfig +from nemo_platform.local.process import ( + check_port_available_for_start, + compute_scope, + format_port_conflict, + log_path_for, + start_background, + stop_instance, +) from nemo_platform.ui.prompts import ( UserCancelled, is_interactive, @@ -635,8 +635,10 @@ def _start_services_background(base_url: str, data_dir: str | None = None) -> su exported it). """ port = _resolve_services_port(base_url) - scope = compute_scope(port=port) - return start_background(scope=scope, port=port, data_dir=data_dir) + return start_background( + PlatformAppConfig(scope=compute_scope(port=port), port=port), + data_dir=data_dir, + ) def _last_startup_service(log_path: Path | None) -> str: @@ -686,16 +688,13 @@ def _kill_existing_services(base_url: str) -> None: Delegates to the shared process lifecycle module. """ - port = _resolve_services_port(base_url) - scope = compute_scope(port=port) - stop_instance(scope, timeout=2.0, force=True) + stop_instance(compute_scope(port=_resolve_services_port(base_url)), timeout=2.0, force=True) def _ensure_port_available_for_start(base_url: str) -> None: """Fail fast when the services port cannot be bound.""" port = _resolve_services_port(base_url) - scope = compute_scope(port=port) - conflict = check_port_available_for_start(DEFAULT_SERVICES_BIND_HOST, port, scope) + conflict = check_port_available_for_start(DEFAULT_LOCAL_SERVICES_BIND_HOST, port, compute_scope(port=port)) if conflict is None: return lines = format_port_conflict(conflict) @@ -777,8 +776,7 @@ def _maybe_start_services( _ensure_port_available_for_start(base_url) proc = _start_services_background(base_url, data_dir=data_dir) - port = _resolve_services_port(base_url) - log = log_path_for(compute_scope(port=port)) + log = log_path_for(compute_scope(port=_resolve_services_port(base_url))) if not _wait_for_platform(base_url, timeout=timeout, log_path=log): exit_code = proc.poll() diff --git a/sdk/python/nemo-platform/src/nemo_platform/local/_service_child.py b/sdk/python/nemo-platform/src/nemo_platform/local/_service_child.py new file mode 100644 index 0000000000..1ff27df811 --- /dev/null +++ b/sdk/python/nemo-platform/src/nemo_platform/local/_service_child.py @@ -0,0 +1,30 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Child entrypoint for SDK-started local services daemons.""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +from nemo_platform.local.services import ServiceRunConfig, run_services + + +def main(argv: list[str] | None = None) -> int: + args = list(sys.argv[1:] if argv is None else argv) + if len(args) != 1: + sys.stderr.write("usage: python -m nemo_platform.local._service_child \n") + return 2 + request_path = Path(args[0]) + try: + payload = json.loads(request_path.read_text(encoding="utf-8")) + finally: + request_path.unlink(missing_ok=True) + run_services(ServiceRunConfig(**payload), _mode="daemon") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/sdk/python/nemo-platform/src/nemo_platform/cli/commands/services/_process.py b/sdk/python/nemo-platform/src/nemo_platform/local/process.py similarity index 76% rename from sdk/python/nemo-platform/src/nemo_platform/cli/commands/services/_process.py rename to sdk/python/nemo-platform/src/nemo_platform/local/process.py index 8ce8cda699..7ee0cd4a9b 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/cli/commands/services/_process.py +++ b/sdk/python/nemo-platform/src/nemo_platform/local/process.py @@ -3,8 +3,13 @@ """Local process lifecycle for ``nemo services``. -Uses per-instance scoped directories under ``$XDG_STATE_HOME/nmp/instances/`` -with flock-based liveness tracking. Each instance directory contains: +In this module, "instance" is a local services process/resource, and "scope" +is the stable key used for that instance's lock, descriptor, socket, and log +paths. The CLI exposes this key as ``--instance`` for compatibility, but +internal code should use "scope" when referring to the key. + +Uses per-scope directories under ``$XDG_STATE_HOME/nmp/instances/`` +with flock-based liveness tracking. Each scope directory contains: - ``services.lock`` -- exclusive flock held for the process lifetime - ``instance.json`` -- descriptor with PID, port, services, etc. @@ -23,7 +28,6 @@ import json import logging import os -import re import shutil import signal import socket @@ -34,10 +38,16 @@ from dataclasses import dataclass, field from datetime import datetime, timezone from pathlib import Path -from typing import Literal +from typing import Literal, Self import psutil -from pydantic import BaseModel, Field +from nmp.platform_runner.config import ( + DEFAULT_LOCAL_SERVICES_BIND_HOST, + PlatformAppConfig, + default_state_root, + validate_scope, +) +from pydantic import BaseModel, Field, model_validator logger = logging.getLogger(__name__) @@ -45,10 +55,10 @@ DESCRIPTOR_FILENAME = "instance.json" LOG_FILENAME = "services.log" -DEFAULT_SERVICES_BIND_HOST = "127.0.0.1" SUGGESTED_ALT_PORT = 9090 _SIGTERM_POLL_INTERVAL = 0.25 +_SIGKILL_WAIT_TIMEOUT = 5.0 _DEFAULT_STOP_TIMEOUT = 30.0 @@ -62,10 +72,7 @@ def _pause(seconds: float) -> None: def _base_state_dir() -> Path: - xdg = os.environ.get("XDG_STATE_HOME") - if xdg: - return Path(xdg) / "nmp" - return Path.home() / ".local" / "state" / "nmp" + return default_state_root() def _instances_dir(*, base_dir: Path | None = None) -> Path: @@ -73,7 +80,7 @@ def _instances_dir(*, base_dir: Path | None = None) -> Path: def _find_git_root() -> str: - """Walk up from cwd looking for a ``.git`` directory. Falls back to cwd.""" + """Walk up from cwd looking for a ``.git`` directory. Falls back to cwd.""" cur = Path.cwd().resolve() for parent in (cur, *cur.parents): if (parent / ".git").exists(): @@ -84,24 +91,19 @@ def _find_git_root() -> str: _scope_prefix_cache: str | None = None -_SCOPE_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,127}$") - - -def _validate_scope(scope: str) -> str: - """Ensure *scope* is safe to use as a directory name.""" - if not _SCOPE_RE.fullmatch(scope): - raise ValueError(f"Invalid instance scope: {scope!r}") - return scope +def compute_scope(*, port: int, explicit_scope: str | None = None) -> str: + """Compute the local services scope. + The default scope is ``sha1(git_toplevel_or_cwd)[:8]-``. Including + the port is intentional: it lets two local services instances from the same + checkout use different TCP ports without sharing a lock, descriptor, or log + directory. -def compute_scope(*, port: int, instance_name: str | None = None) -> str: - """Compute a scope identifier for this working directory + port. - - Default: ``sha1(git_toplevel_or_cwd)[:8]-``. - Override with an explicit *instance_name*. + Explicit scopes are validated and returned as-is, so they do not encode the + port. Callers that pass an explicit scope own its uniqueness. """ - if instance_name: - return _validate_scope(instance_name) + if explicit_scope: + return validate_scope(explicit_scope) global _scope_prefix_cache # noqa: PLW0603 if _scope_prefix_cache is None: root = _find_git_root() @@ -110,7 +112,8 @@ def compute_scope(*, port: int, instance_name: str | None = None) -> str: def instance_dir(scope: str, *, base_dir: Path | None = None) -> Path: - d = _instances_dir(base_dir=base_dir) / _validate_scope(scope) + """Return the state directory for *scope*, creating it if needed.""" + d = _instances_dir(base_dir=base_dir) / validate_scope(scope) d.mkdir(parents=True, exist_ok=True) return d @@ -222,7 +225,7 @@ def _instance_owns_listener( desc = read_descriptor(scope, base_dir=base_dir) if desc is None: return False - return desc.port == port and _normalize_bind_host(desc.host) == _normalize_bind_host(host) + return desc.config.port == port and _normalize_bind_host(desc.config.host) == _normalize_bind_host(host) def is_port_bindable(host: str, port: int) -> bool: @@ -272,8 +275,9 @@ def format_port_conflict(err: PortConflict) -> list[str]: Message text depends on ``err.kind`` (foreign process vs NeMo instance). """ if err.kind == "nemo_instance": + owner = f" '{err.scope}'" if err.scope else "" return [ - f"Port {err.port} is in use by a NeMo Platform instance for this directory.", + f"Port {err.port} is in use by NeMo Platform instance{owner}.", "Stop it first with: nemo services stop", "Or restart with: nemo services restart", ] @@ -292,23 +296,39 @@ def format_port_conflict(err: PortConflict) -> list[str]: class InstanceDescriptor(BaseModel): pid: int - scope: str - host: str = "127.0.0.1" - port: int = 8080 - mode: Literal["foreground", "background"] = "background" + config: PlatformAppConfig = Field(default_factory=PlatformAppConfig) + transport: Literal["tcp", "uds"] = "tcp" + mode: Literal["foreground", "background", "daemon"] = "background" create_time: float = 0.0 started_at: str = Field(default_factory=lambda: datetime.now(timezone.utc).isoformat()) - services: list[str] | None = None - controllers: list[str] | None = None - service_group: str | None = None - controller_group: str | None = None - sidecars: list[str] | None = None - config_path: str | None = None - log_path: str | None = None + + @model_validator(mode="after") + def _validate_client_transport(self) -> Self: + if self.transport == "uds" and self.config.socket_path is None: + raise ValueError("UDS client transport requires config.socket_path") + return self + + @classmethod + def from_config( + cls, + config: PlatformAppConfig, + *, + mode: Literal["foreground", "background", "daemon"], + transport: Literal["uds", "tcp"] = "tcp", + pid: int | None = None, + ) -> Self: + resolved_pid = os.getpid() if pid is None else pid + return cls( + pid=resolved_pid, + config=config, + transport=transport, + mode=mode, + create_time=get_create_time(resolved_pid), + ) def write_descriptor(desc: InstanceDescriptor, *, base_dir: Path | None = None) -> Path: - d = instance_dir(desc.scope, base_dir=base_dir) + d = instance_dir(desc.config.scope, base_dir=base_dir) path = d / DESCRIPTOR_FILENAME payload = desc.model_dump() fd, tmp = tempfile.mkstemp(dir=str(d), suffix=".tmp") @@ -335,10 +355,19 @@ def read_descriptor(scope: str, *, base_dir: Path | None = None) -> InstanceDesc return None try: data = json.loads(path.read_text()) - return InstanceDescriptor.model_validate(data) + desc = InstanceDescriptor.model_validate(data) except (json.JSONDecodeError, KeyError, TypeError, ValueError): logger.debug("Corrupt descriptor at %s, ignoring", path, exc_info=True) return None + if desc.config.scope != scope: + logger.debug( + "Descriptor at %s has scope=%r but lives under %r, ignoring", + path, + desc.config.scope, + scope, + ) + return None + return desc def remove_descriptor(scope: str, *, base_dir: Path | None = None) -> None: @@ -351,24 +380,24 @@ def remove_descriptor(scope: str, *, base_dir: Path | None = None) -> None: def _scope_dir(scope: str, *, base_dir: Path | None = None) -> Path: - return _instances_dir(base_dir=base_dir) / _validate_scope(scope) + return _instances_dir(base_dir=base_dir) / validate_scope(scope) def _is_log_file(path: Path) -> bool: return path.name == LOG_FILENAME or path.name.startswith(f"{LOG_FILENAME}.") -def _iter_log_files(scope_dir: Path): - if not scope_dir.is_dir(): +def _iter_log_files(scope_dir_path: Path): + if not scope_dir_path.is_dir(): return - for path in scope_dir.iterdir(): + for path in scope_dir_path.iterdir(): if path.is_file() and _is_log_file(path): yield path -def _has_preservable_logs(scope_dir: Path) -> bool: - """Return True if *scope_dir* contains non-empty service log files.""" - return any(path.stat().st_size > 0 for path in _iter_log_files(scope_dir)) +def _has_preservable_logs(scope_dir_path: Path) -> bool: + """Return True if *scope_dir_path* contains non-empty service log files.""" + return any(path.stat().st_size > 0 for path in _iter_log_files(scope_dir_path)) def is_removable_ghost( @@ -377,17 +406,17 @@ def is_removable_ghost( base_dir: Path | None = None, descriptor: InstanceDescriptor | None = None, ) -> bool: - """True when a dead scope dir has no descriptor and no non-empty logs.""" + """True when a dead scope directory has no descriptor and no non-empty logs.""" if is_instance_alive(scope, base_dir=base_dir): return False if descriptor is not None: return False - scope_dir = _scope_dir(scope, base_dir=base_dir) - if not scope_dir.is_dir(): + scope_dir_path = _scope_dir(scope, base_dir=base_dir) + if not scope_dir_path.is_dir(): return False - if (scope_dir / DESCRIPTOR_FILENAME).exists(): + if (scope_dir_path / DESCRIPTOR_FILENAME).exists(): return False - return not _has_preservable_logs(scope_dir) + return not _has_preservable_logs(scope_dir_path) # --------------------------------------------------------------------------- @@ -425,7 +454,7 @@ class InstanceInfo: def list_instances(*, base_dir: Path | None = None) -> list[InstanceInfo]: - """Scan all instance directories and return their status. + """Scan all scope directories and return their status. Side effects: - Removes stale descriptors for dead instances. @@ -448,7 +477,7 @@ def list_instances(*, base_dir: Path | None = None) -> list[InstanceInfo]: try: shutil.rmtree(child) except OSError: - logger.debug("Could not remove ghost instance dir %s", child, exc_info=True) + logger.debug("Could not remove ghost scope directory %s", child, exc_info=True) else: continue results.append(InstanceInfo(scope=scope, alive=alive, descriptor=desc)) @@ -456,28 +485,28 @@ def list_instances(*, base_dir: Path | None = None) -> list[InstanceInfo]: def remove_instance(scope: str, *, base_dir: Path | None = None) -> bool: - """Remove an instance scope directory. + """Remove a scope directory. - Returns False if the scope did not exist or could not be removed. + Returns False if the scope directory did not exist or could not be removed. """ - scope = _validate_scope(scope) + scope = validate_scope(scope) if is_instance_alive(scope, base_dir=base_dir): raise InstanceStillRunningError(scope) - scope_dir = _scope_dir(scope, base_dir=base_dir) - if not scope_dir.is_dir(): + scope_dir_path = _scope_dir(scope, base_dir=base_dir) + if not scope_dir_path.is_dir(): return False with contextlib.suppress(OSError): - shutil.rmtree(scope_dir) - return not scope_dir.is_dir() + shutil.rmtree(scope_dir_path) + return not scope_dir_path.is_dir() def list_stopped_scopes(*, base_dir: Path | None = None) -> list[str]: - """Return scope names for instances that are not alive.""" + """Return scopes for instances that are not alive.""" return [info.scope for info in list_instances(base_dir=base_dir) if not info.alive] def prune_instances(*, base_dir: Path | None = None) -> list[str]: - """Remove all stopped instance directories. Returns removed scope names.""" + """Remove all stopped scope directories. Returns removed scopes.""" removed: list[str] = [] for scope in list_stopped_scopes(base_dir=base_dir): if remove_instance(scope, base_dir=base_dir): @@ -497,11 +526,15 @@ def instance_log_bytes(scope: str, *, base_dir: Path | None = None) -> int: def rotate_log(scope: str, *, base_dir: Path | None = None) -> Path: """Rotate the existing log and return the path for the new one.""" - d = instance_dir(scope, base_dir=base_dir) - log_path = d / LOG_FILENAME + return rotate_log_path(log_path_for(scope, base_dir=base_dir)) + + +def rotate_log_path(log_path: Path) -> Path: + """Rotate the existing log at *log_path* and return the path for the new one.""" + log_path.parent.mkdir(parents=True, exist_ok=True) if log_path.exists() and log_path.stat().st_size > 0: ts = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S%fZ") - rotated = d / f"{LOG_FILENAME}.{ts}" + rotated = log_path.with_name(f"{log_path.name}.{ts}") log_path.rename(rotated) return log_path @@ -633,6 +666,10 @@ def stop_instance( return StopResult(stopped_pids=[], swept_children=swept) except OSError: logger.debug("Failed to send SIGKILL to pid %d", pid, exc_info=True) + if not _wait_for_pid_exit(pid, timeout=_SIGKILL_WAIT_TIMEOUT): + logger.warning("PID %d is still alive after SIGKILL; preserving descriptor", pid) + swept = _sweep_orphans(children) if children else [] + return StopResult(stopped_pids=[], swept_children=swept) swept = _sweep_orphans(children) if children else [] @@ -640,13 +677,21 @@ def stop_instance( return StopResult(stopped_pids=[pid], swept_children=swept) +def _wait_for_pid_exit(pid: int, *, timeout: float) -> bool: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if not _pid_alive(pid): + return True + _pause(_SIGTERM_POLL_INTERVAL) + return not _pid_alive(pid) + + def _pid_alive(pid: int) -> bool: try: - os.kill(pid, 0) - return True - except ProcessLookupError: + return psutil.Process(pid).status() != psutil.STATUS_ZOMBIE + except psutil.NoSuchProcess: return False - except PermissionError: + except psutil.AccessDenied: return True except OSError: return False @@ -658,17 +703,8 @@ def _pid_alive(pid: int) -> bool: def start_background( + config: PlatformAppConfig | None = None, *, - scope: str, - services: list[str] | None = None, - service_group: str | None = None, - controllers: list[str] | None = None, - controller_group: str | None = None, - sidecars: list[str] | None = None, - config_path: str | None = None, - host: str = DEFAULT_SERVICES_BIND_HOST, - port: int = 8080, - base_dir: Path | None = None, data_dir: str | None = None, ) -> subprocess.Popen: """Launch ``nemo services run`` as a detached background subprocess. @@ -676,31 +712,32 @@ def start_background( The child acquires the flock and writes its own descriptor. The parent returns the ``Popen`` handle for health polling. """ - log_file_path = rotate_log(scope, base_dir=base_dir) + config = config or PlatformAppConfig(host=DEFAULT_LOCAL_SERVICES_BIND_HOST) + log_file_path = rotate_log_path(config.log_file_path()) log_file = open(log_file_path, "a") # noqa: SIM115 nemo_bin = str(Path(sys.executable).parent / "nemo") args: list[str] = [nemo_bin, "services", "run"] - if services: - args += ["--services", ",".join(services)] - if service_group: - args += ["--service-group", service_group] - if controllers: - args += ["--controllers", ",".join(controllers)] - if controller_group: - args += ["--controller-group", controller_group] - if sidecars: - args += ["--sidecars", ",".join(sidecars)] - if config_path: - args += ["--config", config_path] - args += ["--host", host, "--port", str(port)] - args += ["--instance", scope] + if config.services: + args += ["--services", ",".join(config.services)] + if config.service_group: + args += ["--service-group", config.service_group] + if config.controllers: + args += ["--controllers", ",".join(config.controllers)] + if config.controller_group: + args += ["--controller-group", config.controller_group] + if config.sidecars: + args += ["--sidecars", ",".join(config.sidecars)] + if config.config_path: + args += ["--config", config.config_path] + args += ["--host", config.host, "--port", str(config.port)] + args += ["--instance", config.scope] env = os.environ.copy() if data_dir and "NMP_DATA_DIR" not in env: env["NMP_DATA_DIR"] = data_dir - if base_dir: - env["_NMP_STATE_DIR"] = str(base_dir) + if config.state_root is not None: + env["_NMP_STATE_DIR"] = str(config.state_root) # Tell the child ``run`` process it was launched by ``start`` so it # records mode="background" in its descriptor. This is internal # parent-to-child signaling -- not a public API surface -- following the diff --git a/sdk/python/nemo-platform/src/nemo_platform/local/services.py b/sdk/python/nemo-platform/src/nemo_platform/local/services.py new file mode 100644 index 0000000000..bd904c9cfd --- /dev/null +++ b/sdk/python/nemo-platform/src/nemo_platform/local/services.py @@ -0,0 +1,728 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Programmatic local lifecycle API for NeMo Platform services.""" + +from __future__ import annotations + +import asyncio +import importlib.util +import json +import os +import subprocess +import sys +import tempfile +import time +from collections.abc import MutableMapping, Sequence +from dataclasses import dataclass +from enum import StrEnum +from pathlib import Path +from typing import Any, Literal, Protocol, Self, runtime_checkable + +from nemo_platform import AsyncNeMoPlatform, NeMoPlatform +from nemo_platform.local import process +from nemo_platform.local.transport import ( + EMBEDDED_BASE_URL, + UDS_BASE_URL, + build_async_asgi_http_client, + build_async_http_client, + build_sync_asgi_http_client, + build_sync_http_client, + probe_status, + tcp_base_url, + wait_for_status, + wait_for_status_async, +) +from nmp.platform_runner.config import ( + DEFAULT_SCOPE, + PlatformAppConfig, + default_runtime_root, + default_state_root, + validate_scope, +) + +_AF_UNIX_PATH_MAX_BYTES = 103 if sys.platform.startswith(("darwin", "freebsd", "openbsd", "netbsd")) else 107 + + +class ServicesError(RuntimeError): + """Base class for local services lifecycle errors.""" + + +class ServicesExtraRequiredError(ServicesError): + """Raised when local service dependencies are not installed.""" + + +class ServicesAlreadyRunningError(ServicesError): + """Raised when a requested local instance is already running.""" + + +class ServicesNotRunningError(ServicesError): + """Raised when a requested local instance is not running.""" + + +class ServicesPortInUseError(ServicesError): + """Raised when TCP startup targets an unavailable port.""" + + +class ServicesStartupTimeoutError(ServicesError): + """Raised when startup does not become healthy before the timeout.""" + + +class ServicesStartupExitedError(ServicesError): + """Raised when a daemon child exits before becoming healthy.""" + + +class ServicesSocketStaleError(ServicesError): + """Raised when a stale socket cannot be removed.""" + + +def _as_tuple(value: Sequence[str] | None) -> tuple[str, ...] | None: + if value is None: + return None + return tuple(value) + + +def _optional_str(value: str | Path | None) -> str | None: + if value is None: + return None + return str(value) + + +def _optional_list(value: Sequence[str] | None) -> list[str] | None: + if value is None: + return None + return list(value) + + +class ServiceMode(StrEnum): + EMBEDDED = "embedded" + DAEMON = "daemon" + + +@dataclass(frozen=True) +class StartServicesResult: + requested: list[str] + started: list[str] + already_active: list[str] + active: list[str] + + +@runtime_checkable +class LocalServiceHandle(Protocol): + """Shared lifecycle/client contract for local services handles.""" + + def is_running(self) -> bool: ... + + def wait_until_ready(self, timeout: float | None = None) -> None: ... + + async def wait_until_ready_async(self, timeout: float | None = None) -> None: ... + + def client(self, **kwargs: Any) -> NeMoPlatform: ... + + def async_client(self, **kwargs: Any) -> AsyncNeMoPlatform: ... + + def start_services(self, service_names: Sequence[str]) -> StartServicesResult: ... + + async def start_services_async(self, service_names: Sequence[str]) -> StartServicesResult: ... + + def stop(self, *, timeout: float = 30.0, force: bool = False) -> process.StopResult: ... + + async def stop_async(self, *, timeout: float = 30.0, force: bool = False) -> process.StopResult: ... + + +@dataclass +class ServiceRunConfig: + services: Sequence[str] | None = None + service_group: str | None = None + controllers: Sequence[str] | None = None + controller_group: str | None = None + sidecars: Sequence[str] | None = None + config_path: str | Path | None = None + transport: Literal["uds", "tcp"] = "uds" + socket_path: str | Path | None = None + http_gateway: Literal["enabled", "disabled"] = "disabled" + http_gateway_host: str = "127.0.0.1" + http_gateway_port: int | None = None + host: str = "127.0.0.1" + port: int = 8080 + scope: str = DEFAULT_SCOPE + state_dir: str | Path | None = None + runtime_dir: str | Path | None = None + data_dir: str | Path | None = None + readiness_timeout: float = 60.0 + readiness_poll_interval: float = 0.5 + mode: ServiceMode | str = ServiceMode.DAEMON + + def __post_init__(self) -> None: + self.services = _as_tuple(self.services) + self.controllers = _as_tuple(self.controllers) + self.sidecars = _as_tuple(self.sidecars) + try: + self.mode = ServiceMode(self.mode) + except ValueError as error: + raise ValueError("mode must be 'embedded' or 'daemon'") from error + + if self.services and self.service_group: + raise ValueError("services cannot be combined with service_group") + if self.controllers and self.controller_group: + raise ValueError("controllers cannot be combined with controller_group") + if self.transport not in {"uds", "tcp"}: + raise ValueError("transport must be 'uds' or 'tcp'") + if self.http_gateway not in {"enabled", "disabled"}: + raise ValueError("http_gateway must be 'enabled' or 'disabled'") + if self.http_gateway == "enabled" and self.transport != "uds": + raise ValueError("gateway can only be enabled for UDS transport") + if self.readiness_timeout <= 0: + raise ValueError("readiness_timeout must be greater than 0") + if self.readiness_poll_interval <= 0: + raise ValueError("readiness_poll_interval must be greater than 0") + self.scope = validate_scope(self.scope) + + @property + def state_root(self) -> Path: + return Path(self.state_dir).expanduser() if self.state_dir is not None else default_state_root() + + @property + def runtime_root(self) -> Path: + return Path(self.runtime_dir).expanduser() if self.runtime_dir is not None else default_runtime_root() + + @property + def resolved_socket_path(self) -> Path | None: + if self.socket_path is not None: + socket_path = Path(self.socket_path).expanduser() + elif self.transport == "uds": + socket_path = PlatformAppConfig( + scope=self.scope, + runtime_root=self.runtime_root, + ).socket_file_path() + else: + return None + if not socket_path.is_absolute(): + raise ValueError(f"UDS socket path must be absolute: {socket_path}") + return socket_path + + def to_platform_app_config(self) -> PlatformAppConfig: + return PlatformAppConfig( + services=self.services, + service_group=self.service_group, + controllers=self.controllers, + controller_group=self.controller_group, + sidecars=self.sidecars, + config_path=_optional_str(self.config_path), + scope=self.scope, + host=self.host, + port=self.port, + socket_path=_optional_str(self.resolved_socket_path), + state_root=_optional_str(self.state_root), + runtime_root=_optional_str(self.runtime_dir), + ) + + def to_child_payload(self) -> dict[str, object]: + return { + "mode": ServiceMode(self.mode).value, + "services": _optional_list(self.services), + "service_group": self.service_group, + "controllers": _optional_list(self.controllers), + "controller_group": self.controller_group, + "sidecars": _optional_list(self.sidecars), + "config_path": _optional_str(self.config_path), + "transport": self.transport, + "socket_path": _optional_str(self.socket_path), + "http_gateway": self.http_gateway, + "http_gateway_host": self.http_gateway_host, + "http_gateway_port": self.http_gateway_port, + "host": self.host, + "port": self.port, + "scope": self.scope, + "state_dir": _optional_str(self.state_dir), + "runtime_dir": _optional_str(self.runtime_dir), + "data_dir": _optional_str(self.data_dir), + "readiness_timeout": self.readiness_timeout, + "readiness_poll_interval": self.readiness_poll_interval, + } + + +@dataclass(frozen=True) +class DaemonServiceHandle: + scope: str + transport: Literal["uds", "tcp"] + socket_path: Path | None + gateway_base_url: str | None + host: str + port: int + pid: int | None + mode: Literal["foreground", "daemon"] + log_path: Path | None + state_dir: Path | None + runtime_dir: Path | None + + @classmethod + def from_descriptor(cls, desc: process.InstanceDescriptor) -> Self: + socket_path = Path(desc.config.socket_path) if desc.config.socket_path else None + runtime_dir = desc.config.runtime_dir() if socket_path else None + return cls( + scope=desc.config.scope, + transport=desc.transport, + socket_path=socket_path, + gateway_base_url=None, + host=desc.config.host, + port=desc.config.port, + pid=desc.pid, + mode="daemon" if desc.mode == "daemon" else "foreground", + log_path=desc.config.log_file_path(), + state_dir=desc.config.state_dir(), + runtime_dir=runtime_dir, + ) + + @classmethod + def from_config( + cls, + config: ServiceRunConfig, + *, + pid: int | None = None, + ) -> Self: + app_config = config.to_platform_app_config() + socket_path = config.resolved_socket_path + runtime_dir = app_config.runtime_dir() if socket_path else None + return cls( + scope=config.scope, + transport=config.transport, + socket_path=socket_path, + gateway_base_url=None, + host=config.host, + port=config.port, + pid=pid, + mode="daemon", + log_path=app_config.log_file_path(), + state_dir=app_config.state_dir(), + runtime_dir=runtime_dir, + ) + + @property + def base_url(self) -> str: + if self.transport == "uds": + return UDS_BASE_URL + return tcp_base_url(self.host, self.port) + + def _state_root(self) -> Path | None: + if self.state_dir is None: + return None + if self.state_dir.parent.name == "instances": + return self.state_dir.parent.parent + return self.state_dir.parent + + def is_running(self) -> bool: + state_root = self._state_root() + return process.is_instance_alive(self.scope, base_dir=state_root) + + def wait_until_ready(self, timeout: float | None = None) -> None: + if not wait_for_status( + base_url=self.base_url, + socket_path=self.socket_path if self.transport == "uds" else None, + timeout=60.0 if timeout is None else timeout, + ): + raise ServicesStartupTimeoutError(f"Timed out waiting for services instance {self.scope!r}") + + async def wait_until_ready_async(self, timeout: float | None = None) -> None: + if not await wait_for_status_async( + base_url=self.base_url, + socket_path=self.socket_path if self.transport == "uds" else None, + timeout=60.0 if timeout is None else timeout, + ): + raise ServicesStartupTimeoutError(f"Timed out waiting for services instance {self.scope!r}") + + def stop(self, *, timeout: float = 30.0, force: bool = False) -> process.StopResult: + state_root = self._state_root() + return process.stop_instance(self.scope, base_dir=state_root, timeout=timeout, force=force) + + async def stop_async(self, *, timeout: float = 30.0, force: bool = False) -> process.StopResult: + return await asyncio.to_thread(self.stop, timeout=timeout, force=force) + + def start_services(self, service_names: Sequence[str]) -> StartServicesResult: + raise ServicesError("Staged service start is not implemented for daemon mode yet") + + async def start_services_async(self, service_names: Sequence[str]) -> StartServicesResult: + return await asyncio.to_thread(self.start_services, service_names) + + def client(self, **kwargs: Any) -> NeMoPlatform: + if self.transport == "uds": + if self.socket_path is None: + raise ServicesError("UDS service handle is missing socket_path") + kwargs.setdefault("http_client", build_sync_http_client(self.socket_path)) + kwargs.setdefault("base_url", self.base_url) + return NeMoPlatform(**kwargs) + + def async_client(self, **kwargs: Any) -> AsyncNeMoPlatform: + if self.transport == "uds": + if self.socket_path is None: + raise ServicesError("UDS service handle is missing socket_path") + kwargs.setdefault("http_client", build_async_http_client(self.socket_path)) + kwargs.setdefault("base_url", self.base_url) + return AsyncNeMoPlatform(**kwargs) + + +@dataclass(frozen=True) +class EmbeddedServiceHandle: + app: Any + runtime: object + + def is_running(self) -> bool: + return True + + def wait_until_ready(self, timeout: float | None = None) -> None: + return None + + async def wait_until_ready_async(self, timeout: float | None = None) -> None: + return None + + def client(self, **kwargs: Any) -> NeMoPlatform: + kwargs.setdefault("http_client", build_sync_asgi_http_client(self.app)) + kwargs.setdefault("base_url", EMBEDDED_BASE_URL) + return NeMoPlatform(**kwargs) + + def async_client(self, **kwargs: Any) -> AsyncNeMoPlatform: + kwargs.setdefault("http_client", build_async_asgi_http_client(self.app)) + kwargs.setdefault("base_url", EMBEDDED_BASE_URL) + return AsyncNeMoPlatform(**kwargs) + + def start_services(self, service_names: Sequence[str]) -> StartServicesResult: + raise ServicesError("Staged service start is not implemented for embedded mode yet") + + async def start_services_async(self, service_names: Sequence[str]) -> StartServicesResult: + return await asyncio.to_thread(self.start_services, service_names) + + def stop(self, *, timeout: float = 30.0, force: bool = False) -> process.StopResult: + return process.StopResult(stopped_pids=[], swept_children=[]) + + async def stop_async(self, *, timeout: float = 30.0, force: bool = False) -> process.StopResult: + return self.stop(timeout=timeout, force=force) + + +def require_services_extra() -> None: + if importlib.util.find_spec("pyleak") is not None: + return + raise ServicesExtraRequiredError("Install service dependencies with `pip install 'nemo-platform[all]'`.") + + +def _validate_socket_path_length(socket_path: Path) -> None: + encoded_length = len(os.fsencode(socket_path)) + if encoded_length > _AF_UNIX_PATH_MAX_BYTES: + raise ValueError( + "UDS socket path is too long for AF_UNIX " + f"({encoded_length} bytes; maximum is {_AF_UNIX_PATH_MAX_BYTES} bytes): {socket_path}" + ) + + +def _validated_socket_path(config: ServiceRunConfig) -> Path | None: + socket_path = config.resolved_socket_path + if socket_path is None: + return None + _validate_socket_path_length(socket_path) + return socket_path + + +def _prepare_socket(config: ServiceRunConfig) -> Path | None: + socket_path = _validated_socket_path(config) + if socket_path is None: + return None + socket_path.parent.mkdir(parents=True, exist_ok=True) + if not socket_path.exists(): + return socket_path + if probe_status(base_url=UDS_BASE_URL, socket_path=socket_path, timeout=0.5): + raise ServicesAlreadyRunningError(f"UDS socket is live at {socket_path}") + try: + socket_path.unlink() + except OSError as error: + raise ServicesSocketStaleError(f"Could not remove stale socket at {socket_path}") from error + return socket_path + + +def _check_tcp_available(config: ServiceRunConfig) -> None: + conflict = process.check_port_available_for_start( + config.host, + config.port, + config.scope, + base_dir=config.state_root, + ) + if conflict is not None: + raise ServicesPortInUseError("\n".join(process.format_port_conflict(conflict))) + + +def _write_run_request(config: ServiceRunConfig) -> Path: + state_dir = config.to_platform_app_config().state_dir(create=True) + fd, tmp = tempfile.mkstemp(dir=state_dir, suffix=".json") + path = Path(tmp) + try: + with os.fdopen(fd, "w", encoding="utf-8") as file: + fd = -1 + json.dump(config.to_child_payload(), file, indent=2) + file.write("\n") + except BaseException: + if fd >= 0: + os.close(fd) + fd = -1 + path.unlink(missing_ok=True) + raise + finally: + if fd >= 0: + os.close(fd) + return path + + +def _terminate_startup_process(proc: subprocess.Popen) -> None: + if proc.poll() is not None: + return + proc.terminate() + try: + proc.wait(timeout=5) + except subprocess.TimeoutExpired: + proc.kill() + proc.wait(timeout=5) + + +def serve_embedded_app(app: Any, cfg: ServiceRunConfig, socket_path: Path | None) -> None: + import uvicorn + + if socket_path is not None: + from nmp.platform_runner.server import _run_server_on_bound_sockets + + _run_server_on_bound_sockets(app, host=cfg.host, port=cfg.port, socket_path=str(socket_path)) + else: + uvicorn.run(app, host=cfg.host, port=cfg.port, log_config=None) + + +def run_services( + config: ServiceRunConfig | None = None, + *, + _mode: Literal["foreground", "daemon"] = "foreground", + env: MutableMapping[str, str] | None = None, +) -> None: + cfg = config or ServiceRunConfig() + app_config = cfg.to_platform_app_config() + require_services_extra() + if cfg.http_gateway == "enabled": + raise ServicesError("HTTP gateway support is not implemented yet") + if process.is_instance_alive(cfg.scope, base_dir=cfg.state_root): + raise ServicesAlreadyRunningError(f"Instance {cfg.scope!r} is already running") + _check_tcp_available(cfg) + lock_fd = process.acquire_lock(cfg.scope, base_dir=cfg.state_root) + original_data_dir = os.environ.get("NMP_DATA_DIR") + try: + socket_path = _prepare_socket(cfg) + app_config.log_file_path(create_parent=True) + if cfg.data_dir is not None and "NMP_DATA_DIR" not in os.environ: + os.environ["NMP_DATA_DIR"] = str(cfg.data_dir) + desc = process.InstanceDescriptor.from_config( + app_config, + mode=_mode, + transport=cfg.transport, + ) + process.write_descriptor(desc, base_dir=cfg.state_root) + embedded_handle = start_embedded_services(cfg, env=env) + serve_embedded_app(embedded_handle.app, cfg, socket_path) + finally: + try: + process.remove_descriptor(cfg.scope, base_dir=cfg.state_root) + finally: + if original_data_dir is None: + os.environ.pop("NMP_DATA_DIR", None) + else: + os.environ["NMP_DATA_DIR"] = original_data_dir + os.close(lock_fd) + + +def daemonize_services(config: ServiceRunConfig | None = None) -> DaemonServiceHandle: + cfg = config or ServiceRunConfig() + app_config = cfg.to_platform_app_config() + require_services_extra() + if cfg.http_gateway == "enabled": + raise ServicesError("HTTP gateway support is not implemented yet") + if process.is_instance_alive(cfg.scope, base_dir=cfg.state_root): + raise ServicesAlreadyRunningError(f"Instance {cfg.scope!r} is already running") + _check_tcp_available(cfg) + socket_path = _validated_socket_path(cfg) + if ( + socket_path is not None + and socket_path.exists() + and probe_status(base_url=UDS_BASE_URL, socket_path=socket_path, timeout=0.5) + ): + raise ServicesAlreadyRunningError(f"UDS socket is live at {socket_path}") + + request_path = _write_run_request(cfg) + log_path = process.rotate_log_path(app_config.log_file_path()) + log_file = open(log_path, "a") # noqa: SIM115 + env = os.environ.copy() + if cfg.data_dir is not None and "NMP_DATA_DIR" not in env: + env["NMP_DATA_DIR"] = str(cfg.data_dir) + proc: subprocess.Popen | None = None + ownership_transferred = False + try: + try: + child_module = f"{__package__}._service_child" + proc = subprocess.Popen( + [sys.executable, "-m", child_module, str(request_path)], + stdout=log_file, + stderr=subprocess.STDOUT, + env=env, + start_new_session=True, + stdin=subprocess.DEVNULL, + close_fds=True, + ) + finally: + log_file.close() + assert proc is not None + handle = DaemonServiceHandle.from_config(cfg, pid=proc.pid) + deadline = time.monotonic() + cfg.readiness_timeout + while True: + remaining = deadline - time.monotonic() + if remaining <= 0: + break + if proc.poll() is not None: + raise ServicesStartupExitedError(f"Services daemon exited with code {proc.returncode}; log: {log_path}") + if probe_status( + base_url=handle.base_url, + socket_path=handle.socket_path if handle.transport == "uds" else None, + timeout=remaining, + ): + ownership_transferred = True + return handle + remaining = deadline - time.monotonic() + if remaining <= 0: + break + time.sleep(min(cfg.readiness_poll_interval, remaining)) + raise ServicesStartupTimeoutError(f"Timed out waiting for services daemon {cfg.scope!r}; log: {log_path}") + finally: + if proc is not None and not ownership_transferred: + _terminate_startup_process(proc) + + +async def daemonize_services_async(config: ServiceRunConfig | None = None) -> DaemonServiceHandle: + return await asyncio.to_thread(daemonize_services, config) + + +def start_embedded_services( + config: ServiceRunConfig | None = None, + *, + env: MutableMapping[str, str] | None = None, +) -> EmbeddedServiceHandle: + """Start platform services in the current process. + + Args: + env: Environment mapping passed to :func:`build_platform_app`. + Defaults to ``None`` which writes to ``os.environ``. Tests can + pass an empty dict to avoid polluting the process environment. + """ + cfg = config or ServiceRunConfig(mode=ServiceMode.EMBEDDED) + from nmp.platform_runner.server import build_platform_app + + app = build_platform_app( + config=cfg.to_platform_app_config(), + env=env, + ) + runtime = getattr(app.state, "platform_runtime", None) + return EmbeddedServiceHandle(app=app, runtime=runtime) + + +async def start_embedded_services_async(config: ServiceRunConfig | None = None) -> EmbeddedServiceHandle: + return start_embedded_services(config) + + +def get_service_handle(config: ServiceRunConfig | None = None) -> DaemonServiceHandle | None: + cfg = config or ServiceRunConfig() + desc = process.read_descriptor(cfg.scope, base_dir=cfg.state_root) + if desc is None or not process.is_instance_alive(cfg.scope, base_dir=cfg.state_root): + return None + return DaemonServiceHandle.from_descriptor(desc) + + +def list_service_handles(state_dir: str | Path | None = None) -> list[DaemonServiceHandle]: + state_root = Path(state_dir).expanduser() if state_dir is not None else default_state_root() + handles: list[DaemonServiceHandle] = [] + for info in process.list_instances(base_dir=state_root): + if info.descriptor is not None and info.alive: + handles.append(DaemonServiceHandle.from_descriptor(info.descriptor)) + return handles + + +def ensure_services( + config: ServiceRunConfig | None = None, + *, + daemonize: bool | None = None, +) -> LocalServiceHandle: + cfg = config or ServiceRunConfig() + if cfg.mode is ServiceMode.EMBEDDED: + return start_embedded_services(cfg) + + handle = get_service_handle(cfg) + if handle is not None: + return handle + if daemonize is False: + raise ServicesNotRunningError(f"Instance {cfg.scope!r} is not running") + return daemonize_services(cfg) + + +async def ensure_services_async( + config: ServiceRunConfig | None = None, + *, + daemonize: bool | None = None, +) -> LocalServiceHandle: + cfg = config or ServiceRunConfig() + if cfg.mode is ServiceMode.EMBEDDED: + return await start_embedded_services_async(cfg) + + handle = get_service_handle(cfg) + if handle is not None: + return handle + if daemonize is False: + raise ServicesNotRunningError(f"Instance {cfg.scope!r} is not running") + return await daemonize_services_async(cfg) + + +def connect_services( + config: ServiceRunConfig | None = None, + *, + daemonize: bool | None = None, + start_if_needed: bool = True, + **client_kwargs: Any, +) -> NeMoPlatform: + cfg = config or ServiceRunConfig() + if not start_if_needed and cfg.mode is ServiceMode.DAEMON and get_service_handle(cfg) is None: + raise ServicesNotRunningError(f"Instance {cfg.scope!r} is not running") + handle = ensure_services(cfg, daemonize=daemonize) + return handle.client(**client_kwargs) + + +async def connect_services_async( + config: ServiceRunConfig | None = None, + *, + daemonize: bool | None = None, + start_if_needed: bool = True, + **client_kwargs: Any, +) -> AsyncNeMoPlatform: + cfg = config or ServiceRunConfig() + if not start_if_needed and cfg.mode is ServiceMode.DAEMON and get_service_handle(cfg) is None: + raise ServicesNotRunningError(f"Instance {cfg.scope!r} is not running") + handle = await ensure_services_async(cfg, daemonize=daemonize) + return handle.async_client(**client_kwargs) + + +def stop_services( + config: ServiceRunConfig | None = None, + *, + timeout: float = 30.0, + force: bool = False, +) -> process.StopResult: + cfg = config or ServiceRunConfig() + handle = get_service_handle(cfg) + if handle is None: + raise ServicesNotRunningError(f"Instance {cfg.scope!r} is not running") + return handle.stop(timeout=timeout, force=force) + + +async def stop_services_async( + config: ServiceRunConfig | None = None, + *, + timeout: float = 30.0, + force: bool = False, +) -> process.StopResult: + return await asyncio.to_thread(stop_services, config, timeout=timeout, force=force) diff --git a/sdk/python/nemo-platform/src/nemo_platform/local/transport.py b/sdk/python/nemo-platform/src/nemo_platform/local/transport.py new file mode 100644 index 0000000000..da3c08344a --- /dev/null +++ b/sdk/python/nemo-platform/src/nemo_platform/local/transport.py @@ -0,0 +1,154 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Local service transport helpers for TCP and Unix domain sockets.""" + +from __future__ import annotations + +import asyncio +import time +from pathlib import Path +from typing import Any, TypeAlias + +import httpx +from fastapi.testclient import TestClient +from nmp.common.platform_endpoint import UDS_BASE_URL + +HttpxTimeout: TypeAlias = float | httpx.Timeout | None +_DEFAULT_TIMEOUT: float = 5.0 +EMBEDDED_BASE_URL = "http://nemo-platform.local" + +__all__ = [ + "EMBEDDED_BASE_URL", + "UDS_BASE_URL", + "build_async_asgi_http_client", + "build_async_http_client", + "build_sync_asgi_http_client", + "build_sync_http_client", + "probe_status", + "probe_status_async", + "tcp_base_url", + "wait_for_status", + "wait_for_status_async", +] + + +def build_sync_asgi_http_client(app: Any, *, timeout: HttpxTimeout = _DEFAULT_TIMEOUT) -> Any: + _ = timeout + return TestClient( + app, + base_url=EMBEDDED_BASE_URL, + follow_redirects=True, + ) + + +def build_async_asgi_http_client(app: Any, *, timeout: HttpxTimeout = _DEFAULT_TIMEOUT) -> httpx.AsyncClient: + return httpx.AsyncClient( + transport=httpx.ASGITransport(app=app), + base_url=EMBEDDED_BASE_URL, + follow_redirects=True, + timeout=timeout, + ) + + +def build_sync_http_client(socket_path: Path, *, timeout: HttpxTimeout = _DEFAULT_TIMEOUT) -> httpx.Client: + return httpx.Client( + transport=httpx.HTTPTransport(uds=str(socket_path)), + follow_redirects=True, + timeout=timeout, + ) + + +def build_async_http_client(socket_path: Path, *, timeout: HttpxTimeout = _DEFAULT_TIMEOUT) -> httpx.AsyncClient: + return httpx.AsyncClient( + transport=httpx.AsyncHTTPTransport(uds=str(socket_path)), + follow_redirects=True, + timeout=timeout, + ) + + +def tcp_base_url(host: str, port: int) -> str: + connect_host = "localhost" if host in {"0.0.0.0", "::"} else host # noqa: S104 + normalized = connect_host.strip("[]") + url_host = f"[{normalized}]" if ":" in normalized else normalized + return str(httpx.URL(scheme="http", host=url_host, port=port)) + + +def probe_status( + *, + base_url: str, + socket_path: Path | None = None, + timeout: float = 2.0, +) -> bool: + client = ( + build_sync_http_client(socket_path, timeout=timeout) + if socket_path is not None + else httpx.Client(timeout=timeout) + ) + try: + response = client.get(f"{base_url.rstrip('/')}/status") + return response.status_code == 200 + except httpx.RequestError: + return False + finally: + client.close() + + +async def probe_status_async( + *, + base_url: str, + socket_path: Path | None = None, + timeout: float = 2.0, +) -> bool: + client = ( + build_async_http_client(socket_path, timeout=timeout) + if socket_path is not None + else httpx.AsyncClient(timeout=timeout) + ) + try: + response = await client.get(f"{base_url.rstrip('/')}/status") + return response.status_code == 200 + except httpx.RequestError: + return False + finally: + await client.aclose() + + +def wait_for_status( + *, + base_url: str, + socket_path: Path | None = None, + timeout: float = 60.0, + poll_interval: float = 0.5, +) -> bool: + deadline = time.monotonic() + timeout + while True: + remaining = deadline - time.monotonic() + if remaining <= 0: + return False + if probe_status(base_url=base_url, socket_path=socket_path, timeout=remaining): + return True + remaining = deadline - time.monotonic() + if remaining <= 0: + return False + time.sleep(min(poll_interval, remaining)) + + +async def wait_for_status_async( + *, + base_url: str, + socket_path: Path | None = None, + timeout: float = 60.0, + poll_interval: float = 0.5, +) -> bool: + deadline = time.monotonic() + timeout + while True: + remaining = deadline - time.monotonic() + if remaining <= 0: + return False + if await probe_status_async(base_url=base_url, socket_path=socket_path, timeout=remaining): + return True + remaining = deadline - time.monotonic() + if remaining <= 0: + return False + await asyncio.sleep(min(poll_interval, remaining)) diff --git a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/conftest.py b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/conftest.py index a5889a0b4a..bc1a86a21c 100644 --- a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/conftest.py +++ b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/conftest.py @@ -3,7 +3,7 @@ from __future__ import annotations -import nemo_platform.cli.commands.services._process as _process_mod +import nemo_platform.local.process as _process_mod import pytest diff --git a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_services.py b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_services.py index a33ed5213a..ee18d2a9ee 100644 --- a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_services.py +++ b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_services.py @@ -13,11 +13,11 @@ import socket from pathlib import Path from types import ModuleType -from unittest.mock import ANY, MagicMock, patch +from unittest.mock import MagicMock, patch import pytest from nemo_platform.cli.app import app -from nemo_platform.cli.commands.services._process import ( +from nemo_platform.local.process import ( ForegroundInstanceError, InstanceDescriptor, StopResult, @@ -26,16 +26,17 @@ read_descriptor, write_descriptor, ) +from nmp.platform_runner.config import PlatformAppConfig from typer.testing import CliRunner runner = CliRunner() -_PROCESS_MODULE = "nemo_platform.cli.commands.services._process" +_PROCESS_MODULE = "nemo_platform.local.process" _CLI_MODULE = "nemo_platform.cli.commands.services.cli" def _seed_stopped_scope(base_dir: Path, scope: str, *, log_content: str = "x\n") -> Path: - """Create a stopped instance directory with service logs.""" + """Create a stopped scope directory with service logs.""" d = instance_dir(scope, base_dir=base_dir) (d / "services.log").write_text(log_content) return d @@ -118,17 +119,18 @@ def test_run_invokes_runner(base_dir: Path): ) assert result.exit_code == 0, result.stderr - mock_run_platform.assert_called_once_with( - services=["auth", "entities"], - service_group=None, - controllers=["jobs", "models"], - controller_group=None, - sidecars=None, - config_path=None, - host="127.0.0.1", - port=9000, - on_shutdown=ANY, - ) + mock_run_platform.assert_called_once() + _, kwargs = mock_run_platform.call_args + config = kwargs["config"] + assert config.services == ["auth", "entities"] + assert config.service_group is None + assert config.controllers == ["jobs", "models"] + assert config.controller_group is None + assert config.sidecars is None + assert config.config_path is None + assert config.host == "127.0.0.1" + assert config.port == 9000 + assert kwargs["on_shutdown"] is not None def test_run_refuses_when_already_running(base_dir: Path): @@ -163,7 +165,7 @@ def test_run_writes_descriptor(base_dir: Path): desc = read_descriptor("desc-test", base_dir=base_dir) assert desc is not None assert desc.mode == "foreground" - assert desc.port == 9999 + assert desc.config.port == 9999 def test_run_records_background_mode_when_launched_by_start(base_dir: Path): @@ -382,7 +384,7 @@ def test_restart_errors_when_no_prior_instance(self, base_dir: Path): ["services", "restart", "--instance", "ghost"], ) assert result.exit_code == 1 - assert "No instance found" in result.stderr + assert "No instance found for scope" in result.stderr assert "nemo services start" in result.stderr def test_restart_stops_and_starts(self, base_dir: Path): @@ -390,9 +392,7 @@ def test_restart_stops_and_starts(self, base_dir: Path): fd = acquire_lock(scope, base_dir=base_dir) desc = InstanceDescriptor( pid=os.getpid(), - scope=scope, - host="127.0.0.1", - port=8080, + config=PlatformAppConfig(scope=scope), mode="background", create_time=1.0, ) @@ -429,9 +429,7 @@ def test_restart_exits_early_when_port_occupied_by_foreign_process(self, base_di desc = InstanceDescriptor( pid=99999, - scope=scope, - host="127.0.0.1", - port=port, + config=PlatformAppConfig(scope=scope, host="127.0.0.1", port=port), mode="background", create_time=1.0, ) @@ -459,13 +457,15 @@ def test_restart_preserves_previous_args(self, base_dir: Path): fd = acquire_lock(scope, base_dir=base_dir) desc = InstanceDescriptor( pid=os.getpid(), - scope=scope, - host="127.0.0.1", - port=9000, + config=PlatformAppConfig( + scope=scope, + services=["entities", "models"], + controllers=["jobs"], + host="127.0.0.1", + port=9000, + ), mode="background", create_time=1.0, - services=["entities", "models"], - controllers=["jobs"], ) write_descriptor(desc, base_dir=base_dir) @@ -488,11 +488,12 @@ def test_restart_preserves_previous_args(self, base_dir: Path): os.close(fd) assert result.exit_code == 0 - _, kwargs = mock_start.call_args - assert kwargs["services"] == ["entities", "models"] - assert kwargs["controllers"] == ["jobs"] - assert kwargs["host"] == "127.0.0.1" - assert kwargs["port"] == 9000 + args, _kwargs = mock_start.call_args + config = args[0] + assert config.services == ["entities", "models"] + assert config.controllers == ["jobs"] + assert config.host == "127.0.0.1" + assert config.port == 9000 # --------------------------------------------------------------------------- @@ -504,16 +505,14 @@ class TestServicesStatus: def test_not_running(self, base_dir: Path): result = runner.invoke(app, ["services", "status", "--instance", "none"]) assert result.exit_code == 0 - assert "No running instance" in result.stdout + assert "No running instance for scope" in result.stdout def test_running_instance(self, base_dir: Path): scope = "status-test" fd = acquire_lock(scope, base_dir=base_dir) desc = InstanceDescriptor( pid=os.getpid(), - scope=scope, - host="127.0.0.1", - port=8080, + config=PlatformAppConfig(scope=scope), mode="foreground", create_time=1.0, ) @@ -548,9 +547,7 @@ def test_lists_running_instance(self, base_dir: Path): fd = acquire_lock(scope, base_dir=base_dir) desc = InstanceDescriptor( pid=os.getpid(), - scope=scope, - host="127.0.0.1", - port=8080, + config=PlatformAppConfig(scope=scope), mode="background", create_time=1.0, ) @@ -594,9 +591,7 @@ def test_mixed_running_and_stopped(self, base_dir: Path): write_descriptor( InstanceDescriptor( pid=os.getpid(), - scope=running_scope, - host="127.0.0.1", - port=8080, + config=PlatformAppConfig(scope=running_scope), mode="background", create_time=1.0, ), @@ -658,7 +653,7 @@ def test_rm_requires_scope(self, base_dir: Path): def test_rm_rejects_invalid_scope(self, base_dir: Path): result = runner.invoke(app, ["services", "rm", "../escape"]) assert result.exit_code == 1 - assert "Invalid instance scope" in result.stderr + assert "Invalid scope" in result.stderr def test_rm_rejects_conflicting_scope_args(self, base_dir: Path): result = runner.invoke(app, ["services", "rm", "scope-a", "--instance", "scope-b"]) @@ -747,7 +742,7 @@ def test_default_host_is_loopback(base_dir: Path): assert result.exit_code == 0, result.stderr _, kwargs = mock_run_platform.call_args - assert kwargs["host"] == "127.0.0.1" + assert kwargs["config"].host == "127.0.0.1" def test_bind_all_warning(base_dir: Path): diff --git a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_services_lifecycle.py b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_services_lifecycle.py index 2d035e1355..1ee89693bf 100644 --- a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_services_lifecycle.py +++ b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_services_lifecycle.py @@ -27,7 +27,7 @@ import pytest from nemo_platform.cli.app import app -from nemo_platform.cli.commands.services._process import ( +from nemo_platform.local.process import ( InstanceDescriptor, PortConflict, acquire_lock, @@ -43,6 +43,7 @@ stop_instance, write_descriptor, ) +from nmp.platform_runner.config import PlatformAppConfig from typer.testing import CliRunner _runner = CliRunner() @@ -70,18 +71,14 @@ import psutil as _psutil desc = { "pid": os.getpid(), - "scope": scope, - "host": "127.0.0.1", - "port": 8080, + "config": { + "scope": scope, + "host": "127.0.0.1", + "port": 8080, + }, "mode": "background", "create_time": _psutil.Process(os.getpid()).create_time(), "started_at": "test", - "services": None, - "controllers": None, - "service_group": None, - "controller_group": None, - "sidecars": None, - "config_path": None, "log_path": None, } desc_path = os.path.join(inst_dir, "instance.json") @@ -291,9 +288,7 @@ def test_stale_descriptor_with_reused_pid(self, tmp_path: Path) -> None: # Write a descriptor with the sleeper's PID but wrong create_time desc = InstanceDescriptor( pid=sleeper.pid, - scope=scope, - host="127.0.0.1", - port=8080, + config=PlatformAppConfig(scope=scope), mode="background", create_time=0.0, # intentionally wrong ) @@ -363,7 +358,7 @@ def test_log_preserved_across_restart(self, tmp_path: Path) -> None: log = d / "services.log" log.write_text("first boot log content\n") - from nemo_platform.cli.commands.services._process import rotate_log + from nemo_platform.local.process import rotate_log new_log = rotate_log(scope, base_dir=base_dir) new_log.write_text("second boot log content\n") @@ -399,15 +394,15 @@ def test_log_preserved_across_restart(self, tmp_path: Path) -> None: import psutil as _psutil desc = { "pid": os.getpid(), - "scope": scope, - "host": "127.0.0.1", - "port": port, + "config": { + "scope": scope, + "host": "127.0.0.1", + "port": port, + }, "mode": "background", "create_time": _psutil.Process(os.getpid()).create_time(), "started_at": "test", - "services": None, "controllers": None, - "service_group": None, "controller_group": None, - "sidecars": None, "config_path": None, "log_path": None, + "log_path": None, } desc_path = os.path.join(inst_dir, "instance.json") with open(desc_path, "w") as f: @@ -523,7 +518,7 @@ def test_stop_after_health_check(self, tmp_path: Path) -> None: assert is_instance_alive(scope, base_dir=base_dir) desc = read_descriptor(scope, base_dir=base_dir) assert desc is not None - assert desc.port == port + assert desc.config.port == port result = stop_instance(scope, base_dir=base_dir, timeout=5.0) assert proc.pid in result.stopped_pids @@ -536,7 +531,7 @@ def test_stop_after_health_check(self, tmp_path: Path) -> None: class TestInstanceCleanup: - """Integration tests for rm/prune and post-stop instance directories.""" + """Integration tests for rm/prune and post-stop scope directories.""" def test_stop_leaves_record_until_rm(self, tmp_path: Path, monkeypatch) -> None: base_dir = tmp_path / "state" @@ -668,9 +663,7 @@ def test_check_port_returns_nemo_instance_when_lock_held_and_port_blocked(self, write_descriptor( InstanceDescriptor( pid=os.getpid(), - scope=scope, - host="127.0.0.1", - port=port, + config=PlatformAppConfig(scope=scope, host="127.0.0.1", port=port), mode="background", create_time=1.0, ), @@ -704,9 +697,7 @@ def test_check_port_returns_foreign_when_alive_instance_uses_different_port(self write_descriptor( InstanceDescriptor( pid=os.getpid(), - scope=scope, - host="127.0.0.1", - port=nemo_port, + config=PlatformAppConfig(scope=scope, host="127.0.0.1", port=nemo_port), mode="background", create_time=1.0, ), diff --git a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_services_process.py b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_services_process.py index ae86b6b9f7..6008aed5c9 100644 --- a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_services_process.py +++ b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_services_process.py @@ -6,6 +6,7 @@ from __future__ import annotations import os +import signal import subprocess import sys import time @@ -14,7 +15,8 @@ import psutil import pytest -from nemo_platform.cli.commands.services._process import ( +from nemo_platform.local import process as process_module +from nemo_platform.local.process import ( ForegroundInstanceError, InstanceAlreadyRunningError, InstanceDescriptor, @@ -40,6 +42,7 @@ validate_pid, write_descriptor, ) +from nmp.platform_runner.config import DEFAULT_LOCAL_SERVICES_BIND_HOST, PlatformAppConfig @pytest.fixture() @@ -48,67 +51,23 @@ def base_dir(tmp_path: Path) -> Path: # --------------------------------------------------------------------------- -# Scope computation +# Scope resolution # --------------------------------------------------------------------------- class TestComputeScope: - def test_explicit_instance_name(self) -> None: - assert compute_scope(port=8080, instance_name="myapp") == "myapp" + def test_explicit_scope(self) -> None: + assert compute_scope(port=1234, explicit_scope="myapp") == "myapp" - def test_default_scope_includes_port(self) -> None: - scope = compute_scope(port=9090) - assert scope.endswith("-9090") - - def test_default_scope_is_deterministic(self) -> None: - a = compute_scope(port=8080) - b = compute_scope(port=8080) - assert a == b - - def test_different_ports_different_scopes(self) -> None: - a = compute_scope(port=8080) - b = compute_scope(port=9090) - assert a != b - - def test_hash_prefix_is_8_chars(self) -> None: + def test_default_scope_is_stable_for_port(self) -> None: scope = compute_scope(port=8080) - prefix = scope.rsplit("-", 1)[0] - assert len(prefix) == 8 - - def test_git_failure_falls_back_to_cwd(self) -> None: - import nemo_platform.cli.commands.services._process as proc_mod - - proc_mod._scope_prefix_cache = None - try: - with patch.object(proc_mod, "_find_git_root", return_value="/no/git/here"): - scope = compute_scope(port=8080) - assert scope.endswith("-8080") - assert len(scope.rsplit("-", 1)[0]) == 8 - finally: - proc_mod._scope_prefix_cache = None - def test_different_git_roots_produce_different_prefixes(self) -> None: - """Two different working directories (worktrees) produce distinct scopes.""" - import nemo_platform.cli.commands.services._process as proc_mod - - with patch.object(proc_mod, "_find_git_root", return_value="/workspace/project-a"): - scope_a = compute_scope(port=8080) - - proc_mod._scope_prefix_cache = None - - with patch.object(proc_mod, "_find_git_root", return_value="/workspace/project-b"): - scope_b = compute_scope(port=8080) - - assert scope_a != scope_b - assert scope_a.endswith("-8080") - assert scope_b.endswith("-8080") - prefix_a = scope_a.rsplit("-", 1)[0] - prefix_b = scope_b.rsplit("-", 1)[0] - assert prefix_a != prefix_b + assert scope == compute_scope(port=8080) + assert scope.endswith("-8080") # --------------------------------------------------------------------------- -# Instance directory +# Scope directory # --------------------------------------------------------------------------- @@ -168,26 +127,27 @@ class TestDescriptorRoundTrip: def test_write_and_read(self, base_dir: Path) -> None: desc = InstanceDescriptor( pid=12345, - scope="test-8080", - host="127.0.0.1", - port=8080, + config=PlatformAppConfig( + scope="test-8080", + services=["entities", "models"], + controllers=["jobs"], + host="127.0.0.1", + ), mode="background", create_time=1000.0, - services=["entities", "models"], - controllers=["jobs"], ) write_descriptor(desc, base_dir=base_dir) recovered = read_descriptor("test-8080", base_dir=base_dir) assert recovered is not None assert recovered.pid == 12345 - assert recovered.scope == "test-8080" - assert recovered.host == "127.0.0.1" - assert recovered.port == 8080 + assert recovered.config.scope == "test-8080" + assert recovered.config.host == "127.0.0.1" + assert recovered.config.port == 8080 assert recovered.mode == "background" assert recovered.create_time == 1000.0 - assert recovered.services == ["entities", "models"] - assert recovered.controllers == ["jobs"] + assert recovered.config.services == ["entities", "models"] + assert recovered.config.controllers == ["jobs"] def test_read_missing_returns_none(self, base_dir: Path) -> None: assert read_descriptor("no-such-scope", base_dir=base_dir) is None @@ -200,9 +160,7 @@ def test_read_corrupt_returns_none(self, base_dir: Path) -> None: def test_remove_descriptor(self, base_dir: Path) -> None: desc = InstanceDescriptor( pid=1, - scope="rm-test", - host="127.0.0.1", - port=8080, + config=PlatformAppConfig(scope="rm-test"), mode="background", create_time=1.0, ) @@ -245,9 +203,7 @@ def test_lists_alive_instance(self, base_dir: Path) -> None: fd = acquire_lock("alive-one", base_dir=base_dir) desc = InstanceDescriptor( pid=os.getpid(), - scope="alive-one", - host="127.0.0.1", - port=8080, + config=PlatformAppConfig(scope="alive-one"), mode="foreground", create_time=1.0, ) @@ -265,9 +221,7 @@ def test_cleans_up_dead_descriptor(self, base_dir: Path) -> None: d = instance_dir("dead-scope", base_dir=base_dir) desc = InstanceDescriptor( pid=999999, - scope="dead-scope", - host="127.0.0.1", - port=8080, + config=PlatformAppConfig(scope="dead-scope"), mode="background", create_time=1.0, ) @@ -281,9 +235,7 @@ def test_stale_descriptor_with_logs_stays_listed(self, base_dir: Path) -> None: d = instance_dir("dead-with-logs", base_dir=base_dir) desc = InstanceDescriptor( pid=999999, - scope="dead-with-logs", - host="127.0.0.1", - port=8080, + config=PlatformAppConfig(scope="dead-with-logs"), mode="background", create_time=1.0, ) @@ -370,7 +322,7 @@ def test_refuses_running_instance(self, base_dir: Path) -> None: os.close(fd) def test_rejects_invalid_scope(self, base_dir: Path) -> None: - with pytest.raises(ValueError, match="Invalid instance scope"): + with pytest.raises(ValueError, match="Invalid scope"): remove_instance("../escape", base_dir=base_dir) def test_returns_false_when_rmtree_fails(self, base_dir: Path) -> None: @@ -378,7 +330,7 @@ def test_returns_false_when_rmtree_fails(self, base_dir: Path) -> None: (d / "services.log").write_text("logs\n") with patch( - "nemo_platform.cli.commands.services._process.shutil.rmtree", + "nemo_platform.local.process.shutil.rmtree", side_effect=OSError("permission denied"), ): assert remove_instance("rmtree-fail", base_dir=base_dir) is False @@ -489,9 +441,7 @@ def test_stops_running_process(self, base_dir: Path) -> None: fd = acquire_lock(scope, base_dir=base_dir) desc = InstanceDescriptor( pid=proc.pid, - scope=scope, - host="127.0.0.1", - port=8080, + config=PlatformAppConfig(scope=scope), mode="background", create_time=psutil.Process(proc.pid).create_time(), ) @@ -512,9 +462,7 @@ def test_cleans_up_stale_descriptor(self, base_dir: Path) -> None: scope = "stale" desc = InstanceDescriptor( pid=999999999, - scope=scope, - host="127.0.0.1", - port=8080, + config=PlatformAppConfig(scope=scope), mode="background", create_time=0.0, ) @@ -529,9 +477,7 @@ def test_refuses_to_stop_foreground_instance(self, base_dir: Path) -> None: try: desc = InstanceDescriptor( pid=os.getpid(), - scope=scope, - host="127.0.0.1", - port=8080, + config=PlatformAppConfig(scope=scope), mode="foreground", create_time=1.0, ) @@ -555,9 +501,7 @@ def test_force_stops_foreground_instance(self, base_dir: Path) -> None: try: desc = InstanceDescriptor( pid=proc.pid, - scope=scope, - host="127.0.0.1", - port=8080, + config=PlatformAppConfig(scope=scope), mode="foreground", create_time=psutil.Process(proc.pid).create_time(), ) @@ -571,6 +515,34 @@ def test_force_stops_foreground_instance(self, base_dir: Path) -> None: proc.kill() proc.wait(timeout=5) + def test_preserves_descriptor_when_sigkill_does_not_stop_parent(self, base_dir: Path, monkeypatch) -> None: + scope = "sigkill-still-alive" + desc = InstanceDescriptor( + pid=12345, + config=PlatformAppConfig(scope=scope), + mode="background", + create_time=1.0, + ) + write_descriptor(desc, base_dir=base_dir) + kill_signals: list[int] = [] + + def fake_kill(_pid: int, sig: int) -> None: + kill_signals.append(sig) + + monkeypatch.setattr(process_module, "validate_pid", lambda _pid, _create_time: True) + monkeypatch.setattr(process_module, "_pid_alive", lambda _pid: True) + monkeypatch.setattr(process_module, "_snapshot_children", lambda _pid: [object()]) + monkeypatch.setattr(process_module, "_sweep_orphans", lambda _children: [222]) + monkeypatch.setattr(process_module, "_SIGKILL_WAIT_TIMEOUT", 0.0) + monkeypatch.setattr(process_module.os, "kill", fake_kill) + + result = stop_instance(scope, base_dir=base_dir, timeout=0.0) + + assert kill_signals == [signal.SIGTERM, signal.SIGKILL] + assert result.stopped_pids == [] + assert result.swept_children == [222] + assert read_descriptor(scope, base_dir=base_dir) is not None + # --------------------------------------------------------------------------- # start_background @@ -578,21 +550,49 @@ def test_force_stops_foreground_instance(self, base_dir: Path) -> None: class TestStartBackground: + def test_uses_default_platform_app_config(self, tmp_path: Path, monkeypatch) -> None: + monkeypatch.setenv("XDG_STATE_HOME", str(tmp_path)) + mock_proc = MagicMock() + mock_proc.pid = 99998 + captured_args: list[str] = [] + captured_env: dict[str, str] = {} + + def fake_popen(args, **kwargs): + captured_args.extend(args) + captured_env.update(kwargs["env"]) + return mock_proc + + with patch( + "nemo_platform.local.process.subprocess.Popen", + side_effect=fake_popen, + ): + proc = start_background() + + assert proc.pid == 99998 + assert captured_args[captured_args.index("--instance") + 1] == "default" + assert captured_args[captured_args.index("--host") + 1] == "127.0.0.1" + assert captured_args[captured_args.index("--port") + 1] == "8080" + assert captured_env["XDG_STATE_HOME"] == str(tmp_path) + assert "_NMP_STATE_DIR" not in captured_env + assert (tmp_path / "nmp" / "instances" / "default" / "services.log").exists() + def test_launches_detached_subprocess(self, base_dir: Path) -> None: mock_proc = MagicMock() mock_proc.pid = 99999 with patch( - "nemo_platform.cli.commands.services._process.subprocess.Popen", + "nemo_platform.local.process.subprocess.Popen", return_value=mock_proc, ) as mock_popen: proc = start_background( - scope="bg-test", - services=["entities", "models"], - controllers=["jobs"], - host="127.0.0.1", - port=8080, - base_dir=base_dir, + PlatformAppConfig( + scope="bg-test", + services=["entities", "models"], + controllers=["jobs"], + host="127.0.0.1", + port=8080, + state_root=base_dir, + ), ) assert proc.pid == 99999 @@ -612,13 +612,16 @@ def fake_popen(args, **kwargs): return mock_proc with patch( - "nemo_platform.cli.commands.services._process.subprocess.Popen", + "nemo_platform.local.process.subprocess.Popen", side_effect=fake_popen, ): start_background( - scope="data-dir-test", + PlatformAppConfig( + scope="data-dir-test", + host=DEFAULT_LOCAL_SERVICES_BIND_HOST, + state_root=base_dir, + ), data_dir="/chosen/data/dir", - base_dir=base_dir, ) assert captured_env.get("NMP_DATA_DIR") == "/chosen/data/dir" @@ -634,13 +637,16 @@ def fake_popen(args, **kwargs): return mock_proc with patch( - "nemo_platform.cli.commands.services._process.subprocess.Popen", + "nemo_platform.local.process.subprocess.Popen", side_effect=fake_popen, ): start_background( - scope="shell-env-test", + PlatformAppConfig( + scope="shell-env-test", + host=DEFAULT_LOCAL_SERVICES_BIND_HOST, + state_root=base_dir, + ), data_dir="/chosen/data/dir", - base_dir=base_dir, ) assert captured_env.get("NMP_DATA_DIR") == "/shell/wins" @@ -654,16 +660,22 @@ def test_rotates_log_before_start(self, base_dir: Path) -> None: mock_proc.pid = 5555 with patch( - "nemo_platform.cli.commands.services._process.subprocess.Popen", + "nemo_platform.local.process.subprocess.Popen", return_value=mock_proc, ): - start_background(scope="rotate-test", base_dir=base_dir) + start_background( + PlatformAppConfig( + scope="rotate-test", + host=DEFAULT_LOCAL_SERVICES_BIND_HOST, + state_root=base_dir, + ), + ) rotated = list(d.glob("services.log.*")) assert len(rotated) == 1 assert rotated[0].read_text() == "old log content\n" - def test_forwards_instance_scope_to_child(self, base_dir: Path) -> None: + def test_forwards_scope_to_child(self, base_dir: Path) -> None: mock_proc = MagicMock() mock_proc.pid = 7777 captured_args: list[str] = [] @@ -673,20 +685,22 @@ def fake_popen(args, **kwargs): return mock_proc with patch( - "nemo_platform.cli.commands.services._process.subprocess.Popen", + "nemo_platform.local.process.subprocess.Popen", side_effect=fake_popen, ): start_background( - scope="custom-scope", - services=["entities"], - host="127.0.0.1", - port=9090, - base_dir=base_dir, + PlatformAppConfig( + scope="custom-key", + services=["entities"], + host="127.0.0.1", + port=9090, + state_root=base_dir, + ), ) assert "--instance" in captured_args idx = captured_args.index("--instance") - assert captured_args[idx + 1] == "custom-scope" + assert captured_args[idx + 1] == "custom-key" def test_sets_launch_mode_background_in_child_env(self, base_dir: Path) -> None: mock_proc = MagicMock() @@ -698,10 +712,16 @@ def fake_popen(args, **kwargs): return mock_proc with patch( - "nemo_platform.cli.commands.services._process.subprocess.Popen", + "nemo_platform.local.process.subprocess.Popen", side_effect=fake_popen, ): - start_background(scope="mode-test", base_dir=base_dir) + start_background( + PlatformAppConfig( + scope="mode-test", + host=DEFAULT_LOCAL_SERVICES_BIND_HOST, + state_root=base_dir, + ), + ) assert captured_env.get("_NMP_LAUNCH_MODE") == "background" @@ -846,9 +866,7 @@ def test_sweeps_surviving_children(self, base_dir: Path) -> None: fd = acquire_lock(scope, base_dir=base_dir) desc = InstanceDescriptor( pid=parent.pid, - scope=scope, - host="127.0.0.1", - port=8080, + config=PlatformAppConfig(scope=scope), mode="background", create_time=psutil.Process(parent.pid).create_time(), ) @@ -883,9 +901,7 @@ def test_swept_children_empty_when_no_children(self, base_dir: Path) -> None: fd = acquire_lock(scope, base_dir=base_dir) desc = InstanceDescriptor( pid=proc.pid, - scope=scope, - host="127.0.0.1", - port=8080, + config=PlatformAppConfig(scope=scope), mode="background", create_time=psutil.Process(proc.pid).create_time(), ) @@ -917,9 +933,7 @@ def test_restart_path_sweeps_children(self, base_dir: Path) -> None: fd = acquire_lock(scope, base_dir=base_dir) desc = InstanceDescriptor( pid=parent.pid, - scope=scope, - host="127.0.0.1", - port=8080, + config=PlatformAppConfig(scope=scope), mode="foreground", create_time=psutil.Process(parent.pid).create_time(), ) diff --git a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_setup.py b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_setup.py index 0998f9e6f5..84a99ee90d 100644 --- a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_setup.py +++ b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/cli/commands/test_setup.py @@ -14,7 +14,6 @@ import typer from click.exceptions import Exit as ClickExit from nemo_platform.resources.inference.providers import ProvidersResource -from nemo_platform.cli.commands.services._process import PortConflict from nemo_platform.cli.commands.setup import ( _AGENT_API_READINESS_POLL_INTERVAL, _AGENT_DEPLOY_POLL_INTERVAL, @@ -77,6 +76,7 @@ Context, ContextDefinition, ) +from nemo_platform.local.process import PortConflict from nemo_platform_plugin.client.errors import NotFoundError from nemo_platform_plugin.secrets.types import PlatformSecretCreateRequest, PlatformSecretUpdateRequest @@ -644,9 +644,10 @@ def test_start_services_background_forwards_data_dir(self): mock_start.return_value = MagicMock(pid=42) _start_services_background("http://localhost:9090", data_dir="/chosen/data/dir") mock_start.assert_called_once() - _, kwargs = mock_start.call_args + args, kwargs = mock_start.call_args + config = args[0] assert kwargs["data_dir"] == "/chosen/data/dir" - assert kwargs["port"] == 9090 + assert config.port == 9090 def test_auto_mode_skips_prompt_and_uses_persisted(self, tmp_path, monkeypatch): """`--auto` must not prompt but should still honor any persisted data dir.""" diff --git a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/local/__init__.py b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/local/__init__.py new file mode 100644 index 0000000000..1275d78dff --- /dev/null +++ b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/local/__init__.py @@ -0,0 +1,15 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + diff --git a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/local/test_config_environment.py b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/local/test_config_environment.py new file mode 100644 index 0000000000..8255b3b6be --- /dev/null +++ b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/local/test_config_environment.py @@ -0,0 +1,196 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Integration tests for configuration and environment resolution. + +These tests exercise the real ``apply_run_environment`` code path with +actual YAML config files, verifying that environment variables are set +correctly for different host, port, and base_url scenarios. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest +from nmp.platform_runner.config import ( + ResolvedRunConfiguration, + apply_run_environment, + default_config_path, +) + + +def _resolved( + *, + services: set[str] | None = None, + controllers: set[str] | None = None, + sidecars: set[str] | None = None, + host: str = "127.0.0.1", + port: int = 8080, + config_path: str | None = None, + socket_path: str | None = None, +) -> ResolvedRunConfiguration: + return ResolvedRunConfiguration( + services=services or set(), + controllers=controllers or set(), + sidecars=sidecars or set(), + host=host, + port=port, + config_path=config_path or default_config_path(), + socket_path=socket_path, + available_services={}, + available_controllers={}, + ) + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +@pytest.mark.integration +def test_apply_run_environment_preserves_external_base_url() -> None: + """Pre-set NMP_BASE_URL (e.g. from k8s/Helm) must not be overwritten.""" + env: dict[str, str] = {"NMP_BASE_URL": "https://platform.k8s.internal:443"} + config = _resolved(host="0.0.0.0", port=9090) + + apply_run_environment(config, env=env) + + assert env["NMP_BASE_URL"] == "https://platform.k8s.internal:443" + + +@pytest.mark.integration +def test_apply_run_environment_wildcard_host_becomes_loopback(tmp_path: Path) -> None: + """A wildcard bind host (0.0.0.0) in the config file should resolve to + 127.0.0.1 for the base URL, using the actual bind port.""" + config_file = tmp_path / "config.yaml" + config_file.write_text("platform:\n base_url: http://0.0.0.0:8080\n") + + env: dict[str, str] = {} + config = _resolved(host="0.0.0.0", port=9090, config_path=str(config_file)) + + apply_run_environment(config, env=env) + + assert env["NMP_BASE_URL"] == "http://127.0.0.1:9090" + assert env["NMP_SERVICE_HOST"] == "127.0.0.1" + assert env["NMP_SERVICE_PORT"] == "9090" + + +@pytest.mark.integration +def test_apply_run_environment_ipv6_literal_bracketed(tmp_path: Path) -> None: + """An IPv6 config base_url should produce a bracketed host in the resolved URL.""" + config_file = tmp_path / "config.yaml" + config_file.write_text("platform:\n base_url: http://[::1]:8080\n") + + env: dict[str, str] = {} + config = _resolved(host="::1", port=9090, config_path=str(config_file)) + + apply_run_environment(config, env=env) + + assert env["NMP_BASE_URL"] == "http://[::1]:9090" + + +@pytest.mark.integration +def test_config_file_base_url_malformed_yaml_falls_back(tmp_path: Path) -> None: + """A corrupt config file should fall back to the bind-derived URL.""" + config_file = tmp_path / "config.yaml" + config_file.write_text("{{{{not valid yaml at all") + + env: dict[str, str] = {} + config = _resolved(host="127.0.0.1", port=7777, config_path=str(config_file)) + + apply_run_environment(config, env=env) + + # Falls back to bind-derived: http://: + assert env["NMP_BASE_URL"] == "http://127.0.0.1:7777" + + +@pytest.mark.integration +def test_config_file_missing_falls_back(tmp_path: Path) -> None: + """A missing config file should fall back to the bind-derived URL.""" + env: dict[str, str] = {} + config = _resolved(host="127.0.0.1", port=5555, config_path=str(tmp_path / "nonexistent.yaml")) + + apply_run_environment(config, env=env) + + assert env["NMP_BASE_URL"] == "http://127.0.0.1:5555" + + +@pytest.mark.integration +def test_apply_run_environment_clears_empty_service_lists() -> None: + """When services/controllers/sidecars are empty sets, their env vars + should be removed (popped) rather than set to empty strings.""" + env: dict[str, str] = { + "NMP_SERVICES": "old-service", + "NMP_CONTROLLERS": "old-controller", + "NMP_SIDECARS": "old-sidecar", + } + config = _resolved(services=set(), controllers=set(), sidecars=set()) + + apply_run_environment(config, env=env) + + assert "NMP_SERVICES" not in env + assert "NMP_CONTROLLERS" not in env + assert "NMP_SIDECARS" not in env + + +@pytest.mark.integration +def test_apply_run_environment_sets_service_lists() -> None: + """Non-empty service/controller/sidecar sets should be written as + comma-separated, sorted env var values.""" + env: dict[str, str] = {} + config = _resolved( + services={"models", "auth", "secrets"}, + controllers={"beta-controller"}, + sidecars={"adapters"}, + ) + + apply_run_environment(config, env=env) + + assert env["NMP_SERVICES"] == "auth,models,secrets" + assert env["NMP_CONTROLLERS"] == "beta-controller" + assert env["NMP_SIDECARS"] == "adapters" + + +@pytest.mark.integration +def test_apply_run_environment_uds_transport_uses_unix_base_url() -> None: + """When a socket_path is set (UDS transport), the base URL should use + the ``unix://`` scheme.""" + env: dict[str, str] = {} + config = _resolved(socket_path="/tmp/nemo.sock") + + apply_run_environment(config, env=env) + + assert env["NMP_BASE_URL"] == "unix:///tmp/nemo.sock" + + +@pytest.mark.integration +def test_apply_run_environment_preserves_external_host_and_port() -> None: + """Pre-set NMP_SERVICE_HOST and NMP_SERVICE_PORT should not be overwritten.""" + env: dict[str, str] = { + "NMP_SERVICE_HOST": "10.0.0.1", + "NMP_SERVICE_PORT": "443", + } + config = _resolved(host="0.0.0.0", port=9090) + + apply_run_environment(config, env=env) + + assert env["NMP_SERVICE_HOST"] == "10.0.0.1" + assert env["NMP_SERVICE_PORT"] == "443" + + +@pytest.mark.integration +def test_apply_run_environment_ipv6_wildcard_becomes_loopback(tmp_path: Path) -> None: + """The IPv6 wildcard ``::`` should resolve to ``::1`` for internal clients.""" + # Use a config file without platform.base_url so the bind host drives the URL. + config_file = tmp_path / "config.yaml" + config_file.write_text("platform:\n seed_on_startup: false\n") + + env: dict[str, str] = {} + config = _resolved(host="::", port=8080, config_path=str(config_file)) + + apply_run_environment(config, env=env) + + assert env["NMP_SERVICE_HOST"] == "::1" + # Base URL should have bracketed IPv6. + assert env["NMP_BASE_URL"] == "http://[::1]:8080" diff --git a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/local/test_daemon_lifecycle.py b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/local/test_daemon_lifecycle.py new file mode 100644 index 0000000000..d38ac7c8eb --- /dev/null +++ b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/local/test_daemon_lifecycle.py @@ -0,0 +1,503 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Integration tests for daemon subprocess lifecycle. + +These tests spawn REAL child processes via ``daemonize_services()``, exercise +real lock acquisition, descriptor file I/O, HTTP readiness probing, and +graceful shutdown via ``stop_instance()``. Nothing is monkeypatched away — +the child runs a real uvicorn server with the ``hello-world`` service. + +Requirements: +- All packages installed (``uv sync --all-packages``) so entry-point + discovery finds hello-world. +- ``pyleak`` importable (from the ``[all]`` extra). +""" + +from __future__ import annotations + +import os +import socket +import subprocess +import sys +import tempfile +import time +from pathlib import Path + +import psutil +import pytest +from nemo_platform.local import process, services +from nemo_platform.local.process import ForegroundInstanceError +from nemo_platform.local.services import ( + ServiceRunConfig, + ServicesAlreadyRunningError, + ServicesStartupExitedError, +) +from nmp.platform_runner.config import PlatformAppConfig + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _free_tcp_port() -> int: + """Bind to port 0, let the OS pick, then release and return the port number.""" + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +def _daemon_config( + tmp_path: Path, + *, + scope: str = "integ-daemon", + port: int | None = None, +) -> ServiceRunConfig: + """Build a ServiceRunConfig that is fully isolated under ``tmp_path``.""" + return ServiceRunConfig( + mode=services.ServiceMode.DAEMON, + services=("hello-world",), + controllers=(), + sidecars=(), + transport="tcp", + host="127.0.0.1", + port=port or _free_tcp_port(), + scope=scope, + state_dir=tmp_path / "state", + runtime_dir=tmp_path / "runtime", + readiness_timeout=30.0, + readiness_poll_interval=0.3, + ) + + +def _ensure_stopped(cfg: ServiceRunConfig) -> None: + """Best-effort cleanup: stop any instance left running by a test.""" + try: + process.stop_instance(cfg.scope, base_dir=cfg.state_root, timeout=10, force=True) + except Exception: + pass + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +@pytest.mark.integration +def test_daemonize_services_spawns_child_that_becomes_ready(tmp_path: Path) -> None: + """Spawn a real daemon subprocess, verify readiness via HTTP, then + gracefully shut down with ``stop_instance``.""" + cfg = _daemon_config(tmp_path) + handle = None + try: + handle = services.daemonize_services(cfg) + + # -- The handle should report the child's PID and transport details. + assert handle.pid is not None + assert handle.port == cfg.port + assert handle.transport == "tcp" + + # -- The lock file should be held by the child. + assert process.is_instance_alive(cfg.scope, base_dir=cfg.state_root) + + # -- The descriptor should have been written by the child. + desc = process.read_descriptor(cfg.scope, base_dir=cfg.state_root) + assert desc is not None + assert desc.pid == handle.pid + assert desc.mode == "daemon" + assert "hello-world" in (desc.config.services or []) + + # -- The child should still be running and respond to /status. + assert services.probe_status(base_url=f"http://127.0.0.1:{cfg.port}", timeout=5.0) + + # -- Graceful shutdown. + result = process.stop_instance(cfg.scope, base_dir=cfg.state_root, timeout=15) + assert handle.pid in result.stopped_pids + + # -- After stop, the lock should be released and the descriptor removed. + assert not process.is_instance_alive(cfg.scope, base_dir=cfg.state_root) + assert process.read_descriptor(cfg.scope, base_dir=cfg.state_root) is None + finally: + _ensure_stopped(cfg) + + +@pytest.mark.integration +def test_daemonize_services_child_exit_before_readiness(tmp_path: Path) -> None: + """When the child exits before becoming ready, ``daemonize_services`` + should raise ``ServicesStartupExitedError`` with the log path.""" + # Spawn a child that will exit immediately: give it a bogus service name + # that will fail validation in resolve_run_configuration. + bad_cfg = ServiceRunConfig( + mode=services.ServiceMode.DAEMON, + services=("nonexistent-service-xyz",), + controllers=(), + sidecars=(), + transport="tcp", + host="127.0.0.1", + port=_free_tcp_port(), + scope="integ-early-exit", + state_dir=tmp_path / "state", + runtime_dir=tmp_path / "runtime", + readiness_timeout=15.0, + readiness_poll_interval=0.2, + ) + with pytest.raises(ServicesStartupExitedError, match="exited with code"): + services.daemonize_services(bad_cfg) + + # -- The lock should not be held after the failed startup. + assert not process.is_instance_alive(bad_cfg.scope, base_dir=bad_cfg.state_root) + + +@pytest.mark.integration +def test_stale_socket_cleanup_after_process_crash(tmp_path: Path) -> None: + """If a previous daemon crashed and left a UDS socket file, a new daemon + startup should clean it up and succeed.""" + scope = "stale" + # Use a short temp directory to stay within AF_UNIX path limits (103 bytes on macOS). + short_tmp = Path(tempfile.mkdtemp(prefix="nemo-")) + runtime_dir = short_tmp / "run" + + # Create a stale UDS socket file (no process listening). + socket_dir = runtime_dir / scope + socket_dir.mkdir(parents=True, exist_ok=True) + stale_socket = socket_dir / "nemo-platform.sock" + # Bind a real UDS socket to create the file, then close immediately. + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as s: + s.bind(str(stale_socket)) + assert stale_socket.exists() + + cfg = ServiceRunConfig( + mode=services.ServiceMode.DAEMON, + services=("hello-world",), + controllers=(), + sidecars=(), + transport="uds", + host="127.0.0.1", + port=_free_tcp_port(), + scope=scope, + state_dir=short_tmp / "state", + runtime_dir=runtime_dir, + readiness_timeout=30.0, + readiness_poll_interval=0.3, + ) + try: + services.daemonize_services(cfg) + + # -- The daemon should be ready. + assert process.is_instance_alive(cfg.scope, base_dir=cfg.state_root) + + # -- The stale socket should have been replaced with the new one. + assert stale_socket.exists() + finally: + _ensure_stopped(cfg) + import shutil + + shutil.rmtree(short_tmp, ignore_errors=True) + + +@pytest.mark.integration +def test_concurrent_daemonize_rejects_duplicate_instance(tmp_path: Path) -> None: + """Starting a second daemon with the same instance scope should fail + with ``ServicesAlreadyRunningError`` while the first is running.""" + cfg = _daemon_config(tmp_path, scope="integ-dup") + try: + services.daemonize_services(cfg) + assert process.is_instance_alive(cfg.scope, base_dir=cfg.state_root) + + # -- A second daemonize with the same scope should fail. + dup_cfg = ServiceRunConfig( + mode=services.ServiceMode.DAEMON, + services=("hello-world",), + controllers=(), + sidecars=(), + transport="tcp", + host="127.0.0.1", + port=_free_tcp_port(), # Different port, same scope. + scope="integ-dup", + state_dir=tmp_path / "state", + runtime_dir=tmp_path / "runtime", + readiness_timeout=5.0, + readiness_poll_interval=0.2, + ) + with pytest.raises(ServicesAlreadyRunningError): + services.daemonize_services(dup_cfg) + + # -- Original instance should still be alive. + assert process.is_instance_alive(cfg.scope, base_dir=cfg.state_root) + finally: + _ensure_stopped(cfg) + + +@pytest.mark.integration +def test_stop_instance_escalates_sigterm_to_sigkill(tmp_path: Path) -> None: + """If the daemon child ignores SIGTERM, ``stop_instance`` should escalate + to SIGKILL after the timeout and successfully terminate the process.""" + # Instead of using daemonize_services (which starts a uvicorn server that + # handles SIGTERM), we manually simulate a daemon process that ignores SIGTERM + # using the process module primitives directly. + scope = "integ-sigkill" + state_dir = tmp_path / "state" + + # Spawn a child process that ignores SIGTERM. + child = subprocess.Popen( + [ + sys.executable, + "-c", + "import signal, time; signal.signal(signal.SIGTERM, signal.SIG_IGN); " + "open('/dev/null', 'w'); time.sleep(300)", + ], + start_new_session=True, + ) + try: + # Write a descriptor so stop_instance can find the process. + desc = process.InstanceDescriptor( + pid=child.pid, + config=PlatformAppConfig(scope=scope, host="127.0.0.1", port=0, state_root=state_dir), + transport="tcp", + mode="daemon", + create_time=psutil.Process(child.pid).create_time(), + ) + process.write_descriptor(desc, base_dir=state_dir) + + # Also create a lock file the process "holds" — but since it's a + # different process, we simulate by NOT acquiring a real flock (the + # test exercises PID-based stop, not flock-based liveness). + + # Stop with a very short timeout so it escalates quickly. + result = process.stop_instance(scope, base_dir=state_dir, timeout=1.0, force=True) + assert child.pid in result.stopped_pids + + # The child should be dead now. + child.wait(timeout=5) + assert child.returncode is not None + finally: + try: + child.kill() + child.wait(timeout=3) + except Exception: + pass + + +@pytest.mark.integration +def test_daemonize_services_cleans_up_on_child_exception(tmp_path: Path) -> None: + """When the child process crashes during init (e.g. corrupted request JSON), + the parent detects the exit, raises, and the lock is not left held.""" + scope = "integ-crash" + state_dir = tmp_path / "state" + instance_dir = process.instance_dir(scope, base_dir=state_dir) + + # Write a corrupted request file that will make _service_child crash + # during JSON deserialization. + fd, tmp_req = tempfile.mkstemp(dir=str(instance_dir), suffix=".json") + os.write(fd, b"NOT VALID JSON {{{") + os.close(fd) + + log_path = process.log_path_for(scope, base_dir=state_dir) + log_file = open(log_path, "a") # noqa: SIM115 + child_module = "nemo_platform.local._service_child" + proc = subprocess.Popen( + [sys.executable, "-m", child_module, tmp_req], + stdout=log_file, + stderr=subprocess.STDOUT, + start_new_session=True, + stdin=subprocess.DEVNULL, + close_fds=True, + ) + log_file.close() + + # Wait for the child to exit (it should crash quickly on bad JSON). + proc.wait(timeout=10) + assert proc.returncode != 0 + + # The lock should not be held — the child never acquired it. + assert not process.is_instance_alive(scope, base_dir=state_dir) + + # The request file should have been cleaned up by _service_child. + assert not Path(tmp_req).exists() + + +# --------------------------------------------------------------------------- +# Priority 2: Process Lifecycle & Cleanup +# --------------------------------------------------------------------------- + + +@pytest.mark.integration +def test_stop_instance_sweeps_orphaned_children(tmp_path: Path) -> None: + """When a daemon parent is stopped, any grandchild processes that survive + should be swept by ``_sweep_orphans``.""" + scope = "integ-orphans" + state_dir = tmp_path / "state" + + # Spawn a parent that spawns a long-lived grandchild, then sleeps. + parent = subprocess.Popen( + [ + sys.executable, + "-c", + "import subprocess, sys, time; " + "gc = subprocess.Popen([sys.executable, '-c', 'import time; time.sleep(300)']); " + "time.sleep(300)", + ], + start_new_session=True, + ) + try: + # Give the parent time to spawn the grandchild. + time.sleep(0.5) + grandchildren = psutil.Process(parent.pid).children(recursive=True) + assert len(grandchildren) >= 1, "grandchild was not spawned" + + desc = process.InstanceDescriptor( + pid=parent.pid, + config=PlatformAppConfig(scope=scope, host="127.0.0.1", port=0, state_root=state_dir), + transport="tcp", + mode="daemon", + create_time=psutil.Process(parent.pid).create_time(), + ) + process.write_descriptor(desc, base_dir=state_dir) + + result = process.stop_instance(scope, base_dir=state_dir, timeout=10, force=True) + assert parent.pid in result.stopped_pids + assert len(result.swept_children) >= 1 + + # Both parent and grandchild should be dead. + parent.wait(timeout=5) + for gc in grandchildren: + gc.wait(timeout=5) + finally: + try: + parent.kill() + parent.wait(timeout=3) + except Exception: + pass + for gc in grandchildren: + try: + gc.kill() + gc.wait(timeout=3) + except Exception: + pass + + +@pytest.mark.integration +def test_stop_instance_foreground_mode_requires_force(tmp_path: Path) -> None: + """Stopping a foreground-mode instance without ``force=True`` should raise + ``ForegroundInstanceError``. With ``force=True`` it should proceed.""" + scope = "integ-foreground" + state_dir = tmp_path / "state" + + child = subprocess.Popen( + [sys.executable, "-c", "import time; time.sleep(300)"], + start_new_session=True, + ) + try: + desc = process.InstanceDescriptor( + pid=child.pid, + config=PlatformAppConfig(scope=scope, host="127.0.0.1", port=0, state_root=state_dir), + transport="tcp", + mode="foreground", + create_time=psutil.Process(child.pid).create_time(), + ) + process.write_descriptor(desc, base_dir=state_dir) + + # Without force, should raise. + with pytest.raises(ForegroundInstanceError): + process.stop_instance(scope, base_dir=state_dir, timeout=5) + + # Process should still be alive after the rejected stop. + assert child.poll() is None + + # With force, should succeed. + result = process.stop_instance(scope, base_dir=state_dir, timeout=5, force=True) + assert child.pid in result.stopped_pids + child.wait(timeout=5) + finally: + try: + child.kill() + child.wait(timeout=3) + except Exception: + pass + + +@pytest.mark.integration +def test_is_instance_alive_with_stale_lock(tmp_path: Path) -> None: + """If the lock file exists but no process holds the flock, + ``is_instance_alive`` should return False.""" + scope = "integ-stale-lock" + state_dir = tmp_path / "state" + + # Create the lock file without holding a flock on it. + inst_dir = process.instance_dir(scope, base_dir=state_dir) + lock_path = inst_dir / process.LOCK_FILENAME + lock_path.touch() + + assert not process.is_instance_alive(scope, base_dir=state_dir) + + +@pytest.mark.integration +def test_is_instance_alive_with_held_lock(tmp_path: Path) -> None: + """If a process holds the flock, ``is_instance_alive`` should return True.""" + scope = "integ-held-lock" + state_dir = tmp_path / "state" + + fd = process.acquire_lock(scope, base_dir=state_dir) + try: + assert process.is_instance_alive(scope, base_dir=state_dir) + finally: + os.close(fd) + + # After releasing the fd (which releases the flock), should be false. + assert not process.is_instance_alive(scope, base_dir=state_dir) + + +@pytest.mark.integration +def test_validate_pid_detects_recycled_process(tmp_path: Path) -> None: + """After a process dies, ``validate_pid`` should return False if the PID is + reused by a different process (detected via create_time mismatch).""" + # Spawn and immediately kill a short-lived process to get a PID + create_time. + child = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(300)"]) + pid = child.pid + create_time = psutil.Process(pid).create_time() + + # The PID is alive and create_time matches. + assert process.validate_pid(pid, create_time) + + # Kill it. + child.kill() + child.wait(timeout=5) + + # Now validate_pid should return False — the process is dead. + assert not process.validate_pid(pid, create_time) + + # Even with a wildly wrong create_time, should be False for a dead PID. + assert not process.validate_pid(pid, 0.0) + + +@pytest.mark.integration +def test_rotate_log_preserves_existing_content(tmp_path: Path) -> None: + """``rotate_log`` should rename the existing log and return the path for + the new (empty) log. The old content must be preserved.""" + scope = "integ-rotate" + state_dir = tmp_path / "state" + + # Write initial log content. + log_path = process.log_path_for(scope, base_dir=state_dir) + log_path.parent.mkdir(parents=True, exist_ok=True) + log_path.write_text("original log content\n") + + # Rotate. + new_log = process.rotate_log(scope, base_dir=state_dir) + assert new_log == log_path + assert not log_path.exists() # Original was renamed. + + # Find the rotated file. + rotated_files = [f for f in log_path.parent.iterdir() if f.name.startswith("services.log.")] + assert len(rotated_files) == 1 + assert rotated_files[0].read_text() == "original log content\n" + + # Write new content, rotate again. + log_path.write_text("second run\n") + process.rotate_log(scope, base_dir=state_dir) + + rotated_files = sorted(f for f in log_path.parent.iterdir() if f.name.startswith("services.log.")) + assert len(rotated_files) == 2 + contents = {f.read_text() for f in rotated_files} + assert "original log content\n" in contents + assert "second run\n" in contents diff --git a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/local/test_health_child.py b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/local/test_health_child.py new file mode 100644 index 0000000000..683e7b3720 --- /dev/null +++ b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/local/test_health_child.py @@ -0,0 +1,340 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Integration tests for health/readiness probing, lifespan, and child process module. + +Covers Priorities 5 (lifespan), 6 (health), and 7 (child process) from the +integration test plan. +""" + +from __future__ import annotations + +import json +import os +import socket +import subprocess +import sys +import tempfile +import threading +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest +from nemo_platform.local import process +from nemo_platform.local.services import ServiceRunConfig +from nemo_platform.local.transport import probe_status, wait_for_status + + +def _free_tcp_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +# --------------------------------------------------------------------------- +# Priority 5: Multi-Service Startup & Lifespan +# --------------------------------------------------------------------------- + + +@pytest.mark.integration +def test_create_app_starts_and_joins_controller_threads() -> None: + """A controller registered via ``create_app`` should have its thread + started during lifespan and stopped on exit.""" + started = threading.Event() + stopped = threading.Event() + + def controller_run(stop_signal: threading.Event) -> None: + started.set() + stop_signal.wait(timeout=5.0) + stopped.set() + + with ( + patch("nmp.platform_runner.server.get_platform_config") as mock_pc, + patch("nmp.platform_runner.server.get_auth_config") as mock_ac, + patch("nmp.common.auth.middleware.get_auth_config") as mock_ac2, + ): + mock_pc.return_value = MagicMock(seed_on_startup=False, redirect_root_to_studio=False) + mock_ac.return_value = MagicMock(enabled=False, policy_decision_point_provider="embedded") + mock_ac2.return_value = MagicMock(enabled=False, policy_decision_point_provider="embedded") + + from nmp.platform_runner.server import create_app + + app = create_app(services=[], controller_run_funcs={"test-ctrl": controller_run}) + + from fastapi.testclient import TestClient + + with TestClient(app): + assert started.wait(timeout=2.0), "controller thread did not start" + + assert stopped.wait(timeout=2.0), "controller thread did not stop after lifespan exit" + + +@pytest.mark.integration +def test_create_app_controller_thread_join_timeout() -> None: + """A controller that ignores the stop signal should not hang shutdown — + ``thread.join(timeout=5)`` should return even if the controller is still running.""" + started = threading.Event() + + def stubborn_controller(stop_signal: threading.Event) -> None: + started.set() + # Ignore stop_signal — simulate a controller that hangs. + import time + + time.sleep(300) + + with ( + patch("nmp.platform_runner.server.get_platform_config") as mock_pc, + patch("nmp.platform_runner.server.get_auth_config") as mock_ac, + patch("nmp.common.auth.middleware.get_auth_config") as mock_ac2, + ): + mock_pc.return_value = MagicMock(seed_on_startup=False, redirect_root_to_studio=False) + mock_ac.return_value = MagicMock(enabled=False, policy_decision_point_provider="embedded") + mock_ac2.return_value = MagicMock(enabled=False, policy_decision_point_provider="embedded") + + from nmp.platform_runner.server import create_app + + app = create_app(services=[], controller_run_funcs={"stubborn": stubborn_controller}) + + from fastapi.testclient import TestClient + + # The TestClient __exit__ triggers lifespan exit, which calls thread.join(timeout=5). + # This should NOT hang forever — the 5s timeout should let shutdown proceed. + with TestClient(app): + assert started.wait(timeout=2.0), "controller thread did not start" + + # If we got here, shutdown didn't hang. The stubborn thread is still running + # but as a daemon thread it will be cleaned up when the test process exits. + + +@pytest.mark.integration +def test_lifespan_cleanup_runs_on_app_shutdown() -> None: + """``close_shared_http_clients`` should be called during lifespan teardown.""" + cleanup_called = threading.Event() + + with ( + patch("nmp.platform_runner.server.get_platform_config") as mock_pc, + patch("nmp.platform_runner.server.get_auth_config") as mock_ac, + patch("nmp.common.auth.middleware.get_auth_config") as mock_ac2, + patch("nmp.platform_runner.server.close_shared_http_clients") as mock_close, + ): + mock_pc.return_value = MagicMock(seed_on_startup=False, redirect_root_to_studio=False) + mock_ac.return_value = MagicMock(enabled=False, policy_decision_point_provider="embedded") + mock_ac2.return_value = MagicMock(enabled=False, policy_decision_point_provider="embedded") + + async def fake_close(): + cleanup_called.set() + + mock_close.side_effect = fake_close + + from nmp.platform_runner.server import create_app + + app = create_app(services=[]) + + from fastapi.testclient import TestClient + + with TestClient(app): + pass + + assert cleanup_called.is_set(), "close_shared_http_clients was not called during shutdown" + + +# --------------------------------------------------------------------------- +# Priority 6: Health & Readiness +# --------------------------------------------------------------------------- + + +@pytest.mark.integration +def test_wait_for_status_retries_on_transient_errors(tmp_path: Path) -> None: + """``wait_for_status`` should retry on connection refused and eventually + return True once the server starts responding.""" + from nemo_platform.local import services + + cfg = ServiceRunConfig( + mode=services.ServiceMode.DAEMON, + services=("hello-world",), + controllers=(), + sidecars=(), + transport="tcp", + host="127.0.0.1", + port=_free_tcp_port(), + scope="integ-wait-retry", + state_dir=tmp_path / "state", + runtime_dir=tmp_path / "runtime", + readiness_timeout=30.0, + readiness_poll_interval=0.3, + ) + + # Start the daemon — wait_for_status should retry until it's ready. + services.daemonize_services(cfg) + try: + # The daemon is already ready (daemonize_services waits for readiness). + # Verify wait_for_status succeeds with a fresh probe. + assert wait_for_status( + base_url=f"http://127.0.0.1:{cfg.port}", + timeout=5.0, + poll_interval=0.2, + ) + finally: + process.stop_instance(cfg.scope, base_dir=cfg.state_root, timeout=10, force=True) + + +@pytest.mark.integration +def test_wait_for_status_times_out_on_no_server() -> None: + """``wait_for_status`` should return False when no server is listening.""" + port = _free_tcp_port() + result = wait_for_status( + base_url=f"http://127.0.0.1:{port}", + timeout=1.0, + poll_interval=0.2, + ) + assert result is False + + +@pytest.mark.integration +def test_probe_status_with_missing_uds_socket() -> None: + """Probing a non-existent UDS socket should return False.""" + result = probe_status( + base_url="http+unix:///nonexistent/path/nemo.sock", + socket_path=Path("/nonexistent/path/nemo.sock"), + timeout=1.0, + ) + assert result is False + + +@pytest.mark.integration +def test_probe_status_against_real_daemon(tmp_path: Path) -> None: + """``probe_status`` should return True against a running daemon.""" + from nemo_platform.local import services + + cfg = ServiceRunConfig( + mode=services.ServiceMode.DAEMON, + services=("hello-world",), + controllers=(), + sidecars=(), + transport="tcp", + host="127.0.0.1", + port=_free_tcp_port(), + scope="integ-probe-real", + state_dir=tmp_path / "state", + runtime_dir=tmp_path / "runtime", + readiness_timeout=30.0, + readiness_poll_interval=0.3, + ) + services.daemonize_services(cfg) + try: + assert probe_status(base_url=f"http://127.0.0.1:{cfg.port}", timeout=5.0) + finally: + process.stop_instance(cfg.scope, base_dir=cfg.state_root, timeout=10, force=True) + + +# --------------------------------------------------------------------------- +# Priority 7: Child Process Module +# --------------------------------------------------------------------------- + + +@pytest.mark.integration +def test_service_child_loads_config_and_starts(tmp_path: Path) -> None: + """Write valid JSON config, run ``_service_child`` in a subprocess, + verify it starts and accepts HTTP connections.""" + port = _free_tcp_port() + state_dir = tmp_path / "state" + runtime_dir = tmp_path / "runtime" + scope = "integ-child-real" + + payload = ServiceRunConfig( + mode="daemon", + services=("hello-world",), + controllers=(), + sidecars=(), + transport="tcp", + host="127.0.0.1", + port=port, + scope=scope, + state_dir=str(state_dir), + runtime_dir=str(runtime_dir), + ).to_child_payload() + + # Write the request file the way daemonize_services does. + instance_dir = process.instance_dir(scope, base_dir=state_dir) + fd, req_path = tempfile.mkstemp(dir=str(instance_dir), suffix=".json") + os.write(fd, (json.dumps(payload) + "\n").encode()) + os.close(fd) + + log_path = process.log_path_for(scope, base_dir=state_dir) + log_file = open(log_path, "a") # noqa: SIM115 + subprocess.Popen( + [sys.executable, "-m", "nemo_platform.local._service_child", req_path], + stdout=log_file, + stderr=subprocess.STDOUT, + start_new_session=True, + stdin=subprocess.DEVNULL, + close_fds=True, + ) + log_file.close() + + try: + # Wait for the child to become ready. + assert wait_for_status( + base_url=f"http://127.0.0.1:{port}", + timeout=30.0, + poll_interval=0.3, + ), "child process did not become ready" + + # The request file should have been cleaned up. + assert not Path(req_path).exists() + + # The child should have acquired the lock and written a descriptor. + assert process.is_instance_alive(scope, base_dir=state_dir) + desc = process.read_descriptor(scope, base_dir=state_dir) + assert desc is not None + assert desc.mode == "daemon" + finally: + process.stop_instance(scope, base_dir=state_dir, timeout=10, force=True) + + +@pytest.mark.integration +def test_service_child_corrupted_payload(tmp_path: Path) -> None: + """Bad JSON in the request file should cause the child to exit non-zero.""" + scope = "integ-child-bad" + state_dir = tmp_path / "state" + instance_dir = process.instance_dir(scope, base_dir=state_dir) + + fd, req_path = tempfile.mkstemp(dir=str(instance_dir), suffix=".json") + os.write(fd, b"<<>>") + os.close(fd) + + proc = subprocess.Popen( + [sys.executable, "-m", "nemo_platform.local._service_child", req_path], + stdout=subprocess.DEVNULL, + stderr=subprocess.STDOUT, + start_new_session=True, + stdin=subprocess.DEVNULL, + ) + proc.wait(timeout=15) + assert proc.returncode != 0 + + +@pytest.mark.integration +def test_service_child_cleans_up_request_file(tmp_path: Path) -> None: + """The request file should be unlinked even when the child crashes.""" + scope = "integ-child-cleanup" + state_dir = tmp_path / "state" + instance_dir = process.instance_dir(scope, base_dir=state_dir) + + fd, req_path = tempfile.mkstemp(dir=str(instance_dir), suffix=".json") + os.write(fd, b"<<>>") + os.close(fd) + + assert Path(req_path).exists() + + proc = subprocess.Popen( + [sys.executable, "-m", "nemo_platform.local._service_child", req_path], + stdout=subprocess.DEVNULL, + stderr=subprocess.STDOUT, + start_new_session=True, + stdin=subprocess.DEVNULL, + ) + proc.wait(timeout=15) + + # The request file should have been cleaned up regardless of the error. + assert not Path(req_path).exists() diff --git a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/local/test_port_socket.py b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/local/test_port_socket.py new file mode 100644 index 0000000000..243b65b24d --- /dev/null +++ b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/local/test_port_socket.py @@ -0,0 +1,165 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Integration tests for TCP/UDS port and socket management. + +These tests exercise real port binding, socket creation, and conflict +detection using actual OS resources. +""" + +from __future__ import annotations + +import socket +import tempfile +from pathlib import Path + +import pytest +from nemo_platform.local import process, services +from nemo_platform.local.services import ( + ServiceRunConfig, + ServicesPortInUseError, +) +from nmp.platform_runner.config import PlatformAppConfig + + +def _free_tcp_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +@pytest.mark.integration +def test_tcp_port_conflict_with_foreign_process(tmp_path: Path) -> None: + """When a foreign (non-NeMo) process holds a port, ``_check_tcp_available`` + should raise ``ServicesPortInUseError`` with a helpful suggestion.""" + # Bind a TCP port and hold it open. + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as blocker: + blocker.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + blocker.bind(("127.0.0.1", 0)) + blocker.listen(1) + port = blocker.getsockname()[1] + + cfg = ServiceRunConfig( + mode=services.ServiceMode.DAEMON, + services=("hello-world",), + transport="tcp", + host="127.0.0.1", + port=port, + scope="integ-port-foreign", + state_dir=tmp_path / "state", + runtime_dir=tmp_path / "runtime", + ) + + with pytest.raises(ServicesPortInUseError, match="already in use by another process"): + services._check_tcp_available(cfg) + + +@pytest.mark.integration +def test_tcp_port_conflict_with_nemo_instance(tmp_path: Path) -> None: + """When a NeMo instance holds a port, the error should distinguish it + from a foreign process.""" + scope = "integ-port-nemo" + state_dir = tmp_path / "state" + + # Bind a port and also create a descriptor matching the scope/host/port. + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as blocker: + blocker.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + blocker.bind(("127.0.0.1", 0)) + blocker.listen(1) + port = blocker.getsockname()[1] + + # Create a live lock and descriptor so it looks like a NeMo instance. + lock_fd = process.acquire_lock(scope, base_dir=state_dir) + try: + desc = process.InstanceDescriptor( + pid=1, # Dummy PID — the flock is what matters. + config=PlatformAppConfig(scope=scope, host="127.0.0.1", port=port), + transport="tcp", + mode="daemon", + create_time=0.0, + ) + process.write_descriptor(desc, base_dir=state_dir) + + conflict = process.check_port_available_for_start("127.0.0.1", port, scope, base_dir=state_dir) + assert conflict is not None + assert conflict.kind == "nemo_instance" + assert conflict.port == port + + lines = process.format_port_conflict(conflict) + assert any("NeMo Platform" in line for line in lines) + finally: + import os + + os.close(lock_fd) + + +@pytest.mark.integration +def test_tcp_port_available_when_free(tmp_path: Path) -> None: + """When a port is free, ``check_port_available_for_start`` returns None.""" + port = _free_tcp_port() + conflict = process.check_port_available_for_start("127.0.0.1", port, "integ-free", base_dir=tmp_path / "state") + assert conflict is None + + +@pytest.mark.integration +def test_uds_socket_path_max_validation() -> None: + """A socket path exceeding AF_UNIX_PATH_MAX should raise ValueError.""" + # Build a path that is exactly one byte over the limit. + max_bytes = services._AF_UNIX_PATH_MAX_BYTES + # Create a path that exceeds the limit. + long_path = "/" + "x" * max_bytes # len("/") + max_bytes > max_bytes + assert len(long_path.encode()) > max_bytes + + cfg = ServiceRunConfig( + mode=services.ServiceMode.DAEMON, + services=("hello-world",), + transport="uds", + socket_path=long_path, + scope="integ-long-sock", + state_dir="/tmp/state", + runtime_dir="/tmp/run", + ) + # _validated_socket_path calls _validate_socket_path_length internally. + with pytest.raises(ValueError, match="too long for AF_UNIX"): + services._validated_socket_path(cfg) + + +@pytest.mark.integration +def test_prepare_socket_removes_stale_socket(tmp_path: Path) -> None: + """``_prepare_socket`` should remove a stale (unreachable) socket file + and allow a new daemon to bind.""" + scope = "stale2" + short_tmp = Path(tempfile.mkdtemp(prefix="nemo-")) + runtime_dir = short_tmp / "run" + + # Create a stale socket file. + socket_dir = runtime_dir / scope + socket_dir.mkdir(parents=True, exist_ok=True) + stale_socket = socket_dir / "nemo-platform.sock" + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as s: + s.bind(str(stale_socket)) + assert stale_socket.exists() + + cfg = ServiceRunConfig( + mode=services.ServiceMode.DAEMON, + services=("hello-world",), + transport="uds", + scope=scope, + state_dir=short_tmp / "state", + runtime_dir=runtime_dir, + ) + + # _prepare_socket should probe, find it stale, remove it, and return the path. + result = services._prepare_socket(cfg) + assert result is not None + # The stale socket should have been removed (the new server hasn't bound yet). + assert not stale_socket.exists() + + import shutil + + shutil.rmtree(short_tmp, ignore_errors=True) diff --git a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/local/test_services.py b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/local/test_services.py new file mode 100644 index 0000000000..e2b2bca695 --- /dev/null +++ b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/local/test_services.py @@ -0,0 +1,1039 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +import os +import socket +import subprocess +import sys +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from nemo_platform.local import _service_child, services +from nemo_platform.local.process import ( + DESCRIPTOR_FILENAME, + InstanceDescriptor, +) +from nemo_platform.local.services import ServiceRunConfig +from nemo_platform.local.transport import UDS_BASE_URL +from nmp.platform_runner.config import ( + PlatformAppConfig, + default_runtime_root, + default_state_root, + validate_scope, +) + + +def _allow_tmp_path_socket_paths(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(services, "_AF_UNIX_PATH_MAX_BYTES", 4096) + + +def _free_tcp_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + return sock.getsockname()[1] + + +def _embedded_handle() -> services.EmbeddedServiceHandle: + return services.EmbeddedServiceHandle(app=object(), runtime=object()) + + +def test_service_run_config_normalizes_lists_to_tuples() -> None: + cfg = ServiceRunConfig(services=["entities", "models"], controllers=["jobs"]) + + assert cfg.services == ("entities", "models") + assert cfg.controllers == ("jobs",) + + +def test_service_run_config_converts_to_platform_app_config(tmp_path: Path) -> None: + cfg = ServiceRunConfig( + services=["entities", "models"], + controllers=[], + sidecars=["adapters"], + config_path=tmp_path / "local.yaml", + socket_path=tmp_path / "nemo.sock", + mode="embedded", + ) + + app_config = cfg.to_platform_app_config() + + assert app_config.services == ("entities", "models") + assert app_config.controllers == () + assert app_config.sidecars == ("adapters",) + assert app_config.config_path == str(tmp_path / "local.yaml") + assert app_config.socket_path == str(tmp_path / "nemo.sock") + assert app_config.runtime_root is None + assert app_config.runtime_dir() == tmp_path + assert app_config.host == "127.0.0.1" + assert app_config.port == 8080 + + +def test_instance_descriptor_converts_from_service_run_config( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + cfg = ServiceRunConfig( + services=["entities", "models"], + controllers=[], + sidecars=["adapters"], + config_path=tmp_path / "local.yaml", + socket_path=tmp_path / "nemo.sock", + state_dir=tmp_path / "state", + runtime_dir=tmp_path / "run", + ) + monkeypatch.setattr(services.process, "get_create_time", lambda _pid: 123.0) + app_config = cfg.to_platform_app_config() + app_config.log_path = str(tmp_path / "nemo.log") + + desc = InstanceDescriptor.from_config( + app_config, + pid=4242, + mode="daemon", + transport=cfg.transport, + ) + + assert desc.pid == 4242 + assert desc.config.scope == "default" + assert desc.config.host == "127.0.0.1" + assert desc.config.port == 8080 + assert desc.transport == "uds" + assert desc.config.socket_path == str(tmp_path / "nemo.sock") + assert desc.config.state_root == str(tmp_path / "state") + assert desc.config.runtime_root == str(tmp_path / "run") + assert desc.config.state_dir() == tmp_path / "state" / "instances" / "default" + assert desc.config.runtime_dir() == tmp_path / "run" / "default" + assert desc.mode == "daemon" + assert desc.create_time == 123.0 + assert desc.config.services == ("entities", "models") + assert desc.config.controllers == () + assert desc.config.sidecars == ("adapters",) + assert desc.config.config_path == str(tmp_path / "local.yaml") + assert desc.config.log_path == str(tmp_path / "nemo.log") + assert desc.config.log_file_path() == tmp_path / "nemo.log" + payload = desc.model_dump() + assert "services" not in payload + assert "host" not in payload + assert "state_dir" not in payload + assert "runtime_dir" not in payload + assert "log_path" not in payload + assert payload["config"]["services"] == ("entities", "models") + assert payload["config"]["socket_path"] == str(tmp_path / "nemo.sock") + assert payload["config"]["state_root"] == str(tmp_path / "state") + assert payload["config"]["runtime_root"] == str(tmp_path / "run") + assert payload["config"]["log_path"] == str(tmp_path / "nemo.log") + + +def test_service_mode_enum_values() -> None: + assert services.ServiceMode.EMBEDDED.value == "embedded" + assert services.ServiceMode.DAEMON.value == "daemon" + + +def test_service_run_config_defaults_to_daemon_mode() -> None: + cfg = ServiceRunConfig() + + assert cfg.mode is services.ServiceMode.DAEMON + + +def test_service_run_config_accepts_mode_strings() -> None: + cfg = ServiceRunConfig(mode="embedded") + + assert cfg.mode is services.ServiceMode.EMBEDDED + + +def test_service_run_config_rejects_unknown_mode() -> None: + with pytest.raises(ValueError, match="mode must be 'embedded' or 'daemon'"): + ServiceRunConfig(mode="foreground") + + +def test_embedded_and_daemon_handles_implement_local_service_handle(tmp_path: Path) -> None: + embedded = services.EmbeddedServiceHandle(app=object(), runtime=object()) + daemon = services.DaemonServiceHandle( + scope="dev", + transport="tcp", + socket_path=None, + gateway_base_url=None, + host="127.0.0.1", + port=8080, + pid=123, + mode="daemon", + log_path=None, + state_dir=tmp_path / "state" / "instances" / "dev", + runtime_dir=None, + ) + + assert isinstance(embedded, services.LocalServiceHandle) + assert isinstance(daemon, services.LocalServiceHandle) + + +def test_start_services_result_is_shared_result_type() -> None: + result = services.StartServicesResult( + requested=["jobs"], + started=["auth", "jobs"], + already_active=[], + active=["secrets", "auth", "jobs"], + ) + + assert result.requested == ["jobs"] + assert result.started == ["auth", "jobs"] + assert result.active == ["secrets", "auth", "jobs"] + + +def test_service_run_config_rejects_services_with_service_group() -> None: + with pytest.raises(ValueError, match="services cannot be combined with service_group"): + ServiceRunConfig(services=("entities",), service_group="all") + + +def test_service_run_config_defaults_to_named_uds_instance() -> None: + cfg = ServiceRunConfig() + + assert cfg.transport == "uds" + assert cfg.http_gateway == "disabled" + assert cfg.scope == "default" + assert cfg.socket_path is None + + +@pytest.mark.parametrize("instance", ["has space", "../bad"]) +def test_service_run_config_rejects_invalid_scope_names(instance: str) -> None: + with pytest.raises(ValueError, match="scope"): + ServiceRunConfig(scope=instance) + + +def test_service_run_config_rejects_gateway_for_tcp_transport() -> None: + with pytest.raises(ValueError, match="gateway.*UDS"): + ServiceRunConfig(transport="tcp", http_gateway="enabled") + + +@pytest.mark.parametrize( + ("field", "value", "message"), + [ + ("readiness_timeout", 0.0, "readiness_timeout"), + ("readiness_timeout", -1.0, "readiness_timeout"), + ("readiness_poll_interval", 0.0, "readiness_poll_interval"), + ("readiness_poll_interval", -1.0, "readiness_poll_interval"), + ], +) +def test_service_run_config_rejects_non_positive_readiness_values(field: str, value: float, message: str) -> None: + with pytest.raises(ValueError, match=message): + if field == "readiness_timeout": + ServiceRunConfig(readiness_timeout=value) + else: + ServiceRunConfig(readiness_poll_interval=value) + + +def test_process_paths_follow_existing_nmp_state_convention(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + monkeypatch.setenv("XDG_STATE_HOME", str(tmp_path / "state")) + + assert default_state_root() == tmp_path / "state" / "nmp" + assert default_runtime_root() == tmp_path / "state" / "nmp" / "run" + assert ( + PlatformAppConfig(scope="dev").socket_file_path() + == tmp_path / "state" / "nmp" / "run" / "dev" / "nemo-platform.sock" + ) + assert validate_scope("dev_1-2") == "dev_1-2" + + +def test_resolved_socket_path_rejects_relative_explicit_path() -> None: + cfg = ServiceRunConfig(socket_path="relative.sock") + + with pytest.raises(ValueError, match="UDS socket path must be absolute"): + _ = cfg.resolved_socket_path + + +def test_resolved_socket_path_rejects_relative_runtime_dir() -> None: + cfg = ServiceRunConfig(runtime_dir="relative-run") + + with pytest.raises(ValueError, match="runtime root must be absolute"): + _ = cfg.resolved_socket_path + + +def test_resolved_socket_path_rejects_relative_socket_path_with_tcp_client() -> None: + cfg = ServiceRunConfig(transport="tcp", socket_path="relative.sock") + + with pytest.raises(ValueError, match="UDS socket path must be absolute"): + _ = cfg.resolved_socket_path + + +def test_tcp_client_can_still_configure_uds_listener(tmp_path: Path) -> None: + cfg = ServiceRunConfig(transport="tcp", socket_path=tmp_path / "nemo.sock") + + app_config = cfg.to_platform_app_config() + + assert app_config.socket_path == str(tmp_path / "nemo.sock") + assert app_config.runtime_dir() == tmp_path + + +def test_instance_descriptor_rejects_uds_client_without_socket_path() -> None: + with pytest.raises(ValueError, match="UDS client transport requires config.socket_path"): + InstanceDescriptor(pid=1, config=PlatformAppConfig(scope="dev"), transport="uds") + + +def test_prepare_socket_rejects_long_generated_path_before_creating_runtime_dir( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + monkeypatch.setattr(services, "_AF_UNIX_PATH_MAX_BYTES", 1, raising=False) + runtime_root = tmp_path / "runtime" + cfg = ServiceRunConfig(scope="dev", state_dir=tmp_path / "state", runtime_dir=runtime_root) + + with pytest.raises(ValueError, match="UDS socket path is too long.*AF_UNIX"): + services._prepare_socket(cfg) + + assert not runtime_root.exists() + + +def test_validate_socket_path_length_reserves_trailing_nul(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(services, "_AF_UNIX_PATH_MAX_BYTES", 3, raising=False) + + services._validate_socket_path_length(Path("abc")) + with pytest.raises(ValueError, match=r"4 bytes; maximum is 3 bytes"): + services._validate_socket_path_length(Path("abcd")) + + +def test_prepare_socket_rejects_long_explicit_path_before_filesystem_or_probe( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + monkeypatch.setattr(services, "_AF_UNIX_PATH_MAX_BYTES", 1, raising=False) + socket_parent = tmp_path / "explicit" + cfg = ServiceRunConfig(socket_path=socket_parent / "nemo-platform.sock") + + with patch("nemo_platform.local.services.probe_status") as probe_status: + with pytest.raises(ValueError, match="UDS socket path is too long.*AF_UNIX"): + services._prepare_socket(cfg) + + probe_status.assert_not_called() + assert not socket_parent.exists() + + +def test_run_services_prepares_socket_after_acquiring_lock(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + _allow_tmp_path_socket_paths(monkeypatch) + cfg = ServiceRunConfig( + scope="dev", port=_free_tcp_port(), state_dir=tmp_path / "state", runtime_dir=tmp_path / "run" + ) + events: list[str] = [] + real_acquire_lock = services.process.acquire_lock + + def acquire_lock(scope: str, *, base_dir: Path | None = None) -> int: + events.append("lock") + return real_acquire_lock(scope, base_dir=base_dir) + + def prepare_socket(config: ServiceRunConfig) -> Path | None: + events.append("prepare") + lock_path = ( + services.process.instance_dir(config.scope, base_dir=config.state_root) / services.process.LOCK_FILENAME + ) + assert lock_path.exists() + return config.resolved_socket_path + + with ( + patch("nemo_platform.local.services.require_services_extra"), + patch("nemo_platform.local.services.process.acquire_lock", side_effect=acquire_lock), + patch("nemo_platform.local.services._prepare_socket", side_effect=prepare_socket), + patch("nemo_platform.local.services.start_embedded_services", return_value=_embedded_handle()), + patch("nemo_platform.local.services.serve_embedded_app"), + ): + services.run_services(cfg) + + assert events == ["lock", "prepare"] + + +def test_run_services_cleans_lock_when_socket_prepare_fails(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + _allow_tmp_path_socket_paths(monkeypatch) + cfg = ServiceRunConfig( + scope="dev", port=_free_tcp_port(), state_dir=tmp_path / "state", runtime_dir=tmp_path / "run" + ) + + with ( + patch("nemo_platform.local.services.require_services_extra"), + patch( + "nemo_platform.local.services._prepare_socket", + side_effect=services.ServicesSocketStaleError("boom"), + ), + ): + with pytest.raises(services.ServicesSocketStaleError, match="boom"): + services.run_services(cfg) + + assert not services.process.is_instance_alive(cfg.scope, base_dir=cfg.state_root) + + +def test_run_services_restores_env_and_closes_lock_when_descriptor_cleanup_fails( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + _allow_tmp_path_socket_paths(monkeypatch) + cfg = ServiceRunConfig( + scope="dev", + port=_free_tcp_port(), + state_dir=tmp_path / "state", + runtime_dir=tmp_path / "run", + data_dir=tmp_path / "data", + ) + monkeypatch.delenv("NMP_DATA_DIR", raising=False) + real_acquire_lock = services.process.acquire_lock + locked_fd: int | None = None + + def acquire_lock(scope: str, *, base_dir: Path | None = None) -> int: + nonlocal locked_fd + locked_fd = real_acquire_lock(scope, base_dir=base_dir) + return locked_fd + + real_close = os.close + closed_fds: list[int] = [] + + def close(fd: int) -> None: + closed_fds.append(fd) + real_close(fd) + + monkeypatch.setattr(os, "close", close) + with ( + patch("nemo_platform.local.services.require_services_extra"), + patch("nemo_platform.local.services.process.acquire_lock", side_effect=acquire_lock), + patch( + "nemo_platform.local.services.process.remove_descriptor", + side_effect=RuntimeError("descriptor cleanup failed"), + ), + patch("nemo_platform.local.services.start_embedded_services", return_value=_embedded_handle()), + patch("nemo_platform.local.services.serve_embedded_app"), + ): + with pytest.raises(RuntimeError, match="descriptor cleanup failed"): + services.run_services(cfg) + + assert "NMP_DATA_DIR" not in os.environ + assert locked_fd is not None + assert locked_fd in closed_fds + assert closed_fds[-1] == locked_fd + + +def test_run_services_foreground_serves_embedded_app(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + _allow_tmp_path_socket_paths(monkeypatch) + cfg = ServiceRunConfig( + mode=services.ServiceMode.EMBEDDED, + scope="dev", + port=_free_tcp_port(), + state_dir=tmp_path / "state", + runtime_dir=tmp_path / "run", + ) + app = object() + handle = services.EmbeddedServiceHandle(app=app, runtime=object()) + + with ( + patch("nemo_platform.local.services.require_services_extra"), + patch("nemo_platform.local.services.start_embedded_services", return_value=handle) as start_embedded, + patch("nemo_platform.local.services.serve_embedded_app") as serve_embedded, + ): + services.run_services(cfg) + + start_embedded.assert_called_once_with(cfg, env=None) + serve_embedded.assert_called_once() + assert serve_embedded.call_args.args[0] is app + + +def test_daemon_service_handle_uds_client_uses_socket_transport(tmp_path: Path) -> None: + socket_path = tmp_path / "nemo-platform.sock" + handle = services.DaemonServiceHandle( + scope="dev", + transport="uds", + socket_path=socket_path, + gateway_base_url="http://127.0.0.1:9999", + host="127.0.0.1", + port=8080, + pid=123, + mode="daemon", + log_path=None, + state_dir=tmp_path / "state" / "instances" / "dev", + runtime_dir=tmp_path, + ) + + client = handle.client() + try: + assert str(client.base_url).rstrip("/") == UDS_BASE_URL + assert handle.gateway_base_url == "http://127.0.0.1:9999" + finally: + client.close() + + +def test_embedded_handle_async_client_uses_asgi_transport() -> None: + app = MagicMock() + runtime = MagicMock() + http_client = object() + client_value = object() + handle = services.EmbeddedServiceHandle(app=app, runtime=runtime) + + with ( + patch( + "nemo_platform.local.services.build_async_asgi_http_client", return_value=http_client + ) as build_client, + patch("nemo_platform.local.services.AsyncNeMoPlatform", return_value=client_value) as platform_cls, + ): + client = handle.async_client(access_token="test-token") + + build_client.assert_called_once_with(app) + platform_cls.assert_called_once_with( + access_token="test-token", + http_client=http_client, + base_url=services.EMBEDDED_BASE_URL, + ) + assert client is client_value + + +def test_ensure_services_dispatches_to_embedded_mode() -> None: + cfg = ServiceRunConfig(mode=services.ServiceMode.EMBEDDED) + embedded_handle = MagicMock(spec=services.EmbeddedServiceHandle) + + with patch("nemo_platform.local.services.start_embedded_services", return_value=embedded_handle): + handle = services.ensure_services(cfg) + + assert handle is embedded_handle + + +def test_ensure_services_dispatches_to_daemon_mode() -> None: + cfg = ServiceRunConfig(mode=services.ServiceMode.DAEMON) + daemon_handle = MagicMock(spec=services.DaemonServiceHandle) + + with ( + patch("nemo_platform.local.services.get_service_handle", return_value=None), + patch("nemo_platform.local.services.daemonize_services", return_value=daemon_handle), + ): + handle = services.ensure_services(cfg) + + assert handle is daemon_handle + + +def test_connect_services_uses_selected_mode_handle_client() -> None: + cfg = ServiceRunConfig(mode=services.ServiceMode.EMBEDDED) + handle = MagicMock(spec=services.EmbeddedServiceHandle) + client = object() + handle.client.return_value = client + + with patch("nemo_platform.local.services.ensure_services", return_value=handle): + result = services.connect_services(cfg, access_token="test") + + assert result is client + handle.client.assert_called_once_with(access_token="test") + + +@pytest.mark.parametrize("mode", [services.ServiceMode.EMBEDDED, services.ServiceMode.DAEMON]) +def test_ensure_services_returns_handle_with_parity_methods(mode: services.ServiceMode) -> None: + cfg = ServiceRunConfig(mode=mode) + if mode is services.ServiceMode.EMBEDDED: + handle = MagicMock(spec=services.EmbeddedServiceHandle) + patch_target = "nemo_platform.local.services.start_embedded_services" + else: + handle = MagicMock(spec=services.DaemonServiceHandle) + patch_target = "nemo_platform.local.services.daemonize_services" + + with ( + patch("nemo_platform.local.services.get_service_handle", return_value=None), + patch(patch_target, return_value=handle), + ): + result = services.ensure_services(cfg) + + assert result is handle + for method_name in ( + "is_running", + "wait_until_ready", + "wait_until_ready_async", + "client", + "async_client", + "start_services", + "start_services_async", + "stop", + "stop_async", + ): + assert hasattr(result, method_name), method_name + + +def test_daemonize_services_starts_child_module(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + _allow_tmp_path_socket_paths(monkeypatch) + cfg = ServiceRunConfig( + service_group="all", + scope="dev", + state_dir=tmp_path / "state", + runtime_dir=tmp_path / "run", + readiness_timeout=0.1, + readiness_poll_interval=0.01, + ) + proc = MagicMock() + proc.pid = 4242 + proc.poll.return_value = None + + with ( + patch("nemo_platform.local.services.require_services_extra"), + patch("nemo_platform.local.services._check_tcp_available"), + patch("nemo_platform.local.services.probe_status", return_value=True), + patch("nemo_platform.local.services.subprocess.Popen", return_value=proc) as popen, + ): + handle = services.daemonize_services(cfg) + + args = popen.call_args.args[0] + assert args[:3] == [sys.executable, "-m", f"{services.__package__}._service_child"] + request_path = Path(args[3]) + assert request_path.parent == tmp_path / "state" / "instances" / "dev" + assert request_path.suffix == ".json" + assert request_path.name != "run-request.json" + assert handle.transport == "uds" + assert handle.socket_path == tmp_path / "run" / "dev" / "nemo-platform.sock" + assert handle.pid == 4242 + proc.terminate.assert_not_called() + proc.kill.assert_not_called() + + +def test_daemonize_services_leaves_socket_preparation_to_child(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + _allow_tmp_path_socket_paths(monkeypatch) + cfg = ServiceRunConfig( + scope="dev", + state_dir=tmp_path / "state", + runtime_dir=tmp_path / "run", + readiness_timeout=0.1, + readiness_poll_interval=0.01, + ) + socket_path = cfg.resolved_socket_path + assert socket_path is not None + socket_path.parent.mkdir(parents=True) + socket_path.write_text("stale", encoding="utf-8") + proc = MagicMock() + proc.pid = 4242 + proc.poll.return_value = None + + with ( + patch("nemo_platform.local.services.require_services_extra"), + patch("nemo_platform.local.services._prepare_socket", side_effect=AssertionError("parent prepared socket")), + patch("nemo_platform.local.services.probe_status", side_effect=[False, True]), + patch("nemo_platform.local.services.subprocess.Popen", return_value=proc), + ): + handle = services.daemonize_services(cfg) + + assert handle.socket_path == socket_path + assert socket_path.read_text(encoding="utf-8") == "stale" + + +def test_write_run_request_writes_complete_payload_when_os_write_is_short( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + cfg = ServiceRunConfig(scope="dev", state_dir=tmp_path / "state", runtime_dir=tmp_path / "run") + real_write = os.write + + def short_write(fd: int, data: bytes) -> int: + return real_write(fd, data[: max(1, len(data) // 2)]) + + monkeypatch.setattr(services.os, "write", short_write) + + request_path = services._write_run_request(cfg) + + expected_payload = json.dumps(cfg.to_child_payload(), indent=2) + "\n" + assert request_path.read_text(encoding="utf-8") == expected_payload + + +def test_service_child_unlinks_request_after_read(tmp_path: Path) -> None: + request_path = tmp_path / "run-request.json" + cfg = ServiceRunConfig(scope="dev", state_dir=tmp_path / "state", runtime_dir=tmp_path / "run") + request_path.write_text(json.dumps(cfg.to_child_payload()), encoding="utf-8") + + with patch("nemo_platform.local._service_child.run_services") as run_services: + result = _service_child.main([str(request_path)]) + + assert result == 0 + assert not request_path.exists() + child_cfg = run_services.call_args.args[0] + assert child_cfg.scope == "dev" + assert child_cfg.state_dir == str(tmp_path / "state") + assert child_cfg.runtime_dir == str(tmp_path / "run") + assert run_services.call_args.kwargs == {"_mode": "daemon"} + + +def test_daemonize_services_terminates_child_on_timeout(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + _allow_tmp_path_socket_paths(monkeypatch) + cfg = ServiceRunConfig( + scope="dev", + state_dir=tmp_path / "state", + runtime_dir=tmp_path / "run", + readiness_timeout=0.01, + readiness_poll_interval=0.001, + ) + proc = MagicMock() + proc.pid = 4242 + proc.poll.return_value = None + + with ( + patch("nemo_platform.local.services.require_services_extra"), + patch("nemo_platform.local.services.probe_status", return_value=False), + patch("nemo_platform.local.services.subprocess.Popen", return_value=proc), + ): + with pytest.raises(services.ServicesStartupTimeoutError): + services.daemonize_services(cfg) + + proc.terminate.assert_called_once_with() + proc.kill.assert_not_called() + + +def test_daemonize_services_bounds_probe_and_sleep_by_remaining_deadline( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + _allow_tmp_path_socket_paths(monkeypatch) + cfg = ServiceRunConfig( + scope="dev", + state_dir=tmp_path / "state", + runtime_dir=tmp_path / "run", + readiness_timeout=5.0, + readiness_poll_interval=10.0, + ) + proc = MagicMock() + proc.pid = 4242 + proc.poll.return_value = None + + with ( + patch("nemo_platform.local.services.require_services_extra"), + patch("nemo_platform.local.services.probe_status", return_value=False) as probe_status, + patch("nemo_platform.local.services.subprocess.Popen", return_value=proc), + patch("nemo_platform.local.services.time.monotonic", side_effect=[0.0, 4.0, 4.5, 5.0]), + patch("nemo_platform.local.services.time.sleep") as sleep, + ): + with pytest.raises(services.ServicesStartupTimeoutError): + services.daemonize_services(cfg) + + assert probe_status.call_args.kwargs["timeout"] == pytest.approx(1.0) + sleep.assert_called_once() + assert sleep.call_args.args[0] == pytest.approx(0.5) + proc.terminate.assert_called_once_with() + + +def test_daemonize_services_terminates_child_on_handle_error(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + _allow_tmp_path_socket_paths(monkeypatch) + cfg = ServiceRunConfig( + scope="dev", + state_dir=tmp_path / "state", + runtime_dir=tmp_path / "run", + transport="tcp", + ) + proc = MagicMock() + proc.pid = 4242 + proc.poll.return_value = None + + with ( + patch("nemo_platform.local.services.require_services_extra"), + patch("nemo_platform.local.services._check_tcp_available"), + patch("nemo_platform.local.services.subprocess.Popen", return_value=proc), + patch( + "nemo_platform.local.services.DaemonServiceHandle.from_config", + side_effect=RuntimeError("handle failed"), + ), + ): + with pytest.raises(RuntimeError, match="handle failed"): + services.daemonize_services(cfg) + + proc.terminate.assert_called_once_with() + proc.kill.assert_not_called() + + +def test_daemonize_services_terminates_child_on_probe_error(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + _allow_tmp_path_socket_paths(monkeypatch) + cfg = ServiceRunConfig( + scope="dev", + state_dir=tmp_path / "state", + runtime_dir=tmp_path / "run", + transport="tcp", + ) + proc = MagicMock() + proc.pid = 4242 + proc.poll.return_value = None + + with ( + patch("nemo_platform.local.services.require_services_extra"), + patch("nemo_platform.local.services._check_tcp_available"), + patch("nemo_platform.local.services.probe_status", side_effect=RuntimeError("probe failed")), + patch("nemo_platform.local.services.subprocess.Popen", return_value=proc), + ): + with pytest.raises(RuntimeError, match="probe failed"): + services.daemonize_services(cfg) + + proc.terminate.assert_called_once_with() + proc.kill.assert_not_called() + + +def test_daemonize_services_terminates_child_on_sleep_interruption( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + _allow_tmp_path_socket_paths(monkeypatch) + cfg = ServiceRunConfig( + scope="dev", + state_dir=tmp_path / "state", + runtime_dir=tmp_path / "run", + transport="tcp", + ) + proc = MagicMock() + proc.pid = 4242 + proc.poll.return_value = None + + with ( + patch("nemo_platform.local.services.require_services_extra"), + patch("nemo_platform.local.services._check_tcp_available"), + patch("nemo_platform.local.services.probe_status", return_value=False), + patch("nemo_platform.local.services.subprocess.Popen", return_value=proc), + patch("nemo_platform.local.services.time.sleep", side_effect=KeyboardInterrupt), + ): + with pytest.raises(KeyboardInterrupt): + services.daemonize_services(cfg) + + proc.terminate.assert_called_once_with() + proc.kill.assert_not_called() + + +def test_daemonize_services_kills_child_when_terminate_times_out( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + _allow_tmp_path_socket_paths(monkeypatch) + cfg = ServiceRunConfig( + scope="dev", + state_dir=tmp_path / "state", + runtime_dir=tmp_path / "run", + transport="tcp", + readiness_timeout=0.01, + readiness_poll_interval=0.001, + ) + proc = MagicMock() + proc.pid = 4242 + proc.poll.return_value = None + proc.wait.side_effect = [subprocess.TimeoutExpired("nemo services", 5), None] + + with ( + patch("nemo_platform.local.services.require_services_extra"), + patch("nemo_platform.local.services._check_tcp_available"), + patch("nemo_platform.local.services.probe_status", return_value=False), + patch("nemo_platform.local.services.subprocess.Popen", return_value=proc), + ): + with pytest.raises(services.ServicesStartupTimeoutError): + services.daemonize_services(cfg) + + proc.terminate.assert_called_once_with() + proc.kill.assert_called_once_with() + + +async def test_daemonize_services_async_uses_thread(tmp_path: Path) -> None: + cfg = ServiceRunConfig(scope="dev", state_dir=tmp_path / "state", runtime_dir=tmp_path / "run") + handle = MagicMock() + + with patch("nemo_platform.local.services.asyncio.to_thread", new=AsyncMock(return_value=handle)) as to_thread: + result = await services.daemonize_services_async(cfg) + + assert result is handle + to_thread.assert_awaited_once_with(services.daemonize_services, cfg) + + +def test_run_services_serves_embedded_app_with_socket_path(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + _allow_tmp_path_socket_paths(monkeypatch) + cfg = ServiceRunConfig( + services=["entities"], + controllers=["jobs"], + scope="dev", + port=_free_tcp_port(), + state_dir=tmp_path / "state", + runtime_dir=tmp_path / "run", + ) + handle = _embedded_handle() + + with ( + patch("nemo_platform.local.services.require_services_extra"), + patch("nemo_platform.local.services.start_embedded_services", return_value=handle) as start_embedded, + patch("nemo_platform.local.services.serve_embedded_app") as serve_embedded, + ): + services.run_services(cfg, _mode="daemon") + + start_embedded.assert_called_once_with(cfg, env=None) + serve_embedded.assert_called_once_with(handle.app, cfg, tmp_path / "run" / "dev" / "nemo-platform.sock") + assert not (tmp_path / "state" / "instances" / "dev" / DESCRIPTOR_FILENAME).exists() + + +def test_serve_embedded_app_with_socket_path_listens_on_tcp_and_uds(tmp_path: Path) -> None: + cfg = ServiceRunConfig(transport="tcp", host="127.0.0.1", port=9090) + app = object() + socket_path = tmp_path / "nemo.sock" + + with patch("nmp.platform_runner.server._run_server_on_bound_sockets") as run_bound_sockets: + services.serve_embedded_app(app, cfg, socket_path) + + run_bound_sockets.assert_called_once_with(app, host="127.0.0.1", port=9090, socket_path=str(socket_path)) + + +def test_run_services_cleans_lock_when_log_path_resolution_fails( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + _allow_tmp_path_socket_paths(monkeypatch) + cfg = ServiceRunConfig( + scope="dev", port=_free_tcp_port(), state_dir=tmp_path / "state", runtime_dir=tmp_path / "run" + ) + + with ( + patch("nemo_platform.local.services.require_services_extra"), + patch.object(PlatformAppConfig, "log_file_path", side_effect=RuntimeError("boom")), + ): + with pytest.raises(RuntimeError, match="boom"): + services.run_services(cfg) + + assert not services.process.is_instance_alive(cfg.scope, base_dir=cfg.state_root) + + +def test_run_services_restores_data_dir_and_lock_when_descriptor_write_fails( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + _allow_tmp_path_socket_paths(monkeypatch) + cfg = ServiceRunConfig( + scope="dev", + port=_free_tcp_port(), + state_dir=tmp_path / "state", + runtime_dir=tmp_path / "run", + data_dir=tmp_path / "data", + ) + monkeypatch.delenv("NMP_DATA_DIR", raising=False) + + with ( + patch("nemo_platform.local.services.require_services_extra"), + patch("nemo_platform.local.services.process.write_descriptor", side_effect=RuntimeError("boom")), + ): + with pytest.raises(RuntimeError, match="boom"): + services.run_services(cfg) + + assert "NMP_DATA_DIR" not in os.environ + assert not services.process.is_instance_alive(cfg.scope, base_dir=cfg.state_root) + + +def test_run_services_restores_existing_data_dir(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + _allow_tmp_path_socket_paths(monkeypatch) + cfg = ServiceRunConfig( + scope="dev", port=_free_tcp_port(), state_dir=tmp_path / "state", runtime_dir=tmp_path / "run" + ) + monkeypatch.setenv("NMP_DATA_DIR", "/shell/data") + + with ( + patch("nemo_platform.local.services.require_services_extra"), + patch("nemo_platform.local.services.start_embedded_services", return_value=_embedded_handle()), + patch("nemo_platform.local.services.serve_embedded_app"), + ): + services.run_services(cfg) + + assert os.environ["NMP_DATA_DIR"] == "/shell/data" + + +def test_daemon_service_handle_tcp_client_uses_tcp_base_url(tmp_path: Path) -> None: + handle = services.DaemonServiceHandle( + scope="dev", + transport="tcp", + socket_path=None, + gateway_base_url=None, + host="0.0.0.0", + port=9090, + pid=123, + mode="daemon", + log_path=None, + state_dir=tmp_path / "state" / "instances" / "dev", + runtime_dir=None, + ) + + with patch("nemo_platform.local.services.NeMoPlatform") as sdk: + handle.client(timeout=12) + + sdk.assert_called_once_with(timeout=12, base_url="http://localhost:9090") + + +def test_daemon_service_handle_uds_client_requires_socket_path(tmp_path: Path) -> None: + handle = services.DaemonServiceHandle( + scope="dev", + transport="uds", + socket_path=None, + gateway_base_url=None, + host="127.0.0.1", + port=8080, + pid=123, + mode="daemon", + log_path=None, + state_dir=tmp_path / "state" / "instances" / "dev", + runtime_dir=tmp_path / "run", + ) + + with pytest.raises(services.ServicesError, match="missing socket_path"): + handle.client() + + +def test_ensure_services_returns_existing_handle(tmp_path: Path) -> None: + cfg = ServiceRunConfig(scope="dev", state_dir=tmp_path / "state", runtime_dir=tmp_path / "run") + handle = MagicMock() + + with ( + patch("nemo_platform.local.services.get_service_handle", return_value=handle), + patch("nemo_platform.local.services.daemonize_services") as daemonize, + ): + result = services.ensure_services(cfg) + + assert result is handle + daemonize.assert_not_called() + + +def test_connect_services_respects_start_if_needed_false(tmp_path: Path) -> None: + cfg = ServiceRunConfig(scope="dev", state_dir=tmp_path / "state", runtime_dir=tmp_path / "run") + + with patch("nemo_platform.local.services.get_service_handle", return_value=None): + with pytest.raises(services.ServicesNotRunningError, match="not running"): + services.connect_services(cfg, start_if_needed=False) + + +def test_stop_services_delegates_to_handle(tmp_path: Path) -> None: + cfg = ServiceRunConfig(scope="dev", state_dir=tmp_path / "state", runtime_dir=tmp_path / "run") + handle = MagicMock() + stop_result = MagicMock() + handle.stop.return_value = stop_result + + with patch("nemo_platform.local.services.get_service_handle", return_value=handle): + result = services.stop_services(cfg, timeout=3.0, force=True) + + assert result is stop_result + handle.stop.assert_called_once_with(timeout=3.0, force=True) + + +def test_get_service_handle_returns_none_without_live_descriptor(tmp_path: Path) -> None: + cfg = ServiceRunConfig(scope="dev", state_dir=tmp_path / "state", runtime_dir=tmp_path / "run") + + with patch("nemo_platform.local.services.process.read_descriptor", return_value=None): + assert services.get_service_handle(cfg) is None + + +def test_list_service_handles_filters_dead_or_descriptorless_instances(tmp_path: Path) -> None: + live_desc = InstanceDescriptor( + pid=123, + transport="tcp", + config=PlatformAppConfig(scope="live", state_root=tmp_path / "state"), + mode="daemon", + ) + infos = [ + MagicMock(descriptor=live_desc, alive=True), + MagicMock(descriptor=None, alive=True), + MagicMock(descriptor=live_desc, alive=False), + ] + + with patch("nemo_platform.local.services.process.list_instances", return_value=infos): + handles = services.list_service_handles(tmp_path / "state") + + assert [handle.scope for handle in handles] == ["live"] + + +def test_get_service_handle_reads_live_descriptor(tmp_path: Path) -> None: + cfg = ServiceRunConfig(scope="dev", state_dir=tmp_path / "state", runtime_dir=tmp_path / "run") + state_dir = tmp_path / "state" / "instances" / "dev" + state_dir.mkdir(parents=True) + desc = InstanceDescriptor( + pid=123, + config=PlatformAppConfig( + scope="dev", + socket_path=str(tmp_path / "run" / "dev" / "nemo-platform.sock"), + state_root=tmp_path / "state", + runtime_root=tmp_path / "run", + ), + transport="uds", + mode="daemon", + ) + (state_dir / DESCRIPTOR_FILENAME).write_text(desc.model_dump_json(), encoding="utf-8") + + with patch("nemo_platform.local.process.is_instance_alive", return_value=True): + handle = services.get_service_handle(cfg) + + assert handle is not None + assert handle.scope == "dev" + assert handle.transport == "uds" diff --git a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/local/test_services_contract.py b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/local/test_services_contract.py new file mode 100644 index 0000000000..747a3b0850 --- /dev/null +++ b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/local/test_services_contract.py @@ -0,0 +1,334 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import threading +from collections.abc import Callable +from dataclasses import dataclass, field +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest +from nemo_platform.local import services +from nemo_platform.local.process import StopResult +from nemo_platform.local.services import ServiceRunConfig +from nmp.platform_runner.config import PlatformAppConfig + + +@dataclass(frozen=True) +class ModeContractCase: + mode: services.ServiceMode + launcher_patch: str + existing_handle_patch_value: object | None + + +@dataclass +class ContractHandle: + mode: services.ServiceMode + calls: list[tuple[str, object]] = field(default_factory=list) + + def is_running(self) -> bool: + self.calls.append(("is_running", None)) + return True + + def wait_until_ready(self, timeout: float | None = None) -> None: + self.calls.append(("wait_until_ready", timeout)) + + async def wait_until_ready_async(self, timeout: float | None = None) -> None: + self.calls.append(("wait_until_ready_async", timeout)) + + def client(self, **kwargs: object) -> tuple[str, services.ServiceMode, dict[str, object]]: + self.calls.append(("client", kwargs)) + return ("client", self.mode, kwargs) + + def async_client(self, **kwargs: object) -> tuple[str, services.ServiceMode, dict[str, object]]: + self.calls.append(("async_client", kwargs)) + return ("async_client", self.mode, kwargs) + + def start_services(self, service_names: list[str] | tuple[str, ...]) -> services.StartServicesResult: + requested = list(service_names) + self.calls.append(("start_services", requested)) + return services.StartServicesResult( + requested=requested, + started=["auth", *requested], + already_active=[], + active=["secrets", "auth", *requested], + ) + + async def start_services_async(self, service_names: list[str] | tuple[str, ...]) -> services.StartServicesResult: + requested = list(service_names) + self.calls.append(("start_services_async", requested)) + return services.StartServicesResult( + requested=requested, + started=["auth", *requested], + already_active=[], + active=["secrets", "auth", *requested], + ) + + def stop(self, *, timeout: float = 30.0, force: bool = False) -> StopResult: + self.calls.append(("stop", {"timeout": timeout, "force": force})) + return StopResult(stopped_pids=[], swept_children=[]) + + async def stop_async(self, *, timeout: float = 30.0, force: bool = False) -> StopResult: + self.calls.append(("stop_async", {"timeout": timeout, "force": force})) + return StopResult(stopped_pids=[], swept_children=[]) + + +MODE_CONTRACT_CASES = [ + ModeContractCase( + mode=services.ServiceMode.EMBEDDED, + launcher_patch="nemo_platform.local.services.start_embedded_services", + existing_handle_patch_value=None, + ), + ModeContractCase( + mode=services.ServiceMode.DAEMON, + launcher_patch="nemo_platform.local.services.daemonize_services", + existing_handle_patch_value=None, + ), +] + + +@pytest.fixture(params=MODE_CONTRACT_CASES, ids=lambda case: case.mode.value) +def mode_case(request: pytest.FixtureRequest) -> ModeContractCase: + return request.param + + +def _config_for(case: ModeContractCase, tmp_path: Path) -> ServiceRunConfig: + return ServiceRunConfig( + mode=case.mode, + services=("secrets",), + scope=f"{case.mode.value}-contract", + state_dir=tmp_path / case.mode.value / "state", + runtime_dir=tmp_path / case.mode.value / "runtime", + ) + + +def test_contract_ensure_services_returns_running_mode_handle( + mode_case: ModeContractCase, + tmp_path: Path, +) -> None: + cfg = _config_for(mode_case, tmp_path) + handle = ContractHandle(mode_case.mode) + + with ( + patch( + "nemo_platform.local.services.get_service_handle", return_value=mode_case.existing_handle_patch_value + ), + patch(mode_case.launcher_patch, return_value=handle), + ): + result = services.ensure_services(cfg) + + assert result is handle + assert result.is_running() is True + assert result.calls == [("is_running", None)] + + +def test_contract_connect_services_returns_client_from_selected_mode( + mode_case: ModeContractCase, + tmp_path: Path, +) -> None: + cfg = _config_for(mode_case, tmp_path) + handle = ContractHandle(mode_case.mode) + + with patch("nemo_platform.local.services.ensure_services", return_value=handle): + client = services.connect_services(cfg, api_key="test-key") + + assert client == ("client", mode_case.mode, {"api_key": "test-key"}) + assert handle.calls == [("client", {"api_key": "test-key"})] + + +@pytest.mark.asyncio +async def test_contract_handle_lifecycle_methods_have_same_semantics( + mode_case: ModeContractCase, +) -> None: + handle = ContractHandle(mode_case.mode) + + handle.wait_until_ready(timeout=1.5) + await handle.wait_until_ready_async(timeout=2.5) + sync_start = handle.start_services(["jobs"]) + async_start = await handle.start_services_async(["jobs"]) + stop_result = handle.stop(timeout=3.0, force=True) + async_stop_result = await handle.stop_async(timeout=4.0, force=False) + + assert sync_start == services.StartServicesResult( + requested=["jobs"], + started=["auth", "jobs"], + already_active=[], + active=["secrets", "auth", "jobs"], + ) + assert async_start == sync_start + assert stop_result == StopResult(stopped_pids=[], swept_children=[]) + assert async_stop_result == StopResult(stopped_pids=[], swept_children=[]) + assert handle.calls == [ + ("wait_until_ready", 1.5), + ("wait_until_ready_async", 2.5), + ("start_services", ["jobs"]), + ("start_services_async", ["jobs"]), + ("stop", {"timeout": 3.0, "force": True}), + ("stop_async", {"timeout": 4.0, "force": False}), + ] + + +def test_contract_real_handles_report_same_staged_start_status_before_staged_start_lands(tmp_path: Path) -> None: + embedded = services.EmbeddedServiceHandle(app=object(), runtime=object()) + daemon = services.DaemonServiceHandle( + scope="daemon-contract", + transport="tcp", + socket_path=None, + gateway_base_url=None, + host="127.0.0.1", + port=8080, + pid=None, + mode="daemon", + log_path=None, + state_dir=tmp_path / "state", + runtime_dir=None, + ) + + for handle in (embedded, daemon): + with pytest.raises(services.ServicesError, match="Staged service start is not implemented"): + handle.start_services(["jobs"]) + + +def test_contract_embedded_and_daemon_child_both_delegate_models_to_platform_builder( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls: list[dict[str, object]] = [] + + def fake_build_platform_app( + config: PlatformAppConfig | None = None, + *, + env: object = None, + http_client: object = None, + ) -> MagicMock: + calls.append({"config": config, "env": env, "http_client": http_client}) + return MagicMock() + + def service_config(mode: services.ServiceMode) -> ServiceRunConfig: + return ServiceRunConfig( + mode=mode, + services=("models",), + controllers=(), + transport="tcp", + scope="sc-test", + state_dir=tmp_path / "state", + runtime_dir=tmp_path / "runtime", + ) + + with patch("nmp.platform_runner.server.build_platform_app", side_effect=fake_build_platform_app): + services.start_embedded_services(service_config(services.ServiceMode.EMBEDDED)) + + with ( + patch("nmp.platform_runner.server.build_platform_app", side_effect=fake_build_platform_app), + patch("nemo_platform.local.services.require_services_extra"), + patch("nemo_platform.local.services.process.is_instance_alive", return_value=False), + patch("nemo_platform.local.services._check_tcp_available"), + patch("nemo_platform.local.services.process.acquire_lock", return_value=123), + patch("nemo_platform.local.services.process.log_path_for", return_value=tmp_path / "nemo.log"), + patch("nemo_platform.local.services.process.write_descriptor"), + patch("nemo_platform.local.services.process.remove_descriptor"), + patch("nemo_platform.local.services.serve_embedded_app"), + patch("nemo_platform.local.services.os.close"), + ): + services.run_services(service_config(services.ServiceMode.DAEMON), _mode="daemon") + + configs: list[PlatformAppConfig] = [] + for call in calls: + config = call["config"] + assert isinstance(config, PlatformAppConfig) + configs.append(config) + assert [config.services for config in configs] == [("models",), ("models",)] + assert [config.controllers for config in configs] == [(), ()] + assert [config.sidecars for config in configs] == [None, None] + + +def _sidecar_with_events(started: threading.Event, stopped: threading.Event) -> Callable[[threading.Event], None]: + def run(stop_signal: threading.Event) -> None: + started.set() + stop_signal.wait(timeout=5.0) + stopped.set() + + return run + + +def _patch_runner_registry( + monkeypatch: pytest.MonkeyPatch, + *, + sidecar_run_func: Callable[[threading.Event], None], +) -> None: + """Patch the platform runner registry so only a dummy 'models' service + and a test sidecar are available, avoiding real service imports.""" + from nmp.common.config import AuthConfig + from nmp.common.config.base import OIDCConfig + from nmp.common.service import Service + from nmp.platform_runner import config as runner_config + from nmp.platform_runner import registry, server + + class _DummyService(Service): + def __init__(self) -> None: + super().__init__(name="models", module_name="test.contract") + + def get_routers(self): + return [] + + dummy_services: dict[str, Service] = {"models": _DummyService()} + dummy_sidecars: dict[str, Callable] = {"adapters": sidecar_run_func} + + monkeypatch.setattr(runner_config, "get_available_services", lambda: dummy_services) + monkeypatch.setattr(runner_config, "get_available_controllers", lambda: {}) + monkeypatch.setattr( + runner_config, + "get_service_groups", + lambda _available: {"all": ["models"], "core": ["models"], "api": []}, + ) + monkeypatch.setattr(runner_config, "get_controller_groups", lambda _available: {"all": [], "core": []}) + monkeypatch.setattr(runner_config, "get_default_controllers", lambda _groups: []) + monkeypatch.setattr(runner_config, "AVAILABLE_SIDECARS", dummy_sidecars) + monkeypatch.setattr(registry, "AVAILABLE_SIDECARS", dummy_sidecars) + monkeypatch.setattr(server, "AVAILABLE_SIDECARS", dummy_sidecars, raising=False) + monkeypatch.setattr(server, "order_services_by_dependencies", lambda svc: svc) + + auth_cfg = AuthConfig( + enabled=False, + policy_decision_point_base_url="http://localhost:8181", + oidc=OIDCConfig(enabled=False), + ) + monkeypatch.setattr(server, "get_auth_config", lambda: auth_cfg) + monkeypatch.setattr("nmp.common.auth.middleware.get_auth_config", lambda: auth_cfg) + platform_cfg = MagicMock() + platform_cfg.seed_on_startup = False + platform_cfg.redirect_root_to_studio = False + monkeypatch.setattr(server, "get_platform_config", lambda: platform_cfg) + + +def test_embedded_mode_starts_sidecar_thread_via_full_resolution_path( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """End-to-end: start_embedded_services(models) resolves the adapters sidecar + and the sidecar thread actually runs when the app lifespan starts.""" + started = threading.Event() + stopped = threading.Event() + + _patch_runner_registry(monkeypatch, sidecar_run_func=_sidecar_with_events(started, stopped)) + + cfg = ServiceRunConfig( + mode=services.ServiceMode.EMBEDDED, + services=("models",), + controllers=(), + transport="tcp", + scope="sidecar-e2e", + state_dir=tmp_path / "state", + runtime_dir=tmp_path / "runtime", + ) + handle = services.start_embedded_services(cfg, env={}) + + from fastapi.testclient import TestClient + + with TestClient(handle.app) as client: + assert started.wait(timeout=2.0), "sidecar thread did not start" + assert client.get("/").status_code == 200 + + assert stopped.wait(timeout=2.0), "sidecar thread did not stop" diff --git a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/local/test_sidecar_integration.py b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/local/test_sidecar_integration.py new file mode 100644 index 0000000000..ae813ac3f4 --- /dev/null +++ b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/local/test_sidecar_integration.py @@ -0,0 +1,231 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Integration tests for sidecar lifecycle in embedded and daemon modes. + +These tests let the real ``build_platform_app`` → ``resolve_run_configuration`` → +``create_app`` chain run with a lightweight test sidecar registered in the +platform runner registry. They verify that sidecar threads actually start and +stop during the FastAPI app lifespan, covering the full resolution path without +mocking away the core wiring. +""" + +from __future__ import annotations + +import threading +from collections.abc import Callable +from pathlib import Path +from unittest.mock import MagicMock + +import pytest +from nemo_platform.local import services +from nemo_platform.local.services import ServiceRunConfig +from nmp.common.config import AuthConfig +from nmp.common.config.base import OIDCConfig +from nmp.common.service import Service +from nmp.platform_runner import config as runner_config +from nmp.platform_runner import registry, server + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +class _DummyService(Service): + """Minimal service that registers no routers.""" + + def __init__(self, name: str = "models") -> None: + super().__init__(name=name, module_name="test.sidecar_integration") + + def get_routers(self): + return [] + + +def _sidecar_with_events(started: threading.Event, stopped: threading.Event) -> Callable[[threading.Event], None]: + """Return a sidecar ``run(stop_signal)`` that signals start/stop via events.""" + + def run(stop_signal: threading.Event) -> None: + started.set() + stop_signal.wait(timeout=5.0) + stopped.set() + + return run + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture +def sidecar_events() -> tuple[threading.Event, threading.Event]: + return threading.Event(), threading.Event() + + +@pytest.fixture +def patched_registry( + monkeypatch: pytest.MonkeyPatch, + sidecar_events: tuple[threading.Event, threading.Event], +) -> tuple[threading.Event, threading.Event]: + """Patch the platform runner registry with a dummy models service and a + test sidecar, plus minimal auth/platform config stubs.""" + started, stopped = sidecar_events + dummy_services: dict[str, Service] = {"models": _DummyService()} + dummy_sidecars: dict[str, Callable] = {"adapters": _sidecar_with_events(started, stopped)} + + monkeypatch.setattr(runner_config, "get_available_services", lambda: dummy_services) + monkeypatch.setattr(runner_config, "get_available_controllers", lambda: {}) + monkeypatch.setattr( + runner_config, + "get_service_groups", + lambda _available: {"all": ["models"], "core": ["models"], "api": []}, + ) + monkeypatch.setattr(runner_config, "get_controller_groups", lambda _available: {"all": [], "core": []}) + monkeypatch.setattr(runner_config, "get_default_controllers", lambda _groups: []) + monkeypatch.setattr(runner_config, "AVAILABLE_SIDECARS", dummy_sidecars) + monkeypatch.setattr(registry, "AVAILABLE_SIDECARS", dummy_sidecars) + monkeypatch.setattr(server, "AVAILABLE_SIDECARS", dummy_sidecars, raising=False) + monkeypatch.setattr(server, "order_services_by_dependencies", lambda svc: svc) + + auth_cfg = AuthConfig( + enabled=False, + policy_decision_point_base_url="http://localhost:8181", + oidc=OIDCConfig(enabled=False), + ) + monkeypatch.setattr(server, "get_auth_config", lambda: auth_cfg) + monkeypatch.setattr("nmp.common.auth.middleware.get_auth_config", lambda: auth_cfg) + platform_cfg = MagicMock() + platform_cfg.seed_on_startup = False + platform_cfg.redirect_root_to_studio = False + monkeypatch.setattr(server, "get_platform_config", lambda: platform_cfg) + + return started, stopped + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +@pytest.mark.integration +def test_embedded_sidecar_auto_resolved_from_service_dependency( + patched_registry: tuple[threading.Event, threading.Event], + tmp_path: Path, +) -> None: + """start_embedded_services(models) auto-resolves the adapters sidecar via + SERVICE_SIDECAR_DEPENDENCIES and starts it during app lifespan.""" + started, stopped = patched_registry + + cfg = ServiceRunConfig( + mode=services.ServiceMode.EMBEDDED, + services=("models",), + controllers=(), + # sidecars=None triggers auto-resolution + transport="tcp", + scope="integ-embedded-auto", + state_dir=tmp_path / "state", + runtime_dir=tmp_path / "runtime", + ) + handle = services.start_embedded_services(cfg, env={}) + + from fastapi.testclient import TestClient + + with TestClient(handle.app) as client: + assert started.wait(timeout=2.0), "sidecar thread did not start" + assert client.get("/").status_code == 200 + + assert stopped.wait(timeout=2.0), "sidecar thread did not stop after lifespan exit" + + +@pytest.mark.integration +def test_embedded_explicit_sidecar_without_services( + patched_registry: tuple[threading.Event, threading.Event], + tmp_path: Path, +) -> None: + """An explicitly requested sidecar runs even when no services are selected.""" + started, stopped = patched_registry + + cfg = ServiceRunConfig( + mode=services.ServiceMode.EMBEDDED, + services=(), + controllers=(), + sidecars=("adapters",), + transport="tcp", + scope="integ-embedded-explicit", + state_dir=tmp_path / "state", + runtime_dir=tmp_path / "runtime", + ) + handle = services.start_embedded_services(cfg, env={}) + + from fastapi.testclient import TestClient + + with TestClient(handle.app) as client: + assert started.wait(timeout=2.0), "sidecar thread did not start" + assert client.get("/").status_code == 200 + + assert stopped.wait(timeout=2.0), "sidecar thread did not stop after lifespan exit" + + +@pytest.mark.integration +def test_run_services_daemon_mode_starts_sidecar_in_process( + patched_registry: tuple[threading.Event, threading.Event], + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """run_services(_mode='daemon') exercises the daemon code path in-process. + It calls start_embedded_services then serve_embedded_app. We intercept + serve_embedded_app to capture the app and exercise its lifespan, proving + the daemon path wires sidecars identically to embedded mode.""" + started, stopped = patched_registry + captured_app = {} + + def fake_serve(app, cfg, socket_path): + captured_app["app"] = app + + monkeypatch.setattr("nemo_platform.local.services.require_services_extra", lambda: None) + monkeypatch.setattr("nemo_platform.local.services.process.is_instance_alive", lambda *a, **kw: False) + monkeypatch.setattr("nemo_platform.local.services._check_tcp_available", lambda *a: None) + monkeypatch.setattr("nemo_platform.local.services.process.acquire_lock", lambda *a, **kw: 123) + monkeypatch.setattr("nemo_platform.local.services.process.log_path_for", lambda *a, **kw: tmp_path / "nemo.log") + monkeypatch.setattr("nemo_platform.local.services.process.write_descriptor", lambda *a, **kw: None) + monkeypatch.setattr("nemo_platform.local.services.process.remove_descriptor", lambda *a, **kw: None) + monkeypatch.setattr("nemo_platform.local.services.serve_embedded_app", fake_serve) + monkeypatch.setattr("nemo_platform.local.services.os.close", lambda fd: None) + + cfg = ServiceRunConfig( + mode=services.ServiceMode.DAEMON, + services=("models",), + controllers=(), + transport="tcp", + scope="integ-daemon", + state_dir=tmp_path / "state", + runtime_dir=tmp_path / "runtime", + ) + services.run_services(cfg, _mode="daemon", env={}) + + assert "app" in captured_app, "serve_embedded_app was not called" + + from fastapi.testclient import TestClient + + with TestClient(captured_app["app"]) as client: + assert started.wait(timeout=2.0), "sidecar thread did not start in daemon mode" + assert client.get("/").status_code == 200 + + assert stopped.wait(timeout=2.0), "sidecar thread did not stop after lifespan exit" + + +@pytest.mark.integration +def test_embedded_rejects_unknown_sidecar_name(tmp_path: Path, patched_registry) -> None: + """Requesting a sidecar not in the registry raises ValueError with a clear message.""" + cfg = ServiceRunConfig( + mode=services.ServiceMode.EMBEDDED, + services=(), + controllers=(), + sidecars=("nonexistent",), + transport="tcp", + scope="integ-unknown", + state_dir=tmp_path / "state", + runtime_dir=tmp_path / "runtime", + ) + with pytest.raises(ValueError, match="Unknown sidecars: nonexistent"): + services.start_embedded_services(cfg, env={}) diff --git a/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/local/test_transport.py b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/local/test_transport.py new file mode 100644 index 0000000000..79251dece3 --- /dev/null +++ b/sdk/python/nemo-platform/tests/vendored/nemo_platform_ext/local/test_transport.py @@ -0,0 +1,167 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from unittest.mock import AsyncMock, patch + +import httpx +import pytest +from fastapi import FastAPI +from nemo_platform.local import transport + + +def _assert_timeout_values(timeout: httpx.Timeout, expected: float | None) -> None: + assert timeout.connect == expected + assert timeout.read == expected + assert timeout.write == expected + assert timeout.pool == expected + + +def test_build_sync_http_client_uses_finite_default_timeout(tmp_path) -> None: + client = transport.build_sync_http_client(tmp_path / "nemo.sock") + try: + _assert_timeout_values(client.timeout, 5.0) + finally: + client.close() + + +def test_build_sync_http_client_preserves_explicit_timeout_values(tmp_path) -> None: + no_timeout_client = transport.build_sync_http_client(tmp_path / "nemo.sock", timeout=None) + finite_timeout_client = transport.build_sync_http_client(tmp_path / "nemo.sock", timeout=12.0) + try: + _assert_timeout_values(no_timeout_client.timeout, None) + _assert_timeout_values(finite_timeout_client.timeout, 12.0) + finally: + no_timeout_client.close() + finite_timeout_client.close() + + +@pytest.mark.asyncio +async def test_build_async_http_client_uses_finite_default_timeout(tmp_path) -> None: + client = transport.build_async_http_client(tmp_path / "nemo.sock") + try: + _assert_timeout_values(client.timeout, 5.0) + finally: + await client.aclose() + + +@pytest.mark.asyncio +async def test_build_async_http_client_preserves_explicit_timeout_values(tmp_path) -> None: + no_timeout_client = transport.build_async_http_client(tmp_path / "nemo.sock", timeout=None) + finite_timeout_client = transport.build_async_http_client(tmp_path / "nemo.sock", timeout=12.0) + try: + _assert_timeout_values(no_timeout_client.timeout, None) + _assert_timeout_values(finite_timeout_client.timeout, 12.0) + finally: + await no_timeout_client.aclose() + await finite_timeout_client.aclose() + + +def test_build_sync_asgi_http_client_reaches_app() -> None: + app = FastAPI() + + @app.get("/status") + async def status() -> dict[str, str]: + return {"status": "healthy"} + + client = transport.build_sync_asgi_http_client(app) + try: + response = client.get("http://nemo-platform.local/status") + finally: + client.close() + + assert response.status_code == 200 + assert response.json() == {"status": "healthy"} + + +@pytest.mark.asyncio +async def test_build_async_asgi_http_client_reaches_app() -> None: + app = FastAPI() + + @app.get("/status") + async def status() -> dict[str, str]: + return {"status": "healthy"} + + client = transport.build_async_asgi_http_client(app) + try: + response = await client.get("http://nemo-platform.local/status") + finally: + await client.aclose() + + assert response.status_code == 200 + assert response.json() == {"status": "healthy"} + + +def test_wait_for_status_bounds_probe_and_sleep_by_remaining_deadline() -> None: + with ( + patch("nemo_platform.local.transport.probe_status", return_value=False) as probe_status, + patch("nemo_platform.local.transport.time.monotonic", side_effect=[0.0, 4.0, 4.5, 5.0]), + patch("nemo_platform.local.transport.time.sleep") as sleep, + ): + result = transport.wait_for_status(base_url="http://127.0.0.1:8080", timeout=5.0, poll_interval=10.0) + + assert result is False + assert probe_status.call_args.kwargs["timeout"] == pytest.approx(1.0) + sleep.assert_called_once() + assert sleep.call_args.args[0] == pytest.approx(0.5) + + +@pytest.mark.asyncio +async def test_wait_for_status_async_bounds_probe_and_sleep_by_remaining_deadline() -> None: + with ( + patch( + "nemo_platform.local.transport.probe_status_async", new=AsyncMock(return_value=False) + ) as probe_status, + patch("nemo_platform.local.transport.time.monotonic", side_effect=[0.0, 4.0, 4.5, 5.0]), + patch("nemo_platform.local.transport.asyncio.sleep", new=AsyncMock()) as sleep, + ): + result = await transport.wait_for_status_async( + base_url="http://127.0.0.1:8080", timeout=5.0, poll_interval=10.0 + ) + + assert result is False + assert probe_status.await_args.kwargs["timeout"] == pytest.approx(1.0) + sleep.assert_awaited_once() + assert sleep.await_args.args[0] == pytest.approx(0.5) + + +def test_probe_status_returns_true_for_status_200() -> None: + def handler(request: httpx.Request) -> httpx.Response: + assert str(request.url) == "http://127.0.0.1:8080/status" + return httpx.Response(200) + + with patch("nemo_platform.local.transport.httpx.Client") as client_factory: + client = client_factory.return_value + client.get.side_effect = lambda url: handler(httpx.Request("GET", url)) + assert transport.probe_status(base_url="http://127.0.0.1:8080") is True + client.close.assert_called_once_with() + + +def test_probe_status_returns_false_for_request_error() -> None: + with patch("nemo_platform.local.transport.httpx.Client") as client_factory: + client = client_factory.return_value + client.get.side_effect = httpx.ConnectError("boom") + assert transport.probe_status(base_url="http://127.0.0.1:8080") is False + client.close.assert_called_once_with() + + +@pytest.mark.asyncio +async def test_probe_status_async_returns_false_for_request_error() -> None: + with patch("nemo_platform.local.transport.httpx.AsyncClient") as client_factory: + client = client_factory.return_value + client.get = AsyncMock(side_effect=httpx.ConnectError("boom")) + client.aclose = AsyncMock() + assert await transport.probe_status_async(base_url="http://127.0.0.1:8080") is False + client.aclose.assert_awaited_once_with() + + +def test_wait_for_status_returns_true_without_sleep_when_probe_succeeds() -> None: + with ( + patch("nemo_platform.local.transport.probe_status", return_value=True) as probe_mock, + patch("nemo_platform.local.transport.time.sleep") as sleep, + ): + assert transport.wait_for_status(base_url="http://127.0.0.1:8080", timeout=5.0) is True + + probe_mock.assert_called_once() + sleep.assert_not_called() diff --git a/services/core/jobs/jobs-launcher/cmd/otel.go b/services/core/jobs/jobs-launcher/cmd/otel.go index 92b03cdabc..1e3664a2e7 100644 --- a/services/core/jobs/jobs-launcher/cmd/otel.go +++ b/services/core/jobs/jobs-launcher/cmd/otel.go @@ -9,12 +9,15 @@ import ( "fmt" "log/slog" "net/http" + "net/url" "os" + "strconv" + "strings" "sync/atomic" "time" + "github.com/NVIDIA-NeMo/nemo-platform/services/core/jobs/jobs-launcher/nmpclient" "go.opentelemetry.io/contrib/bridges/otelslog" - "go.opentelemetry.io/contrib/exporters/autoexport" "go.opentelemetry.io/otel/attribute" "go.opentelemetry.io/otel/exporters/otlp/otlplog/otlploghttp" "go.opentelemetry.io/otel/exporters/stdout/stdoutlog" @@ -24,14 +27,22 @@ import ( ) const ( - name = "nmp.nvidia.com/nemo-platform/jobs-launcher" - NEMO_JOB_WORKSPACE = "NEMO_JOB_WORKSPACE" - NEMO_JOB_ID_ENV = "NEMO_JOB_ID" - NEMO_JOB_ATTEMPT_ID_ENV = "NEMO_JOB_ATTEMPT_ID" - NEMO_JOB_STEP_NAME_ENV = "NEMO_JOB_STEP" - NEMO_JOB_TASK_ID_ENV = "NEMO_JOB_TASK" - nmpJobLogsEndpointEnv = "NMP_JOB_LOGS_ENDPOINT" - otlpHTTPLogExportTimeout = 10 * time.Second + name = "nmp.nvidia.com/nemo-platform/jobs-launcher" + NEMO_JOB_WORKSPACE = "NEMO_JOB_WORKSPACE" + NEMO_JOB_ID_ENV = "NEMO_JOB_ID" + NEMO_JOB_ATTEMPT_ID_ENV = "NEMO_JOB_ATTEMPT_ID" + NEMO_JOB_STEP_NAME_ENV = "NEMO_JOB_STEP" + NEMO_JOB_TASK_ID_ENV = "NEMO_JOB_TASK" + + launcherLogsExporterEnv = "NMP_JOB_LAUNCHER_LOGS_EXPORTER" + launcherOTLPLogsEndpointEnv = "NMP_JOB_LAUNCHER_OTLP_LOGS_ENDPOINT" + launcherOTLPLogsHeadersEnv = "NMP_JOB_LAUNCHER_OTLP_LOGS_HEADERS" + launcherOTLPLogsProtocolEnv = "NMP_JOB_LAUNCHER_OTLP_LOGS_PROTOCOL" + launcherOTLPLogsTimeoutEnv = "NMP_JOB_LAUNCHER_OTLP_LOGS_TIMEOUT" + launcherOTLPLogsCompressEnv = "NMP_JOB_LAUNCHER_OTLP_LOGS_COMPRESSION" + launcherOTLPHTTPProto = "http/protobuf" + defaultLauncherLogsExporter = "console" + otlpHTTPLogExportTimeout = 10 * time.Second ) var ( @@ -111,26 +122,64 @@ func newLoggerProvider(ctx context.Context, res *resource.Resource) (*log.Logger } func newLogExporter(ctx context.Context) (log.Exporter, error) { - if endpoint := os.Getenv(nmpJobLogsEndpointEnv); endpoint != "" { - if os.Getenv(workloadIdentityTokenFileEnv) != "" { - tokenSource, err := newOTLPLogWorkloadAuthTokenSource(ctx) - if err != nil { - return nil, fmt.Errorf("configure workload identity auth for OTLP logs: %w", err) - } - return newRefreshableAuthLogExporter(ctx, endpoint, tokenSource, otlpHTTPLogExporter) + switch strings.ToLower(strings.TrimSpace(os.Getenv(launcherLogsExporterEnv))) { + case "", defaultLauncherLogsExporter, "stdout": + return stdoutlog.New() + case "none": + return noopLogExporter{}, nil + case "otlp": + return newLauncherOTLPLogExporter(ctx) + default: + return nil, fmt.Errorf("unsupported %s value %q", launcherLogsExporterEnv, os.Getenv(launcherLogsExporterEnv)) + } +} + +func newLauncherOTLPLogExporter(ctx context.Context) (log.Exporter, error) { + protocol := strings.TrimSpace(os.Getenv(launcherOTLPLogsProtocolEnv)) + if protocol == "" { + protocol = launcherOTLPHTTPProto + } + if protocol != launcherOTLPHTTPProto { + return nil, fmt.Errorf("%s must be %q, got %q", launcherOTLPLogsProtocolEnv, launcherOTLPHTTPProto, protocol) + } + + endpointURL := strings.TrimSpace(os.Getenv(launcherOTLPLogsEndpointEnv)) + if endpointURL == "" { + return nil, fmt.Errorf("%s is required when %s=otlp", launcherOTLPLogsEndpointEnv, launcherLogsExporterEnv) + } + + // If a workload identity token file is available, use a refreshable auth + // transport that exchanges the projected SA token for platform credentials. + if os.Getenv(workloadIdentityTokenFileEnv) != "" { + tokenSource, err := newOTLPLogWorkloadAuthTokenSource(ctx) + if err != nil { + return nil, fmt.Errorf("configure workload identity auth for OTLP logs: %w", err) } - return otlploghttp.New(ctx, otlploghttp.WithEndpointURL(endpoint)) + return newRefreshableAuthLogExporter(ctx, endpointURL, tokenSource, otlpHTTPLogExporter) } - return autoexport.NewLogExporter( - ctx, - // Default to a stdout log exporter if autoexport fails to configure one. - autoexport.WithFallbackLogExporter( - func(ctx context.Context) (log.Exporter, error) { - return stdoutlog.New() - }, - ), - ) + options := []otlploghttp.Option{ + otlploghttp.WithEndpointURL(endpointURL), + otlploghttp.WithHeaders(parseLauncherOTLPHeaders()), + } + + timeout, timeoutSet, err := parseLauncherOTLPTimeout() + if err != nil { + return nil, err + } + if httpClient := launcherOTLPHTTPClient(timeout, timeoutSet); httpClient != nil { + options = append(options, otlploghttp.WithHTTPClient(httpClient)) + } + if timeoutSet { + options = append(options, otlploghttp.WithTimeout(timeout)) + } + if compression, ok, err := parseLauncherOTLPCompression(); err != nil { + return nil, err + } else if ok { + options = append(options, otlploghttp.WithCompression(compression)) + } + + return otlploghttp.New(ctx, options...) } type authHeaderSource interface { @@ -214,3 +263,77 @@ func (t *authHeaderTransport) baseTransport() http.RoundTripper { func otlpHTTPLogExporter(ctx context.Context, opts ...otlploghttp.Option) (log.Exporter, error) { return otlploghttp.New(ctx, opts...) } + +func launcherOTLPHTTPClient(timeout time.Duration, timeoutSet bool) *http.Client { + endpoint, err := nmpclient.ResolvePlatformEndpointFromEnv() + if err != nil || endpoint.Transport != nmpclient.TransportUDS { + return nil + } + httpClient := endpoint.HTTPClient() + if timeoutSet { + httpClient.Timeout = timeout + } + return httpClient +} + +func parseLauncherOTLPHeaders() map[string]string { + raw := strings.TrimSpace(os.Getenv(launcherOTLPLogsHeadersEnv)) + if raw == "" { + return nil + } + headers := map[string]string{} + for _, item := range strings.Split(raw, ",") { + key, value, ok := strings.Cut(strings.TrimSpace(item), "=") + if !ok || key == "" { + continue + } + if decoded, err := url.PathUnescape(value); err == nil { + value = decoded + } + headers[key] = value + } + return headers +} + +func parseLauncherOTLPTimeout() (time.Duration, bool, error) { + raw := strings.TrimSpace(os.Getenv(launcherOTLPLogsTimeoutEnv)) + if raw == "" { + return 0, false, nil + } + duration, err := time.ParseDuration(raw) + if err == nil { + return duration, true, nil + } + milliseconds, intErr := strconv.Atoi(raw) + if intErr == nil { + return time.Duration(milliseconds) * time.Millisecond, true, nil + } + return 0, false, fmt.Errorf("invalid %s value %q: %w", launcherOTLPLogsTimeoutEnv, raw, err) +} + +func parseLauncherOTLPCompression() (otlploghttp.Compression, bool, error) { + switch strings.ToLower(strings.TrimSpace(os.Getenv(launcherOTLPLogsCompressEnv))) { + case "": + return otlploghttp.NoCompression, false, nil + case "none": + return otlploghttp.NoCompression, true, nil + case "gzip": + return otlploghttp.GzipCompression, true, nil + default: + return otlploghttp.NoCompression, false, fmt.Errorf("unsupported %s value %q", launcherOTLPLogsCompressEnv, os.Getenv(launcherOTLPLogsCompressEnv)) + } +} + +type noopLogExporter struct{} + +func (noopLogExporter) Export(context.Context, []log.Record) error { + return nil +} + +func (noopLogExporter) Shutdown(context.Context) error { + return nil +} + +func (noopLogExporter) ForceFlush(context.Context) error { + return nil +} diff --git a/services/core/jobs/jobs-launcher/cmd/otel_test.go b/services/core/jobs/jobs-launcher/cmd/otel_test.go new file mode 100644 index 0000000000..8cd1b5b9c8 --- /dev/null +++ b/services/core/jobs/jobs-launcher/cmd/otel_test.go @@ -0,0 +1,41 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package cmd + +import ( + "testing" + "time" +) + +func TestLauncherOTLPHTTPClientSetsUDSTimeout(t *testing.T) { + t.Setenv("NMP_BASE_URL", "unix:///tmp/nemo-platform.sock") + + client := launcherOTLPHTTPClient(250*time.Millisecond, true) + if client == nil { + t.Fatal("expected UDS HTTP client") + } + if client.Timeout != 250*time.Millisecond { + t.Fatalf("expected UDS HTTP client timeout 250ms, got %s", client.Timeout) + } +} + +func TestLauncherOTLPHTTPClientLeavesUDSTimeoutUnset(t *testing.T) { + t.Setenv("NMP_BASE_URL", "unix:///tmp/nemo-platform.sock") + + client := launcherOTLPHTTPClient(250*time.Millisecond, false) + if client == nil { + t.Fatal("expected UDS HTTP client") + } + if client.Timeout != 0 { + t.Fatalf("expected UDS HTTP client timeout to remain unset, got %s", client.Timeout) + } +} + +func TestLauncherOTLPHTTPClientSkipsTCP(t *testing.T) { + t.Setenv("NMP_BASE_URL", "http://127.0.0.1:8080") + + if client := launcherOTLPHTTPClient(250*time.Millisecond, true); client != nil { + t.Fatal("expected no custom HTTP client for TCP endpoint") + } +} diff --git a/services/core/jobs/jobs-launcher/cmd/run.go b/services/core/jobs/jobs-launcher/cmd/run.go index f5ac028767..eef947cec3 100644 --- a/services/core/jobs/jobs-launcher/cmd/run.go +++ b/services/core/jobs/jobs-launcher/cmd/run.go @@ -10,17 +10,22 @@ import ( "fmt" "io" "log/slog" + "net/http" + "net/url" "os" "os/exec" "os/signal" "strings" "sync" "syscall" + "time" "github.com/NVIDIA-NeMo/nemo-platform/services/core/jobs/jobs-launcher/nmpclient" "github.com/spf13/cobra" ) +const secretFetchTimeout = 30 * time.Second + var runCmd = &cobra.Command{ Use: "run [args...]", Short: "Run a subprocess and tail its logs", @@ -108,11 +113,27 @@ func parseSecretReferences(secretsEnv string) ([]secretReference, error) { // fetchSecrets retrieves secrets using the NeMo Platform API client and returns them as environment variables func fetchSecrets(apiBaseURL string, principal *nmpclient.Principal, secretRefs []secretReference) ([]string, error) { + return fetchSecretsWithClient(nmpclient.NewSecretClient(apiBaseURL, principal), secretRefs) +} + +func fetchSecretsWithEndpoint(endpoint nmpclient.Endpoint, principal *nmpclient.Principal, secretRefs []secretReference) ([]string, error) { + return fetchSecretsWithClient( + nmpclient.NewSecretClientWithHTTPClient(endpoint.ConnectBaseURL, principal, secretEndpointHTTPClient(endpoint)), + secretRefs, + ) +} + +func secretEndpointHTTPClient(endpoint nmpclient.Endpoint) *http.Client { + httpClient := *endpoint.HTTPClient() + httpClient.Timeout = secretFetchTimeout + return &httpClient +} + +func fetchSecretsWithClient(client nmpclient.SecretClient, secretRefs []secretReference) ([]string, error) { if len(secretRefs) == 0 { return nil, nil } - client := nmpclient.NewSecretClient(apiBaseURL, principal) envVars := make([]string, 0, len(secretRefs)) for _, ref := range secretRefs { @@ -131,6 +152,19 @@ func fetchSecrets(apiBaseURL string, principal *nmpclient.Principal, secretRefs return envVars, nil } +func workloadEnvFromParent() []string { + env := os.Environ() + filtered := make([]string, 0, len(env)) + for _, item := range env { + key, _, _ := strings.Cut(item, "=") + if strings.HasPrefix(key, "NMP_JOB_LAUNCHER_") { + continue + } + filtered = append(filtered, item) + } + return filtered +} + // runExecWithStdin sets up OTEL and runs the specified command with stdin func runExecWithStdin(args []string) (exitCode int, err error) { ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) @@ -148,6 +182,30 @@ func runExecWithStdin(args []string) (exitCode int, err error) { return runExec(args, os.Stdin) } +func configureOTELHeadersFromWorkloadToken() { + token := os.Getenv("NEMO_WORKLOAD_TOKEN") + if token == "" { + return + } + + const headersEnv = launcherOTLPLogsHeadersEnv + headers := os.Getenv(headersEnv) + for _, item := range strings.Split(headers, ",") { + key, _, _ := strings.Cut(strings.TrimSpace(item), "=") + if strings.EqualFold(key, "authorization") { + return + } + } + + authHeader := "Authorization=" + url.PathEscape("Bearer "+token) + if headers == "" { + os.Setenv(headersEnv, authHeader) + return + } + os.Setenv(headersEnv, headers+","+authHeader) +} + + // runExec runs the specified command with arguments, injecting secrets as environment variables if specified func runExec(args []string, stdinReader io.Reader) (int, error) { // Command and arguments @@ -160,8 +218,8 @@ func runExec(args []string, stdinReader io.Reader) (int, error) { // Prepare the subprocess cmd := exec.Command(cmdName, cmdArgs...) - // Inherit parent environment - cmd.Env = os.Environ() + // Inherit parent environment, excluding launcher-private control variables. + cmd.Env = workloadEnvFromParent() // Parse and fetch secrets if NEMO_JOB_SECRETS is set secretsEnv := os.Getenv("NEMO_JOB_SECRETS") @@ -173,18 +231,16 @@ func runExec(args []string, stdinReader io.Reader) (int, error) { } if len(secretRefs) > 0 { - // Get API configuration from environment - apiBaseURL := os.Getenv("NMP_SECRETS_URL") - - if apiBaseURL == "" { - logger.Printf("Error: NMP_SECRETS_URL environment variable is required when NEMO_JOB_SECRETS is set\n") - return 1, fmt.Errorf("NMP_SECRETS_URL is not set") + secretEndpoint, err := nmpclient.ResolveServiceEndpointFromEnv("secrets") + if err != nil { + logger.Printf("Error: NMP_SECRETS_URL or NMP_BASE_URL is required when NEMO_JOB_SECRETS is set: %v\n", err) + return 1, fmt.Errorf("secrets endpoint is not configured: %w", err) } // Build auth context from NMP_PRINCIPAL JSON env var set by the jobs controller principal := nmpclient.PrincipalFromEnv() - secretEnvVars, err := fetchSecrets(apiBaseURL, principal, secretRefs) + secretEnvVars, err := fetchSecretsWithEndpoint(secretEndpoint, principal, secretRefs) if err != nil { logger.Printf("Error fetching secrets: %v\n", err) return 1, err diff --git a/services/core/jobs/jobs-launcher/cmd/run_test.go b/services/core/jobs/jobs-launcher/cmd/run_test.go index f1e696a8aa..4e6fa4a420 100644 --- a/services/core/jobs/jobs-launcher/cmd/run_test.go +++ b/services/core/jobs/jobs-launcher/cmd/run_test.go @@ -153,6 +153,7 @@ func TestRunExecWithSecrets(t *testing.T) { name string secretsEnv string apiURL string + baseURL string principalJSON string expectedExitCode int expectError bool @@ -173,6 +174,14 @@ func TestRunExecWithSecrets(t *testing.T) { expectedExitCode: 0, expectError: false, }, + { + name: "uses_base_url_fallback", + secretsEnv: "TEST_SECRET=default/test-secret", + baseURL: mockServer.URL, + principalJSON: `{"id":"test-principal"}`, + expectedExitCode: 0, + expectError: false, + }, { name: "missing_api_url", secretsEnv: "TEST_SECRET=default/test-secret", @@ -197,6 +206,7 @@ func TestRunExecWithSecrets(t *testing.T) { origEnvVars := map[string]envVarState{ "NEMO_JOB_SECRETS": getEnvState("NEMO_JOB_SECRETS"), "NMP_SECRETS_URL": getEnvState("NMP_SECRETS_URL"), + "NMP_BASE_URL": getEnvState("NMP_BASE_URL"), "NMP_PRINCIPAL": getEnvState("NMP_PRINCIPAL"), } defer restoreEnvVars(origEnvVars) @@ -210,6 +220,11 @@ func TestRunExecWithSecrets(t *testing.T) { } else { os.Unsetenv("NMP_SECRETS_URL") } + if tc.baseURL != "" { + os.Setenv("NMP_BASE_URL", tc.baseURL) + } else { + os.Unsetenv("NMP_BASE_URL") + } if tc.principalJSON != "" { os.Setenv("NMP_PRINCIPAL", tc.principalJSON) } else { @@ -258,6 +273,7 @@ func TestRunExecWithSecretsNotFound(t *testing.T) { origEnvVars := map[string]envVarState{ "NEMO_JOB_SECRETS": getEnvState("NEMO_JOB_SECRETS"), "NMP_SECRETS_URL": getEnvState("NMP_SECRETS_URL"), + "NMP_BASE_URL": getEnvState("NMP_BASE_URL"), "NMP_PRINCIPAL": getEnvState("NMP_PRINCIPAL"), } defer restoreEnvVars(origEnvVars) @@ -265,6 +281,7 @@ func TestRunExecWithSecretsNotFound(t *testing.T) { // Set test environment variables os.Setenv("NEMO_JOB_SECRETS", "NONEXISTENT_SECRET=default/nonexistent") os.Setenv("NMP_SECRETS_URL", mockServer.URL) + os.Unsetenv("NMP_BASE_URL") os.Setenv("NMP_PRINCIPAL", `{"id":"test-principal"}`) exitCode, err := runExec([]string{"echo", "test"}, nil) @@ -282,6 +299,37 @@ func TestRunExecWithSecretsNotFound(t *testing.T) { } } +func TestSecretEndpointHTTPClientSetsBoundedTimeout(t *testing.T) { + originalDefaultTimeout := http.DefaultClient.Timeout + + tcpEndpoint, err := nmpclient.ParseEndpoint("http://127.0.0.1:8080") + if err != nil { + t.Fatalf("ParseEndpoint returned error: %v", err) + } + tcpClient := secretEndpointHTTPClient(tcpEndpoint) + if tcpClient == http.DefaultClient { + t.Fatal("expected bounded TCP client to avoid mutating http.DefaultClient") + } + if tcpClient.Timeout != secretFetchTimeout { + t.Fatalf("expected TCP client timeout %s, got %s", secretFetchTimeout, tcpClient.Timeout) + } + if http.DefaultClient.Timeout != originalDefaultTimeout { + t.Fatalf("expected http.DefaultClient timeout to remain %s, got %s", originalDefaultTimeout, http.DefaultClient.Timeout) + } + + udsEndpoint, err := nmpclient.ParseEndpoint("unix:///tmp/nemo-platform.sock") + if err != nil { + t.Fatalf("ParseEndpoint returned error: %v", err) + } + udsClient := secretEndpointHTTPClient(udsEndpoint) + if udsClient.Timeout != secretFetchTimeout { + t.Fatalf("expected UDS client timeout %s, got %s", secretFetchTimeout, udsClient.Timeout) + } + if udsClient.Transport == nil { + t.Fatal("expected UDS client to preserve custom transport") + } +} + func TestRunExecWithoutSecrets(t *testing.T) { // Ensure no secrets environment variables are set origSecretsEnv, wasSet := os.LookupEnv("NEMO_JOB_SECRETS") @@ -306,6 +354,90 @@ func TestRunExecWithoutSecrets(t *testing.T) { } } +func TestConfigureOTELHeadersFromWorkloadToken(t *testing.T) { + testCases := []struct { + name string + token string + existingHeaders string + expectedHeaders string + }{ + { + name: "adds_authorization_header", + token: "token.with-symbols_123", + expectedHeaders: "Authorization=Bearer%20token.with-symbols_123", + }, + { + name: "preserves_existing_headers", + token: "abc.def", + existingHeaders: "X-NMP-Principal-Id=nemo-user", + expectedHeaders: "X-NMP-Principal-Id=nemo-user,Authorization=Bearer%20abc.def", + }, + { + name: "keeps_existing_authorization_header", + token: "abc.def", + existingHeaders: "authorization=Bearer+explicit", + expectedHeaders: "authorization=Bearer+explicit", + }, + { + name: "does_nothing_without_token", + existingHeaders: "X-Test=value", + expectedHeaders: "X-Test=value", + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + origEnvVars := map[string]envVarState{ + "NEMO_WORKLOAD_TOKEN": getEnvState("NEMO_WORKLOAD_TOKEN"), + launcherOTLPLogsHeadersEnv: getEnvState(launcherOTLPLogsHeadersEnv), + "OTEL_EXPORTER_OTLP_LOGS_HEADERS": getEnvState("OTEL_EXPORTER_OTLP_LOGS_HEADERS"), + } + defer restoreEnvVars(origEnvVars) + + if tc.token != "" { + os.Setenv("NEMO_WORKLOAD_TOKEN", tc.token) + } else { + os.Unsetenv("NEMO_WORKLOAD_TOKEN") + } + if tc.existingHeaders != "" { + os.Setenv(launcherOTLPLogsHeadersEnv, tc.existingHeaders) + } else { + os.Unsetenv(launcherOTLPLogsHeadersEnv) + } + os.Setenv("OTEL_EXPORTER_OTLP_LOGS_HEADERS", "user-owned=value") + + configureOTELHeadersFromWorkloadToken() + + got := os.Getenv(launcherOTLPLogsHeadersEnv) + if got != tc.expectedHeaders { + t.Errorf("Expected launcher OTLP headers %q, got %q", tc.expectedHeaders, got) + } + if got := os.Getenv("OTEL_EXPORTER_OTLP_LOGS_HEADERS"); got != "user-owned=value" { + t.Errorf("Expected user OTEL headers to be preserved, got %q", got) + } + }) + } +} + +func TestWorkloadEnvFromParentFiltersLauncherPrivateVars(t *testing.T) { + origEnvVars := map[string]envVarState{ + "NMP_JOB_LAUNCHER_LOGS_EXPORTER": getEnvState("NMP_JOB_LAUNCHER_LOGS_EXPORTER"), + "OTEL_LOGS_EXPORTER": getEnvState("OTEL_LOGS_EXPORTER"), + } + defer restoreEnvVars(origEnvVars) + + os.Setenv("NMP_JOB_LAUNCHER_LOGS_EXPORTER", "otlp") + os.Setenv("OTEL_LOGS_EXPORTER", "otlp") + + env := strings.Join(workloadEnvFromParent(), "\n") + if strings.Contains(env, "NMP_JOB_LAUNCHER_LOGS_EXPORTER=") { + t.Fatal("expected launcher-private env var to be filtered") + } + if !strings.Contains(env, "OTEL_LOGS_EXPORTER=otlp") { + t.Fatal("expected user OTEL env var to be preserved") + } +} + func TestParseSecretReferences(t *testing.T) { testCases := []struct { name string diff --git a/services/core/jobs/jobs-launcher/cmd/workload_auth_test.go b/services/core/jobs/jobs-launcher/cmd/workload_auth_test.go index cb4a72d74f..d05a8623fc 100644 --- a/services/core/jobs/jobs-launcher/cmd/workload_auth_test.go +++ b/services/core/jobs/jobs-launcher/cmd/workload_auth_test.go @@ -149,7 +149,8 @@ func TestNewLogExporterCachesWorkloadAuthAcrossExports(t *testing.T) { t.Setenv(nmpBaseURLEnv, server.URL) t.Setenv(workloadIdentityTokenFileEnv, subjectTokenPath) - t.Setenv(nmpJobLogsEndpointEnv, server.URL+"/v1/logs") + t.Setenv(launcherLogsExporterEnv, "otlp") + t.Setenv(launcherOTLPLogsEndpointEnv, server.URL+"/v1/logs") exporter, err := newLogExporter(context.Background()) if err != nil { @@ -352,7 +353,8 @@ func TestNewLogExporterPropagatesInitialWorkloadAuthFailure(t *testing.T) { t.Setenv(nmpBaseURLEnv, server.URL) t.Setenv(workloadIdentityTokenFileEnv, subjectTokenPath) - t.Setenv(nmpJobLogsEndpointEnv, server.URL+"/v1/logs") + t.Setenv(launcherLogsExporterEnv, "otlp") + t.Setenv(launcherOTLPLogsEndpointEnv, server.URL+"/v1/logs") _, err := newLogExporter(context.Background()) if err == nil { diff --git a/services/core/jobs/jobs-launcher/go.mod b/services/core/jobs/jobs-launcher/go.mod index 5f1dd590bb..eaaaf83b03 100644 --- a/services/core/jobs/jobs-launcher/go.mod +++ b/services/core/jobs/jobs-launcher/go.mod @@ -15,7 +15,6 @@ require ( ) require ( - github.com/beorn7/perks v1.0.1 // indirect github.com/cenkalti/backoff/v5 v5.0.3 // indirect github.com/cespare/xxhash/v2 v2.3.0 // indirect github.com/go-logr/logr v1.4.3 // indirect diff --git a/services/core/jobs/jobs-launcher/go.sum b/services/core/jobs/jobs-launcher/go.sum index f9b7cb2248..46143c776c 100644 --- a/services/core/jobs/jobs-launcher/go.sum +++ b/services/core/jobs/jobs-launcher/go.sum @@ -1,5 +1,3 @@ -github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM= -github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw= github.com/cenkalti/backoff/v5 v5.0.3 h1:ZN+IMa753KfX5hd8vVaMixjnqRZ3y8CuJKRKj1xcsSM= github.com/cenkalti/backoff/v5 v5.0.3/go.mod h1:rkhZdG3JZukswDf7f0cwqPNk4K0sa+F97BxZthm/crw= github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= @@ -22,12 +20,6 @@ github.com/grpc-ecosystem/grpc-gateway/v2 v2.29.0 h1:5VipnvEpbqr2gA2VbM+nYVbkIF2 github.com/grpc-ecosystem/grpc-gateway/v2 v2.29.0/go.mod h1:Hyl3n6Twe1hvtd9XUXDec4pTvgMSEixRuQKPTMH2bNs= github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8= github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw= -github.com/klauspost/compress v1.18.0 h1:c/Cqfb0r+Yi+JtIEq73FWXVkRonBlf0CRNYc8Zttxdo= -github.com/klauspost/compress v1.18.0/go.mod h1:2Pp+KzxcywXVXMr50+X0Q/Lsb43OQHYWRCY2AiWywWQ= -github.com/kylelemons/godebug v1.1.0 h1:RPNrshWIDI6G2gRW9EHilWtl7Z6Sb1BR0xunSBf0SNc= -github.com/kylelemons/godebug v1.1.0/go.mod h1:9/0rRGxNHcop5bhtWyNeEfOS8JIWk580+fNqagV/RAw= -github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 h1:C3w9PqII01/Oq1c1nUAm88MOHcQC9l5mIlSMApZMrHA= -github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822/go.mod h1:+n7T8mK8HuQTcFwEeznm/DIxMOiR9yIdICNftLE1DvQ= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/prometheus/client_golang v1.23.2 h1:Je96obch5RDVy3FDMndoUsjAhG5Edi49h0RJWRi/o0o= diff --git a/services/core/jobs/jobs-launcher/nmpclient/client.go b/services/core/jobs/jobs-launcher/nmpclient/client.go index 041457213f..b95aa247d3 100644 --- a/services/core/jobs/jobs-launcher/nmpclient/client.go +++ b/services/core/jobs/jobs-launcher/nmpclient/client.go @@ -67,13 +67,24 @@ type secretClient struct { } func NewSecretClient(apiBaseURL string, principal *Principal) SecretClient { + return NewSecretClientWithHTTPClient(apiBaseURL, principal, http.DefaultClient) +} + +func NewSecretClientWithHTTPClient(apiBaseURL string, principal *Principal, httpClient *http.Client) SecretClient { + if httpClient == nil { + httpClient = http.DefaultClient + } return &secretClient{ - httpClient: http.DefaultClient, + httpClient: httpClient, principal: principal, apiBaseURL: apiBaseURL, } } +func NewSecretClientForEndpoint(endpoint Endpoint, principal *Principal) SecretClient { + return NewSecretClientWithHTTPClient(endpoint.ConnectBaseURL, principal, endpoint.HTTPClient()) +} + func (c *secretClient) GetSecret(workspaceID, secretName string) (*PlatformSecretAccessResponse, error) { secretURL := getSecretURL(c.apiBaseURL, workspaceID, secretName) diff --git a/services/core/jobs/jobs-launcher/nmpclient/client_test.go b/services/core/jobs/jobs-launcher/nmpclient/client_test.go index bcdca23379..4a56517202 100644 --- a/services/core/jobs/jobs-launcher/nmpclient/client_test.go +++ b/services/core/jobs/jobs-launcher/nmpclient/client_test.go @@ -5,6 +5,7 @@ package nmpclient import ( "fmt" + "net" "net/http" "net/http/httptest" "os" @@ -162,6 +163,55 @@ func TestSecretClient_GetSecret(t *testing.T) { } } +func TestSecretClient_GetSecretOverUDS(t *testing.T) { + socketFile, err := os.CreateTemp("", "nmp-*.sock") + if err != nil { + t.Fatalf("failed to create temp socket path: %v", err) + } + socketPath := socketFile.Name() + socketFile.Close() + os.Remove(socketPath) + listener, err := net.Listen("unix", socketPath) + if err != nil { + t.Fatalf("failed to listen on unix socket: %v", err) + } + + server := &http.Server{ + Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/apis/secrets/v2/workspaces/default/secrets/api-key/access" { + t.Errorf("unexpected path: %s", r.URL.Path) + http.Error(w, "unexpected path", http.StatusNotFound) + return + } + w.WriteHeader(http.StatusOK) + fmt.Fprintln(w, `{"value":"secret-over-uds"}`) + }), + } + defer server.Close() + t.Cleanup(func() { + if err := os.Remove(socketPath); err != nil && !os.IsNotExist(err) { + t.Errorf("failed to remove unix socket path: %v", err) + } + }) + go func() { + _ = server.Serve(listener) + }() + + endpoint, err := ParseEndpoint("unix://" + socketPath) + if err != nil { + t.Fatalf("ParseEndpoint returned error: %v", err) + } + client := NewSecretClientForEndpoint(endpoint, &Principal{ID: "test-principal"}) + + secret, err := client.GetSecret("default", "api-key") + if err != nil { + t.Fatalf("expected no error, got %v", err) + } + if secret.Value != "secret-over-uds" { + t.Fatalf("unexpected secret value: %s", secret.Value) + } +} + func TestSecretClient_AuthHeaders(t *testing.T) { testCases := []struct { name string diff --git a/services/core/jobs/jobs-launcher/nmpclient/endpoint.go b/services/core/jobs/jobs-launcher/nmpclient/endpoint.go new file mode 100644 index 0000000000..67beea265b --- /dev/null +++ b/services/core/jobs/jobs-launcher/nmpclient/endpoint.go @@ -0,0 +1,88 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package nmpclient + +import ( + "context" + "fmt" + "net" + "net/http" + "net/url" + "os" + "strings" +) + +const UDSBaseURL = "http://nemo-platform.local" + +type Transport string + +const ( + TransportTCP Transport = "tcp" + TransportUDS Transport = "uds" +) + +type Endpoint struct { + ConnectBaseURL string + SocketPath string + Transport Transport +} + +func ParseEndpoint(raw string) (Endpoint, error) { + if raw == "" { + return Endpoint{}, fmt.Errorf("platform endpoint URL is not configured") + } + if strings.HasPrefix(raw, "http://") || strings.HasPrefix(raw, "https://") { + parsed, err := url.Parse(raw) + if err != nil || parsed.Host == "" { + return Endpoint{}, fmt.Errorf("invalid platform endpoint URL %q", raw) + } + return Endpoint{ + ConnectBaseURL: strings.TrimRight(raw, "/"), + Transport: TransportTCP, + }, nil + } + if strings.HasPrefix(raw, "unix://") { + socketPath := strings.TrimPrefix(raw, "unix://") + if !strings.HasPrefix(socketPath, "/") { + return Endpoint{}, fmt.Errorf("UDS endpoint must use an absolute socket path, got %q", raw) + } + return Endpoint{ + ConnectBaseURL: UDSBaseURL, + SocketPath: socketPath, + Transport: TransportUDS, + }, nil + } + if strings.HasPrefix(raw, "/") { + return Endpoint{}, fmt.Errorf("raw socket paths are not valid endpoint URLs; use unix://%s", raw) + } + return Endpoint{}, fmt.Errorf("unsupported platform endpoint URL %q; expected http://, https://, or unix://", raw) +} + +func ResolvePlatformEndpointFromEnv() (Endpoint, error) { + return ParseEndpoint(os.Getenv("NMP_BASE_URL")) +} + +func ResolveServiceEndpointFromEnv(service string) (Endpoint, error) { + if serviceEnv := os.Getenv(serviceURLEnvName(service)); serviceEnv != "" { + return ParseEndpoint(serviceEnv) + } + return ResolvePlatformEndpointFromEnv() +} + +func serviceURLEnvName(service string) string { + normalized := strings.ToUpper(strings.ReplaceAll(service, "-", "_")) + return "NMP_" + normalized + "_URL" +} + +func (e Endpoint) HTTPClient() *http.Client { + if e.Transport != TransportUDS { + return http.DefaultClient + } + transport := &http.Transport{ + DialContext: func(ctx context.Context, _, _ string) (net.Conn, error) { + return (&net.Dialer{}).DialContext(ctx, "unix", e.SocketPath) + }, + } + return &http.Client{Transport: transport} +} diff --git a/services/core/jobs/jobs-launcher/nmpclient/endpoint_test.go b/services/core/jobs/jobs-launcher/nmpclient/endpoint_test.go new file mode 100644 index 0000000000..c892ae61ab --- /dev/null +++ b/services/core/jobs/jobs-launcher/nmpclient/endpoint_test.go @@ -0,0 +1,122 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package nmpclient + +import ( + "fmt" + "net" + "net/http" + "os" + "testing" +) + +func TestParseEndpointTCP(t *testing.T) { + endpoint, err := ParseEndpoint("http://127.0.0.1:8080/") + if err != nil { + t.Fatalf("ParseEndpoint returned error: %v", err) + } + if endpoint.ConnectBaseURL != "http://127.0.0.1:8080" { + t.Fatalf("unexpected ConnectBaseURL: %s", endpoint.ConnectBaseURL) + } + if endpoint.Transport != TransportTCP { + t.Fatalf("unexpected transport: %s", endpoint.Transport) + } +} + +func TestParseEndpointUDS(t *testing.T) { + endpoint, err := ParseEndpoint("unix:///tmp/nemo-platform.sock") + if err != nil { + t.Fatalf("ParseEndpoint returned error: %v", err) + } + if endpoint.ConnectBaseURL != UDSBaseURL { + t.Fatalf("unexpected ConnectBaseURL: %s", endpoint.ConnectBaseURL) + } + if endpoint.SocketPath != "/tmp/nemo-platform.sock" { + t.Fatalf("unexpected socket path: %s", endpoint.SocketPath) + } + if endpoint.Transport != TransportUDS { + t.Fatalf("unexpected transport: %s", endpoint.Transport) + } +} + +func TestParseEndpointRejectsRawSocketPath(t *testing.T) { + if _, err := ParseEndpoint("/tmp/nemo-platform.sock"); err == nil { + t.Fatal("expected raw socket path to be rejected") + } +} + +func TestResolveServiceEndpointPrefersServiceURL(t *testing.T) { + t.Setenv("NMP_BASE_URL", "http://platform:8080") + t.Setenv("NMP_SECRETS_URL", "unix:///tmp/secrets.sock") + + endpoint, err := ResolveServiceEndpointFromEnv("secrets") + if err != nil { + t.Fatalf("ResolveServiceEndpointFromEnv returned error: %v", err) + } + if endpoint.Transport != TransportUDS { + t.Fatalf("expected UDS endpoint, got %s", endpoint.Transport) + } + if endpoint.SocketPath != "/tmp/secrets.sock" { + t.Fatalf("unexpected socket path: %s", endpoint.SocketPath) + } +} + +func TestEndpointContractDoesNotReadEndpointEnvFamily(t *testing.T) { + t.Setenv("NMP_PLATFORM_ENDPOINT", "unix:///tmp/platform.sock") + t.Setenv("NMP_SECRETS_ENDPOINT", "unix:///tmp/secrets.sock") + + if _, err := ResolvePlatformEndpointFromEnv(); err == nil { + t.Fatal("expected missing NMP_BASE_URL to fail") + } + if _, err := ResolveServiceEndpointFromEnv("secrets"); err == nil { + t.Fatal("expected missing NMP_SECRETS_URL and NMP_BASE_URL to fail") + } +} + +func TestUDSHTTPClient(t *testing.T) { + socketFile, err := os.CreateTemp("", "nmp-*.sock") + if err != nil { + t.Fatalf("failed to create temp socket path: %v", err) + } + socketPath := socketFile.Name() + socketFile.Close() + os.Remove(socketPath) + listener, err := net.Listen("unix", socketPath) + if err != nil { + t.Fatalf("failed to listen on unix socket: %v", err) + } + + receivedPath := make(chan string, 1) + server := &http.Server{ + Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + receivedPath <- r.URL.Path + if r.URL.Path != "/status" { + http.Error(w, "unexpected path", http.StatusNotFound) + return + } + fmt.Fprintln(w, `{"ok":true}`) + }), + } + defer server.Close() + defer os.Remove(socketPath) + go func() { + _ = server.Serve(listener) + }() + + endpoint, err := ParseEndpoint("unix://" + socketPath) + if err != nil { + t.Fatalf("ParseEndpoint returned error: %v", err) + } + resp, err := endpoint.HTTPClient().Get(endpoint.ConnectBaseURL + "/status") + if err != nil { + t.Fatalf("UDS request failed: %v", err) + } + defer resp.Body.Close() + if path := <-receivedPath; path != "/status" { + t.Fatalf("unexpected path: %s", path) + } + if resp.StatusCode != http.StatusOK { + t.Fatalf("unexpected status: %d", resp.StatusCode) + } +} diff --git a/services/core/jobs/src/nmp/core/jobs/controllers/backends/base.py b/services/core/jobs/src/nmp/core/jobs/controllers/backends/base.py index 64432ba840..e689ef64f5 100644 --- a/services/core/jobs/src/nmp/core/jobs/controllers/backends/base.py +++ b/services/core/jobs/src/nmp/core/jobs/controllers/backends/base.py @@ -7,8 +7,9 @@ import logging from abc import ABC, abstractmethod from collections.abc import Iterable +from dataclasses import dataclass from enum import Enum -from typing import Generic, Optional, TypeVar +from typing import Generic, Literal, Optional, TypeVar from urllib.parse import SplitResult, urlsplit from nemo_platform import NeMoPlatform @@ -40,6 +41,7 @@ PERSISTENT_JOB_STORAGE_PATH_ENVVAR, TASK_CONFIG_ENVVAR, ) +from nmp.common.platform_endpoint import parse_platform_endpoint from nmp.common.sdk_factory import get_entity_parts from nmp.core.jobs.app.providers import ComputeResources from pydantic import BaseModel, model_validator @@ -155,6 +157,31 @@ def get_workload_identity_token_audience() -> str: return "nemo-platform" +NMP_JOB_LAUNCHER_LOGS_EXPORTER_ENVVAR = "NMP_JOB_LAUNCHER_LOGS_EXPORTER" +NMP_JOB_LAUNCHER_OTLP_LOGS_ENDPOINT_ENVVAR = "NMP_JOB_LAUNCHER_OTLP_LOGS_ENDPOINT" +NMP_JOB_LAUNCHER_OTLP_LOGS_HEADERS_ENVVAR = "NMP_JOB_LAUNCHER_OTLP_LOGS_HEADERS" +NMP_JOB_LAUNCHER_OTLP_LOGS_PROTOCOL_ENVVAR = "NMP_JOB_LAUNCHER_OTLP_LOGS_PROTOCOL" +NMP_JOB_LAUNCHER_OTLP_LOGS_SOCKET_PATH_ENVVAR = "NMP_JOB_LAUNCHER_OTLP_LOGS_SOCKET_PATH" +NMP_JOB_LAUNCHER_OTLP_LOGS_TRANSPORT_ENVVAR = "NMP_JOB_LAUNCHER_OTLP_LOGS_TRANSPORT" +NMP_JOB_LAUNCHER_OTLP_LOGS_PROTOCOL = "http/protobuf" + + +@dataclass(frozen=True) +class OtlpLogsEndpointConfig: + endpoint: str + transport: Literal["tcp", "uds"] + socket_path: str | None = None + + def to_env(self) -> dict[str, str]: + env = { + NMP_JOB_LAUNCHER_OTLP_LOGS_ENDPOINT_ENVVAR: self.endpoint, + NMP_JOB_LAUNCHER_OTLP_LOGS_TRANSPORT_ENVVAR: self.transport, + } + if self.socket_path is not None: + env[NMP_JOB_LAUNCHER_OTLP_LOGS_SOCKET_PATH_ENVVAR] = self.socket_path + return env + + class JobUpdate(BaseModel): status: PlatformJobStatus status_details: dict | None = None @@ -531,14 +558,41 @@ def get_logs_endpoint_from_fileset( Returns: Full OTLP logs endpoint URL with appropriate loopback address applied. """ - # Job telemetry is emitted from a separate process/container/pod. When Files - # runs in-process with the API server, local service URLs are not necessarily - # routable from job runtime networks, so fall back through the same - # workload-facing base URL used for job SDK env vars. - base_url = platform_config.service_discovery.get("files") or _job_runtime_base_url(platform_config) + return get_logs_endpoint_config_from_fileset( + platform_config, + workspace, + fileset_id, + loopback_address=loopback_address, + ).endpoint + + +def get_logs_endpoint_config_from_fileset( + platform_config: PlatformConfig, workspace: str, fileset_id: str, loopback_address: str | None = None +) -> OtlpLogsEndpointConfig: + """Get OTLP logs endpoint config, preserving transport metadata for local UDS runtimes. + + Job telemetry is emitted from a separate process/container/pod. When Files + runs in-process with the API server, local service URLs are not necessarily + routable from job runtime networks, so fall back through the same + workload-facing base URL used for job SDK env vars. + """ + # Check service_discovery for a files-specific URL first, then fall back to + # the job runtime base URL (service_discovery["platform"] / base_url). + files_discovery_url = platform_config.service_discovery.get("files") + if files_discovery_url: + platform_endpoint = parse_platform_endpoint(files_discovery_url) + else: + runtime_base = _job_runtime_base_url(platform_config) + platform_endpoint = parse_platform_endpoint(runtime_base) + + base_url = platform_endpoint.connect_base_url # Use configured loopback_address, or fall back to automatic detection effective_override = loopback_address or platform_config.loopback_address or determine_loopback_override() base_url = _replace_loopback_address(base_url, effective_override) - return f"{base_url}/apis/files/v2/workspaces/{workspace}/filesets/{fileset_id}/otlp/v1/logs" + return OtlpLogsEndpointConfig( + endpoint=f"{base_url}/apis/files/v2/workspaces/{workspace}/filesets/{fileset_id}/otlp/v1/logs", + transport=platform_endpoint.transport, + socket_path=str(platform_endpoint.socket_path) if platform_endpoint.socket_path is not None else None, + ) diff --git a/services/core/jobs/src/nmp/core/jobs/controllers/backends/docker.py b/services/core/jobs/src/nmp/core/jobs/controllers/backends/docker.py index b3eefe3640..c0a81ce1b1 100644 --- a/services/core/jobs/src/nmp/core/jobs/controllers/backends/docker.py +++ b/services/core/jobs/src/nmp/core/jobs/controllers/backends/docker.py @@ -89,7 +89,11 @@ GPUExecutionProvider, ) from nmp.core.jobs.controllers.backends.base import ( - JOB_LOGS_ENDPOINT_ENVVAR, + NMP_JOB_LAUNCHER_LOGS_EXPORTER_ENVVAR, + NMP_JOB_LAUNCHER_OTLP_LOGS_ENDPOINT_ENVVAR, + NMP_JOB_LAUNCHER_OTLP_LOGS_HEADERS_ENVVAR, + NMP_JOB_LAUNCHER_OTLP_LOGS_PROTOCOL, + NMP_JOB_LAUNCHER_OTLP_LOGS_PROTOCOL_ENVVAR, WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR, WORKLOAD_IDENTITY_TOKEN_FILE_PATH, WORKLOAD_IDENTITY_VOLUME_PATH, @@ -951,12 +955,14 @@ def schedule_single_container( EPHEMERAL_TASK_STORAGE_PATH_ENVVAR: DEFAULT_TASK_STORAGE_PATH, CONFIG_TASK_STORAGE_PATH_ENVVAR: DEFAULT_CONFIG_STORAGE_PATH, NEMO_JOB_STEP_CONFIG_FILE_PATH_ENVVAR: DEFAULT_NEMO_JOB_STEP_CONFIG_FILE_PATH, - # Endpoint used by jobs-launcher to upload task stdout/stderr logs. - JOB_LOGS_ENDPOINT_ENVVAR: get_logs_endpoint_from_fileset( + # Private env vars for jobs-launcher to export captured logs. + NMP_JOB_LAUNCHER_OTLP_LOGS_ENDPOINT_ENVVAR: get_logs_endpoint_from_fileset( platform_config, step.workspace, step.fileset, ), + NMP_JOB_LAUNCHER_LOGS_EXPORTER_ENVVAR: "otlp", + NMP_JOB_LAUNCHER_OTLP_LOGS_PROTOCOL_ENVVAR: NMP_JOB_LAUNCHER_OTLP_LOGS_PROTOCOL, # Inject secret environment variable mappings for the jobs-launcher to fetch NEMO_JOB_SECRETS_ENVVAR: self.get_secrets_environment_variable_for_injection(step), } @@ -973,6 +979,8 @@ def schedule_single_container( env_var_dict = principal.get_env_var() for name, value in env_var_dict.items(): env[name] = value + # Also set launcher OTLP headers for authenticated platform log export. + env[NMP_JOB_LAUNCHER_OTLP_LOGS_HEADERS_ENVVAR] = principal.get_otlp_headers_value() step_config_json = json.dumps(step.step_spec.config) diff --git a/services/core/jobs/src/nmp/core/jobs/controllers/backends/kubernetes/common.py b/services/core/jobs/src/nmp/core/jobs/controllers/backends/kubernetes/common.py index d6def96124..08a66bae89 100644 --- a/services/core/jobs/src/nmp/core/jobs/controllers/backends/kubernetes/common.py +++ b/services/core/jobs/src/nmp/core/jobs/controllers/backends/kubernetes/common.py @@ -78,7 +78,11 @@ ) from nmp.core.jobs.app.providers import ComputeResources, ContainerSpec from nmp.core.jobs.controllers.backends.base import ( - JOB_LOGS_ENDPOINT_ENVVAR, + NMP_JOB_LAUNCHER_LOGS_EXPORTER_ENVVAR, + NMP_JOB_LAUNCHER_OTLP_LOGS_ENDPOINT_ENVVAR, + NMP_JOB_LAUNCHER_OTLP_LOGS_HEADERS_ENVVAR, + NMP_JOB_LAUNCHER_OTLP_LOGS_PROTOCOL, + NMP_JOB_LAUNCHER_OTLP_LOGS_PROTOCOL_ENVVAR, WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR, WORKLOAD_IDENTITY_TOKEN_FILE_PATH, WORKLOAD_IDENTITY_VOLUME_NAME, @@ -1001,13 +1005,15 @@ def create_pod_template_spec( ), client.V1EnvVar(name=EPHEMERAL_TASK_STORAGE_PATH_ENVVAR, value=DEFAULT_TASK_STORAGE_PATH), client.V1EnvVar( - name=JOB_LOGS_ENDPOINT_ENVVAR, + name=NMP_JOB_LAUNCHER_OTLP_LOGS_ENDPOINT_ENVVAR, value=get_logs_endpoint_from_fileset( platform_config, step.workspace, step.fileset, ), ), + client.V1EnvVar(name=NMP_JOB_LAUNCHER_LOGS_EXPORTER_ENVVAR, value="otlp"), + client.V1EnvVar(name=NMP_JOB_LAUNCHER_OTLP_LOGS_PROTOCOL_ENVVAR, value=NMP_JOB_LAUNCHER_OTLP_LOGS_PROTOCOL), client.V1EnvVar(name=NEMO_JOB_SECRETS_ENVVAR, value=secret_env_var_str), ] ) @@ -1022,6 +1028,10 @@ def create_pod_template_spec( env_var_dict = principal.get_env_var() for name, value in env_var_dict.items(): env.append(client.V1EnvVar(name=name, value=value)) + # Also set launcher OTLP headers for authenticated platform log export. + env.append( + client.V1EnvVar(name=NMP_JOB_LAUNCHER_OTLP_LOGS_HEADERS_ENVVAR, value=principal.get_otlp_headers_value()) + ) # Thread through shared platform envvars to the job shared_envvars = get_job_runtime_shared_envvars(platform_config) diff --git a/services/core/jobs/src/nmp/core/jobs/controllers/backends/subprocess.py b/services/core/jobs/src/nmp/core/jobs/controllers/backends/subprocess.py index 6b71ad0d95..8ed810439d 100644 --- a/services/core/jobs/src/nmp/core/jobs/controllers/backends/subprocess.py +++ b/services/core/jobs/src/nmp/core/jobs/controllers/backends/subprocess.py @@ -38,10 +38,14 @@ from nmp.core.jobs.app.providers import SubprocessExecutionProvider from nmp.core.jobs.app.schemas import BaseExecutionProfile from nmp.core.jobs.controllers.backends.base import ( + NMP_JOB_LAUNCHER_LOGS_EXPORTER_ENVVAR, + NMP_JOB_LAUNCHER_OTLP_LOGS_HEADERS_ENVVAR, + NMP_JOB_LAUNCHER_OTLP_LOGS_PROTOCOL, + NMP_JOB_LAUNCHER_OTLP_LOGS_PROTOCOL_ENVVAR, JobBackend, JobExecutionProfileConfig, JobUpdate, - get_logs_endpoint_from_fileset, + get_logs_endpoint_config_from_fileset, ) from nmp.core.jobs.controllers.backends.subprocess_runtime import ( SubprocessOtelLogger, @@ -428,6 +432,9 @@ def _prepare_runtime(self, step: PlatformJobStepWithContext) -> tuple[dict[str, log_path.touch() platform_config = get_platform_config() + otlp_logs_endpoint = get_logs_endpoint_config_from_fileset( + platform_config, step.workspace, step.fileset, loopback_address="localhost" + ) env = {name: value for name, value in os.environ.items() if name in SUBPROCESS_INHERITED_ENV_ALLOWLIST} env.update(self._execution_profile_config.env) env.update( @@ -442,17 +449,12 @@ def _prepare_runtime(self, step: PlatformJobStepWithContext) -> tuple[dict[str, CONFIG_TASK_STORAGE_PATH_ENVVAR: str(config_dir), PERSISTENT_JOB_STORAGE_PATH_ENVVAR: str(persistent_dir), NEMO_JOB_STEP_CONFIG_FILE_PATH_ENVVAR: str(config_path), - "OTEL_EXPORTER_OTLP_LOGS_ENDPOINT": get_logs_endpoint_from_fileset( - platform_config, - step.workspace, - step.fileset, - loopback_address="localhost", - ), - "OTEL_LOGS_EXPORTER": "otlp", - "OTEL_SERVICE_NAME": "nmp-job-task", + NMP_JOB_LAUNCHER_LOGS_EXPORTER_ENVVAR: "otlp", + NMP_JOB_LAUNCHER_OTLP_LOGS_PROTOCOL_ENVVAR: NMP_JOB_LAUNCHER_OTLP_LOGS_PROTOCOL, NEMO_JOB_SECRETS_ENVVAR: self.get_secrets_environment_variable_for_injection(step), } ) + env.update(otlp_logs_endpoint.to_env()) if spec and spec.environment: for envvar in spec.environment: @@ -479,7 +481,7 @@ def _prepare_runtime(self, step: PlatformJobStepWithContext) -> tuple[dict[str, auth_context = AuthContext.model_validate(step.auth_context.model_dump(mode="python", exclude_none=True)) principal = auth_context.to_principal() env.update(principal.get_env_var()) - env["OTEL_EXPORTER_OTLP_LOGS_HEADERS"] = principal.get_otlp_headers_value() + env[NMP_JOB_LAUNCHER_OTLP_LOGS_HEADERS_ENVVAR] = principal.get_otlp_headers_value() inject_secret_env_vars(env) return env, task_id, work_dir, log_path, persistent_dir diff --git a/services/core/jobs/src/nmp/core/jobs/controllers/backends/subprocess_runtime.py b/services/core/jobs/src/nmp/core/jobs/controllers/backends/subprocess_runtime.py index 6d3c2ee931..5865d8b3e3 100644 --- a/services/core/jobs/src/nmp/core/jobs/controllers/backends/subprocess_runtime.py +++ b/services/core/jobs/src/nmp/core/jobs/controllers/backends/subprocess_runtime.py @@ -6,15 +6,18 @@ import json import logging import threading +from collections.abc import Mapping from dataclasses import dataclass from datetime import datetime, timezone from pathlib import Path from time import time_ns -from typing import IO +from typing import IO, Any from urllib.error import HTTPError, URLError from urllib.parse import quote, unquote, urlparse from urllib.request import Request, urlopen +import httpx +import requests from nmp.common.auth.models import NMP_PRINCIPAL_ENVVAR, Principal from nmp.common.jobs.constants import NEMO_JOB_SECRETS_ENVVAR from opentelemetry._logs import Logger @@ -26,6 +29,11 @@ logger = logging.getLogger(__name__) +NMP_JOB_LAUNCHER_OTLP_LOGS_ENDPOINT_ENVVAR = "NMP_JOB_LAUNCHER_OTLP_LOGS_ENDPOINT" +NMP_JOB_LAUNCHER_OTLP_LOGS_HEADERS_ENVVAR = "NMP_JOB_LAUNCHER_OTLP_LOGS_HEADERS" +NMP_JOB_LAUNCHER_OTLP_LOGS_SOCKET_PATH_ENVVAR = "NMP_JOB_LAUNCHER_OTLP_LOGS_SOCKET_PATH" +NMP_JOB_LAUNCHER_OTLP_LOGS_TRANSPORT_ENVVAR = "NMP_JOB_LAUNCHER_OTLP_LOGS_TRANSPORT" + @dataclass(frozen=True) class SecretReference: @@ -55,6 +63,79 @@ def close(self) -> None: self.provider.shutdown() +class _UnixSocketOTLPSession(requests.Session): + def __init__(self, socket_path: str) -> None: + super().__init__() + self._client = httpx.Client( + transport=httpx.HTTPTransport(uds=socket_path), + follow_redirects=True, + ) + + def request( + self, + method: Any, + url: Any, + params: Any = None, + data: Any = None, + headers: Any = None, + cookies: Any = None, + files: Any = None, + auth: Any = None, + timeout: Any = None, + allow_redirects: bool = True, + proxies: Any = None, + hooks: Any = None, + stream: Any = None, + verify: Any = None, + cert: Any = None, + json: Any = None, + ) -> requests.Response: + try: + response = self._client.request( + str(method), + str(url), + params=params, + content=data, + headers=_merge_headers(self.headers, headers), + json=json, + timeout=timeout, + follow_redirects=allow_redirects, + ) + except httpx.TimeoutException as error: + raise requests.exceptions.Timeout(str(error)) from error + except httpx.TransportError as error: + raise requests.exceptions.ConnectionError(str(error)) from error + except httpx.HTTPError as error: + raise requests.exceptions.RequestException(str(error)) from error + return _to_requests_response(response) + + def close(self) -> None: + self._client.close() + super().close() + + +def _merge_headers(base_headers: Mapping[Any, Any], extra_headers: Any) -> dict[str, str]: + headers = {str(key): str(value) for key, value in base_headers.items()} + if extra_headers is None: + return headers + if isinstance(extra_headers, Mapping): + headers.update({str(key): str(value) for key, value in extra_headers.items()}) + return headers + headers.update({str(key): str(value) for key, value in extra_headers}) + return headers + + +def _to_requests_response(response: httpx.Response) -> requests.Response: + requests_response = requests.Response() + requests_response.status_code = response.status_code + requests_response.reason = response.reason_phrase + requests_response.url = str(response.url) + requests_response.headers.update(response.headers) + requests_response._content = response.content + requests_response.encoding = response.encoding + return requests_response + + def parse_secret_references(secrets_env: str) -> list[SecretReference]: if not secrets_env: return [] @@ -155,11 +236,11 @@ def create_otel_logger( step: str, task_id: str, ) -> SubprocessOtelLogger | None: - endpoint = env.get("OTEL_EXPORTER_OTLP_LOGS_ENDPOINT") or env.get("OTEL_EXPORTER_OTLP_ENDPOINT") + endpoint = env.get(NMP_JOB_LAUNCHER_OTLP_LOGS_ENDPOINT_ENVVAR) if not endpoint: return None - headers = _parse_otel_headers(env.get("OTEL_EXPORTER_OTLP_LOGS_HEADERS", "")) + headers = _parse_otel_headers(env.get(NMP_JOB_LAUNCHER_OTLP_LOGS_HEADERS_ENVVAR, "")) resource = Resource.create( { "workspace": workspace, @@ -170,13 +251,28 @@ def create_otel_logger( } ) logger_provider = LoggerProvider(resource=resource) - logger_provider.add_log_record_processor( - BatchLogRecordProcessor(OTLPLogExporter(endpoint=endpoint, headers=headers or None)) - ) + logger_provider.add_log_record_processor(BatchLogRecordProcessor(_build_otlp_log_exporter(env, endpoint, headers))) logger.info("Created local OTEL logger", extra={"endpoint": endpoint, "job": job, "step": step}) return SubprocessOtelLogger(logger_provider.get_logger("nmp.jobs.subprocess"), logger_provider) +def _build_otlp_log_exporter(env: dict[str, str], endpoint: str, headers: dict[str, str]) -> OTLPLogExporter: + transport = env.get(NMP_JOB_LAUNCHER_OTLP_LOGS_TRANSPORT_ENVVAR, "tcp") + if transport == "tcp": + return OTLPLogExporter(endpoint=endpoint, headers=headers or None) + if transport != "uds": + raise ValueError(f"unsupported OTLP logs transport: {transport!r}") + + socket_path = env.get(NMP_JOB_LAUNCHER_OTLP_LOGS_SOCKET_PATH_ENVVAR) + if not socket_path: + raise ValueError(f"{NMP_JOB_LAUNCHER_OTLP_LOGS_SOCKET_PATH_ENVVAR} is required for UDS OTLP logs") + return OTLPLogExporter( + endpoint=endpoint, + headers=headers or None, + session=_UnixSocketOTLPSession(socket_path), + ) + + def _parse_otel_headers(headers_env: str) -> dict[str, str]: if not headers_env: return {} diff --git a/services/core/jobs/tests/controllers/test_base.py b/services/core/jobs/tests/controllers/test_base.py index d5ba6f2ace..8ccefa73d3 100644 --- a/services/core/jobs/tests/controllers/test_base.py +++ b/services/core/jobs/tests/controllers/test_base.py @@ -18,12 +18,16 @@ from nmp.core.jobs.app.schemas import PlatformJobStepSpec, StepLifecycle from nmp.core.jobs.controllers.backends.base import ( JOB_LOGS_ENDPOINT_ENVVAR, + NMP_JOB_LAUNCHER_OTLP_LOGS_ENDPOINT_ENVVAR, + NMP_JOB_LAUNCHER_OTLP_LOGS_SOCKET_PATH_ENVVAR, + NMP_JOB_LAUNCHER_OTLP_LOGS_TRANSPORT_ENVVAR, WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR, JobExecutionProfileConfig, _contains_loopback_address, _replace_loopback_address, find_reserved_managed_job_environment_variable_names, get_job_runtime_shared_envvars, + get_logs_endpoint_config_from_fileset, get_logs_endpoint_from_fileset, get_workload_identity_token_audience, resolve_task_image, @@ -341,6 +345,38 @@ def test_service_discovery_files_with_loopback_replacement(self): "http://host.docker.internal:3000/apis/files/v2/workspaces/default/filesets/job-logs-123/otlp/v1/logs" ) + def test_uds_files_url_uses_placeholder_http_origin(self): + """UDS files URL is normalized to a valid HTTP URL for OTLP request construction.""" + config = PlatformConfig( # type: ignore[abstract] + service_discovery={"files": "unix:///tmp/nemo-platform.sock"}, + loopback_address="host.docker.internal", + ) + + result = get_logs_endpoint_from_fileset(config, workspace="default", fileset_id="job-logs-123") + + assert result == ( + "http://nemo-platform.local/apis/files/v2/workspaces/default/filesets/job-logs-123/otlp/v1/logs" + ) + + def test_uds_files_url_preserves_transport_metadata(self): + config = PlatformConfig( # type: ignore[abstract] + service_discovery={"files": "unix:///tmp/nemo-platform.sock"}, + loopback_address="host.docker.internal", + ) + + result = get_logs_endpoint_config_from_fileset(config, workspace="default", fileset_id="job-logs-123") + + assert result.endpoint == ( + "http://nemo-platform.local/apis/files/v2/workspaces/default/filesets/job-logs-123/otlp/v1/logs" + ) + assert result.transport == "uds" + assert result.socket_path == "/tmp/nemo-platform.sock" + assert result.to_env() == { + NMP_JOB_LAUNCHER_OTLP_LOGS_ENDPOINT_ENVVAR: result.endpoint, + NMP_JOB_LAUNCHER_OTLP_LOGS_TRANSPORT_ENVVAR: "uds", + NMP_JOB_LAUNCHER_OTLP_LOGS_SOCKET_PATH_ENVVAR: "/tmp/nemo-platform.sock", + } + class TestGetJobRuntimeSharedEnvvars: def test_uses_service_discovery_gateway_urls_for_job_runtimes(self): diff --git a/services/core/jobs/tests/controllers/test_docker_backend.py b/services/core/jobs/tests/controllers/test_docker_backend.py index 08102731d7..4a52fd2564 100644 --- a/services/core/jobs/tests/controllers/test_docker_backend.py +++ b/services/core/jobs/tests/controllers/test_docker_backend.py @@ -53,7 +53,7 @@ PlatformJobStepSpec, ) from nmp.core.jobs.controllers.backends.base import ( - JOB_LOGS_ENDPOINT_ENVVAR, + NMP_JOB_LAUNCHER_OTLP_LOGS_ENDPOINT_ENVVAR, WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR, WORKLOAD_IDENTITY_TOKEN_FILE_PATH, WORKLOAD_IDENTITY_VOLUME_PATH, @@ -769,7 +769,7 @@ def test_docker_job_uses_service_discovery_urls_for_job_runtime(mock_nmp_client, assert env_vars["NMP_FILES_URL"] == "https://nemo-gateway:8080" assert env_vars["NMP_MODELS_URL"] == "https://nemo-gateway:8080" assert env_vars["NMP_SECRETS_URL"] == "https://nemo-gateway:8080" - assert env_vars[JOB_LOGS_ENDPOINT_ENVVAR].startswith("https://nemo-gateway:8080/apis/files/") + assert env_vars[NMP_JOB_LAUNCHER_OTLP_LOGS_ENDPOINT_ENVVAR].startswith("https://nemo-gateway:8080/apis/files/") def test_docker_job_execution_profile_config_rejects_reserved_env_vars(): @@ -2471,12 +2471,11 @@ def test_job_step_with_auth_context(): def test_docker_job_schedule_with_auth_context(docker_job, docker_client_mock, test_job_step_with_auth_context): - """Test that scheduling sets NMP_PRINCIPAL without injecting OTEL log headers. + """Test that scheduling sets NMP_PRINCIPAL and launcher OTLP headers when auth_context is present. Verifies GitLab issue #3390 Gap 2: job tasks should run with the creating - user's auth context, propagated via the NMP_PRINCIPAL environment variable. - Job log upload auth is handled by jobs-launcher workload token exchange, - not by globally scoped OTEL header environment variables. + user's auth context, propagated via the NMP_PRINCIPAL environment variable + and private launcher OTLP headers for authenticated telemetry export. """ step_spec = test_job_step_with_auth_context.step_spec executor_config = step_spec.executor @@ -2505,12 +2504,16 @@ def test_docker_job_schedule_with_auth_context(docker_job, docker_client_mock, t "groups": ["engineering", "ml-team"], } - assert env[JOB_LOGS_ENDPOINT_ENVVAR].endswith( - "/apis/files/v2/workspaces/default/filesets/test-logs-fileset/otlp/v1/logs" - ) + # Verify launcher OTLP headers env var is set for authenticated telemetry + assert "NMP_JOB_LAUNCHER_OTLP_LOGS_HEADERS" in env + otlp_headers = env["NMP_JOB_LAUNCHER_OTLP_LOGS_HEADERS"] + # URL-encoded: @ -> %40, , -> %2C + assert "X-NMP-Principal-Id=creator%40example.com" in otlp_headers + assert "X-NMP-Principal-Email=creator%40example.com" in otlp_headers + assert "X-NMP-Principal-Groups=engineering%2Cml-team" in otlp_headers + + # Verify no globally scoped OTEL header environment variables are set assert "OTEL_EXPORTER_OTLP_LOGS_ENDPOINT" not in env - assert "OTEL_EXPORTER_OTLP_PROTOCOL" not in env - assert "OTEL_EXPORTER_OTLP_LOGS_PROTOCOL" not in env assert "OTEL_LOGS_EXPORTER" not in env assert "OTEL_SERVICE_NAME" not in env assert "OTEL_EXPORTER_OTLP_LOGS_HEADERS" not in env @@ -2534,7 +2537,7 @@ def test_docker_job_schedule_without_auth_context(docker_job, docker_client_mock # Verify auth env vars are NOT set env = kwargs["environment"] assert NMP_PRINCIPAL_ENVVAR not in env - assert "OTEL_EXPORTER_OTLP_LOGS_HEADERS" not in env + assert "NMP_JOB_LAUNCHER_OTLP_LOGS_HEADERS" not in env def test_docker_job_schedule_with_auth_context_empty_groups(): diff --git a/services/core/jobs/tests/controllers/test_kubernetes_backend.py b/services/core/jobs/tests/controllers/test_kubernetes_backend.py index 728bb0354c..6f82e6a21c 100644 --- a/services/core/jobs/tests/controllers/test_kubernetes_backend.py +++ b/services/core/jobs/tests/controllers/test_kubernetes_backend.py @@ -38,7 +38,7 @@ PlatformJobStepSpec, ) from nmp.core.jobs.controllers.backends.base import ( - JOB_LOGS_ENDPOINT_ENVVAR, + NMP_JOB_LAUNCHER_OTLP_LOGS_ENDPOINT_ENVVAR, WORKLOAD_IDENTITY_TOKEN_FILE_ENVVAR, WORKLOAD_IDENTITY_TOKEN_FILE_PATH, WORKLOAD_IDENTITY_VOLUME_NAME, @@ -884,7 +884,7 @@ def test_kubernetes_job_uses_service_discovery_urls_for_job_runtime( assert env_vars["NMP_FILES_URL"] == "https://nemo-gateway:8080" assert env_vars["NMP_MODELS_URL"] == "https://nemo-gateway:8080" assert env_vars["NMP_SECRETS_URL"] == "https://nemo-gateway:8080" - assert env_vars[JOB_LOGS_ENDPOINT_ENVVAR].startswith("https://nemo-gateway:8080/apis/files/") + assert env_vars[NMP_JOB_LAUNCHER_OTLP_LOGS_ENDPOINT_ENVVAR].startswith("https://nemo-gateway:8080/apis/files/") def test_kubernetes_job_execution_profile_config_rejects_reserved_env_vars(): @@ -2145,12 +2145,11 @@ def test_step_pending_with_auth_context() -> PlatformJobStepWithContext: def test_kubernetes_job_schedule_with_auth_context( kubernetes_job, cpu_execution_provider, test_step_pending_with_auth_context ): - """Test that scheduling sets NMP_PRINCIPAL without injecting OTEL log headers. + """Test that scheduling sets NMP_PRINCIPAL and launcher OTLP headers when auth_context is present. Verifies GitLab issue #3390 Gap 2: job tasks should run with the creating - user's auth context, propagated via the NMP_PRINCIPAL environment variable. - Job log upload auth is handled by jobs-launcher workload token exchange, - not by globally scoped OTEL header environment variables. + user's auth context, propagated via the NMP_PRINCIPAL environment variable + and private launcher OTLP headers for authenticated telemetry export. """ import json @@ -2184,12 +2183,16 @@ def test_kubernetes_job_schedule_with_auth_context( "groups": ["engineering", "ml-team"], } - assert env_vars[JOB_LOGS_ENDPOINT_ENVVAR].endswith( - "/apis/files/v2/workspaces/default/filesets/test-logs-fileset/otlp/v1/logs" - ) + # Verify launcher OTLP headers env var is set for authenticated telemetry + assert "NMP_JOB_LAUNCHER_OTLP_LOGS_HEADERS" in env_vars + otlp_headers = env_vars["NMP_JOB_LAUNCHER_OTLP_LOGS_HEADERS"] + # URL-encoded: @ -> %40, , -> %2C + assert "X-NMP-Principal-Id=creator%40example.com" in otlp_headers + assert "X-NMP-Principal-Email=creator%40example.com" in otlp_headers + assert "X-NMP-Principal-Groups=engineering%2Cml-team" in otlp_headers + + # Verify no globally scoped OTEL header environment variables are set assert "OTEL_EXPORTER_OTLP_LOGS_ENDPOINT" not in env_var_names - assert "OTEL_EXPORTER_OTLP_PROTOCOL" not in env_var_names - assert "OTEL_EXPORTER_OTLP_LOGS_PROTOCOL" not in env_var_names assert "OTEL_LOGS_EXPORTER" not in env_var_names assert "OTEL_SERVICE_NAME" not in env_var_names assert "OTEL_EXPORTER_OTLP_LOGS_HEADERS" not in env_var_names @@ -2217,7 +2220,7 @@ def test_kubernetes_job_schedule_without_auth_context(kubernetes_job, cpu_execut # Verify auth env vars are NOT set assert NMP_PRINCIPAL_ENVVAR not in env_vars - assert "OTEL_EXPORTER_OTLP_LOGS_HEADERS" not in env_vars + assert "NMP_JOB_LAUNCHER_OTLP_LOGS_HEADERS" not in env_vars def test_cleanup_steps_with_multi_step_job_only_first_step_complete(kubernetes_job): diff --git a/services/core/jobs/tests/controllers/test_subprocess_backend.py b/services/core/jobs/tests/controllers/test_subprocess_backend.py index a32d8d220b..9aed7313bb 100644 --- a/services/core/jobs/tests/controllers/test_subprocess_backend.py +++ b/services/core/jobs/tests/controllers/test_subprocess_backend.py @@ -8,8 +8,14 @@ from types import SimpleNamespace from unittest.mock import patch +from nmp.common.config import PlatformConfig from nmp.common.jobs.schemas import PlatformJobStatus from nmp.core.jobs.app.providers import SubprocessExecutionProvider +from nmp.core.jobs.controllers.backends.base import ( + NMP_JOB_LAUNCHER_OTLP_LOGS_ENDPOINT_ENVVAR, + NMP_JOB_LAUNCHER_OTLP_LOGS_SOCKET_PATH_ENVVAR, + NMP_JOB_LAUNCHER_OTLP_LOGS_TRANSPORT_ENVVAR, +) from nmp.core.jobs.controllers.backends.subprocess import ( SubprocessJobBackend, SubprocessJobExecutionProfileConfig, @@ -157,6 +163,38 @@ def test_schedule_uses_allowlisted_host_environment(mock_nmp_client, tmp_path, m assert metadata.process.wait(timeout=5) == 0 +def test_schedule_preserves_uds_otlp_metadata_in_runtime_env(mock_nmp_client, tmp_path, test_step_pending): + platform_config = PlatformConfig( # type: ignore[abstract] + service_discovery={"files": "unix:///tmp/nemo-platform.sock"}, + loopback_address=None, + ) + backend = _subprocess_backend(mock_nmp_client, tmp_path, platform_config) + step = _step_with_command(test_step_pending, ["/bin/sh", "-c", "true"]) + captured_env = {} + + def fake_create_otel_logger(*, env, **_kwargs): + captured_env.update(env) + return None + + with ( + patch("nmp.core.jobs.controllers.backends.subprocess.get_platform_config", return_value=platform_config), + patch("nmp.core.jobs.controllers.backends.subprocess.create_otel_logger", side_effect=fake_create_otel_logger), + ): + update = backend.schedule(step.step_spec.executor, step) + + assert update.status == PlatformJobStatus.PENDING + metadata = backend._process_registry.get( + SubprocessProcessKey(step.workspace, step.job, str(step.attempt_id), step.name) + ) + assert metadata is not None + assert metadata.process.wait(timeout=5) == 0 + assert captured_env[NMP_JOB_LAUNCHER_OTLP_LOGS_ENDPOINT_ENVVAR] == ( + "http://nemo-platform.local/apis/files/v2/workspaces/default/filesets/test-logs-fileset/otlp/v1/logs" + ) + assert captured_env[NMP_JOB_LAUNCHER_OTLP_LOGS_TRANSPORT_ENVVAR] == "uds" + assert captured_env[NMP_JOB_LAUNCHER_OTLP_LOGS_SOCKET_PATH_ENVVAR] == "/tmp/nemo-platform.sock" + + def test_schedule_terminates_process_when_post_popen_setup_fails( mock_nmp_client, tmp_path, mock_platform_config, test_step_pending ): diff --git a/services/core/jobs/tests/controllers/test_subprocess_runtime.py b/services/core/jobs/tests/controllers/test_subprocess_runtime.py index 2b6a6b7ff6..3f3e7fad61 100644 --- a/services/core/jobs/tests/controllers/test_subprocess_runtime.py +++ b/services/core/jobs/tests/controllers/test_subprocess_runtime.py @@ -11,7 +11,11 @@ from nmp.common.auth.models import NMP_PRINCIPAL_ENVVAR from nmp.common.jobs.constants import NEMO_JOB_SECRETS_ENVVAR from nmp.core.jobs.controllers.backends.subprocess_runtime import ( + NMP_JOB_LAUNCHER_OTLP_LOGS_SOCKET_PATH_ENVVAR, + NMP_JOB_LAUNCHER_OTLP_LOGS_TRANSPORT_ENVVAR, SubprocessOtelLogger, + _build_otlp_log_exporter, + _UnixSocketOTLPSession, inject_secret_env_vars, parse_secret_references, start_log_capture, @@ -122,3 +126,37 @@ def test_local_otel_logger_close_flushes_and_shuts_down(): mock_otel_logger.emit.assert_called_once() mock_provider.force_flush.assert_called_once() mock_provider.shutdown.assert_called_once() + + +def test_build_otlp_log_exporter_keeps_default_http_exporter_for_tcp(): + with patch("nmp.core.jobs.controllers.backends.subprocess_runtime.OTLPLogExporter") as exporter: + result = _build_otlp_log_exporter({}, "http://files.example/otlp/v1/logs", {"x-test": "yes"}) + + assert result is exporter.return_value + exporter.assert_called_once_with(endpoint="http://files.example/otlp/v1/logs", headers={"x-test": "yes"}) + + +def test_build_otlp_log_exporter_uses_unix_socket_session_for_uds(): + env = { + NMP_JOB_LAUNCHER_OTLP_LOGS_TRANSPORT_ENVVAR: "uds", + NMP_JOB_LAUNCHER_OTLP_LOGS_SOCKET_PATH_ENVVAR: "/tmp/nemo-platform.sock", + } + + with patch("nmp.core.jobs.controllers.backends.subprocess_runtime.OTLPLogExporter") as exporter: + result = _build_otlp_log_exporter(env, "http://nemo-platform.local/otlp/v1/logs", {}) + + assert result is exporter.return_value + kwargs = exporter.call_args.kwargs + assert kwargs["endpoint"] == "http://nemo-platform.local/otlp/v1/logs" + assert kwargs["headers"] is None + assert isinstance(kwargs["session"], _UnixSocketOTLPSession) + kwargs["session"].close() + + +def test_build_otlp_log_exporter_rejects_uds_without_socket_path(): + with pytest.raises(ValueError, match=NMP_JOB_LAUNCHER_OTLP_LOGS_SOCKET_PATH_ENVVAR): + _build_otlp_log_exporter( + {NMP_JOB_LAUNCHER_OTLP_LOGS_TRANSPORT_ENVVAR: "uds"}, + "http://nemo-platform.local/otlp/v1/logs", + {}, + ) From 7b02be823cf306f117855e2522f5ce31a3fe265a Mon Sep 17 00:00:00 2001 From: Sean Teramae Date: Tue, 21 Jul 2026 18:52:04 -0400 Subject: [PATCH 07/24] feat(studio): Dedicated Agent Details Page (#806) * feat(studio): Dedicated Agent Details Page Signed-off-by: Sean Teramae * fix agent suggestions route Signed-off-by: Sean Teramae * Remove unused tabs Signed-off-by: Sean Teramae * fix overflow Signed-off-by: Sean Teramae * fix type error Signed-off-by: Sean Teramae --------- Signed-off-by: Sean Teramae Signed-off-by: Sam Oluwalana --- .../AgentPanel/DeploymentLogsView.tsx | 46 ++-- .../WorkspaceLayout/WorkspaceSideNav.tsx | 44 +-- .../AgentDetailRoute/ConfigurationTab.tsx | 38 +++ .../AgentDetailRoute/DeploymentsTab.tsx | 96 +++++++ .../AgentDetailRoute/EvaluationsTab.tsx | 70 +++++ .../AgentDetailRoute/TabPlaceholder.tsx | 34 +++ .../agents/AgentDetailRoute/index.test.tsx | 49 ++++ .../routes/agents/AgentDetailRoute/index.tsx | 258 ++++++++++++++++++ .../AgentDetailRoute/overview/DetailPanel.tsx | 26 ++ .../routes/agents/AgentsListRoute/index.tsx | 46 +--- .../studio/src/routes/groups/agentRoutes.tsx | 44 ++- 11 files changed, 649 insertions(+), 102 deletions(-) create mode 100644 web/packages/studio/src/routes/agents/AgentDetailRoute/ConfigurationTab.tsx create mode 100644 web/packages/studio/src/routes/agents/AgentDetailRoute/DeploymentsTab.tsx create mode 100644 web/packages/studio/src/routes/agents/AgentDetailRoute/EvaluationsTab.tsx create mode 100644 web/packages/studio/src/routes/agents/AgentDetailRoute/TabPlaceholder.tsx create mode 100644 web/packages/studio/src/routes/agents/AgentDetailRoute/index.test.tsx create mode 100644 web/packages/studio/src/routes/agents/AgentDetailRoute/index.tsx create mode 100644 web/packages/studio/src/routes/agents/AgentDetailRoute/overview/DetailPanel.tsx diff --git a/web/packages/studio/src/components/sidePanels/AgentPanels/AgentPanel/DeploymentLogsView.tsx b/web/packages/studio/src/components/sidePanels/AgentPanels/AgentPanel/DeploymentLogsView.tsx index eb3ff1eed3..9f69d1955f 100644 --- a/web/packages/studio/src/components/sidePanels/AgentPanels/AgentPanel/DeploymentLogsView.tsx +++ b/web/packages/studio/src/components/sidePanels/AgentPanels/AgentPanel/DeploymentLogsView.tsx @@ -19,9 +19,18 @@ interface DeploymentLogsViewProps { workspace: string; /** All known deployments for the active agent — any status. */ deployments: AgentDeployment[]; + /** Controlled selection — when provided, the picker reflects this deployment. */ + selectedDeploymentName?: string; + /** Called when the user (or a caller) changes the selected deployment. */ + onSelectDeployment?: (name: string) => void; } -export const DeploymentLogsView: FC = ({ workspace, deployments }) => { +export const DeploymentLogsView: FC = ({ + workspace, + deployments, + selectedDeploymentName, + onSelectDeployment, +}) => { const sortedDeployments = useMemo( () => [...deployments].sort((a, b) => { @@ -31,18 +40,25 @@ export const DeploymentLogsView: FC = ({ workspace, dep }), [deployments] ); - const [selectedName, setSelectedName] = useState( - () => sortedDeployments[0]?.name + const [internalName, setInternalName] = useState( + () => selectedDeploymentName ?? sortedDeployments[0]?.name ); + const isControlled = selectedDeploymentName !== undefined; + const selectedName = isControlled ? selectedDeploymentName : internalName; + const setSelectedName = (name: string) => { + if (!isControlled) setInternalName(name); + onSelectDeployment?.(name); + }; useEffect(() => { - if (!selectedName) { - setSelectedName(sortedDeployments[0]?.name); + if (isControlled) return; + if (!internalName) { + setInternalName(sortedDeployments[0]?.name); return; } - const stillPresent = sortedDeployments.some((d) => d.name === selectedName); - if (!stillPresent) setSelectedName(sortedDeployments[0]?.name); - }, [sortedDeployments, selectedName]); + const stillPresent = sortedDeployments.some((d) => d.name === internalName); + if (!stillPresent) setInternalName(sortedDeployments[0]?.name); + }, [sortedDeployments, internalName, isControlled]); if (sortedDeployments.length === 0) { return ( @@ -53,11 +69,12 @@ export const DeploymentLogsView: FC = ({ workspace, dep ); } + const pOffset = '2'; return ( - + {sortedDeployments.length > 1 && ( - + - onPatch({ - [SEED_FILE_PATH_KEY]: value ?? '', - [SEED_AVAILABLE_COLUMNS_KEY]: '', - }) - } - placeholder={ - !filesetRef - ? 'Select a fileset first' - : isLoadingFiles - ? 'Loading files…' - : 'Select a file' - } - /> - + setValue(availableColumnsPath, '')} + placeholder={ + !filesetRef + ? 'Select a fileset first' + : isLoadingFiles + ? 'Loading files…' + : 'Select a file' + } + /> {filePath && ( = ({ values, onPatch )} - -