change: Update data wrangler images

src/sagemaker/image_uri_config/data-wrangler.json

@@ -7,6 +7,35 @@
           "ap-east-1": "707077482487",
           "ap-northeast-1": "649008135260",
           "ap-northeast-2": "131546521161",
+          "ap-northeast-3": "913387583493",
+          "ap-south-1": "089933028263",
+          "ap-southeast-1": "119527597002",
+          "ap-southeast-2": "422173101802",
+          "ca-central-1": "557239378090",
+          "eu-central-1": "024640144536",
+          "eu-north-1": "054986407534",
+          "eu-south-1": "488287956546",
+          "eu-west-1": "245179582081",
+          "eu-west-2": "894491911112",
+          "eu-west-3": "807237891255",
+          "me-south-1": "376037874950",
+          "sa-east-1": "424196993095",
+          "us-east-1": "663277389841",
+          "us-east-2": "415577184552",
+          "us-west-1": "926135532090",
+          "us-west-2": "174368400705",
+          "cn-north-1": "245909111842",
+          "cn-northwest-1": "249157047649"
+        },
+        "repository": "sagemaker-data-wrangler-container"
+      },
+      "2.x": {
+        "registries": {
+          "af-south-1": "143210264188",
+          "ap-east-1": "707077482487",
+          "ap-northeast-1": "649008135260",
+          "ap-northeast-2": "131546521161",
+          "ap-northeast-3": "913387583493",
           "ap-south-1": "089933028263",
           "ap-southeast-1": "119527597002",
           "ap-southeast-2": "422173101802",

src/sagemaker/image_uris.py

@@ -36,6 +36,7 @@
 SKLEARN_FRAMEWORK = "sklearn"
 TRAINIUM_ALLOWED_FRAMEWORKS = "pytorch"
 INFERENCE_GRAVITON = "inference_graviton"
+DATA_WRANGLER_FRAMEWORK = "data-wrangler"
 
 
 @override_pipeline_parameter_var
@@ -461,6 +462,9 @@ def _validate_version_and_set_if_needed(version, config, framework):
 
         return available_versions[0]
 
+    if version is None and framework in [DATA_WRANGLER_FRAMEWORK]:
+        version = _get_latest_versions(available_versions)
+
     _validate_arg(version, available_versions + aliased_versions, "{} version".format(framework))
     return version
 

tests/unit/sagemaker/image_uris/test_data_wrangler.py

@@ -20,6 +20,7 @@
     "ap-east-1": "707077482487",
     "ap-northeast-1": "649008135260",
     "ap-northeast-2": "131546521161",
+    "ap-northeast-3": "913387583493",
     "ap-south-1": "089933028263",
     "ap-southeast-1": "119527597002",
     "ap-southeast-2": "422173101802",
@@ -39,15 +40,29 @@
     "cn-north-1": "245909111842",
     "cn-northwest-1": "249157047649",
 }
+VERSIONS = ["1.x", "2.x"]
 
 
 def test_data_wrangler_ecr_uri():
-    for region in DATA_WRANGLER_ACCOUNTS.keys():
-        actual_uri = image_uris.retrieve("data-wrangler", region=region)
-        expected_uri = expected_uris.algo_uri(
-            "sagemaker-data-wrangler-container",
-            DATA_WRANGLER_ACCOUNTS[region],
-            region,
-            version="1.x",
-        )
-        assert expected_uri == actual_uri
+    for version in VERSIONS:
+        for region in DATA_WRANGLER_ACCOUNTS.keys():
+            actual_uri = image_uris.retrieve("data-wrangler", region=region, version="1.x")
+            expected_uri = expected_uris.algo_uri(
+                "sagemaker-data-wrangler-container",
+                DATA_WRANGLER_ACCOUNTS[region],
+                region,
+                version="1.x",
+            )
+            assert expected_uri == actual_uri
+
+
+def test_data_wrangler_ecr_uri_none():
+    region = "us-west-2"
+    actual_uri = image_uris.retrieve("data-wrangler", region=region)
+    expected_uri = expected_uris.algo_uri(
+        "sagemaker-data-wrangler-container",
+        DATA_WRANGLER_ACCOUNTS[region],
+        region,
+        version=VERSIONS[-1],
+    )
+    assert expected_uri == actual_uri

tests/unit/sagemaker/wrangler/test_processing.py

@@ -23,7 +23,7 @@
 REGION = "us-west-2"
 DATA_WRANGLER_RECIPE_SOURCE = "s3://data_wrangler_flows/flow-26-18-43-16-0b48ac2e.flow"
 DATA_WRANGLER_CONTAINER_URI = (
-    "174368400705.dkr.ecr.us-west-2.amazonaws.com/sagemaker-data-wrangler-container:1.x"
+    "174368400705.dkr.ecr.us-west-2.amazonaws.com/sagemaker-data-wrangler-container:2.x"
 )
 MOCK_S3_URI = "s3://mock_data/mock.csv"