def test_default_pool_open_slots(self): set_default_pool_slots(5) assert 5 == Pool.get_default_pool().open_slots() dag = DAG( dag_id='test_default_pool_open_slots', start_date=DEFAULT_DATE, ) op1 = DummyOperator(task_id='dummy1', dag=dag) op2 = DummyOperator(task_id='dummy2', dag=dag, pool_slots=2) ti1 = TI(task=op1, execution_date=DEFAULT_DATE) ti2 = TI(task=op2, execution_date=DEFAULT_DATE) ti1.state = State.RUNNING ti2.state = State.QUEUED session = settings.Session session.add(ti1) session.add(ti2) session.commit() session.close() assert 2 == Pool.get_default_pool().open_slots() assert { "default_pool": { "open": 2, "queued": 2, "total": 5, "running": 1, } } == Pool.slots_stats()
def test_backfill_respect_default_pool_limit(self, mock_log): default_pool_slots = 2 set_default_pool_slots(default_pool_slots) dag = self._get_dummy_dag('test_backfill_with_no_pool_limit') executor = MockExecutor() job = BackfillJob( dag=dag, executor=executor, start_date=DEFAULT_DATE, end_date=DEFAULT_DATE + datetime.timedelta(days=7), ) job.run() self.assertTrue(0 < len(executor.history)) default_pool_task_slot_count_reached_at_least_once = False num_running_task_instances = 0 # if no pool is specified, the number of tasks running in # parallel per backfill should be less than # default_pool slots at any point of time. for running_task_instances in executor.history: self.assertLessEqual( len(running_task_instances), default_pool_slots, ) num_running_task_instances += len(running_task_instances) if len(running_task_instances) == default_pool_slots: default_pool_task_slot_count_reached_at_least_once = True self.assertEquals(8, num_running_task_instances) self.assertTrue(default_pool_task_slot_count_reached_at_least_once) times_dag_concurrency_limit_reached_in_debug = self._times_called_with( mock_log.debug, DagConcurrencyLimitReached, ) times_pool_limit_reached_in_debug = self._times_called_with( mock_log.debug, NoAvailablePoolSlot, ) times_task_concurrency_limit_reached_in_debug = self._times_called_with( mock_log.debug, TaskConcurrencyLimitReached, ) self.assertEqual(0, times_dag_concurrency_limit_reached_in_debug) self.assertEqual(0, times_task_concurrency_limit_reached_in_debug) self.assertGreater(times_pool_limit_reached_in_debug, 0)
def test_default_pool_open_slots(self): set_default_pool_slots(5) self.assertEqual(5, Pool.get_default_pool().open_slots()) dag = DAG( dag_id='test_default_pool_open_slots', start_date=DEFAULT_DATE, ) t1 = DummyOperator(task_id='dummy1', dag=dag) t2 = DummyOperator(task_id='dummy2', dag=dag) ti1 = TI(task=t1, execution_date=DEFAULT_DATE) ti2 = TI(task=t2, execution_date=DEFAULT_DATE) ti1.state = State.RUNNING ti2.state = State.QUEUED session = settings.Session session.add(ti1) session.add(ti2) session.commit() session.close() self.assertEqual(3, Pool.get_default_pool().open_slots())