Merge branch 'master' into feat/telemetry-attribution

Author: mollyheamazon

migration.md

@@ -25,7 +25,7 @@
 
 ## Migration Tool (MCP Server)
 
-An AI-powered migration tool is available as an MCP server that can analyze your V2 code, transform it to V3, validate the results, and answer migration questions interactively through your IDE.
+This AI-powered migration tool is available as an MCP server to analyze your V2 code, transform it to V3, validate the results, and answer migration questions interactively through your IDE.
 
 ### Installation
 
@@ -61,32 +61,22 @@ Add the following to your MCP configuration file:
 {
   "mcpServers": {
     "sagemaker-sdk-helper": {
-      "command": "sagemaker-sdk-helper",
-      "args": ["--log-level", "INFO"]
+      "command": "/path/to/installation",
+      "args": ["--log-level", "INFO"],
+      "autoApprove": [
+        "ask_question",
+        "transform_code",
+        "validate_code",
+        "analyze_code"
+      ]
+    }
     }
-  }
 }
+
 ```
 
 > **Note**: If you installed in a virtual environment, use the full path to the executable (find it with `which sagemaker-sdk-helper`).
 
-**With SDK source artifacts (recommended, 20-30% better accuracy):**
-
-```json
-{
-  "mcpServers": {
-    "sagemaker-sdk-helper": {
-      "command": "/path/to/.venv/bin/sagemaker-sdk-helper",
-      "args": [
-        "--log-level", "INFO",
-        "--v2-artifacts", "/path/to/sdk_v2/sagemaker-python-sdk",
-        "--v3-artifacts", "/path/to/sdk_v3/sagemaker-python-sdk"
-      ]
-    }
-  }
-}
-```
-
 After updating the config, restart your IDE or reconnect MCP servers (in Kiro: Command Palette → "MCP: Reconnect Servers").
 
 ### Kiro CLI
@@ -97,10 +87,16 @@ For Kiro CLI, add the same configuration to `~/.kiro/settings/mcp.json`:
 {
   "mcpServers": {
     "sagemaker-sdk-helper": {
-      "command": "sagemaker-sdk-helper",
-      "args": ["--log-level", "INFO"]
+      "command": "/path/to/installation",
+      "args": ["--log-level", "INFO"],
+      "autoApprove": [
+        "ask_question",
+        "transform_code",
+        "validate_code",
+        "analyze_code"
+      ]
+    }
     }
-  }
 }
 ```
 

sagemaker-core/src/sagemaker/core/image_uri_config/djl-lmi.json

@@ -46,7 +46,7 @@
                 "mx-central-1": "637423239942"
             },
             "repository": "djl-inference",
-            "tag_prefix": "0.36.0-lmi21.0.0-cu129"
+            "tag_prefix": "0.36.0-lmi22.0.0-cu129"
         },
         "0.35.0": {
             "registries": {

sagemaker-core/tests/unit/image_uris/test_djl.py

@@ -86,7 +86,7 @@ def _test_djl_uris(account, region, version, tag, djl_framework):
 # Known missing framework:version:region combinations that don't exist in ECR
 KNOWN_MISSING_COMBINATIONS = {
     "djl-lmi": {
-        "0.36.0-lmi21.0.0-cu129": {"ap-east-2"},
+        "0.36.0-lmi22.0.0-cu129": {"ap-east-2"},
         "0.35.0-lmi17.0.0-cu128": {"ap-east-2"},
         "0.34.0-lmi16.0.0-cu128": {"ap-east-2"},
         "0.33.0-lmi15.0.0-cu128": {"ap-east-2"},

sagemaker-train/src/sagemaker/train/aws_batch/batch_api_helper.py

@@ -32,6 +32,8 @@ def _submit_service_job(
     timeout: Optional[Dict] = None,
     share_identifier: Optional[str] = None,
     tags: Optional[Dict] = None,
+    quota_share_name: Optional[str] = None,
+    preemption_config: Optional[Dict] = None,
 ) -> Dict:
     """Batch submit_service_job API helper function.
 
@@ -44,6 +46,8 @@ def _submit_service_job(
         timeout: Set with value of timeout if specified, else default to 1 day.
         share_identifier: value of shareIdentifier if specified.
         tags: A dict of string to string representing Batch tags.
+        quota_share_name: Quota Share name for the Batch job.
+        preemption_config: Preemption configuration.
 
     Returns:
         A dict containing jobArn, jobName and jobId.
@@ -68,6 +72,10 @@ def _submit_service_job(
         payload["shareIdentifier"] = share_identifier
     if tags or training_payload_tags:
         payload["tags"] = __merge_tags(tags, training_payload_tags)
+    if quota_share_name:
+        payload["quotaShareName"] = quota_share_name
+    if preemption_config:
+        payload["preemptionConfiguration"] = preemption_config
     return client.submit_service_job(**payload)
 
 
@@ -96,21 +104,45 @@ def _describe_service_job(job_id: str) -> Dict:
         'jobId': 'string',
         'jobName': 'string',
         'jobQueue': 'string',
+        'latestAttempt': {
+            'serviceResourceId': {
+                'name': 'string',
+                'value': 'string'
+            }
+        },
+        'preemptionSummary': {
+            'preemptedAttemptCount': 123,
+            'recentPreemptedAttempts': [
+                {
+                    'serviceResourceId': {
+                        'name': 'string',
+                        'value': 'string'
+                    },
+                    'startedAt': 123,
+                    'stoppedAt': 123,
+                    'statusReason': 'string'
+                },
+            ]
+        },
         'retryStrategy': {
             'attempts': 123
         },
         'schedulingPriority': 123,
         'serviceRequestPayload': 'string',
-        'serviceJobType': 'EKS'|'ECS'|'ECS_FARGATE'|'SAGEMAKER_TRAINING',
+        'serviceJobType': 'SAGEMAKER_TRAINING',
         'shareIdentifier': 'string',
+        'quotaShareName': 'string',
+        'preemptionConfiguration': {
+            'preemptionRetriesBeforeTermination': 123
+        },
         'startedAt': 123,
-        'status': 'SUBMITTED'|'PENDING'|'RUNNABLE'|'STARTING'|'RUNNING'|'SUCCEEDED'|'FAILED',
+        'status': 'SUBMITTED'|'PENDING'|'RUNNABLE'|'SCHEDULED'|'STARTING'|'RUNNING'|'SUCCEEDED'|'FAILED',
         'statusReason': 'string',
         'stoppedAt': 123,
         'tags': {
             'string': 'string'
         },
-        'timeout': {
+        'timeoutConfig': {
             'attemptDurationSeconds': 123
         }
     }
@@ -132,6 +164,19 @@ def _terminate_service_job(job_id: str, reason: Optional[str] = "default termina
     return client.terminate_service_job(jobId=job_id, reason=reason)
 
 
+def _update_service_job(job_id: str, scheduling_priority: int) -> Dict:
+    """Batch update_service_job API helper function.
+
+    Args:
+        job_id: Job ID or Job Arn
+        scheduling_priority: An integer representing scheduling priority.
+
+    Returns: a dict containing jobArn, jobId and jobName.
+    """
+    client = get_batch_boto_client()
+    return client.update_service_job(jobId=job_id, schedulingPriority=scheduling_priority)
+
+
 def _list_service_job(
     job_queue: str,
     job_status: Optional[str] = None,

sagemaker-train/src/sagemaker/train/aws_batch/training_queue.py

@@ -41,6 +41,8 @@ def submit(
         share_identifier: Optional[str] = None,
         timeout: Optional[Dict] = None,
         tags: Optional[Dict] = None,
+        quota_share_name: Optional[str] = None,
+        preemption_config: Optional[Dict] = None,
     ) -> TrainingQueuedJob:
         """Submit a queued job and return a QueuedJob object.
 
@@ -53,6 +55,8 @@ def submit(
             share_identifier: Share identifier for Batch job.
             timeout: Timeout configuration for Batch job.
             tags: Tags apply to Batch job. These tags are for Batch job only.
+            quota_share_name: Quota Share name for the Batch job.
+            preemption_config: Preemption configuration.
 
         Returns: a TrainingQueuedJob object with Batch job ARN and job name.
 
@@ -85,6 +89,8 @@ def submit(
             timeout,
             share_identifier,
             tags,
+            quota_share_name,
+            preemption_config,
         )
         if "jobArn" not in resp or "jobName" not in resp:
             raise MissingRequiredArgument(

sagemaker-train/src/sagemaker/train/aws_batch/training_queued_job.py

@@ -29,7 +29,7 @@
     SourceCode,
     TrainingImageConfig,
 )
-from .batch_api_helper import _terminate_service_job, _describe_service_job
+from .batch_api_helper import _terminate_service_job, _describe_service_job, _update_service_job
 from .exception import NoTrainingJob, MissingRequiredArgument
 from ..utils import _get_training_job_name_from_training_job_arn
 from .constants import JOB_STATUS_COMPLETED, JOB_STATUS_FAILED, POLL_IN_SECONDS
@@ -85,6 +85,17 @@ def terminate(self, reason: Optional[str] = "Default terminate reason") -> None:
         """
         _terminate_service_job(self.job_arn, reason)
 
+    def update(self, scheduling_priority: int) -> Dict:
+        """Update Batch job.
+
+        Args:
+            scheduling_priority: An integer representing scheduling priority.
+
+        Returns: A dict which includes jobArn, jobName and jobId.
+
+        """
+        return _update_service_job(self.job_arn, scheduling_priority)
+
     def describe(self) -> Dict:
         """Describe Batch job.