Report NotFound when a CUDA delegate's weights blob is missing

backends/cuda/runtime/cuda_backend.cpp

@@ -1234,7 +1234,17 @@ class ET_EXPERIMENTAL CudaBackend final
       const NamedDataMap* named_data_map,
       const std::string& weights_blob_key) const {
     auto buffer_res = named_data_map->get_data(weights_blob_key.c_str());
-    if (buffer_res.ok() && handle->update_constants_from_blob != nullptr) {
+    // A library built before external weights exports no bind function, so
+    // there is nothing to do with a blob whether or not one was supplied.
+    if (handle->update_constants_from_blob == nullptr) {
+      ET_LOG(
+          Info,
+          "weights_blob '%s' is not used: this library exports no bind function",
+          weights_blob_key.c_str());
+      return Error::Ok;
+    }
+
+    if (buffer_res.ok()) {
       ET_LOG(Info, "Found %s in named data map", weights_blob_key.c_str());
       const void* weights_blob = buffer_res->data();
       auto update_err = handle->update_constants_from_blob(
@@ -1246,10 +1256,39 @@ class ET_EXPERIMENTAL CudaBackend final
       ET_CUDA_CHECK_OR_RETURN_ERROR(cudaDeviceSynchronize());
       buffer_res->Free();
     } else {
+      // The container expects a blob and did not get one, so its constant
+      // pointers stay null and the failure resurfaces much later as an illegal
+      // access inside a generated kernel. Report it while the cause is still
+      // identifiable, as the cached path already does. A model with no
+      // constants needs nothing bound and stays valid.
+      // Both symbols are optional, and the count one has been present for
+      // longer, so a library exporting the update function without it is not
+      // something any released toolchain produces. Refuse rather than guess,
+      // since guessing wrong either rejects a valid model or loads one with
+      // unbound constants.
+      ET_CHECK_OR_RETURN_ERROR(
+          handle->get_num_constants != nullptr,
+          NotSupported,
+          "weights_blob '%s' is unavailable and this library cannot report its constant count, so whether it needs the blob cannot be established",
+          weights_blob_key.c_str());
+      size_t num_constants = 0;
+      ET_CHECK_OK_OR_RETURN_ERROR(
+          handle->get_num_constants(handle->container_handle, &num_constants),
+          "Failed to enumerate CUDA AOTI constants");
+      ET_CHECK_OR_RETURN_ERROR(
+          num_constants == 0,
+          NotFound,
+          "weights_blob '%s' is unavailable (0x%" PRIx32
+          "), but the model has %zu constant(s) to bind",
+          weights_blob_key.c_str(),
+          static_cast<uint32_t>(buffer_res.error()),
+          num_constants);
       ET_LOG(
           Info,
-          "weights_blob '%s' not found or update fn is null",
-          weights_blob_key.c_str());
+          "weights_blob '%s' is unavailable (0x%" PRIx32
+          "), and the model has no constants",
+          weights_blob_key.c_str(),
+          static_cast<uint32_t>(buffer_res.error()));
     }
     return Error::Ok;
   }

backends/cuda/tests/test_missing_weights_blob.py

@@ -0,0 +1,171 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""A model whose weights blob is absent must fail to load, not fail later.
+
+The CUDA backend hands a model's constants to its generated library through a
+sidecar blob. When that blob is missing the container keeps null constant
+pointers, and the load used to succeed: the failure surfaced from the first
+execute as an illegal memory access inside a generated kernel, with nothing in
+the message naming the blob.
+
+Two payload shapes reach that code. A current export carries per-name weight
+metadata and goes through the weight cache, which reports a missing blob itself.
+A library built before external weights carries only the two blob keys, newline
+separated, and goes through the legacy path. That legacy path is the one the
+check was added to, so it is the one this file exercises, by rewriting the
+payload of a real export into the older shape in place.
+"""
+
+import os
+import tempfile
+import unittest
+
+import torch
+from executorch.backends.cuda.cuda_backend import CudaBackend
+from executorch.backends.cuda.cuda_partitioner import CudaPartitioner
+
+from executorch.backends.cuda.cuda_weight_collector import CUDA_WEIGHT_CACHE_MAGIC
+from executorch.exir import to_edge_transform_and_lower
+from executorch.exir._serialize._program import deserialize_pte_binary
+from torch.export import export
+
+
+class ModuleWithConstants(torch.nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.register_buffer("a", torch.randn(4, 64))
+        self.register_buffer("b", torch.randn(64))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        y = x * self.a
+        y = y + self.b
+        y = torch.relu(y)
+        return y * self.a
+
+
+@unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
+@unittest.skipIf(
+    torch.version.hip is not None,
+    "loads through the Python runtime, which the ROCm build does not include",
+)
+class TestMissingWeightsBlob(unittest.TestCase):
+    def _lower(self, outdir: str) -> str:
+        torch.manual_seed(0)
+        module = ModuleWithConstants().eval().cuda()
+        inputs = (torch.randn(4, 64, device="cuda"),)
+        exported = export(module, inputs)
+        spec = CudaBackend.generate_method_name_compile_spec("forward")
+        lowered = to_edge_transform_and_lower(
+            exported, partitioner=[CudaPartitioner([spec])]
+        )
+        program = lowered.to_executorch()
+        path = os.path.join(outdir, "model.pte")
+        with open(path, "wb") as f:
+            program.write_to_file(f)
+        program.write_tensor_data_to_file(outdir)
+        return path
+
+    def _rewrite_payload_as_legacy(self, path: str) -> str:
+        """Replace the weight metadata payload with the older two-key payload.
+
+        Rewriting in place keeps every offset in the file valid, and modelling the
+        older shape this way avoids needing an old toolchain to produce one. The
+        payload is read from the program rather than located by scanning, because a
+        program with more than one delegate holds more than one payload and a scan
+        cannot tell where one ends.
+        """
+        with open(path, "rb") as f:
+            raw = f.read()
+        program = deserialize_pte_binary(raw).program
+
+        payloads = []
+        for plan in program.execution_plan:
+            for delegate in plan.delegates:
+                if delegate.id != "CudaBackend":
+                    continue
+                stored = program.backend_delegate_data[delegate.processed.index].data
+                payloads.append(
+                    b"".join(stored.chunks())
+                    if hasattr(stored, "chunks")
+                    else bytes(stored)
+                )
+        self.assertEqual(len(payloads), 1, "expected one CUDA delegate")
+        payload = payloads[0]
+        self.assertTrue(
+            payload.startswith(CUDA_WEIGHT_CACHE_MAGIC),
+            "expected the weight metadata payload this rewrite consumes",
+        )
+
+        # The payload carries a content hash, so it occurs once.
+        offset = raw.find(payload)
+        self.assertGreaterEqual(offset, 0, "payload not found in the file")
+        self.assertEqual(
+            raw.find(payload, offset + 1), -1, "payload is not unique in the file"
+        )
+
+        marker = b"_so_blob"
+        end = payload.index(marker) + len(marker)
+        so_key = payload[:end].rsplit(b"\x00", 1)[-1].decode("utf-8")
+        weights_key = so_key.replace("_so_blob", "_weights_blob")
+
+        # The two keys, then zeros to keep the payload its original length. The
+        # shared library key ends at the newline and stays exact; the filler lands
+        # in the blob key, which is still a key the data map does not hold.
+        legacy = (so_key + "\n" + weights_key).encode("utf-8")
+        self.assertLessEqual(len(legacy), len(payload), "legacy payload does not fit")
+        legacy += b"\x00" * (len(payload) - len(legacy))
+
+        blob = bytearray(raw)
+        blob[offset : offset + len(payload)] = legacy
+        with open(path, "wb") as f:
+            f.write(bytes(blob))
+        return weights_key
+
+    def test_load_reports_not_found_when_blob_is_absent(self) -> None:
+        from executorch.runtime import Runtime
+
+        with tempfile.TemporaryDirectory() as outdir:
+            path = self._lower(outdir)
+
+            blobs = [f for f in os.listdir(outdir) if f.endswith(".ptd")]
+            self.assertTrue(blobs, "expected an externalized weights blob")
+            # A blob holding no constants is a bare header, and a missing blob
+            # really is harmless then, so the model has to carry real data for
+            # this to be testing anything.
+            largest = max(os.path.getsize(os.path.join(outdir, f)) for f in blobs)
+            self.assertGreater(largest, 256, "expected non-empty constants")
+
+            # A positive control first: the same program loads when its blob is
+            # supplied, so a later failure is about the missing blob and not about
+            # the program, the backend registration or this rewrite.
+            sidecar = os.path.join(outdir, blobs[0])
+            Runtime.get().load_program(path, data_path=sidecar).load_method("forward")
+
+            self._rewrite_payload_as_legacy(path)
+
+            # Without this the test would still pass if the rewrite stopped
+            # working, by exercising the weight cache path instead, which reports
+            # the same error number for the same program.
+            with open(path, "rb") as f:
+                self.assertNotIn(
+                    CUDA_WEIGHT_CACHE_MAGIC,
+                    f.read(),
+                    "the rewrite left the metadata payload in place",
+                )
+
+            # The blob is never supplied, so the load must fail. The runtime's
+            # exception carries only the method name and the error number, so the
+            # cause is asserted through the rewrite check above rather than here.
+            program = Runtime.get().load_program(path)
+            with self.assertRaisesRegex(
+                RuntimeError, r"Failed to load method forward, error: 0x:?20"
+            ):
+                program.load_method("forward")
+
+
+if __name__ == "__main__":
+    unittest.main()