Fix worker lambda status error handling & data corruption

tensorflow_quantum/core/ops/BUILD

@@ -350,6 +350,18 @@ cc_library(
     ],
 )
 
+cc_test(
+    name = "parse_context_test",
+    srcs = ["parse_context_test.cc"],
+    deps = [
+        ":parse_context",
+        "@com_google_absl//absl/synchronization",
+        "@com_google_googletest//:gtest_main",
+        "@local_config_tf//:libtensorflow_framework",
+        "@local_config_tf//:tf_header_lib",
+    ],
+)
+
 cc_binary(
     name = "_tfq_calculate_unitary_op.so",
     srcs = [

tensorflow_quantum/core/ops/parse_context.cc

@@ -20,6 +20,7 @@ limitations under the License.
 #include <string>
 #include <vector>
 
+#include "absl/synchronization/mutex.h"
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/lib/core/error_codes.pb.h"
 #include "tensorflow/core/lib/core/status.h"
@@ -77,9 +78,14 @@ Status ParsePrograms(OpKernelContext* context, const std::string& input_name,
   const int num_programs = program_strings.dimension(0);
   programs->assign(num_programs, Program());
 
+  tensorflow::Status parse_status;
+  absl::Mutex p_lock;
+
   auto DoWork = [&](int64_t start, int64_t end) {
     for (int64_t i = start; i < end; i++) {
-      OP_REQUIRES_OK(context, ParseProto(program_strings(i), &programs->at(i)));
+      NESTED_FN_STATUS_SYNC(parse_status,
+                            ParseProto(program_strings(i), &programs->at(i)),
+                            p_lock);
     }
   };
 
@@ -88,7 +94,7 @@ Status ParsePrograms(OpKernelContext* context, const std::string& input_name,
   context->device()->tensorflow_cpu_worker_threads()->workers->ParallelFor(
       num_programs, cycle_estimate, DoWork);
 
-  return ::tensorflow::Status();
+  return parse_status;
 }
 
 Status ParsePrograms2D(OpKernelContext* context, const std::string& input_name,
@@ -112,12 +118,16 @@ Status ParsePrograms2D(OpKernelContext* context, const std::string& input_name,
   const int num_entries = program_strings.dimension(1);
   programs->assign(num_programs, std::vector<Program>(num_entries, Program()));
 
+  tensorflow::Status parse_status;
+  absl::Mutex p_lock;
+
   auto DoWork = [&](int64_t start, int64_t end) {
     for (int64_t i = start; i < end; i++) {
-      OP_REQUIRES_OK(
-          context,
+      NESTED_FN_STATUS_SYNC(
+          parse_status,
           ParseProto(program_strings(i / num_entries, i % num_entries),
-                     &programs->at(i / num_entries).at(i % num_entries)));
+                     &programs->at(i / num_entries).at(i % num_entries)),
+          p_lock);
     }
   };
 
@@ -126,7 +136,7 @@ Status ParsePrograms2D(OpKernelContext* context, const std::string& input_name,
   context->device()->tensorflow_cpu_worker_threads()->workers->ParallelFor(
       num_programs * num_entries, cycle_estimate, DoWork);
 
-  return ::tensorflow::Status();
+  return parse_status;
 }
 
 Status GetProgramsAndProgramsToAppend(
@@ -181,17 +191,23 @@ Status GetProgramsAndNumQubits(
 
   // Resolve qubit ID's in parallel.
   num_qubits->assign(programs->size(), -1);
+  tensorflow::Status parse_status;
+  absl::Mutex p_lock;
   auto DoWork = [&](int64_t start, int64_t end) {
     for (int64_t i = start; i < end; i++) {
       Program& program = (*programs)[i];
-      unsigned int this_num_qubits;
+      unsigned int this_num_qubits = 0;
       if (p_sums) {
-        OP_REQUIRES_OK(context,
-                       ResolveQubitIds(&program, &this_num_qubits,
-                                       &(p_sums->at(i)), swap_endianness));
+        NESTED_FN_STATUS_SYNC(
+            parse_status,
+            ResolveQubitIds(&program, &this_num_qubits, &(p_sums->at(i)),
+                            swap_endianness),
+            p_lock);
       } else {
-        OP_REQUIRES_OK(context, ResolveQubitIds(&program, &this_num_qubits,
-                                                nullptr, swap_endianness));
+        NESTED_FN_STATUS_SYNC(parse_status,
+                              ResolveQubitIds(&program, &this_num_qubits,
+                                              nullptr, swap_endianness),
+                              p_lock);
       }
       (*num_qubits)[i] = this_num_qubits;
     }
@@ -202,7 +218,7 @@ Status GetProgramsAndNumQubits(
   context->device()->tensorflow_cpu_worker_threads()->workers->ParallelFor(
       num_qubits->size(), cycle_estimate, DoWork);
 
-  return ::tensorflow::Status();
+  return parse_status;
 }
 
 tensorflow::Status GetProgramsAndNumQubits(
@@ -232,12 +248,16 @@ tensorflow::Status GetProgramsAndNumQubits(
 
   // Resolve qubit ID's in parallel.
   num_qubits->assign(programs->size(), -1);
+  tensorflow::Status parse_status;
+  absl::Mutex p_lock;
   auto DoWork = [&](int64_t start, int64_t end) {
     for (int64_t i = start; i < end; i++) {
       Program& program = (*programs)[i];
-      unsigned int this_num_qubits;
-      OP_REQUIRES_OK(context, ResolveQubitIds(&program, &this_num_qubits,
-                                              &(*other_programs)[i]));
+      unsigned int this_num_qubits = 0;
+      NESTED_FN_STATUS_SYNC(
+          parse_status,
+          ResolveQubitIds(&program, &this_num_qubits, &(*other_programs)[i]),
+          p_lock);
       (*num_qubits)[i] = this_num_qubits;
     }
   };
@@ -247,7 +267,7 @@ tensorflow::Status GetProgramsAndNumQubits(
   context->device()->tensorflow_cpu_worker_threads()->workers->ParallelFor(
       num_qubits->size(), cycle_estimate, DoWork);
 
-  return ::tensorflow::Status();
+  return parse_status;
 }
 
 Status GetPauliSums(OpKernelContext* context,
@@ -270,12 +290,15 @@ Status GetPauliSums(OpKernelContext* context,
   p_sums->assign(sum_specs.dimension(0),
                  std::vector<PauliSum>(sum_specs.dimension(1), PauliSum()));
   const int op_dim = sum_specs.dimension(1);
+  tensorflow::Status parse_status;
+  absl::Mutex p_lock;
   auto DoWork = [&](int64_t start, int64_t end) {
     for (int64_t ii = start; ii < end; ii++) {
       const int i = ii / op_dim;
       const int j = ii % op_dim;
       PauliSum p;
-      OP_REQUIRES_OK(context, ParseProto(sum_specs(i, j), &p));
+      NESTED_FN_STATUS_SYNC(parse_status, ParseProto(sum_specs(i, j), &p),
+                            p_lock);
       (*p_sums)[i][j] = p;
     }
   };
@@ -285,7 +308,7 @@ Status GetPauliSums(OpKernelContext* context,
   context->device()->tensorflow_cpu_worker_threads()->workers->ParallelFor(
       sum_specs.dimension(0) * sum_specs.dimension(1), cycle_estimate, DoWork);
 
-  return ::tensorflow::Status();
+  return parse_status;
 }
 
 Status GetSymbolMaps(OpKernelContext* context, std::vector<SymbolMap>* maps) {

tensorflow_quantum/core/ops/parse_context.h

@@ -17,22 +17,24 @@ limitations under the License.
 #define TFQ_CORE_OPS_PARSE_CONTEXT
 
 #include "absl/synchronization/mutex.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/lib/core/status.h"
 
 // Syncs a threads work status with some global status.
 #define NESTED_FN_STATUS_SYNC(global_status, local_status, global_lock) \
-  if (TF_PREDICT_FALSE(!local_status.ok())) {                           \
-    global_lock.Lock();                                                 \
-    global_status = local_status;                                       \
-    global_lock.Unlock();                                               \
-    return;                                                             \
-  }
+  do {                                                                  \
+    const ::tensorflow::Status _s = (local_status);                     \
+    if (TF_PREDICT_FALSE(!_s.ok())) {                                   \
+      ::absl::MutexLock _l(&(global_lock));                             \
+      (global_status) = _s;                                             \
+      return;                                                           \
+    }                                                                   \
+  } while (0)
 
 #include <string>
 #include <vector>
 
 #include "absl/container/flat_hash_map.h"
-#include "tensorflow/core/framework/op_kernel.h"
-#include "tensorflow/core/lib/core/status.h"
 #include "tensorflow_quantum/core/proto/pauli_sum.pb.h"
 #include "tensorflow_quantum/core/proto/program.pb.h"
 

tensorflow_quantum/core/ops/parse_context_test.cc

@@ -0,0 +1,63 @@
+/* Copyright 2020 The TensorFlow Quantum Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow_quantum/core/ops/parse_context.h"
+
+#include "absl/synchronization/mutex.h"
+#include "gtest/gtest.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/platform/errors.h"
+
+namespace tfq {
+namespace {
+
+TEST(ParseContextTest, NestedFnStatusSyncEvaluatesOnce) {
+  tensorflow::Status global_status;
+  absl::Mutex global_lock;
+
+  int ok_evals = 0;
+  bool ok_completed = false;
+  auto ok_fn = [&]() {
+    NESTED_FN_STATUS_SYNC(global_status, ([&]() {
+                            ++ok_evals;
+                            return tensorflow::Status();
+                          })(),
+                          global_lock);
+    ok_completed = true;
+  };
+  ok_fn();
+  EXPECT_EQ(ok_evals, 1);
+  EXPECT_TRUE(global_status.ok());
+  EXPECT_TRUE(ok_completed);
+
+  int err_evals = 0;
+  bool err_completed = false;
+  auto err_fn = [&]() {
+    NESTED_FN_STATUS_SYNC(
+        global_status, ([&]() {
+          ++err_evals;
+          return tensorflow::errors::InvalidArgument("test error");
+        })(),
+        global_lock);
+    err_completed = true;
+  };
+  err_fn();
+  EXPECT_EQ(err_evals, 1);
+  EXPECT_TRUE(tensorflow::errors::IsInvalidArgument(global_status));
+  EXPECT_FALSE(err_completed);
+}
+
+}  // namespace
+}  // namespace tfq

tensorflow_quantum/core/ops/tfq_ps_util_ops_test.py

@@ -65,6 +65,41 @@ def _complex_test_circuit():
     return circuit_batch
 
 
+class PsWeightsWorkerStatusSyncTest(tf.test.TestCase):
+    """Tests worker status sync in tfq_ps_weights_from_symbols."""
+
+    def test_missing_symbols_batch(self):
+        """Tests multi-threaded batches with missing symbols fail."""
+        bit = cirq.LineQubit(0)
+        sym_alpha = sympy.Symbol('alpha')
+        sym_missing = sympy.Symbol('missing')
+        valid_circuit = cirq.Circuit(cirq.X(bit)**(sym_alpha * 1.5))
+        missing_circuit = cirq.Circuit(cirq.X(bit)**(sym_missing * 2.0))
+
+        valid_raw = util.convert_to_tensor([valid_circuit]).numpy()[0]
+        missing_raw = util.convert_to_tensor([missing_circuit]).numpy()[0]
+        symbols = tf.constant(['alpha'])
+
+        # 1. Fully invalid batch: 256 circuits with missing symbol.
+        all_missing_tensor = tf.constant([missing_raw] * 256)
+        for _ in range(20):
+            with self.assertRaisesRegex(tf.errors.InvalidArgumentError,
+                                        'not found in symbols'):
+                tfq_ps_util_ops.tfq_ps_weights_from_symbols(
+                    all_missing_tensor, symbols)
+
+        # 2. Mixed batch: 250 valid circuits and 6 missing-symbol entries.
+        mixed_batch = [valid_raw] * 256
+        for idx in [0, 42, 100, 150, 200, 255]:
+            mixed_batch[idx] = missing_raw
+        mixed_tensor = tf.constant(mixed_batch)
+        for _ in range(20):
+            with self.assertRaisesRegex(tf.errors.InvalidArgumentError,
+                                        'not found in symbols'):
+                tfq_ps_util_ops.tfq_ps_weights_from_symbols(
+                    mixed_tensor, symbols)
+
+
 class PSDecomposeTest(tf.test.TestCase):
     """Tests on tfq_ps_decompose"""
 

tensorflow_quantum/core/ops/tfq_ps_weights_from_symbols_op.cc

@@ -18,6 +18,7 @@ limitations under the License.
 #include "absl/container/flat_hash_map.h"
 #include "absl/container/flat_hash_set.h"
 #include "absl/strings/numbers.h"
+#include "absl/synchronization/mutex.h"
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/framework/shape_inference.h"
 #include "tensorflow/core/framework/tensor_shape.h"
@@ -78,6 +79,8 @@ class TfqPsWeightsFromSymbolOp : public tensorflow::OpKernel {
                                                         ignore_list.end());
 
     std::vector<int> n_single_symbol(programs.size(), 0);
+    tensorflow::Status parse_status;
+    absl::Mutex p_lock;
 
     auto DoWork = [&](int64_t start, int64_t end) {
       for (int64_t i = start; i < end; i++) {
@@ -95,10 +98,12 @@ class TfqPsWeightsFromSymbolOp : public tensorflow::OpKernel {
               const absl::string_view symbol_name = exponent.symbol();
               if (!symbols_map.contains(symbol_name)) {
                 // Should never happen. raise error.
-                OP_REQUIRES(context, false,
-                            tensorflow::errors::InvalidArgument(
-                                "A circuit contains a sympy.Symbol not found "
-                                "in symbols!"));
+                NESTED_FN_STATUS_SYNC(
+                    parse_status,
+                    tensorflow::errors::InvalidArgument(
+                        "A circuit contains a sympy.Symbol not found "
+                        "in symbols!"),
+                    p_lock);
               }
               output_results.at(i)
                   .at(symbols_map.at(symbol_name))
@@ -125,6 +130,8 @@ class TfqPsWeightsFromSymbolOp : public tensorflow::OpKernel {
         ->workers->TransformRangeConcurrently(block_size, programs.size(),
                                               DoWork);
 
+    OP_REQUIRES_OK(context, parse_status);
+
     int largest_single_symbol = 0;
     for (size_t i = 0; i < n_single_symbol.size(); i++) {
       largest_single_symbol =