Exit the process on a TF runtime error from the TF_SessionRun() call, instead of check failing. (#18248)

Exit the process on a TF runtime error from the TF_SessionRun() call, instead of
check failing.

Author: Mingsheng Hong

stdlib/public/TensorFlow/CompilerRuntime.swift

@@ -87,7 +87,8 @@ public enum _RuntimeConfig {
 
   /// When true, run the entire tensor computation in
   /// _TFCStartTensorComputation(), instead of running it on a separate thread.
-  /// - Note: Set to true only for debugging purposes.
+  /// - Note: Set to true only for debugging purposes, as it has limited
+  ///   functionality (e.g. no sends/recvs support).
   static public var usesSynchronousExecution = false
 
   /// For CPU and GPU execution without XLA, use the auto mode. For XLA and/or
@@ -466,6 +467,8 @@ extension TFState {
     inputTensors.append(cTensor!)
   }
 
+  /// Runs the tensor program. Aborts the process on error, and emits an error
+  /// string to STDERR.
   /// See the comment on _TensorComputation.helperFunctionCount on the concept
   /// of a "helper function".
   func execute(_ entryFunctionBaseName: String,
@@ -547,7 +550,10 @@ extension TFState {
       targetNodeSpecs, Int32(targetNodeSpecs.count),
       /*run_metadata*/nil, status
     )
-    checkOk(status)
+    if (TF_GetCode(status) != TF_OK) {
+      _ = fputs(TF_Message(status), stderr)
+      exit(-1)
+    }
     debugLog("Done running TF computation.")
 
     // Delete input tensors.
@@ -715,7 +721,7 @@ public final class _TensorComputation {
       // TODO(hongm): do error handling.
       internalConsistencyCheck(creationStatus == 0)
     }
-    // If it's asynchronous, we call execute() on the main thread directly.
+    // If it's synchronous, we call execute() on the main thread directly.
     else {
       // Log a debug message to differentiate from async computation.
       debugLog("Running tensor computation synchronously.")
@@ -731,6 +737,8 @@ public final class _TensorComputation {
 }
 
 private extension _TensorComputation {
+  /// Runs the tensor program. Aborts the process on error, and emits an error
+  /// string to STDERR.
   // NOTE: This is to be called by the initializer. The computation gets
   // executed on initialization, thus this method will not be exposed to users.
   private func execute() {
@@ -758,6 +766,7 @@ public extension _TensorComputation {
 
   /// Waits for completion the computation as given by 'program', and returns
   /// output handles, whose underlying tensors may live on CPU or GPU.
+  /// Aborts the process on error, and emits an error string to STDERR.
   func finish() -> [CTensorHandle] {
     debugLog("Calling _TensorComputation.finish().")
     if let pthread = pthread {
@@ -842,6 +851,7 @@ public func _TFCStartTensorComputation(
 
 /// Waits for completion of the computation as given by `computation`, and
 /// returns results.
+/// Aborts the process on error, and emits an error string to STDERR.
 ///
 /// - Parameters:
 ///   - computation: The tensor computation to finish.

test/TensorFlowRuntime/dataset_1.swift

@@ -75,11 +75,14 @@ public func model() {
   _hostOp(three)
   expectNearlyEqualWithScalarTensor(3.0, three)
 
-  // TODO: do not crash when TF emits "Fatal error: End of sequence"
-  // let error: TensorHandle<Float> = #tfop("IteratorGetNext",
-  //                                      iterator,
-  //                                      output_types: [Float.self],
-  //                                      output_shapes: [TensorShape()])
+  // Running the commented-out code below will cause the process to exit, with
+  // TF error message "End of sequence" printed on STDERR. The code is commented
+  // out because running it will unfortunately cause the test to fail.
+
+  // let _: TensorHandle<Float> = #tfop("IteratorGetNext",
+  //                                    iterator,
+  //                                    output_types: [Float.self],
+  //                                    output_shapes: [TensorShape()])
 }
 
 DatasetTests.testAllBackends("Basic") {