added check of the s3 URIs

Author: aws-byeldos

src/sagemaker/clarify.py

@@ -29,14 +29,26 @@
 
 from schema import Schema, And, Use, Or, Optional, Regex
 
-from sagemaker import image_uris, s3, utils, Session
+from sagemaker import image_uris, s3, utils
+from sagemaker.session import Session
 from sagemaker.network import NetworkConfig
 from sagemaker.processing import ProcessingInput, ProcessingOutput, Processor
 
 logger = logging.getLogger(__name__)
 
 
-ENDPOINT_NAME_PREFIX_PATTERN = "^[a-zA-Z0-9](-*[a-zA-Z0-9])"
+ENDPOINT_NAME_PREFIX_PATTERN = r"^[a-zA-Z0-9](-*[a-zA-Z0-9])"
+MODEL_NAME_PATTERN = r"^[a-zA-Z0-9]([\-a-zA-Z0-9]*[a-zA-Z0-9])?"
+
+
+def _validate_s3_path(path: str) -> str:
+    """Validates s3 path is correct"""
+    prefix = "s3://"
+    assert path.startswith(prefix)
+
+    assert "//" not in path[len(prefix) :]
+    assert not path.startswith(f"{prefix}/")
+    return path
 
 
 ANALYSIS_CONFIG_SCHEMA_V1_0 = Schema(
@@ -54,7 +66,7 @@
                 "application/x-image",
             ),
         ),
-        Optional("dataset_uri"): str,
+        Optional("dataset_uri"): And(str, _validate_s3_path),
         Optional("headers"): [str],
         Optional("label"): Or(str, int),
         # this field indicates user provides predicted_label in dataset
@@ -65,9 +77,9 @@
         Optional("facet"): [
             {"name_or_index": Or(str, int), Optional("value_or_threshold"): [Or(int, float, str)]}
         ],
-        Optional("facet_dataset_uri"): str,
+        Optional("facet_dataset_uri"): And(str, _validate_s3_path),
         Optional("facet_headers"): [str],
-        Optional("predicted_label_dataset_uri"): str,
+        Optional("predicted_label_dataset_uri"): And(str, _validate_s3_path),
         Optional("predicted_label_headers"): [str],
         Optional("excluded_columns"): [Or(int, str)],
         Optional("joinsource_name_or_index"): Or(str, int),
@@ -82,11 +94,11 @@
                         Or(
                             # CSV row
                             [Or(int, float, str, None)],
-                # JSON row (any JSON object). As I write this only
-                # SageMaker JSONLines Dense Format ([1])
-                # is supported and the validation is NOT done
-                # by the schema but by the data loader.
-                # [1] https://docs.aws.amazon.com/sagemaker/latest/dg/cdf-inference.html#cm-jsonlines
+                            # JSON row (any JSON object). As I write this only
+                            # SageMaker JSONLines Dense Format ([1])
+                            # is supported and the validation is NOT done
+                            # by the schema but by the data loader.
+                            # [1] https://docs.aws.amazon.com/sagemaker/latest/dg/cdf-inference.html#cm-jsonlines
                             {object: object},
                         )
                     ],
@@ -266,7 +278,7 @@
         Optional("predictor"): {
             Optional("endpoint_name"): str,
             Optional("endpoint_name_prefix"): And(str, Regex(ENDPOINT_NAME_PREFIX_PATTERN)),
-            Optional("model_name"): str,
+            Optional("model_name"): And(str, Regex(MODEL_NAME_PATTERN)),
             Optional("target_model"): str,
             Optional("instance_type"): str,
             Optional("initial_instance_count"): int,

tests/unit/test_clarify.py

@@ -30,6 +30,7 @@
     TextConfig,
     ImageConfig,
     _AnalysisConfigGenerator,
+    _validate_s3_path,
 )
 
 JOB_NAME_PREFIX = "my-prefix"
@@ -42,6 +43,21 @@ def test_uri():
     assert "306415355426.dkr.ecr.us-west-2.amazonaws.com/sagemaker-clarify-processing:1.0" == uri
 
 
+def test_validated_s3_path():
+    # Success
+    path = "s3://some_path/key"
+    assert _validate_s3_path(path) == path
+
+    with pytest.raises(AssertionError):
+        _validate_s3_path("wrong-prefix://some_path/key")
+
+    with pytest.raises(AssertionError):
+        _validate_s3_path("s3://some_path//key")
+
+    with pytest.raises(AssertionError):
+        _validate_s3_path("s3:///some_path/key")
+
+
 def test_data_config():
     # facets in input dataset
     s3_data_input_path = "s3://path/to/input.csv"