Merge remote-tracking branch 'origin' into feat/jumpstart-model-estimator-classes

Author: evakravi

src/sagemaker/base_deserializers.py

@@ -189,7 +189,12 @@ def deserialize(self, stream, content_type):
 
 
 class NumpyDeserializer(SimpleBaseDeserializer):
-    """Deserialize a stream of data in .npy or UTF-8 CSV/JSON format to a numpy array."""
+    """Deserialize a stream of data in .npy, .npz or UTF-8 CSV/JSON format to a numpy array.
+
+    Note that when using application/x-npz archive format, the result will usually be a
+    dictionary-like object containing multiple arrays (as per ``numpy.load()``) - instead of a
+    single array.
+    """
 
     def __init__(self, dtype=None, accept="application/x-npy", allow_pickle=True):
         """Initialize a ``NumpyDeserializer`` instance.
@@ -223,6 +228,11 @@ def deserialize(self, stream, content_type):
                 return np.array(json.load(codecs.getreader("utf-8")(stream)), dtype=self.dtype)
             if content_type == "application/x-npy":
                 return np.load(io.BytesIO(stream.read()), allow_pickle=self.allow_pickle)
+            if content_type == "application/x-npz":
+                try:
+                    return np.load(io.BytesIO(stream.read()), allow_pickle=self.allow_pickle)
+                finally:
+                    stream.close()
         finally:
             stream.close()
 

src/sagemaker/lambda_helper.py

@@ -37,7 +37,6 @@ def __init__(
         memory_size: int = 128,
         runtime: str = "python3.8",
         vpc_config: dict = None,
-        architectures: list = None,
         environment: dict = None,
         layers: list = None,
     ):
@@ -71,8 +70,6 @@ def __init__(
             memory_size (int): Memory of the Lambda function in megabytes. Default is 128 MB.
             runtime (str): Runtime of the Lambda function. Default is set to python3.8.
             vpc_config (dict): VPC to deploy the Lambda function to. Default is None.
-            architectures (list): Which architecture to deploy to. Valid Values are
-                'x86_64' and 'arm64', default is None.
             environment (dict): Environment Variables for the Lambda function. Default is None.
             layers (list): List of Lambda layers for the Lambda function. Default is None.
         """
@@ -87,10 +84,9 @@ def __init__(
         self.timeout = timeout
         self.memory_size = memory_size
         self.runtime = runtime
-        self.vpc_config = vpc_config
-        self.environment = environment
-        self.architectures = architectures
-        self.layers = layers
+        self.vpc_config = vpc_config or {}
+        self.environment = environment or {}
+        self.layers = layers or []
 
         if function_arn is None and function_name is None:
             raise ValueError("Either function_arn or function_name must be provided.")
@@ -142,7 +138,6 @@ def create(self):
                 MemorySize=self.memory_size,
                 VpcConfig=self.vpc_config,
                 Environment=self.environment,
-                Architectures=self.architectures,
                 Layers=self.layers,
             )
             return response
@@ -163,7 +158,6 @@ def update(self):
                     response = lambda_client.update_function_code(
                         FunctionName=self.function_name or self.function_arn,
                         ZipFile=_zip_lambda_code(self.script),
-                        Architectures=self.architectures,
                     )
                 else:
                     bucket = self.s3_bucket or self.session.default_bucket()
@@ -186,7 +180,6 @@ def update(self):
                             zipped_code_dir=self.zipped_code_dir,
                             s3_bucket=bucket,
                         ),
-                        Architectures=self.architectures,
                     )
                 return response
             except ClientError as e:

src/sagemaker/multidatamodel.py

@@ -99,6 +99,7 @@ def __init__(
         self.model = model
         self.container_mode = MULTI_MODEL_CONTAINER_MODE
         self.sagemaker_session = sagemaker_session or Session()
+        self.endpoint_name = None
 
         if self.sagemaker_session.s3_client is None:
             self.s3_client = self.sagemaker_session.boto_session.client(

tests/integ/sagemaker/experiments/test_run.py

@@ -134,11 +134,15 @@ def test_run_name_vs_trial_component_name_edge_cases(sagemaker_session, input_na
         ) as run1:
             assert not run1._experiment.tags
             assert not run1._trial.tags
-            is_run_tc = is_run_trial_component(
-                trial_component_name=run1._trial_component.trial_component_name,
-                sagemaker_session=sagemaker_session,
-            )
-            assert is_run_tc
+
+            def verify_is_run():
+                is_run_tc = is_run_trial_component(
+                    trial_component_name=run1._trial_component.trial_component_name,
+                    sagemaker_session=sagemaker_session,
+                )
+                assert is_run_tc
+
+            retry_with_backoff(verify_is_run, 4)
 
         with load_run(
             experiment_name=exp_name,

tests/unit/sagemaker/deserializers/test_deserializers.py

@@ -164,6 +164,23 @@ def test_numpy_deserializer_from_npy_object_array_with_allow_pickle_false():
         numpy_deserializer.deserialize(stream, "application/x-npy")
 
 
+def test_numpy_deserializer_from_npz(numpy_deserializer):
+    arrays = {
+        "alpha": np.ones((2, 3)),
+        "beta": np.zeros((3, 2)),
+    }
+    stream = io.BytesIO()
+    np.savez_compressed(stream, **arrays)
+    stream.seek(0)
+
+    result = numpy_deserializer.deserialize(stream, "application/x-npz")
+
+    assert isinstance(result, np.lib.npyio.NpzFile)
+    assert set(arrays.keys()) == set(result.keys())
+    for key, arr in arrays.items():
+        assert np.array_equal(arr, result[key])
+
+
 @pytest.fixture
 def json_deserializer():
     return JSONDeserializer()

tests/unit/test_lambda_helper.py

@@ -187,10 +187,9 @@ def test_create_lambda_happycase1(sagemaker_session):
         Code=code,
         Timeout=120,
         MemorySize=128,
-        Architectures=None,
-        VpcConfig=None,
-        Environment=None,
-        Layers=None,
+        VpcConfig={},
+        Environment={},
+        Layers=[],
     )
 
 
@@ -216,10 +215,9 @@ def test_create_lambda_happycase2(sagemaker_session):
         Code=code,
         Timeout=120,
         MemorySize=128,
-        Architectures=None,
-        VpcConfig=None,
-        Environment=None,
-        Layers=None,
+        VpcConfig={},
+        Environment={},
+        Layers=[],
     )
 
 
@@ -231,7 +229,6 @@ def test_create_lambda_happycase3(sagemaker_session):
         script=SCRIPT,
         handler=HANDLER,
         session=sagemaker_session,
-        architectures=["x86_64"],
         environment={"Name": "my-test-lambda"},
         vpc_config={
             "SubnetIds": ["test-subnet-1"],
@@ -251,7 +248,6 @@ def test_create_lambda_happycase3(sagemaker_session):
         Code=code,
         Timeout=120,
         MemorySize=128,
-        Architectures=["x86_64"],
         VpcConfig={"SubnetIds": ["test-subnet-1"], "SecurityGroupIds": ["sec-group-1"]},
         Environment={"Name": "my-test-lambda"},
         Layers=["my-test-layer-1", "my-test-layer-2"],
@@ -314,7 +310,6 @@ def test_update_lambda_happycase1(sagemaker_session):
     sagemaker_session.lambda_client.update_function_code.assert_called_with(
         FunctionName=FUNCTION_NAME,
         ZipFile=ZIPPED_CODE,
-        Architectures=None,
     )
 
 
@@ -335,7 +330,6 @@ def test_update_lambda_happycase2(sagemaker_session):
         FunctionName=LAMBDA_ARN,
         S3Bucket=S3_BUCKET,
         S3Key=S3_KEY,
-        Architectures=None,
     )
 
 
@@ -347,7 +341,6 @@ def test_update_lambda_happycase3(sagemaker_session):
         script=SCRIPT,
         handler=HANDLER,
         session=sagemaker_session,
-        architectures=["x86_64"],
         environment={"Name": "my-test-lambda"},
         vpc_config={
             "SubnetIds": ["test-subnet-1"],
@@ -360,7 +353,6 @@ def test_update_lambda_happycase3(sagemaker_session):
     sagemaker_session.lambda_client.update_function_code.assert_called_with(
         FunctionName=FUNCTION_NAME,
         ZipFile=ZIPPED_CODE,
-        Architectures=["x86_64"],
     )
 
 
@@ -380,7 +372,6 @@ def test_update_lambda_s3bucket_not_provided(s3_upload, sagemaker_session):
         FunctionName=LAMBDA_ARN,
         S3Bucket=sagemaker_session.default_bucket(),
         S3Key=s3_upload.return_value,
-        Architectures=None,
     )
 
 
@@ -425,10 +416,9 @@ def test_upsert_lambda_happycase1(sagemaker_session):
         Code=code,
         Timeout=120,
         MemorySize=128,
-        Architectures=None,
-        VpcConfig=None,
-        Environment=None,
-        Layers=None,
+        VpcConfig={},
+        Environment={},
+        Layers=[],
     )
 
 
@@ -455,7 +445,7 @@ def test_upsert_lambda_happycase2(sagemaker_session):
     lambda_obj.upsert()
 
     sagemaker_session.lambda_client.update_function_code.assert_called_once_with(
-        FunctionName=FUNCTION_NAME, ZipFile=ZIPPED_CODE, Architectures=None
+        FunctionName=FUNCTION_NAME, ZipFile=ZIPPED_CODE
     )