From ed988ff145b524c2549fad14fa03950c15cbeaa5 Mon Sep 17 00:00:00 2001 From: Kevin Crouse Date: Wed, 22 Jun 2022 09:22:54 -0400 Subject: [PATCH 1/6] updating local_task_job to better handle cases in which the task instance pid is None --- airflow/jobs/local_task_job.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/airflow/jobs/local_task_job.py b/airflow/jobs/local_task_job.py index 5711342e04d97..67cfda5866e67 100644 --- a/airflow/jobs/local_task_job.py +++ b/airflow/jobs/local_task_job.py @@ -190,16 +190,20 @@ def heartbeat_callback(self, session=None): ti.hostname, fqdn, ) - raise AirflowException("Hostname of job runner does not match") + raise AirflowException("Hostname of job runner does not match") current_pid = self.task_runner.process.pid recorded_pid = ti.pid same_process = recorded_pid == current_pid - if ti.run_as_user or self.task_runner.run_as_user: + if recorded_pid is not None and ti.run_as_user or self.task_runner.run_as_user: + # when running as another user, compare the task runner pid to the parent of + # the recorded pid because user delegation becomes an extra process level. + + # However, if recorded_pid is None, pass that through as it signals the task runner process has already completed and been cleared out. `psutil.Process` uses the current process if the parameter is None, which is not what is intended for comparison. recorded_pid = psutil.Process(ti.pid).ppid() same_process = recorded_pid == current_pid - if recorded_pid is not None and not same_process: + if recorded_pid is not None and not same_process: self.log.warning( "Recorded pid %s does not match the current pid %s", recorded_pid, current_pid ) From bdf07df1b66edd5f8ea8165b76523a22c05e2c56 Mon Sep 17 00:00:00 2001 From: Kevin Crouse Date: Wed, 22 Jun 2022 09:52:20 -0400 Subject: [PATCH 2/6] remove whitespace --- airflow/jobs/local_task_job.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/airflow/jobs/local_task_job.py b/airflow/jobs/local_task_job.py index 67cfda5866e67..91617f6c5e678 100644 --- a/airflow/jobs/local_task_job.py +++ b/airflow/jobs/local_task_job.py @@ -190,12 +190,12 @@ def heartbeat_callback(self, session=None): ti.hostname, fqdn, ) - raise AirflowException("Hostname of job runner does not match") + raise AirflowException("Hostname of job runner does not match") current_pid = self.task_runner.process.pid recorded_pid = ti.pid same_process = recorded_pid == current_pid - if recorded_pid is not None and ti.run_as_user or self.task_runner.run_as_user: + if recorded_pid is not None and ti.run_as_user or self.task_runner.run_as_user: # when running as another user, compare the task runner pid to the parent of # the recorded pid because user delegation becomes an extra process level. @@ -203,7 +203,7 @@ def heartbeat_callback(self, session=None): recorded_pid = psutil.Process(ti.pid).ppid() same_process = recorded_pid == current_pid - if recorded_pid is not None and not same_process: + if recorded_pid is not None and not same_process: self.log.warning( "Recorded pid %s does not match the current pid %s", recorded_pid, current_pid ) From 7fe1c9a1afa61a8b20242fb766a8b81a7a52cd93 Mon Sep 17 00:00:00 2001 From: Kevin Crouse Date: Sun, 26 Jun 2022 04:43:00 -0400 Subject: [PATCH 3/6] updating comment to adhere to max line width --- airflow/jobs/local_task_job.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/airflow/jobs/local_task_job.py b/airflow/jobs/local_task_job.py index 91617f6c5e678..c77560a28fda1 100644 --- a/airflow/jobs/local_task_job.py +++ b/airflow/jobs/local_task_job.py @@ -198,8 +198,10 @@ def heartbeat_callback(self, session=None): if recorded_pid is not None and ti.run_as_user or self.task_runner.run_as_user: # when running as another user, compare the task runner pid to the parent of # the recorded pid because user delegation becomes an extra process level. - - # However, if recorded_pid is None, pass that through as it signals the task runner process has already completed and been cleared out. `psutil.Process` uses the current process if the parameter is None, which is not what is intended for comparison. + # However, if recorded_pid is None, pass that through as it signals the task + # runner process has already completed and been cleared out. `psutil.Process` + # uses the current process if the parameter is None, which is not what is intended + # for comparison. recorded_pid = psutil.Process(ti.pid).ppid() same_process = recorded_pid == current_pid From d4a44bda1b2cb1012075686efe0f026b317f4063 Mon Sep 17 00:00:00 2001 From: Kevin Crouse Date: Thu, 7 Jul 2022 02:56:26 -0400 Subject: [PATCH 4/6] adding parens to force order of operations --- airflow/jobs/local_task_job.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/jobs/local_task_job.py b/airflow/jobs/local_task_job.py index c77560a28fda1..3cb362f660efe 100644 --- a/airflow/jobs/local_task_job.py +++ b/airflow/jobs/local_task_job.py @@ -195,7 +195,7 @@ def heartbeat_callback(self, session=None): recorded_pid = ti.pid same_process = recorded_pid == current_pid - if recorded_pid is not None and ti.run_as_user or self.task_runner.run_as_user: + if recorded_pid is not None and (ti.run_as_user or self.task_runner.run_as_user): # when running as another user, compare the task runner pid to the parent of # the recorded pid because user delegation becomes an extra process level. # However, if recorded_pid is None, pass that through as it signals the task From 2ba30e36357c1c45dd6275c47bda64a9a494a001 Mon Sep 17 00:00:00 2001 From: Kevin Crouse Date: Thu, 7 Jul 2022 03:02:07 -0400 Subject: [PATCH 5/6] adding in tests for when the ti.pid is None --- tests/jobs/test_local_task_job.py | 30 ++++++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/tests/jobs/test_local_task_job.py b/tests/jobs/test_local_task_job.py index 362bba18287b2..74348f8f43645 100644 --- a/tests/jobs/test_local_task_job.py +++ b/tests/jobs/test_local_task_job.py @@ -154,6 +154,16 @@ def test_localtaskjob_heartbeat(self, dag_maker): with pytest.raises(AirflowException): job1.heartbeat_callback() + # Now, set the ti.pid to None and test that no error + # is raised. + ti.pid = None + session.merge(ti) + session.commit() + assert ti.pid != job1.task_runner.process.pid + assert not ti.run_as_user + assert not job1.task_runner.run_as_user + job1.heartbeat_callback() + @mock.patch('subprocess.check_call') @mock.patch('airflow.jobs.local_task_job.psutil') def test_localtaskjob_heartbeat_with_run_as_user(self, psutil_mock, _, dag_maker): @@ -196,6 +206,16 @@ def test_localtaskjob_heartbeat_with_run_as_user(self, psutil_mock, _, dag_maker with pytest.raises(AirflowException, match='PID of job runner does not match'): job1.heartbeat_callback() + # Here we set the ti.pid to None and test that no error is + # raised + ti.pid = None + session.merge(ti) + session.commit() + assert ti.run_as_user + assert job1.task_runner.run_as_user == ti.run_as_user + assert ti.pid != job1.task_runner.process.pid + job1.heartbeat_callback() + @conf_vars({('core', 'default_impersonation'): 'testuser'}) @mock.patch('subprocess.check_call') @mock.patch('airflow.jobs.local_task_job.psutil') @@ -239,6 +259,16 @@ def test_localtaskjob_heartbeat_with_default_impersonation(self, psutil_mock, _, with pytest.raises(AirflowException, match='PID of job runner does not match'): job1.heartbeat_callback() + # Now, set the ti.pid to None and test that no error + # is raised. + ti.pid = None + session.merge(ti) + session.commit() + assert job1.task_runner.run_as_user == 'testuser' + assert ti.run_as_user is None + assert ti.pid != job1.task_runner.process.pid + job1.heartbeat_callback() + def test_heartbeat_failed_fast(self): """ Test that task heartbeat will sleep when it fails fast From 1898966927e0b3cd8d78866026fb1e66ae0ade4f Mon Sep 17 00:00:00 2001 From: Kevin Crouse Date: Thu, 7 Jul 2022 03:06:23 -0400 Subject: [PATCH 6/6] formatting comment text --- airflow/jobs/local_task_job.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/airflow/jobs/local_task_job.py b/airflow/jobs/local_task_job.py index 3cb362f660efe..b73c8992d8a73 100644 --- a/airflow/jobs/local_task_job.py +++ b/airflow/jobs/local_task_job.py @@ -196,12 +196,12 @@ def heartbeat_callback(self, session=None): same_process = recorded_pid == current_pid if recorded_pid is not None and (ti.run_as_user or self.task_runner.run_as_user): - # when running as another user, compare the task runner pid to the parent of + # when running as another user, compare the task runner pid to the parent of # the recorded pid because user delegation becomes an extra process level. - # However, if recorded_pid is None, pass that through as it signals the task - # runner process has already completed and been cleared out. `psutil.Process` - # uses the current process if the parameter is None, which is not what is intended - # for comparison. + # However, if recorded_pid is None, pass that through as it signals the task + # runner process has already completed and been cleared out. `psutil.Process` + # uses the current process if the parameter is None, which is not what is intended + # for comparison. recorded_pid = psutil.Process(ti.pid).ppid() same_process = recorded_pid == current_pid