Running TensorFlow tests in Google Colab - tensorflow

Question

I want to run tests on Google Colab to ensure reproducibility but I get a system error at the end, which I do not on my local machine.
I set up TensorFlow in Google Colab with
!pip install tensorflow==1.12.0
import tensorflow as tf
print(tf.__version__)
which, after some lines of installation, prints:
1.12.0
I then want to run a simple test:
import tensorflow as tf
class Tests(tf.test.TestCase):
def test_gpu(self):
self.assertEqual(False, tf.test.is_gpu_available())
tf.test.main()
The test passes (along with a default session test) on my local machine, and also on Colab, but after that the kernel returns a system error:
..
----------------------------------------------------------------------
Ran 2 tests in 0.005s
OK
An exception has occurred, use %tb to see the full traceback.
SystemExit: False
After calling %tb, I get a long stack trace pasted below, which gives little indication. How can I fix it?
The stacktrace is:
SystemExit Traceback (most recent call last)
<ipython-input-20-6a87bf6320f2> in <module>()
7 self.assertEqual(False, tf.test.is_gpu_available())
8
----> 9 tf.test.main()
10
11
/usr/local/lib/python3.6/dist-packages/tensorflow/python/platform/test.py in main(argv)
62 """Runs all unit tests."""
63 _test_util.InstallStackTraceHandler()
---> 64 return _googletest.main(argv)
65
66
/usr/local/lib/python3.6/dist-packages/tensorflow/python/platform/googletest.py in main(argv)
98 args = sys.argv
99 return app.run(main=g_main, argv=args)
--> 100 benchmark.benchmarks_main(true_main=main_wrapper)
101
102
/usr/local/lib/python3.6/dist-packages/tensorflow/python/platform/benchmark.py in benchmarks_main(true_main, argv)
342 app.run(lambda _: _run_benchmarks(regex), argv=argv)
343 else:
--> 344 true_main()
/usr/local/lib/python3.6/dist-packages/tensorflow/python/platform/googletest.py in main_wrapper()
97 if args is None:
98 args = sys.argv
---> 99 return app.run(main=g_main, argv=args)
100 benchmark.benchmarks_main(true_main=main_wrapper)
101
/usr/local/lib/python3.6/dist-packages/tensorflow/python/platform/app.py in run(main, argv)
123 # Call the main function, passing through any arguments
124 # to the final program.
--> 125 _sys.exit(main(argv))
126
/usr/local/lib/python3.6/dist-packages/tensorflow/python/platform/googletest.py in g_main(argv)
68 if ('TEST_TOTAL_SHARDS' not in os.environ or
69 'TEST_SHARD_INDEX' not in os.environ):
---> 70 return unittest_main(argv=argv)
71
72 total_shards = int(os.environ['TEST_TOTAL_SHARDS'])
/usr/lib/python3.6/unittest/main.py in __init__(self, module, defaultTest, argv, testRunner, testLoader, exit, verbosity, failfast, catchbreak, buffer, warnings, tb_locals)
93 self.progName = os.path.basename(argv[0])
94 self.parseArgs(argv)
---> 95 self.runTests()
96
97 def usageExit(self, msg=None):
/usr/lib/python3.6/unittest/main.py in runTests(self)
256 self.result = testRunner.run(self.test)
257 if self.exit:
--> 258 sys.exit(not self.result.wasSuccessful())
259
260 main = TestProgram
SystemExit: False

Answer 1

The error you're seeing is coming from unittest trying to exit the python process, which Jupyter prevents on your behalf. You can avoid that with e.g.:
import tensorflow as tf
class Tests(tf.test.TestCase):
def test_gpu(self):
self.assertEqual(False, tf.test.is_gpu_available())
import unittest
unittest.main(argv=['first-arg-is-ignored'], exit=False)
(note the last line is different to yours and is lifted from https://github.com/jupyter/notebook/issues/2746)