From 217a24ec273f04f28e2f2c488bfd5b9b60f07631 Mon Sep 17 00:00:00 2001
From: jzhaoqwa <jzhaoqwa@amazon.com>
Date: Wed, 16 Jul 2025 17:21:30 -0700
Subject: [PATCH 1/2] Fix import error in integ tests

---
 test/integration_tests/conftest.py                       | 4 ++--
 test/integration_tests/training/sdk/test_sdk_training.py | 7 -------
 2 files changed, 2 insertions(+), 9 deletions(-)

diff --git a/test/integration_tests/conftest.py b/test/integration_tests/conftest.py
index b9cedea2..e926c087 100644
--- a/test/integration_tests/conftest.py
+++ b/test/integration_tests/conftest.py
@@ -4,7 +4,7 @@
 from test.integration_tests.utils import execute_command
 from sagemaker.hyperpod.training import (
     HyperPodPytorchJob,
-    Container,
+    Containers,
     ReplicaSpec,
     Resources,
     RunPolicy,
@@ -49,7 +49,7 @@ def pytorch_job(test_job_name, image_uri):
             template=Template(
                 spec=Spec(
                     containers=[
-                        Container(
+                        Containers(
                             name="container-name",
                             image=image_uri,
                             image_pull_policy="Always",
diff --git a/test/integration_tests/training/sdk/test_sdk_training.py b/test/integration_tests/training/sdk/test_sdk_training.py
index 16998ad4..c92d3fdf 100644
--- a/test/integration_tests/training/sdk/test_sdk_training.py
+++ b/test/integration_tests/training/sdk/test_sdk_training.py
@@ -14,15 +14,8 @@
 import pytest
 import time
 import yaml
-
 from sagemaker.hyperpod.training import (
     HyperPodPytorchJob,
-    Container,
-    ReplicaSpec,
-    Resources,
-    RunPolicy,
-    Spec,
-    Template,
 )
 from sagemaker.hyperpod.common.config import Metadata
 from sagemaker.hyperpod.cli.utils import setup_logger

From 33fb92b4a93ca0e67bdbf6874e48e85c2421a14a Mon Sep 17 00:00:00 2001
From: jzhaoqwa <jzhaoqwa@amazon.com>
Date: Wed, 16 Jul 2025 17:35:10 -0700
Subject: [PATCH 2/2] Reduce PR check integ test to only run on Python 3.11

---
 .github/workflows/codebuild-ci.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/codebuild-ci.yml b/.github/workflows/codebuild-ci.yml
index 8ccf28ac..518d5686 100644
--- a/.github/workflows/codebuild-ci.yml
+++ b/.github/workflows/codebuild-ci.yml
@@ -67,7 +67,7 @@ jobs:
     strategy:
         fail-fast: false
         matrix:
-          python-version: ["39", "310", "311"]
+          python-version: ["311"]
     steps:
       - name: Configure AWS Credentials
         uses: aws-actions/configure-aws-credentials@v3