fix: codestyle-docs-test

mufaddal-rohawala · mufaddal-rohawala · commit 211b2edb5829 · 2024-03-13T15:32:21.000-07:00
diff --git a/src/sagemaker/base_predictor.py b/src/sagemaker/base_predictor.py
@@ -52,7 +52,7 @@
     JSONSerializer,
     NumpySerializer,
 )
-from sagemaker.iterators import LineIterator
+from sagemaker.iterators import ByteIterator
 from sagemaker.session import production_variant, Session
 from sagemaker.utils import name_from_base, stringify_object, format_tags
 
@@ -303,7 +303,7 @@ def predict_stream(
         custom_attributes=None,
         component_name: Optional[str] = None,
         target_container_hostname=None,
-        iterator=LineIterator,
+        iterator=ByteIterator,
     ):
         """Return the inference from the specified endpoint.
 
@@ -315,14 +315,14 @@ def predict_stream(
                 predict method then sends the bytes in the request body as is.
             initial_args (dict[str,str]): Optional. Default arguments for boto3
                 ``invoke_endpoint_with_response_stream`` call. Default is None (no default
-                arguments).
-            target_variant (str): The name of the production variant to run an inference
+                arguments). (Default: None)
+            target_variant (str): Optional. The name of the production variant to run an inference
                 request on (Default: None). Note that the ProductionVariant identifies the
                 model you want to host and the resources you want to deploy for hosting it.
-            inference_id (str): If you provide a value, it is added to the captured data
+            inference_id (str): Optional. If you provide a value, it is added to the captured data
                 when you enable data capture on the endpoint (Default: None).
-            custom_attributes (str): Provides additional information about a request for an
-                inference submitted to a model hosted at an Amazon SageMaker endpoint.
+            custom_attributes (str): Optional. Provides additional information about a request for
+                an inference submitted to a model hosted at an Amazon SageMaker endpoint.
                 The information is an opaque value that is forwarded verbatim. You could use this
                 value, for example, to provide an ID that you can use to track a request or to
                 provide other metadata that a service endpoint was programmed to process. The value
@@ -334,14 +334,14 @@ def predict_stream(
                 model can prepend the custom attribute with Trace ID: in your post-processing
                 function (Default: None).
             component_name (str): Optional. Name of the Amazon SageMaker inference component
-                corresponding the predictor.
-            target_container_hostname (str): If the endpoint hosts multiple containers and is
-                configured to use direct invocation, this parameter specifies the host name of the
-                container to invoke. (Default: None).
+                corresponding the predictor. (Default: None)
+            target_container_hostname (str): Optional. If the endpoint hosts multiple containers
+                and is configured to use direct invocation, this parameter specifies the host name
+                of the container to invoke. (Default: None).
             iterator (:class:`~sagemaker.iterators.BaseIterator`): An iterator class which provides
-                an iterable interface to deserialize a stream response from Inference Endpoint.
+                an iterable interface to iterate Event stream response from Inference Endpoint.
                 An object of the iterator class provided will be returned by the predict_stream
-                method (Default::class:`~sagemaker.iterators.LineIterator`). Iterators defined in
+                method (Default::class:`~sagemaker.iterators.ByteIterator`). Iterators defined in
                 :class:`~sagemaker.iterators` or custom iterators (needs to inherit
                 :class:`~sagemaker.iterators.BaseIterator`) can be specified as an input.
 
diff --git a/src/sagemaker/exceptions.py b/src/sagemaker/exceptions.py
@@ -89,6 +89,8 @@ def __init__(self, message):
 
 
 class ModelStreamError(Exception):
+    """Raised when invoke_endpoint_with_response_stream Response returns ModelStreamError"""
+
     def __init__(self, message="An error occurred", code=None):
         self.message = message
         self.code = code
@@ -99,6 +101,8 @@ def __init__(self, message="An error occurred", code=None):
 
 
 class InternalStreamFailure(Exception):
+    """Raised when invoke_endpoint_with_response_stream Response returns InternalStreamFailure"""
+
     def __init__(self, message="An error occurred"):
         self.message = message
         super().__init__(self.message)
diff --git a/src/sagemaker/iterators.py b/src/sagemaker/iterators.py
@@ -41,7 +41,7 @@ def handle_stream_errors(chunk):
 
 
 class BaseIterator(ABC):
-    """Abstract base class for creation of new iterators.
+    """Abstract base class for Inference Streaming iterators.
 
     Provides a skeleton for customization requiring the overriding of iterator methods
     __iter__ and __next__.
@@ -53,45 +53,75 @@ class BaseIterator(ABC):
     2. Needs to implement logic in __next__ to:
         2.1. Concatenate and provide next chunk of response from botocore.eventstream.EventStream.
             While doing so parse the response_chunk["PayloadPart"]["Bytes"].
-        2.2. Perform deserialization of response chunk based on expected response type.
-        2.3. If PayloadPart not in EventStream response, handle Errors.
+        2.2. If PayloadPart not in EventStream response, handle Errors
+            [Recommended to use `iterators.handle_stream_errors` method].
     """
 
-    def __init__(self, stream):
+    def __init__(self, event_stream):
         """Initialises a Iterator object to help parse the byte event stream input.
 
         Args:
-            stream: (botocore.eventstream.EventStream): Event Stream object to be iterated.
+            event_stream: (botocore.eventstream.EventStream): Event Stream object to be iterated.
         """
-        self.stream = stream
+        self.event_stream = event_stream
 
     @abstractmethod
     def __iter__(self):
-        """Abstract __iter__ method, returns an iterator object itself"""
+        """Abstract method, returns an iterator object itself"""
         return self
 
     @abstractmethod
     def __next__(self):
-        """Abstract __next__ method, is responsible for returning the next element in the
-        iteration"""
-        pass
+        """Abstract method, is responsible for returning the next element in the iteration"""
+
+
+class ByteIterator(BaseIterator):
+    """A helper class for parsing the byte Event Stream input to provide Byte iteration."""
+
+    def __init__(self, event_stream):
+        """Initialises a BytesIterator Iterator object
+
+        Args:
+            event_stream: (botocore.eventstream.EventStream): Event Stream object to be iterated.
+        """
+        super().__init__(event_stream)
+        self.byte_iterator = iter(event_stream)
+
+    def __iter__(self):
+        """Returns an iterator object itself, which allows the object to be iterated.
+
+        Returns:
+            iter : object
+                    An iterator object representing the iterable.
+        """
+        return self
+
+    def __next__(self):
+        """Returns the next chunk of Byte directly."""
+        # Even with "while True" loop the function still behaves like a generator
+        # and sends the next new byte chunk.
+        while True:
+            chunk = next(self.byte_iterator)
+            if "PayloadPart" not in chunk:
+                # handle API response errors and force terminate.
+                handle_stream_errors(chunk)
+                # print and move on to next response byte
+                print("Unknown event type:" + chunk)
+                continue
+            return chunk["PayloadPart"]["Bytes"]
 
 
 class LineIterator(BaseIterator):
-    """
-    A helper class for parsing the byte stream input and provide iteration on lines with
-    '\n' separators.
-    """
+    """A helper class for parsing the byte Event Stream input to provide Line iteration."""
 
-    def __init__(self, stream):
-        """Initialises a Iterator object to help parse the byte stream input and
-        provide iteration on lines with '\n' separators
+    def __init__(self, event_stream):
+        """Initialises a LineIterator Iterator object
 
         Args:
-            stream: (botocore.eventstream.EventStream): Event Stream object to be iterated.
+            event_stream: (botocore.eventstream.EventStream): Event Stream object to be iterated.
         """
-        super().__init__(stream)
-        self.byte_iterator = iter(self.stream)
+        super().__init__(event_stream)
+        self.byte_iterator = iter(self.event_stream)
         self.buffer = io.BytesIO()
         self.read_pos = 0
 
@@ -105,7 +135,8 @@ def __iter__(self):
         return self
 
     def __next__(self):
-        """
+        r"""Returns the next Line for an Line iterable.
+
         The output of the event stream will be in the following format:
 
         ```
@@ -146,8 +177,9 @@ def __next__(self):
                     continue
                 raise
             if "PayloadPart" not in chunk:
-                # handle errors within API Response if any.
+                # handle API response errors and force terminate.
                 handle_stream_errors(chunk)
+                # print and move on to next response byte
                 print("Unknown event type:" + chunk)
                 continue
             self.buffer.seek(0, io.SEEK_END)
diff --git a/tests/integ/test_predict_stream.py b/tests/integ/test_predict_stream.py
@@ -86,3 +86,24 @@ def test_predict_stream(sagemaker_session, endpoint_name):
         response += resp.get("outputs")[0]
 
     assert "AWS stands for Amazon Web Services." in response
+
+    data = {"inputs": "what does AWS stand for?", "parameters": {"max_new_tokens": 400}}
+    initial_args = {"ContentType": "application/json"}
+    predictor = Predictor(
+        endpoint_name=endpoint_name,
+        sagemaker_session=sagemaker_session,
+    )
+
+    # Validate that no exception is raised when the target_variant is specified.
+    # uses the default `sagemaker.iterator.ByteIterator`
+    stream_iterator = predictor.predict_stream(
+        data=json.dumps(data),
+        initial_args=initial_args,
+    )
+
+    response = ""
+    for line in stream_iterator:
+        resp = json.loads(line)
+        response += resp.get("outputs")[0]
+
+    assert "AWS stands for Amazon Web Services." in response
diff --git a/tests/unit/sagemaker/iterators/test_iterators.py b/tests/unit/sagemaker/iterators/test_iterators.py
@@ -1,11 +1,25 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+from __future__ import absolute_import
+
 import unittest
 from unittest.mock import MagicMock
 
 from sagemaker.exceptions import ModelStreamError, InternalStreamFailure
-from sagemaker.iterators import LineIterator
+from sagemaker.iterators import ByteIterator, LineIterator
 
 
-class TestLineIterator(unittest.TestCase):
+class TestByteIterator(unittest.TestCase):
     def test_iteration_with_payload_parts(self):
         # Mocking the stream object
         self.stream = MagicMock()
@@ -14,6 +28,60 @@ def test_iteration_with_payload_parts(self):
             {"PayloadPart": {"Bytes": b'{"outputs": [" challenging"]}\n'}},
             {"PayloadPart": {"Bytes": b'{"outputs": [" problem"]}\n'}},
         ]
+        self.iterator = ByteIterator(self.stream)
+
+        lines = list(self.iterator)
+        expected_lines = [
+            b'{"outputs": [" a"]}\n',
+            b'{"outputs": [" challenging"]}\n',
+            b'{"outputs": [" problem"]}\n',
+        ]
+        self.assertEqual(lines, expected_lines)
+
+    def test_iteration_with_model_stream_error(self):
+        # Mocking the stream object
+        self.stream = MagicMock()
+        self.stream.__iter__.return_value = [
+            {"PayloadPart": {"Bytes": b'{"outputs": [" a"]}\n'}},
+            {"PayloadPart": {"Bytes": b'{"outputs": [" challenging"]}\n'}},
+            {"ModelStreamError": {"Message": "Error message", "ErrorCode": "500"}},
+            {"PayloadPart": {"Bytes": b'{"outputs": [" problem"]}\n'}},
+        ]
+        self.iterator = ByteIterator(self.stream)
+
+        with self.assertRaises(ModelStreamError) as e:
+            list(self.iterator)
+
+        self.assertEqual(str(e.exception.message), "Error message")
+        self.assertEqual(str(e.exception.code), "500")
+
+    def test_iteration_with_internal_stream_failure(self):
+        # Mocking the stream object
+        self.stream = MagicMock()
+        self.stream.__iter__.return_value = [
+            {"PayloadPart": {"Bytes": b'{"outputs": [" a"]}\n'}},
+            {"PayloadPart": {"Bytes": b'{"outputs": [" challenging"]}\n'}},
+            {"InternalStreamFailure": {"Message": "Error internal stream failure"}},
+            {"PayloadPart": {"Bytes": b'{"outputs": [" problem"]}\n'}},
+        ]
+        self.iterator = ByteIterator(self.stream)
+
+        with self.assertRaises(InternalStreamFailure) as e:
+            list(self.iterator)
+
+        self.assertEqual(str(e.exception.message), "Error internal stream failure")
+
+
+class TestLineIterator(unittest.TestCase):
+    def test_iteration_with_payload_parts(self):
+        # Mocking the stream object
+        self.stream = MagicMock()
+        self.stream.__iter__.return_value = [
+            {"PayloadPart": {"Bytes": b'{"outputs": [" a"]}\n'}},
+            {"PayloadPart": {"Bytes": b'{"outputs": [" challenging"]}\n'}},
+            {"PayloadPart": {"Bytes": b'{"outputs": '}},
+            {"PayloadPart": {"Bytes": b'[" problem"]}\n'}},
+        ]
         self.iterator = LineIterator(self.stream)
 
         lines = list(self.iterator)
