You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by GitBox <gi...@apache.org> on 2020/09/15 09:32:13 UTC
[GitHub] [airflow] ashb opened a new pull request #10949: Allow CeleryExecutor to "adopt" an orphaned queued or running task
ashb opened a new pull request #10949:
URL: https://github.com/apache/airflow/pull/10949
This can happen when a task is enqueued by one executor, and then that
scheduler dies/exits.
The default fallback behaviour is unchanged -- that queued tasks are
cleared and then and then later rescheduled.
But for Celery we can do better -- if we record the Celery-generated
task_id, we can then re-create the AsyncResult objects for orphaned
tasks at a later date.
However since Celery just reports all AsyncResult as "PENDING", even if
they aren't tasks currently in the broker queue, we need to apply a
timeout to "unblock" these tasks in case they never actually made it to
the Celery broker.
This all means that we can adopt tasks that have been enqueued another
CeleryExecutor if it dies, without having to clear the task and slow
down. This is especially useful as the task may have already started
running, and while clearing it would stop it, it's better if we don't
have to reset it!
<!--
Thank you for contributing! Please make sure that your code changes
are covered with tests. And in case of new features or big changes
remember to adjust the documentation.
Feel free to ping committers for the review!
In case of existing issue, reference it using one of the following:
closes: #ISSUE
related: #ISSUE
How to write a good git commit message:
http://chris.beams.io/posts/git-commit/
-->
---
**^ Add meaningful description above**
Read the **[Pull Request Guidelines](https://github.com/apache/airflow/blob/master/CONTRIBUTING.rst#pull-request-guidelines)** for more information.
In case of fundamental code change, Airflow Improvement Proposal ([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvements+Proposals)) is needed.
In case of a new dependency, check compliance with the [ASF 3rd Party License Policy](https://www.apache.org/legal/resolved.html#category-x).
In case of backwards incompatible changes please leave a note in [UPDATING.md](https://github.com/apache/airflow/blob/master/UPDATING.md).
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] ashb commented on a change in pull request #10949: Allow CeleryExecutor to "adopt" an orphaned queued or running task
Posted by GitBox <gi...@apache.org>.
ashb commented on a change in pull request #10949:
URL: https://github.com/apache/airflow/pull/10949#discussion_r489542363
##########
File path: airflow/executors/celery_executor.py
##########
@@ -274,6 +331,61 @@ def execute_async(self,
def terminate(self):
pass
+ def try_adopt_task_instances(self, tis: List[TaskInstance]) -> List[TaskInstance]:
+ # See which of the TIs are still alive (or have finished even!)
+ #
+ # Since Celery doesn't store "SENT" state for queued commands (if we create an AsyncResult with a made
+ # up id it just returns PENDING state for it), we have to store Celery's task_id against the TI row to
+ # look at in future.
+ #
+ # This process is not perfect -- we could have sent the task to celery, and crashed before we were
+ # able to record the AsyncResult.task_id in the TaskInstance table, in which case we won't adopt the
+ # task (it'll either run and update the TI state, or the scheduler will clear and re-queue it. Either
+ # way it won't get executed more than once)
+ #
+ # (If we swapped it around, and generated a task_id for Celery, stored that in TI and enqueued that
+ # there is also still a race condition where we could generate and store the task_id, but die before
+ # we managed to enqueue the command. Since neither way is perfect we always have to deal with this
+ # process not being perfect.)
+
+ celery_tasks = {
+ ti.external_executor_id: (AsyncResult(ti.external_executor_id), ti)
+ for ti in tis
+ if ti.external_executor_id is not None
+ }
+
+ if not celery_tasks:
+ # Nothing to adopt
+ return tis
+
+ states_by_celery_task_id = self.bulk_state_fetcher.get_many(
+ map(operator.itemgetter(0), celery_tasks.values())
+ )
+
+ not_adopted = [ti for ti in tis if ti.external_executor_id is None]
+ adopted = []
+ cached_celery_backend = next(iter(celery_tasks.values()))[0].backend
+
+ for celery_task_id, (state, info) in states_by_celery_task_id.items():
+ result, ti = celery_tasks[celery_task_id]
+ result.backend = cached_celery_backend
+ result.backend = cached_celery_backend
+
+ # Set the correct elements of the state dicts, then update this
+ # like we just queried it.
+ self.adopted_task_timeouts[ti.key] = ti.queued_dttm + self.adoption_timeout
+ self.tasks[ti.key] = result
+ self.running.add(ti.key)
+ self.update_task_state(ti.key, state, info)
+ adopted.append(repr(ti) + " in state " + state)
Review comment:
:+1:
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] turbaszek commented on a change in pull request #10949: Allow CeleryExecutor to "adopt" an orphaned queued or running task
Posted by GitBox <gi...@apache.org>.
turbaszek commented on a change in pull request #10949:
URL: https://github.com/apache/airflow/pull/10949#discussion_r489508144
##########
File path: airflow/executors/celery_executor.py
##########
@@ -274,6 +331,61 @@ def execute_async(self,
def terminate(self):
pass
+ def try_adopt_task_instances(self, tis: List[TaskInstance]) -> List[TaskInstance]:
+ # See which of the TIs are still alive (or have finished even!)
+ #
+ # Since Celery doesn't store "SENT" state for queued commands (if we create an AsyncResult with a made
+ # up id it just returns PENDING state for it), we have to store Celery's task_id against the TI row to
+ # look at in future.
+ #
+ # This process is not perfect -- we could have sent the task to celery, and crashed before we were
+ # able to record the AsyncResult.task_id in the TaskInstance table, in which case we won't adopt the
+ # task (it'll either run and update the TI state, or the scheduler will clear and re-queue it. Either
+ # way it won't get executed more than once)
+ #
+ # (If we swapped it around, and generated a task_id for Celery, stored that in TI and enqueued that
+ # there is also still a race condition where we could generate and store the task_id, but die before
+ # we managed to enqueue the command. Since neither way is perfect we always have to deal with this
+ # process not being perfect.)
+
+ celery_tasks = {
+ ti.external_executor_id: (AsyncResult(ti.external_executor_id), ti)
+ for ti in tis
+ if ti.external_executor_id is not None
+ }
Review comment:
Not a big thing but we can do one `for` loop instead of this dict comprehension and list in L365, right?
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] XD-DENG commented on a change in pull request #10949: Allow CeleryExecutor to "adopt" an orphaned queued or running task
Posted by GitBox <gi...@apache.org>.
XD-DENG commented on a change in pull request #10949:
URL: https://github.com/apache/airflow/pull/10949#discussion_r488932380
##########
File path: airflow/migrations/versions/e1a11ece99cc_add_external_executor_id_to_ti.py
##########
@@ -0,0 +1,46 @@
+#
+# 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.
+
+"""Add external executor ID to TI
+
+Revision ID: e1a11ece99cc
+Revises: b247b1e3d1ed
+Create Date: 2020-09-12 08:23:45.698865
+
+"""
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = 'e1a11ece99cc'
+down_revision = 'b247b1e3d1ed'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+ """Apply Add external executor ID to TI"""
+ with op.batch_alter_table('task_instance', schema=None) as batch_op:
+ batch_op.add_column(sa.Column('external_executor_id', sa.String(length=250), nullable=True))
Review comment:
`external_executor_id` is the [Celery task's UUID](https://docs.celeryproject.org/en/stable/reference/celery.result.html#celery.result.AsyncResult.id) if I'm not wrong. Is 250 too big for it?
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] XD-DENG commented on a change in pull request #10949: Allow CeleryExecutor to "adopt" an orphaned queued or running task
Posted by GitBox <gi...@apache.org>.
XD-DENG commented on a change in pull request #10949:
URL: https://github.com/apache/airflow/pull/10949#discussion_r489571015
##########
File path: airflow/executors/celery_executor.py
##########
@@ -142,7 +146,11 @@ def __init__(self):
self._sync_parallelism = max(1, cpu_count() - 1)
self.bulk_state_fetcher = BulkStateFetcher(self._sync_parallelism)
self.tasks = {}
- self.last_state = {}
+ # Mapping of tasks we've adopted, ordered by the earliest date they timeout
+ self.adopted_task_timeouts: Dict[TaskInstanceKey, datetime.datetime] = OrderedDict()
+ self.adoption_timeout = datetime.timedelta(
+ seconds=conf.getint('celery', 'adopted_task_timeout', fallback=600)
Review comment:
I leave this to an English speaker 😉 I will be happy as long as they distinguish well from each other.
But my preference goes for `adopted_task_timeout`-> `task_adoption_timeout`
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] XD-DENG commented on a change in pull request #10949: Allow CeleryExecutor to "adopt" an orphaned queued or running task
Posted by GitBox <gi...@apache.org>.
XD-DENG commented on a change in pull request #10949:
URL: https://github.com/apache/airflow/pull/10949#discussion_r488925717
##########
File path: airflow/executors/celery_executor.py
##########
@@ -274,6 +330,61 @@ def execute_async(self,
def terminate(self):
pass
+ def try_adopt_task_instances(self, tis: List[TaskInstance]) -> List[TaskInstance]:
+ # See which if the tis are still alive (or have finished even!)
+ #
+ # Since Celery doesn't store "SENT" state for queued commands (if we create an AsyncResult with a made
+ # up id it just returns PENDING state for it), we have to store Celery's task_id against the TI row to
+ # look at in future.
+ #
+ # This process is not perfect -- we could have sent the task to celery, and crashed before we were
+ # able to record the AsyncResult.task_id in the TaskInstance table, in which case we won't adopt the
+ # task (it'll either run and update the TI state, or the scheduler will clear and re-queue it. Either
+ # way it won't get executed more than once)
+ #
+ # (If we swapped it around, and generated a task_id for Celery, stored that in TI and enqueued that
+ # there is also still a race condition where we could generate and store the task_id, but die before
+ # we managed to enqueue the command. Since neither way is perfect we always have to deal with this
+ # process not being perfect.)
+
+ celery_tasks = {
+ ti.external_executor_id: AsyncResult(ti.external_executor_id)
+ for ti in tis
+ if ti.external_executor_id is not None
+ }
+
+ if not celery_tasks:
+ # Nothing to adopt
+ return tis
+
+ tis_by_celery_task_id = {ti.external_executor_id: ti for ti in tis}
+
+ states_by_celery_task_id = self.bulk_state_fetcher.get_many(celery_tasks.values())
+
+ not_adopted = [ti for ti in tis if ti.external_executor_id is None]
+ adopted = []
+ cached_celery_backend = next(iter(celery_tasks.values())).backend
+
+ for celery_task_id, (state, info) in states_by_celery_task_id.items():
+ ti = tis_by_celery_task_id[celery_task_id]
+ result = celery_tasks[celery_task_id]
+ result.backend = cached_celery_backend
Review comment:
Maybe we don't need to have `tis_by_celery_task_id = {ti.external_executor_id: ti for ti in tis}` (line 360) separately, instead, we can have
```python
celery_tasks = {
ti.external_executor_id: (AsyncResult(ti.external_executor_id), ti)
for ti in tis
if ti.external_executor_id is not None
}
```
Then lines 369-370 can also be merged into
```python
result, ti = celery_tasks[celery_task_id]
```
By doing this, we can avoid traversing `tis` for two times (only one time would suffice), also save one query on the dict.
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] ashb commented on a change in pull request #10949: Allow CeleryExecutor to "adopt" an orphaned queued or running task
Posted by GitBox <gi...@apache.org>.
ashb commented on a change in pull request #10949:
URL: https://github.com/apache/airflow/pull/10949#discussion_r489537937
##########
File path: airflow/utils/sqlalchemy.py
##########
@@ -118,3 +120,27 @@ def process_result_value(self, value, dialect):
type_map = {key.__name__: key for key in self.attr_keys}
return type_map[data['type']](**data['attrs'])
return data
+
+
+def skip_locked(session: Session, of: Any = None) -> Dict[str, Any]:
Review comment:
Oh, right. That is mirroring the SQLA argument https://docs.sqlalchemy.org/en/13/orm/query.html?highlight=with_for_update#sqlalchemy.orm.query.Query.with_for_update:
> sqlalchemy.orm.query.Query.with_for_update(read=False, nowait=False, of=None, skip_locked=False, key_share=False)
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] ashb edited a comment on pull request #10949: Allow CeleryExecutor to "adopt" an orphaned queued or running task
Posted by GitBox <gi...@apache.org>.
ashb edited a comment on pull request #10949:
URL: https://github.com/apache/airflow/pull/10949#issuecomment-693294141
The doc tests are now failing with:
> /opt/airflow/docs/_api/airflow/executors/base_executor/index.rst:198:more than one target found for cross-reference 'TaskInstance': airflow.models.TaskInstance, airflow.models.taskinstance.TaskInstance
But I didn't change the imports in that file :/
Oh I missed Kaxil's earlier fix in my rebase+force-push
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] XD-DENG commented on a change in pull request #10949: Allow CeleryExecutor to "adopt" an orphaned queued or running task
Posted by GitBox <gi...@apache.org>.
XD-DENG commented on a change in pull request #10949:
URL: https://github.com/apache/airflow/pull/10949#discussion_r488929900
##########
File path: airflow/executors/celery_executor.py
##########
@@ -274,6 +330,61 @@ def execute_async(self,
def terminate(self):
pass
+ def try_adopt_task_instances(self, tis: List[TaskInstance]) -> List[TaskInstance]:
+ # See which if the tis are still alive (or have finished even!)
Review comment:
`see which if the tis...` may contain a typo?
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] ashb commented on a change in pull request #10949: Allow CeleryExecutor to "adopt" an orphaned queued or running task
Posted by GitBox <gi...@apache.org>.
ashb commented on a change in pull request #10949:
URL: https://github.com/apache/airflow/pull/10949#discussion_r489522276
##########
File path: airflow/utils/sqlalchemy.py
##########
@@ -118,3 +120,27 @@ def process_result_value(self, value, dialect):
type_map = {key.__name__: key for key in self.attr_keys}
return type_map[data['type']](**data['attrs'])
return data
+
+
+def skip_locked(session: Session, of: Any = None) -> Dict[str, Any]:
Review comment:
I wasn't aware of any more concrete type -- it's "any SQLA Model class". I'll see what I can find.
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] XD-DENG commented on a change in pull request #10949: Allow CeleryExecutor to "adopt" an orphaned queued or running task
Posted by GitBox <gi...@apache.org>.
XD-DENG commented on a change in pull request #10949:
URL: https://github.com/apache/airflow/pull/10949#discussion_r488925717
##########
File path: airflow/executors/celery_executor.py
##########
@@ -274,6 +330,61 @@ def execute_async(self,
def terminate(self):
pass
+ def try_adopt_task_instances(self, tis: List[TaskInstance]) -> List[TaskInstance]:
+ # See which if the tis are still alive (or have finished even!)
+ #
+ # Since Celery doesn't store "SENT" state for queued commands (if we create an AsyncResult with a made
+ # up id it just returns PENDING state for it), we have to store Celery's task_id against the TI row to
+ # look at in future.
+ #
+ # This process is not perfect -- we could have sent the task to celery, and crashed before we were
+ # able to record the AsyncResult.task_id in the TaskInstance table, in which case we won't adopt the
+ # task (it'll either run and update the TI state, or the scheduler will clear and re-queue it. Either
+ # way it won't get executed more than once)
+ #
+ # (If we swapped it around, and generated a task_id for Celery, stored that in TI and enqueued that
+ # there is also still a race condition where we could generate and store the task_id, but die before
+ # we managed to enqueue the command. Since neither way is perfect we always have to deal with this
+ # process not being perfect.)
+
+ celery_tasks = {
+ ti.external_executor_id: AsyncResult(ti.external_executor_id)
+ for ti in tis
+ if ti.external_executor_id is not None
+ }
+
+ if not celery_tasks:
+ # Nothing to adopt
+ return tis
+
+ tis_by_celery_task_id = {ti.external_executor_id: ti for ti in tis}
+
+ states_by_celery_task_id = self.bulk_state_fetcher.get_many(celery_tasks.values())
+
+ not_adopted = [ti for ti in tis if ti.external_executor_id is None]
+ adopted = []
+ cached_celery_backend = next(iter(celery_tasks.values())).backend
+
+ for celery_task_id, (state, info) in states_by_celery_task_id.items():
+ ti = tis_by_celery_task_id[celery_task_id]
+ result = celery_tasks[celery_task_id]
+ result.backend = cached_celery_backend
Review comment:
Maybe we do not need to have `tis_by_celery_task_id = {ti.external_executor_id: ti for ti in tis}` (line 360) separately. Instead, we can have
```python
celery_tasks = {
ti.external_executor_id: (AsyncResult(ti.external_executor_id), ti)
for ti in tis
if ti.external_executor_id is not None
}
```
Then lines 369-370 can also be merged into
```python
result, ti = celery_tasks[celery_task_id]
```
By doing this, we can avoid traversing `tis` for two times (only one time would suffice), also save one query on the dict.
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] turbaszek commented on a change in pull request #10949: Allow CeleryExecutor to "adopt" an orphaned queued or running task
Posted by GitBox <gi...@apache.org>.
turbaszek commented on a change in pull request #10949:
URL: https://github.com/apache/airflow/pull/10949#discussion_r489517351
##########
File path: airflow/executors/celery_executor.py
##########
@@ -274,6 +331,61 @@ def execute_async(self,
def terminate(self):
pass
+ def try_adopt_task_instances(self, tis: List[TaskInstance]) -> List[TaskInstance]:
+ # See which of the TIs are still alive (or have finished even!)
+ #
+ # Since Celery doesn't store "SENT" state for queued commands (if we create an AsyncResult with a made
+ # up id it just returns PENDING state for it), we have to store Celery's task_id against the TI row to
+ # look at in future.
+ #
+ # This process is not perfect -- we could have sent the task to celery, and crashed before we were
+ # able to record the AsyncResult.task_id in the TaskInstance table, in which case we won't adopt the
+ # task (it'll either run and update the TI state, or the scheduler will clear and re-queue it. Either
+ # way it won't get executed more than once)
+ #
+ # (If we swapped it around, and generated a task_id for Celery, stored that in TI and enqueued that
+ # there is also still a race condition where we could generate and store the task_id, but die before
+ # we managed to enqueue the command. Since neither way is perfect we always have to deal with this
+ # process not being perfect.)
+
+ celery_tasks = {
+ ti.external_executor_id: (AsyncResult(ti.external_executor_id), ti)
+ for ti in tis
+ if ti.external_executor_id is not None
+ }
Review comment:
Yes, I don't think this changes performance
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] ashb commented on pull request #10949: Allow CeleryExecutor to "adopt" an orphaned queued or running task
Posted by GitBox <gi...@apache.org>.
ashb commented on pull request #10949:
URL: https://github.com/apache/airflow/pull/10949#issuecomment-693294141
The doc tests are now failing with:
> /opt/airflow/docs/_api/airflow/executors/base_executor/index.rst:198:more than one target found for cross-reference 'TaskInstance': airflow.models.TaskInstance, airflow.models.taskinstance.TaskInstance
But I didn't change the imports in that file :/
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] turbaszek commented on a change in pull request #10949: Allow CeleryExecutor to "adopt" an orphaned queued or running task
Posted by GitBox <gi...@apache.org>.
turbaszek commented on a change in pull request #10949:
URL: https://github.com/apache/airflow/pull/10949#discussion_r489513412
##########
File path: airflow/jobs/scheduler_job.py
##########
@@ -1557,6 +1560,10 @@ def _process_executor_events(self, simple_dag_bag: SimpleDagBag, session: Sessio
state, info = event_buffer.pop(buffer_key)
# TODO: should we fail RUNNING as well, as we do in Backfills?
Review comment:
Not related, but is this comment still valid?
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] ashb commented on a change in pull request #10949: Allow CeleryExecutor to "adopt" an orphaned queued or running task
Posted by GitBox <gi...@apache.org>.
ashb commented on a change in pull request #10949:
URL: https://github.com/apache/airflow/pull/10949#discussion_r489512615
##########
File path: airflow/executors/celery_executor.py
##########
@@ -274,6 +331,61 @@ def execute_async(self,
def terminate(self):
pass
+ def try_adopt_task_instances(self, tis: List[TaskInstance]) -> List[TaskInstance]:
+ # See which of the TIs are still alive (or have finished even!)
+ #
+ # Since Celery doesn't store "SENT" state for queued commands (if we create an AsyncResult with a made
+ # up id it just returns PENDING state for it), we have to store Celery's task_id against the TI row to
+ # look at in future.
+ #
+ # This process is not perfect -- we could have sent the task to celery, and crashed before we were
+ # able to record the AsyncResult.task_id in the TaskInstance table, in which case we won't adopt the
+ # task (it'll either run and update the TI state, or the scheduler will clear and re-queue it. Either
+ # way it won't get executed more than once)
+ #
+ # (If we swapped it around, and generated a task_id for Celery, stored that in TI and enqueued that
+ # there is also still a race condition where we could generate and store the task_id, but die before
+ # we managed to enqueue the command. Since neither way is perfect we always have to deal with this
+ # process not being perfect.)
+
+ celery_tasks = {
+ ti.external_executor_id: (AsyncResult(ti.external_executor_id), ti)
+ for ti in tis
+ if ti.external_executor_id is not None
+ }
Review comment:
We could, but they have different conditions (`if` and `if not`) - we'd need to do something like:
```python
celery_tasks = []
not_adopted = []
for ti in its:
if ti.external_executor_id:
celery_tasks.append(...)
else:
not_adopted.append(...)
```
Is that what you meant?
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] ashb commented on a change in pull request #10949: Allow CeleryExecutor to "adopt" an orphaned queued or running task
Posted by GitBox <gi...@apache.org>.
ashb commented on a change in pull request #10949:
URL: https://github.com/apache/airflow/pull/10949#discussion_r489544387
##########
File path: airflow/executors/celery_executor.py
##########
@@ -274,6 +331,61 @@ def execute_async(self,
def terminate(self):
pass
+ def try_adopt_task_instances(self, tis: List[TaskInstance]) -> List[TaskInstance]:
+ # See which of the TIs are still alive (or have finished even!)
+ #
+ # Since Celery doesn't store "SENT" state for queued commands (if we create an AsyncResult with a made
+ # up id it just returns PENDING state for it), we have to store Celery's task_id against the TI row to
+ # look at in future.
+ #
+ # This process is not perfect -- we could have sent the task to celery, and crashed before we were
+ # able to record the AsyncResult.task_id in the TaskInstance table, in which case we won't adopt the
+ # task (it'll either run and update the TI state, or the scheduler will clear and re-queue it. Either
+ # way it won't get executed more than once)
+ #
+ # (If we swapped it around, and generated a task_id for Celery, stored that in TI and enqueued that
+ # there is also still a race condition where we could generate and store the task_id, but die before
+ # we managed to enqueue the command. Since neither way is perfect we always have to deal with this
+ # process not being perfect.)
+
+ celery_tasks = {
+ ti.external_executor_id: (AsyncResult(ti.external_executor_id), ti)
+ for ti in tis
+ if ti.external_executor_id is not None
+ }
+
+ if not celery_tasks:
+ # Nothing to adopt
+ return tis
+
+ states_by_celery_task_id = self.bulk_state_fetcher.get_many(
+ map(operator.itemgetter(0), celery_tasks.values())
+ )
+
+ not_adopted = [ti for ti in tis if ti.external_executor_id is None]
+ adopted = []
+ cached_celery_backend = next(iter(celery_tasks.values()))[0].backend
+
+ for celery_task_id, (state, info) in states_by_celery_task_id.items():
+ result, ti = celery_tasks[celery_task_id]
+ result.backend = cached_celery_backend
+ result.backend = cached_celery_backend
+
+ # Set the correct elements of the state dicts, then update this
+ # like we just queried it.
+ self.adopted_task_timeouts[ti.key] = ti.queued_dttm + self.adoption_timeout
+ self.tasks[ti.key] = result
+ self.running.add(ti.key)
+ self.update_task_state(ti.key, state, info)
+ adopted.append(repr(ti) + " in state " + state)
Review comment:
Done in https://github.com/apache/airflow/pull/10949/commits/e4f915efb2b1bd4165d33425799f3bbffd88caa3
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] kaxil merged pull request #10949: Allow CeleryExecutor to "adopt" an orphaned queued or running task
Posted by GitBox <gi...@apache.org>.
kaxil merged pull request #10949:
URL: https://github.com/apache/airflow/pull/10949
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] ashb commented on a change in pull request #10949: Allow CeleryExecutor to "adopt" an orphaned queued or running task
Posted by GitBox <gi...@apache.org>.
ashb commented on a change in pull request #10949:
URL: https://github.com/apache/airflow/pull/10949#discussion_r489546690
##########
File path: airflow/jobs/scheduler_job.py
##########
@@ -1557,6 +1560,10 @@ def _process_executor_events(self, simple_dag_bag: SimpleDagBag, session: Sessio
state, info = event_buffer.pop(buffer_key)
# TODO: should we fail RUNNING as well, as we do in Backfills?
Review comment:
Not sure off the top of my head.
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] ashb commented on a change in pull request #10949: Allow CeleryExecutor to "adopt" an orphaned queued or running task
Posted by GitBox <gi...@apache.org>.
ashb commented on a change in pull request #10949:
URL: https://github.com/apache/airflow/pull/10949#discussion_r489272586
##########
File path: airflow/executors/celery_executor.py
##########
@@ -274,6 +330,61 @@ def execute_async(self,
def terminate(self):
pass
+ def try_adopt_task_instances(self, tis: List[TaskInstance]) -> List[TaskInstance]:
+ # See which if the tis are still alive (or have finished even!)
+ #
+ # Since Celery doesn't store "SENT" state for queued commands (if we create an AsyncResult with a made
+ # up id it just returns PENDING state for it), we have to store Celery's task_id against the TI row to
+ # look at in future.
+ #
+ # This process is not perfect -- we could have sent the task to celery, and crashed before we were
+ # able to record the AsyncResult.task_id in the TaskInstance table, in which case we won't adopt the
+ # task (it'll either run and update the TI state, or the scheduler will clear and re-queue it. Either
+ # way it won't get executed more than once)
+ #
+ # (If we swapped it around, and generated a task_id for Celery, stored that in TI and enqueued that
+ # there is also still a race condition where we could generate and store the task_id, but die before
+ # we managed to enqueue the command. Since neither way is perfect we always have to deal with this
+ # process not being perfect.)
+
+ celery_tasks = {
+ ti.external_executor_id: AsyncResult(ti.external_executor_id)
+ for ti in tis
+ if ti.external_executor_id is not None
+ }
+
+ if not celery_tasks:
+ # Nothing to adopt
+ return tis
+
+ tis_by_celery_task_id = {ti.external_executor_id: ti for ti in tis}
+
+ states_by_celery_task_id = self.bulk_state_fetcher.get_many(celery_tasks.values())
+
+ not_adopted = [ti for ti in tis if ti.external_executor_id is None]
+ adopted = []
+ cached_celery_backend = next(iter(celery_tasks.values())).backend
+
+ for celery_task_id, (state, info) in states_by_celery_task_id.items():
+ ti = tis_by_celery_task_id[celery_task_id]
+ result = celery_tasks[celery_task_id]
+ result.backend = cached_celery_backend
Review comment:
Nice!
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] XD-DENG commented on a change in pull request #10949: Allow CeleryExecutor to "adopt" an orphaned queued or running task
Posted by GitBox <gi...@apache.org>.
XD-DENG commented on a change in pull request #10949:
URL: https://github.com/apache/airflow/pull/10949#discussion_r488928231
##########
File path: airflow/executors/celery_executor.py
##########
@@ -274,6 +330,61 @@ def execute_async(self,
def terminate(self):
pass
+ def try_adopt_task_instances(self, tis: List[TaskInstance]) -> List[TaskInstance]:
+ # See which if the tis are still alive (or have finished even!)
+ #
+ # Since Celery doesn't store "SENT" state for queued commands (if we create an AsyncResult with a made
+ # up id it just returns PENDING state for it), we have to store Celery's task_id against the TI row to
+ # look at in future.
+ #
+ # This process is not perfect -- we could have sent the task to celery, and crashed before we were
+ # able to record the AsyncResult.task_id in the TaskInstance table, in which case we won't adopt the
+ # task (it'll either run and update the TI state, or the scheduler will clear and re-queue it. Either
+ # way it won't get executed more than once)
+ #
+ # (If we swapped it around, and generated a task_id for Celery, stored that in TI and enqueued that
+ # there is also still a race condition where we could generate and store the task_id, but die before
+ # we managed to enqueue the command. Since neither way is perfect we always have to deal with this
+ # process not being perfect.)
+
+ celery_tasks = {
+ ti.external_executor_id: AsyncResult(ti.external_executor_id)
+ for ti in tis
+ if ti.external_executor_id is not None
+ }
+
+ if not celery_tasks:
+ # Nothing to adopt
+ return tis
+
+ tis_by_celery_task_id = {ti.external_executor_id: ti for ti in tis}
+
+ states_by_celery_task_id = self.bulk_state_fetcher.get_many(celery_tasks.values())
+
+ not_adopted = [ti for ti in tis if ti.external_executor_id is None]
+ adopted = []
+ cached_celery_backend = next(iter(celery_tasks.values())).backend
+
+ for celery_task_id, (state, info) in states_by_celery_task_id.items():
+ ti = tis_by_celery_task_id[celery_task_id]
+ result = celery_tasks[celery_task_id]
+ result.backend = cached_celery_backend
Review comment:
In addition, `tis_by_celery_task_id` covers all `tis`, which is not necessary, because it will only be used to query `tis` which present in `celery_tasks` (essentially means `tis` whose `external_executor_id` is not `None`).
So my suggested revision above may save a little memory as well ;-)
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] turbaszek commented on a change in pull request #10949: Allow CeleryExecutor to "adopt" an orphaned queued or running task
Posted by GitBox <gi...@apache.org>.
turbaszek commented on a change in pull request #10949:
URL: https://github.com/apache/airflow/pull/10949#discussion_r489526061
##########
File path: airflow/utils/sqlalchemy.py
##########
@@ -118,3 +120,27 @@ def process_result_value(self, value, dialect):
type_map = {key.__name__: key for key in self.attr_keys}
return type_map[data['type']](**data['attrs'])
return data
+
+
+def skip_locked(session: Session, of: Any = None) -> Dict[str, Any]:
Review comment:
My concern is not about typing but about the name of argument, the `of` is not meaningful to me
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] ashb commented on pull request #10949: Allow CeleryExecutor to "adopt" an orphaned queued or running task
Posted by GitBox <gi...@apache.org>.
ashb commented on pull request #10949:
URL: https://github.com/apache/airflow/pull/10949#issuecomment-693528168
Right, final rebase done, will merge once tests are green.
(I wish Github had a button for that)
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] XD-DENG commented on a change in pull request #10949: Allow CeleryExecutor to "adopt" an orphaned queued or running task
Posted by GitBox <gi...@apache.org>.
XD-DENG commented on a change in pull request #10949:
URL: https://github.com/apache/airflow/pull/10949#discussion_r489567951
##########
File path: airflow/executors/celery_executor.py
##########
@@ -142,7 +146,11 @@ def __init__(self):
self._sync_parallelism = max(1, cpu_count() - 1)
self.bulk_state_fetcher = BulkStateFetcher(self._sync_parallelism)
self.tasks = {}
- self.last_state = {}
+ # Mapping of tasks we've adopted, ordered by the earliest date they timeout
+ self.adopted_task_timeouts: Dict[TaskInstanceKey, datetime.datetime] = OrderedDict()
+ self.adoption_timeout = datetime.timedelta(
+ seconds=conf.getint('celery', 'adopted_task_timeout', fallback=600)
Review comment:
Only one more minor comment relating to naming: `self. adopted_task_timeouts` and `adopted_task_timeout` may make this less comprehensive.
I understand they are linked via `self.adopted_task_timeouts[ti.key] = ti.queued_dttm + self.adoption_timeout`, but the extremely similar names made me think they are "directly" linked to each other.
I would suggest to rename either one.
Otherwise all look good to me.
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] turbaszek commented on a change in pull request #10949: Allow CeleryExecutor to "adopt" an orphaned queued or running task
Posted by GitBox <gi...@apache.org>.
turbaszek commented on a change in pull request #10949:
URL: https://github.com/apache/airflow/pull/10949#discussion_r489510616
##########
File path: airflow/executors/celery_executor.py
##########
@@ -274,6 +331,61 @@ def execute_async(self,
def terminate(self):
pass
+ def try_adopt_task_instances(self, tis: List[TaskInstance]) -> List[TaskInstance]:
+ # See which of the TIs are still alive (or have finished even!)
+ #
+ # Since Celery doesn't store "SENT" state for queued commands (if we create an AsyncResult with a made
+ # up id it just returns PENDING state for it), we have to store Celery's task_id against the TI row to
+ # look at in future.
+ #
+ # This process is not perfect -- we could have sent the task to celery, and crashed before we were
+ # able to record the AsyncResult.task_id in the TaskInstance table, in which case we won't adopt the
+ # task (it'll either run and update the TI state, or the scheduler will clear and re-queue it. Either
+ # way it won't get executed more than once)
+ #
+ # (If we swapped it around, and generated a task_id for Celery, stored that in TI and enqueued that
+ # there is also still a race condition where we could generate and store the task_id, but die before
+ # we managed to enqueue the command. Since neither way is perfect we always have to deal with this
+ # process not being perfect.)
+
+ celery_tasks = {
+ ti.external_executor_id: (AsyncResult(ti.external_executor_id), ti)
+ for ti in tis
+ if ti.external_executor_id is not None
+ }
+
+ if not celery_tasks:
+ # Nothing to adopt
+ return tis
+
+ states_by_celery_task_id = self.bulk_state_fetcher.get_many(
+ map(operator.itemgetter(0), celery_tasks.values())
+ )
+
+ not_adopted = [ti for ti in tis if ti.external_executor_id is None]
+ adopted = []
+ cached_celery_backend = next(iter(celery_tasks.values()))[0].backend
+
+ for celery_task_id, (state, info) in states_by_celery_task_id.items():
+ result, ti = celery_tasks[celery_task_id]
+ result.backend = cached_celery_backend
+ result.backend = cached_celery_backend
Review comment:
```suggestion
result.backend = cached_celery_backend
```
Duplicated line
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] ashb commented on a change in pull request #10949: Allow CeleryExecutor to "adopt" an orphaned queued or running task
Posted by GitBox <gi...@apache.org>.
ashb commented on a change in pull request #10949:
URL: https://github.com/apache/airflow/pull/10949#discussion_r489280171
##########
File path: airflow/executors/celery_executor.py
##########
@@ -274,6 +330,61 @@ def execute_async(self,
def terminate(self):
pass
+ def try_adopt_task_instances(self, tis: List[TaskInstance]) -> List[TaskInstance]:
+ # See which if the tis are still alive (or have finished even!)
+ #
+ # Since Celery doesn't store "SENT" state for queued commands (if we create an AsyncResult with a made
+ # up id it just returns PENDING state for it), we have to store Celery's task_id against the TI row to
+ # look at in future.
+ #
+ # This process is not perfect -- we could have sent the task to celery, and crashed before we were
+ # able to record the AsyncResult.task_id in the TaskInstance table, in which case we won't adopt the
+ # task (it'll either run and update the TI state, or the scheduler will clear and re-queue it. Either
+ # way it won't get executed more than once)
+ #
+ # (If we swapped it around, and generated a task_id for Celery, stored that in TI and enqueued that
+ # there is also still a race condition where we could generate and store the task_id, but die before
+ # we managed to enqueue the command. Since neither way is perfect we always have to deal with this
+ # process not being perfect.)
+
+ celery_tasks = {
+ ti.external_executor_id: AsyncResult(ti.external_executor_id)
+ for ti in tis
+ if ti.external_executor_id is not None
+ }
+
+ if not celery_tasks:
+ # Nothing to adopt
+ return tis
+
+ tis_by_celery_task_id = {ti.external_executor_id: ti for ti in tis}
+
+ states_by_celery_task_id = self.bulk_state_fetcher.get_many(celery_tasks.values())
+
+ not_adopted = [ti for ti in tis if ti.external_executor_id is None]
+ adopted = []
+ cached_celery_backend = next(iter(celery_tasks.values())).backend
+
+ for celery_task_id, (state, info) in states_by_celery_task_id.items():
+ ti = tis_by_celery_task_id[celery_task_id]
+ result = celery_tasks[celery_task_id]
+ result.backend = cached_celery_backend
Review comment:
Done.
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] turbaszek commented on a change in pull request #10949: Allow CeleryExecutor to "adopt" an orphaned queued or running task
Posted by GitBox <gi...@apache.org>.
turbaszek commented on a change in pull request #10949:
URL: https://github.com/apache/airflow/pull/10949#discussion_r489511497
##########
File path: airflow/executors/celery_executor.py
##########
@@ -274,6 +331,61 @@ def execute_async(self,
def terminate(self):
pass
+ def try_adopt_task_instances(self, tis: List[TaskInstance]) -> List[TaskInstance]:
+ # See which of the TIs are still alive (or have finished even!)
+ #
+ # Since Celery doesn't store "SENT" state for queued commands (if we create an AsyncResult with a made
+ # up id it just returns PENDING state for it), we have to store Celery's task_id against the TI row to
+ # look at in future.
+ #
+ # This process is not perfect -- we could have sent the task to celery, and crashed before we were
+ # able to record the AsyncResult.task_id in the TaskInstance table, in which case we won't adopt the
+ # task (it'll either run and update the TI state, or the scheduler will clear and re-queue it. Either
+ # way it won't get executed more than once)
+ #
+ # (If we swapped it around, and generated a task_id for Celery, stored that in TI and enqueued that
+ # there is also still a race condition where we could generate and store the task_id, but die before
+ # we managed to enqueue the command. Since neither way is perfect we always have to deal with this
+ # process not being perfect.)
+
+ celery_tasks = {
+ ti.external_executor_id: (AsyncResult(ti.external_executor_id), ti)
+ for ti in tis
+ if ti.external_executor_id is not None
+ }
+
+ if not celery_tasks:
+ # Nothing to adopt
+ return tis
+
+ states_by_celery_task_id = self.bulk_state_fetcher.get_many(
+ map(operator.itemgetter(0), celery_tasks.values())
+ )
+
+ not_adopted = [ti for ti in tis if ti.external_executor_id is None]
+ adopted = []
+ cached_celery_backend = next(iter(celery_tasks.values()))[0].backend
+
+ for celery_task_id, (state, info) in states_by_celery_task_id.items():
+ result, ti = celery_tasks[celery_task_id]
+ result.backend = cached_celery_backend
+ result.backend = cached_celery_backend
+
+ # Set the correct elements of the state dicts, then update this
+ # like we just queried it.
+ self.adopted_task_timeouts[ti.key] = ti.queued_dttm + self.adoption_timeout
+ self.tasks[ti.key] = result
+ self.running.add(ti.key)
+ self.update_task_state(ti.key, state, info)
+ adopted.append(repr(ti) + " in state " + state)
Review comment:
```suggestion
adopted.append(f"{ti} in state {state}")
```
How about this?
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] XD-DENG commented on a change in pull request #10949: Allow CeleryExecutor to "adopt" an orphaned queued or running task
Posted by GitBox <gi...@apache.org>.
XD-DENG commented on a change in pull request #10949:
URL: https://github.com/apache/airflow/pull/10949#discussion_r488925717
##########
File path: airflow/executors/celery_executor.py
##########
@@ -274,6 +330,61 @@ def execute_async(self,
def terminate(self):
pass
+ def try_adopt_task_instances(self, tis: List[TaskInstance]) -> List[TaskInstance]:
+ # See which if the tis are still alive (or have finished even!)
+ #
+ # Since Celery doesn't store "SENT" state for queued commands (if we create an AsyncResult with a made
+ # up id it just returns PENDING state for it), we have to store Celery's task_id against the TI row to
+ # look at in future.
+ #
+ # This process is not perfect -- we could have sent the task to celery, and crashed before we were
+ # able to record the AsyncResult.task_id in the TaskInstance table, in which case we won't adopt the
+ # task (it'll either run and update the TI state, or the scheduler will clear and re-queue it. Either
+ # way it won't get executed more than once)
+ #
+ # (If we swapped it around, and generated a task_id for Celery, stored that in TI and enqueued that
+ # there is also still a race condition where we could generate and store the task_id, but die before
+ # we managed to enqueue the command. Since neither way is perfect we always have to deal with this
+ # process not being perfect.)
+
+ celery_tasks = {
+ ti.external_executor_id: AsyncResult(ti.external_executor_id)
+ for ti in tis
+ if ti.external_executor_id is not None
+ }
+
+ if not celery_tasks:
+ # Nothing to adopt
+ return tis
+
+ tis_by_celery_task_id = {ti.external_executor_id: ti for ti in tis}
+
+ states_by_celery_task_id = self.bulk_state_fetcher.get_many(celery_tasks.values())
+
+ not_adopted = [ti for ti in tis if ti.external_executor_id is None]
+ adopted = []
+ cached_celery_backend = next(iter(celery_tasks.values())).backend
+
+ for celery_task_id, (state, info) in states_by_celery_task_id.items():
+ ti = tis_by_celery_task_id[celery_task_id]
+ result = celery_tasks[celery_task_id]
+ result.backend = cached_celery_backend
Review comment:
Maybe we don't need to have `tis_by_celery_task_id = {ti.external_executor_id: ti for ti in tis}` (line 360) separately, instead, we can have
```python
celery_tasks = {
ti.external_executor_id: (AsyncResult(ti.external_executor_id), ti)
for ti in tis
if ti.external_executor_id is not None
}
```
Then lines 369-370 can also be merged into `result, ti = celery_tasks[celery_task_id]`
By doing this, we can avoid traversing `tis` two times, also save one query on the dict.
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] ashb commented on a change in pull request #10949: Allow CeleryExecutor to "adopt" an orphaned queued or running task
Posted by GitBox <gi...@apache.org>.
ashb commented on a change in pull request #10949:
URL: https://github.com/apache/airflow/pull/10949#discussion_r489539844
##########
File path: airflow/utils/sqlalchemy.py
##########
@@ -118,3 +120,27 @@ def process_result_value(self, value, dialect):
type_map = {key.__name__: key for key in self.attr_keys}
return type_map[data['type']](**data['attrs'])
return data
+
+
+def skip_locked(session: Session, of: Any = None) -> Dict[str, Any]:
Review comment:
Oh wait, I see a way to remove this option from here:
`.with_for_update(**skip_locked(of=DagRun, session=session)` -> `with_for_update(of=DagRun, **skip_locked(session))`
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] ashb commented on a change in pull request #10949: Allow CeleryExecutor to "adopt" an orphaned queued or running task
Posted by GitBox <gi...@apache.org>.
ashb commented on a change in pull request #10949:
URL: https://github.com/apache/airflow/pull/10949#discussion_r489569069
##########
File path: airflow/executors/celery_executor.py
##########
@@ -142,7 +146,11 @@ def __init__(self):
self._sync_parallelism = max(1, cpu_count() - 1)
self.bulk_state_fetcher = BulkStateFetcher(self._sync_parallelism)
self.tasks = {}
- self.last_state = {}
+ # Mapping of tasks we've adopted, ordered by the earliest date they timeout
+ self.adopted_task_timeouts: Dict[TaskInstanceKey, datetime.datetime] = OrderedDict()
+ self.adoption_timeout = datetime.timedelta(
+ seconds=conf.getint('celery', 'adopted_task_timeout', fallback=600)
Review comment:
Or perhaps `adopted_task_timeout` to `task_adoption_timeout` (and leave `adopted_task_timeouts`)
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] ashb commented on a change in pull request #10949: Allow CeleryExecutor to "adopt" an orphaned queued or running task
Posted by GitBox <gi...@apache.org>.
ashb commented on a change in pull request #10949:
URL: https://github.com/apache/airflow/pull/10949#discussion_r488937450
##########
File path: airflow/migrations/versions/e1a11ece99cc_add_external_executor_id_to_ti.py
##########
@@ -0,0 +1,46 @@
+#
+# 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.
+
+"""Add external executor ID to TI
+
+Revision ID: e1a11ece99cc
+Revises: b247b1e3d1ed
+Create Date: 2020-09-12 08:23:45.698865
+
+"""
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = 'e1a11ece99cc'
+down_revision = 'b247b1e3d1ed'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+ """Apply Add external executor ID to TI"""
+ with op.batch_alter_table('task_instance', schema=None) as batch_op:
+ batch_op.add_column(sa.Column('external_executor_id', sa.String(length=250), nullable=True))
Review comment:
To big for Celery, but this gives a bit more freedom for other executors to store things in there.
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] XD-DENG commented on a change in pull request #10949: Allow CeleryExecutor to "adopt" an orphaned queued or running task
Posted by GitBox <gi...@apache.org>.
XD-DENG commented on a change in pull request #10949:
URL: https://github.com/apache/airflow/pull/10949#discussion_r488939428
##########
File path: airflow/migrations/versions/e1a11ece99cc_add_external_executor_id_to_ti.py
##########
@@ -0,0 +1,46 @@
+#
+# 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.
+
+"""Add external executor ID to TI
+
+Revision ID: e1a11ece99cc
+Revises: b247b1e3d1ed
+Create Date: 2020-09-12 08:23:45.698865
+
+"""
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = 'e1a11ece99cc'
+down_revision = 'b247b1e3d1ed'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+ """Apply Add external executor ID to TI"""
+ with op.batch_alter_table('task_instance', schema=None) as batch_op:
+ batch_op.add_column(sa.Column('external_executor_id', sa.String(length=250), nullable=True))
Review comment:
Get your point.
IMHO possibly somewhere around 100 would suffice given its nature as `id`. But feel free to ignore this one.
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] ashb commented on a change in pull request #10949: Allow CeleryExecutor to "adopt" an orphaned queued or running task
Posted by GitBox <gi...@apache.org>.
ashb commented on a change in pull request #10949:
URL: https://github.com/apache/airflow/pull/10949#discussion_r488937065
##########
File path: airflow/executors/celery_executor.py
##########
@@ -274,6 +330,61 @@ def execute_async(self,
def terminate(self):
pass
+ def try_adopt_task_instances(self, tis: List[TaskInstance]) -> List[TaskInstance]:
+ # See which if the tis are still alive (or have finished even!)
Review comment:
Yup, `which of` it should be
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] ashb commented on a change in pull request #10949: Allow CeleryExecutor to "adopt" an orphaned queued or running task
Posted by GitBox <gi...@apache.org>.
ashb commented on a change in pull request #10949:
URL: https://github.com/apache/airflow/pull/10949#discussion_r489544495
##########
File path: airflow/executors/celery_executor.py
##########
@@ -274,6 +331,61 @@ def execute_async(self,
def terminate(self):
pass
+ def try_adopt_task_instances(self, tis: List[TaskInstance]) -> List[TaskInstance]:
+ # See which of the TIs are still alive (or have finished even!)
+ #
+ # Since Celery doesn't store "SENT" state for queued commands (if we create an AsyncResult with a made
+ # up id it just returns PENDING state for it), we have to store Celery's task_id against the TI row to
+ # look at in future.
+ #
+ # This process is not perfect -- we could have sent the task to celery, and crashed before we were
+ # able to record the AsyncResult.task_id in the TaskInstance table, in which case we won't adopt the
+ # task (it'll either run and update the TI state, or the scheduler will clear and re-queue it. Either
+ # way it won't get executed more than once)
+ #
+ # (If we swapped it around, and generated a task_id for Celery, stored that in TI and enqueued that
+ # there is also still a race condition where we could generate and store the task_id, but die before
+ # we managed to enqueue the command. Since neither way is perfect we always have to deal with this
+ # process not being perfect.)
+
+ celery_tasks = {
+ ti.external_executor_id: (AsyncResult(ti.external_executor_id), ti)
+ for ti in tis
+ if ti.external_executor_id is not None
+ }
+
+ if not celery_tasks:
+ # Nothing to adopt
+ return tis
+
+ states_by_celery_task_id = self.bulk_state_fetcher.get_many(
+ map(operator.itemgetter(0), celery_tasks.values())
+ )
+
+ not_adopted = [ti for ti in tis if ti.external_executor_id is None]
+ adopted = []
+ cached_celery_backend = next(iter(celery_tasks.values()))[0].backend
+
+ for celery_task_id, (state, info) in states_by_celery_task_id.items():
+ result, ti = celery_tasks[celery_task_id]
+ result.backend = cached_celery_backend
+ result.backend = cached_celery_backend
Review comment:
https://github.com/apache/airflow/pull/10949/commits/e4f915efb2b1bd4165d33425799f3bbffd88caa3
##########
File path: airflow/executors/celery_executor.py
##########
@@ -274,6 +331,61 @@ def execute_async(self,
def terminate(self):
pass
+ def try_adopt_task_instances(self, tis: List[TaskInstance]) -> List[TaskInstance]:
+ # See which of the TIs are still alive (or have finished even!)
+ #
+ # Since Celery doesn't store "SENT" state for queued commands (if we create an AsyncResult with a made
+ # up id it just returns PENDING state for it), we have to store Celery's task_id against the TI row to
+ # look at in future.
+ #
+ # This process is not perfect -- we could have sent the task to celery, and crashed before we were
+ # able to record the AsyncResult.task_id in the TaskInstance table, in which case we won't adopt the
+ # task (it'll either run and update the TI state, or the scheduler will clear and re-queue it. Either
+ # way it won't get executed more than once)
+ #
+ # (If we swapped it around, and generated a task_id for Celery, stored that in TI and enqueued that
+ # there is also still a race condition where we could generate and store the task_id, but die before
+ # we managed to enqueue the command. Since neither way is perfect we always have to deal with this
+ # process not being perfect.)
+
+ celery_tasks = {
+ ti.external_executor_id: (AsyncResult(ti.external_executor_id), ti)
+ for ti in tis
+ if ti.external_executor_id is not None
+ }
Review comment:
https://github.com/apache/airflow/pull/10949/commits/e4f915efb2b1bd4165d33425799f3bbffd88caa3
##########
File path: airflow/config_templates/config.yml
##########
@@ -1373,6 +1373,23 @@
type: int
example: ~
default: "2"
+ - name: task_track_started
+ description: |
+ Celery task will report its status as 'started' when the task is executed by a worker.
+ This is used in Airflow to keep track of the running tasks and if a Scheduler is restarted
+ or run in HA mode, it can adopt the orphan tasks launched by previous SchedulerJob.
+ version_added: 2.0.0
+ type: boolean
+ example: ~
+ default: "True"
+ - name: adopted_task_timeout
+ description: |
+ Time after which Adopted tasks are cleared by CeleryExecutor. This is helpful to clear
Review comment:
https://github.com/apache/airflow/pull/10949/commits/e4f915efb2b1bd4165d33425799f3bbffd88caa3
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] XD-DENG commented on a change in pull request #10949: Allow CeleryExecutor to "adopt" an orphaned queued or running task
Posted by GitBox <gi...@apache.org>.
XD-DENG commented on a change in pull request #10949:
URL: https://github.com/apache/airflow/pull/10949#discussion_r489332072
##########
File path: airflow/executors/base_executor.py
##########
@@ -265,6 +265,20 @@ def terminate(self):
"""
raise NotImplementedError()
+ def try_adopt_task_instances(self, tis: List[TaskInstance]) -> List[TaskInstance]:
+ """
+ Try to adopt running task instances that have been abandoned by a SchedulerJob dying.
+
+ Anything that is not adopted will be cleared by the scheduler (and then become eligible for
+ re-scheduling)
+
+ :return: any TaskInstances that were unable to be adopted
+ :rtype: list[airflow.models.TaskInstance]
+ """
+ # By default, assume Executors cannot addopt tasks, so just say we failed to adopt anything.
Review comment:
Minor: another typo spotted, `addopt`.
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] ashb commented on a change in pull request #10949: Allow CeleryExecutor to "adopt" an orphaned queued or running task
Posted by GitBox <gi...@apache.org>.
ashb commented on a change in pull request #10949:
URL: https://github.com/apache/airflow/pull/10949#discussion_r489543818
##########
File path: airflow/utils/sqlalchemy.py
##########
@@ -118,3 +120,27 @@ def process_result_value(self, value, dialect):
type_map = {key.__name__: key for key in self.attr_keys}
return type_map[data['type']](**data['attrs'])
return data
+
+
+def skip_locked(session: Session, of: Any = None) -> Dict[str, Any]:
Review comment:
Fixed in https://github.com/apache/airflow/pull/10949/commits/3f4b2ea0c45541f462ad102646d685a74416aba2
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] XD-DENG commented on a change in pull request #10949: Allow CeleryExecutor to "adopt" an orphaned queued or running task
Posted by GitBox <gi...@apache.org>.
XD-DENG commented on a change in pull request #10949:
URL: https://github.com/apache/airflow/pull/10949#discussion_r489331721
##########
File path: airflow/migrations/versions/e1a11ece99cc_add_external_executor_id_to_ti.py
##########
@@ -0,0 +1,46 @@
+#
+# 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.
+
+"""Add external executor ID to TI
+
+Revision ID: e1a11ece99cc
+Revises: b247b1e3d1ed
+Create Date: 2020-09-12 08:23:45.698865
+
+"""
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = 'e1a11ece99cc'
+down_revision = 'b247b1e3d1ed'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+ """Apply Add external executor ID to TI"""
+ with op.batch_alter_table('task_instance', schema=None) as batch_op:
+ batch_op.add_column(sa.Column('external_executor_id', sa.String(length=250), nullable=True))
Review comment:
Just realise `ID_LEN` is specified in `airflow.models.base` as a constant. Yes, please ignore this comment
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] ashb commented on a change in pull request #10949: Allow CeleryExecutor to "adopt" an orphaned queued or running task
Posted by GitBox <gi...@apache.org>.
ashb commented on a change in pull request #10949:
URL: https://github.com/apache/airflow/pull/10949#discussion_r489568804
##########
File path: airflow/executors/celery_executor.py
##########
@@ -142,7 +146,11 @@ def __init__(self):
self._sync_parallelism = max(1, cpu_count() - 1)
self.bulk_state_fetcher = BulkStateFetcher(self._sync_parallelism)
self.tasks = {}
- self.last_state = {}
+ # Mapping of tasks we've adopted, ordered by the earliest date they timeout
+ self.adopted_task_timeouts: Dict[TaskInstanceKey, datetime.datetime] = OrderedDict()
+ self.adoption_timeout = datetime.timedelta(
+ seconds=conf.getint('celery', 'adopted_task_timeout', fallback=600)
Review comment:
Good shout. I'll rename `adopted_task_timeouts` to `adopted_task_by_timeout`? `abandon_adopted_tasks_after`?
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] turbaszek commented on a change in pull request #10949: Allow CeleryExecutor to "adopt" an orphaned queued or running task
Posted by GitBox <gi...@apache.org>.
turbaszek commented on a change in pull request #10949:
URL: https://github.com/apache/airflow/pull/10949#discussion_r489508144
##########
File path: airflow/executors/celery_executor.py
##########
@@ -274,6 +331,61 @@ def execute_async(self,
def terminate(self):
pass
+ def try_adopt_task_instances(self, tis: List[TaskInstance]) -> List[TaskInstance]:
+ # See which of the TIs are still alive (or have finished even!)
+ #
+ # Since Celery doesn't store "SENT" state for queued commands (if we create an AsyncResult with a made
+ # up id it just returns PENDING state for it), we have to store Celery's task_id against the TI row to
+ # look at in future.
+ #
+ # This process is not perfect -- we could have sent the task to celery, and crashed before we were
+ # able to record the AsyncResult.task_id in the TaskInstance table, in which case we won't adopt the
+ # task (it'll either run and update the TI state, or the scheduler will clear and re-queue it. Either
+ # way it won't get executed more than once)
+ #
+ # (If we swapped it around, and generated a task_id for Celery, stored that in TI and enqueued that
+ # there is also still a race condition where we could generate and store the task_id, but die before
+ # we managed to enqueue the command. Since neither way is perfect we always have to deal with this
+ # process not being perfect.)
+
+ celery_tasks = {
+ ti.external_executor_id: (AsyncResult(ti.external_executor_id), ti)
+ for ti in tis
+ if ti.external_executor_id is not None
+ }
Review comment:
Not a big thing but we can do one `for` loop instead of this list comprehension and list in L365, right?
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] turbaszek commented on a change in pull request #10949: Allow CeleryExecutor to "adopt" an orphaned queued or running task
Posted by GitBox <gi...@apache.org>.
turbaszek commented on a change in pull request #10949:
URL: https://github.com/apache/airflow/pull/10949#discussion_r489515666
##########
File path: airflow/utils/sqlalchemy.py
##########
@@ -118,3 +120,27 @@ def process_result_value(self, value, dialect):
type_map = {key.__name__: key for key in self.attr_keys}
return type_map[data['type']](**data['attrs'])
return data
+
+
+def skip_locked(session: Session, of: Any = None) -> Dict[str, Any]:
Review comment:
```suggestion
def skip_locked(session: Session, of: Any = None) -> Dict[str, Any]:
```
Can you describe the `of` or use something more informative? I'm not sure if I got it 😄
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] XD-DENG commented on a change in pull request #10949: Allow CeleryExecutor to "adopt" an orphaned queued or running task
Posted by GitBox <gi...@apache.org>.
XD-DENG commented on a change in pull request #10949:
URL: https://github.com/apache/airflow/pull/10949#discussion_r488925717
##########
File path: airflow/executors/celery_executor.py
##########
@@ -274,6 +330,61 @@ def execute_async(self,
def terminate(self):
pass
+ def try_adopt_task_instances(self, tis: List[TaskInstance]) -> List[TaskInstance]:
+ # See which if the tis are still alive (or have finished even!)
+ #
+ # Since Celery doesn't store "SENT" state for queued commands (if we create an AsyncResult with a made
+ # up id it just returns PENDING state for it), we have to store Celery's task_id against the TI row to
+ # look at in future.
+ #
+ # This process is not perfect -- we could have sent the task to celery, and crashed before we were
+ # able to record the AsyncResult.task_id in the TaskInstance table, in which case we won't adopt the
+ # task (it'll either run and update the TI state, or the scheduler will clear and re-queue it. Either
+ # way it won't get executed more than once)
+ #
+ # (If we swapped it around, and generated a task_id for Celery, stored that in TI and enqueued that
+ # there is also still a race condition where we could generate and store the task_id, but die before
+ # we managed to enqueue the command. Since neither way is perfect we always have to deal with this
+ # process not being perfect.)
+
+ celery_tasks = {
+ ti.external_executor_id: AsyncResult(ti.external_executor_id)
+ for ti in tis
+ if ti.external_executor_id is not None
+ }
+
+ if not celery_tasks:
+ # Nothing to adopt
+ return tis
+
+ tis_by_celery_task_id = {ti.external_executor_id: ti for ti in tis}
+
+ states_by_celery_task_id = self.bulk_state_fetcher.get_many(celery_tasks.values())
+
+ not_adopted = [ti for ti in tis if ti.external_executor_id is None]
+ adopted = []
+ cached_celery_backend = next(iter(celery_tasks.values())).backend
+
+ for celery_task_id, (state, info) in states_by_celery_task_id.items():
+ ti = tis_by_celery_task_id[celery_task_id]
+ result = celery_tasks[celery_task_id]
+ result.backend = cached_celery_backend
Review comment:
Maybe we don't need to have `tis_by_celery_task_id = {ti.external_executor_id: ti for ti in tis}` (line 360) separately, instead, we can have
```python
celery_tasks = {
ti.external_executor_id: (AsyncResult(ti.external_executor_id), ti)
for ti in tis
if ti.external_executor_id is not None
}
```
Then lines 369-370 can also be merged into `result, ti = celery_tasks[celery_task_id]`
By doing this, we can avoid traversing `tis` for two times (only one time would suffice), also save one query on the dict.
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] turbaszek commented on a change in pull request #10949: Allow CeleryExecutor to "adopt" an orphaned queued or running task
Posted by GitBox <gi...@apache.org>.
turbaszek commented on a change in pull request #10949:
URL: https://github.com/apache/airflow/pull/10949#discussion_r489503134
##########
File path: airflow/config_templates/config.yml
##########
@@ -1373,6 +1373,23 @@
type: int
example: ~
default: "2"
+ - name: task_track_started
+ description: |
+ Celery task will report its status as 'started' when the task is executed by a worker.
+ This is used in Airflow to keep track of the running tasks and if a Scheduler is restarted
+ or run in HA mode, it can adopt the orphan tasks launched by previous SchedulerJob.
+ version_added: 2.0.0
+ type: boolean
+ example: ~
+ default: "True"
+ - name: adopted_task_timeout
+ description: |
+ Time after which Adopted tasks are cleared by CeleryExecutor. This is helpful to clear
Review comment:
```suggestion
Time in seconds after which Adopted tasks are cleared by CeleryExecutor. This is helpful to clear
```
Just to clarify
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org