Ensure we use ti.queued_by_job_id

[atrbgithub]

Ensure that we use ti.queued_by_job_id when searching for pods.

When a task is adopted by a new scheduler, the id of the current task is used:

airflow/airflow/executors/kubernetes_executor.py

Lines 620 to 623 in feb6b81

	self.log.info("attempting to adopt pod %s", pod.metadata.name)
	pod.metadata.labels['airflow-worker'] = pod_generator.make_safe_label_value(
	str(self.scheduler_job_id)
	)

When this is successful the task instance is updated, and queued_by_job_id is updated with the id of the current scheduler:

airflow/airflow/jobs/scheduler_job.py

Lines 1877 to 1878 in feb6b81

	for ti in set(tis_to_reset_or_adopt) - set(to_reset):
	ti.queued_by_job_id = self.id

Therefore when we search for the pod labels on subsequent scheduler relaunches, we must search for pods using the queued_by_job_id and not external_executor_id, as we are currently doing:

airflow/airflow/executors/kubernetes_executor.py

Line 592 in feb6b81

scheduler_job_ids = [ti.external_executor_id for ti in tis]

external_executor_id is static and never appears to be updated when tasks are adopted.

This relates to #13808

[ashb]

Paging @dimberman

[samwedge]

It would be great to get a steer on whether this change is acceptable/safe, even if the PR isn't approved at this stage.

[eejbyfeldt]

The celery_executor also uses external_executor_id in its implementation of try_adopt_task_instances:

airflow/airflow/executors/celery_executor.py

Lines 469 to 470 in 2a2adb3

	if ti.external_executor_id is not None:
	celery_tasks[ti.external_executor_id] = (AsyncResult(ti.external_executor_id), ti)

Does it also experience these problems?

[ashb]

The usage in Celery is correct. ti.external_exeuctor_id there is the Celery Task ID (a UUID) and it's how we keep track of what the celery task ID is.

[ashb]

Okay, I understand the problem, but this doesn't quite feel like the right fix. I'm digging in a problem to look at what a better fix might be.

Additionally you will need to add unit tests the cover this case please.

[samwedge]

Hi @ashb,

Not sure we want to spend time adding unit tests if you don't feel this is the right fix.

Do we think a better solution would be to merge queued_by_job_id and external_executor_id? As far as I can see, they serve the same purpose.

[ashb]

Hi @samwedge There could be value in the unit tests anyway, as they should show the problem is fixed (i.e. that the task is not adopted tiwce/orphaned) without getting in to implementation details.

But it's also cool if you want to hold off on doing this for now.

[samwedge]

Thanks. I agree if we can get a nice generic test in place that doesn't have any implementation detail, then it's worth adding in.

That said, I've been looking at where to make the change. My first thought was in test_kubernetes_executor.py as this is the executor we have changed. But really, the bug only shows itself when calling SchedulerJob.adopt_or_reset_orphaned_tasks(). Adding a test here (in test_scheduler_job.py) means setting up a KubernetesExecutor with a lot of mocking/patching. And I'm not sure this is the best place for it.

Alternatively, do you think this is a candidate for a system test?

[samwedge]

@ashb Just wondered if you had any further thoughts on my previous message. In particular, where the unit test might live. It will need to test the integration between SchedulerJob and the KubernetesExecutor. I don't want to muddy the existing test files, which seem to test each in isolation.

[dimberman]

@samwedge are you able to reproduce the error in a system test? I think if you can make it fail then a system test should be sufficient here.

[samwedge]

@samwedge are you able to reproduce the error in a system test? I think if you can make it fail then a system test should be sufficient here.

Sorry for the silence @dimberman, trying to find some time to work on this. I'm fine with the unit and integration tests, but have never run the system tests before. I'll take a look and drop a message on Slack if I have any issues.

[ashb]

I was so wrong on this. I didn't realise that we are already re-setting ti.queued_by_job_id on adoption.

(Given I wrote that code I probably should do. But, well, 2020 was looooong.)

[kaxil]

@samwedge @atrbgithub Can we merge autotraderuk#2 (from @jedcunningham ) in your branch and rebase the PR on master please so that we can merge this PR

Stale review

[github-actions]

The PR most likely needs to run full matrix of tests because it modifies parts of the core of Airflow. However, committers might decide to merge it quickly and take the risk. If they don't merge it quickly - please rebase it to the latest master at your convenience, or amend the last commit of the PR, and push it with --force-with-lease.