Merge branch 'master' into feat/jumpstart-private-model-artifacts

Author: evakravi

CHANGELOG.md

@@ -1,5 +1,18 @@
 # Changelog
 
+## v2.194.0 (2023-10-19)
+
+### Features
+
+ * Added register step in Jumpstart model
+ * jumpstart instance specific metric definitions
+
+### Bug Fixes and Other Changes
+
+ * Updates for DJL 0.24.0 Release
+ * use getter for resource-metadata dict
+ * add method to Model class to check if repack is needed
+
 ## v2.193.0 (2023-10-18)
 
 ### Features

VERSION

@@ -1 +1 @@
-2.193.1.dev0
+2.194.1.dev0

src/sagemaker/feature_store/feature_processor/feature_scheduler.py

@@ -281,12 +281,13 @@ def schedule(
     Args:
         pipeline_name (str): The SageMaker Pipeline name that will be scheduled.
         schedule_expression (str): The expression that defines when the schedule runs. It supports
-            at expression, rate expression and cron expression. See '''https://docs.aws.amazon.com\
-                /scheduler/latest/APIReference/API_CreateSchedule.html#scheduler-CreateSchedule-\
-                request-ScheduleExpression''' for more details.
+            at expression, rate expression and cron expression. See the
+            `CreateSchedule API
+            <https://docs.aws.amazon.com/scheduler/latest/APIReference/API_CreateSchedule.html#scheduler-CreateSchedulerequest-ScheduleExpression>`_
+            for more details.
         state (str): Specifies whether the schedule is enabled or disabled. Valid values are
-            ENABLED and DISABLED. See '''https://docs.aws.amazon.com/scheduler/latest/APIReference\
-                /API_CreateSchedule.html#scheduler-CreateSchedule-request-State'''
+            ENABLED and DISABLED. See the `State request parameter
+            <https://docs.aws.amazon.com/scheduler/latest/APIReference/API_CreateSchedule.html#scheduler-CreateSchedule-request-State>`_
             for more details. If not specified, it will default to ENABLED.
         start_date (Optional[datetime]): The date, in UTC, after which the schedule can begin
             invoking its target. Depending on the schedule’s recurrence expression, invocations

src/sagemaker/jumpstart/payload_utils.py

@@ -14,20 +14,181 @@
 from __future__ import absolute_import
 import base64
 import json
-from typing import Optional, Union
+from typing import Any, Dict, List, Optional, Union
 import re
 import boto3
 
 from sagemaker.jumpstart.accessors import JumpStartS3PayloadAccessor
-from sagemaker.jumpstart.constants import JUMPSTART_DEFAULT_REGION_NAME
+from sagemaker.jumpstart.artifacts.payloads import _retrieve_example_payloads
+from sagemaker.jumpstart.constants import (
+    DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
+    JUMPSTART_DEFAULT_REGION_NAME,
+)
 from sagemaker.jumpstart.enums import MIMEType
 from sagemaker.jumpstart.types import JumpStartSerializablePayload
-from sagemaker.jumpstart.utils import get_jumpstart_content_bucket
+from sagemaker.jumpstart.utils import (
+    get_jumpstart_content_bucket,
+)
+from sagemaker.session import Session
+
 
 S3_BYTES_REGEX = r"^\$s3<(?P<s3_key>[a-zA-Z0-9-_/.]+)>$"
 S3_B64_STR_REGEX = r"\$s3_b64<(?P<s3_key>[a-zA-Z0-9-_/.]+)>"
 
 
+def _extract_field_from_json(
+    json_input: dict,
+    keys: List[str],
+) -> Any:
+    """Given a dictionary, returns value at specified keys.
+
+    Raises:
+        KeyError: If a key cannot be found in the json input.
+    """
+    curr_json = json_input
+    for idx, key in enumerate(keys):
+        if idx < len(keys) - 1:
+            curr_json = curr_json[key]
+            continue
+        return curr_json[key]
+
+
+def _construct_payload(
+    prompt: str,
+    model_id: str,
+    model_version: str,
+    region: Optional[str] = None,
+    tolerate_vulnerable_model: bool = False,
+    tolerate_deprecated_model: bool = False,
+    sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
+) -> Optional[JumpStartSerializablePayload]:
+    """Returns example payload from prompt.
+
+    Args:
+        prompt (str): String-valued prompt to embed in payload.
+        model_id (str): JumpStart model ID of the JumpStart model for which to construct
+            the payload.
+        model_version (str): Version of the JumpStart model for which to retrieve the
+            payload.
+        region (Optional[str]): Region for which to retrieve the
+            payload. (Default: None).
+        tolerate_vulnerable_model (bool): True if vulnerable versions of model
+            specifications should be tolerated (exception not raised). If False, raises an
+            exception if the script used by this version of the model has dependencies with known
+            security vulnerabilities. (Default: False).
+        tolerate_deprecated_model (bool): True if deprecated versions of model
+            specifications should be tolerated (exception not raised). If False, raises
+            an exception if the version of the model is deprecated. (Default: False).
+        sagemaker_session (sagemaker.session.Session): A SageMaker Session
+            object, used for SageMaker interactions. If not
+            specified, one is created using the default AWS configuration
+            chain. (Default: sagemaker.jumpstart.constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION).
+    Returns:
+        Optional[JumpStartSerializablePayload]: serializable payload with prompt, or None if
+            this feature is unavailable for the specified model.
+    """
+    payloads: Optional[Dict[str, JumpStartSerializablePayload]] = _retrieve_example_payloads(
+        model_id=model_id,
+        model_version=model_version,
+        region=region,
+        tolerate_vulnerable_model=tolerate_vulnerable_model,
+        tolerate_deprecated_model=tolerate_deprecated_model,
+        sagemaker_session=sagemaker_session,
+    )
+    if payloads is None or len(payloads) == 0:
+        return None
+
+    payload_to_use: JumpStartSerializablePayload = list(payloads.values())[0]
+
+    prompt_key: Optional[str] = payload_to_use.prompt_key
+    if prompt_key is None:
+        return None
+
+    payload_body = payload_to_use.body
+    prompt_key_split = prompt_key.split(".")
+    for idx, prompt_key in enumerate(prompt_key_split):
+        if idx < len(prompt_key_split) - 1:
+            payload_body = payload_body[prompt_key]
+        else:
+            payload_body[prompt_key] = prompt
+
+    return payload_to_use
+
+
+def _extract_generated_text_from_response(
+    response: dict,
+    model_id: str,
+    model_version: str,
+    region: Optional[str] = None,
+    tolerate_vulnerable_model: bool = False,
+    tolerate_deprecated_model: bool = False,
+    sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
+    accept_type: Optional[str] = None,
+) -> str:
+    """Returns generated text extracted from full response payload.
+
+    Args:
+        response (dict): Dictionary-valued response from which to extract
+            generated text.
+        model_id (str): JumpStart model ID of the JumpStart model from which to extract
+            generated text.
+        model_version (str): Version of the JumpStart model for which to extract generated
+            text.
+        region (Optional[str]): Region for which to extract generated
+            text. (Default: None).
+        tolerate_vulnerable_model (bool): True if vulnerable versions of model
+            specifications should be tolerated (exception not raised). If False, raises an
+            exception if the script used by this version of the model has dependencies with known
+            security vulnerabilities. (Default: False).
+        tolerate_deprecated_model (bool): True if deprecated versions of model
+            specifications should be tolerated (exception not raised). If False, raises
+            an exception if the version of the model is deprecated. (Default: False).
+        sagemaker_session (sagemaker.session.Session): A SageMaker Session
+            object, used for SageMaker interactions. If not
+            specified, one is created using the default AWS configuration
+            chain. (Default: sagemaker.jumpstart.constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION).
+        accept_type (Optional[str]): The accept type to optionally specify for the response.
+            (Default: None).
+
+    Returns:
+        str: extracted generated text from the endpoint response payload.
+
+    Raises:
+        ValueError: If the model is invalid, the model does not support generated text extraction,
+        or if the response is malformed.
+    """
+
+    if not isinstance(response, dict):
+        raise ValueError(f"Response must be dictionary. Instead, got: {type(response)}")
+
+    payloads: Optional[Dict[str, JumpStartSerializablePayload]] = _retrieve_example_payloads(
+        model_id=model_id,
+        model_version=model_version,
+        region=region,
+        tolerate_vulnerable_model=tolerate_vulnerable_model,
+        tolerate_deprecated_model=tolerate_deprecated_model,
+        sagemaker_session=sagemaker_session,
+    )
+    if payloads is None or len(payloads) == 0:
+        raise ValueError(f"Model ID '{model_id}' does not support generated text extraction.")
+
+    for payload in payloads.values():
+        if accept_type is None or payload.accept == accept_type:
+            generated_text_response_key: Optional[str] = payload.generated_text_response_key
+            if generated_text_response_key is None:
+                raise ValueError(
+                    f"Model ID '{model_id}' does not support generated text extraction."
+                )
+
+            generated_text_response_key_split = generated_text_response_key.split(".")
+            try:
+                return _extract_field_from_json(response, generated_text_response_key_split)
+            except KeyError:
+                raise ValueError(f"Response is malformed: {response}")
+
+    raise ValueError(f"Model ID '{model_id}' does not support generated text extraction.")
+
+
 class PayloadSerializer:
     """Utility class for serializing payloads associated with JumpStart models.
 

src/sagemaker/jumpstart/types.py

@@ -339,9 +339,11 @@ class JumpStartSerializablePayload(JumpStartDataHolderType):
         "content_type",
         "accept",
         "body",
+        "generated_text_response_key",
+        "prompt_key",
     ]
 
-    _non_serializable_slots = ["raw_payload"]
+    _non_serializable_slots = ["raw_payload", "prompt_key"]
 
     def __init__(self, spec: Optional[Dict[str, Any]]):
         """Initializes a JumpStartSerializablePayload object from its json representation.
@@ -369,6 +371,8 @@ def from_json(self, json_obj: Optional[Dict[str, Any]]) -> None:
         self.content_type = json_obj["content_type"]
         self.body = json_obj["body"]
         accept = json_obj.get("accept")
+        self.generated_text_response_key = json_obj.get("generated_text_response_key")
+        self.prompt_key = json_obj.get("prompt_key")
         if accept:
             self.accept = accept