Taragolis commented on code in PR #37618: URL: https://github.com/apache/airflow/pull/37618#discussion_r1502885932
########## airflow/providers/amazon/aws/executors/batch/batch_executor_config.py: ########## @@ -0,0 +1,78 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License 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. + +""" +AWS Batch Executor configuration. + +This is the configuration for calling the Batch ``submit_job`` function. The AWS Batch Executor calls +Boto3's ``submit_job(**kwargs)`` function with the kwargs templated by this dictionary. See the URL +below for documentation on the parameters accepted by the Boto3 submit_job function. + +.. seealso:: + https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/batch/client/submit_job.html + +""" + +from __future__ import annotations + +import json + +from airflow.configuration import conf +from airflow.providers.amazon.aws.executors.batch.utils import ( + CONFIG_GROUP_NAME, + AllBatchConfigKeys, + BatchSubmitJobKwargsConfigKeys, +) +from airflow.providers.amazon.aws.executors.ecs.utils import camelize_dict_keys +from airflow.utils.helpers import prune_dict + + +def _fetch_templated_kwargs() -> dict[str, str]: + submit_job_kwargs_value = conf.get( + CONFIG_GROUP_NAME, AllBatchConfigKeys.SUBMIT_JOB_KWARGS, fallback=dict() + ) + return json.loads(str(submit_job_kwargs_value)) + + +def _fetch_config_values() -> dict[str, str]: + return prune_dict( + {key: conf.get(CONFIG_GROUP_NAME, key, fallback=None) for key in BatchSubmitJobKwargsConfigKeys()} + ) + + +def build_submit_kwargs() -> dict: + job_kwargs = _fetch_config_values() + job_kwargs.update(_fetch_templated_kwargs()) + + if "containerOverrides" not in job_kwargs: + job_kwargs["containerOverrides"] = {} # type: ignore + job_kwargs["containerOverrides"]["command"] = [] # type: ignore + + if "nodeOverrides" in job_kwargs: + raise KeyError("Multi-node jobs are not currently supported.") + if "eksPropertiesOverride" in job_kwargs: + raise KeyError("Eks jobs are not currently supported.") + + # some checks with some helpful errors + assert isinstance(job_kwargs, dict) Review Comment: We should avoid to use assert outside of the tests: https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#don-t-use-asserts-outside-tests If it for make mypy happy, then it fine, just need keep it under the TYPE_CHECKING block ```suggestion if TYPE_CHECKING: assert isinstance(job_kwargs, dict) ``` if it intends to check in runtime then better raise appropriate error, e.g. TypeError ########## airflow/providers/amazon/aws/executors/batch/utils.py: ########## @@ -0,0 +1,157 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License 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. +from __future__ import annotations + +from collections import defaultdict +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Dict, List + +if TYPE_CHECKING: + from airflow.models.taskinstance import TaskInstanceKey +from airflow.providers.amazon.aws.executors.utils.base_config_keys import BaseConfigKeys +from airflow.utils.state import State Review Comment: ```suggestion from airflow.providers.amazon.aws.executors.utils.base_config_keys import BaseConfigKeys from airflow.utils.state import State if TYPE_CHECKING: from airflow.models.taskinstance import TaskInstanceKey ``` -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: commits-unsubscr...@airflow.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org
