Queue and QueuedJob implementation for AWS Batch Service (aws#1511)

* feat: Queue and QueuedJob implementation for AWS Batch Service (aws#991)

* fix: Update Batch endpoint; Function renaming (aws#1230)

Co-authored-by: Jenna Zhao <[email protected]>

* chore: minor formatting

---------

Co-authored-by: haoxinwa <[email protected]>
Co-authored-by: JennaZhao <[email protected]>
Co-authored-by: Jenna Zhao <[email protected]>
Co-authored-by: Gregory Katkov <[email protected]>

Author: 5 people

src/sagemaker/batch_queueing/__init__.py

@@ -0,0 +1,146 @@
+"""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 .constants import SAGEMAKER_TRAINING, DEFAULT_TIMEOUT
+from .boto_client import get_batch_boto_client
+
+
+def submit_service_job(
+    training_payload: Dict,
+    job_name: str,
+    job_queue: str,
+    retry_attempts: Optional[int] = None,
+    scheduling_priority: Optional[int] = None,
+    timeout: Optional[Dict] = None,
+    share_identifier: Optional[str] = None,
+    tags: Optional[List] = 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_attempts: The number of retry expected.
+        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 list of string representing Batch tags.
+
+    Returns:
+        A dict containing jobArn, jobName and jobId.
+    """
+    if timeout is None:
+        timeout = DEFAULT_TIMEOUT
+    client = get_batch_boto_client()
+    payload = {
+        "jobName": job_name,
+        "jobQueue": job_queue,
+        "serviceJobType": SAGEMAKER_TRAINING,
+        "serviceRequestPayload": json.dumps(training_payload),
+        "timeoutConfig": timeout,
+    }
+    if retry_attempts:
+        payload["retryStrategy"] = {"attempts": retry_attempts}
+    if scheduling_priority:
+        payload["schedulingPriority"] = scheduling_priority
+    if share_identifier:
+        payload["shareIdentifier"] = share_identifier
+    if tags:
+        payload["tags"] = 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)

src/sagemaker/batch_queueing/batch_api_helper.py

@@ -0,0 +1,21 @@
+"""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("sm_batch", region_name=region, endpoint_url=endpoint)

src/sagemaker/batch_queueing/boto_client.py

@@ -0,0 +1,8 @@
+"""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}
+JOB_STATUS_RUNNING = "RUNNING"

src/sagemaker/batch_queueing/constants.py

@@ -0,0 +1,41 @@
+"""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)

src/sagemaker/batch_queueing/exception.py

@@ -0,0 +1,136 @@
+"""Define Queue class for AWS Batch service"""
+
+from __future__ import absolute_import
+from typing import Dict, Optional, List
+import logging
+from sagemaker.estimator import Estimator, _TrainingJob
+from .queue_job import QueuedJob
+from .batch_api_helper import submit_service_job, list_service_job
+from .exception import MissingRequiredArgument
+from .constants import DEFAULT_TIMEOUT, JOB_STATUS_RUNNING
+
+
+class Queue:
+    """Queue class for AWS Batch service
+
+    With this class, customers are able to create a new queue and submit jobs to AWS Batch Service.
+    """
+
+    def __init__(self, queue_name: str):
+        self.queue_name = queue_name
+
+    def submit(
+        self,
+        estimator: Estimator,
+        inputs,
+        job_name: Optional[str] = None,
+        max_retry_attempts: Optional[int] = None,
+        priority: Optional[int] = None,
+        share_identifier: Optional[str] = None,
+        timeout: Optional[Dict] = None,
+        tags: Optional[Dict] = None,
+        experiment_config: Optional[Dict] = None,
+    ) -> QueuedJob:
+        """Submit a queued job and return a QueuedJob object.
+
+        Args:
+            estimator: Training job estimator object.
+            inputs: Training job inputs.
+            job_name: Batch job name.
+            max_retry_attempts: Max retry attempts for Batch job.
+            priority: Scheduling priority for Batch job.
+            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.
+            experiment_config: Experiment management configuration.
+                Optionally, the dict can contain four keys:
+                'ExperimentName', 'TrialName', 'TrialComponentDisplayName' and 'RunName'.
+
+        Returns: a QueuedJob object with Batch job ARN and job name.
+
+        """
+        if experiment_config is None:
+            experiment_config = {}
+        estimator.prepare_workflow_for_training(job_name)
+        training_args = _TrainingJob.get_train_args(estimator, inputs, experiment_config)
+        training_payload = estimator.sagemaker_session.get_train_request(**training_args)
+        training_payload = _convert_first_letter_to_lowercase_for_all_keys(training_payload)
+
+        if timeout is None:
+            timeout = DEFAULT_TIMEOUT
+        if job_name is None:
+            job_name = training_payload["trainingJobName"]
+
+        resp = submit_service_job(
+            training_payload,
+            job_name,
+            self.queue_name,
+            max_retry_attempts,
+            priority,
+            timeout,
+            share_identifier,
+            tags,
+        )
+        if "jobArn" not in resp or "jobName" not in resp:
+            raise MissingRequiredArgument(
+                "jobArn or jobName is missing in response from Batch submit_service_job API"
+            )
+        return QueuedJob(resp["jobArn"], resp["jobName"])
+
+    def list_jobs(
+        self, job_name: Optional[str] = None, status: Optional[str] = JOB_STATUS_RUNNING
+    ) -> List[QueuedJob]:
+        """List Batch jobs according to job_name or status.
+
+        Args:
+            job_name: Batch job name.
+            status: Batch job status.
+
+        Returns: A list of QueuedJob.
+
+        """
+        filters = None
+        if job_name:
+            filters = [{"name": "JOB_NAME", "values": [job_name]}]
+            status = None  # job_status is ignored when job_name is specified.
+        jobs_to_return = []
+        next_token = None
+        for job_result_dict in list_service_job(self.queue_name, status, filters, next_token):
+            for job_result in job_result_dict.get("jobSummaryList", []):
+                if "jobArn" in job_result and "jobName" in job_result:
+                    jobs_to_return.append(QueuedJob(job_result["jobArn"], job_result["jobName"]))
+                else:
+                    logging.warning("Missing JobArn or JobName in Batch ListJobs API")
+                    continue
+        return jobs_to_return
+
+
+def _first_letter_to_lowercase(key):
+    """Convert first letter in string to lowercase.
+
+    Args:
+        key: a string.
+
+    Returns: a string with first letter being lowercase.
+
+    """
+    return str.lower(key[0]) + key[1:]
+
+
+def _convert_first_letter_to_lowercase_for_all_keys(payload):
+    """Convert first letter of all keys to lowercase in a Dict.
+
+    Args:
+        payload: a Dict or a List.
+
+    Returns: converted Dict or List
+
+    """
+    if isinstance(payload, dict):
+        return {
+            _first_letter_to_lowercase(key): _convert_first_letter_to_lowercase_for_all_keys(value)
+            for key, value in payload.items()
+        }
+    if isinstance(payload, list):
+        return [_convert_first_letter_to_lowercase_for_all_keys(item) for item in payload]
+    return payload