Add bash file and modify docker

miscellaneous_scripts/bash_telemetry.sh

@@ -0,0 +1,12 @@
+# telemetry.sh
+#!/bin/bash
+if [ -f /usr/local/bin/deep_learning_container.py ] && [[ -z "${OPT_OUT_TRACKING}" || "${OPT_OUT_TRACKING,,}" != "true" ]]; then
+    (
+        python /usr/local/bin/deep_learning_container.py \
+            --framework "${FRAMEWORK}" \
+            --framework-version "${FRAMEWORK_VERSION}" \
+            --container-type "${CONTAINER_TYPE}" \
+            &>/dev/null &
+    )
+fi
+

pytorch/inference/docker/build_artifacts/torchserve-ec2-entrypoint.py

@@ -17,6 +17,16 @@
 import sys
 
 
+try:
+    subprocess.call(
+        ["/bin/bash", "/usr/local/bin/bash_telemetry.sh"],
+        stdout=subprocess.DEVNULL,
+        stderr=subprocess.DEVNULL,
+    )
+except:
+    pass
+
+
 subprocess.check_call(shlex.split(" ".join(sys.argv[1:])))
 
 # prevent docker exit

pytorch/inference/docker/build_artifacts/torchserve-entrypoint.py

@@ -19,6 +19,15 @@
 import signal
 import torch
 
+try:
+    subprocess.call(
+        ["/bin/bash", "/usr/local/bin/bash_telemetry.sh"],
+        stdout=subprocess.DEVNULL,
+        stderr=subprocess.DEVNULL,
+    )
+except:
+    pass
+
 if torch.cuda.is_available():
     # run compat mounting by default
     try:

pytorch/training/buildspec-2-5-ec2.yml

@@ -42,7 +42,7 @@ images:
     os_version: &OS_VERSION ubuntu22.04
     tag: !join [ *VERSION, "-", *DEVICE_TYPE, "-", *TAG_PYTHON_VERSION, "-", *OS_VERSION, "-ec2" ]
     latest_release_tag: !join [ *VERSION, "-", *DEVICE_TYPE, "-", *TAG_PYTHON_VERSION, "-", *OS_VERSION, "-ec2" ]
-    # build_tag_override: "beta:2.5.1-cpu-py311-ubuntu22.04-ec2"
+    # build_tag_override: "true"
     docker_file: !join [ docker/, *SHORT_VERSION, /, *DOCKER_PYTHON_VERSION, /Dockerfile., *DEVICE_TYPE ]
     target: ec2
     context:
@@ -58,7 +58,7 @@ images:
     os_version: &OS_VERSION ubuntu22.04
     tag: !join [ *VERSION, "-", *DEVICE_TYPE, "-", *TAG_PYTHON_VERSION, "-", *CUDA_VERSION, "-", *OS_VERSION, "-ec2" ]
     latest_release_tag: !join [ *VERSION, "-", *DEVICE_TYPE, "-", *TAG_PYTHON_VERSION, "-", *CUDA_VERSION, "-", *OS_VERSION, "-ec2" ]
-    # build_tag_override: "beta:2.5.1-gpu-py311-cu121-ubuntu22.04-ec2"
+    # build_tag_override: "true"
     docker_file: !join [ docker/, *SHORT_VERSION, /, *DOCKER_PYTHON_VERSION, /, *CUDA_VERSION, /Dockerfile.,
                          *DEVICE_TYPE ]
     target: ec2

pytorch/training/docker/build_artifacts/dockerd_entrypoint.sh

@@ -1,4 +1,7 @@
 #!/usr/bin/env bash
+# Check if telemetry file exists before executing
+# Execute telemetry script if it exists, suppress errors
+bash /usr/local/bin/bash_telemetry.sh >/dev/null 2>&1 || true
 
 CUDA_AVAILABLE=$(python -c "import torch; print(torch.cuda.is_available())")
 if [ "$CUDA_AVAILABLE" = "True" ]; then

pytorch/training/docker/build_artifacts/start_with_right_hostname.sh

@@ -1,5 +1,10 @@
 #!/usr/bin/env bash
 
+# telemetry.sh
+# Execute telemetry script if it exists, suppress errors
+bash /usr/local/bin/bash_telemetry.sh >/dev/null 2>&1 || true
+
+
 CUDA_AVAILABLE=$(python -c "import torch; print(torch.cuda.is_available())")
 if [ "$CUDA_AVAILABLE" = "True" ]; then
     bash /usr/local/bin/start_cuda_compat.sh

src/image_builder.py

@@ -255,26 +255,42 @@ def image_builder(buildspec, image_types=[], device_types=[]):
                 f"This is required to set job_type label."
             )
 
-        template_file = os.path.join(
+        sitecustomize_template_file = os.path.join(
             os.sep, get_cloned_folder_path(), "miscellaneous_scripts", "dlc_template.py"
         )
+        bash_template_file = os.path.join(
+            os.sep, get_cloned_folder_path(), "miscellaneous_scripts", "bash_telemetry.sh"
+        )
 
         template_fw_version = (
             str(image_config["framework_version"])
             if image_config.get("framework_version")
             else str(BUILDSPEC["version"])
         )
         template_fw = str(BUILDSPEC["framework"])
-        post_template_file = utils.generate_dlc_cmd(
-            template_path=template_file,
+        sitecustomize_post_template_file = utils.generate_dlc_cmd(
+            template_path=sitecustomize_template_file,
             output_path=os.path.join(image_config["root"], "out.py"),
             framework=template_fw,
             framework_version=template_fw_version,
             container_type=label_job_type,
         )
+        bash_post_template_file = utils.generate_dlc_cmd(
+            template_path=bash_template_file,
+            output_path=os.path.join(image_config["root"], "telemetry.sh"),
+            framework=template_fw,
+            framework_version=template_fw_version,
+            container_type=label_job_type,
+        )
 
         ARTIFACTS.update(
-            {"customize": {"source": post_template_file, "target": "sitecustomize.py"}}
+            {
+                "customize": {
+                    "source": sitecustomize_post_template_file,
+                    "target": "sitecustomize.py",
+                },
+                "bash": {"source": bash_post_template_file, "target": "bash_telemetry.sh"},
+            }
         )
 
         context = Context(ARTIFACTS, f"build/{image_name}.tar.gz", image_config["root"])

src/utils.py

@@ -688,7 +688,8 @@ def generate_dlc_cmd(template_path, output_path, framework, framework_version, c
     }
 
     for anchor, value in replacements.items():
-        content = content.replace(f"{{{anchor}}}", value)
+        content = content.replace(f"${{{anchor}}}", value)  # replace ${VARIABLE} with value
+        content = content.replace(f"{{{anchor}}}", value)  # replace {VARIABLE} with value
 
     with open(output_path, "w") as out_f:
         out_f.write(content)

tensorflow/inference/buildspec-2-18-ec2.yml

@@ -47,7 +47,7 @@ images:
                                 "-ec2" ]
     docker_file: !join [ docker/, *SHORT_VERSION, /, *DOCKER_PYTHON_VERSION, /Dockerfile., *DEVICE_TYPE ]
     target: ec2
-    # build_tag_override: "beta:2.16.1-cpu-py310-ubuntu20.04-ec2"
+    # build_tag_override: "true"
     context:
       <<: *INFERENCE_CONTEXT
   BuildEC2TensorflowGPUInferencePy3DockerImage:
@@ -67,6 +67,6 @@ images:
     docker_file: !join [ docker/, *SHORT_VERSION, /, *DOCKER_PYTHON_VERSION, /, *CUDA_VERSION, /Dockerfile.,
                          *DEVICE_TYPE ]
     target: ec2
-    # build_tag_override: "beta:2.16.1-gpu-py310-cu122-ubuntu20.04-ec2"
+    # build_tag_override: "true"
     context:
       <<: *INFERENCE_CONTEXT

tensorflow/inference/buildspec-arm64-2-18-ec2.yml

@@ -45,6 +45,6 @@ images:
     latest_release_tag: !join [ *VERSION, "-", *DEVICE_TYPE, "-", *TAG_PYTHON_VERSION, "-", *OS_VERSION, "-ec2" ]
     docker_file: !join [ docker/, *SHORT_VERSION, /, *DOCKER_PYTHON_VERSION, /Dockerfile.arm64., *DEVICE_TYPE ]
     target: ec2
-    # build_tag_override: "beta:2.16.1-cpu-py310-ubuntu20.04-ec2"
+    # build_tag_override: "true"
     context:
       <<: *INFERENCE_CONTEXT

tensorflow/inference/docker/build_artifacts/dockerd_entrypoint.sh

@@ -1,5 +1,7 @@
 #!/usr/bin/env bash
 
+# Execute telemetry script if it exists, suppress errors
+bash /usr/local/bin/bash_telemetry.sh >/dev/null 2>&1 || true
 
 TF_SERVING_PACKAGE=$(pip list | grep tensorflow-serving | cut -d ' ' -f 1)
 

tensorflow/training/buildspec-2-18-ec2.yml

@@ -22,6 +22,9 @@ context:
     dockerd-entrypoint:
       source: docker/build_artifacts/dockerd-entrypoint.py
       target: dockerd-entrypoint.py
+    dockerd_ec2_entrypoint:
+      source: docker/build_artifacts/dockerd_ec2_entrypoint.sh
+      target: dockerd_ec2_entrypoint.sh
     deep_learning_container:
       source: ../../src/deep_learning_container.py
       target: deep_learning_container.py

tensorflow/training/buildspec-2-18-sm.yml

@@ -25,6 +25,9 @@ context:
     dockerd-entrypoint:
       source: docker/build_artifacts/dockerd-entrypoint.py
       target: dockerd-entrypoint.py
+    dockerd_ec2_entrypoint:
+      source: docker/build_artifacts/dockerd_ec2_entrypoint.sh
+      target: dockerd_ec2_entrypoint.sh
     deep_learning_container:
       source: ../../src/deep_learning_container.py
       target: deep_learning_container.py

tensorflow/training/docker/build_artifacts/dockerd-entrypoint.py

@@ -18,14 +18,21 @@
 import sys
 import tensorflow as tf
 
+try:
+    subprocess.call(
+        ["/bin/bash", "/usr/local/bin/bash_telemetry.sh"],
+        stdout=subprocess.DEVNULL,
+        stderr=subprocess.DEVNULL,
+    )
+except:
+    pass
+
+
 if tf.test.is_built_with_cuda():
     # run compat mounting by default
     try:
         subprocess.run(["bash", "-m", "/usr/local/bin/start_cuda_compat.sh"])
     except Exception as e:
         print(f"Error running script: {e}")
 
-if not os.path.exists("/opt/ml/input/config"):
-    subprocess.call(["python", "/usr/local/bin/deep_learning_container.py", "&>/dev/null", "&"])
-
 subprocess.check_call(shlex.split(" ".join(sys.argv[1:])))

tensorflow/training/docker/build_artifacts/dockerd_ec2_entrypoint.sh

@@ -0,0 +1,6 @@
+#!/usr/bin/env bash
+
+# Execute telemetry script if it exists, suppress errors
+bash /usr/local/bin/bash_telemetry.sh >/dev/null 2>&1 || true
+
+eval "$@"

test/dlc_tests/conftest.py

@@ -144,6 +144,13 @@
     "feature_s3_plugin_present": {NightlyFeatureLabel.AWS_S3_PLUGIN_INSTALLED.value},
 }
 
+# Skip telemetry tests for specific versions
+TELEMETRY_SKIP_VERSIONS = {
+    "entrypoint": {"pytorch": ["2.4.0", "2.5.1", "2.6.0"], "tensorflow": ["2.18.0"]},
+    "bashrc": {"pytorch": ["2.4.0", "2.5.1", "2.6.0"], "tensorflow": ["2.18.0"]},
+    "framework": {"pytorch": [""], "tensorflow": [""]},
+}
+
 
 # Nightly fixtures
 @pytest.fixture(scope="session")
@@ -1004,6 +1011,56 @@ def skip_serialized_release_pt_test(request):
         )
 
 
+@pytest.fixture(autouse=True)
+def skip_telemetry_tests(request):
+    """Skip specific telemetry tests based on test name and image version"""
+    test_name = request.node.name.lower()
+
+    if "telemetry_entrypoint" in test_name:
+        _check_telemetry_skip(request, "entrypoint")
+    elif "telemetry_bashrc" in test_name:
+        _check_telemetry_skip(request, "bashrc")
+    elif "telemetry_framework" in test_name:
+        _check_telemetry_skip(request, "framework")
+
+
+def _get_telemetry_image_info(request):
+    """Helper function to get image URI and framework info from fixtures."""
+    telemetry_framework_fixtures = [
+        "pytorch_training",
+        "tensorflow_training",
+        "tensorflow_inference",
+        "pytorch_inference",
+        "pytorch_inference_arm64",
+        "pytorch_training_arm64",
+        "tensorflow_inference_arm64",
+    ]
+
+    for fixture_name in telemetry_framework_fixtures:
+        if fixture_name in request.fixturenames:
+            img_uri = request.getfixturevalue(fixture_name)
+            image_framework, image_framework_version = get_framework_and_version_from_tag(img_uri)
+            return image_framework, image_framework_version
+    return None, None
+
+
+def _check_telemetry_skip(request, test_type):
+    """Common logic for skipping telemetry tests."""
+    if test_type not in TELEMETRY_SKIP_VERSIONS:
+        return
+    image_framework, image_framework_version = _get_telemetry_image_info(request)
+    if not image_framework:
+        return
+    if image_framework not in TELEMETRY_SKIP_VERSIONS[test_type]:
+        return
+
+    if image_framework_version in TELEMETRY_SKIP_VERSIONS[test_type][image_framework]:
+        pytest.skip(
+            f"Telemetry {test_type} test is not supported for "
+            f"{image_framework} version {image_framework_version}"
+        )
+
+
 def _validate_pytorch_framework_version(request, image_uri, test_name, skip_dict):
     """
     Expected format of skip_dic: