Merge remote-tracking branch 'origin' into fix/raise-exception-no-gated-env-var

Author: evakravi

.github/workflows/codebuild-ci.yml

@@ -0,0 +1,48 @@
+name: PR Checks
+on:
+    pull_request_target:
+
+concurrency:
+    group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.head_ref }}
+    cancel-in-progress: true
+
+permissions:
+    id-token: write # This is required for requesting the JWT
+
+jobs:
+  codestyle-doc-tests:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Configure AWS Credentials
+        uses: aws-actions/configure-aws-credentials@v4
+        with:
+          role-to-assume: ${{ secrets.CI_AWS_ROLE_ARN }}
+          aws-region: us-west-2
+          role-duration-seconds: 10800
+      - name: Run Codestyle & Doc Tests
+        uses: aws-actions/aws-codebuild-run-build@v1
+        with:
+          project-name: sagemaker-python-sdk-ci-codestyle-doc-tests
+          source-version-override: 'pr/${{ github.event.pull_request.number }}'
+  unit-tests:
+    runs-on: ubuntu-latest
+    strategy:
+        fail-fast: false
+        matrix:
+          python-version: ["py38", "py39", "py310"]
+    steps:
+      - name: Configure AWS Credentials
+        uses: aws-actions/configure-aws-credentials@v4
+        with:
+          role-to-assume: ${{ secrets.CI_AWS_ROLE_ARN }}
+          aws-region: us-west-2
+          role-duration-seconds: 10800
+      - name: Run Unit Tests
+        uses: aws-actions/aws-codebuild-run-build@v1
+        with:
+          project-name: sagemaker-python-sdk-ci-unit-tests
+          source-version-override: 'pr/${{ github.event.pull_request.number }}'
+          env-vars-for-codebuild: |
+            PY_VERSION
+        env:
+          PY_VERSION: ${{ matrix.python-version }}

doc/doc_utils/jumpstart_doc_utils.py

@@ -74,9 +74,12 @@ class Frameworks(str, Enum):
 
 JUMPSTART_REGION = "eu-west-2"
 SDK_MANIFEST_FILE = "models_manifest.json"
+PROPRIETARY_SDK_MANIFEST_FILE = "proprietary-sdk-manifest.json"
 JUMPSTART_BUCKET_BASE_URL = "https://jumpstart-cache-prod-{}.s3.{}.amazonaws.com".format(
     JUMPSTART_REGION, JUMPSTART_REGION
 )
+PROPRIETARY_DOC_BUCKET = "https://jumpstart-cache-prod-us-west-2.s3.us-west-2.amazonaws.com"
+
 TASK_MAP = {
     Tasks.IC: ProblemTypes.IMAGE_CLASSIFICATION,
     Tasks.IC_EMBEDDING: ProblemTypes.IMAGE_EMBEDDING,
@@ -152,18 +155,26 @@ class Frameworks(str, Enum):
 }
 
 
-def get_jumpstart_sdk_manifest():
-    url = "{}/{}".format(JUMPSTART_BUCKET_BASE_URL, SDK_MANIFEST_FILE)
+def get_public_s3_json_object(url):
     with request.urlopen(url) as f:
         models_manifest = f.read().decode("utf-8")
     return json.loads(models_manifest)
 
 
-def get_jumpstart_sdk_spec(key):
-    url = "{}/{}".format(JUMPSTART_BUCKET_BASE_URL, key)
-    with request.urlopen(url) as f:
-        model_spec = f.read().decode("utf-8")
-    return json.loads(model_spec)
+def get_jumpstart_sdk_manifest():
+    return get_public_s3_json_object(f"{JUMPSTART_BUCKET_BASE_URL}/{SDK_MANIFEST_FILE}")
+
+
+def get_proprietary_sdk_manifest():
+    return get_public_s3_json_object(f"{PROPRIETARY_DOC_BUCKET}/{PROPRIETARY_SDK_MANIFEST_FILE}")
+
+
+def get_jumpstart_sdk_spec(s3_key: str):
+    return get_public_s3_json_object(f"{JUMPSTART_BUCKET_BASE_URL}/{s3_key}")
+
+
+def get_proprietary_sdk_spec(s3_key: str):
+    return get_public_s3_json_object(f"{PROPRIETARY_DOC_BUCKET}/{s3_key}")
 
 
 def get_model_task(id):
@@ -196,6 +207,45 @@ def get_model_source(url):
         return "Source"
 
 
+def create_proprietary_model_table():
+    proprietary_content_intro = []
+    proprietary_content_intro.append("\n")
+    proprietary_content_intro.append(".. list-table:: Available Proprietary Models\n")
+    proprietary_content_intro.append("   :widths: 50 20 20 20 20\n")
+    proprietary_content_intro.append("   :header-rows: 1\n")
+    proprietary_content_intro.append("   :class: datatable\n")
+    proprietary_content_intro.append("\n")
+    proprietary_content_intro.append("   * - Model ID\n")
+    proprietary_content_intro.append("     - Fine Tunable?\n")
+    proprietary_content_intro.append("     - Supported Version\n")
+    proprietary_content_intro.append("     - Min SDK Version\n")
+    proprietary_content_intro.append("     - Source\n")
+
+    sdk_manifest = get_proprietary_sdk_manifest()
+    sdk_manifest_top_versions_for_models = {}
+
+    for model in sdk_manifest:
+        if model["model_id"] not in sdk_manifest_top_versions_for_models:
+            sdk_manifest_top_versions_for_models[model["model_id"]] = model
+        else:
+            if str(sdk_manifest_top_versions_for_models[model["model_id"]]["version"]) < str(
+                model["version"]
+            ):
+                sdk_manifest_top_versions_for_models[model["model_id"]] = model
+
+    proprietary_content_entries = []
+    for model in sdk_manifest_top_versions_for_models.values():
+        model_spec = get_proprietary_sdk_spec(model["spec_key"])
+        proprietary_content_entries.append("   * - {}\n".format(model_spec["model_id"]))
+        proprietary_content_entries.append("     - {}\n".format(False))  # TODO: support training
+        proprietary_content_entries.append("     - {}\n".format(model["version"]))
+        proprietary_content_entries.append("     - {}\n".format(model["min_version"]))
+        proprietary_content_entries.append(
+            "     - `{} <{}>`__ |external-link|\n".format("Source", model_spec.get("url"))
+        )
+    return proprietary_content_intro + proprietary_content_entries + ["\n"]
+
+
 def create_jumpstart_model_table():
     sdk_manifest = get_jumpstart_sdk_manifest()
     sdk_manifest_top_versions_for_models = {}
@@ -249,19 +299,19 @@ def create_jumpstart_model_table():
     file_content_intro.append("     - Source\n")
 
     dynamic_table_files = []
-    file_content_entries = []
+    open_weight_content_entries = []
 
     for model in sdk_manifest_top_versions_for_models.values():
         model_spec = get_jumpstart_sdk_spec(model["spec_key"])
         model_task = get_model_task(model_spec["model_id"])
         string_model_task = get_string_model_task(model_spec["model_id"])
         model_source = get_model_source(model_spec["url"])
-        file_content_entries.append("   * - {}\n".format(model_spec["model_id"]))
-        file_content_entries.append("     - {}\n".format(model_spec["training_supported"]))
-        file_content_entries.append("     - {}\n".format(model["version"]))
-        file_content_entries.append("     - {}\n".format(model["min_version"]))
-        file_content_entries.append("     - {}\n".format(model_task))
-        file_content_entries.append(
+        open_weight_content_entries.append("   * - {}\n".format(model_spec["model_id"]))
+        open_weight_content_entries.append("     - {}\n".format(model_spec["training_supported"]))
+        open_weight_content_entries.append("     - {}\n".format(model["version"]))
+        open_weight_content_entries.append("     - {}\n".format(model["min_version"]))
+        open_weight_content_entries.append("     - {}\n".format(model_task))
+        open_weight_content_entries.append(
             "     - `{} <{}>`__ |external-link|\n".format(model_source, model_spec["url"])
         )
 
@@ -299,7 +349,10 @@ def create_jumpstart_model_table():
             f.writelines(file_content_single_entry)
             f.close()
 
+    proprietary_content_entries = create_proprietary_model_table()
+
     f = open("doc_utils/pretrainedmodels.rst", "a")
     f.writelines(file_content_intro)
-    f.writelines(file_content_entries)
+    f.writelines(open_weight_content_entries)
+    f.writelines(proprietary_content_entries)
     f.close()

src/sagemaker/accept_types.py

@@ -16,6 +16,7 @@
 
 from sagemaker.jumpstart import artifacts, utils as jumpstart_utils
 from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
+from sagemaker.jumpstart.enums import JumpStartModelType
 from sagemaker.session import Session
 
 
@@ -75,6 +76,7 @@ def retrieve_default(
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
+    model_type: JumpStartModelType = JumpStartModelType.OPEN_WEIGHTS,
 ) -> str:
     """Retrieves the default accept type for the model matching the given arguments.
 
@@ -114,4 +116,5 @@ def retrieve_default(
         tolerate_vulnerable_model,
         tolerate_deprecated_model,
         sagemaker_session=sagemaker_session,
+        model_type=model_type,
     )

src/sagemaker/base_predictor.py

@@ -58,7 +58,9 @@
 from sagemaker.model_monitor.model_monitoring import DEFAULT_REPOSITORY_NAME
 
 from sagemaker.lineage.context import EndpointContext
-from sagemaker.compute_resource_requirements.resource_requirements import ResourceRequirements
+from sagemaker.compute_resource_requirements.resource_requirements import (
+    ResourceRequirements,
+)
 
 LOGGER = logging.getLogger("sagemaker")
 

src/sagemaker/chainer/model.py

@@ -174,6 +174,7 @@ def register(
         nearest_model_name: Optional[Union[str, PipelineVariable]] = None,
         data_input_configuration: Optional[Union[str, PipelineVariable]] = None,
         skip_model_validation: Optional[Union[str, PipelineVariable]] = None,
+        source_uri: Optional[Union[str, PipelineVariable]] = None,
     ):
         """Creates a model package for creating SageMaker models or listing on Marketplace.
 
@@ -223,6 +224,8 @@ def register(
                 (default: None).
             skip_model_validation (str or PipelineVariable): Indicates if you want to skip model
                 validation. Values can be "All" or "None" (default: None).
+            source_uri (str or PipelineVariable): The URI of the source for the model package
+                (default: None).
 
         Returns:
             str: A string of SageMaker Model Package ARN.
@@ -262,6 +265,7 @@ def register(
             nearest_model_name=nearest_model_name,
             data_input_configuration=data_input_configuration,
             skip_model_validation=skip_model_validation,
+            source_uri=source_uri,
         )
 
     def prepare_container_def(

src/sagemaker/content_types.py

@@ -16,6 +16,7 @@
 
 from sagemaker.jumpstart import artifacts, utils as jumpstart_utils
 from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
+from sagemaker.jumpstart.enums import JumpStartModelType
 from sagemaker.session import Session
 
 
@@ -75,6 +76,7 @@ def retrieve_default(
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
+    model_type: JumpStartModelType = JumpStartModelType.OPEN_WEIGHTS,
 ) -> str:
     """Retrieves the default content type for the model matching the given arguments.
 
@@ -114,6 +116,7 @@ def retrieve_default(
         tolerate_vulnerable_model,
         tolerate_deprecated_model,
         sagemaker_session=sagemaker_session,
+        model_type=model_type,
     )
 
 

src/sagemaker/deserializers.py

@@ -35,6 +35,7 @@
 
 from sagemaker.jumpstart import artifacts, utils as jumpstart_utils
 from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
+from sagemaker.jumpstart.enums import JumpStartModelType
 from sagemaker.session import Session
 
 
@@ -95,6 +96,7 @@ def retrieve_default(
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
+    model_type: JumpStartModelType = JumpStartModelType.OPEN_WEIGHTS,
 ) -> BaseDeserializer:
     """Retrieves the default deserializer for the model matching the given arguments.
 
@@ -135,4 +137,5 @@ def retrieve_default(
         tolerate_vulnerable_model,
         tolerate_deprecated_model,
         sagemaker_session=sagemaker_session,
+        model_type=model_type,
     )

src/sagemaker/estimator.py

@@ -1718,6 +1718,7 @@ def register(
         nearest_model_name=None,
         data_input_configuration=None,
         skip_model_validation=None,
+        source_uri=None,
         **kwargs,
     ):
         """Creates a model package for creating SageMaker models or listing on Marketplace.
@@ -1765,6 +1766,7 @@ def register(
             data_input_configuration (str): Input object for the model (default: None).
             skip_model_validation (str): Indicates if you want to skip model validation.
                 Values can be "All" or "None" (default: None).
+            source_uri (str): The URI of the source for the model package (default: None).
             **kwargs: Passed to invocation of ``create_model()``. Implementations may customize
                 ``create_model()`` to accept ``**kwargs`` to customize model creation during
                 deploy. For more, see the implementation docs.
@@ -1809,6 +1811,7 @@ def register(
             nearest_model_name=nearest_model_name,
             data_input_configuration=data_input_configuration,
             skip_model_validation=skip_model_validation,
+            source_uri=source_uri,
         )
 
     @property

src/sagemaker/huggingface/model.py

@@ -360,6 +360,7 @@ def register(
         nearest_model_name: Optional[Union[str, PipelineVariable]] = None,
         data_input_configuration: Optional[Union[str, PipelineVariable]] = None,
         skip_model_validation: Optional[Union[str, PipelineVariable]] = None,
+        source_uri: Optional[Union[str, PipelineVariable]] = None,
     ):
         """Creates a model package for creating SageMaker models or listing on Marketplace.
 
@@ -410,6 +411,8 @@ def register(
                 (default: None).
             skip_model_validation (str or PipelineVariable): Indicates if you want to skip model
                 validation. Values can be "All" or "None" (default: None).
+            source_uri (str or PipelineVariable): The URI of the source for the model package
+                (default: None).
 
         Returns:
             A `sagemaker.model.ModelPackage` instance.
@@ -457,6 +460,7 @@ def register(
             nearest_model_name=nearest_model_name,
             data_input_configuration=data_input_configuration,
             skip_model_validation=skip_model_validation,
+            source_uri=source_uri,
         )
 
     def prepare_container_def(

src/sagemaker/instance_types.py

@@ -20,6 +20,7 @@
 from sagemaker.jumpstart import utils as jumpstart_utils
 from sagemaker.jumpstart import artifacts
 from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
+from sagemaker.jumpstart.enums import JumpStartModelType
 from sagemaker.session import Session
 
 logger = logging.getLogger(__name__)
@@ -34,6 +35,7 @@ def retrieve_default(
     tolerate_deprecated_model: bool = False,
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
     training_instance_type: Optional[str] = None,
+    model_type: JumpStartModelType = JumpStartModelType.OPEN_WEIGHTS,
 ) -> str:
     """Retrieves the default instance type for the model matching the given arguments.
 
@@ -85,6 +87,7 @@ def retrieve_default(
         tolerate_deprecated_model,
         sagemaker_session=sagemaker_session,
         training_instance_type=training_instance_type,
+        model_type=model_type,
     )