Cherry picking the changes from PR#138 to this branch

Author: leleamol

smdebug/core/utils.py

@@ -343,21 +343,28 @@ def get_distributed_worker():
         except (ModuleNotFoundError, ValueError, ImportError):
             pass
 
-        try:
-            import smdistributed.dataparallel.torch.distributed as smdataparallel
-
-            if smdataparallel.get_world_size():
-                rank = smdataparallel.get_rank()
-        except (ModuleNotFoundError, ValueError, ImportError):
-            pass
-
-        try:
-            import smdistributed.dataparallel.tensorflow as smdataparallel
-
-            if smdataparallel.size():
-                rank = smdataparallel.rank()
-        except (ModuleNotFoundError, ValueError, ImportError):
-            pass
+        # smdistributed.dataparallel should be invoked via `mpirun`.
+        # It supports EC2 machines with 8 GPUs per machine.
+        _is_invoked_via_mpi = (
+            os.getenv("OMPI_COMM_WORLD_SIZE") is not None
+            and int(os.getenv("OMPI_COMM_WORLD_SIZE")) >= 8
+        )
+        if _is_invoked_via_mpi:
+            try:
+                import smdistributed.dataparallel.torch.distributed as smdataparallel
+
+                if smdataparallel.get_world_size():
+                    rank = smdataparallel.get_rank()
+            except (ModuleNotFoundError, ValueError, ImportError):
+                pass
+
+            try:
+                import smdistributed.dataparallel.tensorflow as smdataparallel
+
+                if smdataparallel.size():
+                    rank = smdataparallel.rank()
+            except (ModuleNotFoundError, ValueError, ImportError):
+                pass
     return rank
 
 

smdebug/pytorch/hook.py

@@ -23,10 +23,16 @@
 from smdebug.pytorch.singleton_utils import set_hook
 from smdebug.pytorch.utils import get_reduction_of_data
 
-try:
-    import smdistributed.dataparallel.torch.distributed as smdataparallel
-except ImportError:
-    smdataparallel = None
+# smdistributed.dataparallel should be invoked via `mpirun`.
+# It supports EC2 machines with 8 GPUs per machine.
+_is_invoked_via_mpi = (
+    os.getenv("OMPI_COMM_WORLD_SIZE") is not None and int(os.getenv("OMPI_COMM_WORLD_SIZE")) >= 8
+)
+if _is_invoked_via_mpi:
+    try:
+        import smdistributed.dataparallel.torch.distributed as smdataparallel
+    except ImportError:
+        smdataparallel = None
 
 
 DEFAULT_INCLUDE_COLLECTIONS = [CollectionKeys.LOSSES]
@@ -185,13 +191,20 @@ def _get_num_workers(self):
                 pass
 
             # Try smdataparallel
-            try:
-                import smdistributed.dataparallel.torch.distributed as smdataparallel
-
-                if smdataparallel.get_world_size():
-                    return smdataparallel.get_world_size()
-            except (ModuleNotFoundError, ValueError, ImportError):
-                pass
+            # smdistributed.dataparallel should be invoked via `mpirun`.
+            # It supports EC2 machines with 8 GPUs per machine.
+            _is_invoked_via_mpi = (
+                os.getenv("OMPI_COMM_WORLD_SIZE") is not None
+                and int(os.getenv("OMPI_COMM_WORLD_SIZE")) >= 8
+            )
+            if _is_invoked_via_mpi:
+                try:
+                    import smdistributed.dataparallel.torch.distributed as smdataparallel
+
+                    if smdataparallel.get_world_size():
+                        return smdataparallel.get_world_size()
+                except (ModuleNotFoundError, ValueError, ImportError):
+                    pass
         # Return default
         return 1
 
@@ -212,13 +225,20 @@ def _get_worker_name(self):
                 pass
 
             # Try smdataparallel
-            try:
-                import smdistributed.dataparallel.torch.distributed as smdataparallel
-
-                if smdataparallel.get_world_size():
-                    return f"worker_{smdataparallel.get_rank()}"
-            except (ModuleNotFoundError, ValueError, ImportError):
-                pass
+            # smdistributed.dataparallel should be invoked via `mpirun`.
+            # It supports EC2 machines with 8 GPUs per machine.
+            _is_invoked_via_mpi = (
+                os.getenv("OMPI_COMM_WORLD_SIZE") is not None
+                and int(os.getenv("OMPI_COMM_WORLD_SIZE")) >= 8
+            )
+            if _is_invoked_via_mpi:
+                try:
+                    import smdistributed.dataparallel.torch.distributed as smdataparallel
+
+                    if smdataparallel.get_world_size():
+                        return f"worker_{smdataparallel.get_rank()}"
+                except (ModuleNotFoundError, ValueError, ImportError):
+                    pass
         # Return default
         return DEFAULT_WORKER_NAME
 

smdebug/tensorflow/base_hook.py

@@ -133,14 +133,21 @@ def _get_distribution_strategy(self) -> TFDistributionStrategy:
         except (ModuleNotFoundError, ValueError, ImportError):
             pass
 
-        try:
-            import smdistributed.dataparallel.tensorflow as smdataparallel
-
-            # The total number of GPUs across all the nodes in the cluster
-            if smdataparallel.size():
-                return TFDistributionStrategy.SMDATAPARALLEL
-        except (ModuleNotFoundError, ValueError, ImportError):
-            pass
+        # smdistributed.dataparallel should be invoked via `mpirun`.
+        # It supports EC2 machines with 8 GPUs per machine.
+        _is_invoked_via_mpi = (
+            os.getenv("OMPI_COMM_WORLD_SIZE") is not None
+            and int(os.getenv("OMPI_COMM_WORLD_SIZE")) >= 8
+        )
+        if _is_invoked_via_mpi:
+            try:
+                import smdistributed.dataparallel.tensorflow as smdataparallel
+
+                # The total number of GPUs across all the nodes in the cluster
+                if smdataparallel.size():
+                    return TFDistributionStrategy.SMDATAPARALLEL
+            except (ModuleNotFoundError, ValueError, ImportError):
+                pass
 
         strat = tf.distribute.get_strategy()
         if is_mirrored_strategy(strat):