Repository: apache/airflow (Python)

Issue:
Scheduler Crashes When Passing Invalid Value to Argument in default_args
### Apache Airflow version

2.9.2

### If "Other Airflow 2 version" selected, which one?

_No response_

### What happened?

When an invalid value is passed to any argument in default_args, such as passing a string value "2" to max_active_tis_per_dag which expects an integer, the scheduler crashes. The error message received is:

```
[2024-06-12T05:04:39.041+0000] {scheduler_command.py:54} ERROR - Exception when running scheduler job
Traceback (most recent call last):
  File "/usr/local/lib/python3.11/site-packages/airflow/cli/commands/scheduler_command.py", line 52, in _run_scheduler_job
    run_job(job=job_runner.job, execute_callable=job_runner._execute)
  File "/usr/local/lib/python3.11/site-packages/airflow/utils/session.py", line 79, in wrapper
    return func(*args, session=session, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/airflow/jobs/job.py", line 395, in run_job
    return execute_job(job, execute_callable=execute_callable)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/airflow/jobs/job.py", line 424, in execute_job
    ret = execute_callable()
          ^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/astronomer/airflow/version_check/plugin.py", line 30, in run_before
    fn(*args, **kwargs)
  File "/usr/local/lib/python3.11/site-packages/airflow/jobs/scheduler_job_runner.py", line 854, in _execute
    self._run_scheduler_loop()
  File "/usr/local/lib/python3.11/site-packages/airflow/jobs/scheduler_job_runner.py", line 986, in _run_scheduler_loop
    num_queued_tis = self._do_scheduling(session)
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/airflow/jobs/scheduler_job_runner.py", line 1099, in _do_scheduling
    num_queued_tis = self._critical_section_enqueue_task_instances(session=session)
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/airflow/jobs/scheduler_job_runner.py", line 677, in _critical_section_enqueue_task_instances
    queued_tis = self._executable_task_instances_to_queued(max_tis, session=session)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/airflow/jobs/scheduler_job_runner.py", line 533, in _executable_task_instances_to_queued
    if current_task_concurrency >= task_concurrency_limit:
       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: '>=' not supported between instances of 'int' and 'str'
```

### What you think should happen instead?

Instead of crashing the scheduler, Airflow should throw an import error on the Airflow UI for the problematic DAG. This would help in identifying the problematic DAG easily. The current behaviour makes it difficult to identify which DAG caused the issue, especially when managing hundreds of DAGs. This small mistake in DAG can also cause an outage in the production environment.

### How to reproduce

- Define a DAG with default_args that includes an invalid value for max_active_tis_per_dag, such as:

```
default_args = {
    'max_active_tis_per_dag': "2",  # invalid value, should be an integer
}
```
- Deploy and start the DAG.
- Observe the scheduler crash with the TypeError.

### Operating System

Linux ad25902d8cef 6.7.12-orbstack-00201-g2ddb8f197a46 #1 SMP Tue May 21 04:38:26 UTC 2024 aarch64 GNU/Linux

### Versions of Apache Airflow Providers

_No response_

### Deployment

Astronomer

### Deployment details

_No response_

### Anything else?

This problem occurs every time an invalid value is passed to an argument in default_args.

### Are you willing to submit PR?

- [ ] Yes I am willing to submit a PR!

### Code of Conduct

- [X] I agree to follow this project's [Code of Conduct](https://github.com/apache/airflow/blob/main/CODE_OF_CONDUCT.md)



Discussion: It's not a new issue in 2.9.2, confirmed with @hkc-8010 he is able to repro in the previous version as well. We should handle invalid values gracefully instead of getting the scheduler crashed



Fix (patch):
diff --git a/airflow/models/baseoperator.py b/airflow/models/baseoperator.py
index bbd629cfc156c..5d58cb80dc32c 100644
--- a/airflow/models/baseoperator.py
+++ b/airflow/models/baseoperator.py
@@ -95,7 +95,7 @@
 from airflow.utils.context import Context, context_get_outlet_events
 from airflow.utils.decorators import fixup_decorator_warning_stack
 from airflow.utils.edgemodifier import EdgeModifier
-from airflow.utils.helpers import validate_key
+from airflow.utils.helpers import validate_instance_args, validate_key
 from airflow.utils.operator_helpers import ExecutionCallableRunner
 from airflow.utils.operator_resources import Resources
 from airflow.utils.session import NEW_SESSION, provide_session
@@ -521,6 +521,47 @@ def __new__(cls, name, bases, namespace, **kwargs):
         return new_cls
 
 
+# TODO: The following mapping is used to validate that the arguments passed to the BaseOperator are of the
+#  correct type. This is a temporary solution until we find a more sophisticated method for argument
+#  validation. One potential method is to use `get_type_hints` from the typing module. However, this is not
+#  fully compatible with future annotations for Python versions below 3.10. Once we require a minimum Python
+#  version that supports `get_type_hints` effectively or find a better approach, we can replace this
+#  manual type-checking method.
+BASEOPERATOR_ARGS_EXPECTED_TYPES = {
+    "task_id": str,
+    "email": (str, Iterable),
+    "email_on_retry": bool,
+    "email_on_failure": bool,
+    "retries": int,
+    "retry_exponential_backoff": bool,
+    "depends_on_past": bool,
+    "ignore_first_depends_on_past": bool,
+    "wait_for_past_depends_before_skipping": bool,
+    "wait_for_downstream": bool,
+    "priority_weight": int,
+    "queue": str,
+    "pool": str,
+    "pool_slots": int,
+    "trigger_rule": str,
+    "run_as_user": str,
+    "task_concurrency": int,
+    "map_index_template": str,
+    "max_active_tis_per_dag": int,
+    "max_active_tis_per_dagrun": int,
+    "executor": str,
+    "do_xcom_push": bool,
+    "multiple_outputs": bool,
+    "doc": str,
+    "doc_md": str,
+    "doc_json": str,
+    "doc_yaml": str,
+    "doc_rst": str,
+    "task_display_name": str,
+    "logger_name": str,
+    "allow_nested_operators": bool,
+}
+
+
 @total_ordering
 class BaseOperator(AbstractOperator, metaclass=BaseOperatorMeta):
     r"""
@@ -1078,6 +1119,8 @@ def __init__(
         if SetupTeardownContext.active:
             SetupTeardownContext.update_context_map(self)
 
+        validate_instance_args(self, BASEOPERATOR_ARGS_EXPECTED_TYPES)
+
     def __eq__(self, other):
         if type(self) is type(other):
             # Use getattr() instead of __dict__ as __dict__ doesn't return
diff --git a/airflow/models/dag.py b/airflow/models/dag.py
index b2de299e64646..97e39d4a93f08 100644
--- a/airflow/models/dag.py
+++ b/airflow/models/dag.py
@@ -128,7 +128,7 @@
 from airflow.utils.dag_cycle_tester import check_cycle
 from airflow.utils.dates import cron_presets, date_range as utils_date_range
 from airflow.utils.decorators import fixup_decorator_warning_stack
-from airflow.utils.helpers import at_most_one, exactly_one, validate_key
+from airflow.utils.helpers import at_most_one, exactly_one, validate_instance_args, validate_key
 from airflow.utils.log.logging_mixin import LoggingMixin
 from airflow.utils.session import NEW_SESSION, provide_session
 from airflow.utils.sqlalchemy import (
@@ -341,6 +341,32 @@ def _create_orm_dagrun(
     return run
 
 
+# TODO: The following mapping is used to validate that the arguments passed to the DAG are of the correct
+#  type. This is a temporary solution until we find a more sophisticated method for argument validation.
+#  One potential method is to use `get_type_hints` from the typing module. However, this is not fully
+#  compatible with future annotations for Python versions below 3.10. Once we require a minimum Python
+#  version that supports `get_type_hints` effectively or find a better approach, we can replace this
+#  manual type-checking method.
+DAG_ARGS_EXPECTED_TYPES = {
+    "dag_id": str,
+    "description": str,
+    "max_active_tasks": int,
+    "max_active_runs": int,
+    "max_consecutive_failed_dag_runs": int,
+    "dagrun_timeout": timedelta,
+    "default_view": str,
+    "orientation": str,
+    "catchup": bool,
+    "doc_md": str,
+    "is_paused_upon_creation": bool,
+    "render_template_as_native_obj": bool,
+    "tags": list,
+    "auto_register": bool,
+    "fail_stop": bool,
+    "dag_display_name": str,
+}
+
+
 @functools.total_ordering
 class DAG(LoggingMixin):
     """
@@ -744,6 +770,8 @@ def __init__(
         # fileloc based only on the serialize dag
         self._processor_dags_folder = None
 
+        validate_instance_args(self, DAG_ARGS_EXPECTED_TYPES)
+
     def get_doc_md(self, doc_md: str | None) -> str | None:
         if doc_md is None:
             return doc_md
diff --git a/airflow/utils/helpers.py b/airflow/utils/helpers.py
index 3ad87bd104035..8d14aadc21fa9 100644
--- a/airflow/utils/helpers.py
+++ b/airflow/utils/helpers.py
@@ -60,6 +60,17 @@ def validate_key(k: str, max_length: int = 250):
         )
 
 
+def validate_instance_args(instance: object, expected_arg_types: dict[str, Any]) -> None:
+    """Validate that the instance has the expected types for the arguments."""
+    for arg_name, expected_arg_type in expected_arg_types.items():
+        instance_arg_value = getattr(instance, arg_name, None)
+        if instance_arg_value is not None and not isinstance(instance_arg_value, expected_arg_type):
+            raise TypeError(
+                f"'{arg_name}' has an invalid type {type(instance_arg_value)} with value "
+                f"{instance_arg_value}, expected type is {expected_arg_type}"
+            )
+
+
 def validate_group_key(k: str, max_length: int = 200):
     """Validate value used as a group key."""
     if not isinstance(k, str):
diff --git a/airflow/utils/task_group.py b/airflow/utils/task_group.py
index 37b03e1907f15..4c158498d1570 100644
--- a/airflow/utils/task_group.py
+++ b/airflow/utils/task_group.py
@@ -36,7 +36,7 @@
 )
 from airflow.models.taskmixin import DAGNode
 from airflow.serialization.enums import DagAttributeTypes
-from airflow.utils.helpers import validate_group_key
+from airflow.utils.helpers import validate_group_key, validate_instance_args
 
 if TYPE_CHECKING:
     from sqlalchemy.orm import Session
@@ -49,6 +49,21 @@
     from airflow.models.taskmixin import DependencyMixin
     from airflow.utils.edgemodifier import EdgeModifier
 
+# TODO: The following mapping is used to validate that the arguments passed to the TaskGroup are of the
+#  correct type. This is a temporary solution until we find a more sophisticated method for argument
+#  validation. One potential method is to use get_type_hints from the typing module. However, this is not
+#  fully compatible with future annotations for Python versions below 3.10. Once we require a minimum Python
+#  version that supports `get_type_hints` effectively or find a better approach, we can replace this
+#  manual type-checking method.
+TASKGROUP_ARGS_EXPECTED_TYPES = {
+    "group_id": str,
+    "prefix_group_id": bool,
+    "tooltip": str,
+    "ui_color": str,
+    "ui_fgcolor": str,
+    "add_suffix_on_collision": bool,
+}
+
 
 class TaskGroup(DAGNode):
     """
@@ -160,6 +175,8 @@ def __init__(
         self.upstream_task_ids = set()
         self.downstream_task_ids = set()
 
+        validate_instance_args(self, TASKGROUP_ARGS_EXPECTED_TYPES)
+
     def _check_for_group_id_collisions(self, add_suffix_on_collision: bool):
         if self._group_id is None:
             return