[ML] Harden pytorch_inference with TorchScript model graph validation

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
   )
 

bin/pytorch_inference/CModelGraphValidator.cc

@@ -0,0 +1,115 @@
+/*
+ * 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 <core/CLogger.h>
+
+#include <torch/csrc/jit/passes/inliner.h>
+
+#include <algorithm>
+
+namespace ml {
+namespace torch {
+
+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<std::string_view>& allowedOps,
+                               const std::unordered_set<std::string_view>& 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.
+    for (const auto& op : observedOps) {
+        if (forbiddenOps.contains(op)) {
+            result.s_IsValid = false;
+            result.s_ForbiddenOps.push_back(op);
+        }
+    }
+
+    if (result.s_ForbiddenOps.empty()) {
+        for (const auto& op : observedOps) {
+            if (allowedOps.contains(op) == false) {
+                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;
+            }
+        }
+    }
+}
+
+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;
+        }
+    }
+}
+}
+}

bin/pytorch_inference/CModelGraphValidator.h

@@ -0,0 +1,91 @@
+/*
+ * 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 <torch/script.h>
+
+#include <string>
+#include <string_view>
+#include <unordered_set>
+#include <vector>
+
+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<std::string>;
+    using TStringVec = std::vector<std::string>;
+
+    //! 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<std::string_view>& allowedOps,
+                            const std::unordered_set<std::string_view>& forbiddenOps);
+
+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

bin/pytorch_inference/CSupportedOperations.cc

@@ -0,0 +1,129 @@
+/*
+ * 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,
+    // 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, elastic/splade-v3,
+// elastic/test-elser-v2.
+// Additional ops from Elasticsearch integration test models
+// (PyTorchModelIT, TextExpansionQueryIT, TextEmbeddingQueryIT).
+const CSupportedOperations::TStringViewSet CSupportedOperations::ALLOWED_OPERATIONS = {
+    // aten operations (core tensor computations)
+    "aten::Int"sv,
+    "aten::IntImplicit"sv,
+    "aten::ScalarImplicit"sv,
+    "aten::__and__"sv,
+    "aten::abs"sv,
+    "aten::add"sv,
+    "aten::add_"sv,
+    "aten::arange"sv,
+    "aten::bitwise_not"sv,
+    "aten::cat"sv,
+    "aten::chunk"sv,
+    "aten::clamp"sv,
+    "aten::contiguous"sv,
+    "aten::cumsum"sv,
+    "aten::div"sv,
+    "aten::div_"sv,
+    "aten::dropout"sv,
+    "aten::embedding"sv,
+    "aten::expand"sv,
+    "aten::full_like"sv,
+    "aten::gather"sv,
+    "aten::ge"sv,
+    "aten::gelu"sv,
+    "aten::hash"sv,
+    "aten::index"sv,
+    "aten::index_put_"sv,
+    "aten::layer_norm"sv,
+    "aten::len"sv,
+    "aten::linear"sv,
+    "aten::log"sv,
+    "aten::lt"sv,
+    "aten::manual_seed"sv,
+    "aten::masked_fill"sv,
+    "aten::matmul"sv,
+    "aten::max"sv,
+    "aten::mean"sv,
+    "aten::min"sv,
+    "aten::mul"sv,
+    "aten::ne"sv,
+    "aten::neg"sv,
+    "aten::new_ones"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::rsub"sv,
+    "aten::scaled_dot_product_attention"sv,
+    "aten::select"sv,
+    "aten::size"sv,
+    "aten::slice"sv,
+    "aten::softmax"sv,
+    "aten::sqrt"sv,
+    "aten::squeeze"sv,
+    "aten::str"sv,
+    "aten::sub"sv,
+    "aten::tanh"sv,
+    "aten::tensor"sv,
+    "aten::to"sv,
+    "aten::transpose"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,
+};
+}
+}