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 <aagonzales@nvidia.com>

Author: binaryaaron

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
 

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).

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
 

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(

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,

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,15 +336,31 @@ 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
     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["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)
 

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:

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",

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(),