diff --git a/dev-tools/extract_model_ops/README.md b/dev-tools/extract_model_ops/README.md index f7b7f2f39..fa27b4cfe 100644 --- a/dev-tools/extract_model_ops/README.md +++ b/dev-tools/extract_model_ops/README.md @@ -116,6 +116,25 @@ Each file maps a short architecture name to a HuggingFace model identifier: } ``` +An entry may instead be an object with `model_id` plus optional fields: +`quantized`, `auto_class`, `config_overrides`, and `trust_remote_code`. + +```json +{ + "jina-embeddings-v5-text-nano": { + "model_id": "jinaai/jina-embeddings-v5-text-nano", + "trust_remote_code": true + } +} +``` + +`trust_remote_code` defaults to `false` so an untrusted or compromised model +repo cannot execute arbitrary Python on the build machine during load (see +CVE-2026-5241, where a nested config could override the caller's setting). +Enable it only for a vetted model that genuinely ships custom modeling code +(e.g. the Jina v5 embeddings model above). The bundled reference/validation +models are otherwise native architectures that load without remote code. + To add a new architecture, append an entry to `reference_models.json`, re-run `extract_model_ops.py --cpp`, and update `CSupportedOperations.cc`. Then add the same entry (plus any task-specific variants) to diff --git a/dev-tools/extract_model_ops/extract_model_ops.py b/dev-tools/extract_model_ops/extract_model_ops.py index 2a070d1cc..f8229a26e 100644 --- a/dev-tools/extract_model_ops/extract_model_ops.py +++ b/dev-tools/extract_model_ops/extract_model_ops.py @@ -49,7 +49,8 @@ def extract_ops_for_model(model_name: str, quantize: bool = False, auto_class: str | None = None, - config_overrides: dict | None = None) -> Optional[set[str]]: + config_overrides: dict | None = None, + trust_remote_code: bool = False) -> Optional[set[str]]: """Trace a HuggingFace model and return its TorchScript op set. Returns None if the model could not be loaded or traced. @@ -58,7 +59,8 @@ def extract_ops_for_model(model_name: str, print(f" Loading {label}...", file=sys.stderr) traced = load_and_trace_hf_model(model_name, quantize=quantize, auto_class=auto_class, - config_overrides=config_overrides) + config_overrides=config_overrides, + trust_remote_code=trust_remote_code) if traced is None: return None return collect_inlined_ops(traced) @@ -99,7 +101,8 @@ def main(): ops = extract_ops_for_model(spec["model_id"], quantize=spec["quantized"], auto_class=spec["auto_class"], - config_overrides=spec["config_overrides"]) + config_overrides=spec["config_overrides"], + trust_remote_code=spec["trust_remote_code"]) if ops is None: failed.append(arch) print(f" {arch}: FAILED", file=sys.stderr) diff --git a/dev-tools/extract_model_ops/reference_models.json b/dev-tools/extract_model_ops/reference_models.json index 23368f053..32747a268 100644 --- a/dev-tools/extract_model_ops/reference_models.json +++ b/dev-tools/extract_model_ops/reference_models.json @@ -30,7 +30,8 @@ "elastic-eis-elser-v2-quantized": {"model_id": "elastic/eis-elser-v2", "quantized": true}, "elastic-test-elser-v2-quantized": {"model_id": "elastic/test-elser-v2", "quantized": true}, - "jina-embeddings-v5-text-nano": "jinaai/jina-embeddings-v5-text-nano", + "_comment:trust-remote-code": "trust_remote_code defaults to false for safety (see CVE-2026-5241). Enable it only for vetted models that ship custom modeling code; jina-embeddings-v5 is one such model.", + "jina-embeddings-v5-text-nano": {"model_id": "jinaai/jina-embeddings-v5-text-nano", "trust_remote_code": true}, "_comment:qa-models": "Models from the Appex QA pytorch_tests suite. BART models require auto_class and config_overrides to trace correctly.", "qa-tinyroberta-squad2": {"model_id": "deepset/tinyroberta-squad2", "auto_class": "AutoModelForQuestionAnswering"}, diff --git a/dev-tools/extract_model_ops/requirements.txt b/dev-tools/extract_model_ops/requirements.txt index 0a2c117e9..0e49d0e49 100644 --- a/dev-tools/extract_model_ops/requirements.txt +++ b/dev-tools/extract_model_ops/requirements.txt @@ -1,4 +1,10 @@ torch==2.7.1 +# Pinned to the 4.x line on purpose: transformers 5.0 removed TorchScript +# support (the torchscript=True config + torch.jit tracing this tool relies on). +# CVE-2026-5241 (RCE via the LightGlue loading path) is not reachable here -- the +# configured models are curated NLP encoders, never LightGlue, and this is a +# dev/CI-only tool. We additionally default trust_remote_code=False (see +# torchscript_utils.load_and_trace_hf_model) to neutralise that class of issue. transformers>=4.40.0,<5.0.0 sentencepiece>=0.2.0 protobuf>=5.0.0 diff --git a/dev-tools/extract_model_ops/torchscript_utils.py b/dev-tools/extract_model_ops/torchscript_utils.py index a8560c421..d72323fa2 100644 --- a/dev-tools/extract_model_ops/torchscript_utils.py +++ b/dev-tools/extract_model_ops/torchscript_utils.py @@ -31,6 +31,10 @@ def load_model_config(config_path: Path) -> dict[str, dict]: of ``AutoModel`` (e.g. ``"AutoModelForSequenceClassification"``). - ``config_overrides`` (dict) — extra kwargs passed to ``AutoConfig.from_pretrained`` (e.g. ``{"use_cache": false}``). + - ``trust_remote_code`` (bool, default False) — allow the model repo to + execute custom Python during load. Off by default so an untrusted repo + cannot run arbitrary code (see CVE-2026-5241); enable it only for a + vetted model that genuinely ships a custom architecture. Keys starting with ``_comment`` are silently skipped. @@ -46,17 +50,28 @@ def load_model_config(config_path: Path) -> dict[str, dict]: continue if isinstance(value, str): models[key] = {"model_id": value, "quantized": False, - "auto_class": None, "config_overrides": {}} + "auto_class": None, "config_overrides": {}, + "trust_remote_code": False} elif isinstance(value, dict): if "model_id" not in value: raise ValueError( f"Config entry {key!r} is a dict but missing required " f"'model_id' key: {value!r}") + trust_remote_code = value.get("trust_remote_code", False) + # Validate strictly: this flag gates remote code execution, so a + # non-bool (e.g. the string "false", which is truthy) must never be + # silently coerced into enabling it. + if not isinstance(trust_remote_code, bool): + raise ValueError( + f"Config entry {key!r} has non-boolean 'trust_remote_code' " + f"{trust_remote_code!r} ({type(trust_remote_code).__name__}); " + f"expected true or false.") models[key] = { "model_id": value["model_id"], "quantized": value.get("quantized", False), "auto_class": value.get("auto_class"), "config_overrides": value.get("config_overrides", {}), + "trust_remote_code": trust_remote_code, } else: raise ValueError( @@ -96,7 +111,8 @@ def _resolve_auto_class(class_name: str | None): def load_and_trace_hf_model(model_name: str, quantize: bool = False, auto_class: str | None = None, - config_overrides: dict | None = None): + config_overrides: dict | None = None, + trust_remote_code: bool = False): """Load a HuggingFace model, tokenize sample input, and trace to TorchScript. When *quantize* is True the model is dynamically quantized (nn.Linear @@ -109,6 +125,12 @@ def load_and_trace_hf_model(model_name: str, quantize: bool = False, *config_overrides* supplies extra kwargs to ``AutoConfig.from_pretrained`` (e.g. ``{"use_cache": False}`` for encoder-decoder models like BART). + *trust_remote_code* allows the model repo to execute custom Python during + load. It defaults to False so that an untrusted/compromised repo cannot run + arbitrary code on the build machine (see CVE-2026-5241, where a nested + config could override the caller's setting). Enable it per-model only for a + vetted architecture that genuinely ships custom modeling code. + Returns the traced module, or None if the model could not be loaded or traced. """ token = os.environ.get("HF_TOKEN") or None @@ -117,13 +139,13 @@ def load_and_trace_hf_model(model_name: str, quantize: bool = False, try: tokenizer = AutoTokenizer.from_pretrained( - model_name, token=token, trust_remote_code=True) + model_name, token=token, trust_remote_code=trust_remote_code) config = AutoConfig.from_pretrained( model_name, torchscript=True, token=token, - trust_remote_code=True, **overrides) + trust_remote_code=trust_remote_code, **overrides) model = model_cls.from_pretrained( model_name, config=config, token=token, - trust_remote_code=True) + trust_remote_code=trust_remote_code) model.eval() except Exception as exc: print(f" LOAD ERROR: {exc}", file=sys.stderr) diff --git a/dev-tools/extract_model_ops/validate_allowlist.py b/dev-tools/extract_model_ops/validate_allowlist.py index a1dbffc6b..fc5515c4b 100644 --- a/dev-tools/extract_model_ops/validate_allowlist.py +++ b/dev-tools/extract_model_ops/validate_allowlist.py @@ -123,6 +123,7 @@ def validate_model(model_name: str, quantize: bool = False, auto_class: str | None = None, config_overrides: dict | None = None, + trust_remote_code: bool = False, timeout: int = MODEL_TIMEOUT_SECONDS) -> str: """Validate one HuggingFace model. @@ -139,7 +140,8 @@ def validate_model(model_name: str, try: traced = load_and_trace_hf_model(model_name, quantize=quantize, auto_class=auto_class, - config_overrides=config_overrides) + config_overrides=config_overrides, + trust_remote_code=trust_remote_code) except ModelTimeoutError: print(f" SKIPPED (timed out after {timeout}s)", file=sys.stderr) return "skip" @@ -212,6 +214,7 @@ def main(): quantize=spec["quantized"], auto_class=spec.get("auto_class"), config_overrides=spec.get("config_overrides"), + trust_remote_code=spec.get("trust_remote_code", False), timeout=args.model_timeout) if args.pt_dir and args.pt_dir.is_dir(): diff --git a/dev-tools/extract_model_ops/validation_models.json b/dev-tools/extract_model_ops/validation_models.json index 20aaf98d1..09bcfa315 100644 --- a/dev-tools/extract_model_ops/validation_models.json +++ b/dev-tools/extract_model_ops/validation_models.json @@ -31,7 +31,8 @@ "es-cross-encoder-ms-marco": "cross-encoder/ms-marco-MiniLM-L-6-v2", "es-dpr-question-encoder": "facebook/dpr-question_encoder-single-nq-base", - "jina-embeddings-v5-text-nano": "jinaai/jina-embeddings-v5-text-nano", + "_comment:trust-remote-code": "trust_remote_code defaults to false for safety (see CVE-2026-5241). Enable it only for vetted models that ship custom modeling code; jina-embeddings-v5 is one such model.", + "jina-embeddings-v5-text-nano": {"model_id": "jinaai/jina-embeddings-v5-text-nano", "trust_remote_code": true}, "_comment:qa-models": "Models from the Appex QA pytorch_tests suite. BART models require auto_class and config_overrides to trace correctly.", "qa-tinyroberta-squad2": {"model_id": "deepset/tinyroberta-squad2", "auto_class": "AutoModelForQuestionAnswering"}, diff --git a/dev-tools/unittest/test_extract_model_ops_config.py b/dev-tools/unittest/test_extract_model_ops_config.py new file mode 100644 index 000000000..242f212df --- /dev/null +++ b/dev-tools/unittest/test_extract_model_ops_config.py @@ -0,0 +1,78 @@ +#!/usr/bin/env python3 +# Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one +# or more contributor license agreements. Licensed under the Elastic License +# 2.0 and the following additional limitation. Functionality enabled by the +# files subject to the Elastic License 2.0 may only be used in production when +# invoked by an Elasticsearch process with a license key installed that permits +# use of machine learning features. You may not use this file except in +# compliance with the Elastic License 2.0 and the foregoing additional +# limitation. + +"""Tests for dev-tools/extract_model_ops config parsing (trust_remote_code). + +The extract_model_ops tool defaults ``trust_remote_code`` to False so an +untrusted model repo cannot run arbitrary code during load (CVE-2026-5241); +it must be opt-in per model. ``torchscript_utils`` imports torch/transformers +at module load, so these tests are skipped where those deps are absent (e.g. +the lean dev-tools CI unittest env); they run in the extract_model_ops venv. +""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +import pytest + +pytest.importorskip("torch", reason="extract_model_ops requires torch") +pytest.importorskip("transformers", reason="extract_model_ops requires transformers") + +_EXTRACT_DIR = Path(__file__).resolve().parents[1] / "extract_model_ops" +if str(_EXTRACT_DIR) not in sys.path: + sys.path.insert(0, str(_EXTRACT_DIR)) + +import torchscript_utils as tsu # noqa: E402 + + +def _write(tmp_path: Path, data: dict) -> Path: + p = tmp_path / "models.json" + p.write_text(json.dumps(data), encoding="utf-8") + return p + + +def test_string_entry_defaults_trust_remote_code_false(tmp_path: Path) -> None: + cfg = _write(tmp_path, {"bert": "bert-base-uncased"}) + models = tsu.load_model_config(cfg) + assert models["bert"]["trust_remote_code"] is False + + +def test_dict_entry_defaults_trust_remote_code_false(tmp_path: Path) -> None: + cfg = _write(tmp_path, {"m": {"model_id": "foo/bar", "quantized": True}}) + models = tsu.load_model_config(cfg) + assert models["m"]["trust_remote_code"] is False + + +def test_dict_entry_opt_in_trust_remote_code(tmp_path: Path) -> None: + cfg = _write(tmp_path, {"jina": {"model_id": "jinaai/x", "trust_remote_code": True}}) + models = tsu.load_model_config(cfg) + assert models["jina"]["trust_remote_code"] is True + + +@pytest.mark.parametrize("bad_value", ["false", "true", 1, 0, None, ["true"]]) +def test_non_bool_trust_remote_code_is_rejected(tmp_path: Path, bad_value) -> None: + # A non-bool must never be silently coerced into enabling remote code: + # e.g. the string "false" is truthy and would otherwise turn it on. + cfg = _write(tmp_path, {"m": {"model_id": "foo/bar", "trust_remote_code": bad_value}}) + with pytest.raises(ValueError, match="trust_remote_code"): + tsu.load_model_config(cfg) + + +def test_bundled_configs_only_trust_vetted_models() -> None: + # The real configs must not silently trust arbitrary repos: only models + # explicitly known to ship custom code may opt in. + allowed_trusted = {"jina-embeddings-v5-text-nano"} + for name in ("reference_models.json", "validation_models.json"): + models = tsu.load_model_config(_EXTRACT_DIR / name) + trusted = {k for k, v in models.items() if v["trust_remote_code"]} + assert trusted <= allowed_trusted, f"{name} trusts unexpected models: {trusted - allowed_trusted}"