feature: AWS Batch for SageMaker Training jobs (#5249)

---------

Co-authored-by: Greg Katkov <[email protected]>
Co-authored-by: haoxinwa <[email protected]>
Co-authored-by: JennaZhao <[email protected]>
Co-authored-by: Jessica Zhu <[email protected]>
Co-authored-by: David Lindskog <[email protected]>

Author: 6 people

src/sagemaker/aws_batch/__init__.py

@@ -0,0 +1,186 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""The module provides helper function for Batch Submit/Describe/Terminal job APIs."""
+from __future__ import absolute_import
+
+import json
+from typing import List, Dict, Optional
+from sagemaker.aws_batch.constants import (
+    SAGEMAKER_TRAINING,
+    DEFAULT_TIMEOUT,
+    DEFAULT_SAGEMAKER_TRAINING_RETRY_CONFIG,
+)
+from sagemaker.aws_batch.boto_client import get_batch_boto_client
+
+
+def submit_service_job(
+    training_payload: Dict,
+    job_name: str,
+    job_queue: str,
+    retry_config: Optional[Dict] = None,
+    scheduling_priority: Optional[int] = None,
+    timeout: Optional[Dict] = None,
+    share_identifier: Optional[str] = None,
+    tags: Optional[Dict] = None,
+) -> Dict:
+    """Batch submit_service_job API helper function.
+
+    Args:
+        training_payload: a dict containing a dict of arguments for Training job.
+        job_name: Batch job name.
+        job_queue: Batch job queue ARN.
+        retry_config: Batch job retry configuration.
+        scheduling_priority: An integer representing scheduling priority.
+        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.
+
+    Returns:
+        A dict containing jobArn, jobName and jobId.
+    """
+    if timeout is None:
+        timeout = DEFAULT_TIMEOUT
+    client = get_batch_boto_client()
+    training_payload_tags = training_payload.pop("Tags", None)
+    payload = {
+        "jobName": job_name,
+        "jobQueue": job_queue,
+        "retryStrategy": DEFAULT_SAGEMAKER_TRAINING_RETRY_CONFIG,
+        "serviceJobType": SAGEMAKER_TRAINING,
+        "serviceRequestPayload": json.dumps(training_payload),
+        "timeoutConfig": timeout,
+    }
+    if retry_config:
+        payload["retryStrategy"] = retry_config
+    if scheduling_priority:
+        payload["schedulingPriority"] = scheduling_priority
+    if share_identifier:
+        payload["shareIdentifier"] = share_identifier
+    if tags or training_payload_tags:
+        payload["tags"] = __merge_tags(tags, training_payload_tags)
+    return client.submit_service_job(**payload)
+
+
+def describe_service_job(job_id: str) -> Dict:
+    """Batch describe_service_job API helper function.
+
+    Args:
+        job_id: Job ID used.
+
+    Returns: a dict. See the sample below
+    {
+        'attempts': [
+            {
+                'serviceResourceId': {
+                    'name': 'string',
+                    'value': 'string'
+                },
+                'startedAt': 123,
+                'stoppedAt': 123,
+                'statusReason': 'string'
+            },
+        ],
+        'createdAt': 123,
+        'isTerminated': True|False,
+        'jobArn': 'string',
+        'jobId': 'string',
+        'jobName': 'string',
+        'jobQueue': 'string',
+        'retryStrategy': {
+            'attempts': 123
+        },
+        'schedulingPriority': 123,
+        'serviceRequestPayload': 'string',
+        'serviceJobType': 'EKS'|'ECS'|'ECS_FARGATE'|'SAGEMAKER_TRAINING',
+        'shareIdentifier': 'string',
+        'startedAt': 123,
+        'status': 'SUBMITTED'|'PENDING'|'RUNNABLE'|'STARTING'|'RUNNING'|'SUCCEEDED'|'FAILED',
+        'statusReason': 'string',
+        'stoppedAt': 123,
+        'tags': {
+            'string': 'string'
+        },
+        'timeout': {
+            'attemptDurationSeconds': 123
+        }
+    }
+    """
+    client = get_batch_boto_client()
+    return client.describe_service_job(jobId=job_id)
+
+
+def terminate_service_job(job_id: str, reason: Optional[str] = "default terminate reason") -> Dict:
+    """Batch terminate_service_job API helper function.
+
+    Args:
+        job_id: Job ID
+        reason: A string representing terminate reason.
+
+    Returns: an empty dict
+    """
+    client = get_batch_boto_client()
+    return client.terminate_service_job(jobId=job_id, reason=reason)
+
+
+def list_service_job(
+    job_queue: str,
+    job_status: Optional[str] = None,
+    filters: Optional[List] = None,
+    next_token: Optional[str] = None,
+) -> Dict:
+    """Batch list_service_job API helper function.
+
+    Args:
+        job_queue: Batch job queue ARN.
+        job_status: Batch job status.
+        filters: A list of Dict. Each contains a filter.
+        next_token: Used to retrieve data in next page.
+
+    Returns: A generator containing list results.
+
+    """
+    client = get_batch_boto_client()
+    payload = {"jobQueue": job_queue}
+    if filters:
+        payload["filters"] = filters
+    if next_token:
+        payload["nextToken"] = next_token
+    if job_status:
+        payload["jobStatus"] = job_status
+    part_of_jobs = client.list_service_jobs(**payload)
+    next_token = part_of_jobs.get("nextToken")
+    yield part_of_jobs
+    if next_token:
+        yield from list_service_job(job_queue, job_status, filters, next_token)
+
+
+def __merge_tags(batch_tags: Optional[Dict], training_tags: Optional[List]) -> Optional[Dict]:
+    """Merges Batch and training payload tags.
+
+    Returns a copy of Batch tags merged with training payload tags.  Training payload tags take
+    precedence in the case of key conflicts.
+
+    :param batch_tags: A dict of string to string representing Batch tags.
+    :param training_tags: A list of `{"Key": "string", "Value": "string"}` objects representing
+    training payload tags.
+    :return: A dict of string to string representing batch tags merged with training tags.
+    batch_tags is returned unaltered if training_tags is None or empty.
+    """
+    if not training_tags:
+        return batch_tags
+
+    training_tags_to_merge = {tag["Key"]: tag["Value"] for tag in training_tags}
+    batch_tags_copy = batch_tags.copy() if batch_tags else {}
+    batch_tags_copy.update(training_tags_to_merge)
+
+    return batch_tags_copy

src/sagemaker/aws_batch/batch_api_helper.py

@@ -0,0 +1,33 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""The file provides helper function for getting Batch boto client."""
+from __future__ import absolute_import
+
+from typing import Optional
+import boto3
+
+
+def get_batch_boto_client(
+    region: Optional[str] = None,
+    endpoint: Optional[str] = None,
+) -> boto3.session.Session.client:
+    """Helper function for getting Batch boto3 client.
+
+    Args:
+        region: Region specified
+        endpoint: Batch API endpoint.
+
+    Returns: Batch boto3 client.
+
+    """
+    return boto3.client("batch", region_name=region, endpoint_url=endpoint)

src/sagemaker/aws_batch/boto_client.py

@@ -0,0 +1,34 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""The file defines constants used for Batch API helper functions."""
+
+from __future__ import absolute_import
+
+SAGEMAKER_TRAINING = "SAGEMAKER_TRAINING"
+DEFAULT_ATTEMPT_DURATION_IN_SECONDS = 86400  # 1 day in seconds.
+DEFAULT_TIMEOUT = {"attemptDurationSeconds": DEFAULT_ATTEMPT_DURATION_IN_SECONDS}
+POLL_IN_SECONDS = 5
+JOB_STATUS_RUNNING = "RUNNING"
+JOB_STATUS_COMPLETED = "SUCCEEDED"
+JOB_STATUS_FAILED = "FAILED"
+DEFAULT_SAGEMAKER_TRAINING_RETRY_CONFIG = {
+    "attempts": 1,
+    "evaluateOnExit": [
+        {
+            "action": "RETRY",
+            "onStatusReason": "Received status from SageMaker:InternalServerError: "
+            "We encountered an internal error. Please try again.",
+        },
+        {"action": "EXIT", "onStatusReason": "*"},
+    ],
+}

src/sagemaker/aws_batch/constants.py

@@ -0,0 +1,52 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""The file Defines customized exception for Batch queueing"""
+from __future__ import absolute_import
+
+
+class NoTrainingJob(Exception):
+    """Define NoTrainingJob Exception.
+
+    It means no Training job has been created by AWS Batch service.
+    """
+
+    def __init__(self, value):
+        super().__init__(value)
+        self.value = value
+
+    def __str__(self):
+        """Convert Exception to string.
+
+        Returns: a String containing exception error messages.
+
+        """
+        return repr(self.value)
+
+
+class MissingRequiredArgument(Exception):
+    """Define MissingRequiredArgument exception.
+
+    It means some required arguments are missing.
+    """
+
+    def __init__(self, value):
+        super().__init__(value)
+        self.value = value
+
+    def __str__(self):
+        """Convert Exception to string.
+
+        Returns: a String containing exception error messages.
+
+        """
+        return repr(self.value)