Add back efa configs to SM tests

src/config/build_config.py

@@ -5,9 +5,9 @@
 # Do remember to revert it back to False before merging any PR (including NEURON dedicated PR)
 ENABLE_NEURON_MODE = False
 # Frameworks for which you want to disable both builds and tests
-DISABLE_FRAMEWORK_TESTS = []
+DISABLE_FRAMEWORK_TESTS = ["pytorch", "mxnet", "huggingface_pytorch", "huggingface_tensorflow"]
 # Disable new builds or build without datetime tag
-DISABLE_DATETIME_TAG = False
+DISABLE_DATETIME_TAG = True
 # Note: Need to build the images at least once with DISABLE_DATETIME_TAG = True
 # before disabling new builds or tests will fail
 DISABLE_NEW_BUILDS = False

src/config/test_config.py

@@ -9,7 +9,7 @@
 
 DISABLE_SANITY_TESTS = False
 DISABLE_SAGEMAKER_TESTS = False
-DISABLE_ECS_TESTS = False
-DISABLE_EKS_TESTS = False
-DISABLE_EC2_TESTS = False
+DISABLE_ECS_TESTS = True
+DISABLE_EKS_TESTS = True
+DISABLE_EC2_TESTS = True
 USE_SCHEDULER = False

test/sagemaker_tests/pytorch/training/conftest.py

@@ -105,6 +105,20 @@ def pytest_addoption(parser):
     parser.addoption('--tag', default=None)
     parser.addoption('--generate-coverage-doc', default=False, action='store_true',
                      help='use this option to generate test coverage doc')
+    parser.addoption(
+        "--efa", action="store_true", default=False, help="Run only efa tests",
+    )
+
+
+def pytest_configure(config):
+    config.addinivalue_line("markers", "efa(): explicitly mark to run efa tests")
+
+
+def pytest_runtest_setup(item):
+    if item.config.getoption("--efa"):
+        efa_tests = [mark for mark in item.iter_markers(name="efa")]
+        if not efa_tests:
+            pytest.skip("Skipping non-efa tests")
 
 
 def pytest_collection_modifyitems(session, config, items):

test/sagemaker_tests/tensorflow/tensorflow2_training/integration/conftest.py

@@ -42,6 +42,16 @@ def pytest_addoption(parser):
     parser.addoption('--instance-type', default=None)
     parser.addoption('--generate-coverage-doc', default=False, action='store_true',
                      help='use this option to generate test coverage doc')
+    parser.addoption(
+        "--efa", action="store_true", default=False, help="Run only efa tests",
+    )
+
+
+def pytest_runtest_setup(item):
+    if item.config.getoption("--efa"):
+        efa_tests = [mark for mark in item.iter_markers("efa")]
+        if not efa_tests:
+            pytest.skip("Skipping non-efa tests")
 
 
 def pytest_collection_modifyitems(session, config, items):
@@ -54,6 +64,7 @@ def pytest_collection_modifyitems(session, config, items):
 def pytest_configure(config):
     os.environ['TEST_PY_VERSIONS'] = config.getoption('--py-version')
     os.environ['TEST_PROCESSORS'] = config.getoption('--processor')
+    config.addinivalue_line("markers", "efa(): explicitly mark to run efa tests")
 
 
 @pytest.fixture(scope='session')

test/test_utils/sagemaker.py

@@ -140,9 +140,10 @@ def generate_sagemaker_pytest_cmd(image, sagemaker_test_type):
     is_py3 = " python3 -m "
 
     efa_flag = ""
-    if job_type == "training" and (framework == "tensorflow" or framework == "pytorch"):
-        efa_dedicated = os.getenv("EFA_DEDICATED", "False").lower() == "true"
-        efa_flag = '--efa' if efa_dedicated else '-m not efa'
+    if job_type == "training":
+        if (framework == "tensorflow" and framework_major_version >= "2") or framework == "pytorch":
+            efa_dedicated = os.getenv("EFA_DEDICATED", "False").lower() == "true"
+            efa_flag = '--efa' if efa_dedicated else '-m \"not efa\"'
 
     remote_pytest_cmd = (
         f"pytest -rA {integration_path} --region {region} --processor {processor} {docker_base_arg} "