Deprecation warning for TF2.12 ant PT 2.0 for framework profiling

Author: mariumof

src/sagemaker/debugger/profiler_config.py

@@ -18,6 +18,7 @@
 
 from sagemaker.debugger.framework_profile import FrameworkProfile
 from sagemaker.workflow.entities import PipelineVariable
+from sagemaker.deprecations import deprecation_warn_base
 
 logger = logging.getLogger(__name__)
 
@@ -85,7 +86,7 @@ class and SageMaker Framework estimators.
         self.disable_profiler = disable_profiler
 
         if self.framework_profile_params is not None:
-            logger.warning(
+            deprecation_warn_base(
                 "Framework profiling will be deprecated from tensorflow 2.12 and pytorch 2.0"
             )
 

src/sagemaker/deprecations.py

@@ -64,6 +64,15 @@ def deprecation_warn(name, date, msg=None):
     _warn(f"{name} will be deprecated on {date}.{msg}")
 
 
+def deprecation_warn_base(msg):
+    """Raise a warning for soon to be deprecated feature in sagemaker>=2
+
+    Args:
+        msg (str): the warning message.
+    """
+    _warn(msg)
+
+
 def deprecation_warning(date, msg=None):
     """Decorator for raising deprecation warning for a feature in sagemaker>=2
 

src/sagemaker/fw_utils.py

@@ -31,6 +31,7 @@
 
 from sagemaker.deprecations import renamed_warning, renamed_kwargs
 from sagemaker.workflow.entities import PipelineVariable
+from sagemaker.deprecations import deprecation_warn_base
 
 logger = logging.getLogger(__name__)
 
@@ -661,11 +662,10 @@ def profiler_config_deprecation_warning(
         )
         framework_profile = version.parse(framework_version)
         if framework_profile >= framework_profile_thresh:
-            logger.warning(
-                "Framework profiling is deprecated from %s version %s."
-                "No framework metrics will be collected.",
-                framework_name,
-                framework_version,
+            deprecation_warn_base(
+                f"Framework profiling is deprecated from\
+                 {framework_name} version {framework_version}.\
+                 No framework metrics will be collected"
             )
 
 

tests/unit/test_profiler_config.py

@@ -16,8 +16,11 @@
 import pytest
 import re
 import time
+import warnings
+from packaging import version
 
 from sagemaker import image_uris
+import sagemaker.fw_utils as fw
 from sagemaker.pytorch import PyTorch
 from sagemaker.tensorflow import TensorFlow
 from sagemaker.debugger.profiler_config import ProfilerConfig, FrameworkProfile
@@ -656,6 +659,21 @@ def test_validation():
 REGION = "us-west-2"
 
 
+def _check_framework_profile_deprecation_warning(framework_version, framework_name, warn_list):
+    """Check the collected warnings for a framework fromfile DeprecationWarning"""
+
+    thresh = version.parse("2.12") if framework_name == "tensorflow" else version.parse("2.0")
+    actual = version.parse(framework_version)
+
+    if actual >= thresh:
+        # should find a Framework profiling deprecation warning
+        for w in warn_list:
+            if issubclass(w.category, DeprecationWarning):
+                if "Framework profiling" in str(w.message):
+                    return
+        assert 0  # Should have found a deprecation and exited above
+
+
 def test_create_pytorch_estimator_with_framework_profile(
     sagemaker_session,
     pytorch_inference_version,
@@ -664,18 +682,24 @@ def test_create_pytorch_estimator_with_framework_profile(
 ):
     profiler_config = ProfilerConfig(framework_profile_params=default_framework_profile)
 
-    pytorch = PyTorch(
-        entry_point=SCRIPT_PATH,
-        framework_version=pytorch_inference_version,
-        py_version=pytorch_inference_py_version,
-        role=ROLE,
-        sagemaker_session=sagemaker_session,
-        instance_count=INSTANCE_COUNT,
-        instance_type=INSTANCE_TYPE,
-        base_job_name="job",
-        profiler_config=profiler_config,
-    )
-    assert pytorch._framework_name == "pytorch"
+    with warnings.catch_warnings(record=True) as warn_list:
+        warnings.simplefilter("always")
+        framework_version = pytorch_inference_version
+        pytorch = PyTorch(
+            entry_point=SCRIPT_PATH,
+            framework_version=framework_version,
+            py_version=pytorch_inference_py_version,
+            role=ROLE,
+            sagemaker_session=sagemaker_session,
+            instance_count=INSTANCE_COUNT,
+            instance_type=INSTANCE_TYPE,
+            base_job_name="job",
+            profiler_config=profiler_config,
+        )
+
+        _check_framework_profile_deprecation_warning(
+            framework_version, pytorch._framework_name, warn_list
+        )
 
 
 def test_create_pytorch_estimator_w_image_with_framework_profile(
@@ -696,35 +720,53 @@ def test_create_pytorch_estimator_w_image_with_framework_profile(
 
     profiler_config = ProfilerConfig(framework_profile_params=default_framework_profile)
 
-    pytorch = PyTorch(
-        entry_point=SCRIPT_PATH,
-        role=ROLE,
-        sagemaker_session=sagemaker_session,
-        instance_count=INSTANCE_COUNT,
-        instance_type=gpu_pytorch_instance_type,
-        image_uri=image_uri,
-        profiler_config=profiler_config,
-    )
-    assert pytorch._framework_name == "pytorch"
+    with warnings.catch_warnings(record=True) as warn_list:
+        warnings.simplefilter("always")
+        pytorch = PyTorch(
+            entry_point=SCRIPT_PATH,
+            role=ROLE,
+            sagemaker_session=sagemaker_session,
+            instance_count=INSTANCE_COUNT,
+            instance_type=gpu_pytorch_instance_type,
+            image_uri=image_uri,
+            profiler_config=profiler_config,
+        )
+
+        framework_version = None
+        _, _, image_tag, _ = fw.framework_name_from_image(image_uri)
 
+        if image_tag is not None:
+            framework_version = fw.framework_version_from_tag(image_tag)
 
-def test_create_tf_estimator_with_framework_profile(
+        if framework_version is not None:
+            _check_framework_profile_deprecation_warning(
+                framework_version, pytorch._framework_name, warn_list
+            )
+
+
+def test_create_tf_estimator_with_framework_profile_212(
     sagemaker_session,
     default_framework_profile,
 ):
     profiler_config = ProfilerConfig(framework_profile_params=default_framework_profile)
 
-    tf = TensorFlow(
-        entry_point=SCRIPT_PATH,
-        role=ROLE,
-        framework_version="2.8",
-        py_version="py39",
-        sagemaker_session=sagemaker_session,
-        instance_count=INSTANCE_COUNT,
-        instance_type=INSTANCE_TYPE,
-        profiler_config=profiler_config,
-    )
-    assert tf._framework_name == "tensorflow"
+    with warnings.catch_warnings(record=True) as warn_list:
+        warnings.simplefilter("always")
+        framework_version = "2.12"
+        tf = TensorFlow(
+            entry_point=SCRIPT_PATH,
+            role=ROLE,
+            framework_version=framework_version,
+            py_version="py39",
+            sagemaker_session=sagemaker_session,
+            instance_count=INSTANCE_COUNT,
+            instance_type=INSTANCE_TYPE,
+            profiler_config=profiler_config,
+        )
+
+        _check_framework_profile_deprecation_warning(
+            framework_version, tf._framework_name, warn_list
+        )
 
 
 def test_create_tf_estimator_w_image_with_framework_profile(
@@ -744,13 +786,25 @@ def test_create_tf_estimator_w_image_with_framework_profile(
 
     profiler_config = ProfilerConfig(framework_profile_params=default_framework_profile)
 
-    tf = TensorFlow(
-        entry_point=SCRIPT_PATH,
-        role=ROLE,
-        sagemaker_session=sagemaker_session,
-        instance_count=INSTANCE_COUNT,
-        instance_type=INSTANCE_TYPE,
-        image_uri=image_uri,
-        profiler_config=profiler_config,
-    )
-    assert tf._framework_name == "tensorflow"
+    with warnings.catch_warnings(record=True) as warn_list:
+        warnings.simplefilter("always")
+        tf = TensorFlow(
+            entry_point=SCRIPT_PATH,
+            role=ROLE,
+            sagemaker_session=sagemaker_session,
+            instance_count=INSTANCE_COUNT,
+            instance_type=INSTANCE_TYPE,
+            image_uri=image_uri,
+            profiler_config=profiler_config,
+        )
+
+        framework_version = None
+        _, _, image_tag, _ = fw.framework_name_from_image(image_uri)
+
+        if image_tag is not None:
+            framework_version = fw.framework_version_from_tag(image_tag)
+
+        if framework_version is not None:
+            _check_framework_profile_deprecation_warning(
+                framework_version, tf._framework_name, warn_list
+            )