sagemaker-core 1.0.3__py3-none-any.whl → 1.0.5__py3-none-any.whl

sagemaker_core/main/resources.py

@@ -27,7 +27,6 @@ from sagemaker_core.main.code_injection.codec import transform
 from sagemaker_core.main.code_injection.constants import Color
 from sagemaker_core.main.utils import (
     SageMakerClient,
-    SageMakerRuntimeClient,
     ResourceIterator,
     Unassigned,
     get_textual_rich_logger,
@@ -35,7 +34,6 @@ from sagemaker_core.main.utils import (
     pascal_to_snake,
     is_not_primitive,
     is_not_str_dict,
-    is_snake_case,
     is_primitive_list,
     serialize,
 )
@@ -55,9 +53,9 @@ class Base(BaseModel):
 
     @classmethod
     def get_sagemaker_client(cls, session=None, region_name=None, service_name="sagemaker"):
-        return SageMakerClient(
-            session=session, region_name=region_name, service_name=service_name
-        ).client
+        return SageMakerClient(session=session, region_name=region_name).get_client(
+            service_name=service_name
+        )
 
     @staticmethod
     def get_updated_kwargs_with_configured_attributes(
@@ -3198,6 +3196,8 @@ class Cluster(Base):
         creation_time: The time when the SageMaker Cluster is created.
         failure_message: The failure message of the SageMaker HyperPod cluster.
         vpc_config:
+        orchestrator: The type of orchestrator used for the SageMaker HyperPod cluster.
+        node_recovery: The node recovery mode configured for the SageMaker HyperPod cluster.
 
     """
 
@@ -3208,6 +3208,8 @@ class Cluster(Base):
     failure_message: Optional[str] = Unassigned()
     instance_groups: Optional[List[ClusterInstanceGroupDetails]] = Unassigned()
     vpc_config: Optional[VpcConfig] = Unassigned()
+    orchestrator: Optional[ClusterOrchestrator] = Unassigned()
+    node_recovery: Optional[str] = Unassigned()
 
     def get_name(self) -> str:
         attributes = vars(self)
@@ -3252,6 +3254,8 @@ class Cluster(Base):
         instance_groups: List[ClusterInstanceGroupSpecification],
         vpc_config: Optional[VpcConfig] = Unassigned(),
         tags: Optional[List[Tag]] = Unassigned(),
+        orchestrator: Optional[ClusterOrchestrator] = Unassigned(),
+        node_recovery: Optional[str] = Unassigned(),
         session: Optional[Session] = None,
         region: Optional[str] = None,
     ) -> Optional["Cluster"]:
@@ -3263,6 +3267,8 @@ class Cluster(Base):
             instance_groups: The instance groups to be created in the SageMaker HyperPod cluster.
             vpc_config:
             tags: Custom tags for managing the SageMaker HyperPod cluster as an Amazon Web Services resource. You can add tags to your cluster in the same way you add them in other Amazon Web Services services that support tagging. To learn more about tagging Amazon Web Services resources in general, see Tagging Amazon Web Services Resources User Guide.
+            orchestrator: The type of orchestrator to use for the SageMaker HyperPod cluster. Currently, the only supported value is "eks", which is to use an Amazon Elastic Kubernetes Service (EKS) cluster as the orchestrator.
+            node_recovery: The node recovery mode for the SageMaker HyperPod cluster. When set to Automatic, SageMaker HyperPod will automatically reboot or replace faulty nodes when issues are detected. When set to None, cluster administrators will need to manually manage any faulty cluster instances.
             session: Boto3 session.
             region: Region name.
 
@@ -3296,6 +3302,8 @@ class Cluster(Base):
             "InstanceGroups": instance_groups,
             "VpcConfig": vpc_config,
             "Tags": tags,
+            "Orchestrator": orchestrator,
+            "NodeRecovery": node_recovery,
         }
 
         operation_input_args = Base.populate_chained_attributes(
@@ -3406,6 +3414,7 @@ class Cluster(Base):
     def update(
         self,
         instance_groups: List[ClusterInstanceGroupSpecification],
+        node_recovery: Optional[str] = Unassigned(),
     ) -> Optional["Cluster"]:
         """
         Update a Cluster resource
@@ -3434,6 +3443,7 @@ class Cluster(Base):
         operation_input_args = {
             "ClusterName": self.cluster_name,
             "InstanceGroups": instance_groups,
+            "NodeRecovery": node_recovery,
         }
         logger.debug(f"Input request: {operation_input_args}")
         # serialize the input request
@@ -8095,6 +8105,125 @@ class Endpoint(Base):
                     raise e
                 time.sleep(poll)
 
+    @classmethod
+    @Base.add_validate_call
+    def get_all(
+        cls,
+        sort_by: Optional[str] = Unassigned(),
+        sort_order: Optional[str] = Unassigned(),
+        name_contains: Optional[str] = Unassigned(),
+        creation_time_before: Optional[datetime.datetime] = Unassigned(),
+        creation_time_after: Optional[datetime.datetime] = Unassigned(),
+        last_modified_time_before: Optional[datetime.datetime] = Unassigned(),
+        last_modified_time_after: Optional[datetime.datetime] = Unassigned(),
+        status_equals: Optional[str] = Unassigned(),
+        session: Optional[Session] = None,
+        region: Optional[str] = None,
+    ) -> ResourceIterator["Endpoint"]:
+        """
+        Get all Endpoint resources
+
+        Parameters:
+            sort_by: Sorts the list of results. The default is CreationTime.
+            sort_order: The sort order for results. The default is Descending.
+            next_token: If the result of a ListEndpoints request was truncated, the response includes a NextToken. To retrieve the next set of endpoints, use the token in the next request.
+            max_results: The maximum number of endpoints to return in the response. This value defaults to 10.
+            name_contains: A string in endpoint names. This filter returns only endpoints whose name contains the specified string.
+            creation_time_before: A filter that returns only endpoints that were created before the specified time (timestamp).
+            creation_time_after: A filter that returns only endpoints with a creation time greater than or equal to the specified time (timestamp).
+            last_modified_time_before:  A filter that returns only endpoints that were modified before the specified timestamp.
+            last_modified_time_after:  A filter that returns only endpoints that were modified after the specified timestamp.
+            status_equals:  A filter that returns only endpoints with the specified status.
+            session: Boto3 session.
+            region: Region name.
+
+        Returns:
+            Iterator for listed Endpoint resources.
+
+        Raises:
+            botocore.exceptions.ClientError: This exception is raised for AWS service related errors.
+                The error message and error code can be parsed from the exception as follows:
+                ```
+                try:
+                    # AWS service call here
+                except botocore.exceptions.ClientError as e:
+                    error_message = e.response['Error']['Message']
+                    error_code = e.response['Error']['Code']
+                ```
+        """
+
+        client = Base.get_sagemaker_client(
+            session=session, region_name=region, service_name="sagemaker"
+        )
+
+        operation_input_args = {
+            "SortBy": sort_by,
+            "SortOrder": sort_order,
+            "NameContains": name_contains,
+            "CreationTimeBefore": creation_time_before,
+            "CreationTimeAfter": creation_time_after,
+            "LastModifiedTimeBefore": last_modified_time_before,
+            "LastModifiedTimeAfter": last_modified_time_after,
+            "StatusEquals": status_equals,
+        }
+
+        # serialize the input request
+        operation_input_args = serialize(operation_input_args)
+        logger.debug(f"Serialized input request: {operation_input_args}")
+
+        return ResourceIterator(
+            client=client,
+            list_method="list_endpoints",
+            summaries_key="Endpoints",
+            summary_name="EndpointSummary",
+            resource_cls=Endpoint,
+            list_method_kwargs=operation_input_args,
+        )
+
+    @Base.add_validate_call
+    def update_weights_and_capacities(
+        self,
+        desired_weights_and_capacities: List[DesiredWeightAndCapacity],
+        session: Optional[Session] = None,
+        region: Optional[str] = None,
+    ) -> None:
+        """
+        Updates variant weight of one or more variants associated with an existing endpoint, or capacity of one variant associated with an existing endpoint.
+
+        Parameters:
+            desired_weights_and_capacities: An object that provides new capacity and weight values for a variant.
+            session: Boto3 session.
+            region: Region name.
+
+        Raises:
+            botocore.exceptions.ClientError: This exception is raised for AWS service related errors.
+                The error message and error code can be parsed from the exception as follows:
+                ```
+                try:
+                    # AWS service call here
+                except botocore.exceptions.ClientError as e:
+                    error_message = e.response['Error']['Message']
+                    error_code = e.response['Error']['Code']
+                ```
+            ResourceLimitExceeded: You have exceeded an SageMaker resource limit. For example, you might have too many training jobs created.
+        """
+
+        operation_input_args = {
+            "EndpointName": self.endpoint_name,
+            "DesiredWeightsAndCapacities": desired_weights_and_capacities,
+        }
+        # serialize the input request
+        operation_input_args = serialize(operation_input_args)
+        logger.debug(f"Serialized input request: {operation_input_args}")
+
+        client = Base.get_sagemaker_client(
+            session=session, region_name=region, service_name="sagemaker"
+        )
+
+        logger.debug(f"Calling update_endpoint_weights_and_capacities API")
+        response = client.update_endpoint_weights_and_capacities(**operation_input_args)
+        logger.debug(f"Response: {response}")
+
     @Base.add_validate_call
     def invoke(
         self,
@@ -8108,9 +8237,12 @@ class Endpoint(Base):
         inference_id: Optional[str] = Unassigned(),
         enable_explanations: Optional[str] = Unassigned(),
         inference_component_name: Optional[str] = Unassigned(),
-    ) -> Optional[object]:
+        session_id: Optional[str] = Unassigned(),
+        session: Optional[Session] = None,
+        region: Optional[str] = None,
+    ) -> Optional[InvokeEndpointOutput]:
         """
-        Invoke a Endpoint resource
+        After you deploy a model into production using Amazon SageMaker hosting services, your client applications use this API to get inferences from the model hosted at the specified endpoint.
 
         Parameters:
             body: Provides input data, in the format specified in the ContentType request header. Amazon SageMaker passes all of the data in the body to the model.  For information about the format of the request body, see Common Data Formats-Inference.
@@ -8123,9 +8255,12 @@ class Endpoint(Base):
             inference_id: If you provide a value, it is added to the captured data when you enable data capture on the endpoint. For information about data capture, see Capture Data.
             enable_explanations: An optional JMESPath expression used to override the EnableExplanations parameter of the ClarifyExplainerConfig API. See the EnableExplanations section in the developer guide for more information.
             inference_component_name: If the endpoint hosts one or more inference components, this parameter specifies the name of inference component to invoke.
+            session_id: Creates a stateful session or identifies an existing one. You can do one of the following:   Create a stateful session by specifying the value NEW_SESSION.   Send your request to an existing stateful session by specifying the ID of that session.   With a stateful session, you can send multiple requests to a stateful model. When you create a session with a stateful model, the model must create the session ID and set the expiration time. The model must also provide that information in the response to your request. You can get the ID and timestamp from the NewSessionId response parameter. For any subsequent request where you specify that session ID, SageMaker routes the request to the same instance that supports the session.
+            session: Boto3 session.
+            region: Region name.
 
         Returns:
-            The Invoke response.
+            InvokeEndpointOutput
 
         Raises:
             botocore.exceptions.ClientError: This exception is raised for AWS service related errors.
@@ -8138,15 +8273,13 @@ class Endpoint(Base):
                     error_code = e.response['Error']['Code']
                 ```
             InternalDependencyException: Your request caused an exception with an internal dependency. Contact customer support.
-            InternalFailure: An internal failure occurred.
+            InternalFailure: An internal failure occurred. Try your request again. If the problem persists, contact Amazon Web Services customer support.
             ModelError: Model (owned by the customer in the container) returned 4xx or 5xx error code.
             ModelNotReadyException: Either a serverless endpoint variant's resources are still being provisioned, or a multi-model endpoint is still downloading or loading the target model. Wait and try your request again.
-            ServiceUnavailable: The service is unavailable. Try your call again.
-            ValidationError: Inspect your request and try again.
+            ServiceUnavailable: The service is currently unavailable.
+            ValidationError: There was an error validating your request.
         """
 
-        logger.info(f"Invoking endpoint resource.")
-        client = SageMakerRuntimeClient(service_name="sagemaker-runtime").client
         operation_input_args = {
             "EndpointName": self.endpoint_name,
             "Body": body,
@@ -8159,17 +8292,22 @@ class Endpoint(Base):
             "InferenceId": inference_id,
             "EnableExplanations": enable_explanations,
             "InferenceComponentName": inference_component_name,
+            "SessionId": session_id,
         }
-        logger.debug(f"Input request: {operation_input_args}")
         # serialize the input request
         operation_input_args = serialize(operation_input_args)
         logger.debug(f"Serialized input request: {operation_input_args}")
 
-        # create the resource
+        client = Base.get_sagemaker_client(
+            session=session, region_name=region, service_name="sagemaker-runtime"
+        )
+
+        logger.debug(f"Calling invoke_endpoint API")
         response = client.invoke_endpoint(**operation_input_args)
         logger.debug(f"Response: {response}")
 
-        return response
+        transformed_response = transform(response, "InvokeEndpointOutput")
+        return InvokeEndpointOutput(**transformed_response)
 
     @Base.add_validate_call
     def invoke_async(
@@ -8181,9 +8319,11 @@ class Endpoint(Base):
         inference_id: Optional[str] = Unassigned(),
         request_ttl_seconds: Optional[int] = Unassigned(),
         invocation_timeout_seconds: Optional[int] = Unassigned(),
-    ) -> Optional[object]:
+        session: Optional[Session] = None,
+        region: Optional[str] = None,
+    ) -> Optional[InvokeEndpointAsyncOutput]:
         """
-        Invoke Async a Endpoint resource
+        After you deploy a model into production using Amazon SageMaker hosting services, your client applications use this API to get inferences from the model hosted at the specified endpoint in an asynchronous manner.
 
         Parameters:
             input_location: The Amazon S3 URI where the inference request payload is stored.
@@ -8193,9 +8333,11 @@ class Endpoint(Base):
             inference_id: The identifier for the inference request. Amazon SageMaker will generate an identifier for you if none is specified.
             request_ttl_seconds: Maximum age in seconds a request can be in the queue before it is marked as expired. The default is 6 hours, or 21,600 seconds.
             invocation_timeout_seconds: Maximum amount of time in seconds a request can be processed before it is marked as expired. The default is 15 minutes, or 900 seconds.
+            session: Boto3 session.
+            region: Region name.
 
         Returns:
-            The Invoke response.
+            InvokeEndpointAsyncOutput
 
         Raises:
             botocore.exceptions.ClientError: This exception is raised for AWS service related errors.
@@ -8207,14 +8349,11 @@ class Endpoint(Base):
                     error_message = e.response['Error']['Message']
                     error_code = e.response['Error']['Code']
                 ```
-            InternalFailure: An internal failure occurred.
-            ServiceUnavailable: The service is unavailable. Try your call again.
-            ValidationError: Inspect your request and try again.
+            InternalFailure: An internal failure occurred. Try your request again. If the problem persists, contact Amazon Web Services customer support.
+            ServiceUnavailable: The service is currently unavailable.
+            ValidationError: There was an error validating your request.
         """
 
-        logger.info(f"Invoking endpoint resource Async.")
-        client = SageMakerRuntimeClient(service_name="sagemaker-runtime").client
-
         operation_input_args = {
             "EndpointName": self.endpoint_name,
             "ContentType": content_type,
@@ -8225,16 +8364,20 @@ class Endpoint(Base):
             "RequestTTLSeconds": request_ttl_seconds,
             "InvocationTimeoutSeconds": invocation_timeout_seconds,
         }
-        logger.debug(f"Input request: {operation_input_args}")
         # serialize the input request
         operation_input_args = serialize(operation_input_args)
         logger.debug(f"Serialized input request: {operation_input_args}")
 
-        # create the resource
+        client = Base.get_sagemaker_client(
+            session=session, region_name=region, service_name="sagemaker-runtime"
+        )
+
+        logger.debug(f"Calling invoke_endpoint_async API")
         response = client.invoke_endpoint_async(**operation_input_args)
         logger.debug(f"Response: {response}")
 
-        return response
+        transformed_response = transform(response, "InvokeEndpointAsyncOutput")
+        return InvokeEndpointAsyncOutput(**transformed_response)
 
     @Base.add_validate_call
     def invoke_with_response_stream(
@@ -8247,9 +8390,12 @@ class Endpoint(Base):
         target_container_hostname: Optional[str] = Unassigned(),
         inference_id: Optional[str] = Unassigned(),
         inference_component_name: Optional[str] = Unassigned(),
-    ) -> Optional[object]:
+        session_id: Optional[str] = Unassigned(),
+        session: Optional[Session] = None,
+        region: Optional[str] = None,
+    ) -> Optional[InvokeEndpointWithResponseStreamOutput]:
         """
-        Invoke with response stream a Endpoint resource
+        Invokes a model at the specified endpoint to return the inference response as a stream.
 
         Parameters:
             body: Provides input data, in the format specified in the ContentType request header. Amazon SageMaker passes all of the data in the body to the model.  For information about the format of the request body, see Common Data Formats-Inference.
@@ -8260,9 +8406,12 @@ class Endpoint(Base):
             target_container_hostname: If the endpoint hosts multiple containers and is configured to use direct invocation, this parameter specifies the host name of the container to invoke.
             inference_id: An identifier that you assign to your request.
             inference_component_name: If the endpoint hosts one or more inference components, this parameter specifies the name of inference component to invoke for a streaming response.
+            session_id: The ID of a stateful session to handle your request. You can't create a stateful session by using the InvokeEndpointWithResponseStream action. Instead, you can create one by using the  InvokeEndpoint  action. In your request, you specify NEW_SESSION for the SessionId request parameter. The response to that request provides the session ID for the NewSessionId response parameter.
+            session: Boto3 session.
+            region: Region name.
 
         Returns:
-            The Invoke response.
+            InvokeEndpointWithResponseStreamOutput
 
         Raises:
             botocore.exceptions.ClientError: This exception is raised for AWS service related errors.
@@ -8274,17 +8423,14 @@ class Endpoint(Base):
                     error_message = e.response['Error']['Message']
                     error_code = e.response['Error']['Code']
                 ```
-            InternalFailure: An internal failure occurred.
+            InternalFailure: An internal failure occurred. Try your request again. If the problem persists, contact Amazon Web Services customer support.
             InternalStreamFailure: The stream processing failed because of an unknown error, exception or failure. Try your request again.
             ModelError: Model (owned by the customer in the container) returned 4xx or 5xx error code.
             ModelStreamError: An error occurred while streaming the response body. This error can have the following error codes:  ModelInvocationTimeExceeded  The model failed to finish sending the response within the timeout period allowed by Amazon SageMaker.  StreamBroken  The Transmission Control Protocol (TCP) connection between the client and the model was reset or closed.
-            ServiceUnavailable: The service is unavailable. Try your call again.
-            ValidationError: Inspect your request and try again.
+            ServiceUnavailable: The service is currently unavailable.
+            ValidationError: There was an error validating your request.
         """
 
-        logger.info(f"Invoking endpoint resource with Response Stream.")
-        client = SageMakerRuntimeClient(service_name="sagemaker-runtime").client
-
         operation_input_args = {
             "EndpointName": self.endpoint_name,
             "Body": body,
@@ -8295,137 +8441,23 @@ class Endpoint(Base):
             "TargetContainerHostname": target_container_hostname,
             "InferenceId": inference_id,
             "InferenceComponentName": inference_component_name,
-        }
-        logger.debug(f"Input request: {operation_input_args}")
-        # serialize the input request
-        operation_input_args = serialize(operation_input_args)
-        logger.debug(f"Serialized input request: {operation_input_args}")
-
-        # create the resource
-        response = client.invoke_endpoint_with_response_stream(**operation_input_args)
-        logger.debug(f"Response: {response}")
-
-        return response
-
-    @classmethod
-    @Base.add_validate_call
-    def get_all(
-        cls,
-        sort_by: Optional[str] = Unassigned(),
-        sort_order: Optional[str] = Unassigned(),
-        name_contains: Optional[str] = Unassigned(),
-        creation_time_before: Optional[datetime.datetime] = Unassigned(),
-        creation_time_after: Optional[datetime.datetime] = Unassigned(),
-        last_modified_time_before: Optional[datetime.datetime] = Unassigned(),
-        last_modified_time_after: Optional[datetime.datetime] = Unassigned(),
-        status_equals: Optional[str] = Unassigned(),
-        session: Optional[Session] = None,
-        region: Optional[str] = None,
-    ) -> ResourceIterator["Endpoint"]:
-        """
-        Get all Endpoint resources
-
-        Parameters:
-            sort_by: Sorts the list of results. The default is CreationTime.
-            sort_order: The sort order for results. The default is Descending.
-            next_token: If the result of a ListEndpoints request was truncated, the response includes a NextToken. To retrieve the next set of endpoints, use the token in the next request.
-            max_results: The maximum number of endpoints to return in the response. This value defaults to 10.
-            name_contains: A string in endpoint names. This filter returns only endpoints whose name contains the specified string.
-            creation_time_before: A filter that returns only endpoints that were created before the specified time (timestamp).
-            creation_time_after: A filter that returns only endpoints with a creation time greater than or equal to the specified time (timestamp).
-            last_modified_time_before:  A filter that returns only endpoints that were modified before the specified timestamp.
-            last_modified_time_after:  A filter that returns only endpoints that were modified after the specified timestamp.
-            status_equals:  A filter that returns only endpoints with the specified status.
-            session: Boto3 session.
-            region: Region name.
-
-        Returns:
-            Iterator for listed Endpoint resources.
-
-        Raises:
-            botocore.exceptions.ClientError: This exception is raised for AWS service related errors.
-                The error message and error code can be parsed from the exception as follows:
-                ```
-                try:
-                    # AWS service call here
-                except botocore.exceptions.ClientError as e:
-                    error_message = e.response['Error']['Message']
-                    error_code = e.response['Error']['Code']
-                ```
-        """
-
-        client = Base.get_sagemaker_client(
-            session=session, region_name=region, service_name="sagemaker"
-        )
-
-        operation_input_args = {
-            "SortBy": sort_by,
-            "SortOrder": sort_order,
-            "NameContains": name_contains,
-            "CreationTimeBefore": creation_time_before,
-            "CreationTimeAfter": creation_time_after,
-            "LastModifiedTimeBefore": last_modified_time_before,
-            "LastModifiedTimeAfter": last_modified_time_after,
-            "StatusEquals": status_equals,
-        }
-
-        # serialize the input request
-        operation_input_args = serialize(operation_input_args)
-        logger.debug(f"Serialized input request: {operation_input_args}")
-
-        return ResourceIterator(
-            client=client,
-            list_method="list_endpoints",
-            summaries_key="Endpoints",
-            summary_name="EndpointSummary",
-            resource_cls=Endpoint,
-            list_method_kwargs=operation_input_args,
-        )
-
-    @Base.add_validate_call
-    def update_weights_and_capacities(
-        self,
-        desired_weights_and_capacities: List[DesiredWeightAndCapacity],
-        session: Optional[Session] = None,
-        region: Optional[str] = None,
-    ) -> None:
-        """
-        Updates variant weight of one or more variants associated with an existing endpoint, or capacity of one variant associated with an existing endpoint.
-
-        Parameters:
-            desired_weights_and_capacities: An object that provides new capacity and weight values for a variant.
-            session: Boto3 session.
-            region: Region name.
-
-        Raises:
-            botocore.exceptions.ClientError: This exception is raised for AWS service related errors.
-                The error message and error code can be parsed from the exception as follows:
-                ```
-                try:
-                    # AWS service call here
-                except botocore.exceptions.ClientError as e:
-                    error_message = e.response['Error']['Message']
-                    error_code = e.response['Error']['Code']
-                ```
-            ResourceLimitExceeded: You have exceeded an SageMaker resource limit. For example, you might have too many training jobs created.
-        """
-
-        operation_input_args = {
-            "EndpointName": self.endpoint_name,
-            "DesiredWeightsAndCapacities": desired_weights_and_capacities,
+            "SessionId": session_id,
         }
         # serialize the input request
         operation_input_args = serialize(operation_input_args)
         logger.debug(f"Serialized input request: {operation_input_args}")
 
         client = Base.get_sagemaker_client(
-            session=session, region_name=region, service_name="sagemaker"
+            session=session, region_name=region, service_name="sagemaker-runtime"
         )
 
-        logger.debug(f"Calling update_endpoint_weights_and_capacities API")
-        response = client.update_endpoint_weights_and_capacities(**operation_input_args)
+        logger.debug(f"Calling invoke_endpoint_with_response_stream API")
+        response = client.invoke_endpoint_with_response_stream(**operation_input_args)
         logger.debug(f"Response: {response}")
 
+        transformed_response = transform(response, "InvokeEndpointWithResponseStreamOutput")
+        return InvokeEndpointWithResponseStreamOutput(**transformed_response)
+
 
 class EndpointConfig(Base):
     """
@@ -9665,6 +9697,230 @@ class FeatureGroup(Base):
             list_method_kwargs=operation_input_args,
         )
 
+    @Base.add_validate_call
+    def get_record(
+        self,
+        record_identifier_value_as_string: str,
+        feature_names: Optional[List[str]] = Unassigned(),
+        expiration_time_response: Optional[str] = Unassigned(),
+        session: Optional[Session] = None,
+        region: Optional[str] = None,
+    ) -> Optional[GetRecordResponse]:
+        """
+        Use for OnlineStore serving from a FeatureStore.
+
+        Parameters:
+            record_identifier_value_as_string: The value that corresponds to RecordIdentifier type and uniquely identifies the record in the FeatureGroup.
+            feature_names: List of names of Features to be retrieved. If not specified, the latest value for all the Features are returned.
+            expiration_time_response: Parameter to request ExpiresAt in response. If Enabled, GetRecord will return the value of ExpiresAt, if it is not null. If Disabled and null, GetRecord will return null.
+            session: Boto3 session.
+            region: Region name.
+
+        Returns:
+            GetRecordResponse
+
+        Raises:
+            botocore.exceptions.ClientError: This exception is raised for AWS service related errors.
+                The error message and error code can be parsed from the exception as follows:
+                ```
+                try:
+                    # AWS service call here
+                except botocore.exceptions.ClientError as e:
+                    error_message = e.response['Error']['Message']
+                    error_code = e.response['Error']['Code']
+                ```
+            AccessForbidden: You do not have permission to perform an action.
+            InternalFailure: An internal failure occurred. Try your request again. If the problem persists, contact Amazon Web Services customer support.
+            ResourceNotFound: Resource being access is not found.
+            ServiceUnavailable: The service is currently unavailable.
+            ValidationError: There was an error validating your request.
+        """
+
+        operation_input_args = {
+            "FeatureGroupName": self.feature_group_name,
+            "RecordIdentifierValueAsString": record_identifier_value_as_string,
+            "FeatureNames": feature_names,
+            "ExpirationTimeResponse": expiration_time_response,
+        }
+        # serialize the input request
+        operation_input_args = serialize(operation_input_args)
+        logger.debug(f"Serialized input request: {operation_input_args}")
+
+        client = Base.get_sagemaker_client(
+            session=session, region_name=region, service_name="sagemaker-featurestore-runtime"
+        )
+
+        logger.debug(f"Calling get_record API")
+        response = client.get_record(**operation_input_args)
+        logger.debug(f"Response: {response}")
+
+        transformed_response = transform(response, "GetRecordResponse")
+        return GetRecordResponse(**transformed_response)
+
+    @Base.add_validate_call
+    def put_record(
+        self,
+        record: List[FeatureValue],
+        target_stores: Optional[List[str]] = Unassigned(),
+        ttl_duration: Optional[TtlDuration] = Unassigned(),
+        session: Optional[Session] = None,
+        region: Optional[str] = None,
+    ) -> None:
+        """
+        The PutRecord API is used to ingest a list of Records into your feature group.
+
+        Parameters:
+            record: List of FeatureValues to be inserted. This will be a full over-write. If you only want to update few of the feature values, do the following:   Use GetRecord to retrieve the latest record.   Update the record returned from GetRecord.    Use PutRecord to update feature values.
+            target_stores: A list of stores to which you're adding the record. By default, Feature Store adds the record to all of the stores that you're using for the FeatureGroup.
+            ttl_duration: Time to live duration, where the record is hard deleted after the expiration time is reached; ExpiresAt = EventTime + TtlDuration. For information on HardDelete, see the DeleteRecord API in the Amazon SageMaker API Reference guide.
+            session: Boto3 session.
+            region: Region name.
+
+        Raises:
+            botocore.exceptions.ClientError: This exception is raised for AWS service related errors.
+                The error message and error code can be parsed from the exception as follows:
+                ```
+                try:
+                    # AWS service call here
+                except botocore.exceptions.ClientError as e:
+                    error_message = e.response['Error']['Message']
+                    error_code = e.response['Error']['Code']
+                ```
+            AccessForbidden: You do not have permission to perform an action.
+            InternalFailure: An internal failure occurred. Try your request again. If the problem persists, contact Amazon Web Services customer support.
+            ServiceUnavailable: The service is currently unavailable.
+            ValidationError: There was an error validating your request.
+        """
+
+        operation_input_args = {
+            "FeatureGroupName": self.feature_group_name,
+            "Record": record,
+            "TargetStores": target_stores,
+            "TtlDuration": ttl_duration,
+        }
+        # serialize the input request
+        operation_input_args = serialize(operation_input_args)
+        logger.debug(f"Serialized input request: {operation_input_args}")
+
+        client = Base.get_sagemaker_client(
+            session=session, region_name=region, service_name="sagemaker-featurestore-runtime"
+        )
+
+        logger.debug(f"Calling put_record API")
+        response = client.put_record(**operation_input_args)
+        logger.debug(f"Response: {response}")
+
+    @Base.add_validate_call
+    def delete_record(
+        self,
+        record_identifier_value_as_string: str,
+        event_time: str,
+        target_stores: Optional[List[str]] = Unassigned(),
+        deletion_mode: Optional[str] = Unassigned(),
+        session: Optional[Session] = None,
+        region: Optional[str] = None,
+    ) -> None:
+        """
+        Deletes a Record from a FeatureGroup in the OnlineStore.
+
+        Parameters:
+            record_identifier_value_as_string: The value for the RecordIdentifier that uniquely identifies the record, in string format.
+            event_time: Timestamp indicating when the deletion event occurred. EventTime can be used to query data at a certain point in time.
+            target_stores: A list of stores from which you're deleting the record. By default, Feature Store deletes the record from all of the stores that you're using for the FeatureGroup.
+            deletion_mode: The name of the deletion mode for deleting the record. By default, the deletion mode is set to SoftDelete.
+            session: Boto3 session.
+            region: Region name.
+
+        Raises:
+            botocore.exceptions.ClientError: This exception is raised for AWS service related errors.
+                The error message and error code can be parsed from the exception as follows:
+                ```
+                try:
+                    # AWS service call here
+                except botocore.exceptions.ClientError as e:
+                    error_message = e.response['Error']['Message']
+                    error_code = e.response['Error']['Code']
+                ```
+            AccessForbidden: You do not have permission to perform an action.
+            InternalFailure: An internal failure occurred. Try your request again. If the problem persists, contact Amazon Web Services customer support.
+            ServiceUnavailable: The service is currently unavailable.
+            ValidationError: There was an error validating your request.
+        """
+
+        operation_input_args = {
+            "FeatureGroupName": self.feature_group_name,
+            "RecordIdentifierValueAsString": record_identifier_value_as_string,
+            "EventTime": event_time,
+            "TargetStores": target_stores,
+            "DeletionMode": deletion_mode,
+        }
+        # serialize the input request
+        operation_input_args = serialize(operation_input_args)
+        logger.debug(f"Serialized input request: {operation_input_args}")
+
+        client = Base.get_sagemaker_client(
+            session=session, region_name=region, service_name="sagemaker-featurestore-runtime"
+        )
+
+        logger.debug(f"Calling delete_record API")
+        response = client.delete_record(**operation_input_args)
+        logger.debug(f"Response: {response}")
+
+    @Base.add_validate_call
+    def batch_get_record(
+        self,
+        identifiers: List[BatchGetRecordIdentifier],
+        expiration_time_response: Optional[str] = Unassigned(),
+        session: Optional[Session] = None,
+        region: Optional[str] = None,
+    ) -> Optional[BatchGetRecordResponse]:
+        """
+        Retrieves a batch of Records from a FeatureGroup.
+
+        Parameters:
+            identifiers: A list containing the name or Amazon Resource Name (ARN) of the FeatureGroup, the list of names of Features to be retrieved, and the corresponding RecordIdentifier values as strings.
+            expiration_time_response: Parameter to request ExpiresAt in response. If Enabled, BatchGetRecord will return the value of ExpiresAt, if it is not null. If Disabled and null, BatchGetRecord will return null.
+            session: Boto3 session.
+            region: Region name.
+
+        Returns:
+            BatchGetRecordResponse
+
+        Raises:
+            botocore.exceptions.ClientError: This exception is raised for AWS service related errors.
+                The error message and error code can be parsed from the exception as follows:
+                ```
+                try:
+                    # AWS service call here
+                except botocore.exceptions.ClientError as e:
+                    error_message = e.response['Error']['Message']
+                    error_code = e.response['Error']['Code']
+                ```
+            AccessForbidden: You do not have permission to perform an action.
+            InternalFailure: An internal failure occurred. Try your request again. If the problem persists, contact Amazon Web Services customer support.
+            ServiceUnavailable: The service is currently unavailable.
+            ValidationError: There was an error validating your request.
+        """
+
+        operation_input_args = {
+            "Identifiers": identifiers,
+            "ExpirationTimeResponse": expiration_time_response,
+        }
+        # serialize the input request
+        operation_input_args = serialize(operation_input_args)
+        logger.debug(f"Serialized input request: {operation_input_args}")
+
+        client = Base.get_sagemaker_client(
+            session=session, region_name=region, service_name="sagemaker-featurestore-runtime"
+        )
+
+        logger.debug(f"Calling batch_get_record API")
+        response = client.batch_get_record(**operation_input_args)
+        logger.debug(f"Response: {response}")
+
+        transformed_response = transform(response, "BatchGetRecordResponse")
+        return BatchGetRecordResponse(**transformed_response)
+
 
 class FeatureMetadata(Base):
     """
@@ -27654,6 +27910,49 @@ class TrialComponent(Base):
         response = client.disassociate_trial_component(**operation_input_args)
         logger.debug(f"Response: {response}")
 
+    @Base.add_validate_call
+    def batch_put_metrics(
+        self,
+        metric_data: List[RawMetricData],
+        session: Optional[Session] = None,
+        region: Optional[str] = None,
+    ) -> None:
+        """
+        Used to ingest training metrics into SageMaker.
+
+        Parameters:
+            metric_data: A list of raw metric values to put.
+            session: Boto3 session.
+            region: Region name.
+
+        Raises:
+            botocore.exceptions.ClientError: This exception is raised for AWS service related errors.
+                The error message and error code can be parsed from the exception as follows:
+                ```
+                try:
+                    # AWS service call here
+                except botocore.exceptions.ClientError as e:
+                    error_message = e.response['Error']['Message']
+                    error_code = e.response['Error']['Code']
+                ```
+        """
+
+        operation_input_args = {
+            "TrialComponentName": self.trial_component_name,
+            "MetricData": metric_data,
+        }
+        # serialize the input request
+        operation_input_args = serialize(operation_input_args)
+        logger.debug(f"Serialized input request: {operation_input_args}")
+
+        client = Base.get_sagemaker_client(
+            session=session, region_name=region, service_name="sagemaker-metrics"
+        )
+
+        logger.debug(f"Calling batch_put_metrics API")
+        response = client.batch_put_metrics(**operation_input_args)
+        logger.debug(f"Response: {response}")
+
 
 class UserProfile(Base):
     """