@@ -569,6 +569,12 @@ public final class _ExecutionContext {
569
569
// Initialize the TF runtime exactly once. Only affects local execution
570
570
// (when _RuntimeConfig.tensorFlowServer is set to "").
571
571
if !_RuntimeConfig. tensorFlowRuntimeInitialized {
572
+ // Install a signal handler to ensure we exit when interrupted.
573
+ signal ( SIGINT) { _ in
574
+ print ( " Caught interrupt signal, exiting... " )
575
+ exit ( 1 )
576
+ }
577
+
572
578
var args = [ " dummyProgramName " ]
573
579
if _RuntimeConfig. printsDebugLog {
574
580
args. append ( " --alsologtostderr " )
@@ -588,24 +594,19 @@ public final class _ExecutionContext {
588
594
589
595
// Calculate the addresses of all the strings within our single buffer, and then call
590
596
// TF_InitMain.
591
- flattenedStringBytes. withUnsafeMutableBufferPointer { flattenedStringBytesBuffer in
597
+ flattenedStringBytes. withUnsafeMutableBufferPointer { buffer in
592
598
var stringAddrs : [ UnsafeMutablePointer < Int8 > ? ] = [ ]
593
- var currentStringAddr = flattenedStringBytesBuffer . baseAddress
599
+ var currentStringAddr = buffer . baseAddress
594
600
. map ( UnsafeMutablePointer . init)
595
601
for length in lengths {
596
602
stringAddrs. append ( currentStringAddr)
597
603
currentStringAddr = currentStringAddr? . advanced ( by: length)
598
604
}
599
605
600
606
stringAddrs. withUnsafeMutableBufferPointer { stringAddrsBuffer in
601
- var cArgs = [ stringAddrsBuffer. baseAddress. map ( UnsafeMutablePointer . init) ]
602
- var cArgsCount = [ Int32 ( args. count) ]
603
-
604
- cArgs. withUnsafeMutableBufferPointer { cArgsBuffer in
605
- cArgsCount. withUnsafeMutableBufferPointer { cArgsCountBuffer in
606
- TF_InitMain ( nil , cArgsCountBuffer. baseAddress, cArgsBuffer. baseAddress)
607
- }
608
- }
607
+ var cArgsCount = Int32 ( args. count)
608
+ var cArgs = stringAddrsBuffer. baseAddress. map ( UnsafeMutablePointer . init)
609
+ TF_InitMain ( nil , & cArgsCount, & cArgs)
609
610
}
610
611
}
611
612
_RuntimeConfig. tensorFlowRuntimeInitialized = true
0 commit comments