fix: training arn support (aws#1494)

* fix: training job ARN

* pylint

---------

Co-authored-by: Jacky Lee <[email protected]>

Author: grenmester

src/sagemaker/serve/builder/jumpstart_builder.py

@@ -546,24 +546,29 @@ def _update_model_data_for_fine_tuned_model(self, pysdk_model: Type[Model]) -> T
             pysdk_model.add_tags(
                 {"key": Tag.FINE_TUNING_MODEL_PATH, "value": fine_tuning_model_path}
             )
+            logger.info(
+                "FINE_TUNING_MODEL_PATH detected. Using fine-tuned model found in %s.",
+                fine_tuning_model_path,
+            )
             return pysdk_model
 
         if fine_tuning_job_name := self.model_metadata.get("FINE_TUNING_JOB_NAME"):
             try:
                 response = self.sagemaker_session.sagemaker_client.describe_training_job(
                     TrainingJobName=fine_tuning_job_name
                 )
-                fine_tuning_model_path = response["OutputDataConfig"]["S3OutputPath"]
+                fine_tuning_model_path = response["ModelArtifacts"]["S3ModelArtifacts"]
                 pysdk_model.model_data["S3DataSource"]["S3Uri"] = fine_tuning_model_path
-                pysdk_model.model_data["S3DataSource"]["CompressionType"] = response[
-                    "OutputDataConfig"
-                ]["CompressionType"]
                 pysdk_model.add_tags(
                     [
                         {"key": Tag.FINE_TUNING_JOB_NAME, "value": fine_tuning_job_name},
                         {"key": Tag.FINE_TUNING_MODEL_PATH, "value": fine_tuning_model_path},
                     ]
                 )
+                logger.info(
+                    "FINE_TUNING_JOB_NAME detected. Using fine-tuned model found in %s.",
+                    fine_tuning_model_path,
+                )
                 return pysdk_model
             except ClientError:
                 raise ValueError(

tests/unit/sagemaker/serve/builder/test_js_builder.py

@@ -1039,9 +1039,8 @@ def test_fine_tuned_model_with_fine_tuning_job_name(
         mock_fine_tuning_model_path = "s3://test"
         mock_sagemaker_session = Mock()
         mock_sagemaker_session.sagemaker_client.describe_training_job.return_value = {
-            "OutputDataConfig": {
-                "S3OutputPath": mock_fine_tuning_model_path,
-                "CompressionType": "None",
+            "ModelArtifacts": {
+                "S3ModelArtifacts": mock_fine_tuning_model_path,
             }
         }
         mock_pre_trained_model.return_value.image_uri = mock_djl_image_uri