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_mark_success_on_success_callback(self):
"""
Test that ensures that where a task is marked suceess in the UI
on_success_callback gets executed
"""
data = {'called': False}
def success_callback(context):
self.assertEqual(context['dag_run'].dag_id, 'test_mark_success')
data['called'] = True
dag = DAG(dag_id='test_mark_success',
start_date=DEFAULT_DATE,
default_args={'owner': 'owner1'})
task = DummyOperator(task_id='test_state_succeeded1',
dag=dag,
on_success_callback=success_callback)
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 = TaskInstance(task=task, execution_date=DEFAULT_DATE)
ti.refresh_from_db()
job1 = LocalTaskJob(task_instance=ti,
ignore_ti_state=True,
executor=SequentialExecutor())
from airflow.task.task_runner.standard_task_runner import StandardTaskRunner
job1.task_runner = StandardTaskRunner(job1)
process = multiprocessing.Process(target=job1.run)
process.start()
ti.refresh_from_db()
for _ in range(0, 50):
if ti.state == State.RUNNING:
break
time.sleep(0.1)
ti.refresh_from_db()
self.assertEqual(State.RUNNING, ti.state)
ti.state = State.SUCCESS
session.merge(ti)
session.commit()
job1.heartbeat_callback(session=None)
self.assertTrue(data['called'])
process.join(timeout=10)
self.assertFalse(process.is_alive())
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 get_task_runner(local_task_job):
"""
Get the task runner that can be used to run the given job.
:param local_task_job: The LocalTaskJob associated with the TaskInstance
that needs to be executed.
:type local_task_job: airflow.jobs.LocalTaskJob
:return: The task runner to use to run the task.
:rtype: airflow.task.task_runner.base_task_runner.BaseTaskRunner
"""
if _TASK_RUNNER == "StandardTaskRunner":
return StandardTaskRunner(local_task_job)
elif _TASK_RUNNER == "CgroupTaskRunner":
from airflow.task.task_runner.cgroup_task_runner import CgroupTaskRunner
return CgroupTaskRunner(local_task_job)
else:
raise AirflowException(
"Unknown task runner type {}".format(_TASK_RUNNER))
def test_localtaskjob_heartbeat(self):
session = settings.Session()
dag = DAG('test_localtaskjob_heartbeat',
start_date=DEFAULT_DATE,
default_args={'owner': 'owner1'})
with dag:
op1 = DummyOperator(task_id='op1')
dag.clear()
dr = dag.create_dagrun(
run_id="test",
state=State.SUCCESS,
execution_date=DEFAULT_DATE,
start_date=DEFAULT_DATE,
session=session,
)
ti = dr.get_task_instance(task_id=op1.task_id, session=session)
ti.state = State.RUNNING
ti.hostname = "blablabla"
session.commit()
job1 = LocalTaskJob(task_instance=ti,
ignore_ti_state=True,
executor=SequentialExecutor())
ti.task = op1
ti.refresh_from_task(op1)
job1.task_runner = StandardTaskRunner(job1)
job1.task_runner.process = mock.Mock()
with pytest.raises(AirflowException):
job1.heartbeat_callback() # pylint: disable=no-value-for-parameter
job1.task_runner.process.pid = 1
ti.state = State.RUNNING
ti.hostname = get_hostname()
ti.pid = 1
session.merge(ti)
session.commit()
assert ti.pid != os.getpid()
job1.heartbeat_callback(session=None)
job1.task_runner.process.pid = 2
with pytest.raises(AirflowException):
job1.heartbeat_callback() # pylint: disable=no-value-for-parameter
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
def test_mark_success_on_success_callback(self):
"""
Test that ensures that where a task is marked suceess in the UI
on_success_callback gets executed
"""
# use shared memory value so we can properly track value change even if
# it's been updated across processes.
success_callback_called = Value('i', 0)
task_terminated_externally = Value('i', 1)
shared_mem_lock = Lock()
def success_callback(context):
with shared_mem_lock:
success_callback_called.value += 1
assert context['dag_run'].dag_id == 'test_mark_success'
dag = DAG(dag_id='test_mark_success', start_date=DEFAULT_DATE, default_args={'owner': 'owner1'})
def task_function(ti):
# pylint: disable=unused-argument
time.sleep(60)
# This should not happen -- the state change should be noticed and the task should get killed
with shared_mem_lock:
task_terminated_externally.value = 0
task = PythonOperator(
task_id='test_state_succeeded1',
python_callable=task_function,
on_success_callback=success_callback,
dag=dag,
)
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 = TaskInstance(task=task, execution_date=DEFAULT_DATE)
ti.refresh_from_db()
job1 = LocalTaskJob(task_instance=ti, ignore_ti_state=True, executor=SequentialExecutor())
job1.task_runner = StandardTaskRunner(job1)
settings.engine.dispose()
process = multiprocessing.Process(target=job1.run)
process.start()
for _ in range(0, 25):
ti.refresh_from_db()
if ti.state == State.RUNNING:
break
time.sleep(0.2)
assert ti.state == State.RUNNING
ti.state = State.SUCCESS
session.merge(ti)
session.commit()
process.join(timeout=10)
assert success_callback_called.value == 1
assert task_terminated_externally.value == 1
assert not process.is_alive()