Capture tracelog when SIGSEGV happens during remote execution.

PiperOrigin-RevId: 378647010

Author: SinaChavoshi

src/python/tensorflow_cloud/core/preprocess.py

@@ -128,6 +128,18 @@ def get_preprocessed_entry_point(
         'os.environ["TF_KERAS_RUNNING_REMOTELY"]="1"\n',
     ]
 
+    # Capture job trace log when SIGSEGV occures
+
+    if worker_count > 0 and machine_config.is_tpu_config(worker_config):
+      # faulthandler is not supported in TPU v1.x images
+      logger.info("Faulthandler is not supported with TF < v2.2 images.")
+
+    else:
+      script_lines.extend([
+          "import faulthandler\n",
+          "faulthandler.enable()\n",
+      ])
+
     # Setting default Tuner_ID if one is provided in args
     script_lines.extend([
         "import sys\n",

src/python/tensorflow_cloud/core/tests/unit/preprocess_test.py

@@ -64,6 +64,8 @@ def test_auto_one_device_strategy(self):
             "import os\n",
             "import tensorflow as tf\n",
             'os.environ["TF_KERAS_RUNNING_REMOTELY"]="1"\n',
+            "import faulthandler\n",
+            "faulthandler.enable()\n",
             "import sys\n",
             "for flag in sys.argv[1:]:\n",
             '  if flag.startswith("TUNER_ID"):\n',
@@ -83,6 +85,8 @@ def test_auto_mirrored_strategy(self):
             "import os\n",
             "import tensorflow as tf\n",
             'os.environ["TF_KERAS_RUNNING_REMOTELY"]="1"\n',
+            "import faulthandler\n",
+            "faulthandler.enable()\n",
             "import sys\n",
             "for flag in sys.argv[1:]:\n",
             '  if flag.startswith("TUNER_ID"):\n',
@@ -100,6 +104,8 @@ def test_auto_multi_worker_strategy(self):
             "import os\n",
             "import tensorflow as tf\n",
             'os.environ["TF_KERAS_RUNNING_REMOTELY"]="1"\n',
+            "import faulthandler\n",
+            "faulthandler.enable()\n",
             "import sys\n",
             "for flag in sys.argv[1:]:\n",
             '  if flag.startswith("TUNER_ID"):\n',
@@ -188,6 +194,8 @@ def test_ipython_notebook(self, mock_python_exporter):
             "import os\n",
             "import tensorflow as tf\n",
             'os.environ["TF_KERAS_RUNNING_REMOTELY"]="1"\n',
+            "import faulthandler\n",
+            "faulthandler.enable()\n",
             "import sys\n",
             "for flag in sys.argv[1:]:\n",
             '  if flag.startswith("TUNER_ID"):\n',