diff --git a/sagemaker-core/src/sagemaker/core/clarify/__init__.py b/sagemaker-core/src/sagemaker/core/clarify/__init__.py index d2206939e8..6363f5cc8e 100644 --- a/sagemaker-core/src/sagemaker/core/clarify/__init__.py +++ b/sagemaker-core/src/sagemaker/core/clarify/__init__.py @@ -380,12 +380,19 @@ "image/jpeg", "image/png", "application/x-npy", + "application/x-parquet", ), ), SchemaOptional("accept_type"): And( str, Use(str.lower), - lambda s: s in ("text/csv", "application/jsonlines", "application/json"), + lambda s: s + in ( + "text/csv", + "application/jsonlines", + "application/json", + "application/x-parquet", + ), ), SchemaOptional("label"): Or(str, int), SchemaOptional("probability"): Or(str, int), @@ -1021,11 +1028,13 @@ def __init__( Must be set with ``instance_count``, ``model_name`` accept_type (str): The model output format to be used for getting inferences with the shadow endpoint. Valid values are ``"text/csv"`` for CSV, - ``"application/jsonlines"`` for JSON Lines, and ``"application/json"`` for JSON. + ``"application/jsonlines"`` for JSON Lines, ``"application/json"`` for JSON, and + ``"application/x-parquet"`` for Parquet. Default is the same as ``content_type``. content_type (str): The model input format to be used for getting inferences with the shadow endpoint. Valid values are ``"text/csv"`` for CSV, - ``"application/jsonlines"`` for JSON Lines, and ``"application/json"`` for JSON. + ``"application/jsonlines"`` for JSON Lines, ``"application/json"`` for JSON, and + ``"application/x-parquet"`` for Parquet. Default is the same as ``dataset_format``. content_template (str): A template string to be used to construct the model input from dataset instances. It is only used, and required, when ``model_content_type`` is @@ -1161,10 +1170,15 @@ def __init__( ) self.predictor_config["endpoint_name_prefix"] = endpoint_name_prefix if accept_type is not None: - if accept_type not in ["text/csv", "application/jsonlines", "application/json"]: + if accept_type not in [ + "text/csv", + "application/jsonlines", + "application/json", + "application/x-parquet", + ]: raise ValueError( f"Invalid accept_type {accept_type}." - f" Please choose text/csv or application/jsonlines." + f" Please choose text/csv, application/jsonlines, application/json, or application/x-parquet." ) if time_series_model_config and accept_type == "text/csv": raise ValueError( @@ -1180,10 +1194,11 @@ def __init__( "image/jpg", "image/png", "application/x-npy", + "application/x-parquet", ]: raise ValueError( f"Invalid content_type {content_type}." - f" Please choose text/csv or application/jsonlines." + f" Please choose text/csv, application/jsonlines, application/json, image/jpeg, image/jpg, image/png, application/x-npy, or application/x-parquet." ) if content_type == "application/jsonlines": if content_template is None: diff --git a/sagemaker-core/src/sagemaker/core/workflow/job_config_document.py b/sagemaker-core/src/sagemaker/core/workflow/job_config_document.py index 63ed77eee5..09fe329f98 100644 --- a/sagemaker-core/src/sagemaker/core/workflow/job_config_document.py +++ b/sagemaker-core/src/sagemaker/core/workflow/job_config_document.py @@ -28,6 +28,7 @@ number or boolean reaches the service quoted. Parameterise string-valued fields only. """ + from __future__ import absolute_import import json @@ -67,9 +68,7 @@ def _placeholder(obj): if isinstance(obj, PipelineVariable): variables.append(obj) return "%s%d__" % (token_prefix, len(variables) - 1) - raise TypeError( - "Object of type %s is not JSON serializable" % obj.__class__.__name__ - ) + raise TypeError("Object of type %s is not JSON serializable" % obj.__class__.__name__) document = json.dumps(job_config, default=_placeholder) if not variables: diff --git a/sagemaker-core/src/sagemaker/core/workflow/pipeline_capture.py b/sagemaker-core/src/sagemaker/core/workflow/pipeline_capture.py index e48b059c68..8e319b64fb 100644 --- a/sagemaker-core/src/sagemaker/core/workflow/pipeline_capture.py +++ b/sagemaker-core/src/sagemaker/core/workflow/pipeline_capture.py @@ -11,6 +11,7 @@ # ANY KIND, either express or implied. See the License for the specific # language governing permissions and limitations under the License. """Shared capture paths for producers running under a ``PipelineSession``.""" + from __future__ import absolute_import from typing import Any, Dict, List, Optional @@ -87,9 +88,11 @@ def capture_training_request(pipeline_session, create_args: Dict[str, Any]) -> N # Dict-form tags arrive in both shapes: lowercase keys from the SDK's own # tag builders, PascalCase from callers passing wire-form dicts through. request["Tags"] = [ - Tag(key=tag.get("key", tag.get("Key")), value=tag.get("value", tag.get("Value"))) - if isinstance(tag, dict) - else tag + ( + Tag(key=tag.get("key", tag.get("Key")), value=tag.get("value", tag.get("Value"))) + if isinstance(tag, dict) + else tag + ) for tag in request["Tags"] ] pipeline_session._intercept_create_request(serialize(request), None, "train") diff --git a/sagemaker-core/tests/unit/test_clarify.py b/sagemaker-core/tests/unit/test_clarify.py index 5b2595fab2..fb5bddae30 100644 --- a/sagemaker-core/tests/unit/test_clarify.py +++ b/sagemaker-core/tests/unit/test_clarify.py @@ -566,6 +566,20 @@ def test_init_with_endpoint_name(self): config = ModelConfig(endpoint_name="my-endpoint") assert config.predictor_config["endpoint_name"] == "my-endpoint" + def test_init_with_parquet_content_and_accept_type(self): + """Test initialization with application/x-parquet content_type and accept_type.""" + from sagemaker.core.clarify import ModelConfig + + config = ModelConfig( + model_name="my-model", + instance_count=1, + instance_type="ml.m5.xlarge", + content_type="application/x-parquet", + accept_type="application/x-parquet", + ) + assert config.predictor_config["content_type"] == "application/x-parquet" + assert config.predictor_config["accept_type"] == "application/x-parquet" + def test_init_with_invalid_endpoint_prefix_raises_error(self): """Test that invalid endpoint_name_prefix raises ValueError.""" from sagemaker.core.clarify import ModelConfig diff --git a/sagemaker-core/tests/unit/workflow/test_job_config_document.py b/sagemaker-core/tests/unit/workflow/test_job_config_document.py index 5ad36a8746..ef3ffec949 100644 --- a/sagemaker-core/tests/unit/workflow/test_job_config_document.py +++ b/sagemaker-core/tests/unit/workflow/test_job_config_document.py @@ -11,6 +11,7 @@ # ANY KIND, either express or implied. See the License for the specific # language governing permissions and limitations under the License. """Tests for the ``JobConfigDocument`` encoder.""" + from __future__ import absolute_import import json diff --git a/sagemaker-core/tests/unit/workflow/test_pipeline_capture.py b/sagemaker-core/tests/unit/workflow/test_pipeline_capture.py index 3ad2849cae..5e7867682c 100644 --- a/sagemaker-core/tests/unit/workflow/test_pipeline_capture.py +++ b/sagemaker-core/tests/unit/workflow/test_pipeline_capture.py @@ -11,6 +11,7 @@ # ANY KIND, either express or implied. See the License for the specific # language governing permissions and limitations under the License. """Tests for the shared ``PipelineSession`` capture helpers.""" + from __future__ import absolute_import import json diff --git a/sagemaker-core/tests/unit/workflow/test_retrieve_caller_name.py b/sagemaker-core/tests/unit/workflow/test_retrieve_caller_name.py index ff08fa1d50..5de23c6fd9 100644 --- a/sagemaker-core/tests/unit/workflow/test_retrieve_caller_name.py +++ b/sagemaker-core/tests/unit/workflow/test_retrieve_caller_name.py @@ -19,6 +19,7 @@ -- `ProcessingStep` expects `"run"` from here while interception is called with `"process"` -- so a step guard cannot be satisfied by the interception name alone. """ + from __future__ import absolute_import import pytest