Add first-pass changes for enable_determinism function

test/container.sh

@@ -35,6 +35,8 @@ else
   IMAGE=${1}
   if [ "${2}" == "python" ]; then
     ENTRYPOINT="--entrypoint python"
+  elif [ "${2}" == "bash" ]; then
+    ENTRYPOINT="--entrypoint bash"
   else
     ENTRYPOINT="--entrypoint /mnt/test/${2}"
   fi

test/test_utils.py

@@ -0,0 +1,60 @@
+# Copyright 2020 NVIDIA Corporation. All Rights Reserved
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ========================================================================
+
+import sys
+import unittest
+
+sys.path.insert(0, '..')
+from tfdeterminism.utils import _Version as Version
+
+class TestUtils(unittest.TestCase):
+
+  def test_version_class(self):
+    v = Version('23.45.26')
+    self.assertTrue (v.in_list(['1.2', '2.8', '23.45']))
+    self.assertFalse(v.in_list(['4.5', '2.9', '99.4']))
+    self.assertTrue (v.at_least('23.45'))
+    self.assertFalse(v.at_least('23.46'))
+    self.assertTrue (v.at_most('23.45'))
+    self.assertFalse(v.at_most('23.44'))
+    self.assertTrue (v.between('23.44', '23.47'))
+    self.assertFalse(v.between('1.2', '2.4'))
+    self.assertFalse(v.between('100.2', '2.4'))
+    self.assertTrue (v.between('1.0', '200.4'))
+    v = Version('1.2')
+    self.assertTrue (v.between('0.9', '1.4'))
+    v = Version('10.09-tf3')
+    self.assertTrue (v.in_list(['10.02', '10.09', '09.12']))
+    self.assertTrue (v.at_least('10.09'))
+    self.assertFalse(v.at_least('10.10'))
+    self.assertTrue (v.at_most('10.09'))
+    self.assertFalse(v.at_most('10.08'))
+
+  def test_version_class_exceptions(self):
+    self.assertRaises(ValueError, Version, '10')
+    self.assertRaises(TypeError, Version, None)
+    v = Version('2.3')
+    self.assertRaises(ValueError, v.at_least, '1')
+    self.assertRaises(ValueError, v.at_least, '1.2.3')
+    self.assertRaises(ValueError, v.at_most, '012')
+    self.assertRaises(ValueError, v.at_most, '012.004.435')
+    self.assertRaises(ValueError, v.between, '10', '2.2')
+    self.assertRaises(ValueError, v.between, '1.3', '20')
+    # self.assertRaises(ValueError, v.between, '10.2', '2.26.2') # short-circuit
+    self.assertRaises(ValueError, v.between, '1.2', '2.26.2')
+    self.assertRaises(ValueError, v.between, '1.3.6', '20.5')
+
+if __name__ == '__main__':
+  unittest.main()

tfdeterminism/__init__.py

@@ -17,6 +17,7 @@
 from __future__ import division
 from __future__ import print_function
 
+from .enable_determinism import _enable_determinism as enable_determinism
 from .patch import _patch as patch
 
 from .version import __version__

tfdeterminism/enable_determinism.py

@@ -0,0 +1,58 @@
+# Copyright 2020 NVIDIA Corporation. All Rights Reserved
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ========================================================================
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+
+import tensorflow as tf
+
+from .patch import _patch_bias_add
+from .utils import _Version as Version
+
+def _enable_determinism(seed=None):
+  """Provides a best-effort recipe to increase framework determinism when
+    running on GPUs.
+
+    Call this method either before or after explicitly importing TensorFlow,
+    but always before constructing any graphs.
+
+    This function cannot address all possible sources of non-determinism. Please 
+    see further instructions at https://github.com/NVIDIA/tensorflow-determinism
+    to understand how to use it in a larger deterministic context.
+
+    Arguments:
+      seed: <fill in>
+
+    Returns: nothing
+  """
+  tf_vers = Version(tf.version.VERSION)
+  ngc_tf_container_version_string = os.environ.get('NVIDIA_TENSORFLOW_VERSION')
+  if ngc_tf_container_version_string:
+    in_ngc_cont = True
+    ngc_vers = Version(ngc_tf_container_version_string)
+  else:
+    in_ngc_cont = False
+  if not in_ngc_cont and tf_vers.between('1.14', '2.0'):
+    os.environ['TF_CUDNN_DETERMINISTIC'] = '1'
+    _patch_bias_add()
+  if in_ngc_cont and ngc_vers.at_least('19.06') or tf_vers.at_least('2.1'):
+    os.environ['TF_DETERMINISTIC_OPS'] = '1'
+  if in_ngc_cont and ngc_vers.at_least('19.06') or tf_vers.at_least('1.14'):
+    # Apply the fused softmax/cross-entropy patch here
+    pass
+  # TODO: Add other recipe items 

tfdeterminism/patch.py

@@ -70,6 +70,7 @@ def _patch():
   if re.match("(1\.(14|15)|2\.0)", tf_version):
     os.environ['TF_CUDNN_DETERMINISTIC'] = '1'
     _patch_bias_add()
+    # Apply the fused softmax/cross-entropy patch here
     print("TensorFlow version %s has been patched "
           "using tfdeterminism version %s" %
           (tf_version, __version__), file=sys.stderr)

tfdeterminism/utils.py

@@ -0,0 +1,71 @@
+# Copyright 2020 NVIDIA Corporation. All Rights Reserved
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ========================================================================
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import re
+
+class _Version:
+  """Provides version comparison functionality"""
+  def __init__(self, version_string):
+    """Provide a version string containing at least a major and minor version"""
+    version_pieces = re.split('\.|-', version_string)
+    if len(version_pieces) < 2:
+      raise ValueError("The version string must contain at least a major and a minor version")
+    major = version_pieces[0]
+    minor = version_pieces[1]
+    self.version_string = major + '.' + minor
+    self.major = int(major)
+    self.minor = int(minor)
+
+  def in_list(self, list_of_versions):
+    """Is the version in the list of version provided?"""
+    return self.version_string in list_of_versions
+
+  def _only_major_and_minor(self, version):
+    version_pieces = version.split('.')
+    if len(version_pieces) != 2:
+      raise ValueError("The version string must contain a major and a minor version (only)")
+    major = int(version_pieces[0])
+    minor = int(version_pieces[1])
+    return major, minor
+
+  def at_least(self, oldest_version):
+    """Is the version at least the oldest_version provided?"""
+    oldest_major, oldest_minor = self._only_major_and_minor(oldest_version)
+    if (self.major > oldest_major or
+        self.major == oldest_major and self.minor >= oldest_minor):
+      return True
+    else:
+      return False
+
+  def at_most(self, newest_version):
+    """Is the version at most the newest version provided?"""
+    newest_major, newest_minor = self._only_major_and_minor(newest_version)
+    if (self.major < newest_major or
+        self.major == newest_major and self.minor <= newest_minor):
+      return True
+    else:
+      return False
+
+  def between(self, oldest_version, newest_version):
+    """Is the version between the oldest and newest versions
+    provided (inclusive)?"""
+    if self.at_least(oldest_version) and self.at_most(newest_version):
+      return True
+    else:
+      return False