Enable contrib.nccl test.

PiperOrigin-RevId: 167592162

2 files changed, 28 insertions, 22 deletions

diff --git a/tensorflow/contrib/nccl/BUILD b/tensorflow/contrib/nccl/BUILD
index 338181e4ca..d6508362b8 100644
--- a/tensorflow/contrib/nccl/BUILD
+++ b/tensorflow/contrib/nccl/BUILD
@@ -48,6 +48,8 @@ tf_cuda_cc_test(
     # Disabled on jenkins until errors finding nvmlShutdown are found.
     tags = [
         "manual",
+        "no_oss",
+        "noguitar",  # note: is run manually there
         "notap",
     ],
     deps = if_cuda(
@@ -112,25 +114,26 @@ tf_custom_op_py_library(
     ],
 )
 
-# http://b/62064807
-# cuda_py_test(
-#     name = "nccl_ops_test",
-#     size = "small",
-#     srcs = ["python/ops/nccl_ops_test.py"],
-#     additional_deps = [
-#         ":nccl_py",
-#         "//tensorflow/python:array_ops",
-#         "//tensorflow/python:client_testlib",
-#         "//tensorflow/python:framework_for_generated_wrappers",
-#         "//tensorflow/python:framework_test_lib",
-#         "//tensorflow/python:platform_test",
-#     ],
-#     # Disabled on jenkins until errors finding nvmlShutdown are found.
-#     tags = [
-#         "manual",
-#         "notap",
-#     ],
-# )
+cuda_py_test(
+    name = "nccl_ops_test",
+    size = "small",
+    srcs = ["python/ops/nccl_ops_test.py"],
+    additional_deps = [
+        ":nccl_py",
+        "//tensorflow/python:array_ops",
+        "//tensorflow/python:client_testlib",
+        "//tensorflow/python:framework_for_generated_wrappers",
+        "//tensorflow/python:framework_test_lib",
+        "//tensorflow/python:platform_test",
+    ],
+    # Disabled on jenkins until errors finding nvmlShutdown are found.
+    tags = [
+        "manual",
+        "no_oss",
+        "noguitar",  # note: is run manually there
+        "notap",
+    ],
+)
 
 filegroup(
     name = "all_files",
diff --git a/tensorflow/contrib/nccl/python/ops/nccl_ops_test.py b/tensorflow/contrib/nccl/python/ops/nccl_ops_test.py
index ae658e7322..1621e9f28e 100644
--- a/tensorflow/contrib/nccl/python/ops/nccl_ops_test.py
+++ b/tensorflow/contrib/nccl/python/ops/nccl_ops_test.py
@@ -43,7 +43,8 @@ class AllReduceTest(test.TestCase):
         self._testSingleAllReduce(sess, dtype, nccl.all_max, np.maximum)
 
   def _testSingleAllReduce(self, sess, np_type, nccl_fn, numpy_accumulation_fn):
-    for devices in [['/device:GPU:0', '/device:GPU:0', '/device:GPU:0'], ['/device:GPU:0', '/device:GPU:0']]:
+    for devices in [['/device:GPU:1', '/device:GPU:2', '/device:GPU:0'],
+                    ['/device:GPU:1', '/device:GPU:0']]:
       shape = (3, 4)
       np_ans = None
       tensors = []
@@ -84,7 +85,8 @@ class BroadcastTest(test.TestCase):
       # Create session inside outer loop to test use of
       # same communicator across multiple sessions.
       with self.test_session(use_gpu=True) as sess:
-        for devices in [['/device:GPU:0', '/device:GPU:0', '/device:GPU:0'], ['/device:GPU:0', '/device:GPU:0']]:
+        for devices in [['/device:GPU:1', '/device:GPU:0', '/device:GPU:2'],
+                        ['/device:GPU:1', '/device:GPU:0']]:
           shape = (3, 4)
           sender = np.random.randint(0, len(devices) - 1)
           with ops.device(devices[sender]):
@@ -115,7 +117,8 @@ class CombinedTest(test.TestCase):
       # Create session inside outer loop to test use of
       # same communicator across multiple sessions.
       with self.test_session(use_gpu=True) as sess:
-        for devices in [['/device:GPU:0', '/device:GPU:0', '/device:GPU:0'], ['/device:GPU:0', '/device:GPU:0']]:
+        for devices in [['/device:GPU:1', '/device:GPU:2', '/device:GPU:0'],
+                        ['/device:GPU:0', '/device:GPU:1']]:
           shape = (3, 4)
 
           # all-reduce