Add a wrapper module to import Tensorflow.

This is to workaround a known issue with packaging Tensorflow
as a pip dependency in bazel. See:

github.com/bazel-contrib/rules_python/issues/71

Signed-off-by: format 2020.01.12 <github.com/ChrisCummins/format>

Author: ChrisCummins

deeplearning/clgen/models/helper.py

@@ -1,9 +1,10 @@
-import tensorflow as tf
-import tensorflow.contrib.seq2seq as seq2seq
 from tensorflow.python.ops import math_ops
 from tensorflow.python.ops.distributions import categorical
 from tensorflow.python.util import nest
 
+import third_party.py.tensorflow.tf.contrib.seq2seq as seq2seq
+from third_party.py.tensorflow import tf
+
 
 class CustomInferenceHelper(seq2seq.TrainingHelper):
   """An inference helper that takes a seed text"""

deeplearning/clgen/models/tensorflow_backend.py

@@ -87,7 +87,7 @@ def __init__(self, *args, **kwargs):
 
     # Create the summary writer, shared between Train() and
     # _EndOfEpochTestSample().
-    import tensorflow as tf
+    from third_party.py.tensorflow import tf
 
     tensorboard_dir = f"{self.cache.path}/tensorboard"
     app.Log(
@@ -124,8 +124,8 @@ def InitTfGraph(
     os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
 
     # Deferred importing of TensorFlow.
-    import tensorflow as tf
-    import tensorflow.contrib.seq2seq as seq2seq
+    from third_party.py.tensorflow import tf
+    import third_party.py.tensorflow.tf.contrib.seq2seq as seq2seq
     from tensorflow.contrib import rnn
     from deeplearning.clgen.models import helper
 
@@ -483,7 +483,7 @@ def _EndOfEpochTestSample(
     self, corpus, sampler: samplers.Sampler, step: int, epoch_num: int
   ):
     """Run sampler"""
-    import tensorflow as tf
+    from third_party.py.tensorflow import tf
 
     atomizer = corpus.atomizer
     sampler.Specialize(atomizer)
@@ -545,7 +545,7 @@ def InitSampling(
     self, sampler: samplers.Sampler, seed: typing.Optional[int] = None
   ) -> None:
     """Initialize model for sampling."""
-    import tensorflow as tf
+    from third_party.py.tensorflow import tf
 
     # Delete any previous sampling session.
     if self.inference_tf:
@@ -622,7 +622,7 @@ def SampleNextIndices(self, sampler: samplers.Sampler, done: np.ndarray):
     return generated
 
   def RandomizeSampleState(self) -> None:
-    import tensorflow as tf
+    from third_party.py.tensorflow import tf
 
     self.inference_state = [
       tf.nn.rnn_cell.LSTMStateTuple(

third_party/py/tensorflow/tf.py

@@ -0,0 +1,46 @@
+"""Import Tensorflow.
+
+This module is a drop-in replacement for regular tensorflow. Replace:
+
+    import tensorflow as tf
+
+with:
+
+    from third_party.py.tensorflow import tf
+
+This wrapper is required to workaround a known bug with packaging Tensorflow
+as a pip dependency with bazel. See:
+github.com/bazelbuild/rules_python/issues/71
+"""
+import importlib
+import pathlib
+import sys
+
+try:
+  # Try importing Tensorflow the vanilla way. This will succeed once
+  # github.com/bazelbuild/rules_python/issues/71 is fixed.
+  import tensorflow
+except (ImportError, ModuleNotFoundError):
+  # That failed, so see if there is a system install of Tensorflow that we
+  # can trick python into importing. This should succeed in the
+  # chriscummins/phd_base_tf_cpu docker image.
+  PYTHON_SITE_PACKAGES = pathlib.Path("/usr/local/lib/python3.7/site-packages")
+
+  try:
+    if pathlib.Path(PYTHON_SITE_PACKAGES / "tensorflow").is_dir():
+      sys.path.insert(0, "/usr/local/lib/python3.7/site-packages")
+
+      import tensorflow
+    else:
+      raise ModuleNotFoundError
+  except (ImportError, ModuleNotFoundError):
+    # That failed, so a final hail mary let's try importing the module directly.
+    tensorflow = importlib.import_module(
+      "tensorflow", PYTHON_SITE_PACKAGES / "tensorflow",
+    )
+
+# Import Tensorflow into this module's namespace.
+from tensorflow import *
+
+# Pretend that we've imported the regular Tensorflow.
+__file__ = tensorflow.__file__