smdistributed.dataparallel environment check (#140)

* smdistributed.dataparallel environment check

* addressed comments

* Modified check_smdataparallel_env logic

Author: karan6181

smdebug/core/utils.py

@@ -22,6 +22,7 @@
 from smdebug.core.logger import get_logger
 from smdebug.exceptions import IndexReaderException
 
+_is_invoked_via_smddp = None
 logger = get_logger()
 
 
@@ -345,24 +346,20 @@ def get_distributed_worker():
 
         # 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:
+        if check_smdataparallel_env():
             try:
                 import smdistributed.dataparallel.torch.distributed as smdataparallel
 
                 if smdataparallel.get_world_size():
-                    rank = smdataparallel.get_rank()
+                    return smdataparallel.get_rank()
             except (ModuleNotFoundError, ValueError, ImportError):
                 pass
 
             try:
                 import smdistributed.dataparallel.tensorflow as smdataparallel
 
                 if smdataparallel.size():
-                    rank = smdataparallel.rank()
+                    return smdataparallel.rank()
             except (ModuleNotFoundError, ValueError, ImportError):
                 pass
     return rank
@@ -474,3 +471,29 @@ def __exit__(self, *args):
         shutil.rmtree(self.out_dir, ignore_errors=True)
         if self.tensorboard_dir:
             shutil.rmtree(self.tensorboard_dir, ignore_errors=True)
+
+
+def check_smdataparallel_env():
+    # Check to ensure it is invoked by mpi and the SM distribution is `dataparallel`
+    global _is_invoked_via_smddp
+    if _is_invoked_via_smddp is None:
+        _is_invoked_via_mpi = (
+            os.getenv("OMPI_COMM_WORLD_SIZE") is not None
+            and int(os.getenv("OMPI_COMM_WORLD_SIZE")) >= 8
+        )
+        if os.getenv("SM_FRAMEWORK_PARAMS") is None:
+            _is_invoked_via_smddp = False
+        else:
+            try:
+                smddp_flag = json.loads(os.getenv("SM_FRAMEWORK_PARAMS"))
+            except:
+                _is_invoked_via_smddp = False
+                return _is_invoked_via_smddp
+            if (
+                smddp_flag.get("sagemaker_distributed_dataparallel_enabled", False)
+                and _is_invoked_via_mpi
+            ):
+                _is_invoked_via_smddp = True
+            else:
+                _is_invoked_via_smddp = False
+    return _is_invoked_via_smddp

smdebug/pytorch/hook.py

@@ -12,7 +12,7 @@
 from smdebug.core.collection import DEFAULT_PYTORCH_COLLECTIONS, CollectionKeys
 from smdebug.core.hook import CallbackHook
 from smdebug.core.json_config import DEFAULT_WORKER_NAME
-from smdebug.core.utils import make_numpy_array
+from smdebug.core.utils import check_smdataparallel_env, make_numpy_array
 from smdebug.profiler.hvd_trace_file_rotation import HvdTraceFileRotation
 from smdebug.profiler.profiler_config_parser import MetricsCategory, ProfilerConfigParser
 from smdebug.profiler.profiler_constants import CONVERT_TO_MICROSECS
@@ -25,14 +25,12 @@
 
 # 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:
+smdataparallel = None
+if check_smdataparallel_env():
     try:
         import smdistributed.dataparallel.torch.distributed as smdataparallel
     except ImportError:
-        smdataparallel = None
+        pass
 
 
 DEFAULT_INCLUDE_COLLECTIONS = [CollectionKeys.LOSSES]
@@ -193,11 +191,7 @@ def _get_num_workers(self):
             # Try smdataparallel
             # 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:
+            if check_smdataparallel_env():
                 try:
                     import smdistributed.dataparallel.torch.distributed as smdataparallel
 
@@ -227,11 +221,7 @@ def _get_worker_name(self):
             # Try smdataparallel
             # 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:
+            if check_smdataparallel_env():
                 try:
                     import smdistributed.dataparallel.torch.distributed as smdataparallel
 

smdebug/tensorflow/base_hook.py

@@ -14,7 +14,7 @@
 from smdebug.core.hook import BaseHook
 from smdebug.core.modes import ModeKeys
 from smdebug.core.reductions import get_numpy_reduction, get_reduction_tensor_name
-from smdebug.core.utils import make_numpy_array, serialize_tf_device
+from smdebug.core.utils import check_smdataparallel_env, make_numpy_array, serialize_tf_device
 from smdebug.core.writer import FileWriter
 
 # Local
@@ -135,11 +135,7 @@ def _get_distribution_strategy(self) -> TFDistributionStrategy:
 
         # 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:
+        if check_smdataparallel_env():
             try:
                 import smdistributed.dataparallel.tensorflow as smdataparallel