Fix import error in integ tests and update PR check (#134)

* Fix import error in integ tests

* Reduce PR check integ test to only run on Python 3.11

Author: zhaoqizqwang

.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

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

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