[ML] Harden pytorch_inference with TorchScript model graph validation

.buildkite/pipeline.json.py

@@ -52,6 +52,8 @@ def main():
         if config.build_x86_64:
             pipeline_steps.append(pipeline_steps.generate_step("Upload ES tests x86_64 runner pipeline",
                                                                ".buildkite/pipelines/run_es_tests_x86_64.yml.sh"))
+            pipeline_steps.append(pipeline_steps.generate_step("Upload ES inference tests x86_64 runner pipeline",
+                                                               ".buildkite/pipelines/run_es_inference_tests_x86_64.yml.sh"))
             # We only use linux x86_64 builds for QA tests.
             if config.run_qa_tests:
                 pipeline_steps.append(pipeline_steps.generate_step("Upload QA tests runner pipeline",

.buildkite/pipelines/run_es_inference_tests_x86_64.yml.sh

@@ -0,0 +1,34 @@
+#!/bin/bash
+# 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.
+
+cat <<EOL
+steps:
+  - label: "Java :java: Inference Integration Tests for x86_64 :hammer:"
+    key: "java_inference_tests_x86_64"
+    command:
+      - 'sudo rpm --import https://yum.corretto.aws/corretto.key'
+      - 'sudo curl -L -o /etc/yum.repos.d/corretto.repo https://yum.corretto.aws/corretto.repo'
+      - 'sudo dnf install -y java-21-amazon-corretto-devel'
+      - 'buildkite-agent artifact download "build/*" . --step build_test_linux-x86_64-RelWithDebInfo'
+      - '.buildkite/scripts/steps/run_es_inference_tests.sh || (cd ../elasticsearch && find x-pack -name logs | xargs tar cvzf logs.tgz && buildkite-agent artifact upload logs.tgz && false)'
+    depends_on: "build_test_linux-x86_64-RelWithDebInfo"
+    agents:
+      provider: aws
+      instanceType: m6i.2xlarge
+      imagePrefix: core-amazonlinux-2023
+      diskSizeGb: 100
+      diskName: '/dev/xvda'
+    env:
+      IVY_REPO: "../ivy"
+      GRADLE_JVM_OPTS: "-Dorg.gradle.jvmargs=-Xmx16g"
+    notify:
+      - github_commit_status:
+          context: "Java Inference Integration Tests for x86_64"
+EOL

.buildkite/pipelines/run_pytorch_tests.yml.sh

@@ -8,6 +8,8 @@
 # compliance with the Elastic License 2.0 and the foregoing additional
 # limitation.
 
+SAFE_MESSAGE=$(printf '%s' "${BUILDKITE_MESSAGE}" | head -1 | sed -e 's/\\/\\\\/g' -e 's/"/\\"/g')
+
 cat <<EOL
 steps:
   - label: "Trigger Appex PyTorch Tests :test_tube:"
@@ -22,7 +24,7 @@ steps:
   - trigger: appex-qa-stateful-custom-ml-cpp-build-testing
     async: false
     build:
-      message: "${BUILDKITE_MESSAGE}"
+      message: "${SAFE_MESSAGE}"
       env:
         QAF_TESTS_TO_RUN: "pytorch_tests"
 EOL

.buildkite/pipelines/run_qa_tests.yml.sh

@@ -8,6 +8,8 @@
 # compliance with the Elastic License 2.0 and the foregoing additional
 # limitation.
 
+SAFE_MESSAGE=$(printf '%s' "${BUILDKITE_MESSAGE}" | head -1 | sed -e 's/\\/\\\\/g' -e 's/"/\\"/g')
+
 cat <<EOL
 steps:
   - label: "Trigger Appex QA Tests :test_tube:"
@@ -22,13 +24,7 @@ steps:
   - trigger: appex-qa-stateful-custom-ml-cpp-build-testing
     async: false
     build:
-      message: |
-EOL
-
-# Output the message with proper indentation for YAML literal block scalar
-printf '%s\n' "${BUILDKITE_MESSAGE}" | sed 's/^/        /'
-
-cat <<EOL
+      message: "${SAFE_MESSAGE}"
       env:
         QAF_TESTS_TO_RUN: "${QAF_TESTS_TO_RUN:-ml_cpp_pr}"
 EOL

.buildkite/scripts/steps/run_es_inference_tests.sh

@@ -0,0 +1,31 @@
+#!/bin/bash
+# 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.
+
+set -euo pipefail
+
+echo "pwd = $(pwd)"
+
+export HARDWARE_ARCH=$(uname -m | sed 's/arm64/aarch64/')
+
+VERSION=$(cat ${REPO_ROOT}/gradle.properties | grep '^elasticsearchVersion' | awk -F= '{ print $2 }' | xargs echo)
+if [ "${BUILD_SNAPSHOT:=true}" = "true" ] ; then
+    VERSION=${VERSION}-SNAPSHOT
+fi
+export VERSION
+
+export PR_AUTHOR=$(expr "$BUILDKITE_BRANCH" : '\(.*\):.*')
+export PR_SOURCE_BRANCH=$(expr "$BUILDKITE_BRANCH" : '.*:\(.*\)')
+export PR_TARGET_BRANCH=${BUILDKITE_PULL_REQUEST_BASE_BRANCH}
+
+mkdir -p "${IVY_REPO}/maven/org/elasticsearch/ml/ml-cpp/$VERSION"
+cp "build/distributions/ml-cpp-$VERSION-linux-$HARDWARE_ARCH.zip" "${IVY_REPO}/maven/org/elasticsearch/ml/ml-cpp/$VERSION/ml-cpp-$VERSION.zip"
+cp "build/distributions/ml-cpp-$VERSION-linux-$HARDWARE_ARCH.zip" "${IVY_REPO}/maven/org/elasticsearch/ml/ml-cpp/$VERSION/ml-cpp-$VERSION-nodeps.zip"
+cp dev-tools/minimal.zip "${IVY_REPO}/maven/org/elasticsearch/ml/ml-cpp/$VERSION/ml-cpp-$VERSION-deps.zip"
+./dev-tools/run_es_inference_tests.sh ".." "$(cd "${IVY_REPO}" && pwd)"

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