diff --git a/docs/user-guide/docker.md b/docs/user-guide/docker.md index 7c1baf056..7911f850f 100644 --- a/docs/user-guide/docker.md +++ b/docs/user-guide/docker.md @@ -244,7 +244,7 @@ docker run --gpus all --shm-size=1g \ nss-gpu:latest run --config /workspace/data/config.yaml --data-source /workspace/data/input.csv ``` -See [Environment Variables -- Hugging Face Cache](environment.md#hugging-face-cache) +See [Environment Variables -- Hugging Face cache and offline](environment.md#hugging-face-cache-and-offline) for details on `HF_HOME`, `HF_HUB_OFFLINE`, and `VLLM_CACHE_ROOT`. --- diff --git a/docs/user-guide/environment.md b/docs/user-guide/environment.md index b28252269..bb690deff 100644 --- a/docs/user-guide/environment.md +++ b/docs/user-guide/environment.md @@ -3,274 +3,290 @@ # Environment Variables -All environment variables that affect Safe Synthesizer behavior. For runtime -errors and OOM issues, see [Program Runtime](troubleshooting.md). For output -quality and evaluation metrics, see [Synthetic Data Quality](evaluating-data.md). +Reference for infrastructure settings: artifact paths, logging, model caches, +network endpoints, and third-party library behavior. Synthesis parameters (`training.learning_rate`, `generation.num_records`, etc.) are set via YAML, CLI flags, or the Python SDK -- not environment variables. -Environment variables control infrastructure: where artifacts go, how models -are cached, and which network endpoints are used. +See [Configuration Reference](configuration.md) for parameter tables and +[Configuration Precedence](configuration.md#configuration-precedence) for how +YAML, CLI, and SDK layers combine. + +For runtime errors and OOM issues, see [Program Runtime](troubleshooting.md). +For output quality and evaluation metrics, see +[Synthetic Data Quality](evaluating-data.md). --- -## NSS Variables - -| Variable | CLI flag | Purpose | -|----------|----------|---------| -| `NSS_CONFIG` | `--config` | Path to YAML config file | -| `NSS_ARTIFACTS_PATH` | `--artifact-path` | Default artifact path | -| `NSS_LOG_FORMAT` | `--log-format` | Log format (`json` or `plain`) | -| `NSS_LOG_FILE` | `--log-file` | Log file path | -| `NSS_LOG_COLOR` | `--log-color` / `--no-log-color` | Colorize console output (auto-detected from TTY) | -| `NSS_LOG_LEVEL` | -- | Log level: `DEBUG`, `INFO`, `WARNING`, `ERROR`, `CRITICAL`, or `DEBUG_DEPENDENCIES` | -| `NSS_DATASET_REGISTRY` | `--dataset-registry` | Dataset registry YAML path/URL | -| `NSS_WANDB_MODE` | `--wandb-mode` | WandB mode (alias for `WANDB_MODE`) | -| `NSS_WANDB_PROJECT` | `--wandb-project` | WandB project name (alias for `WANDB_PROJECT`) | -| `NSS_INFERENCE_ENDPOINT` | -- | LLM endpoint for PII column classification (default: `https://integrate.api.nvidia.com/v1`) | -| `NSS_INFERENCE_KEY` | -- | API key for the `NSS_INFERENCE_ENDPOINT` is required for column classification in both CLI and SDK. | -| `NIM_MODEL_ID` | -- | Column classification model ID | -| `LOCAL_FILES_ONLY` | -- | Set to `true` for offline mode (GLiNER) | -| `SAFE_SYNTHESIZER_CPU_COUNT` | -- | NER CPU processes | +## At a glance ---- +| Task | Start here | +|------|------------| +| Run offline or air-gapped | [HF cache and offline](#hugging-face-cache-and-offline) · [Running in Offline Environments](running.md#running-in-offline-environments) | +| Docker / container mounts | [Containers](#containers) · [Docker](docker.md) | +| Logging and WandB | [Running -- Logging and Experiment Tracking](running.md#logging-and-experiment-tracking) | +| PII column classification API key | [PII, NER, and column classification](#pii-ner-and-column-classification) · [Running -- LLM Column Classification](running.md#llm-column-classification) | +| Disable telemetry | [Telemetry](#telemetry) | +| Resolve CLI vs env vs defaults | [Precedence](#precedence) | -## Third-Party Variables +--- -| Variable | Read by | Purpose | -|----------|---------|---------| -| `HF_HOME` | Hugging Face Hub | Cache directory for model downloads | -| `HF_HUB_OFFLINE` | Hugging Face Hub | Set to `1` to error instead of downloading | -| `VLLM_ATTENTION_BACKEND` | vLLM | Override attention backend | -| `VLLM_CACHE_ROOT` | vLLM | vLLM internal cache directory (defaults to `~/.cache/vllm`) | -| `WANDB_MODE` | WandB | Mode (`online`, `offline`, `disabled`) | -| `WANDB_PROJECT` | WandB | Project name | -| `WANDB_API_KEY` | WandB | API key for authentication | +## Master reference table + +Grouped by the `Category` column -- `nss`-native settings first, then +`telemetry`, `third-party`, `container`, and `internal`. + +| Variable | Category | CLI flag | Read by | Default | Purpose | Details | +|----------|----------|----------|---------|---------|---------|---------| +| `NSS_CONFIG` | nss | `--config` | CLI | -- | Path to YAML config file | [Configuration Reference](configuration.md) | +| `NSS_ARTIFACTS_PATH` | nss | `--artifact-path` | CLI | `./safe-synthesizer-artifacts` | Base directory for run artifacts | [Running -- Artifacts](running.md#artifacts-and-output) | +| `NSS_LOG_FORMAT` | nss | `--log-format` | CLI / observability | auto (`plain` on TTY, else `json`) | Console log format | [Running -- Log Format](running.md#log-format) | +| `NSS_LOG_FILE` | nss | `--log-file` | CLI / observability | run log under workdir | Path to log file | [Running -- Logging](running.md#logging-and-experiment-tracking) | +| `NSS_LOG_COLOR` | nss | `--log-color` / `--no-log-color` | CLI / observability | auto (TTY) | Colorize console output | [Running -- Log Format](running.md#log-format) | +| `NSS_LOG_LEVEL` | nss | `--verbose` (0–2) | observability | `INFO` | Log level (`DEBUG`, `DEBUG_DEPENDENCIES`, etc.) | Set via verbosity, not a direct CLI flag | +| `NSS_DATASET_REGISTRY` | nss | `--dataset-registry` | CLI | -- | Dataset registry YAML path or URL | [Running -- Dataset Registry](running.md#dataset-registry) | +| `NSS_WANDB_MODE` | nss | `--wandb-mode` | WandB | `disabled` | WandB run mode | Alias for `WANDB_MODE` | +| `NSS_WANDB_PROJECT` | nss | `--wandb-project` | WandB | -- | WandB project name | Alias for `WANDB_PROJECT` | +| `NSS_INFERENCE_ENDPOINT` | nss | `--inference-endpoint-url` | PII column classifier | NVIDIA integrate URL | OpenAI-compatible endpoint for column classification | [PII appendix](#pii-ner-and-column-classification) | +| `NSS_INFERENCE_KEY` | nss | `--inference-api-key` | PII column classifier | -- | API key for `NSS_INFERENCE_ENDPOINT` | Required for LLM column classification | +| `NSS_INFERENCE_MODEL` | nss | `--inference-model-id` | PII column classifier | `qwen/qwen3-next-80b-a3b-instruct` | Model ID sent to the inference endpoint | [PII appendix](#pii-ner-and-column-classification) | +| `NSS_PII_REPLACER_CPU_COUNT` | nss | `--cpu-count` | NER worker pool | `max(1, cpu_count - 1)` | CPU processes for PII NER | [PII appendix](#pii-ner-and-column-classification) | +| `NEMO_TELEMETRY_ENABLED` | telemetry | `--emit_telemetry` | telemetry | `true` | Enable anonymous usage telemetry | Also `emit_telemetry` in YAML; see [Telemetry](#telemetry) | +| `HF_HOME` | third-party | -- | Hugging Face Hub | platform cache dir | Root directory for HF downloads | [HF appendix](#hugging-face-cache-and-offline) | +| `HF_HUB_OFFLINE` | third-party | `--enable-huggingface-remote` / `--disable-huggingface-remote` | Hugging Face Hub | unset | Fail if a model is not cached (covers base model and GLiNER) | Preferred offline gate; CLI flag also sets `TRANSFORMERS_OFFLINE` | +| `VLLM_CACHE_ROOT` | third-party | -- | vLLM | `~/.cache/vllm` | vLLM model cache directory | [vLLM appendix](#vllm-and-attention) | +| `VLLM_ATTENTION_BACKEND` | third-party | -- | vLLM | auto | Override attention implementation | [vLLM appendix](#vllm-and-attention) | +| `WANDB_MODE` | third-party | `--wandb-mode` | WandB | `disabled` | WandB run mode | Same as `NSS_WANDB_MODE` | +| `WANDB_PROJECT` | third-party | `--wandb-project` | WandB | -- | WandB project name | Same as `NSS_WANDB_PROJECT` | +| `WANDB_API_KEY` | third-party | -- | WandB | -- | WandB authentication | Required for online logging | +| `NVIDIA_VISIBLE_DEVICES` | container | -- | NVIDIA runtime | all visible GPUs | Limit GPUs inside a container | [Containers](#containers) · [Docker -- GPU Access](docker.md#gpu-access) | +| `NSS_OPT_BUCKET` | internal | -- | NER optimization | `nss-opt-dev-use2` | S3 bucket for optional NER opt artifacts | [Internal](#internal-and-cluster) | +| `NSS_OPT_CACHE_DIR` | internal | -- | NER optimization | `.optcache` | Local cache for NER optimization downloads | [Internal](#internal-and-cluster) | +| `NEMO_TELEMETRY_ENDPOINT` | internal | -- | telemetry | NVIDIA default | Override telemetry upload URL | [Telemetry](#telemetry) | +| `NEMO_SESSION_PREFIX` | internal | -- | telemetry | -- | Prefix for telemetry session IDs | [Telemetry](#telemetry) | +| `NEMO_JOB_ID` | internal | -- | evaluation reports | -- | Cluster job ID in multimodal reports | [Internal](#internal-and-cluster) | --- ## Precedence -Infrastructure settings (artifact path, logging, WandB): +### Infrastructure (CLISettings) + +For artifact paths, logging, WandB overrides, and the five runtime flags +(`--inference-*`, `--enable-huggingface-remote` / `--disable-huggingface-remote`, +`--cpu-count`): -1. CLI flags (`--artifact-path`, `--log-format`, etc.) -2. Environment variables (`NSS_ARTIFACTS_PATH`, `NSS_LOG_FORMAT`, etc.) +1. CLI flags +2. Environment variables 3. Built-in defaults ---- +WandB accepts both `NSS_WANDB_*` and `WANDB_*` names; CLI `--wandb-mode` and +`--wandb-project` override either. -## Hugging Face Cache +### Synthesis parameters -All model and tokenizer downloads go through -[Hugging Face Hub](https://huggingface.co/docs/huggingface_hub/guides/manage-cache). -The following variables control where downloads are stored and whether the -network is used. For a step-by-step offline setup guide, see -[Running in Offline Environments](running.md#running-in-offline-environments). +YAML fields, CLI `--section__field` overrides, and SDK builder calls follow +[Configuration Precedence](configuration.md#configuration-precedence) -- not +the order above. -### `HF_HOME` +### Telemetry precedence -Sets the root cache directory for all Hugging Face downloads -- model weights, -tokenizers, compiled attention kernels, and the SentenceTransformer used for -evaluation. +`--emit_telemetry` / `emit_telemetry` in YAML override `NEMO_TELEMETRY_ENABLED` +when explicitly set. When unset, the env var defaults to enabled. -```bash -export HF_HOME=/shared/cache/huggingface -``` +--- + +## Hugging Face cache and offline + +Downloads go through +[Hugging Face Hub](https://huggingface.co/docs/huggingface_hub/guides/manage-cache). +For a step-by-step offline workflow, see +[Running in Offline Environments](running.md#running-in-offline-environments) +and [Docker -- Offline and Air-Gapped Environments](docker.md#offline-and-air-gapped-environments). -### Pre-Caching Models +### `HF_HOME` -To avoid runtime downloads, run the pipeline once in an environment with -internet access, then copy or mount the populated cache in your target -environment: +Root cache for model weights, tokenizers, compiled attention kernels, GLiNER, +evaluation SentenceTransformer weights, and other Hub assets. ```bash export HF_HOME=/shared/cache/huggingface -safe-synthesizer run --config config.yaml --data-source data.csv ``` -What gets downloaded on first use: - -- Model weights, config, and tokenizer (all backends, via HF Hub) -- Compiled attention kernels when `training.attn_implementation` starts with - `kernels-community/` -- GLiNER NER model (PII replacement) -- `distiluse-base-multilingual-cased-v2` (evaluation semantic similarity) -- vLLM base model (generation) - -!!! warning "Silent downloads on first use" - All downloads happen silently on first use. If the first run is in an - environment without internet access, connection errors will appear at - whichever pipeline stage tries to download first. - ### `HF_HUB_OFFLINE` -When set, prevents all Hugging Face Hub network requests. Any attempt to -access a model that is not already cached raises an error immediately. +`HF_HUB_OFFLINE=1` tells Hugging Face Hub to refuse network access. It is the +canonical offline switch: huggingface_hub honors it globally, so a single +setting covers both the base model and GLiNER. Pair it with a pre-populated +`HF_HOME`. ```bash export HF_HUB_OFFLINE=1 ``` -Prefer this over `LOCAL_FILES_ONLY` for the most reliable offline experience -- -see the warning under `LOCAL_FILES_ONLY` below. - -### `LOCAL_FILES_ONLY` - -Skips network downloads for GLiNER. Not respected by the HuggingFace training -backend or vLLM. +Set it before the process starts. huggingface_hub reads the value once, when it +is first imported, and caches it -- changing it later has no effect for that +process. For the CLI, export it before launching `safe-synthesizer`. When +driving the pipeline programmatically, set it before importing +`nemo_safe_synthesizer`. -```bash -export LOCAL_FILES_ONLY=true -``` +### `--enable-huggingface-remote` / `--disable-huggingface-remote` -!!! warning "Partial offline support" - `LOCAL_FILES_ONLY` is not consistently supported across all backends. - Set `HF_HUB_OFFLINE=1` combined with a pre-populated `HF_HOME` cache - for the most reliable offline experience. +CLI shorthand for the switch above, with no separate NSS env var: -### `VLLM_CACHE_ROOT` +- `--disable-huggingface-remote` -- offline run; sets `HF_HUB_OFFLINE=1` and + `TRANSFORMERS_OFFLINE=1`. +- `--enable-huggingface-remote` -- online run; sets both to `0`, overriding any + inherited offline environment. +- Default (neither flag) -- the environment is left untouched: the run inherits + `HF_HUB_OFFLINE` / `TRANSFORMERS_OFFLINE` if set, and otherwise allows remote + downloads. The effective default is `--enable-huggingface-remote`. -Sets the vLLM model cache directory. +The CLI applies the flag before huggingface_hub loads, so the flag always wins +over an inherited environment value. For env-based control, set `HF_HUB_OFFLINE` +directly. ```bash -export VLLM_CACHE_ROOT=/shared/cache/vllm +safe-synthesizer run --disable-huggingface-remote ... ``` ---- - -## Attention and Compute - -GPU attention backend selection for the vLLM generation engine. +!!! warning "Models must be cached" + Offline mode requires the base model and GLiNER to already be present in + `HF_HOME`. Loading fails if a required model is not cached. -### `VLLM_ATTENTION_BACKEND` +### Pre-caching models -Controls the attention implementation used by the vLLM generation engine. -Safe Synthesizer sets this automatically when `generation.attention_backend` -is configured. Leave it unset unless you have a specific reason to override -vLLM's auto-detection. +Run once with network access, then copy or mount the populated cache. Typical +first-run downloads include training weights, GLiNER, evaluation embeddings, +and the vLLM base model. -```bash -export VLLM_ATTENTION_BACKEND=FLASH_ATTN -``` +!!! warning "Silent downloads on first use" + Downloads happen on first use. In an air-gapped environment, the first + missing asset fails at the stage that needs it. -Common values: `FLASHINFER`, `FLASH_ATTN`, `TORCH_SDPA`, `TRITON_ATTN`, -`FLEX_ATTENTION`. +See [Running in Offline Environments](running.md#running-in-offline-environments) +for the full pre-cache checklist. --- -## PII and NER +## PII, NER, and column classification -NIM endpoint, API keys, and CPU parallelism for PII detection. +Controls LLM-based column classification and CPU parallelism for NER-based PII +replacement. For setup examples and NER-only fallback behavior, see +[Running -- LLM Column Classification](running.md#llm-column-classification). -### `NSS_INFERENCE_ENDPOINT` +### `NSS_INFERENCE_ENDPOINT` and `NSS_INFERENCE_KEY` -The NIM/OpenAI-compatible endpoint used for PII column classification. Defaults -to `https://integrate.api.nvidia.com/v1` when unset. Override for a custom endpoint: +OpenAI-compatible endpoint and API key for column classification. The endpoint +defaults to `https://integrate.api.nvidia.com/v1` when unset. ```bash export NSS_INFERENCE_ENDPOINT="https://your-llm-inference-endpoint" export NSS_INFERENCE_KEY="your-api-key" # pragma: allowlist secret ``` -When using the CLI or SDK: for column classification to work, set `NSS_INFERENCE_KEY` (and -`NSS_INFERENCE_ENDPOINT` only if you are not using the default URL). +On the CLI, can also use `--inference-api-key` and optionally +`--inference-endpoint-url` instead of exporting these variables. -To disable column classification entirely instead of pointing it at a local -endpoint, use the `replace_pii.globals.classify.enable_classify` config option. -PII classify config is deeply nested -- use YAML or SDK: +To disable column classification entirely, set +`replace_pii.globals.classify.enable_classify: false` in YAML or use the SDK. +See [Configuration Reference -- Replacing PII](configuration.md#replacing-pii). -=== "Config reference" +### `NSS_INFERENCE_MODEL` - ```yaml - replace_pii: - globals: - classify: - enable_classify: false - ``` +Model ID sent to the inference endpoint. Defaults to +`qwen/qwen3-next-80b-a3b-instruct`. Override with `--inference-model-id`. -=== "SDK" +### `NSS_PII_REPLACER_CPU_COUNT` - ```python - from nemo_safe_synthesizer.config.replace_pii import PiiReplacerConfig +Number of CPU worker processes for NER. Override with `--cpu-count`. Defaults +to `max(1, cpu_count - 1)`, capped so each worker handles at least 1,000 +records. - pii_config = PiiReplacerConfig.get_default_config() - pii_config.globals.classify.enable_classify = False +```bash +export NSS_PII_REPLACER_CPU_COUNT=4 +``` - synthesizer = ( - SafeSynthesizer(config) - .with_data_source("data.csv") - .with_replace_pii(config=pii_config) - ) - ``` +--- -### `NSS_INFERENCE_KEY` +## vLLM and attention -API key for the NSS inference endpoint. Required for PII column classification when using the -CLI and SDK (with the default or custom `NSS_INFERENCE_ENDPOINT`). +### `VLLM_CACHE_ROOT` -### `NIM_MODEL_ID` +Directory for vLLM's internal model cache (default `~/.cache/vllm`). -Model ID sent to the NIM endpoint for PII column classification. Defaults to -`qwen/qwen3-next-80b-a3b-instruct`. +```bash +export VLLM_CACHE_ROOT=/shared/cache/vllm +``` -### `SAFE_SYNTHESIZER_CPU_COUNT` +### `VLLM_ATTENTION_BACKEND` -Controls the number of CPU worker processes used for NER (PII replacement). +Override the vLLM attention implementation. Safe Synthesizer sets this from +`generation.attention_backend` when configured; leave unset to use vLLM +auto-detection. ```bash -export SAFE_SYNTHESIZER_CPU_COUNT=4 +export VLLM_ATTENTION_BACKEND=FLASH_ATTN ``` -Defaults to `max(1, cpu_count - 1)` (one CPU left free), further capped so -there are at least 1,000 records per worker. +Common values: `FLASHINFER`, `FLASH_ATTN`, `TORCH_SDPA`, `TRITON_ATTN`, +`FLEX_ATTENTION`. See [Running -- Attention Backends](running.md#attention-backends). -### `NEMO_TELEMETRY_ENABLED` +--- -Controls whether telemetry is sent for train/generate events. +## Telemetry -Defaults to `true`. Set it to `false` to disable telemetry for the current shell: +### `NEMO_TELEMETRY_ENABLED` + +Whether anonymous train/generate telemetry is sent. Defaults to enabled. ```bash export NEMO_TELEMETRY_ENABLED=false ``` -You can also disable telemetry in a Safe Synthesizer config file: +Also disable per run with `--emit_telemetry false` or `emit_telemetry: false` +in YAML. Explicit config/CLI values override the env var. -```yaml -emit_telemetry: false -``` +### `NEMO_TELEMETRY_ENDPOINT` and `NEMO_SESSION_PREFIX` + +Override the telemetry upload endpoint or prefix session IDs. Env-only; no CLI +equivalent. Intended for controlled test environments. --- -## Container Usage +## Containers -When running Safe Synthesizer in a Docker container, these variables are -particularly important: +Common bind-mount targets when running in Docker: -| Variable | Recommended Value | Why | -|----------|-------------------|-----| -| `HF_HOME` | `/workspace/.hf_cache` | Point at a bind-mounted host directory so model downloads persist across container runs | -| `HF_HUB_OFFLINE` | `1` | Set in air-gapped environments after pre-caching models | -| `VLLM_CACHE_ROOT` | `/workspace/.vllm_cache` | Persist vLLM's internal cache if needed | -| `NSS_ARTIFACTS_PATH` | `/workspace/artifacts` | Write artifacts to a mounted volume | -| `NSS_LOG_FORMAT` | `json` | Structured logs for log aggregators; auto-detected in non-TTY containers | -| `NVIDIA_VISIBLE_DEVICES` | `0` or `all` | Select GPUs (set by `--gpus` flag, but can be overridden) | +| Variable | Typical value | Why | +|----------|---------------|-----| +| `HF_HOME` | `/workspace/.hf_cache` | Persist Hub downloads across runs | +| `HF_HUB_OFFLINE` | `1` | Air-gapped runs after pre-caching | +| `VLLM_CACHE_ROOT` | `/workspace/.vllm_cache` | Persist vLLM cache | +| `NSS_ARTIFACTS_PATH` | `/workspace/artifacts` | Write artifacts to a volume | +| `NSS_LOG_FORMAT` | `json` | Structured logs in non-TTY containers | +| `NVIDIA_VISIBLE_DEVICES` | `0` or `all` | GPU selection inside the container | -Example: +See [Docker](docker.md) for mount paths, secrets, GPU flags, and Makefile +shortcuts. -```bash -docker run --gpus all --shm-size=1g \ - -v $(pwd):/workspace \ - -v ~/.cache/huggingface:/workspace/.hf_cache \ - -e HF_HOME=/workspace/.hf_cache \ - -e NSS_ARTIFACTS_PATH=/workspace/artifacts \ - nss-gpu:latest run --config /workspace/config.yaml --data-source /workspace/data.csv -``` +--- -See [Docker](docker.md) for full container setup and Makefile shortcuts. +## Internal and cluster + +Advanced env-only settings without CLI equivalents: + +| Variable | Purpose | +|----------|---------| +| `NSS_OPT_BUCKET` | S3 bucket for optional NER optimization artifacts | +| `NSS_OPT_CACHE_DIR` | Local cache directory for NER optimization downloads | +| `NEMO_JOB_ID` | Cluster job ID attached to multimodal evaluation reports | --- -- [Running Safe Synthesizer](running.md) -- pipeline execution, CLI commands, and artifacts -- [Configuration Reference](configuration.md) -- parameter tables +## Related guides + +- [Running Safe Synthesizer](running.md) -- pipeline execution, CLI commands, offline workflow +- [Configuration Reference](configuration.md) -- synthesis parameter tables and precedence +- [Docker](docker.md) -- container setup, caches, and secrets - [Program Runtime](troubleshooting.md) -- runtime errors and OOM fixes diff --git a/docs/user-guide/running.md b/docs/user-guide/running.md index bcf92d748..2e1594fd2 100644 --- a/docs/user-guide/running.md +++ b/docs/user-guide/running.md @@ -274,7 +274,7 @@ execute in order (`config` → `dataframe` → `metadata` → `advisory`). | Check name | Stage | What it validates | |-------|-------|-------------------| | `gpu.cuda` | config | PyTorch is importable and a CUDA GPU is visible | -| `env.inference_key` | config | `NSS_INFERENCE_KEY` is set when PII classification is enabled (warning only) | +| `env.inference` | config | Inference config for PII classification: `NSS_INFERENCE_KEY` is set, `NSS_INFERENCE_MODEL` is non-empty, and `NSS_INFERENCE_ENDPOINT` is a valid http(s) URL (warnings only) | | `env.hf_model_availability` | config | The pretrained model reference is usable locally or can be fetched from Hugging Face; warns about a missing HF token only when online HF access may be needed | | `dataset.size` | dataframe | Training split meets the hard minimum row count | | `columns.groupby` | dataframe | `group_training_examples_by` column is present and has no nulls | @@ -1230,10 +1230,12 @@ See [`artifacts clean`](#artifacts-clean) in the CLI Commands section for option ## Running in Offline Environments Pre-cache models by running once with internet access, then set -`HF_HUB_OFFLINE=1` in your target environment. For detailed cache setup -and environment variables (`HF_HOME`, `HF_HUB_OFFLINE`, `LOCAL_FILES_ONLY`, -`VLLM_CACHE_ROOT`), see -[Environment Variables -- Hugging Face Cache](environment.md#hugging-face-cache). +`HF_HUB_OFFLINE=1` in your target environment. Export it before launching +`safe-synthesizer` (or pass `--disable-huggingface-remote`) -- huggingface_hub +reads the value once at import time, so setting it after the process starts has +no effect. For detailed cache setup and environment variables (`HF_HOME`, +`HF_HUB_OFFLINE`, `VLLM_CACHE_ROOT`), see +[Environment Variables -- Hugging Face cache and offline](environment.md#hugging-face-cache-and-offline). For offline-specific errors, see [Program Runtime](troubleshooting.md). diff --git a/docs/user-guide/troubleshooting.md b/docs/user-guide/troubleshooting.md index 40639c6fc..ce231f6d7 100644 --- a/docs/user-guide/troubleshooting.md +++ b/docs/user-guide/troubleshooting.md @@ -513,7 +513,9 @@ check of its own. | `torch_missing` | error | `gpu.cuda` | PyTorch not installed; cannot verify GPU availability | | `no_gpu` | error | `gpu.cuda` | No CUDA GPU detected (required for training or generation) | | `low_vram` | warning | `gpu.vram` | Free GPU VRAM may be insufficient | -| `inference_key_missing` | warning | `env.inference_key` | `NSS_INFERENCE_KEY` not set; PII classification degraded | +| `inference_key_missing` | warning | `env.inference` | `NSS_INFERENCE_KEY` not set; PII classification degraded | +| `inference_model_blank` | warning | `env.inference` | `NSS_INFERENCE_MODEL` set but empty; the blank value is ignored and the default model id is used | +| `inference_endpoint_invalid` | error | `env.inference` | `NSS_INFERENCE_ENDPOINT` set but not a valid http(s) URL; classification requests will fail | | `hf_token_missing` | warning | `env.hf_model_availability` | Neither `HF_TOKEN` nor `HUGGING_FACE_HUB_TOKEN` set, and model loading may need online Hugging Face access | | `hf_model_not_cached` | warning/error | `env.hf_model_availability` | Hugging Face model is not present in the local cache; severity is error when HF offline mode is enabled | | `hf_model_cache_incomplete` | warning/error | `env.hf_model_availability` | Cached Hugging Face model snapshot is missing required config, tokenizer, weights, or shards; severity is error when HF offline mode is enabled | @@ -545,7 +547,29 @@ The PII replacer downloads the GLiNER NER model on first use. If the download fails, it raises an exception immediately. Pre-download the model by running PII replacement once in an environment -with internet access, or set `LOCAL_FILES_ONLY=true` after the model is cached. +with internet access. To force offline use after the model is cached, set +`HF_HUB_OFFLINE=1` or pass `--disable-huggingface-remote`. + +### Offline Mode Not Taking Effect + +Symptom: `HF_HUB_OFFLINE=1` (or `--disable-huggingface-remote`) is set, yet the +run still attempts a download, or `--enable-huggingface-remote` does not +re-enable downloads. + +Cause: huggingface_hub reads `HF_HUB_OFFLINE` once, at import time, and caches +it. If the variable is changed after huggingface_hub has been imported in the +process, the change is ignored. + +Fixes: + +- CLI: export `HF_HUB_OFFLINE` before launching `safe-synthesizer`, or use + `--enable-huggingface-remote` / `--disable-huggingface-remote`. The CLI + applies the flag before huggingface_hub loads, so the flag always wins over + an inherited environment value. +- Programmatic / SDK: set `HF_HUB_OFFLINE` before importing + `nemo_safe_synthesizer` (or any library that imports huggingface_hub, such as + `transformers` or `datasets`). Setting it afterward has no effect for that + process. ### NER Processing Timeouts diff --git a/script/slurm/slurm_nss_matrix.sh b/script/slurm/slurm_nss_matrix.sh index 8cd1e1f52..14a5a6f96 100644 --- a/script/slurm/slurm_nss_matrix.sh +++ b/script/slurm/slurm_nss_matrix.sh @@ -129,7 +129,7 @@ echo "[NSS SLURM] nemo-safe-synthesizer version: $(python -c 'from nemo_safe_syn # for column classification export NSS_INFERENCE_ENDPOINT=https://integrate.api.nvidia.com/v1 -export NIM_MODEL_ID=qwen/qwen3-next-80b-a3b-instruct +export NSS_INFERENCE_MODEL=qwen/qwen3-next-80b-a3b-instruct # Extract dataset name for path construction (handles both full paths and simple names) # e.g., "/path/to/adult.csv" -> "adult", "/path/to/data.parquet" -> "data", "adult" -> "adult" diff --git a/src/nemo_safe_synthesizer/cli/run.py b/src/nemo_safe_synthesizer/cli/run.py index a6ffbc14c..17dfc5ef6 100644 --- a/src/nemo_safe_synthesizer/cli/run.py +++ b/src/nemo_safe_synthesizer/cli/run.py @@ -159,6 +159,63 @@ def common_run_options(f: Callable[..., object]) -> Callable[..., object]: "If both env var and CLI option are provided, the CLI option takes precedence.", ) ) + options.append( + click.option( + "--inference-endpoint-url", + type=str, + required=False, + default=None, + help="OpenAI-compatible inference endpoint URL for PII column classification. " + "Can also be set via NSS_INFERENCE_ENDPOINT env var.", + ) + ) + options.append( + click.option( + "--inference-api-key", + type=str, + required=False, + default=None, + help="API key for the inference endpoint used in PII column classification. " + "Can also be set via NSS_INFERENCE_KEY env var.", + ) + ) + options.append( + click.option( + "--inference-model-id", + type=str, + required=False, + default=None, + help="Model ID sent to the inference endpoint for PII column classification. " + "Can also be set via NSS_INFERENCE_MODEL env var. " + "[default: qwen/qwen3-next-80b-a3b-instruct]", + ) + ) + options.append( + click.option( + "--enable-huggingface-remote/--disable-huggingface-remote", + "huggingface_remote", + required=False, + default=None, + help="Allow or block Hugging Face remote downloads for both the base model " + "and GLiNER. --disable-huggingface-remote forces a fully offline run by " + "setting HF_HUB_OFFLINE and TRANSFORMERS_OFFLINE; both must already be " + "cached. Equivalent to setting HF_HUB_OFFLINE in the environment. When " + "neither flag is given, the run inherits HF_HUB_OFFLINE/TRANSFORMERS_OFFLINE " + "from the environment (remote downloads enabled when unset). " + "[default: --enable-huggingface-remote]", + ) + ) + options.append( + click.option( + "--cpu-count", + type=int, + required=False, + default=None, + help="Number of CPU worker processes used for NER (PII replacement). " + "Can also be set via NSS_PII_REPLACER_CPU_COUNT env var. " + "[default: max(1, cpu_count - 1)]", + ) + ) # Apply each option decorator in reverse order (decorators apply bottom-up) for option in reversed(options): f = option(f) @@ -170,6 +227,31 @@ def _parse_run_overrides(kwargs: dict[str, Any]) -> dict[str, Any]: return parse_overrides(kwargs) +# CLISettings fields populated from common_run_options flags. ``synthesis_overrides`` +# is excluded -- it is derived from the leftover pydantic_options kwargs, not bound +# to a single flag. ``observability``/``wandb`` are nested sub-settings with no CLI +# flag, so they never appear in command kwargs. +_CLI_SETTINGS_FIELDS: frozenset[str] = frozenset(CLISettings.model_fields) - {"synthesis_overrides"} + + +def _settings_from_run_kwargs(kwargs: dict[str, Any]) -> CLISettings: + """Build ``CLISettings`` from a run command's ``**kwargs``. + + ``common_run_options`` binds each infrastructure flag to a kwarg whose name + matches a ``CLISettings`` field; those are pulled out here. Everything left + (the ``pydantic_options`` ``--section__field`` options) becomes synthesis + overrides. This keeps the three run commands from re-listing the shared flag + set in both their signature and their settings construction -- adding a flag + now means editing ``common_run_options`` and ``CLISettings`` only. + + ``kwargs`` is mutated: matched settings keys are popped before the remainder + is parsed into overrides. + """ + settings_kwargs = {name: kwargs.pop(name) for name in _CLI_SETTINGS_FIELDS if name in kwargs} + settings_kwargs["synthesis_overrides"] = _parse_run_overrides(kwargs) + return CLISettings.from_cli_kwargs(**settings_kwargs) + + def _set_cli_deployment_type_default() -> None: """Default telemetry deployment type for CLI commands without overriding Slurm or explicit settings.""" os.environ.setdefault("NEMO_DEPLOYMENT_TYPE", DeploymentTypeEnum.CLI.value) @@ -311,20 +393,8 @@ def _build_validate_render_context( ) def run( ctx: click.Context, - config_path: PathT | None, - data_source: str | None, - artifact_path: PathT | None, - run_path: PathT | None, - output_file: PathT | None, - log_file: PathT | None, - log_color: bool | None, - log_format: str | None, - verbose: int = 0, - wandb_mode: str | None = None, - wandb_project: str | None = None, - dataset_registry: str | None = None, validate: bool = False, - **kwargs: object, + **kwargs: Any, ) -> None: """Run the Safe Synthesizer end-to-end pipeline. @@ -337,21 +407,7 @@ def run( _set_cli_deployment_type_default() - settings = CLISettings.from_cli_kwargs( - data_source=data_source, - config_path=config_path, - artifact_path=artifact_path, - run_path=run_path, - output_file=output_file, - log_file=log_file, - log_color=log_color, - log_format=log_format, - verbose=verbose, - wandb_mode=wandb_mode, - wandb_project=wandb_project, - synthesis_overrides=_parse_run_overrides(kwargs), - dataset_registry=dataset_registry, - ) + settings = _settings_from_run_kwargs(kwargs) if validate: os.environ["NSS_PHASE"] = "process_data" @@ -407,20 +463,8 @@ def run( help="Run pre-flight validation only, then exit without training or generating.", ) def run_train( - config_path: PathT, - data_source: str | None, - artifact_path: PathT | None, - run_path: PathT | None, - output_file: PathT | None, - log_format: str | None, - log_color: bool | None, - log_file: PathT | None, - verbose: int, - wandb_mode: str | None = None, - wandb_project: str | None = None, - dataset_registry: str | None = None, validate: bool = False, - **kwargs: object, + **kwargs: Any, ) -> None: """Run the training stage only. @@ -429,21 +473,7 @@ def run_train( """ _set_cli_deployment_type_default() - settings = CLISettings.from_cli_kwargs( - data_source=data_source, - config_path=config_path, - artifact_path=artifact_path, - run_path=run_path, - output_file=output_file, - log_file=log_file, - log_color=log_color, - log_format=log_format, - verbose=verbose, - wandb_mode=wandb_mode, - wandb_project=wandb_project, - synthesis_overrides=_parse_run_overrides(kwargs), - dataset_registry=dataset_registry, - ) + settings = _settings_from_run_kwargs(kwargs) if validate: os.environ["NSS_PHASE"] = "process_data" @@ -498,21 +528,9 @@ def run_train( ) @pydantic_options(SafeSynthesizerParameters, field_separator=CLI_NESTED_FIELD_SEPARATOR) def run_generate( - config_path: PathT, - data_source: str | None, - run_path: PathT | None, - artifact_path: PathT | None, - output_file: PathT | None, - log_format: str | None, - log_color: bool | None, - log_file: PathT | None, - verbose: int, - wandb_mode: str | None = None, - wandb_project: str | None = None, auto_discover_adapter: bool = False, wandb_resume_job_id: str | None = None, - dataset_registry: str | None = None, - **kwargs: object, + **kwargs: Any, ) -> None: """Run the generation stage only. @@ -526,21 +544,7 @@ def run_generate( _set_cli_deployment_type_default() # Create unified settings from CLI kwargs - settings = CLISettings.from_cli_kwargs( - data_source=data_source, - config_path=config_path, - artifact_path=artifact_path, - run_path=run_path, - output_file=output_file, - log_file=log_file, - log_color=log_color, - log_format=log_format, - verbose=verbose, - wandb_mode=wandb_mode, - wandb_project=wandb_project, - synthesis_overrides=_parse_run_overrides(kwargs), - dataset_registry=dataset_registry, - ) + settings = _settings_from_run_kwargs(kwargs) os.environ["NSS_PHASE"] = "generate" # Generation always resumes from an existing workdir with a trained model diff --git a/src/nemo_safe_synthesizer/cli/settings.py b/src/nemo_safe_synthesizer/cli/settings.py index e088bef89..bdcb5fb36 100644 --- a/src/nemo_safe_synthesizer/cli/settings.py +++ b/src/nemo_safe_synthesizer/cli/settings.py @@ -118,9 +118,10 @@ class CLISettings(BaseSettings): log_color: bool | None = Field( default=None, + validation_alias=AliasChoices("log_color", "NSS_LOG_COLOR"), description="Whether to colorize console output", ) - """Whether to colorize console output.""" + """Whether to colorize console output (env variable: ``NSS_LOG_COLOR``).""" log_file: str | None = Field( default=None, @@ -164,6 +165,51 @@ class CLISettings(BaseSettings): ) """URL or path to a dataset registry YAML file (env: ``NSS_DATASET_REGISTRY``).""" + inference_endpoint_url: str | None = Field( + default=None, + validation_alias=AliasChoices("inference_endpoint_url", "NSS_INFERENCE_ENDPOINT"), + description="OpenAI-compatible inference endpoint URL for PII column classification", + ) + """OpenAI-compatible inference endpoint URL for PII column classification + (env: ``NSS_INFERENCE_ENDPOINT``).""" + + inference_api_key: str | None = Field( + default=None, + validation_alias=AliasChoices("inference_api_key", "NSS_INFERENCE_KEY"), + description="API key for the inference endpoint used in PII column classification", + ) + """API key for the inference endpoint used in PII column classification + (env: ``NSS_INFERENCE_KEY``).""" + + inference_model_id: str | None = Field( + default=None, + validation_alias=AliasChoices("inference_model_id", "NSS_INFERENCE_MODEL"), + description="Model ID sent to the inference endpoint for PII column classification", + ) + """Model ID sent to the inference endpoint for PII column classification + (env: ``NSS_INFERENCE_MODEL``).""" + + huggingface_remote: bool | None = Field( + default=None, + validation_alias=AliasChoices("huggingface_remote"), + description="Whether to allow Hugging Face remote downloads (base model and GLiNER)", + ) + """Whether to allow Hugging Face remote downloads for the base model and GLiNER. + + ``None`` leaves the environment untouched. ``True`` / ``False`` is propagated + to the standard ``HF_HUB_OFFLINE`` and ``TRANSFORMERS_OFFLINE`` variables (the + canonical env switch) by ``_propagate_runtime_settings_to_env``; there is no + separate NSS env var.""" + + cpu_count: int | None = Field( + default=None, + ge=1, + validation_alias=AliasChoices("cpu_count", "NSS_PII_REPLACER_CPU_COUNT"), + description="Number of CPU worker processes used for NER (PII replacement)", + ) + """Number of CPU worker processes used for NER (PII replacement) + (env: ``NSS_PII_REPLACER_CPU_COUNT``).""" + @field_validator("wandb_mode", mode="before") @classmethod def validate_wandb_mode(cls, v: str | WandbMode | None) -> WandbMode | None: diff --git a/src/nemo_safe_synthesizer/cli/utils.py b/src/nemo_safe_synthesizer/cli/utils.py index a3ec2ba93..fad664d04 100644 --- a/src/nemo_safe_synthesizer/cli/utils.py +++ b/src/nemo_safe_synthesizer/cli/utils.py @@ -230,6 +230,12 @@ def common_setup( Tuple of (logger, config, dataframe, workdir). For generate-only runs with cached datasets, dataframe may be None (loaded from cached files by SafeSynthesizer). """ + # 0. Propagate CLI-resolved runtime settings back to os.environ. This must + # run before any deferred pii_replacer imports so that module-level reads + # of NSS_INFERENCE_*, HF_HUB_OFFLINE/TRANSFORMERS_OFFLINE, and + # NSS_PII_REPLACER_CPU_COUNT see the CLI-overridden values. + _propagate_runtime_settings_to_env(settings) + # 1. Create workdir FIRST - this establishes all artifact paths workdir = _create_workdir( settings.artifact_path, @@ -314,6 +320,51 @@ def _set_wandb_env_vars( os.environ["WANDB_RUN_NAME"] = wandb_run_name +def _propagate_runtime_settings_to_env(settings: "CLISettings") -> None: + """Materialize CLI-resolved runtime settings back to ``os.environ``. + + The downstream readers for these settings live deep in ``pii_replacer`` + (NER, GLiNER, column classification) and historically read directly from + the process environment. Rather than thread a ``CLISettings`` handle + through every callsite, we propagate the resolved values back to + ``os.environ`` here so that CLI flag precedence -- which ``CLISettings`` + handles via ``from_cli_kwargs`` -- carries through to those readers + unchanged. + + ``CLISettings`` values are already env-aware (via ``AliasChoices``); when + no CLI flag is provided, the field carries the env var's existing value + and writing it back is a no-op. When a CLI flag overrides the env var, + this overwrites ``os.environ`` so the deferred imports in the runtime + pipeline see the CLI value. + + ``huggingface_remote`` is the exception: it has no NSS env var and instead + maps to the standard Hugging Face offline switches (``HF_HUB_OFFLINE`` and + ``TRANSFORMERS_OFFLINE``). ``--disable-huggingface-remote`` sets them to + ``1``; ``--enable-huggingface-remote`` sets them to ``0`` (overriding any + inherited offline env). + + ``huggingface_hub`` caches ``HF_HUB_OFFLINE`` at import time, so this write + is only effective if it runs before the first ``huggingface_hub`` import. + The CLI import chain is kept hub-free for exactly this reason -- + ``telemetry`` defers its ``huggingface_hub`` import (see + ``sanitize_model_for_telemetry``) -- so ``huggingface_hub`` first loads + during the pipeline, after this propagation. ``tests/cli/test_cli_import`` + guards the hub-free import invariant. + """ + if settings.inference_endpoint_url is not None: + os.environ["NSS_INFERENCE_ENDPOINT"] = settings.inference_endpoint_url + if settings.inference_api_key is not None: + os.environ["NSS_INFERENCE_KEY"] = settings.inference_api_key + if settings.inference_model_id is not None: + os.environ["NSS_INFERENCE_MODEL"] = settings.inference_model_id + if settings.huggingface_remote is not None: + offline = "0" if settings.huggingface_remote else "1" + os.environ["HF_HUB_OFFLINE"] = offline + os.environ["TRANSFORMERS_OFFLINE"] = offline + if settings.cpu_count is not None: + os.environ["NSS_PII_REPLACER_CPU_COUNT"] = str(settings.cpu_count) + + def _initialize_logging_for_cli_from_settings( settings: "CLISettings", workdir: Workdir, diff --git a/src/nemo_safe_synthesizer/pii_replacer/data_editor/detect.py b/src/nemo_safe_synthesizer/pii_replacer/data_editor/detect.py index 6981d6cd0..2bba807ee 100644 --- a/src/nemo_safe_synthesizer/pii_replacer/data_editor/detect.py +++ b/src/nemo_safe_synthesizer/pii_replacer/data_editor/detect.py @@ -21,6 +21,7 @@ from pydantic import ConfigDict, TypeAdapter, ValidationError from ...observability import get_logger +from ...utils import hf_offline_enabled from ..ner import ner_mp from ..ner.factory import LabelSetPredictorFilter, NERFactory from ..ner.ner import NERPrediction @@ -32,12 +33,10 @@ class DefaultLLMConfig: """Default settings for the LLM used in column classification. - All attributes are class-level. Used by ``classify_columns`` when calling the - inference API for column-type classification. + Used by ``classify_columns`` when calling the inference API for column-type + classification. Attributes: - CONFIG_ID: Model identifier for the LLM. From env ``NIM_MODEL_ID``, or - ``qwen/qwen3-next-80b-a3b-instruct`` if unset. SYSTEM_PROMPT: System message describing the column-type annotation task sent to the LLM. MAX_OUTPUT_TOKENS: Maximum number of tokens allowed in the LLM response @@ -46,11 +45,25 @@ class DefaultLLMConfig: Lower values give more deterministic output. """ - CONFIG_ID = os.environ.get("NIM_MODEL_ID", "qwen/qwen3-next-80b-a3b-instruct") + DEFAULT_CONFIG_ID = "qwen/qwen3-next-80b-a3b-instruct" SYSTEM_PROMPT = "You are a helpful AI that annotates columns in datasets with their respective types. " MAX_OUTPUT_TOKENS = 2048 TEMPERATURE = 0.2 + @classmethod + def config_id(cls) -> str: + """Model identifier for the LLM, read from env at call-time. + + Reads ``NSS_INFERENCE_MODEL`` on each call (falling back to + ``DEFAULT_CONFIG_ID``) so a value set after this module is imported still + takes effect, matching the call-time env handling used elsewhere in this + module. A blank or whitespace-only value is treated as unset, so it falls + back to ``DEFAULT_CONFIG_ID`` rather than sending an empty model id to the + inference API. + """ + model = os.environ.get("NSS_INFERENCE_MODEL", "").strip() + return model or cls.DEFAULT_CONFIG_ID + DEFAULT_ENTITIES: set[str] = { "name", @@ -249,7 +262,7 @@ def classify_columns( llm_start = timer() response = client.chat.completions.create( - model=DefaultLLMConfig.CONFIG_ID, + model=DefaultLLMConfig.config_id(), messages=[ {"role": "system", "content": DefaultLLMConfig.SYSTEM_PROMPT}, {"role": "user", "content": formatted_prompt}, @@ -575,7 +588,7 @@ def get_entity_extractor( extractor._model = GLiNER.from_pretrained( clsfy_cfg.gliner_model, map_location=map_location, - local_files_only=os.environ.get("LOCAL_FILES_ONLY") in ["true", "True"], + local_files_only=hf_offline_enabled(), ) entity_types = DEFAULT_ENTITIES if clsfy_cfg.ner_entities: diff --git a/src/nemo_safe_synthesizer/pii_replacer/nemo_pii.py b/src/nemo_safe_synthesizer/pii_replacer/nemo_pii.py index b5fa1aef8..9710bcc16 100644 --- a/src/nemo_safe_synthesizer/pii_replacer/nemo_pii.py +++ b/src/nemo_safe_synthesizer/pii_replacer/nemo_pii.py @@ -103,8 +103,11 @@ def build_entity_extractor(clsfy_cfg: ClassifyConfig) -> EntityExtractor: def _get_classify_endpoint_url() -> str: """Resolve the NIM/OpenAI-compatible base URL for PII column classification. - If ``NSS_INFERENCE_ENDPOINT`` is present in the environment, that value is used. - If the variable is unset, uses ``DEFAULT_NSS_INFERENCE_ENDPOINT`` from ``defaults``. + If ``NSS_INFERENCE_ENDPOINT`` holds a non-blank value, that value is used. + If the variable is unset or blank, uses ``DEFAULT_NSS_INFERENCE_ENDPOINT`` + from ``defaults``. A blank value is treated as unset so it never reaches the + OpenAI client as an empty ``base_url`` (which would fail every request), + matching the preflight ``env.inference`` check that ignores a blank endpoint. Note: Emits an INFO log indicating whether the default or configured URL applies. @@ -112,8 +115,8 @@ def _get_classify_endpoint_url() -> str: Returns: inference endpoint for PII column classification. """ - configured = os.environ.get("NSS_INFERENCE_ENDPOINT") - if configured is None: + configured = os.environ.get("NSS_INFERENCE_ENDPOINT", "").strip() + if not configured: url = DEFAULT_NSS_INFERENCE_ENDPOINT logging.info( "PII column classification will call the default NVIDIA inference API at %s. " diff --git a/src/nemo_safe_synthesizer/pii_replacer/ner/factory.py b/src/nemo_safe_synthesizer/pii_replacer/ner/factory.py index 7b4810cec..e7168a648 100644 --- a/src/nemo_safe_synthesizer/pii_replacer/ner/factory.py +++ b/src/nemo_safe_synthesizer/pii_replacer/ner/factory.py @@ -147,7 +147,7 @@ def _create_parallel_ner( # system tries to run another process and not enough memory workers will # start getting killed. So I'm setting an env here that allows an override # of the num CPUs when we need to explicitly control it. - num_proc_env = os.getenv("SAFE_SYNTHESIZER_CPU_COUNT") + num_proc_env = os.getenv("NSS_PII_REPLACER_CPU_COUNT") if num_proc_env: try: num_proc = int(num_proc_env) diff --git a/src/nemo_safe_synthesizer/preflight/__init__.py b/src/nemo_safe_synthesizer/preflight/__init__.py index e1f639f51..576553a72 100644 --- a/src/nemo_safe_synthesizer/preflight/__init__.py +++ b/src/nemo_safe_synthesizer/preflight/__init__.py @@ -19,7 +19,7 @@ DatasetSizeCheck, GroupbyColumnCheck, HFModelAvailabilityCheck, - InferenceKeyCheck, + InferenceModelCheck, OrderbyColumnCheck, OversamplingCheck, PseudoColumnCheck, @@ -61,7 +61,7 @@ "DatasetSizeCheck", "GroupbyColumnCheck", "HFModelAvailabilityCheck", - "InferenceKeyCheck", + "InferenceModelCheck", "IssueCollector", "MetadataCheck", "MetadataView", diff --git a/src/nemo_safe_synthesizer/preflight/checks/__init__.py b/src/nemo_safe_synthesizer/preflight/checks/__init__.py index 1e02b6589..4629c31b8 100644 --- a/src/nemo_safe_synthesizer/preflight/checks/__init__.py +++ b/src/nemo_safe_synthesizer/preflight/checks/__init__.py @@ -27,7 +27,7 @@ from .environment import ( CUDAAvailabilityCheck, HFModelAvailabilityCheck, - InferenceKeyCheck, + InferenceModelCheck, VRAMHeadroomCheck, ) from .metadata import TokenBudgetCheck @@ -39,7 +39,7 @@ "DatasetSizeCheck", "GroupbyColumnCheck", "HFModelAvailabilityCheck", - "InferenceKeyCheck", + "InferenceModelCheck", "OrderbyColumnCheck", "OversamplingCheck", "PseudoColumnCheck", @@ -57,7 +57,7 @@ _CORE_CHECKS: tuple[PreflightCheck, ...] = ( # CONFIG CUDAAvailabilityCheck(), - InferenceKeyCheck(), + InferenceModelCheck(), HFModelAvailabilityCheck(), # DATAFRAME DatasetSizeCheck(), diff --git a/src/nemo_safe_synthesizer/preflight/checks/environment.py b/src/nemo_safe_synthesizer/preflight/checks/environment.py index b9fa7509a..91daffbab 100644 --- a/src/nemo_safe_synthesizer/preflight/checks/environment.py +++ b/src/nemo_safe_synthesizer/preflight/checks/environment.py @@ -8,10 +8,11 @@ import os from pathlib import Path from typing import TYPE_CHECKING, Literal +from urllib.parse import urlparse -from ...config.replace_pii import has_inference_key from ...llm.utils import ModelRef from ...observability import get_logger +from ...utils import hf_offline_enabled from ..base import ConfigCheck, IssueCollector, MetadataCheck from ..helpers import require_import from ..types import ConfigView, MetadataView @@ -26,7 +27,7 @@ __all__ = [ "CUDAAvailabilityCheck", "HFModelAvailabilityCheck", - "InferenceKeyCheck", + "InferenceModelCheck", "VRAMHeadroomCheck", "bytes_per_base_weight", "estimate_base_model_params", @@ -325,33 +326,74 @@ def check(self, ctx: MetadataView, collector: IssueCollector) -> None: ) -class InferenceKeyCheck(ConfigCheck): - """Check NSS_INFERENCE_KEY environment variable.""" +def _is_blank(value: str | None) -> bool: + """Whether ``value`` is set but contains only whitespace (an empty override).""" + return value is not None and not value.strip() - name = "env.inference_key" - label = "Inference key" + +def _is_valid_http_url(value: str | None) -> bool: + """Whether ``value`` parses as an ``http(s)`` URL with a network location.""" + if value is None: + return False + parsed = urlparse(value.strip()) + return parsed.scheme in ("http", "https") and bool(parsed.netloc) + + +class InferenceModelCheck(ConfigCheck): + """Validate the inference configuration used for PII column classification. + + When classification is enabled, the runtime calls an OpenAI-compatible + inference endpoint configured by ``NSS_INFERENCE_KEY``, + ``NSS_INFERENCE_MODEL``, and ``NSS_INFERENCE_ENDPOINT`` (set directly or via + the matching CLI flags, which are propagated to the environment before + preflight runs). This check reads those env vars -- not ``config`` -- because + the inference settings live in ``CLISettings``/the environment rather than in + ``SafeSynthesizerParameters``. + + The body uses a single-dispatch ``match`` over ``(model, key, endpoint)``, + so at most one finding is emitted per run -- the highest-priority problem. + Priority order: invalid endpoint, then missing key, then blank model id. The + invalid endpoint is an error (a non-http(s) endpoint cannot succeed, so the + run must not pass ``--validate``); the key and model findings are warnings + (classification degrades or falls back rather than failing the run). The + error is checked first so a lower-severity warning never masks it. + """ + + name = "env.inference" + label = "Inference configuration" category = "environment" def check(self, ctx: ConfigView, collector: IssueCollector) -> None: config = ctx.config - if config.replace_pii is not None and config.replace_pii.globals.classify.enable_classify is not False: - if not has_inference_key(): + if config.replace_pii is None or config.replace_pii.globals.classify.enable_classify is False: + return + + model = os.environ.get("NSS_INFERENCE_MODEL") + key = os.environ.get("NSS_INFERENCE_KEY") + endpoint = os.environ.get("NSS_INFERENCE_ENDPOINT") + + # Single-dispatch: the first matching case wins, so cases are ordered by + # severity then priority. The invalid endpoint is a hard error and is + # checked first so it is never masked by the missing-key or blank-model + # warnings. + match model, key, endpoint: + case _, _, e if e is not None and e.strip() and not _is_valid_http_url(e): + collector.error( + "inference_endpoint_invalid", + f"NSS_INFERENCE_ENDPOINT '{e}' is not a valid http(s) URL. " + "PII column classification requests will fail.", + ) + case _, k, _ if not (k or "").strip(): collector.warning( "inference_key_missing", "NSS_INFERENCE_KEY is not set. PII column classification will run in degraded mode.", ) - - -_OFFLINE_ENV_VARS = ("HF_HUB_OFFLINE", "TRANSFORMERS_OFFLINE") - - -def _env_flag_enabled(name: str) -> bool: - value = os.environ.get(name) - return value is not None and value.casefold() in {"1", "true", "yes", "on"} - - -def _hf_offline_enabled() -> bool: - return any(_env_flag_enabled(name) for name in _OFFLINE_ENV_VARS) + case m, _, _ if _is_blank(m): + collector.warning( + "inference_model_blank", + "NSS_INFERENCE_MODEL is set but empty. The blank value is ignored and the " + "default model id is used. Set a non-empty model id to override the default.", + ) def _has_hf_token() -> bool: @@ -410,7 +452,7 @@ def check(self, ctx: ConfigView, collector: IssueCollector) -> None: message = ( f"Cached Hugging Face model '{model_ref.repo_id}' at '{snapshot_path}' is missing {', '.join(missing)}." ) - if _hf_offline_enabled(): + if hf_offline_enabled(): collector.error( "hf_model_cache_incomplete", f"{message} Offline Hugging Face mode is enabled, so model loading will fail.", @@ -445,7 +487,7 @@ def _report_missing_cache(model_ref: ModelRef, collector: IssueCollector) -> Non message = ( f"Hugging Face model '{model_ref.repo_id}' is not present in the local cache at '{model_ref.cache_root}'." ) - if _hf_offline_enabled(): + if hf_offline_enabled(): collector.error( "hf_model_not_cached", f"{message} Offline Hugging Face mode is enabled, so model loading will fail.", @@ -470,7 +512,7 @@ def _report_missing_remote_code(model_ref: ModelRef, model_path: Path, collector f"Trusted Hugging Face model '{model_ref.repo_id}' at '{model_path}' references remote code " f"that is not cached locally: {', '.join(missing)}." ) - if _hf_offline_enabled(): + if hf_offline_enabled(): collector.error( "hf_remote_code_not_cached", f"{message} Offline Hugging Face mode is enabled, so Transformers cannot fetch it.", diff --git a/src/nemo_safe_synthesizer/telemetry.py b/src/nemo_safe_synthesizer/telemetry.py index d819af430..f8b055fd5 100644 --- a/src/nemo_safe_synthesizer/telemetry.py +++ b/src/nemo_safe_synthesizer/telemetry.py @@ -25,7 +25,6 @@ from typing import TYPE_CHECKING, Any, ClassVar from urllib.parse import urlsplit, urlunsplit -from huggingface_hub.utils import HFValidationError, validate_repo_id from pydantic import BaseModel, Field from .observability import get_logger @@ -112,6 +111,13 @@ def sanitize_model_for_telemetry(model: str | None) -> str: if Path(model).expanduser().exists(): return LOCAL_MODEL_LABEL + # Imported lazily: huggingface_hub caches HF_HUB_OFFLINE at import time, and + # this module loads during CLI startup (cli.cli -> cli.run -> telemetry), + # before common_setup propagates the --(enable|disable)-huggingface-remote + # flag. Deferring the import keeps cli.cli hub-free so that propagation runs + # first. See cli.utils._propagate_runtime_settings_to_env. + from huggingface_hub.utils import HFValidationError, validate_repo_id + try: validate_repo_id(model) except HFValidationError: diff --git a/src/nemo_safe_synthesizer/utils.py b/src/nemo_safe_synthesizer/utils.py index 2d62f81b9..93ac5bac8 100644 --- a/src/nemo_safe_synthesizer/utils.py +++ b/src/nemo_safe_synthesizer/utils.py @@ -15,18 +15,49 @@ import time from collections.abc import Callable, Generator, Iterable from pathlib import Path -from typing import Any, Protocol +from typing import TYPE_CHECKING, Any, Protocol import numpy as np import pandas as pd -from datasets import Dataset from pandas import DataFrame from .data_processing.stats import Statistics from .observability import get_logger +if TYPE_CHECKING: + # Annotation-only. Imported here to keep the CLI import chain free of + # `datasets` (which pulls huggingface_hub, caching HF_HUB_OFFLINE at import + # time). See cli.utils._propagate_runtime_settings_to_env. + from datasets import Dataset + logger = get_logger(__name__) +_TRUTHY_ENV_VALUES = frozenset({"1", "true", "yes", "on"}) + +_HF_OFFLINE_ENV_VARS = ("HF_HUB_OFFLINE", "TRANSFORMERS_OFFLINE") + + +def env_flag_is_true(name: str, *, default: bool = False) -> bool: + """Return whether ``name`` is set to a truthy env value. + + Accepts common boolean spellings used across NSS and pydantic-settings: + ``1``, ``true``, ``yes``, and ``on`` (case-insensitive). + """ + raw = os.environ.get(name) + if raw is None: + return default + return raw.strip().lower() in _TRUTHY_ENV_VALUES + + +def hf_offline_enabled() -> bool: + """Return whether Hugging Face offline mode is enabled. + + True when ``HF_HUB_OFFLINE`` or ``TRANSFORMERS_OFFLINE`` is set to a truthy + value. huggingface_hub honors these globally, so when enabled both the base + model and GLiNER skip network downloads and resolve from the local cache. + """ + return any(env_flag_is_true(name) for name in _HF_OFFLINE_ENV_VARS) + def _get_num_items_pattern(min_items: int | None, max_items: int | None, whitespace_pattern: str) -> str | None: """Return a regex quantifier for JSON array/object item counts. diff --git a/tests/cli/test_cli_import.py b/tests/cli/test_cli_import.py new file mode 100644 index 000000000..d9486a7b8 --- /dev/null +++ b/tests/cli/test_cli_import.py @@ -0,0 +1,36 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Guard the hub-free CLI import invariant. + +``huggingface_hub`` caches ``HF_HUB_OFFLINE`` at import time. The CLI propagates +the ``--(enable|disable)-huggingface-remote`` flag to that env var inside +``common_setup``; for the propagation to take effect, ``huggingface_hub`` must +not be imported during the ``cli.cli`` import chain. Run in a subprocess because +``sys.modules`` is process-global and other tests import ``huggingface_hub``. +""" + +from __future__ import annotations + +import subprocess +import sys + + +def test_importing_cli_does_not_import_huggingface_hub(): + code = ( + "import sys;" + "import nemo_safe_synthesizer.cli.cli;" + "loaded = 'huggingface_hub' in sys.modules;" + "print('LOADED' if loaded else 'CLEAN');" + "sys.exit(1 if loaded else 0)" + ) + result = subprocess.run( + [sys.executable, "-c", code], + capture_output=True, + text=True, + ) + assert result.returncode == 0, ( + "Importing nemo_safe_synthesizer.cli.cli pulled in huggingface_hub, which " + "caches HF_HUB_OFFLINE at import time and breaks --(enable|disable)-" + f"huggingface-remote propagation.\nstdout: {result.stdout}\nstderr: {result.stderr}" + ) diff --git a/tests/cli/test_run.py b/tests/cli/test_run.py index 6fad15b6f..4978d1d16 100644 --- a/tests/cli/test_run.py +++ b/tests/cli/test_run.py @@ -304,6 +304,19 @@ def test_run_help_shows_run_path_option(self, cli_runner: CliRunner): assert "--run-path" in result.output assert "Explicit path for this run" in result.output + def test_run_help_shows_runtime_settings_options(self, cli_runner: CliRunner): + """Verify runtime PII/NER settings appear in run command help.""" + result = cli_runner.invoke(run, ["--help"]) + + assert result.exit_code == 0 + assert "--inference-endpoint-url" in result.output + assert "--inference-api-key" in result.output + assert "--inference-model-id" in result.output + assert "--disable-huggingface-remote" in result.output + assert "--cpu-count" in result.output + assert "NSS_INFERENCE_ENDPOINT" in result.output + assert "NSS_INFERENCE_KEY" in result.output + def test_run_with_artifact_path_only( self, cli_runner: CliRunner, @@ -1006,3 +1019,23 @@ def test_generate_with_nonexistent_run_path_exits_nonzero( ) assert result.exit_code != 0 + + +def test_common_run_options_map_to_settings_fields() -> None: + """Every shared run flag must be backed by a CLISettings field. + + ``_settings_from_run_kwargs`` splits a command's kwargs by matching names + against ``CLISettings.model_fields``; anything unmatched is routed to + synthesis overrides. A shared flag whose name is not a settings field would + therefore be silently misrouted instead of populating settings. + """ + from nemo_safe_synthesizer.cli.run import common_run_options + + def _target(**kwargs: object) -> None: ... + + decorated = common_run_options(_target) + option_names = {param.name for param in getattr(decorated, "__click_params__", [])} + assert option_names, "common_run_options registered no Click options" + + unmapped = option_names - set(CLISettings.model_fields) + assert not unmapped, f"common_run_options flags not backed by CLISettings fields: {sorted(unmapped)}" diff --git a/tests/cli/test_settings.py b/tests/cli/test_settings.py index ca8151ba6..24b60badf 100644 --- a/tests/cli/test_settings.py +++ b/tests/cli/test_settings.py @@ -5,6 +5,9 @@ from __future__ import annotations +import pytest +from pydantic import ValidationError + from nemo_safe_synthesizer.cli.settings import CLISettings from nemo_safe_synthesizer.cli.wandb_setup import WandbMode @@ -217,6 +220,70 @@ def test_dataset_registry_from_cli(self, monkeypatch): settings = CLISettings.from_cli_kwargs(dataset_registry="path/to/registry.yaml") assert settings.dataset_registry == "path/to/registry.yaml" + def test_inference_endpoint_url_from_nss_inference_env(self, monkeypatch): + """NSS_INFERENCE_ENDPOINT loads into inference_endpoint_url.""" + monkeypatch.setenv("NSS_INFERENCE_ENDPOINT", "https://custom.example/v1") + settings = CLISettings() + assert settings.inference_endpoint_url == "https://custom.example/v1" + + def test_inference_api_key_from_nss_inference_env(self, monkeypatch): + """NSS_INFERENCE_KEY loads into inference_api_key.""" + monkeypatch.setenv("NSS_INFERENCE_KEY", "token-from-env") + settings = CLISettings() + assert settings.inference_api_key == "token-from-env" # pragma: allowlist secret + + def test_inference_endpoint_url_cli_overrides_env(self, monkeypatch): + """CLI --inference-endpoint-url takes precedence over NSS_INFERENCE_ENDPOINT.""" + monkeypatch.setenv("NSS_INFERENCE_ENDPOINT", "https://env.example/v1") + settings = CLISettings.from_cli_kwargs(inference_endpoint_url="https://cli.example/v1") + assert settings.inference_endpoint_url == "https://cli.example/v1" + + def test_inference_api_key_cli_overrides_env(self, monkeypatch): + """CLI --inference-api-key takes precedence over NSS_INFERENCE_KEY.""" + monkeypatch.setenv("NSS_INFERENCE_KEY", "token-from-env") + settings = CLISettings.from_cli_kwargs(inference_api_key="token-from-cli") # pragma: allowlist secret + assert settings.inference_api_key == "token-from-cli" # pragma: allowlist secret + + def test_log_color_from_nss_log_color_env(self, monkeypatch): + """NSS_LOG_COLOR loads into CLISettings.log_color.""" + monkeypatch.setenv("NSS_LOG_COLOR", "false") + settings = CLISettings() + assert settings.log_color is False + assert settings.effective_log_color is False + + def test_log_color_cli_overrides_nss_log_color_env(self, monkeypatch): + """CLI --log-color takes precedence over NSS_LOG_COLOR.""" + monkeypatch.setenv("NSS_LOG_COLOR", "false") + settings = CLISettings.from_cli_kwargs(log_color=True) + assert settings.effective_log_color is True + + def test_runtime_settings_from_env(self, monkeypatch): + """Remaining runtime settings load from their documented env vars.""" + monkeypatch.setenv("NSS_INFERENCE_MODEL", "custom/model") + monkeypatch.setenv("NSS_PII_REPLACER_CPU_COUNT", "4") + + settings = CLISettings() + assert settings.inference_model_id == "custom/model" + assert settings.cpu_count == 4 + + def test_huggingface_remote_is_cli_only(self, monkeypatch): + """huggingface_remote is set via the CLI flag, not a parallel NSS env var.""" + settings = CLISettings.from_cli_kwargs(huggingface_remote=False) + assert settings.huggingface_remote is False + + @pytest.mark.parametrize("bad_value", ["0", "-1"]) + def test_cpu_count_rejects_non_positive(self, monkeypatch, bad_value): + """cpu_count must be >= 1; 0 or negative fails fast at parse time.""" + monkeypatch.setenv("NSS_PII_REPLACER_CPU_COUNT", bad_value) + with pytest.raises(ValidationError): + CLISettings() + + @pytest.mark.parametrize("bad_value", [0, -1]) + def test_cpu_count_rejects_non_positive_from_cli(self, bad_value): + """A non-positive --cpu-count is rejected when passed via CLI kwargs.""" + with pytest.raises(ValidationError): + CLISettings.from_cli_kwargs(cpu_count=bad_value) + class TestCLISettingsIntegration: """Integration tests for CLISettings with env vars.""" diff --git a/tests/cli/test_utils.py b/tests/cli/test_utils.py index 27f518a36..3a694f67e 100644 --- a/tests/cli/test_utils.py +++ b/tests/cli/test_utils.py @@ -5,6 +5,7 @@ from __future__ import annotations +import os from pathlib import Path from unittest.mock import MagicMock, patch @@ -12,7 +13,7 @@ import pytest from nemo_safe_synthesizer.cli.settings import CLISettings -from nemo_safe_synthesizer.cli.utils import common_setup +from nemo_safe_synthesizer.cli.utils import _propagate_runtime_settings_to_env, common_setup @pytest.fixture @@ -326,6 +327,76 @@ def test_apply_cli_overrides_without_registry( assert config.generation.temperature == 0.7 +class TestPropagateRuntimeSettingsToEnv: + """Tests for materializing CLISettings runtime fields back to os.environ.""" + + def test_propagates_nss_inference_settings(self, monkeypatch): + """Endpoint and key propagate to NSS_INFERENCE_* env vars read by pii_replacer.""" + monkeypatch.delenv("NSS_INFERENCE_ENDPOINT", raising=False) + monkeypatch.delenv("NSS_INFERENCE_KEY", raising=False) + + settings = CLISettings.from_cli_kwargs( + inference_endpoint_url="https://cli.example/v1", + inference_api_key="token-propagated-cli", # pragma: allowlist secret + ) + _propagate_runtime_settings_to_env(settings) + + assert os.environ["NSS_INFERENCE_ENDPOINT"] == "https://cli.example/v1" + assert os.environ["NSS_INFERENCE_KEY"] == "token-propagated-cli" + + def test_propagates_remaining_runtime_settings(self, monkeypatch): + """Model ID, offline mode, and CPU count propagate to their runtime env vars.""" + monkeypatch.delenv("NSS_INFERENCE_MODEL", raising=False) + monkeypatch.delenv("HF_HUB_OFFLINE", raising=False) + monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising=False) + monkeypatch.delenv("NSS_PII_REPLACER_CPU_COUNT", raising=False) + + settings = CLISettings.from_cli_kwargs( + inference_model_id="custom/model", + huggingface_remote=False, + cpu_count=3, + ) + _propagate_runtime_settings_to_env(settings) + + assert os.environ["NSS_INFERENCE_MODEL"] == "custom/model" + assert os.environ["HF_HUB_OFFLINE"] == "1" + assert os.environ["TRANSFORMERS_OFFLINE"] == "1" + assert os.environ["NSS_PII_REPLACER_CPU_COUNT"] == "3" + + def test_enabling_huggingface_remote_disables_offline_env(self, monkeypatch): + """--enable-huggingface-remote sets the HF offline vars to 0, overriding inherited offline env.""" + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1") + + settings = CLISettings.from_cli_kwargs(huggingface_remote=True) + _propagate_runtime_settings_to_env(settings) + + assert os.environ["HF_HUB_OFFLINE"] == "0" + assert os.environ["TRANSFORMERS_OFFLINE"] == "0" + + def test_common_setup_propagates_before_workdir(self, monkeypatch, dummy_csv: Path): + """common_setup writes resolved runtime settings before downstream imports.""" + monkeypatch.delenv("NSS_INFERENCE_KEY", raising=False) + + settings = CLISettings.from_cli_kwargs( + data_source=str(dummy_csv), + inference_api_key="token-propagated-setup", # pragma: allowlist secret + ) + + with ( + patch("nemo_safe_synthesizer.cli.utils._create_workdir") as mock_create_workdir, + patch("nemo_safe_synthesizer.cli.utils.initialize_wandb_run"), + patch("nemo_safe_synthesizer.cli.utils._initialize_logging_for_cli_from_settings") as mock_init_logging, + ): + mock_workdir = MagicMock() + mock_create_workdir.return_value = mock_workdir + mock_init_logging.return_value = MagicMock() + + common_setup(settings) + + assert os.environ["NSS_INFERENCE_KEY"] == "token-propagated-setup" + + class TestCommonSetupReturnValues: """Tests for common_setup return values.""" diff --git a/tests/nss_pii_replacer_test.py b/tests/nss_pii_replacer_test.py index 624533f15..2d36969aa 100644 --- a/tests/nss_pii_replacer_test.py +++ b/tests/nss_pii_replacer_test.py @@ -14,7 +14,7 @@ # Currently use env variables to configure the endpoint and model for column classification. # export NSS_INFERENCE_KEY=<...> # export NSS_INFERENCE_ENDPOINT=https://integrate.api.nvidia.com/v1 -# export NIM_MODEL_ID=qwen/qwen3-next-80b-a3b-instruct +# export NSS_INFERENCE_MODEL=qwen/qwen3-next-80b-a3b-instruct def main(): diff --git a/tests/pii_replacer/test_detect.py b/tests/pii_replacer/test_detect.py index e4d2e814e..9b9d6f203 100644 --- a/tests/pii_replacer/test_detect.py +++ b/tests/pii_replacer/test_detect.py @@ -18,6 +18,7 @@ UNKNOWN_ENTITY, ClassifyConfig, ColumnClassifierLLM, + DefaultLLMConfig, EntityExtractorGliner, _format_prompt, merge_subsume, @@ -28,6 +29,25 @@ from nemo_safe_synthesizer.pii_replacer.ner.ner import NERPrediction +class TestDefaultLLMConfigId: + def test_uses_env_override(self, monkeypatch): + monkeypatch.setenv("NSS_INFERENCE_MODEL", "custom/model") + assert DefaultLLMConfig.config_id() == "custom/model" + + def test_falls_back_when_unset(self, monkeypatch): + monkeypatch.delenv("NSS_INFERENCE_MODEL", raising=False) + assert DefaultLLMConfig.config_id() == DefaultLLMConfig.DEFAULT_CONFIG_ID + + @pytest.mark.parametrize("blank", ["", " ", "\t"]) + def test_blank_value_falls_back(self, monkeypatch, blank): + monkeypatch.setenv("NSS_INFERENCE_MODEL", blank) + assert DefaultLLMConfig.config_id() == DefaultLLMConfig.DEFAULT_CONFIG_ID + + def test_strips_surrounding_whitespace(self, monkeypatch): + monkeypatch.setenv("NSS_INFERENCE_MODEL", " custom/model ") + assert DefaultLLMConfig.config_id() == "custom/model" + + def test_gliner_batch_predict_config(): # Test batch_update_cache is short-circuited iff batch mode disabled. cfg = ClassifyConfig( @@ -67,6 +87,30 @@ def test_gliner_batch_predict_config(): entity_extractor._model.batch_predict_entities.assert_called() # ty: ignore[call-non-callable, unresolved-attribute] -- mock object +@pytest.mark.parametrize("offline_var", ["HF_HUB_OFFLINE", "TRANSFORMERS_OFFLINE"]) +@pytest.mark.parametrize("env_value", ["1", "yes", "on"]) +def test_gliner_local_files_only_follows_hf_offline_env(env_value, offline_var, monkeypatch): + """GLiNER offline mode follows the standard Hugging Face offline env vars.""" + cfg = ClassifyConfig( + valid_entities={"name"}, + ner_threshold=0.8, + ner_regexps_enabled=False, + ner_entities=None, + gliner_enabled=True, + gliner_batch_mode_enabled=False, + gliner_batch_mode_chunk_length=10, + gliner_batch_mode_batch_size=20, + gliner_model="nvidia/gliner-PII", + ) + monkeypatch.delenv("HF_HUB_OFFLINE", raising=False) + monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising=False) + monkeypatch.setenv(offline_var, env_value) + + with patch("nemo_safe_synthesizer.pii_replacer.data_editor.detect.GLiNER") as mock_gliner: + EntityExtractorGliner.get_entity_extractor(cfg) + assert mock_gliner.from_pretrained.call_args.kwargs["local_files_only"] is True + + def test_gliner_pii_detection_recall(): # Tests GLiNER’s PII detection on a short text, ensuring it finds a reasonable number of entities without over- or under-detecting. diff --git a/tests/pii_replacer/test_nemo_pii.py b/tests/pii_replacer/test_nemo_pii.py index 3659729d7..ca63ef193 100644 --- a/tests/pii_replacer/test_nemo_pii.py +++ b/tests/pii_replacer/test_nemo_pii.py @@ -10,8 +10,35 @@ import pandas as pd +from nemo_safe_synthesizer.defaults import DEFAULT_NSS_INFERENCE_ENDPOINT from nemo_safe_synthesizer.pii_replacer.data_editor.edit import TransformFnAccounting -from nemo_safe_synthesizer.pii_replacer.nemo_pii import ColumnClassification, NemoPII, _build_column_statistics +from nemo_safe_synthesizer.pii_replacer.nemo_pii import ( + ColumnClassification, + NemoPII, + _build_column_statistics, + _get_classify_endpoint_url, +) + + +class TestGetClassifyEndpointUrl: + def test_configured_value_is_used(self, monkeypatch): + monkeypatch.setenv("NSS_INFERENCE_ENDPOINT", "https://custom.example/v1") + assert _get_classify_endpoint_url() == "https://custom.example/v1" + + def test_configured_value_is_stripped(self, monkeypatch): + monkeypatch.setenv("NSS_INFERENCE_ENDPOINT", " https://custom.example/v1 ") + assert _get_classify_endpoint_url() == "https://custom.example/v1" + + def test_unset_falls_back_to_default(self, monkeypatch): + monkeypatch.delenv("NSS_INFERENCE_ENDPOINT", raising=False) + assert _get_classify_endpoint_url() == DEFAULT_NSS_INFERENCE_ENDPOINT + + @pytest.mark.parametrize("blank", ["", " ", "\t"]) + def test_blank_falls_back_to_default(self, monkeypatch, blank): + # A blank endpoint must resolve to the default, never reach the OpenAI + # client as an empty base_url. Mirrors the preflight blank-endpoint rule. + monkeypatch.setenv("NSS_INFERENCE_ENDPOINT", blank) + assert _get_classify_endpoint_url() == DEFAULT_NSS_INFERENCE_ENDPOINT @pytest.fixture diff --git a/tests/preflight/test_preflight.py b/tests/preflight/test_preflight.py index 656acb632..ef1497912 100644 --- a/tests/preflight/test_preflight.py +++ b/tests/preflight/test_preflight.py @@ -29,7 +29,7 @@ DatasetSizeCheck, GroupbyColumnCheck, HFModelAvailabilityCheck, - InferenceKeyCheck, + InferenceModelCheck, OrderbyColumnCheck, OversamplingCheck, PreflightContext, @@ -243,22 +243,103 @@ def test_meta_tensor_path_is_architecture_exact(self, model_type, fields, expect @pytest.mark.unit -class TestInferenceKeyCheck: - def test_empty_env_emits_warning(self, default_config): +class TestInferenceModelCheck: + def test_empty_env_emits_key_warning(self, default_config): with patch.dict("os.environ", {}, clear=True): - issues = InferenceKeyCheck().run(make_ctx(config=default_config)) + issues = InferenceModelCheck().run(make_ctx(config=default_config)) assert any(i.code == "inference_key_missing" and i.severity == "warning" for i in issues) def test_inference_key_present_is_silent(self, default_config): with patch.dict("os.environ", {"NSS_INFERENCE_KEY": "test-key", "HF_TOKEN": "hf_xxx"}): - issues = InferenceKeyCheck().run(make_ctx(config=default_config)) + issues = InferenceModelCheck().run(make_ctx(config=default_config)) assert not any(i.code == "inference_key_missing" for i in issues) - def test_pii_disabled_skips_key_requirement(self): + def test_pii_disabled_skips_all_checks(self): config = SafeSynthesizerParameters(replace_pii=None) - with patch.dict("os.environ", {"HF_TOKEN": "hf_xxx"}, clear=True): - issues = InferenceKeyCheck().run(make_ctx(config=config)) - assert not any(i.code == "inference_key_missing" for i in issues) + with patch.dict( + "os.environ", + {"NSS_INFERENCE_MODEL": "", "NSS_INFERENCE_ENDPOINT": "not-a-url"}, + clear=True, + ): + issues = InferenceModelCheck().run(make_ctx(config=config)) + assert issues == [] + + def test_blank_model_emits_warning(self, default_config): + with patch.dict( + "os.environ", + {"NSS_INFERENCE_KEY": "test-key", "NSS_INFERENCE_MODEL": " "}, + ): + issues = InferenceModelCheck().run(make_ctx(config=default_config)) + assert any(i.code == "inference_model_blank" and i.severity == "warning" for i in issues) + + def test_unset_model_is_silent(self, default_config): + with patch.dict("os.environ", {"NSS_INFERENCE_KEY": "test-key"}, clear=True): + issues = InferenceModelCheck().run(make_ctx(config=default_config)) + assert not any(i.code == "inference_model_blank" for i in issues) + + def test_valid_model_is_silent(self, default_config): + with patch.dict( + "os.environ", + {"NSS_INFERENCE_KEY": "test-key", "NSS_INFERENCE_MODEL": "qwen/qwen3-next-80b-a3b-instruct"}, + ): + issues = InferenceModelCheck().run(make_ctx(config=default_config)) + assert not any(i.code == "inference_model_blank" for i in issues) + + @pytest.mark.parametrize("endpoint", ["not-a-url", "ftp://example.com", "http://"]) + def test_invalid_endpoint_emits_error(self, default_config, endpoint): + # An invalid endpoint cannot succeed, so it must fail preflight (error), + # not merely warn -- otherwise --validate passes a config that fails on + # the first classification request. + with patch.dict( + "os.environ", + {"NSS_INFERENCE_KEY": "test-key", "NSS_INFERENCE_ENDPOINT": endpoint}, + ): + issues = InferenceModelCheck().run(make_ctx(config=default_config)) + assert any(i.code == "inference_endpoint_invalid" and i.severity == "error" for i in issues) + + def test_valid_endpoint_is_silent(self, default_config): + with patch.dict( + "os.environ", + {"NSS_INFERENCE_KEY": "test-key", "NSS_INFERENCE_ENDPOINT": "https://integrate.api.nvidia.com/v1"}, + ): + issues = InferenceModelCheck().run(make_ctx(config=default_config)) + assert not any(i.code == "inference_endpoint_invalid" for i in issues) + + @pytest.mark.parametrize("blank", ["", " "]) + def test_blank_endpoint_is_silent(self, default_config, blank): + # A blank endpoint is treated as unset (falls back to the default base + # URL), not as an invalid endpoint. + with patch.dict( + "os.environ", + {"NSS_INFERENCE_KEY": "test-key", "NSS_INFERENCE_ENDPOINT": blank}, + ): + issues = InferenceModelCheck().run(make_ctx(config=default_config)) + assert not any(i.code == "inference_endpoint_invalid" for i in issues) + + def test_invalid_endpoint_takes_priority_over_warnings(self, default_config): + # Single-dispatch match: the invalid-endpoint error is checked first, so + # it wins over the missing-key and blank-model warnings. + with patch.dict( + "os.environ", + {"NSS_INFERENCE_MODEL": "", "NSS_INFERENCE_ENDPOINT": "not-a-url"}, + clear=True, + ): + issues = InferenceModelCheck().run(make_ctx(config=default_config)) + codes = {i.code for i in issues} + assert codes == {"inference_endpoint_invalid"} + assert all(i.severity == "error" for i in issues if i.code == "inference_endpoint_invalid") + + def test_missing_key_takes_priority_over_blank_model(self, default_config): + # With a valid endpoint, the missing-key warning outranks the blank-model + # warning. + with patch.dict( + "os.environ", + {"NSS_INFERENCE_MODEL": " ", "NSS_INFERENCE_ENDPOINT": "https://integrate.api.nvidia.com/v1"}, + clear=True, + ): + issues = InferenceModelCheck().run(make_ctx(config=default_config)) + codes = {i.code for i in issues} + assert codes == {"inference_key_missing"} @pytest.mark.unit @@ -725,6 +806,15 @@ def test_extreme_oversampling_is_flagged(self, sample_df): @pytest.mark.unit class TestRunPreflight: + @pytest.fixture(autouse=True) + def _isolate_hf_offline_env(self, monkeypatch): + # run_preflight invokes HFModelAvailabilityCheck, which escalates + # hf_model_not_cached to an error when HF offline mode is enabled. Clear + # the ambient offline vars so these tests do not fail when the developer + # (or CI) has HF_HUB_OFFLINE set and the model is not cached. + monkeypatch.delenv("HF_HUB_OFFLINE", raising=False) + monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising=False) + def test_clean_dataset_has_no_errors(self, sample_df, default_config): resolved_config = default_config.model_copy( update={ diff --git a/tests/test_env_flags.py b/tests/test_env_flags.py new file mode 100644 index 000000000..38ff9f61c --- /dev/null +++ b/tests/test_env_flags.py @@ -0,0 +1,51 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for shared environment-flag parsing helpers.""" + +from __future__ import annotations + +import pytest + +from nemo_safe_synthesizer.utils import env_flag_is_true, hf_offline_enabled + +_PROBE_VAR = "NSS_TEST_FLAG" + + +@pytest.mark.parametrize( + ("value", "expected"), + [ + ("true", True), + ("True", True), + ("1", True), + ("yes", True), + ("on", True), + ("false", False), + ("0", False), + ("no", False), + ("", False), + ], +) +def test_env_flag_is_true(value: str, expected: bool, monkeypatch): + monkeypatch.setenv(_PROBE_VAR, value) + assert env_flag_is_true(_PROBE_VAR) is expected + + +def test_env_flag_is_true_unset_uses_default(monkeypatch): + monkeypatch.delenv(_PROBE_VAR, raising=False) + assert env_flag_is_true(_PROBE_VAR) is False + assert env_flag_is_true(_PROBE_VAR, default=True) is True + + +@pytest.mark.parametrize("offline_var", ["HF_HUB_OFFLINE", "TRANSFORMERS_OFFLINE"]) +def test_hf_offline_enabled_true_for_either_var(offline_var: str, monkeypatch): + monkeypatch.delenv("HF_HUB_OFFLINE", raising=False) + monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising=False) + monkeypatch.setenv(offline_var, "1") + assert hf_offline_enabled() is True + + +def test_hf_offline_enabled_false_when_unset(monkeypatch): + monkeypatch.delenv("HF_HUB_OFFLINE", raising=False) + monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising=False) + assert hf_offline_enabled() is False