Split sanity test into functionality and security for test time reduction

dlc_developer_config.toml

@@ -53,7 +53,8 @@ notify_test_failures = false
 
 [test]
 ### On by default
-sanity_tests = true
+security_sanity_tests = true
+functionality_sanity_tests = true
   safety_check_test = false
   ecr_scan_allowlist_feature = false
 ecs_tests = true

src/config.py

@@ -95,8 +95,12 @@ def is_sm_benchmark_test_enabled():
     return parse_dlc_developer_configs("test", "sagemaker_benchmark_tests")
 
 
-def is_sanity_test_enabled():
-    return parse_dlc_developer_configs("test", "sanity_tests")
+def is_security_sanity_test_enabled():
+    return parse_dlc_developer_configs("test", "security_sanity_tests")
+
+
+def is_functionality_sanity_test_enabled():
+    return parse_dlc_developer_configs("test", "functionality_sanity_tests")
 
 
 def is_sm_local_test_enabled():

src/constants.py

@@ -87,7 +87,8 @@
 SAGEMAKER_EFA_TESTS = "sagemaker-efa"
 SAGEMAKER_RC_TESTS = "sagemaker-rc"
 SAGEMAKER_BENCHMARK_TESTS = "sagemaker-benchmark"
-SANITY_TESTS = "sanity"
+SECURITY_SANITY_TESTS = "security-sanity"
+FUNCTIONALITY_SANITY_TESTS = "functionality-sanity"
 EC2_TESTS = "ec2"
 EC2_BENCHMARK_TESTS = "ec2-benchmark"
 ECS_TESTS = "ecs"

src/prepare_dlc_dev_environment.py

@@ -22,7 +22,8 @@
 
 
 VALID_TEST_TYPES = [
-    "sanity_tests",
+    "security_sanity_tests",
+    "functionality_sanity_tests",
     "ec2_tests",
     "ecs_tests",
     "eks_tests",

src/send_test_notification.py

@@ -58,7 +58,7 @@ def get_dlc_images(build_context):
         with open(test_env_file) as test_env:
             test_images = json.load(test_env)
         for dlc_test_type, images in test_images.items():
-            if dlc_test_type == "sanity":
+            if "sanity" in dlc_test_type:
                 return " ".join(images)
         raise RuntimeError(f"Cannot find any images for in {test_images}")
 

src/start_testbuilds.py

@@ -127,7 +127,12 @@ def is_test_job_enabled(test_type):
         return True
     if test_type == constants.EKS_TESTS and config.is_eks_test_enabled():
         return True
-    if test_type == constants.SANITY_TESTS and config.is_sanity_test_enabled():
+    if test_type == constants.SECURITY_SANITY_TESTS and config.is_security_sanity_test_enabled():
+        return True
+    if (
+        test_type == constants.FUNCTIONALITY_SANITY_TESTS
+        and config.is_functionality_sanity_test_enabled()
+    ):
         return True
 
     return False
@@ -244,7 +249,7 @@ def main():
             pr_test_job = f"dlc-pr-{test_type}-test"
             images_str = " ".join(images)
             # Maintaining separate codebuild project for graviton sanity test
-            if "graviton" in images_str and test_type == "sanity":
+            if "graviton" in images_str and "sanity" in test_type:
                 pr_test_job += "-graviton"
             if is_test_job_enabled(test_type) and is_test_job_implemented_for_framework(
                 images_str, test_type

src/utils.py

@@ -140,7 +140,8 @@ def fetch_dlc_images_for_test_jobs(images, use_latest_additional_tag=False):
         "eks": [],
         "ec2": [],
         "ec2-benchmark": [],
-        "sanity": [],
+        "security-sanity": [],
+        "functionality-sanity": [],
         "autopr": [],
     }
 

test/dlc_tests/conftest.py

@@ -38,6 +38,8 @@
     get_ecr_repo_name,
     UBUNTU_HOME_DIR,
     NightlyFeatureLabel,
+    is_mainline_context,
+    is_pr_context,
 )
 from test.test_utils.imageutils import are_image_labels_matched, are_fixture_labels_enabled
 from test.test_utils.test_reporting import TestReportGenerator
@@ -1029,6 +1031,49 @@ def _validate_pytorch_framework_version(request, image_uri, test_name, skip_dict
     return False
 
 
+@pytest.fixture(scope="session")
+def telemetry():
+    """
+    Telemetry tests are run in ec2 job in PR context but will run in its own job in MAINLINE context.
+    This fixture ensures that only telemetry tests are run in the `telemetry` job in the MAINLINE context.
+    """
+    is_telemetry_test_job = os.getenv("TEST_TYPE") == "telemetry"
+    if is_mainline_context() and not is_telemetry_test_job:
+        pytest.skip(
+            f"Test in not running in telemetry job in the pipeline context, Skipping current test."
+        )
+
+
+@pytest.fixture(scope="session")
+def security_sanity():
+    """
+    Skip test if job type is not `security_sanity` in either PR or Pipeline contexts.
+    Otherwise, sanity tests can run as usual in Canary/Deep Canary contexts.
+    Each sanity tests should only have either `security_sanity` or `functionality_sanity` fixtures.
+    Both should not be used at the same time. If neither are used, the test will run in both jobs.
+    """
+    is_security_sanity_test_job = os.getenv("TEST_TYPE") == "security_sanity"
+    if (is_pr_context() or is_mainline_context()) and not is_security_sanity_test_job:
+        pytest.skip(
+            f"Test in not running in `security_sanity` test type within the pipeline context, Skipping current test."
+        )
+
+
+@pytest.fixture(scope="session")
+def functionality_sanity():
+    """
+    Skip test if job type is not `functionality_sanity` in either PR or Pipeline contexts.
+    Otherwise, sanity tests can run as usual in Canary/Deep Canary contexts.
+    Each sanity tests should only have either `security_sanity` or `functionality_sanity` fixtures.
+    Both should not be used at the same time. If neither are used, the test will run in both jobs.
+    """
+    is_functionality_sanity_test_job = os.getenv("TEST_TYPE") == "functionality_sanity"
+    if (is_pr_context() or is_mainline_context()) and not is_functionality_sanity_test_job:
+        pytest.skip(
+            f"Test in not running in `functionality_sanity` test type within the pipeline context, Skipping current test."
+        )
+
+
 @pytest.fixture(scope="session")
 def dlc_images(request):
     return request.config.getoption("--images")
@@ -1630,7 +1675,7 @@ def pytest_generate_tests(metafunc):
                     if (
                         "stabilityai" not in metafunc.fixturenames
                         and "stabilityai" in image
-                        and os.getenv("TEST_TYPE") != "sanity"
+                        and "sanity" not in os.getenv("TEST_TYPE")
                     ):
                         LOGGER.info(
                             f"Skipping test, as this function is not marked as 'stabilityai'"