fix: added framework profiling version deprecation warning

Author: mariumof

src/sagemaker/fw_utils.py

@@ -22,6 +22,7 @@
 import tempfile
 from collections import namedtuple
 from typing import Optional, Union, Dict
+from packaging import version
 
 import sagemaker.image_uris
 from sagemaker.session_settings import SessionSettings
@@ -638,6 +639,39 @@ def warn_if_parameter_server_with_multi_gpu(training_instance_type, distribution
         logger.warning(PARAMETER_SERVER_MULTI_GPU_WARNING)
 
 
+def profiler_config_deprecation_warning(
+    profiler_config, image_uri, framework_name, framework_version
+):
+    """
+    Put out a deprecation message for if framework profiling is specified TF >= 2.12 and PT >= 2.0
+    """
+    if profiler_config is None or profiler_config.framework_profile_params is None:
+        return
+
+    if framework_name not in ("pytorch", "tensorflow"):
+        return
+
+    if framework_version is None:
+        logger.warning(framework_version)
+        framework_name, _, image_tag, _ = framework_name_from_image(image_uri)
+
+        if image_tag is not None:
+            framework_version = framework_version_from_tag(image_tag)
+
+    if framework_version is not None:
+        framework_profile_thresh = (
+            version.parse("2.0") if framework_name == "pytorch" else version.parse("2.12")
+        )
+        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,
+            )
+
+
 def validate_smdistributed(
     instance_type, framework_name, framework_version, py_version, distribution, image_uri=None
 ):

src/sagemaker/pytorch/estimator.py

@@ -25,6 +25,7 @@
     python_deprecation_warning,
     validate_version_or_image_args,
     validate_distribution,
+    profiler_config_deprecation_warning,
 )
 from sagemaker.pytorch import defaults
 from sagemaker.pytorch.model import PyTorchModel
@@ -298,6 +299,11 @@ def __init__(
             )
         self.compiler_config = compiler_config
 
+        if "profiler_config" in kwargs:
+            profiler_config_deprecation_warning(
+                kwargs["profiler_config"], image_uri, self._framework_name, framework_version
+            )
+
     def _pytorch_distribution_configuration(self, distribution):
         """Returns a dict of distribution config for PyTorch training
 

src/sagemaker/tensorflow/estimator.py

@@ -235,6 +235,11 @@ def __init__(
                 compiler_config.validate(self)
         self.compiler_config = compiler_config
 
+        if "profiler_config" in kwargs:
+            fw.profiler_config_deprecation_warning(
+                kwargs["profiler_config"], image_uri, self._framework_name, framework_version
+            )
+
     def _validate_args(self, py_version):
         """Placeholder docstring"""
 

tests/unit/test_profiler_config.py

@@ -12,11 +12,14 @@
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
 
+import os
 import pytest
 import re
 import time
 
-
+from sagemaker import image_uris
+from sagemaker.pytorch import PyTorch
+from sagemaker.tensorflow import TensorFlow
 from sagemaker.debugger.profiler_config import ProfilerConfig, FrameworkProfile
 
 from sagemaker.debugger.metrics_config import (
@@ -643,3 +646,119 @@ def test_validation():
 
     with pytest.raises(AssertionError, match=ErrorMessages.INVALID_CPROFILE_TIMER.value):
         PythonProfilingConfig(cprofile_timer="bad_cprofile_timer")
+
+
+DATA_DIR = os.path.join(os.path.dirname(__file__), "..", "data")
+SCRIPT_PATH = os.path.join(DATA_DIR, "dummy_script.py")
+INSTANCE_COUNT = 1
+INSTANCE_TYPE = "ml.p3.2xlarge"
+ROLE = "Dummy"
+REGION = "us-west-2"
+
+
+def test_create_pytorch_estimator_with_framework_profile(
+    sagemaker_session,
+    pytorch_inference_version,
+    pytorch_inference_py_version,
+    default_framework_profile,
+):
+    profiler_config = ProfilerConfig(framework_profile_params=default_framework_profile)
+
+    container_log_level = '"logging.INFO"'
+    source_dir = "s3://mybucket/source"
+    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,
+    )
+
+
+def test_create_pytorch_estimator_w_image_with_framework_profile(
+    sagemaker_session,
+    pytorch_inference_version,
+    pytorch_inference_py_version,
+    gpu_pytorch_instance_type,
+    default_framework_profile,
+):
+    image_uri = image_uris.retrieve(
+        "pytorch",
+        REGION,
+        version=pytorch_inference_version,
+        py_version=pytorch_inference_py_version,
+        instance_type=gpu_pytorch_instance_type,
+        image_scope="inference",
+    )
+
+    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,
+    )
+
+
+"""
+def test_create_tf_estimator_with_framework_profile(
+    sagemaker_session,
+    tensorflow_inference_version,
+    tensorflow_inference_py_version,
+    default_framework_profile,
+):
+    profiler_config = ProfilerConfig(framework_profile_params=default_framework_profile)
+
+    tf = TensorFlow(
+        entry_point=SCRIPT_PATH,
+        role=ROLE,
+        framework_version=tensorflow_inference_version,
+        py_version=tensorflow_inference_py_version,
+        sagemaker_session=sagemaker_session,
+        instance_count=INSTANCE_COUNT,
+        instance_type=INSTANCE_TYPE,
+        profiler_config=profiler_config,
+    )
+... ValueError: TF 1.5 supports only legacy mode. 
+Please supply the image URI directly with 
+'image_uri=520713654638.dkr.ecr.us-west-2.amazonaws.com/sagemaker-tensorflow:1.5-cpu-py2'
+ and set 'model_dir=False'   etc etc
+"""
+
+
+def test_create_tf_estimator_w_image_with_framework_profile(
+    sagemaker_session,
+    tensorflow_inference_version,
+    tensorflow_inference_py_version,
+    default_framework_profile,
+):
+    image_uri = image_uris.retrieve(
+        "tensorflow",
+        REGION,
+        version=tensorflow_inference_version,
+        py_version=tensorflow_inference_py_version,
+        instance_type=INSTANCE_TYPE,
+        image_scope="inference",
+    )
+
+    assert image_uri is not None
+
+    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,
+    )