def test_start_and_terminate(self):
local_task_job = mock.Mock()
local_task_job.task_instance = mock.MagicMock()
local_task_job.task_instance.run_as_user = None
local_task_job.task_instance.command_as_list.return_value = [
'airflow', 'tasks', 'test', 'test_on_kill', 'task1', '2016-01-01'
]
runner = StandardTaskRunner(local_task_job)
runner.start()
time.sleep(0.5)
pgid = os.getpgid(runner.process.pid)
self.assertGreater(pgid, 0)
self.assertNotEqual(
pgid, os.getpgid(0),
"Task should be in a different process group to us")
processes = list(self._procs_in_pgroup(pgid))
runner.terminate()
for process in processes:
self.assertFalse(psutil.pid_exists(process.pid),
"{} is still alive".format(process))
self.assertIsNotNone(runner.return_code())
def test_start_and_terminate(self):
local_task_job = mock.Mock()
local_task_job.task_instance = mock.MagicMock()
local_task_job.task_instance.run_as_user = None
local_task_job.task_instance.command_as_list.return_value = [
'airflow',
'tasks',
'test',
'test_on_kill',
'task1',
'2016-01-01',
]
runner = StandardTaskRunner(local_task_job)
runner.start()
time.sleep(0.5)
pgid = os.getpgid(runner.process.pid)
assert pgid > 0
assert pgid != os.getpgid(
0), "Task should be in a different process group to us"
processes = list(self._procs_in_pgroup(pgid))
runner.terminate()
for process in processes:
assert not psutil.pid_exists(
process.pid), f"{process} is still alive"
assert runner.return_code() is not None
def test_on_kill(self):
"""
Test that ensures that clearing in the UI SIGTERMS
the task
"""
path = "/tmp/airflow_on_kill"
try:
os.unlink(path)
except OSError:
pass
dagbag = models.DagBag(
dag_folder=TEST_DAG_FOLDER,
include_examples=False,
)
dag = dagbag.dags.get('test_on_kill')
task = dag.get_task('task1')
session = settings.Session()
dag.clear()
dag.create_dagrun(
run_id="test",
state=State.RUNNING,
execution_date=DEFAULT_DATE,
start_date=DEFAULT_DATE,
session=session,
)
ti = TI(task=task, execution_date=DEFAULT_DATE)
job1 = LocalTaskJob(task_instance=ti, ignore_ti_state=True)
session.commit()
runner = StandardTaskRunner(job1)
runner.start()
# give the task some time to startup
time.sleep(3)
pgid = os.getpgid(runner.process.pid)
self.assertGreater(pgid, 0)
self.assertNotEqual(
pgid, os.getpgid(0),
"Task should be in a different process group to us")
processes = list(self._procs_in_pgroup(pgid))
runner.terminate()
# Wait some time for the result
for _ in range(20):
if os.path.exists(path):
break
time.sleep(2)
with open(path) as f:
self.assertEqual("ON_KILL_TEST", f.readline())
for process in processes:
self.assertFalse(psutil.pid_exists(process.pid),
f"{process} is still alive")
def test_early_reap_exit(self, caplog):
"""
Tests that when a child process running a task is killed externally
(e.g. by an OOM error, which we fake here), then we get return code
-9 and a log message.
"""
# Set up mock task
local_task_job = mock.Mock()
local_task_job.task_instance = mock.MagicMock()
local_task_job.task_instance.run_as_user = getuser()
local_task_job.task_instance.command_as_list.return_value = [
'airflow',
'tasks',
'test',
'test_on_kill',
'task1',
'2016-01-01',
]
# Kick off the runner
runner = StandardTaskRunner(local_task_job)
runner.start()
time.sleep(0.2)
# Kill the child process externally from the runner
# Note that we have to do this from ANOTHER process, as if we just
# call os.kill here we're doing it from the parent process and it
# won't be the same as an external kill in terms of OS tracking.
pgid = os.getpgid(runner.process.pid)
os.system(f"kill -s KILL {pgid}")
time.sleep(0.2)
runner.terminate()
assert runner.return_code() == -9
assert "running out of memory" in caplog.text