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
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