fix: integ tests

Author: evakravi

tests/integ/sagemaker/jumpstart/constants.py

@@ -46,6 +46,7 @@ def _to_s3_path(filename: str, s3_prefix: Optional[str]) -> str:
 TRAINING_DATASET_MODEL_DICT = {
     ("huggingface-spc-bert-base-cased", "1.0.0"): ("training-datasets/QNLI-tiny/"),
     ("huggingface-spc-bert-base-cased", "1.2.3"): ("training-datasets/QNLI-tiny/"),
+    ("huggingface-spc-bert-base-cased", "2.0.3"): ("training-datasets/QNLI-tiny/"),
     ("huggingface-spc-bert-base-cased", "*"): ("training-datasets/QNLI-tiny/"),
     ("js-trainable-model", "*"): ("training-datasets/QNLI-tiny/"),
     ("meta-textgeneration-llama-2-7b", "*"): ("training-datasets/sec_amazon/"),

tests/integ/sagemaker/jumpstart/retrieve_uri/inference.py

@@ -77,6 +77,10 @@ def package_artifacts(self):
 
         self.model_name = self.get_model_name()
 
+        if self.script_uri is None:
+            print("No script uri provided. Not performing prepack")
+            return self.model_uri
+
         cache_bucket_uri = f"s3://{get_test_artifact_bucket()}"
         repacked_model_uri = "/".join(
             [
@@ -147,16 +151,26 @@ def get_model_name(self) -> str:
         return f"{non_timestamped_name}{self.suffix}"
 
     def create_model(self) -> None:
+        primary_container = {
+            "Image": self.image_uri,
+            "Mode": "SingleModel",
+            "Environment": self.environment_variables,
+        }
+        if self.repacked_model_uri.endswith(".tar.gz"):
+            primary_container["ModelDataUrl"] = self.repacked_model_uri
+        else:
+            primary_container["ModelDataSource"] = {
+                "S3DataSource": {
+                    "S3Uri": self.repacked_model_uri,
+                    "S3DataType": "S3Prefix",
+                    "CompressionType": "None",
+                }
+            }
         self.sagemaker_client.create_model(
             ModelName=self.model_name,
             EnableNetworkIsolation=True,
             ExecutionRoleArn=self.execution_role,
-            PrimaryContainer={
-                "Image": self.image_uri,
-                "ModelDataUrl": self.repacked_model_uri,
-                "Mode": "SingleModel",
-                "Environment": self.environment_variables,
-            },
+            PrimaryContainer=primary_container,
         )
 
     def create_endpoint_config(self) -> None:

tests/integ/sagemaker/jumpstart/retrieve_uri/test_inference.py

@@ -46,13 +46,6 @@ def test_jumpstart_inference_retrieve_functions(setup):
         tolerate_vulnerable_model=True,
     )
 
-    script_uri = script_uris.retrieve(
-        model_id=model_id,
-        model_version=model_version,
-        script_scope="inference",
-        tolerate_vulnerable_model=True,
-    )
-
     model_uri = model_uris.retrieve(
         model_id=model_id,
         model_version=model_version,
@@ -68,7 +61,7 @@ def test_jumpstart_inference_retrieve_functions(setup):
 
     inference_job = InferenceJobLauncher(
         image_uri=image_uri,
-        script_uri=script_uri,
+        script_uri=None,
         model_uri=model_uri,
         instance_type=instance_type,
         base_name="catboost",