From 7d39790b250c7522469ebb5f1896209e618ae255 Mon Sep 17 00:00:00 2001 From: Aaron Gonzales Date: Thu, 28 May 2026 01:10:05 +0000 Subject: [PATCH 01/10] feat(cli): add CLI flags for env-only NIM/offline/cpu settings 5 settings were configurable only via environment variables with no CLI flag equivalent, making them undiscoverable through --help. Expose --nim-endpoint-url, --nim-api-key, --nim-model-id, --local-files-only, and --cpu-count as CLISettings fields with AliasChoices so env vars continue to work and CLI flags take precedence. Resolved values are propagated back to os.environ in common_setup before the deferred pii_replacer imports so the existing deep readers honor CLI overrides without threading a settings handle through every callsite. Closes #155 Signed-off-by: Aaron Gonzales --- src/nemo_safe_synthesizer/cli/run.py | 83 +++++++++++++++++++++++ src/nemo_safe_synthesizer/cli/settings.py | 36 ++++++++++ src/nemo_safe_synthesizer/cli/utils.py | 35 ++++++++++ 3 files changed, 154 insertions(+) diff --git a/src/nemo_safe_synthesizer/cli/run.py b/src/nemo_safe_synthesizer/cli/run.py index a6ffbc14c..562e646da 100644 --- a/src/nemo_safe_synthesizer/cli/run.py +++ b/src/nemo_safe_synthesizer/cli/run.py @@ -159,6 +159,59 @@ 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( + "--nim-endpoint-url", + type=str, + required=False, + default=None, + help="NIM/OpenAI-compatible endpoint URL for PII column classification. " + "Can also be set via NIM_ENDPOINT_URL env var.", + ) + ) + options.append( + click.option( + "--nim-api-key", + type=str, + required=False, + default=None, + help="API key for the NIM endpoint used in PII column classification. " + "Can also be set via NIM_API_KEY env var.", + ) + ) + options.append( + click.option( + "--nim-model-id", + type=str, + required=False, + default=None, + help="Model ID sent to the NIM endpoint for PII column classification. " + "Can also be set via NIM_MODEL_ID env var. " + "[default: qwen/qwen3-next-80b-a3b-instruct]", + ) + ) + options.append( + click.option( + "--local-files-only/--no-local-files-only", + "local_files_only", + type=click.BOOL, + required=False, + default=None, + help="If set, GLiNER skips network downloads and uses only local files. " + "Can also be set via LOCAL_FILES_ONLY env var.", + ) + ) + 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 SAFE_SYNTHESIZER_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) @@ -323,6 +376,11 @@ def run( wandb_mode: str | None = None, wandb_project: str | None = None, dataset_registry: str | None = None, + nim_endpoint_url: str | None = None, + nim_api_key: str | None = None, + nim_model_id: str | None = None, + local_files_only: bool | None = None, + cpu_count: int | None = None, validate: bool = False, **kwargs: object, ) -> None: @@ -351,6 +409,11 @@ def run( wandb_project=wandb_project, synthesis_overrides=_parse_run_overrides(kwargs), dataset_registry=dataset_registry, + nim_endpoint_url=nim_endpoint_url, + nim_api_key=nim_api_key, + nim_model_id=nim_model_id, + local_files_only=local_files_only, + cpu_count=cpu_count, ) if validate: @@ -419,6 +482,11 @@ def run_train( wandb_mode: str | None = None, wandb_project: str | None = None, dataset_registry: str | None = None, + nim_endpoint_url: str | None = None, + nim_api_key: str | None = None, + nim_model_id: str | None = None, + local_files_only: bool | None = None, + cpu_count: int | None = None, validate: bool = False, **kwargs: object, ) -> None: @@ -443,6 +511,11 @@ def run_train( wandb_project=wandb_project, synthesis_overrides=_parse_run_overrides(kwargs), dataset_registry=dataset_registry, + nim_endpoint_url=nim_endpoint_url, + nim_api_key=nim_api_key, + nim_model_id=nim_model_id, + local_files_only=local_files_only, + cpu_count=cpu_count, ) if validate: @@ -512,6 +585,11 @@ def run_generate( auto_discover_adapter: bool = False, wandb_resume_job_id: str | None = None, dataset_registry: str | None = None, + nim_endpoint_url: str | None = None, + nim_api_key: str | None = None, + nim_model_id: str | None = None, + local_files_only: bool | None = None, + cpu_count: int | None = None, **kwargs: object, ) -> None: """Run the generation stage only. @@ -540,6 +618,11 @@ def run_generate( wandb_project=wandb_project, synthesis_overrides=_parse_run_overrides(kwargs), dataset_registry=dataset_registry, + nim_endpoint_url=nim_endpoint_url, + nim_api_key=nim_api_key, + nim_model_id=nim_model_id, + local_files_only=local_files_only, + cpu_count=cpu_count, ) os.environ["NSS_PHASE"] = "generate" diff --git a/src/nemo_safe_synthesizer/cli/settings.py b/src/nemo_safe_synthesizer/cli/settings.py index e088bef89..412385750 100644 --- a/src/nemo_safe_synthesizer/cli/settings.py +++ b/src/nemo_safe_synthesizer/cli/settings.py @@ -164,6 +164,42 @@ class CLISettings(BaseSettings): ) """URL or path to a dataset registry YAML file (env: ``NSS_DATASET_REGISTRY``).""" + nim_endpoint_url: str | None = Field( + default=None, + validation_alias=AliasChoices("nim_endpoint_url", "NIM_ENDPOINT_URL"), + description="NIM/OpenAI-compatible endpoint URL for PII column classification", + ) + """NIM/OpenAI-compatible endpoint URL for PII column classification (env: ``NIM_ENDPOINT_URL``).""" + + nim_api_key: str | None = Field( + default=None, + validation_alias=AliasChoices("nim_api_key", "NIM_API_KEY"), + description="API key for the NIM endpoint used in PII column classification", + ) + """API key for the NIM endpoint used in PII column classification (env: ``NIM_API_KEY``).""" + + nim_model_id: str | None = Field( + default=None, + validation_alias=AliasChoices("nim_model_id", "NIM_MODEL_ID"), + description="Model ID sent to the NIM endpoint for PII column classification", + ) + """Model ID sent to the NIM endpoint for PII column classification (env: ``NIM_MODEL_ID``).""" + + local_files_only: bool | None = Field( + default=None, + validation_alias=AliasChoices("local_files_only", "LOCAL_FILES_ONLY"), + description="Whether GLiNER should skip network downloads and use only local files", + ) + """Whether GLiNER should skip network downloads and use only local files (env: ``LOCAL_FILES_ONLY``).""" + + cpu_count: int | None = Field( + default=None, + validation_alias=AliasChoices("cpu_count", "SAFE_SYNTHESIZER_CPU_COUNT"), + description="Number of CPU worker processes used for NER (PII replacement)", + ) + """Number of CPU worker processes used for NER (PII replacement) + (env: ``SAFE_SYNTHESIZER_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..d17986a4c 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 NIM_MODEL_ID, LOCAL_FILES_ONLY, and SAFE_SYNTHESIZER_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,35 @@ 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. + """ + if settings.nim_endpoint_url is not None: + os.environ["NIM_ENDPOINT_URL"] = settings.nim_endpoint_url + if settings.nim_api_key is not None: + os.environ["NIM_API_KEY"] = settings.nim_api_key + if settings.nim_model_id is not None: + os.environ["NIM_MODEL_ID"] = settings.nim_model_id + if settings.local_files_only is not None: + os.environ["LOCAL_FILES_ONLY"] = "true" if settings.local_files_only else "false" + if settings.cpu_count is not None: + os.environ["SAFE_SYNTHESIZER_CPU_COUNT"] = str(settings.cpu_count) + + def _initialize_logging_for_cli_from_settings( settings: "CLISettings", workdir: Workdir, From b92007712b0f180824d03fc94fe34903ac4239e6 Mon Sep 17 00:00:00 2001 From: Aaron Gonzales Date: Thu, 28 May 2026 17:34:05 +0000 Subject: [PATCH 02/10] fix(cli): propagate PII runtime flags via NSS_INFERENCE_* env names Align CLISettings and env propagation with runtime readers in pii_replacer (NSS_INFERENCE_ENDPOINT/KEY), add NSS-over-NIM alias precedence and NSS_LOG_COLOR, share env_flag_is_true for GLiNER offline parsing, and restructure environment docs with CLI flag mappings. Closes #155 Signed-off-by: Aaron Gonzales --- docs/user-guide/docker.md | 2 +- docs/user-guide/environment.md | 343 +++++++++--------- docs/user-guide/running.md | 2 +- src/nemo_safe_synthesizer/cli/run.py | 4 +- src/nemo_safe_synthesizer/cli/settings.py | 53 ++- src/nemo_safe_synthesizer/cli/utils.py | 8 +- .../pii_replacer/data_editor/detect.py | 3 +- src/nemo_safe_synthesizer/utils.py | 14 + tests/cli/test_run.py | 13 + tests/cli/test_settings.py | 62 ++++ tests/cli/test_utils.py | 60 ++- tests/pii_replacer/test_detect.py | 21 ++ tests/test_env_flags.py | 35 ++ 13 files changed, 431 insertions(+), 189 deletions(-) create mode 100644 tests/test_env_flags.py 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..d0c331269 100644 --- a/docs/user-guide/environment.md +++ b/docs/user-guide/environment.md @@ -3,274 +3,271 @@ # 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 + +| 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 | `--nim-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 | `--nim-api-key` | PII column classifier | -- | API key for `NSS_INFERENCE_ENDPOINT` | Required for LLM column classification | +| `NIM_MODEL_ID` | nss | `--nim-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) | +| `LOCAL_FILES_ONLY` | nss | `--local-files-only` / `--no-local-files-only` | GLiNER (PII) | unset | Skip GLiNER network downloads | Partial offline; see [HF appendix](#hugging-face-cache-and-offline) | +| `SAFE_SYNTHESIZER_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 | -- | Hugging Face Hub | unset | Fail if a model is not cached | Preferred offline gate | +| `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 +## Precedence {#precedence} + +### Infrastructure (CLISettings) -Infrastructure settings (artifact path, logging, WandB): +For artifact paths, logging, WandB overrides, and the five runtime flags +(`--nim-*`, `--local-files-only`, `--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 -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 {#hugging-face-cache-and-offline} -### Pre-Caching Models +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). + +### `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. +When set to `1`, Hugging Face Hub refuses network access. Use with a +pre-populated `HF_HOME` for reliable offline runs. ```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. +Prefer this over `LOCAL_FILES_ONLY` for end-to-end offline behavior. ### `LOCAL_FILES_ONLY` -Skips network downloads for GLiNER. Not respected by the HuggingFace training -backend or vLLM. +Skips network downloads for GLiNER only. Not respected by the HuggingFace +training backend or vLLM. Override on the CLI with `--local-files-only` or +`--no-local-files-only`. ```bash export LOCAL_FILES_ONLY=true ``` !!! 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. - -### `VLLM_CACHE_ROOT` - -Sets the vLLM model cache directory. - -```bash -export VLLM_CACHE_ROOT=/shared/cache/vllm -``` - ---- + For the most reliable offline experience, set `HF_HUB_OFFLINE=1` with a + pre-populated `HF_HOME` cache instead of relying on `LOCAL_FILES_ONLY` + alone. -## Attention and Compute +### Pre-caching models {#pre-caching-models} -GPU attention backend selection for the vLLM generation engine. +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. -### `VLLM_ATTENTION_BACKEND` - -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. - -```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 {#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, use `--nim-api-key` and optionally `--nim-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" +### `NIM_MODEL_ID` - ```yaml - replace_pii: - globals: - classify: - enable_classify: false - ``` +Model ID sent to the inference endpoint. Override with `--nim-model-id`. -=== "SDK" +### `SAFE_SYNTHESIZER_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 SAFE_SYNTHESIZER_CPU_COUNT=4 +``` - synthesizer = ( - SafeSynthesizer(config) - .with_data_source("data.csv") - .with_replace_pii(config=pii_config) - ) - ``` +--- -### `NSS_INFERENCE_KEY` +## vLLM and attention {#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). + +--- + +## Telemetry {#telemetry} ### `NEMO_TELEMETRY_ENABLED` -Controls whether telemetry is sent for train/generate events. - -Defaults to `true`. Set it to `false` to disable telemetry for the current shell: +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 {#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 -``` +--- + +## Internal and cluster {#internal-and-cluster} -See [Docker](docker.md) for full container setup and Makefile shortcuts. +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..50133f706 100644 --- a/docs/user-guide/running.md +++ b/docs/user-guide/running.md @@ -1233,7 +1233,7 @@ 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). +[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/src/nemo_safe_synthesizer/cli/run.py b/src/nemo_safe_synthesizer/cli/run.py index 562e646da..3d98ac34d 100644 --- a/src/nemo_safe_synthesizer/cli/run.py +++ b/src/nemo_safe_synthesizer/cli/run.py @@ -166,7 +166,7 @@ def common_run_options(f: Callable[..., object]) -> Callable[..., object]: required=False, default=None, help="NIM/OpenAI-compatible endpoint URL for PII column classification. " - "Can also be set via NIM_ENDPOINT_URL env var.", + "Can also be set via NSS_INFERENCE_ENDPOINT env var.", ) ) options.append( @@ -176,7 +176,7 @@ def common_run_options(f: Callable[..., object]) -> Callable[..., object]: required=False, default=None, help="API key for the NIM endpoint used in PII column classification. " - "Can also be set via NIM_API_KEY env var.", + "Can also be set via NSS_INFERENCE_KEY env var.", ) ) options.append( diff --git a/src/nemo_safe_synthesizer/cli/settings.py b/src/nemo_safe_synthesizer/cli/settings.py index 412385750..e9c7b2d5f 100644 --- a/src/nemo_safe_synthesizer/cli/settings.py +++ b/src/nemo_safe_synthesizer/cli/settings.py @@ -27,10 +27,11 @@ from __future__ import annotations +import os from pathlib import Path from typing import Any, Literal -from pydantic import AliasChoices, Field, field_validator +from pydantic import AliasChoices, Field, field_validator, model_validator from pydantic_settings import BaseSettings, SettingsConfigDict from ..defaults import DEFAULT_ARTIFACTS_PATH @@ -39,6 +40,30 @@ __all__ = ["CLISettings"] +# (settings field, canonical env var, legacy env alias from issue #155) +_INFERENCE_ENV_ALIASES: tuple[tuple[str, str, str], ...] = ( + ("nim_endpoint_url", "NSS_INFERENCE_ENDPOINT", "NIM_ENDPOINT_URL"), + ("nim_api_key", "NSS_INFERENCE_KEY", "NIM_API_KEY"), +) + + +def _apply_inference_env_precedence( + data: dict[str, Any], + field: str, + canonical_env: str, + legacy_env: str, +) -> None: + """Prefer ``canonical_env`` over ``legacy_env`` when both are set.""" + canonical = os.environ.get(canonical_env) + legacy = os.environ.get(legacy_env) + match (field in data, data.get(field), canonical, legacy): + case (True, leg, str() as canon, str() as leg_env) if leg == leg_env and canon != leg_env: + data[field] = canon + case (False, _, str() as canon, _): + data[field] = canon + case (False, _, None, str() as leg_env): + data[field] = leg_env + class CLISettings(BaseSettings): """Unified CLI settings composing all sub-settings. @@ -118,9 +143,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, @@ -166,17 +192,19 @@ class CLISettings(BaseSettings): nim_endpoint_url: str | None = Field( default=None, - validation_alias=AliasChoices("nim_endpoint_url", "NIM_ENDPOINT_URL"), + validation_alias=AliasChoices("nim_endpoint_url", "NSS_INFERENCE_ENDPOINT"), description="NIM/OpenAI-compatible endpoint URL for PII column classification", ) - """NIM/OpenAI-compatible endpoint URL for PII column classification (env: ``NIM_ENDPOINT_URL``).""" + """NIM/OpenAI-compatible endpoint URL for PII column classification + (env: ``NSS_INFERENCE_ENDPOINT``; alias: ``NIM_ENDPOINT_URL``).""" nim_api_key: str | None = Field( default=None, - validation_alias=AliasChoices("nim_api_key", "NIM_API_KEY"), + validation_alias=AliasChoices("nim_api_key", "NSS_INFERENCE_KEY"), description="API key for the NIM endpoint used in PII column classification", ) - """API key for the NIM endpoint used in PII column classification (env: ``NIM_API_KEY``).""" + """API key for the NIM endpoint used in PII column classification + (env: ``NSS_INFERENCE_KEY``; alias: ``NIM_API_KEY``).""" nim_model_id: str | None = Field( default=None, @@ -200,6 +228,19 @@ class CLISettings(BaseSettings): """Number of CPU worker processes used for NER (PII replacement) (env: ``SAFE_SYNTHESIZER_CPU_COUNT``).""" + @model_validator(mode="before") + @classmethod + def resolve_inference_env_aliases(cls, data: Any) -> Any: + """Prefer ``NSS_INFERENCE_*`` over legacy ``NIM_*`` env aliases.""" + match data: + case dict() as payload: + resolved = dict(payload) + case _: + return data + for field, canonical_env, legacy_env in _INFERENCE_ENV_ALIASES: + _apply_inference_env_precedence(resolved, field, canonical_env, legacy_env) + return resolved + @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 d17986a4c..a3094f7ce 100644 --- a/src/nemo_safe_synthesizer/cli/utils.py +++ b/src/nemo_safe_synthesizer/cli/utils.py @@ -232,8 +232,8 @@ def common_setup( """ # 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 NIM_MODEL_ID, LOCAL_FILES_ONLY, and SAFE_SYNTHESIZER_CPU_COUNT see - # the CLI-overridden values. + # of NSS_INFERENCE_*, NIM_MODEL_ID, LOCAL_FILES_ONLY, and + # SAFE_SYNTHESIZER_CPU_COUNT see the CLI-overridden values. _propagate_runtime_settings_to_env(settings) # 1. Create workdir FIRST - this establishes all artifact paths @@ -338,9 +338,9 @@ def _propagate_runtime_settings_to_env(settings: "CLISettings") -> None: pipeline see the CLI value. """ if settings.nim_endpoint_url is not None: - os.environ["NIM_ENDPOINT_URL"] = settings.nim_endpoint_url + os.environ["NSS_INFERENCE_ENDPOINT"] = settings.nim_endpoint_url if settings.nim_api_key is not None: - os.environ["NIM_API_KEY"] = settings.nim_api_key + os.environ["NSS_INFERENCE_KEY"] = settings.nim_api_key if settings.nim_model_id is not None: os.environ["NIM_MODEL_ID"] = settings.nim_model_id if settings.local_files_only is not None: 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..1e97e05cf 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 env_flag_is_true from ..ner import ner_mp from ..ner.factory import LabelSetPredictorFilter, NERFactory from ..ner.ner import NERPrediction @@ -575,7 +576,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=env_flag_is_true("LOCAL_FILES_ONLY"), ) entity_types = DEFAULT_ENTITIES if clsfy_cfg.ner_entities: diff --git a/src/nemo_safe_synthesizer/utils.py b/src/nemo_safe_synthesizer/utils.py index 2d62f81b9..153f7dafe 100644 --- a/src/nemo_safe_synthesizer/utils.py +++ b/src/nemo_safe_synthesizer/utils.py @@ -27,6 +27,20 @@ logger = get_logger(__name__) +_TRUTHY_ENV_VALUES = frozenset({"1", "true", "yes", "on"}) + + +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 _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_run.py b/tests/cli/test_run.py index 6fad15b6f..870944ad1 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 "--nim-endpoint-url" in result.output + assert "--nim-api-key" in result.output + assert "--nim-model-id" in result.output + assert "--local-files-only" 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, diff --git a/tests/cli/test_settings.py b/tests/cli/test_settings.py index ca8151ba6..0122de87f 100644 --- a/tests/cli/test_settings.py +++ b/tests/cli/test_settings.py @@ -217,6 +217,68 @@ 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_nim_endpoint_url_from_nss_inference_env(self, monkeypatch): + """NSS_INFERENCE_ENDPOINT loads into nim_endpoint_url.""" + monkeypatch.setenv("NSS_INFERENCE_ENDPOINT", "https://custom.example/v1") + settings = CLISettings() + assert settings.nim_endpoint_url == "https://custom.example/v1" + + def test_nim_api_key_from_nss_inference_env(self, monkeypatch): + """NSS_INFERENCE_KEY loads into nim_api_key.""" + monkeypatch.setenv("NSS_INFERENCE_KEY", "secret-key") + settings = CLISettings() + assert settings.nim_api_key == "secret-key" + + def test_nim_endpoint_url_cli_overrides_env(self, monkeypatch): + """CLI --nim-endpoint-url takes precedence over NSS_INFERENCE_ENDPOINT.""" + monkeypatch.setenv("NSS_INFERENCE_ENDPOINT", "https://env.example/v1") + settings = CLISettings.from_cli_kwargs(nim_endpoint_url="https://cli.example/v1") + assert settings.nim_endpoint_url == "https://cli.example/v1" + + def test_nim_api_key_cli_overrides_env(self, monkeypatch): + """CLI --nim-api-key takes precedence over NSS_INFERENCE_KEY.""" + monkeypatch.setenv("NSS_INFERENCE_KEY", "env-key") + settings = CLISettings.from_cli_kwargs(nim_api_key="cli-key") + assert settings.nim_api_key == "cli-key" + + def test_nim_api_key_prefers_nss_over_nim_alias_when_both_set(self, monkeypatch): + """Canonical NSS_INFERENCE_KEY wins when both NSS and NIM aliases are set.""" + monkeypatch.setenv("NSS_INFERENCE_KEY", "nss-key") + monkeypatch.setenv("NIM_API_KEY", "nim-key") + settings = CLISettings() + assert settings.nim_api_key == "nss-key" + + def test_nim_api_key_falls_back_to_nim_alias(self, monkeypatch): + """Legacy NIM_API_KEY loads when NSS_INFERENCE_KEY is unset.""" + monkeypatch.delenv("NSS_INFERENCE_KEY", raising=False) + monkeypatch.setenv("NIM_API_KEY", "nim-only") + settings = CLISettings() + assert settings.nim_api_key == "nim-only" + + 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("NIM_MODEL_ID", "custom/model") + monkeypatch.setenv("LOCAL_FILES_ONLY", "true") + monkeypatch.setenv("SAFE_SYNTHESIZER_CPU_COUNT", "4") + + settings = CLISettings() + assert settings.nim_model_id == "custom/model" + assert settings.local_files_only is True + assert settings.cpu_count == 4 + 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..8cdc02734 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,63 @@ 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( + nim_endpoint_url="https://cli.example/v1", + nim_api_key="cli-secret", + ) + _propagate_runtime_settings_to_env(settings) + + assert os.environ["NSS_INFERENCE_ENDPOINT"] == "https://cli.example/v1" + assert os.environ["NSS_INFERENCE_KEY"] == "cli-secret" + + def test_propagates_remaining_runtime_settings(self, monkeypatch): + """Model ID, offline mode, and CPU count propagate to their runtime env vars.""" + monkeypatch.delenv("NIM_MODEL_ID", raising=False) + monkeypatch.delenv("LOCAL_FILES_ONLY", raising=False) + monkeypatch.delenv("SAFE_SYNTHESIZER_CPU_COUNT", raising=False) + + settings = CLISettings.from_cli_kwargs( + nim_model_id="custom/model", + local_files_only=True, + cpu_count=3, + ) + _propagate_runtime_settings_to_env(settings) + + assert os.environ["NIM_MODEL_ID"] == "custom/model" + assert os.environ["LOCAL_FILES_ONLY"] == "true" + assert os.environ["SAFE_SYNTHESIZER_CPU_COUNT"] == "3" + + 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), + nim_api_key="setup-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"] == "setup-secret" + + class TestCommonSetupReturnValues: """Tests for common_setup return values.""" diff --git a/tests/pii_replacer/test_detect.py b/tests/pii_replacer/test_detect.py index e4d2e814e..7ad970f00 100644 --- a/tests/pii_replacer/test_detect.py +++ b/tests/pii_replacer/test_detect.py @@ -67,6 +67,27 @@ 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("env_value", ["1", "yes", "on"]) +def test_gliner_local_files_only_accepts_common_truthy_env_values(env_value, monkeypatch): + """GLiNER offline mode accepts the same truthy spellings as env_flag_is_true.""" + 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.setenv("LOCAL_FILES_ONLY", 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/test_env_flags.py b/tests/test_env_flags.py new file mode 100644 index 000000000..d2d0c0acc --- /dev/null +++ b/tests/test_env_flags.py @@ -0,0 +1,35 @@ +# 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 + + +@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("LOCAL_FILES_ONLY", value) + assert env_flag_is_true("LOCAL_FILES_ONLY") is expected + + +def test_env_flag_is_true_unset_uses_default(monkeypatch): + monkeypatch.delenv("LOCAL_FILES_ONLY", raising=False) + assert env_flag_is_true("LOCAL_FILES_ONLY") is False + assert env_flag_is_true("LOCAL_FILES_ONLY", default=True) is True From d70ac4286fdddf7f32a85407873bdde447e072a4 Mon Sep 17 00:00:00 2001 From: Aaron Gonzales Date: Thu, 28 May 2026 19:26:36 +0000 Subject: [PATCH 03/10] fix(tests): allowlist fake inference tokens in CLI settings tests Signed-off-by: Aaron Gonzales --- tests/cli/test_settings.py | 20 ++++++++++---------- tests/cli/test_utils.py | 8 ++++---- 2 files changed, 14 insertions(+), 14 deletions(-) diff --git a/tests/cli/test_settings.py b/tests/cli/test_settings.py index 0122de87f..b46a43fb9 100644 --- a/tests/cli/test_settings.py +++ b/tests/cli/test_settings.py @@ -225,9 +225,9 @@ def test_nim_endpoint_url_from_nss_inference_env(self, monkeypatch): def test_nim_api_key_from_nss_inference_env(self, monkeypatch): """NSS_INFERENCE_KEY loads into nim_api_key.""" - monkeypatch.setenv("NSS_INFERENCE_KEY", "secret-key") + monkeypatch.setenv("NSS_INFERENCE_KEY", "token-from-env") settings = CLISettings() - assert settings.nim_api_key == "secret-key" + assert settings.nim_api_key == "token-from-env" # pragma: allowlist secret def test_nim_endpoint_url_cli_overrides_env(self, monkeypatch): """CLI --nim-endpoint-url takes precedence over NSS_INFERENCE_ENDPOINT.""" @@ -237,23 +237,23 @@ def test_nim_endpoint_url_cli_overrides_env(self, monkeypatch): def test_nim_api_key_cli_overrides_env(self, monkeypatch): """CLI --nim-api-key takes precedence over NSS_INFERENCE_KEY.""" - monkeypatch.setenv("NSS_INFERENCE_KEY", "env-key") - settings = CLISettings.from_cli_kwargs(nim_api_key="cli-key") - assert settings.nim_api_key == "cli-key" + monkeypatch.setenv("NSS_INFERENCE_KEY", "token-from-env") + settings = CLISettings.from_cli_kwargs(nim_api_key="token-from-cli") # pragma: allowlist secret + assert settings.nim_api_key == "token-from-cli" # pragma: allowlist secret def test_nim_api_key_prefers_nss_over_nim_alias_when_both_set(self, monkeypatch): """Canonical NSS_INFERENCE_KEY wins when both NSS and NIM aliases are set.""" - monkeypatch.setenv("NSS_INFERENCE_KEY", "nss-key") - monkeypatch.setenv("NIM_API_KEY", "nim-key") + monkeypatch.setenv("NSS_INFERENCE_KEY", "token-from-nss") + monkeypatch.setenv("NIM_API_KEY", "token-from-nim") settings = CLISettings() - assert settings.nim_api_key == "nss-key" + assert settings.nim_api_key == "token-from-nss" # pragma: allowlist secret def test_nim_api_key_falls_back_to_nim_alias(self, monkeypatch): """Legacy NIM_API_KEY loads when NSS_INFERENCE_KEY is unset.""" monkeypatch.delenv("NSS_INFERENCE_KEY", raising=False) - monkeypatch.setenv("NIM_API_KEY", "nim-only") + monkeypatch.setenv("NIM_API_KEY", "token-from-nim") settings = CLISettings() - assert settings.nim_api_key == "nim-only" + assert settings.nim_api_key == "token-from-nim" # pragma: allowlist secret def test_log_color_from_nss_log_color_env(self, monkeypatch): """NSS_LOG_COLOR loads into CLISettings.log_color.""" diff --git a/tests/cli/test_utils.py b/tests/cli/test_utils.py index 8cdc02734..0e727375d 100644 --- a/tests/cli/test_utils.py +++ b/tests/cli/test_utils.py @@ -337,12 +337,12 @@ def test_propagates_nss_inference_settings(self, monkeypatch): settings = CLISettings.from_cli_kwargs( nim_endpoint_url="https://cli.example/v1", - nim_api_key="cli-secret", + nim_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"] == "cli-secret" + 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.""" @@ -367,7 +367,7 @@ def test_common_setup_propagates_before_workdir(self, monkeypatch, dummy_csv: Pa settings = CLISettings.from_cli_kwargs( data_source=str(dummy_csv), - nim_api_key="setup-secret", + nim_api_key="token-propagated-setup", # pragma: allowlist secret ) with ( @@ -381,7 +381,7 @@ def test_common_setup_propagates_before_workdir(self, monkeypatch, dummy_csv: Pa common_setup(settings) - assert os.environ["NSS_INFERENCE_KEY"] == "setup-secret" + assert os.environ["NSS_INFERENCE_KEY"] == "token-propagated-setup" class TestCommonSetupReturnValues: From e94e8cce3dbfb98578b209361538e79b1ee8f2ca Mon Sep 17 00:00:00 2001 From: Aaron Gonzales Date: Mon, 1 Jun 2026 17:58:42 +0000 Subject: [PATCH 04/10] refactor(cli)!: rename PII inference flags to --inference-* with NSS_* env Hard cutover of the env-only PII/NER runtime settings to a consistent NSS-prefixed scheme, addressing reviewer feedback on PR #538: - Flags: --nim-endpoint-url/-api-key/-model-id -> --inference-endpoint-url/ --inference-api-key/--inference-model-id. - Env: NSS_INFERENCE_ENDPOINT/_KEY/_MODEL, plus new NSS_LOCAL_FILES_ONLY and NSS_CPU_COUNT (replacing NIM_MODEL_ID, LOCAL_FILES_ONLY, SAFE_SYNTHESIZER_CPU_COUNT). - Drop legacy NIM_* aliases and the value-equality precedence shim; pydantic AliasChoices now handles CLI > env precedence directly. - Update downstream readers (column classifier, GLiNER loader, NER worker pool) to the new env names. Also from review: - Enforce cpu_count >= 1 at settings parse time. - Remove redundant type=click.BOOL on the --local-files-only flag pair. - docs: restore model-id default text, NSS_* correspondence, table ordering note, and drop anchors that duplicate MkDocs auto-slugs. BREAKING CHANGE: NIM_* env vars and --nim-* flags are removed; use the NSS_INFERENCE_*/NSS_LOCAL_FILES_ONLY/NSS_CPU_COUNT names and --inference-* flags instead. Signed-off-by: Aaron Gonzales --- docs/user-guide/environment.md | 52 +++++++------ docs/user-guide/running.md | 2 +- docs/user-guide/troubleshooting.md | 2 +- script/slurm/slurm_nss_matrix.sh | 2 +- src/nemo_safe_synthesizer/cli/run.py | 55 +++++++------ src/nemo_safe_synthesizer/cli/settings.py | 78 +++++-------------- src/nemo_safe_synthesizer/cli/utils.py | 20 ++--- .../pii_replacer/data_editor/detect.py | 6 +- .../pii_replacer/ner/factory.py | 2 +- tests/cli/test_run.py | 6 +- tests/cli/test_settings.py | 66 ++++++++-------- tests/cli/test_utils.py | 20 ++--- tests/nss_pii_replacer_test.py | 2 +- tests/pii_replacer/test_detect.py | 2 +- tests/test_env_flags.py | 10 +-- 15 files changed, 147 insertions(+), 178 deletions(-) diff --git a/docs/user-guide/environment.md b/docs/user-guide/environment.md index d0c331269..210d9eda9 100644 --- a/docs/user-guide/environment.md +++ b/docs/user-guide/environment.md @@ -33,6 +33,9 @@ For output quality and evaluation metrics, see ## 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) | @@ -44,11 +47,11 @@ For output quality and evaluation metrics, see | `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 | `--nim-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 | `--nim-api-key` | PII column classifier | -- | API key for `NSS_INFERENCE_ENDPOINT` | Required for LLM column classification | -| `NIM_MODEL_ID` | nss | `--nim-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) | -| `LOCAL_FILES_ONLY` | nss | `--local-files-only` / `--no-local-files-only` | GLiNER (PII) | unset | Skip GLiNER network downloads | Partial offline; see [HF appendix](#hugging-face-cache-and-offline) | -| `SAFE_SYNTHESIZER_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) | +| `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_LOCAL_FILES_ONLY` | nss | `--local-files-only` / `--no-local-files-only` | GLiNER (PII) | unset | Skip GLiNER network downloads | Partial offline; see [HF appendix](#hugging-face-cache-and-offline) | +| `NSS_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 | -- | Hugging Face Hub | unset | Fail if a model is not cached | Preferred offline gate | @@ -66,12 +69,12 @@ For output quality and evaluation metrics, see --- -## Precedence {#precedence} +## Precedence ### Infrastructure (CLISettings) For artifact paths, logging, WandB overrides, and the five runtime flags -(`--nim-*`, `--local-files-only`, `--cpu-count`): +(`--inference-*`, `--local-files-only`, `--cpu-count`): 1. CLI flags 2. Environment variables @@ -93,7 +96,7 @@ when explicitly set. When unset, the env var defaults to enabled. --- -## Hugging Face cache and offline {#hugging-face-cache-and-offline} +## Hugging Face cache and offline Downloads go through [Hugging Face Hub](https://huggingface.co/docs/huggingface_hub/guides/manage-cache). @@ -119,24 +122,24 @@ pre-populated `HF_HOME` for reliable offline runs. export HF_HUB_OFFLINE=1 ``` -Prefer this over `LOCAL_FILES_ONLY` for end-to-end offline behavior. +Prefer this over `NSS_LOCAL_FILES_ONLY` for end-to-end offline behavior. -### `LOCAL_FILES_ONLY` +### `NSS_LOCAL_FILES_ONLY` Skips network downloads for GLiNER only. Not respected by the HuggingFace training backend or vLLM. Override on the CLI with `--local-files-only` or `--no-local-files-only`. ```bash -export LOCAL_FILES_ONLY=true +export NSS_LOCAL_FILES_ONLY=true ``` !!! warning "Partial offline support" For the most reliable offline experience, set `HF_HUB_OFFLINE=1` with a - pre-populated `HF_HOME` cache instead of relying on `LOCAL_FILES_ONLY` + pre-populated `HF_HOME` cache instead of relying on `NSS_LOCAL_FILES_ONLY` alone. -### Pre-caching models {#pre-caching-models} +### Pre-caching models Run once with network access, then copy or mount the populated cache. Typical first-run downloads include training weights, GLiNER, evaluation embeddings, @@ -151,7 +154,7 @@ for the full pre-cache checklist. --- -## PII, NER, and column classification {#pii-ner-and-column-classification} +## PII, NER, and column classification Controls LLM-based column classification and CPU parallelism for NER-based PII replacement. For setup examples and NER-only fallback behavior, see @@ -167,30 +170,31 @@ export NSS_INFERENCE_ENDPOINT="https://your-llm-inference-endpoint" export NSS_INFERENCE_KEY="your-api-key" # pragma: allowlist secret ``` -On the CLI, use `--nim-api-key` and optionally `--nim-endpoint-url` instead of -exporting these variables. +On the CLI, can also use `--inference-api-key` and optionally +`--inference-endpoint-url` instead of exporting these variables. 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). -### `NIM_MODEL_ID` +### `NSS_INFERENCE_MODEL` -Model ID sent to the inference endpoint. Override with `--nim-model-id`. +Model ID sent to the inference endpoint. Defaults to +`qwen/qwen3-next-80b-a3b-instruct`. Override with `--inference-model-id`. -### `SAFE_SYNTHESIZER_CPU_COUNT` +### `NSS_CPU_COUNT` 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. ```bash -export SAFE_SYNTHESIZER_CPU_COUNT=4 +export NSS_CPU_COUNT=4 ``` --- -## vLLM and attention {#vllm-and-attention} +## vLLM and attention ### `VLLM_CACHE_ROOT` @@ -215,7 +219,7 @@ Common values: `FLASHINFER`, `FLASH_ATTN`, `TORCH_SDPA`, `TRITON_ATTN`, --- -## Telemetry {#telemetry} +## Telemetry ### `NEMO_TELEMETRY_ENABLED` @@ -235,7 +239,7 @@ equivalent. Intended for controlled test environments. --- -## Containers {#containers} +## Containers Common bind-mount targets when running in Docker: @@ -253,7 +257,7 @@ shortcuts. --- -## Internal and cluster {#internal-and-cluster} +## Internal and cluster Advanced env-only settings without CLI equivalents: diff --git a/docs/user-guide/running.md b/docs/user-guide/running.md index 50133f706..434e95cf1 100644 --- a/docs/user-guide/running.md +++ b/docs/user-guide/running.md @@ -1231,7 +1231,7 @@ See [`artifacts clean`](#artifacts-clean) in the CLI Commands section for option 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`, +and environment variables (`HF_HOME`, `HF_HUB_OFFLINE`, `NSS_LOCAL_FILES_ONLY`, `VLLM_CACHE_ROOT`), see [Environment Variables -- Hugging Face cache and offline](environment.md#hugging-face-cache-and-offline). diff --git a/docs/user-guide/troubleshooting.md b/docs/user-guide/troubleshooting.md index 40639c6fc..897b03e2f 100644 --- a/docs/user-guide/troubleshooting.md +++ b/docs/user-guide/troubleshooting.md @@ -545,7 +545,7 @@ 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, or set `NSS_LOCAL_FILES_ONLY=true` after the model is cached. ### 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 3d98ac34d..19ed87208 100644 --- a/src/nemo_safe_synthesizer/cli/run.py +++ b/src/nemo_safe_synthesizer/cli/run.py @@ -161,32 +161,32 @@ def common_run_options(f: Callable[..., object]) -> Callable[..., object]: ) options.append( click.option( - "--nim-endpoint-url", + "--inference-endpoint-url", type=str, required=False, default=None, - help="NIM/OpenAI-compatible endpoint URL for PII column classification. " + help="OpenAI-compatible inference endpoint URL for PII column classification. " "Can also be set via NSS_INFERENCE_ENDPOINT env var.", ) ) options.append( click.option( - "--nim-api-key", + "--inference-api-key", type=str, required=False, default=None, - help="API key for the NIM endpoint used in PII column classification. " + 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( - "--nim-model-id", + "--inference-model-id", type=str, required=False, default=None, - help="Model ID sent to the NIM endpoint for PII column classification. " - "Can also be set via NIM_MODEL_ID env var. " + 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]", ) ) @@ -194,11 +194,10 @@ def common_run_options(f: Callable[..., object]) -> Callable[..., object]: click.option( "--local-files-only/--no-local-files-only", "local_files_only", - type=click.BOOL, required=False, default=None, help="If set, GLiNER skips network downloads and uses only local files. " - "Can also be set via LOCAL_FILES_ONLY env var.", + "Can also be set via NSS_LOCAL_FILES_ONLY env var.", ) ) options.append( @@ -208,7 +207,7 @@ def common_run_options(f: Callable[..., object]) -> Callable[..., object]: required=False, default=None, help="Number of CPU worker processes used for NER (PII replacement). " - "Can also be set via SAFE_SYNTHESIZER_CPU_COUNT env var. " + "Can also be set via NSS_CPU_COUNT env var. " "[default: max(1, cpu_count - 1)]", ) ) @@ -376,9 +375,9 @@ def run( wandb_mode: str | None = None, wandb_project: str | None = None, dataset_registry: str | None = None, - nim_endpoint_url: str | None = None, - nim_api_key: str | None = None, - nim_model_id: str | None = None, + inference_endpoint_url: str | None = None, + inference_api_key: str | None = None, + inference_model_id: str | None = None, local_files_only: bool | None = None, cpu_count: int | None = None, validate: bool = False, @@ -409,9 +408,9 @@ def run( wandb_project=wandb_project, synthesis_overrides=_parse_run_overrides(kwargs), dataset_registry=dataset_registry, - nim_endpoint_url=nim_endpoint_url, - nim_api_key=nim_api_key, - nim_model_id=nim_model_id, + inference_endpoint_url=inference_endpoint_url, + inference_api_key=inference_api_key, + inference_model_id=inference_model_id, local_files_only=local_files_only, cpu_count=cpu_count, ) @@ -482,9 +481,9 @@ def run_train( wandb_mode: str | None = None, wandb_project: str | None = None, dataset_registry: str | None = None, - nim_endpoint_url: str | None = None, - nim_api_key: str | None = None, - nim_model_id: str | None = None, + inference_endpoint_url: str | None = None, + inference_api_key: str | None = None, + inference_model_id: str | None = None, local_files_only: bool | None = None, cpu_count: int | None = None, validate: bool = False, @@ -511,9 +510,9 @@ def run_train( wandb_project=wandb_project, synthesis_overrides=_parse_run_overrides(kwargs), dataset_registry=dataset_registry, - nim_endpoint_url=nim_endpoint_url, - nim_api_key=nim_api_key, - nim_model_id=nim_model_id, + inference_endpoint_url=inference_endpoint_url, + inference_api_key=inference_api_key, + inference_model_id=inference_model_id, local_files_only=local_files_only, cpu_count=cpu_count, ) @@ -585,9 +584,9 @@ def run_generate( auto_discover_adapter: bool = False, wandb_resume_job_id: str | None = None, dataset_registry: str | None = None, - nim_endpoint_url: str | None = None, - nim_api_key: str | None = None, - nim_model_id: str | None = None, + inference_endpoint_url: str | None = None, + inference_api_key: str | None = None, + inference_model_id: str | None = None, local_files_only: bool | None = None, cpu_count: int | None = None, **kwargs: object, @@ -618,9 +617,9 @@ def run_generate( wandb_project=wandb_project, synthesis_overrides=_parse_run_overrides(kwargs), dataset_registry=dataset_registry, - nim_endpoint_url=nim_endpoint_url, - nim_api_key=nim_api_key, - nim_model_id=nim_model_id, + inference_endpoint_url=inference_endpoint_url, + inference_api_key=inference_api_key, + inference_model_id=inference_model_id, local_files_only=local_files_only, cpu_count=cpu_count, ) diff --git a/src/nemo_safe_synthesizer/cli/settings.py b/src/nemo_safe_synthesizer/cli/settings.py index e9c7b2d5f..4b9bc7ded 100644 --- a/src/nemo_safe_synthesizer/cli/settings.py +++ b/src/nemo_safe_synthesizer/cli/settings.py @@ -27,11 +27,10 @@ from __future__ import annotations -import os from pathlib import Path from typing import Any, Literal -from pydantic import AliasChoices, Field, field_validator, model_validator +from pydantic import AliasChoices, Field, field_validator from pydantic_settings import BaseSettings, SettingsConfigDict from ..defaults import DEFAULT_ARTIFACTS_PATH @@ -40,30 +39,6 @@ __all__ = ["CLISettings"] -# (settings field, canonical env var, legacy env alias from issue #155) -_INFERENCE_ENV_ALIASES: tuple[tuple[str, str, str], ...] = ( - ("nim_endpoint_url", "NSS_INFERENCE_ENDPOINT", "NIM_ENDPOINT_URL"), - ("nim_api_key", "NSS_INFERENCE_KEY", "NIM_API_KEY"), -) - - -def _apply_inference_env_precedence( - data: dict[str, Any], - field: str, - canonical_env: str, - legacy_env: str, -) -> None: - """Prefer ``canonical_env`` over ``legacy_env`` when both are set.""" - canonical = os.environ.get(canonical_env) - legacy = os.environ.get(legacy_env) - match (field in data, data.get(field), canonical, legacy): - case (True, leg, str() as canon, str() as leg_env) if leg == leg_env and canon != leg_env: - data[field] = canon - case (False, _, str() as canon, _): - data[field] = canon - case (False, _, None, str() as leg_env): - data[field] = leg_env - class CLISettings(BaseSettings): """Unified CLI settings composing all sub-settings. @@ -190,56 +165,45 @@ class CLISettings(BaseSettings): ) """URL or path to a dataset registry YAML file (env: ``NSS_DATASET_REGISTRY``).""" - nim_endpoint_url: str | None = Field( + inference_endpoint_url: str | None = Field( default=None, - validation_alias=AliasChoices("nim_endpoint_url", "NSS_INFERENCE_ENDPOINT"), - description="NIM/OpenAI-compatible endpoint URL for PII column classification", + validation_alias=AliasChoices("inference_endpoint_url", "NSS_INFERENCE_ENDPOINT"), + description="OpenAI-compatible inference endpoint URL for PII column classification", ) - """NIM/OpenAI-compatible endpoint URL for PII column classification - (env: ``NSS_INFERENCE_ENDPOINT``; alias: ``NIM_ENDPOINT_URL``).""" + """OpenAI-compatible inference endpoint URL for PII column classification + (env: ``NSS_INFERENCE_ENDPOINT``).""" - nim_api_key: str | None = Field( + inference_api_key: str | None = Field( default=None, - validation_alias=AliasChoices("nim_api_key", "NSS_INFERENCE_KEY"), - description="API key for the NIM endpoint used in PII column classification", + 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 NIM endpoint used in PII column classification - (env: ``NSS_INFERENCE_KEY``; alias: ``NIM_API_KEY``).""" + """API key for the inference endpoint used in PII column classification + (env: ``NSS_INFERENCE_KEY``).""" - nim_model_id: str | None = Field( + inference_model_id: str | None = Field( default=None, - validation_alias=AliasChoices("nim_model_id", "NIM_MODEL_ID"), - description="Model ID sent to the NIM endpoint for PII column classification", + 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 NIM endpoint for PII column classification (env: ``NIM_MODEL_ID``).""" + """Model ID sent to the inference endpoint for PII column classification + (env: ``NSS_INFERENCE_MODEL``).""" local_files_only: bool | None = Field( default=None, - validation_alias=AliasChoices("local_files_only", "LOCAL_FILES_ONLY"), + validation_alias=AliasChoices("local_files_only", "NSS_LOCAL_FILES_ONLY"), description="Whether GLiNER should skip network downloads and use only local files", ) - """Whether GLiNER should skip network downloads and use only local files (env: ``LOCAL_FILES_ONLY``).""" + """Whether GLiNER should skip network downloads and use only local files (env: ``NSS_LOCAL_FILES_ONLY``).""" cpu_count: int | None = Field( default=None, - validation_alias=AliasChoices("cpu_count", "SAFE_SYNTHESIZER_CPU_COUNT"), + ge=1, + validation_alias=AliasChoices("cpu_count", "NSS_CPU_COUNT"), description="Number of CPU worker processes used for NER (PII replacement)", ) """Number of CPU worker processes used for NER (PII replacement) - (env: ``SAFE_SYNTHESIZER_CPU_COUNT``).""" - - @model_validator(mode="before") - @classmethod - def resolve_inference_env_aliases(cls, data: Any) -> Any: - """Prefer ``NSS_INFERENCE_*`` over legacy ``NIM_*`` env aliases.""" - match data: - case dict() as payload: - resolved = dict(payload) - case _: - return data - for field, canonical_env, legacy_env in _INFERENCE_ENV_ALIASES: - _apply_inference_env_precedence(resolved, field, canonical_env, legacy_env) - return resolved + (env: ``NSS_CPU_COUNT``).""" @field_validator("wandb_mode", mode="before") @classmethod diff --git a/src/nemo_safe_synthesizer/cli/utils.py b/src/nemo_safe_synthesizer/cli/utils.py index a3094f7ce..d954e8a4a 100644 --- a/src/nemo_safe_synthesizer/cli/utils.py +++ b/src/nemo_safe_synthesizer/cli/utils.py @@ -232,8 +232,8 @@ def common_setup( """ # 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_*, NIM_MODEL_ID, LOCAL_FILES_ONLY, and - # SAFE_SYNTHESIZER_CPU_COUNT see the CLI-overridden values. + # of NSS_INFERENCE_*, NSS_LOCAL_FILES_ONLY, and NSS_CPU_COUNT see the + # CLI-overridden values. _propagate_runtime_settings_to_env(settings) # 1. Create workdir FIRST - this establishes all artifact paths @@ -337,16 +337,16 @@ def _propagate_runtime_settings_to_env(settings: "CLISettings") -> None: this overwrites ``os.environ`` so the deferred imports in the runtime pipeline see the CLI value. """ - if settings.nim_endpoint_url is not None: - os.environ["NSS_INFERENCE_ENDPOINT"] = settings.nim_endpoint_url - if settings.nim_api_key is not None: - os.environ["NSS_INFERENCE_KEY"] = settings.nim_api_key - if settings.nim_model_id is not None: - os.environ["NIM_MODEL_ID"] = settings.nim_model_id + 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.local_files_only is not None: - os.environ["LOCAL_FILES_ONLY"] = "true" if settings.local_files_only else "false" + os.environ["NSS_LOCAL_FILES_ONLY"] = "true" if settings.local_files_only else "false" if settings.cpu_count is not None: - os.environ["SAFE_SYNTHESIZER_CPU_COUNT"] = str(settings.cpu_count) + os.environ["NSS_CPU_COUNT"] = str(settings.cpu_count) def _initialize_logging_for_cli_from_settings( 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 1e97e05cf..75dbb5162 100644 --- a/src/nemo_safe_synthesizer/pii_replacer/data_editor/detect.py +++ b/src/nemo_safe_synthesizer/pii_replacer/data_editor/detect.py @@ -37,7 +37,7 @@ class DefaultLLMConfig: inference API for column-type classification. Attributes: - CONFIG_ID: Model identifier for the LLM. From env ``NIM_MODEL_ID``, or + CONFIG_ID: Model identifier for the LLM. From env ``NSS_INFERENCE_MODEL``, or ``qwen/qwen3-next-80b-a3b-instruct`` if unset. SYSTEM_PROMPT: System message describing the column-type annotation task sent to the LLM. @@ -47,7 +47,7 @@ class DefaultLLMConfig: Lower values give more deterministic output. """ - CONFIG_ID = os.environ.get("NIM_MODEL_ID", "qwen/qwen3-next-80b-a3b-instruct") + CONFIG_ID = os.environ.get("NSS_INFERENCE_MODEL", "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 @@ -576,7 +576,7 @@ def get_entity_extractor( extractor._model = GLiNER.from_pretrained( clsfy_cfg.gliner_model, map_location=map_location, - local_files_only=env_flag_is_true("LOCAL_FILES_ONLY"), + local_files_only=env_flag_is_true("NSS_LOCAL_FILES_ONLY"), ) entity_types = DEFAULT_ENTITIES if clsfy_cfg.ner_entities: diff --git a/src/nemo_safe_synthesizer/pii_replacer/ner/factory.py b/src/nemo_safe_synthesizer/pii_replacer/ner/factory.py index 7b4810cec..e303d4fdf 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_CPU_COUNT") if num_proc_env: try: num_proc = int(num_proc_env) diff --git a/tests/cli/test_run.py b/tests/cli/test_run.py index 870944ad1..6e14285f1 100644 --- a/tests/cli/test_run.py +++ b/tests/cli/test_run.py @@ -309,9 +309,9 @@ def test_run_help_shows_runtime_settings_options(self, cli_runner: CliRunner): result = cli_runner.invoke(run, ["--help"]) assert result.exit_code == 0 - assert "--nim-endpoint-url" in result.output - assert "--nim-api-key" in result.output - assert "--nim-model-id" in result.output + assert "--inference-endpoint-url" in result.output + assert "--inference-api-key" in result.output + assert "--inference-model-id" in result.output assert "--local-files-only" in result.output assert "--cpu-count" in result.output assert "NSS_INFERENCE_ENDPOINT" in result.output diff --git a/tests/cli/test_settings.py b/tests/cli/test_settings.py index b46a43fb9..7b2c566cb 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,43 +220,29 @@ 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_nim_endpoint_url_from_nss_inference_env(self, monkeypatch): - """NSS_INFERENCE_ENDPOINT loads into nim_endpoint_url.""" + 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.nim_endpoint_url == "https://custom.example/v1" + assert settings.inference_endpoint_url == "https://custom.example/v1" - def test_nim_api_key_from_nss_inference_env(self, monkeypatch): - """NSS_INFERENCE_KEY loads into nim_api_key.""" + 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.nim_api_key == "token-from-env" # pragma: allowlist secret + assert settings.inference_api_key == "token-from-env" # pragma: allowlist secret - def test_nim_endpoint_url_cli_overrides_env(self, monkeypatch): - """CLI --nim-endpoint-url takes precedence over NSS_INFERENCE_ENDPOINT.""" + 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(nim_endpoint_url="https://cli.example/v1") - assert settings.nim_endpoint_url == "https://cli.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_nim_api_key_cli_overrides_env(self, monkeypatch): - """CLI --nim-api-key takes precedence over NSS_INFERENCE_KEY.""" + 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(nim_api_key="token-from-cli") # pragma: allowlist secret - assert settings.nim_api_key == "token-from-cli" # pragma: allowlist secret - - def test_nim_api_key_prefers_nss_over_nim_alias_when_both_set(self, monkeypatch): - """Canonical NSS_INFERENCE_KEY wins when both NSS and NIM aliases are set.""" - monkeypatch.setenv("NSS_INFERENCE_KEY", "token-from-nss") - monkeypatch.setenv("NIM_API_KEY", "token-from-nim") - settings = CLISettings() - assert settings.nim_api_key == "token-from-nss" # pragma: allowlist secret - - def test_nim_api_key_falls_back_to_nim_alias(self, monkeypatch): - """Legacy NIM_API_KEY loads when NSS_INFERENCE_KEY is unset.""" - monkeypatch.delenv("NSS_INFERENCE_KEY", raising=False) - monkeypatch.setenv("NIM_API_KEY", "token-from-nim") - settings = CLISettings() - assert settings.nim_api_key == "token-from-nim" # pragma: allowlist secret + 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.""" @@ -270,15 +259,28 @@ def test_log_color_cli_overrides_nss_log_color_env(self, monkeypatch): def test_runtime_settings_from_env(self, monkeypatch): """Remaining runtime settings load from their documented env vars.""" - monkeypatch.setenv("NIM_MODEL_ID", "custom/model") - monkeypatch.setenv("LOCAL_FILES_ONLY", "true") - monkeypatch.setenv("SAFE_SYNTHESIZER_CPU_COUNT", "4") + monkeypatch.setenv("NSS_INFERENCE_MODEL", "custom/model") + monkeypatch.setenv("NSS_LOCAL_FILES_ONLY", "true") + monkeypatch.setenv("NSS_CPU_COUNT", "4") settings = CLISettings() - assert settings.nim_model_id == "custom/model" + assert settings.inference_model_id == "custom/model" assert settings.local_files_only is True assert settings.cpu_count == 4 + @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_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 0e727375d..1868a0fa0 100644 --- a/tests/cli/test_utils.py +++ b/tests/cli/test_utils.py @@ -336,8 +336,8 @@ def test_propagates_nss_inference_settings(self, monkeypatch): monkeypatch.delenv("NSS_INFERENCE_KEY", raising=False) settings = CLISettings.from_cli_kwargs( - nim_endpoint_url="https://cli.example/v1", - nim_api_key="token-propagated-cli", # pragma: allowlist secret + inference_endpoint_url="https://cli.example/v1", + inference_api_key="token-propagated-cli", # pragma: allowlist secret ) _propagate_runtime_settings_to_env(settings) @@ -346,20 +346,20 @@ def test_propagates_nss_inference_settings(self, monkeypatch): def test_propagates_remaining_runtime_settings(self, monkeypatch): """Model ID, offline mode, and CPU count propagate to their runtime env vars.""" - monkeypatch.delenv("NIM_MODEL_ID", raising=False) - monkeypatch.delenv("LOCAL_FILES_ONLY", raising=False) - monkeypatch.delenv("SAFE_SYNTHESIZER_CPU_COUNT", raising=False) + monkeypatch.delenv("NSS_INFERENCE_MODEL", raising=False) + monkeypatch.delenv("NSS_LOCAL_FILES_ONLY", raising=False) + monkeypatch.delenv("NSS_CPU_COUNT", raising=False) settings = CLISettings.from_cli_kwargs( - nim_model_id="custom/model", + inference_model_id="custom/model", local_files_only=True, cpu_count=3, ) _propagate_runtime_settings_to_env(settings) - assert os.environ["NIM_MODEL_ID"] == "custom/model" - assert os.environ["LOCAL_FILES_ONLY"] == "true" - assert os.environ["SAFE_SYNTHESIZER_CPU_COUNT"] == "3" + assert os.environ["NSS_INFERENCE_MODEL"] == "custom/model" + assert os.environ["NSS_LOCAL_FILES_ONLY"] == "true" + assert os.environ["NSS_CPU_COUNT"] == "3" def test_common_setup_propagates_before_workdir(self, monkeypatch, dummy_csv: Path): """common_setup writes resolved runtime settings before downstream imports.""" @@ -367,7 +367,7 @@ def test_common_setup_propagates_before_workdir(self, monkeypatch, dummy_csv: Pa settings = CLISettings.from_cli_kwargs( data_source=str(dummy_csv), - nim_api_key="token-propagated-setup", # pragma: allowlist secret + inference_api_key="token-propagated-setup", # pragma: allowlist secret ) with ( 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 7ad970f00..0acc005e0 100644 --- a/tests/pii_replacer/test_detect.py +++ b/tests/pii_replacer/test_detect.py @@ -81,7 +81,7 @@ def test_gliner_local_files_only_accepts_common_truthy_env_values(env_value, mon gliner_batch_mode_batch_size=20, gliner_model="nvidia/gliner-PII", ) - monkeypatch.setenv("LOCAL_FILES_ONLY", env_value) + monkeypatch.setenv("NSS_LOCAL_FILES_ONLY", env_value) with patch("nemo_safe_synthesizer.pii_replacer.data_editor.detect.GLiNER") as mock_gliner: EntityExtractorGliner.get_entity_extractor(cfg) diff --git a/tests/test_env_flags.py b/tests/test_env_flags.py index d2d0c0acc..fbfa1185a 100644 --- a/tests/test_env_flags.py +++ b/tests/test_env_flags.py @@ -25,11 +25,11 @@ ], ) def test_env_flag_is_true(value: str, expected: bool, monkeypatch): - monkeypatch.setenv("LOCAL_FILES_ONLY", value) - assert env_flag_is_true("LOCAL_FILES_ONLY") is expected + monkeypatch.setenv("NSS_LOCAL_FILES_ONLY", value) + assert env_flag_is_true("NSS_LOCAL_FILES_ONLY") is expected def test_env_flag_is_true_unset_uses_default(monkeypatch): - monkeypatch.delenv("LOCAL_FILES_ONLY", raising=False) - assert env_flag_is_true("LOCAL_FILES_ONLY") is False - assert env_flag_is_true("LOCAL_FILES_ONLY", default=True) is True + monkeypatch.delenv("NSS_LOCAL_FILES_ONLY", raising=False) + assert env_flag_is_true("NSS_LOCAL_FILES_ONLY") is False + assert env_flag_is_true("NSS_LOCAL_FILES_ONLY", default=True) is True From b31405a82b1921cfdcea4a0ed792708faf1a3347 Mon Sep 17 00:00:00 2001 From: Aaron Gonzales Date: Mon, 1 Jun 2026 18:11:06 +0000 Subject: [PATCH 05/10] refactor(cli): rename NSS_CPU_COUNT to NSS_PII_REPLACER_CPU_COUNT Make the NER worker-count env var self-describing and scoped to the PII replacer. The --cpu-count flag and cpu_count field are unchanged. Signed-off-by: Aaron Gonzales --- docs/user-guide/environment.md | 6 +++--- src/nemo_safe_synthesizer/cli/run.py | 2 +- src/nemo_safe_synthesizer/cli/settings.py | 4 ++-- src/nemo_safe_synthesizer/cli/utils.py | 6 +++--- src/nemo_safe_synthesizer/pii_replacer/ner/factory.py | 2 +- tests/cli/test_settings.py | 4 ++-- tests/cli/test_utils.py | 4 ++-- 7 files changed, 14 insertions(+), 14 deletions(-) diff --git a/docs/user-guide/environment.md b/docs/user-guide/environment.md index 210d9eda9..b003bcc7f 100644 --- a/docs/user-guide/environment.md +++ b/docs/user-guide/environment.md @@ -51,7 +51,7 @@ Grouped by the `Category` column -- `nss`-native settings first, then | `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_LOCAL_FILES_ONLY` | nss | `--local-files-only` / `--no-local-files-only` | GLiNER (PII) | unset | Skip GLiNER network downloads | Partial offline; see [HF appendix](#hugging-face-cache-and-offline) | -| `NSS_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) | +| `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 | -- | Hugging Face Hub | unset | Fail if a model is not cached | Preferred offline gate | @@ -182,14 +182,14 @@ See [Configuration Reference -- Replacing PII](configuration.md#replacing-pii). Model ID sent to the inference endpoint. Defaults to `qwen/qwen3-next-80b-a3b-instruct`. Override with `--inference-model-id`. -### `NSS_CPU_COUNT` +### `NSS_PII_REPLACER_CPU_COUNT` 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. ```bash -export NSS_CPU_COUNT=4 +export NSS_PII_REPLACER_CPU_COUNT=4 ``` --- diff --git a/src/nemo_safe_synthesizer/cli/run.py b/src/nemo_safe_synthesizer/cli/run.py index 19ed87208..1ff9c5e67 100644 --- a/src/nemo_safe_synthesizer/cli/run.py +++ b/src/nemo_safe_synthesizer/cli/run.py @@ -207,7 +207,7 @@ def common_run_options(f: Callable[..., object]) -> Callable[..., object]: required=False, default=None, help="Number of CPU worker processes used for NER (PII replacement). " - "Can also be set via NSS_CPU_COUNT env var. " + "Can also be set via NSS_PII_REPLACER_CPU_COUNT env var. " "[default: max(1, cpu_count - 1)]", ) ) diff --git a/src/nemo_safe_synthesizer/cli/settings.py b/src/nemo_safe_synthesizer/cli/settings.py index 4b9bc7ded..021192b8a 100644 --- a/src/nemo_safe_synthesizer/cli/settings.py +++ b/src/nemo_safe_synthesizer/cli/settings.py @@ -199,11 +199,11 @@ class CLISettings(BaseSettings): cpu_count: int | None = Field( default=None, ge=1, - validation_alias=AliasChoices("cpu_count", "NSS_CPU_COUNT"), + 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_CPU_COUNT``).""" + (env: ``NSS_PII_REPLACER_CPU_COUNT``).""" @field_validator("wandb_mode", mode="before") @classmethod diff --git a/src/nemo_safe_synthesizer/cli/utils.py b/src/nemo_safe_synthesizer/cli/utils.py index d954e8a4a..95597feae 100644 --- a/src/nemo_safe_synthesizer/cli/utils.py +++ b/src/nemo_safe_synthesizer/cli/utils.py @@ -232,8 +232,8 @@ def common_setup( """ # 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_*, NSS_LOCAL_FILES_ONLY, and NSS_CPU_COUNT see the - # CLI-overridden values. + # of NSS_INFERENCE_*, NSS_LOCAL_FILES_ONLY, 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 @@ -346,7 +346,7 @@ def _propagate_runtime_settings_to_env(settings: "CLISettings") -> None: if settings.local_files_only is not None: os.environ["NSS_LOCAL_FILES_ONLY"] = "true" if settings.local_files_only else "false" if settings.cpu_count is not None: - os.environ["NSS_CPU_COUNT"] = str(settings.cpu_count) + os.environ["NSS_PII_REPLACER_CPU_COUNT"] = str(settings.cpu_count) def _initialize_logging_for_cli_from_settings( diff --git a/src/nemo_safe_synthesizer/pii_replacer/ner/factory.py b/src/nemo_safe_synthesizer/pii_replacer/ner/factory.py index e303d4fdf..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("NSS_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/tests/cli/test_settings.py b/tests/cli/test_settings.py index 7b2c566cb..82fa36e43 100644 --- a/tests/cli/test_settings.py +++ b/tests/cli/test_settings.py @@ -261,7 +261,7 @@ 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_LOCAL_FILES_ONLY", "true") - monkeypatch.setenv("NSS_CPU_COUNT", "4") + monkeypatch.setenv("NSS_PII_REPLACER_CPU_COUNT", "4") settings = CLISettings() assert settings.inference_model_id == "custom/model" @@ -271,7 +271,7 @@ def test_runtime_settings_from_env(self, monkeypatch): @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_CPU_COUNT", bad_value) + monkeypatch.setenv("NSS_PII_REPLACER_CPU_COUNT", bad_value) with pytest.raises(ValidationError): CLISettings() diff --git a/tests/cli/test_utils.py b/tests/cli/test_utils.py index 1868a0fa0..9dc44fef4 100644 --- a/tests/cli/test_utils.py +++ b/tests/cli/test_utils.py @@ -348,7 +348,7 @@ 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("NSS_LOCAL_FILES_ONLY", raising=False) - monkeypatch.delenv("NSS_CPU_COUNT", raising=False) + monkeypatch.delenv("NSS_PII_REPLACER_CPU_COUNT", raising=False) settings = CLISettings.from_cli_kwargs( inference_model_id="custom/model", @@ -359,7 +359,7 @@ def test_propagates_remaining_runtime_settings(self, monkeypatch): assert os.environ["NSS_INFERENCE_MODEL"] == "custom/model" assert os.environ["NSS_LOCAL_FILES_ONLY"] == "true" - assert os.environ["NSS_CPU_COUNT"] == "3" + assert os.environ["NSS_PII_REPLACER_CPU_COUNT"] == "3" def test_common_setup_propagates_before_workdir(self, monkeypatch, dummy_csv: Path): """common_setup writes resolved runtime settings before downstream imports.""" From 92588e8e4b1342e7700bba0c810e23742fa1cbca Mon Sep 17 00:00:00 2001 From: Aaron Gonzales Date: Mon, 1 Jun 2026 18:20:01 +0000 Subject: [PATCH 06/10] docs(env): rename duplicate Telemetry heading to fix MD024 The Precedence subsection and the main section both rendered as "Telemetry" after the redundant {#anchor} cleanup, tripping markdownlint MD024. Rename the precedence subsection to "Telemetry precedence"; the main section keeps its #telemetry slug so existing links still resolve. Signed-off-by: Aaron Gonzales --- docs/user-guide/environment.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/user-guide/environment.md b/docs/user-guide/environment.md index b003bcc7f..5d8fc26b6 100644 --- a/docs/user-guide/environment.md +++ b/docs/user-guide/environment.md @@ -89,7 +89,7 @@ YAML fields, CLI `--section__field` overrides, and SDK builder calls follow [Configuration Precedence](configuration.md#configuration-precedence) -- not the order above. -### Telemetry +### Telemetry precedence `--emit_telemetry` / `emit_telemetry` in YAML override `NEMO_TELEMETRY_ENABLED` when explicitly set. When unset, the env var defaults to enabled. From c1ab31590fab92daee11a9a9675dc18bb1f38a05 Mon Sep 17 00:00:00 2001 From: Aaron Gonzales Date: Mon, 1 Jun 2026 18:29:53 +0000 Subject: [PATCH 07/10] refactor(cli): fold shared run flags into kwargs-driven settings builder Collapse run/train/generate signatures to command-specific params plus **kwargs, and add _settings_from_run_kwargs to split shared flags from synthesis overrides via CLISettings.model_fields. Adding a shared flag now touches common_run_options and CLISettings only. A guard test asserts every shared flag name maps to a CLISettings field. Signed-off-by: Aaron Gonzales --- src/nemo_safe_synthesizer/cli/run.py | 145 ++++++--------------------- tests/cli/test_run.py | 20 ++++ 2 files changed, 51 insertions(+), 114 deletions(-) diff --git a/src/nemo_safe_synthesizer/cli/run.py b/src/nemo_safe_synthesizer/cli/run.py index 1ff9c5e67..2cec01a98 100644 --- a/src/nemo_safe_synthesizer/cli/run.py +++ b/src/nemo_safe_synthesizer/cli/run.py @@ -222,6 +222,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) @@ -363,25 +388,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, - inference_endpoint_url: str | None = None, - inference_api_key: str | None = None, - inference_model_id: str | None = None, - local_files_only: bool | None = None, - cpu_count: int | None = None, validate: bool = False, - **kwargs: object, + **kwargs: Any, ) -> None: """Run the Safe Synthesizer end-to-end pipeline. @@ -394,26 +402,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, - inference_endpoint_url=inference_endpoint_url, - inference_api_key=inference_api_key, - inference_model_id=inference_model_id, - local_files_only=local_files_only, - cpu_count=cpu_count, - ) + settings = _settings_from_run_kwargs(kwargs) if validate: os.environ["NSS_PHASE"] = "process_data" @@ -469,25 +458,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, - inference_endpoint_url: str | None = None, - inference_api_key: str | None = None, - inference_model_id: str | None = None, - local_files_only: bool | None = None, - cpu_count: int | None = None, validate: bool = False, - **kwargs: object, + **kwargs: Any, ) -> None: """Run the training stage only. @@ -496,26 +468,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, - inference_endpoint_url=inference_endpoint_url, - inference_api_key=inference_api_key, - inference_model_id=inference_model_id, - local_files_only=local_files_only, - cpu_count=cpu_count, - ) + settings = _settings_from_run_kwargs(kwargs) if validate: os.environ["NSS_PHASE"] = "process_data" @@ -570,26 +523,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, - inference_endpoint_url: str | None = None, - inference_api_key: str | None = None, - inference_model_id: str | None = None, - local_files_only: bool | None = None, - cpu_count: int | None = None, - **kwargs: object, + **kwargs: Any, ) -> None: """Run the generation stage only. @@ -603,26 +539,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, - inference_endpoint_url=inference_endpoint_url, - inference_api_key=inference_api_key, - inference_model_id=inference_model_id, - local_files_only=local_files_only, - cpu_count=cpu_count, - ) + 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/tests/cli/test_run.py b/tests/cli/test_run.py index 6e14285f1..8d8d9e4cf 100644 --- a/tests/cli/test_run.py +++ b/tests/cli/test_run.py @@ -1019,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)}" From 883e40a590c8e5d2602a181f5985fe4a86bf7720 Mon Sep 17 00:00:00 2001 From: Aaron Gonzales Date: Mon, 1 Jun 2026 23:14:14 +0000 Subject: [PATCH 08/10] feat(cli): add InferenceModelCheck and keep CLI import hub-free for offline flag Consolidate inference env validation into a single preflight check and make the Hugging Face offline switch reliable end to end. - preflight: rename InferenceKeyCheck to InferenceModelCheck (env.inference); validate NSS_INFERENCE_KEY, NSS_INFERENCE_MODEL, and NSS_INFERENCE_ENDPOINT via single-dispatch match logic. - cli: replace --local-files-only with --enable/--disable-huggingface-remote (CLI-only, no NSS env var); propagate to HF_HUB_OFFLINE/TRANSFORMERS_OFFLINE. - utils: add shared hf_offline_enabled() and env_flag_is_true(); detect.py reads NSS_INFERENCE_MODEL at call time and GLiNER offline from env. - imports: defer huggingface_hub in telemetry and datasets in utils so the cli.cli import chain stays hub-free; HF_HUB_OFFLINE is then propagated before huggingface_hub first loads. Add tests/cli/test_cli_import regression guard. - docs: document the offline switch, CLI flag precedence, and import-time caching of HF_HUB_OFFLINE. Signed-off-by: Aaron Gonzales --- docs/user-guide/environment.md | 42 +++++---- docs/user-guide/running.md | 10 ++- docs/user-guide/troubleshooting.md | 28 +++++- src/nemo_safe_synthesizer/cli/run.py | 10 ++- src/nemo_safe_synthesizer/cli/settings.py | 13 ++- src/nemo_safe_synthesizer/cli/utils.py | 24 ++++- .../pii_replacer/data_editor/detect.py | 25 ++++-- .../preflight/__init__.py | 4 +- .../preflight/checks/__init__.py | 6 +- .../preflight/checks/environment.py | 84 ++++++++++++----- src/nemo_safe_synthesizer/telemetry.py | 8 +- src/nemo_safe_synthesizer/utils.py | 21 ++++- tests/cli/test_cli_import.py | 36 ++++++++ tests/cli/test_run.py | 2 +- tests/cli/test_settings.py | 7 +- tests/cli/test_utils.py | 19 +++- tests/pii_replacer/test_detect.py | 9 +- tests/preflight/test_preflight.py | 90 +++++++++++++++++-- tests/test_env_flags.py | 28 ++++-- 19 files changed, 370 insertions(+), 96 deletions(-) create mode 100644 tests/cli/test_cli_import.py diff --git a/docs/user-guide/environment.md b/docs/user-guide/environment.md index 5d8fc26b6..c4a55d659 100644 --- a/docs/user-guide/environment.md +++ b/docs/user-guide/environment.md @@ -50,11 +50,10 @@ Grouped by the `Category` column -- `nss`-native settings first, then | `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_LOCAL_FILES_ONLY` | nss | `--local-files-only` / `--no-local-files-only` | GLiNER (PII) | unset | Skip GLiNER network downloads | Partial offline; see [HF appendix](#hugging-face-cache-and-offline) | | `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 | -- | Hugging Face Hub | unset | Fail if a model is not cached | Preferred offline gate | +| `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` | @@ -74,7 +73,8 @@ Grouped by the `Category` column -- `nss`-native settings first, then ### Infrastructure (CLISettings) For artifact paths, logging, WandB overrides, and the five runtime flags -(`--inference-*`, `--local-files-only`, `--cpu-count`): +(`--inference-*`, `--enable-huggingface-remote` / `--disable-huggingface-remote`, +`--cpu-count`): 1. CLI flags 2. Environment variables @@ -115,29 +115,41 @@ export HF_HOME=/shared/cache/huggingface ### `HF_HUB_OFFLINE` -When set to `1`, Hugging Face Hub refuses network access. Use with a -pre-populated `HF_HOME` for reliable offline runs. +`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 `NSS_LOCAL_FILES_ONLY` for end-to-end offline behavior. +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`. -### `NSS_LOCAL_FILES_ONLY` +### `--enable-huggingface-remote` / `--disable-huggingface-remote` -Skips network downloads for GLiNER only. Not respected by the HuggingFace -training backend or vLLM. Override on the CLI with `--local-files-only` or -`--no-local-files-only`. +CLI shorthand for the switch above, with no separate NSS env var: + +- `--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. + +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 NSS_LOCAL_FILES_ONLY=true +safe-synthesizer run --disable-huggingface-remote ... ``` -!!! warning "Partial offline support" - For the most reliable offline experience, set `HF_HUB_OFFLINE=1` with a - pre-populated `HF_HOME` cache instead of relying on `NSS_LOCAL_FILES_ONLY` - alone. +!!! 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. ### Pre-caching models diff --git a/docs/user-guide/running.md b/docs/user-guide/running.md index 434e95cf1..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,9 +1230,11 @@ 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`, `NSS_LOCAL_FILES_ONLY`, -`VLLM_CACHE_ROOT`), see +`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 897b03e2f..124a5f346 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; classification would send an empty model id and fail | +| `inference_endpoint_invalid` | warning | `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 `NSS_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/src/nemo_safe_synthesizer/cli/run.py b/src/nemo_safe_synthesizer/cli/run.py index 2cec01a98..ceb554e73 100644 --- a/src/nemo_safe_synthesizer/cli/run.py +++ b/src/nemo_safe_synthesizer/cli/run.py @@ -192,12 +192,14 @@ def common_run_options(f: Callable[..., object]) -> Callable[..., object]: ) options.append( click.option( - "--local-files-only/--no-local-files-only", - "local_files_only", + "--enable-huggingface-remote/--disable-huggingface-remote", + "huggingface_remote", required=False, default=None, - help="If set, GLiNER skips network downloads and uses only local files. " - "Can also be set via NSS_LOCAL_FILES_ONLY env var.", + 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.", ) ) options.append( diff --git a/src/nemo_safe_synthesizer/cli/settings.py b/src/nemo_safe_synthesizer/cli/settings.py index 021192b8a..bdcb5fb36 100644 --- a/src/nemo_safe_synthesizer/cli/settings.py +++ b/src/nemo_safe_synthesizer/cli/settings.py @@ -189,12 +189,17 @@ class CLISettings(BaseSettings): """Model ID sent to the inference endpoint for PII column classification (env: ``NSS_INFERENCE_MODEL``).""" - local_files_only: bool | None = Field( + huggingface_remote: bool | None = Field( default=None, - validation_alias=AliasChoices("local_files_only", "NSS_LOCAL_FILES_ONLY"), - description="Whether GLiNER should skip network downloads and use only local files", + validation_alias=AliasChoices("huggingface_remote"), + description="Whether to allow Hugging Face remote downloads (base model and GLiNER)", ) - """Whether GLiNER should skip network downloads and use only local files (env: ``NSS_LOCAL_FILES_ONLY``).""" + """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, diff --git a/src/nemo_safe_synthesizer/cli/utils.py b/src/nemo_safe_synthesizer/cli/utils.py index 95597feae..fad664d04 100644 --- a/src/nemo_safe_synthesizer/cli/utils.py +++ b/src/nemo_safe_synthesizer/cli/utils.py @@ -232,8 +232,8 @@ def common_setup( """ # 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_*, NSS_LOCAL_FILES_ONLY, and NSS_PII_REPLACER_CPU_COUNT - # see the CLI-overridden values. + # 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 @@ -336,6 +336,20 @@ def _propagate_runtime_settings_to_env(settings: "CLISettings") -> None: 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 @@ -343,8 +357,10 @@ def _propagate_runtime_settings_to_env(settings: "CLISettings") -> 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.local_files_only is not None: - os.environ["NSS_LOCAL_FILES_ONLY"] = "true" if settings.local_files_only else "false" + 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) 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 75dbb5162..9c045ae5a 100644 --- a/src/nemo_safe_synthesizer/pii_replacer/data_editor/detect.py +++ b/src/nemo_safe_synthesizer/pii_replacer/data_editor/detect.py @@ -21,7 +21,7 @@ from pydantic import ConfigDict, TypeAdapter, ValidationError from ...observability import get_logger -from ...utils import env_flag_is_true +from ...utils import hf_offline_enabled from ..ner import ner_mp from ..ner.factory import LabelSetPredictorFilter, NERFactory from ..ner.ner import NERPrediction @@ -33,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 ``NSS_INFERENCE_MODEL``, 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 @@ -47,11 +45,22 @@ class DefaultLLMConfig: Lower values give more deterministic output. """ - CONFIG_ID = os.environ.get("NSS_INFERENCE_MODEL", "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. + """ + return os.environ.get("NSS_INFERENCE_MODEL", cls.DEFAULT_CONFIG_ID) + DEFAULT_ENTITIES: set[str] = { "name", @@ -250,7 +259,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}, @@ -576,7 +585,7 @@ def get_entity_extractor( extractor._model = GLiNER.from_pretrained( clsfy_cfg.gliner_model, map_location=map_location, - local_files_only=env_flag_is_true("NSS_LOCAL_FILES_ONLY"), + local_files_only=hf_offline_enabled(), ) entity_types = DEFAULT_ENTITIES if clsfy_cfg.ner_entities: 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..619e43550 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,70 @@ 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``. All findings are warnings: classification + degrades or fails at call time rather than blocking the run outright. + + The body uses a single-dispatch ``match`` over ``(model, key, endpoint)``, + so at most one warning is emitted per run -- the highest-priority problem. + Priority order: missing key, then blank model id, then invalid endpoint. + """ + + 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 + # priority. A missing key (degraded mode) is reported before a blank + # model id or an invalid endpoint (hard failures at call time). + match model, key, endpoint: + 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. PII column classification will send an " + "empty model id and fail. Unset it to use the default, or provide a model id.", + ) + case _, _, e if e is not None and not _is_valid_http_url(e): + collector.warning( + "inference_endpoint_invalid", + f"NSS_INFERENCE_ENDPOINT '{e}' is not a valid http(s) URL. " + "PII column classification requests will fail.", + ) def _has_hf_token() -> bool: @@ -410,7 +448,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 +483,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 +508,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 153f7dafe..93ac5bac8 100644 --- a/src/nemo_safe_synthesizer/utils.py +++ b/src/nemo_safe_synthesizer/utils.py @@ -15,20 +15,27 @@ 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. @@ -42,6 +49,16 @@ def env_flag_is_true(name: str, *, default: bool = False) -> bool: 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 8d8d9e4cf..4978d1d16 100644 --- a/tests/cli/test_run.py +++ b/tests/cli/test_run.py @@ -312,7 +312,7 @@ def test_run_help_shows_runtime_settings_options(self, cli_runner: CliRunner): assert "--inference-endpoint-url" in result.output assert "--inference-api-key" in result.output assert "--inference-model-id" in result.output - assert "--local-files-only" 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 diff --git a/tests/cli/test_settings.py b/tests/cli/test_settings.py index 82fa36e43..24b60badf 100644 --- a/tests/cli/test_settings.py +++ b/tests/cli/test_settings.py @@ -260,14 +260,17 @@ def test_log_color_cli_overrides_nss_log_color_env(self, monkeypatch): 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_LOCAL_FILES_ONLY", "true") monkeypatch.setenv("NSS_PII_REPLACER_CPU_COUNT", "4") settings = CLISettings() assert settings.inference_model_id == "custom/model" - assert settings.local_files_only is True 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.""" diff --git a/tests/cli/test_utils.py b/tests/cli/test_utils.py index 9dc44fef4..3a694f67e 100644 --- a/tests/cli/test_utils.py +++ b/tests/cli/test_utils.py @@ -347,20 +347,33 @@ def test_propagates_nss_inference_settings(self, monkeypatch): 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("NSS_LOCAL_FILES_ONLY", 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", - local_files_only=True, + huggingface_remote=False, cpu_count=3, ) _propagate_runtime_settings_to_env(settings) assert os.environ["NSS_INFERENCE_MODEL"] == "custom/model" - assert os.environ["NSS_LOCAL_FILES_ONLY"] == "true" + 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) diff --git a/tests/pii_replacer/test_detect.py b/tests/pii_replacer/test_detect.py index 0acc005e0..7bc5e4b02 100644 --- a/tests/pii_replacer/test_detect.py +++ b/tests/pii_replacer/test_detect.py @@ -67,9 +67,10 @@ 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_accepts_common_truthy_env_values(env_value, monkeypatch): - """GLiNER offline mode accepts the same truthy spellings as env_flag_is_true.""" +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, @@ -81,7 +82,9 @@ def test_gliner_local_files_only_accepts_common_truthy_env_values(env_value, mon gliner_batch_mode_batch_size=20, gliner_model="nvidia/gliner-PII", ) - monkeypatch.setenv("NSS_LOCAL_FILES_ONLY", env_value) + 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) diff --git a/tests/preflight/test_preflight.py b/tests/preflight/test_preflight.py index 656acb632..c7d3e169a 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,85 @@ 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_warning(self, default_config, endpoint): + 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 == "warning" 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) + + def test_missing_key_takes_priority_over_other_problems(self, default_config): + # Single-dispatch match: the first matching case wins, so a missing key + # is reported even when the model id and endpoint are also bad. + 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_key_missing"} + + def test_blank_model_takes_priority_over_invalid_endpoint(self, default_config): + with patch.dict( + "os.environ", + {"NSS_INFERENCE_KEY": "test-key", "NSS_INFERENCE_MODEL": " ", "NSS_INFERENCE_ENDPOINT": "not-a-url"}, + ): + issues = InferenceModelCheck().run(make_ctx(config=default_config)) + codes = {i.code for i in issues} + assert codes == {"inference_model_blank"} @pytest.mark.unit @@ -725,6 +788,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 index fbfa1185a..38ff9f61c 100644 --- a/tests/test_env_flags.py +++ b/tests/test_env_flags.py @@ -7,7 +7,9 @@ import pytest -from nemo_safe_synthesizer.utils import env_flag_is_true +from nemo_safe_synthesizer.utils import env_flag_is_true, hf_offline_enabled + +_PROBE_VAR = "NSS_TEST_FLAG" @pytest.mark.parametrize( @@ -25,11 +27,25 @@ ], ) def test_env_flag_is_true(value: str, expected: bool, monkeypatch): - monkeypatch.setenv("NSS_LOCAL_FILES_ONLY", value) - assert env_flag_is_true("NSS_LOCAL_FILES_ONLY") is expected + 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("NSS_LOCAL_FILES_ONLY", raising=False) - assert env_flag_is_true("NSS_LOCAL_FILES_ONLY") is False - assert env_flag_is_true("NSS_LOCAL_FILES_ONLY", default=True) is True + 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 From 2ca6f4a755b32c4c28fe8cd2e9c555a3b9819d47 Mon Sep 17 00:00:00 2001 From: Aaron Gonzales Date: Tue, 2 Jun 2026 00:14:34 +0000 Subject: [PATCH 09/10] fix(inference): normalize blank inference env vars and fail invalid endpoint Align runtime and preflight handling of the PII column-classification inference settings: - detect.config_id(): strip NSS_INFERENCE_MODEL; blank/whitespace falls back to DEFAULT_CONFIG_ID instead of sending an empty model id. - nemo_pii._get_classify_endpoint_url(): strip NSS_INFERENCE_ENDPOINT; blank falls back to DEFAULT_NSS_INFERENCE_ENDPOINT instead of passing an empty base_url to the OpenAI client. - preflight env.inference: a non-http(s) NSS_INFERENCE_ENDPOINT is now an error (must not pass --validate), checked before the missing-key and blank-model warnings; a blank endpoint is ignored. Update troubleshooting table and tests accordingly. Signed-off-by: Aaron Gonzales --- docs/user-guide/troubleshooting.md | 4 +-- .../pii_replacer/data_editor/detect.py | 7 ++-- .../pii_replacer/nemo_pii.py | 11 +++--- .../preflight/checks/environment.py | 32 +++++++++-------- tests/pii_replacer/test_detect.py | 20 +++++++++++ tests/pii_replacer/test_nemo_pii.py | 29 ++++++++++++++- tests/preflight/test_preflight.py | 36 ++++++++++++++----- 7 files changed, 107 insertions(+), 32 deletions(-) diff --git a/docs/user-guide/troubleshooting.md b/docs/user-guide/troubleshooting.md index 124a5f346..ce231f6d7 100644 --- a/docs/user-guide/troubleshooting.md +++ b/docs/user-guide/troubleshooting.md @@ -514,8 +514,8 @@ check of its own. | `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` | `NSS_INFERENCE_KEY` not set; PII classification degraded | -| `inference_model_blank` | warning | `env.inference` | `NSS_INFERENCE_MODEL` set but empty; classification would send an empty model id and fail | -| `inference_endpoint_invalid` | warning | `env.inference` | `NSS_INFERENCE_ENDPOINT` set but not a valid http(s) URL; classification requests will fail | +| `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 | 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 9c045ae5a..2bba807ee 100644 --- a/src/nemo_safe_synthesizer/pii_replacer/data_editor/detect.py +++ b/src/nemo_safe_synthesizer/pii_replacer/data_editor/detect.py @@ -57,9 +57,12 @@ def config_id(cls) -> str: 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. + 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. """ - return os.environ.get("NSS_INFERENCE_MODEL", cls.DEFAULT_CONFIG_ID) + model = os.environ.get("NSS_INFERENCE_MODEL", "").strip() + return model or cls.DEFAULT_CONFIG_ID DEFAULT_ENTITIES: set[str] = { 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/preflight/checks/environment.py b/src/nemo_safe_synthesizer/preflight/checks/environment.py index 619e43550..91daffbab 100644 --- a/src/nemo_safe_synthesizer/preflight/checks/environment.py +++ b/src/nemo_safe_synthesizer/preflight/checks/environment.py @@ -348,12 +348,15 @@ class InferenceModelCheck(ConfigCheck): 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``. All findings are warnings: classification - degrades or fails at call time rather than blocking the run outright. + ``SafeSynthesizerParameters``. The body uses a single-dispatch ``match`` over ``(model, key, endpoint)``, - so at most one warning is emitted per run -- the highest-priority problem. - Priority order: missing key, then blank model id, then invalid 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" @@ -370,9 +373,16 @@ def check(self, ctx: ConfigView, collector: IssueCollector) -> None: endpoint = os.environ.get("NSS_INFERENCE_ENDPOINT") # Single-dispatch: the first matching case wins, so cases are ordered by - # priority. A missing key (degraded mode) is reported before a blank - # model id or an invalid endpoint (hard failures at call time). + # 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", @@ -381,14 +391,8 @@ def check(self, ctx: ConfigView, collector: IssueCollector) -> None: case m, _, _ if _is_blank(m): collector.warning( "inference_model_blank", - "NSS_INFERENCE_MODEL is set but empty. PII column classification will send an " - "empty model id and fail. Unset it to use the default, or provide a model id.", - ) - case _, _, e if e is not None and not _is_valid_http_url(e): - collector.warning( - "inference_endpoint_invalid", - f"NSS_INFERENCE_ENDPOINT '{e}' is not a valid http(s) URL. " - "PII column classification requests will fail.", + "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.", ) diff --git a/tests/pii_replacer/test_detect.py b/tests/pii_replacer/test_detect.py index 7bc5e4b02..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( 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 c7d3e169a..ef1497912 100644 --- a/tests/preflight/test_preflight.py +++ b/tests/preflight/test_preflight.py @@ -286,13 +286,16 @@ def test_valid_model_is_silent(self, 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_warning(self, default_config, endpoint): + 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 == "warning" for i in issues) + 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( @@ -302,9 +305,20 @@ def test_valid_endpoint_is_silent(self, default_config): issues = InferenceModelCheck().run(make_ctx(config=default_config)) assert not any(i.code == "inference_endpoint_invalid" for i in issues) - def test_missing_key_takes_priority_over_other_problems(self, default_config): - # Single-dispatch match: the first matching case wins, so a missing key - # is reported even when the model id and endpoint are also bad. + @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"}, @@ -312,16 +326,20 @@ def test_missing_key_takes_priority_over_other_problems(self, default_config): ): issues = InferenceModelCheck().run(make_ctx(config=default_config)) codes = {i.code for i in issues} - assert codes == {"inference_key_missing"} + assert codes == {"inference_endpoint_invalid"} + assert all(i.severity == "error" for i in issues if i.code == "inference_endpoint_invalid") - def test_blank_model_takes_priority_over_invalid_endpoint(self, default_config): + 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_KEY": "test-key", "NSS_INFERENCE_MODEL": " ", "NSS_INFERENCE_ENDPOINT": "not-a-url"}, + {"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_model_blank"} + assert codes == {"inference_key_missing"} @pytest.mark.unit From ec8738b641728714d60fe9c043524ecd11040a2b Mon Sep 17 00:00:00 2001 From: Aaron Gonzales Date: Tue, 2 Jun 2026 20:28:14 +0000 Subject: [PATCH 10/10] chore: update cli help text Signed-off-by: Aaron Gonzales --- docs/user-guide/environment.md | 3 +++ src/nemo_safe_synthesizer/cli/run.py | 5 ++++- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/docs/user-guide/environment.md b/docs/user-guide/environment.md index c4a55d659..bb690deff 100644 --- a/docs/user-guide/environment.md +++ b/docs/user-guide/environment.md @@ -138,6 +138,9 @@ CLI shorthand for the switch above, with no separate NSS env var: `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`. 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` diff --git a/src/nemo_safe_synthesizer/cli/run.py b/src/nemo_safe_synthesizer/cli/run.py index ceb554e73..17dfc5ef6 100644 --- a/src/nemo_safe_synthesizer/cli/run.py +++ b/src/nemo_safe_synthesizer/cli/run.py @@ -199,7 +199,10 @@ def common_run_options(f: Callable[..., object]) -> Callable[..., object]: 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.", + "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(