Merge branch 'master' into master

Author: malav-shastri

.github/workflows/codeql.yml

@@ -0,0 +1,35 @@
+name: "CodeQL"
+on:
+  push:
+    branches: [ "master" ]
+  pull_request:
+    branches: [ "master" ]
+  schedule:
+    - cron: '30 8 * * *'
+jobs:
+  analyze:
+    name: Analyze (${{ matrix.language }})
+    runs-on: ${{ 'ubuntu-latest' }}
+    permissions:
+      security-events: write
+      packages: read
+
+    strategy:
+      matrix:
+        include:
+          - language: python
+            build-mode: none
+          - language: java-kotlin
+            build-mode: none
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@6ccd57f4c5d15bdc2fef309bd9fb6cc9db2ef1c6
+      - name: Initialize CodeQL
+        uses: github/codeql-action/init@4b1d7da102ff94aca014c0245062b1a463356d72
+        with:
+          languages: ${{ matrix.language }}
+          build-mode: ${{ matrix.build-mode }}
+      - name: Perform CodeQL Analysis
+        uses: github/codeql-action/analyze@4b1d7da102ff94aca014c0245062b1a463356d72
+        with:
+          category: "/language:${{matrix.language}}"

src/sagemaker/estimator.py

@@ -68,6 +68,7 @@
 from sagemaker.interactive_apps import SupportedInteractiveAppTypes
 from sagemaker.interactive_apps.tensorboard import TensorBoardApp
 from sagemaker.instance_group import InstanceGroup
+from sagemaker.model_card.model_card import ModelCard, TrainingDetails
 from sagemaker.utils import instance_supports_kms
 from sagemaker.job import _Job
 from sagemaker.jumpstart.utils import (
@@ -1797,8 +1798,17 @@ def register(
         else:
             if "model_kms_key" not in kwargs:
                 kwargs["model_kms_key"] = self.output_kms_key
-            model = self.create_model(image_uri=image_uri, **kwargs)
+            model = self.create_model(image_uri=image_uri, name=model_name, **kwargs)
         model.name = model_name
+        if self.model_data is not None and model_card is None:
+            training_details = TrainingDetails.from_model_s3_artifacts(
+                model_artifacts=[self.model_data], sagemaker_session=self.sagemaker_session
+            )
+            model_card = ModelCard(
+                name="estimator_card",
+                training_details=training_details,
+                sagemaker_session=self.sagemaker_session,
+            )
         return model.register(
             content_types,
             response_types,

src/sagemaker/fw_utils.py

@@ -152,7 +152,6 @@
     "2.1.0",
     "2.1.2",
     "2.2.0",
-    "2.3.0",
     "2.3.1",
 ]
 

src/sagemaker/image_uri_config/huggingface-llm.json

@@ -717,7 +717,7 @@
                 "tag_prefix": "2.3.0-tgi2.2.0",
                 "repository": "huggingface-pytorch-tgi-inference",
                 "container_version": {
-                    "gpu": "cu121-ubuntu22.04"
+                    "gpu": "cu121-ubuntu22.04-v2.0"
                 }
             }
         }

src/sagemaker/image_uri_config/pytorch-smp.json

@@ -8,7 +8,7 @@
             "2.1": "2.1.2",
             "2.2": "2.3.1",
             "2.2.0": "2.3.1",
-            "2.3": "2.4.0"
+            "2.3.1": "2.4.0"
         },
         "versions": {
             "2.0.1": {

src/sagemaker/image_uri_config/pytorch.json

@@ -1008,6 +1008,7 @@
                     "us-gov-west-1": "442386744353",
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
+                    "us-isof-east-1": "303241398832",
                     "us-isof-south-1": "454834333376",
                     "us-west-1": "763104351884",
                     "us-west-2": "763104351884"
@@ -1052,6 +1053,7 @@
                     "us-gov-west-1": "442386744353",
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
+                    "us-isof-east-1": "303241398832",
                     "us-isof-south-1": "454834333376",
                     "us-west-1": "763104351884",
                     "us-west-2": "763104351884"
@@ -2331,6 +2333,7 @@
                     "us-gov-west-1": "442386744353",
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
+                    "us-isof-east-1": "303241398832",
                     "us-isof-south-1": "454834333376",
                     "us-west-1": "763104351884",
                     "us-west-2": "763104351884"
@@ -2419,6 +2422,7 @@
                     "us-gov-west-1": "442386744353",
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
+                    "us-isof-east-1": "303241398832",
                     "us-isof-south-1": "454834333376",
                     "us-west-1": "763104351884",
                     "us-west-2": "763104351884"

src/sagemaker/image_uri_config/tensorflow.json

@@ -2140,6 +2140,7 @@
                     "us-gov-west-1": "442386744353",
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
+                    "us-isof-east-1": "303241398832",
                     "us-isof-south-1": "454834333376",
                     "us-west-1": "763104351884",
                     "us-west-2": "763104351884"
@@ -2181,6 +2182,7 @@
                     "us-gov-west-1": "442386744353",
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
+                    "us-isof-east-1": "303241398832",
                     "us-isof-south-1": "454834333376",
                     "us-west-1": "763104351884",
                     "us-west-2": "763104351884"
@@ -4354,6 +4356,7 @@
                     "us-gov-west-1": "442386744353",
                     "us-iso-east-1": "886529160074",
                     "us-isob-east-1": "094389454867",
+                    "us-isof-east-1": "303241398832",
                     "us-isof-south-1": "454834333376",
                     "us-west-1": "763104351884",
                     "us-west-2": "763104351884"

src/sagemaker/model.py

@@ -549,7 +549,10 @@ def register(
             model_package_group_name = utils.base_name_from_image(
                 self.image_uri, default_base_name=ModelPackage.__name__
             )
-        if model_package_group_name is not None:
+        if (
+            model_package_group_name is not None
+            and model_type is not JumpStartModelType.PROPRIETARY
+        ):
             container_def = self.prepare_container_def(accept_eula=accept_eula)
             container_def = update_container_with_inference_params(
                 framework=framework,
@@ -2466,32 +2469,55 @@ def update_model_card(self, model_card: Union[ModelCard, ModelPackageModelCard])
         desc_model_package = sagemaker_session.sagemaker_client.describe_model_package(
             ModelPackageName=self.model_package_arn
         )
+        if hasattr(model_card, "model_package_details"):
+            model_card.model_package_details = None
         update_model_card_req = model_card._create_request_args()
-        if update_model_card_req["ModelCardStatus"] is not None:
-            if (
-                desc_model_package["ModelCard"]["ModelCardStatus"]
-                == update_model_card_req["ModelCardStatus"]
-            ):
-                del update_model_card_req["ModelCardStatus"]
-
         if update_model_card_req.get("ModelCardName") is not None:
             del update_model_card_req["ModelCardName"]
-        if update_model_card_req.get("Content") is not None:
-            previous_content_hash = _hash_content_str(
-                desc_model_package["ModelCard"]["ModelCardContent"]
-            )
-            current_content_hash = _hash_content_str(update_model_card_req["Content"])
-            if (
-                previous_content_hash == current_content_hash
-                or update_model_card_req.get("Content") == "{}"
-                or update_model_card_req.get("Content") == "null"
-            ):
-                del update_model_card_req["Content"]
-            else:
-                update_model_card_req["ModelCardContent"] = update_model_card_req["Content"]
-                del update_model_card_req["Content"]
-        update_model_package_args = {
-            "ModelPackageArn": self.model_package_arn,
-            "ModelCard": update_model_card_req,
-        }
-        sagemaker_session.sagemaker_client.update_model_package(**update_model_package_args)
+        if update_model_card_req["Content"] is not None:
+            if "model_package_details" in update_model_card_req["Content"]:
+                update_model_card_req["Content"].pop("model_package_details", None)
+            update_model_card_req["ModelCardContent"] = update_model_card_req["Content"]
+            del update_model_card_req["Content"]
+
+        if "ModelCard" in desc_model_package:
+            if update_model_card_req["ModelCardStatus"] is not None:
+                if (
+                    desc_model_package["ModelCard"]["ModelCardStatus"]
+                    != update_model_card_req["ModelCardStatus"]
+                ):
+                    new_mc_mp_req = update_model_card_req
+                    del new_mc_mp_req["ModelCardContent"]
+                    update_model_package_args = {
+                        "ModelPackageArn": self.model_package_arn,
+                        "ModelCard": new_mc_mp_req,
+                    }
+                    sagemaker_session.sagemaker_client.update_model_package(
+                        **update_model_package_args
+                    )
+
+            if update_model_card_req.get("ModelCardContent") is not None:
+                previous_content_hash = _hash_content_str(
+                    desc_model_package["ModelCard"]["ModelCardContent"]
+                )
+                current_content_hash = _hash_content_str(update_model_card_req["ModelCardContent"])
+                if not (
+                    previous_content_hash == current_content_hash
+                    or update_model_card_req.get("ModelCardContent") == "{}"
+                    or update_model_card_req.get("ModelCardContent") == "null"
+                ):
+                    new_mc_mp_req = update_model_card_req
+                    del new_mc_mp_req["ModelCardStatus"]
+                    update_model_package_args = {
+                        "ModelPackageArn": self.model_package_arn,
+                        "ModelCard": new_mc_mp_req,
+                    }
+                    sagemaker_session.sagemaker_client.update_model_package(
+                        **update_model_package_args
+                    )
+        else:
+            update_model_package_args = {
+                "ModelPackageArn": self.model_package_arn,
+                "ModelCard": update_model_card_req,
+            }
+            sagemaker_session.sagemaker_client.update_model_package(**update_model_package_args)

src/sagemaker/session.py

@@ -8348,7 +8348,9 @@ def _logs_for_job(  # noqa: C901 - suppress complexity warning for this method
     """
     sagemaker_client = sagemaker_session.sagemaker_client
     request_end_time = time.time() + timeout if timeout else None
-    description = sagemaker_client.describe_training_job(TrainingJobName=job_name)
+    description = _wait_until(
+        lambda: sagemaker_client.describe_training_job(TrainingJobName=job_name)
+    )
     print(secondary_training_status_message(description, None), end="")
 
     instance_count, stream_names, positions, client, log_group, dot, color_wrap = _logs_init(

tests/integ/test_byo_estimator.py

@@ -12,14 +12,20 @@
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
 
+import io
 import json
 import os
 
+import numpy as np
+
 import pytest
+import sagemaker.amazon.common as smac
+
 
 import sagemaker
 from sagemaker import image_uris
 from sagemaker.estimator import Estimator
+from sagemaker.s3 import S3Uploader
 from sagemaker.serializers import SimpleBaseSerializer
 from sagemaker.utils import unique_name_from_base
 from tests.integ import DATA_DIR, TRAINING_DEFAULT_TIMEOUT_MINUTES, datasets
@@ -102,6 +108,60 @@ def test_byo_estimator(sagemaker_session, region, cpu_instance_type, training_se
             assert prediction["score"] is not None
 
 
+@pytest.mark.release
+def test_estimator_register_publish_training_details(sagemaker_session, region, cpu_instance_type):
+
+    bucket = sagemaker_session.default_bucket()
+    prefix = "model-card-sample-notebook"
+
+    raw_data = (
+        (0.5, 0),
+        (0.75, 0),
+        (1.0, 0),
+        (1.25, 0),
+        (1.50, 0),
+        (1.75, 0),
+        (2.0, 0),
+        (2.25, 1),
+        (2.5, 0),
+        (2.75, 1),
+        (3.0, 0),
+        (3.25, 1),
+        (3.5, 0),
+        (4.0, 1),
+        (4.25, 1),
+        (4.5, 1),
+        (4.75, 1),
+        (5.0, 1),
+        (5.5, 1),
+    )
+    training_data = np.array(raw_data).astype("float32")
+    labels = training_data[:, 1]
+
+    # upload data to S3 bucket
+    buf = io.BytesIO()
+    smac.write_numpy_to_dense_tensor(buf, training_data, labels)
+    buf.seek(0)
+    s3_train_data = f"s3://{bucket}/{prefix}/train"
+    S3Uploader.upload_bytes(b=buf, s3_uri=s3_train_data, sagemaker_session=sagemaker_session)
+    output_location = f"s3://{bucket}/{prefix}/output"
+    container = image_uris.retrieve("linear-learner", region)
+    estimator = Estimator(
+        container,
+        role="SageMakerRole",
+        instance_count=1,
+        instance_type=cpu_instance_type,
+        output_path=output_location,
+        sagemaker_session=sagemaker_session,
+    )
+    estimator.set_hyperparameters(
+        feature_dim=2, mini_batch_size=10, predictor_type="binary_classifier"
+    )
+    estimator.fit({"train": s3_train_data})
+    print(f"Training job name: {estimator.latest_training_job.name}")
+    estimator.register()
+
+
 def test_async_byo_estimator(sagemaker_session, region, cpu_instance_type, training_set):
     image_uri = image_uris.retrieve("factorization-machines", region)
     endpoint_name = unique_name_from_base("byo")