feature: add PyTorch version environment variable, to facilitate SMTT (#250)

Co-authored-by: Yongyan Rao <[email protected]>

Author: yongyanrao

src/sagemaker_pytorch_container/training.py

@@ -96,6 +96,7 @@ def train(training_environment):
         capture_error = False
         logger.info(f'capture_error is {capture_error}. Default is True')
 
+    _set_torch_version_environment()
     try:
         entry_point.run(uri=training_environment.module_dir,
                         user_entry_point=training_environment.user_entry_point,
@@ -149,5 +150,22 @@ def _set_nccl_environment(network_interface_name):
     os.environ['NCCL_DEBUG'] = 'WARN'
 
 
+def _set_torch_version_environment():
+    """Set PyTorch version environment variable.
+
+    This is the PyTorch version of the DLC.
+    """
+    try:
+        import torch
+
+        os.environ["SM_DLC_TORCH_VERSION"] = torch.__version__
+    except ModuleNotFoundError:
+        logger.warn("PyTorch cannot be found")
+    except ImportError:
+        logger.warn("PyTorch can be found, but cannot be imported")
+    except Exception:
+        logger.warn("Torch version environment variable cannot be set")
+
+
 def main():
     train(environment.Environment())

test/unit/test_train.py

@@ -14,6 +14,7 @@
 
 import os
 import shutil
+import sys
 import tempfile
 
 import pytest
@@ -23,7 +24,14 @@
 from mock import MagicMock, PropertyMock
 from mock import patch
 
-from sagemaker_pytorch_container.training import main, train, _dns_lookup, LAUNCH_PYTORCH_XLA_ENV_NAME, MASTER_PORT
+from sagemaker_pytorch_container.training import (
+    main,
+    train,
+    _dns_lookup,
+    LAUNCH_PYTORCH_XLA_ENV_NAME,
+    MASTER_PORT,
+    _set_torch_version_environment,
+)
 
 
 @pytest.fixture(name='training_env')
@@ -218,3 +226,11 @@ def test_user_script_error_raised(run_entry_point, training_env):
     )
     with pytest.raises(errors.ExecuteUserScriptError):
         train(training_env)
+
+
+def test_set_torch_version_environment():
+    mock_torch = MagicMock()
+    mock_torch.__version__ = '2.0.0'
+    sys.modules['torch'] = mock_torch
+    _set_torch_version_environment()
+    assert os.environ.get("SM_DLC_TORCH_VERSION") == '2.0.0'