From 9e1cac9a9b6836f380e9bdaa673da5ced0b1cf08 Mon Sep 17 00:00:00 2001 From: Bas Harenslak Date: Sat, 22 Feb 2020 16:40:55 +0100 Subject: [PATCH 1/6] [AIRFLOW-6882] Implement SageMakerHook list_training_jobs method + convenience method for iterating list_* requests --- .../providers/amazon/aws/hooks/sagemaker.py | 82 +++++++++++++++++++ 1 file changed, 82 insertions(+) diff --git a/airflow/providers/amazon/aws/hooks/sagemaker.py b/airflow/providers/amazon/aws/hooks/sagemaker.py index 838cb8c401384..d959615f7ed3f 100644 --- a/airflow/providers/amazon/aws/hooks/sagemaker.py +++ b/airflow/providers/amazon/aws/hooks/sagemaker.py @@ -21,6 +21,8 @@ import tempfile import time import warnings +from functools import partial +from typing import List, Dict from botocore.exceptions import ClientError @@ -742,3 +744,83 @@ def check_training_status_with_log(self, job_name, non_terminal_states, failed_s billable_time = (last_description['TrainingEndTime'] - last_description['TrainingStartTime']) \ * instance_count self.log.info('Billable seconds: %d', int(billable_time.total_seconds()) + 1) + + def list_training_jobs(self, name_contains: str = None, max_results: int = None, **kwargs) -> List[Dict]: + """ + This method wraps boto3's list_training_jobs(). The training job name and max results are configurable + via arguments. Other arguments are not, and should be provided via kwargs. Note boto3 expects these in + CamelCase format, for example: + + .. code-block:: python + + list_training_jobs(name_contains="myjob", StatusEquals="Failed") + + .. seealso:: + https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker.html#SageMaker.Client.list_training_jobs + + :param name_contains: (optional) partial name to match + :param max_results: (optional) maximum number of results to return. None returns infinite results + :param kwargs: (optional) kwargs to boto3's list_training_jobs method + :return: results of the list_training_jobs request + """ + + config = dict() + + if name_contains: + if "NameContains" in kwargs: + raise AirflowException("Either name_contains or NameContains can be provided, not both.") + config["NameContains"] = name_contains + + if "MaxResults" in kwargs and kwargs["MaxResults"] is not None: + if max_results: + raise AirflowException("Either max_results or MaxResults can be provided, not both.") + # Unset MaxResults, we'll use the SageMakerHook's internal method for iteratively fetching results + max_results = kwargs["MaxResults"] + del kwargs["MaxResults"] + + config.update(kwargs) + list_training_jobs_request = partial(self.get_conn().list_training_jobs, **config) + results = self._list_request( + list_training_jobs_request, "TrainingJobSummaries", max_results=max_results + ) + return results + + def _list_request(self, partial_func, result_key: str, max_results: int = None) -> List[Dict]: + """ + All AWS boto3 list_* requests return results in batches (if the key "NextToken" is contained in the + result, there are more results to fetch). The default AWS batch size is 10, and configurable up to + 100. This function iteratively loads all results (or up to a given maximum). + + Each boto3 function returns the results in a different dict structure. The key of this structure must + be given to iterate over the results, e.g. "TransformJobSummaries" for list_transform_jobs(). + + :param partial_func: boto3 function with arguments + :param result_key: the result key to iterate over + :param max_results: maximum number of results to return (None = infinite) + :return: Results of the list_* request + """ + + SAGEMAKER_MAX_RESULTS = 100 # Fixed number set by AWS + + results = [] + next_token = None + + while True: + kwargs = dict() + if next_token is not None: + kwargs["NextToken"] = next_token + + if max_results is None: + kwargs["MaxResults"] = SAGEMAKER_MAX_RESULTS + else: + kwargs["MaxResults"] = min(max_results - len(results), SAGEMAKER_MAX_RESULTS) + + response = partial_func(**kwargs) + self.log.debug("Fetched %s results.", len(response[result_key])) + results.extend(response[result_key]) + + if "NextToken" not in response or (max_results is not None and len(results) == max_results): + # Return when there are no results left (no NextToken) or when we've reached max_results. + return results + else: + next_token = response["NextToken"] From 45fc40d7eaab3420ce2a9a1d2f5e7ec4ca17ca53 Mon Sep 17 00:00:00 2001 From: Bas Harenslak Date: Sat, 22 Feb 2020 20:11:00 +0100 Subject: [PATCH 2/6] Add if_jobname_exists argument and check if training job name exists before creation --- .../aws/operators/sagemaker_training.py | 24 ++++++++++++-- .../aws/operators/test_sagemaker_training.py | 33 +++++++++++++++++++ 2 files changed, 54 insertions(+), 3 deletions(-) diff --git a/airflow/providers/amazon/aws/operators/sagemaker_training.py b/airflow/providers/amazon/aws/operators/sagemaker_training.py index 2b9b4edffe697..06004530765b5 100644 --- a/airflow/providers/amazon/aws/operators/sagemaker_training.py +++ b/airflow/providers/amazon/aws/operators/sagemaker_training.py @@ -46,6 +46,9 @@ class SageMakerTrainingOperator(SageMakerBaseOperator): doesn't finish within max_ingestion_time seconds. If you set this parameter to None, the operation does not timeout. :type max_ingestion_time: int + :param if_jobname_exists: Behaviour if the job name already exists. Possible options are "increment" + (default) and "fail". + :type if_jobname_exists: str """ integer_fields = [ @@ -61,14 +64,15 @@ def __init__(self, print_log=True, check_interval=30, max_ingestion_time=None, + if_jobname_exists: str = "increment", # TODO: can use typing.Literal for this in Python 3.8 *args, **kwargs): - super().__init__(config=config, - *args, **kwargs) + super().__init__(config=config, *args, **kwargs) self.wait_for_completion = wait_for_completion self.print_log = print_log self.check_interval = check_interval self.max_ingestion_time = max_ingestion_time + self.if_jobname_exists = if_jobname_exists def expand_role(self): if 'RoleArn' in self.config: @@ -78,8 +82,22 @@ def expand_role(self): def execute(self, context): self.preprocess_config() - self.log.info('Creating SageMaker Training Job %s.', self.config['TrainingJobName']) + training_job_name = self.config["TrainingJobName"] + training_jobs = self.hook.list_training_jobs(name_contains=training_job_name) + # Check if given TrainingJobName already exists + if training_job_name in [tj["TrainingJobName"] for tj in training_jobs]: + if self.if_jobname_exists == "increment": + self.log.info("Found existing training job with name '%s'.", training_job_name) + new_training_job_name = f"{training_job_name}-{len(training_jobs) + 1}" + self.config["TrainingJobName"] = new_training_job_name + self.log.info("Incremented training job name to '%s'.", new_training_job_name) + elif self.if_jobname_exists == "fail": + raise AirflowException( + f"A SageMaker training job with name {training_job_name} already exists." + ) + + self.log.info("Creating SageMaker training job %s.", self.config["TrainingJobName"]) response = self.hook.create_training_job( self.config, wait_for_completion=self.wait_for_completion, diff --git a/tests/providers/amazon/aws/operators/test_sagemaker_training.py b/tests/providers/amazon/aws/operators/test_sagemaker_training.py index 29c506750e92d..2b0162b058a80 100644 --- a/tests/providers/amazon/aws/operators/test_sagemaker_training.py +++ b/tests/providers/amazon/aws/operators/test_sagemaker_training.py @@ -122,6 +122,39 @@ def test_execute_with_failure(self, mock_training, mock_client): self.assertRaises(AirflowException, self.sagemaker.execute, None) # pylint: enable=unused-argument + @mock.patch.object(SageMakerHook, "get_conn") + @mock.patch.object(SageMakerHook, "list_training_jobs") + @mock.patch.object(SageMakerHook, "create_training_job") + def test_execute_with_existing_job_increment( + self, mock_create_training_job, mock_list_training_jobs, mock_client + ): + self.sagemaker.if_jobname_exists = "increment" + mock_create_training_job.return_value = {"ResponseMetadata": {"HTTPStatusCode": 200}} + mock_list_training_jobs.return_value = [{"TrainingJobName": job_name}] + self.sagemaker.execute(None) + + expected_config = create_training_params.copy() + # Expect to see TrainingJobName suffixed with "-2" + expected_config["TrainingJobName"] = f"{job_name}-2" + mock_create_training_job.assert_called_once_with( + expected_config, + wait_for_completion=False, + print_log=True, + check_interval=5, + max_ingestion_time=None, + ) + + @mock.patch.object(SageMakerHook, "get_conn") + @mock.patch.object(SageMakerHook, "list_training_jobs") + @mock.patch.object(SageMakerHook, "create_training_job") + def test_execute_with_existing_job_fail( + self, mock_create_training_job, mock_list_training_jobs, mock_client + ): + self.sagemaker.if_jobname_exists = "fail" + mock_create_training_job.return_value = {"ResponseMetadata": {"HTTPStatusCode": 200}} + mock_list_training_jobs.return_value = [{"TrainingJobName": job_name}] + self.assertRaises(AirflowException, self.sagemaker.execute, None) + if __name__ == '__main__': unittest.main() From a2afc42937ffaed72605e40f0f77fe4af8368c02 Mon Sep 17 00:00:00 2001 From: Bas Harenslak Date: Sat, 29 Feb 2020 13:02:57 +0100 Subject: [PATCH 3/6] Small fixes --- airflow/providers/amazon/aws/hooks/sagemaker.py | 11 ++++++----- .../amazon/aws/operators/sagemaker_training.py | 12 ++++++------ .../amazon/aws/operators/test_sagemaker_training.py | 6 +++--- 3 files changed, 15 insertions(+), 14 deletions(-) diff --git a/airflow/providers/amazon/aws/hooks/sagemaker.py b/airflow/providers/amazon/aws/hooks/sagemaker.py index d959615f7ed3f..6edd1e361a87c 100644 --- a/airflow/providers/amazon/aws/hooks/sagemaker.py +++ b/airflow/providers/amazon/aws/hooks/sagemaker.py @@ -791,8 +791,9 @@ def _list_request(self, partial_func, result_key: str, max_results: int = None) result, there are more results to fetch). The default AWS batch size is 10, and configurable up to 100. This function iteratively loads all results (or up to a given maximum). - Each boto3 function returns the results in a different dict structure. The key of this structure must - be given to iterate over the results, e.g. "TransformJobSummaries" for list_transform_jobs(). + Each boto3 list_* function returns the results in a list with a different name. The key of this + structure must be given to iterate over the results, e.g. "TransformJobSummaries" for + list_transform_jobs(). :param partial_func: boto3 function with arguments :param result_key: the result key to iterate over @@ -800,7 +801,7 @@ def _list_request(self, partial_func, result_key: str, max_results: int = None) :return: Results of the list_* request """ - SAGEMAKER_MAX_RESULTS = 100 # Fixed number set by AWS + sagemaker_max_results = 100 # Fixed number set by AWS results = [] next_token = None @@ -811,9 +812,9 @@ def _list_request(self, partial_func, result_key: str, max_results: int = None) kwargs["NextToken"] = next_token if max_results is None: - kwargs["MaxResults"] = SAGEMAKER_MAX_RESULTS + kwargs["MaxResults"] = sagemaker_max_results else: - kwargs["MaxResults"] = min(max_results - len(results), SAGEMAKER_MAX_RESULTS) + kwargs["MaxResults"] = min(max_results - len(results), sagemaker_max_results) response = partial_func(**kwargs) self.log.debug("Fetched %s results.", len(response[result_key])) diff --git a/airflow/providers/amazon/aws/operators/sagemaker_training.py b/airflow/providers/amazon/aws/operators/sagemaker_training.py index 06004530765b5..ca3d09d4b1d4a 100644 --- a/airflow/providers/amazon/aws/operators/sagemaker_training.py +++ b/airflow/providers/amazon/aws/operators/sagemaker_training.py @@ -46,9 +46,9 @@ class SageMakerTrainingOperator(SageMakerBaseOperator): doesn't finish within max_ingestion_time seconds. If you set this parameter to None, the operation does not timeout. :type max_ingestion_time: int - :param if_jobname_exists: Behaviour if the job name already exists. Possible options are "increment" + :param action_if_job_exists: Behaviour if the job name already exists. Possible options are "increment" (default) and "fail". - :type if_jobname_exists: str + :type action_if_job_exists: str """ integer_fields = [ @@ -64,7 +64,7 @@ def __init__(self, print_log=True, check_interval=30, max_ingestion_time=None, - if_jobname_exists: str = "increment", # TODO: can use typing.Literal for this in Python 3.8 + action_if_job_exists: str = "increment", # TODO use typing.Literal for this in Python 3.8 *args, **kwargs): super().__init__(config=config, *args, **kwargs) @@ -72,7 +72,7 @@ def __init__(self, self.print_log = print_log self.check_interval = check_interval self.max_ingestion_time = max_ingestion_time - self.if_jobname_exists = if_jobname_exists + self.action_if_job_exists = action_if_job_exists def expand_role(self): if 'RoleArn' in self.config: @@ -87,12 +87,12 @@ def execute(self, context): # Check if given TrainingJobName already exists if training_job_name in [tj["TrainingJobName"] for tj in training_jobs]: - if self.if_jobname_exists == "increment": + if self.action_if_job_exists == "increment": self.log.info("Found existing training job with name '%s'.", training_job_name) new_training_job_name = f"{training_job_name}-{len(training_jobs) + 1}" self.config["TrainingJobName"] = new_training_job_name self.log.info("Incremented training job name to '%s'.", new_training_job_name) - elif self.if_jobname_exists == "fail": + elif self.action_if_job_exists == "fail": raise AirflowException( f"A SageMaker training job with name {training_job_name} already exists." ) diff --git a/tests/providers/amazon/aws/operators/test_sagemaker_training.py b/tests/providers/amazon/aws/operators/test_sagemaker_training.py index 2b0162b058a80..33934f9fde771 100644 --- a/tests/providers/amazon/aws/operators/test_sagemaker_training.py +++ b/tests/providers/amazon/aws/operators/test_sagemaker_training.py @@ -128,13 +128,13 @@ def test_execute_with_failure(self, mock_training, mock_client): def test_execute_with_existing_job_increment( self, mock_create_training_job, mock_list_training_jobs, mock_client ): - self.sagemaker.if_jobname_exists = "increment" + self.sagemaker.action_if_job_exists = "increment" mock_create_training_job.return_value = {"ResponseMetadata": {"HTTPStatusCode": 200}} mock_list_training_jobs.return_value = [{"TrainingJobName": job_name}] self.sagemaker.execute(None) expected_config = create_training_params.copy() - # Expect to see TrainingJobName suffixed with "-2" + # Expect to see TrainingJobName suffixed with "-2" because we return one existing job expected_config["TrainingJobName"] = f"{job_name}-2" mock_create_training_job.assert_called_once_with( expected_config, @@ -150,7 +150,7 @@ def test_execute_with_existing_job_increment( def test_execute_with_existing_job_fail( self, mock_create_training_job, mock_list_training_jobs, mock_client ): - self.sagemaker.if_jobname_exists = "fail" + self.sagemaker.action_if_job_exists = "fail" mock_create_training_job.return_value = {"ResponseMetadata": {"HTTPStatusCode": 200}} mock_list_training_jobs.return_value = [{"TrainingJobName": job_name}] self.assertRaises(AirflowException, self.sagemaker.execute, None) From f1850202ec3cba345f1d7e4fc210f7171a7a7900 Mon Sep 17 00:00:00 2001 From: Bas Harenslak Date: Sat, 29 Feb 2020 13:22:26 +0100 Subject: [PATCH 4/6] Fix isort --- airflow/providers/amazon/aws/hooks/sagemaker.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/providers/amazon/aws/hooks/sagemaker.py b/airflow/providers/amazon/aws/hooks/sagemaker.py index 6edd1e361a87c..32dd1bd480c22 100644 --- a/airflow/providers/amazon/aws/hooks/sagemaker.py +++ b/airflow/providers/amazon/aws/hooks/sagemaker.py @@ -22,7 +22,7 @@ import time import warnings from functools import partial -from typing import List, Dict +from typing import Dict, List from botocore.exceptions import ClientError From 4bacba0014295e5ec167e4d544d85a9b69082bab Mon Sep 17 00:00:00 2001 From: Bas Harenslak Date: Sat, 29 Feb 2020 13:44:42 +0100 Subject: [PATCH 5/6] Fix mypy issues --- airflow/providers/amazon/aws/hooks/sagemaker.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/airflow/providers/amazon/aws/hooks/sagemaker.py b/airflow/providers/amazon/aws/hooks/sagemaker.py index 32dd1bd480c22..b99ed8de1cd80 100644 --- a/airflow/providers/amazon/aws/hooks/sagemaker.py +++ b/airflow/providers/amazon/aws/hooks/sagemaker.py @@ -22,7 +22,7 @@ import time import warnings from functools import partial -from typing import Dict, List +from typing import Dict, List, Optional from botocore.exceptions import ClientError @@ -745,7 +745,9 @@ def check_training_status_with_log(self, job_name, non_terminal_states, failed_s * instance_count self.log.info('Billable seconds: %d', int(billable_time.total_seconds()) + 1) - def list_training_jobs(self, name_contains: str = None, max_results: int = None, **kwargs) -> List[Dict]: + def list_training_jobs( + self, name_contains: Optional[str] = None, max_results: Optional[int] = None, **kwargs + ) -> List[Dict]: """ This method wraps boto3's list_training_jobs(). The training job name and max results are configurable via arguments. Other arguments are not, and should be provided via kwargs. Note boto3 expects these in @@ -785,7 +787,7 @@ def list_training_jobs(self, name_contains: str = None, max_results: int = None, ) return results - def _list_request(self, partial_func, result_key: str, max_results: int = None) -> List[Dict]: + def _list_request(self, partial_func, result_key: str, max_results: Optional[int] = None) -> List[Dict]: """ All AWS boto3 list_* requests return results in batches (if the key "NextToken" is contained in the result, there are more results to fetch). The default AWS batch size is 10, and configurable up to @@ -803,7 +805,7 @@ def _list_request(self, partial_func, result_key: str, max_results: int = None) sagemaker_max_results = 100 # Fixed number set by AWS - results = [] + results: List[Dict] = [] next_token = None while True: From 05b51cee1b4ff00c48994a2d8deb7d404cd39dbb Mon Sep 17 00:00:00 2001 From: Bas Harenslak Date: Sat, 29 Feb 2020 14:33:02 +0100 Subject: [PATCH 6/6] Add value check for action_if_job_exists --- .../providers/amazon/aws/operators/sagemaker_training.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/airflow/providers/amazon/aws/operators/sagemaker_training.py b/airflow/providers/amazon/aws/operators/sagemaker_training.py index ca3d09d4b1d4a..afe070f93feff 100644 --- a/airflow/providers/amazon/aws/operators/sagemaker_training.py +++ b/airflow/providers/amazon/aws/operators/sagemaker_training.py @@ -72,7 +72,14 @@ def __init__(self, self.print_log = print_log self.check_interval = check_interval self.max_ingestion_time = max_ingestion_time - self.action_if_job_exists = action_if_job_exists + + if action_if_job_exists in ("increment", "fail"): + self.action_if_job_exists = action_if_job_exists + else: + raise AirflowException( + "Argument action_if_job_exists accepts only 'increment' and 'fail'. " + f"Provided value: '{action_if_job_exists}'." + ) def expand_role(self): if 'RoleArn' in self.config: