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