def test_two_rules_trigger_will_clear_all_event_life(self):
condition = MeetAnyEventCondition() \
.add_event(event_key="key_1", event_value="value_1", sender="1-job-name")
rule1 = SchedulingRule(condition, action=JobAction.START)
condition = MeetAllEventCondition() \
.add_event(event_key="key_1", event_value="value_1", sender="1-job-name") \
.add_event(event_key='key_2', event_value='value_2', sender="1-job-name")
rule2 = SchedulingRule(condition, action=JobAction.STOP)
configs = [rule1, rule2]
config_str = json_utils.dumps(configs)
handler = AIFlowHandler(config=config_str)
event1: BaseEvent = BaseEvent(key='key_1',
value='value_1',
namespace='default',
sender='1-job-name',
create_time=round(time.time() * 1000))
event2: BaseEvent = BaseEvent(key='key_2',
value='value_2',
namespace='default',
sender='1-job-name',
create_time=round(time.time() * 1000))
action, ts = handler.handle_event(event2, None)
self.assertEqual(SchedulingAction.NONE, action)
action, ts = handler.handle_event(event1, ts)
self.assertEqual(SchedulingAction.START, action)
action, ts = handler.handle_event(event2, ts)
self.assertEqual(SchedulingAction.NONE, action)
def test_event_handler_schedule_time(self):
condition = MeetAnyEventCondition() \
.add_event(event_key="key_1", event_value="value_1", sender="1-job-name")
rule1 = SchedulingRule(condition, action=JobAction.START)
condition = MeetAllEventCondition() \
.add_event(event_key="key_1", event_value="value_1", sender="1-job-name") \
.add_event(event_key='key_2', event_value='value_2', sender="1-job-name")
rule2 = SchedulingRule(condition, action=JobAction.STOP)
configs = [rule1, rule2]
config_str = json_utils.dumps(configs)
handler = AIFlowHandler(config=config_str)
event1_time = round(time.time() * 1000)
event1: BaseEvent = BaseEvent(key='key_1',
value='value_1',
namespace='default',
sender='1-job-name',
create_time=event1_time)
action, ts = handler.handle_event(event1, None)
ts: AiFlowTs = ts
self.assertEqual(event1_time, ts.rule_states[0].schedule_time)
self.assertEqual(event1_time, ts.rule_states[0].latest_time)
self.assertNotEqual(event1_time, ts.rule_states[1].schedule_time)
self.assertEqual(event1_time, ts.rule_states[1].latest_time)
def test_handle_task_status_change_event(self):
from airflow.events.scheduler_events import TaskStateChangedEvent
from airflow.utils.state import State
rule1 = SchedulingRule(MeetAnyEventCondition().add_event(
event_key="dag_1.task_1",
event_value="RUNNING",
event_type="TASK_STATUS_CHANGED",
namespace="test",
sender="task_1"), action=JobAction.START)
rule2 = SchedulingRule(MeetAnyEventCondition().add_event(
event_key="dag_1.task_1",
event_value="FINISHED",
event_type="TASK_STATUS_CHANGED",
namespace="test",
sender="task_1"), action=JobAction.STOP)
configs = [rule1, rule2]
config_str = json_utils.dumps(configs)
handler = AIFlowHandler(config=config_str)
event = TaskStateChangedEvent("task_1", "test.dag_1", datetime.utcnow(), State.RUNNING).to_event()
action, ts = handler.handle_event(event, None)
self.assertEqual(SchedulingAction.START, action)
time.sleep(0.01)
event = TaskStateChangedEvent("task_1", "test.dag_1", datetime.utcnow(), State.KILLING).to_event()
action, ts = handler.handle_event(event, ts)
self.assertEqual(SchedulingAction.NONE, action)
time.sleep(0.01)
event = TaskStateChangedEvent("task_1", "test.dag_1", datetime.utcnow(), State.SUCCESS).to_event()
action, ts = handler.handle_event(event, ts)
self.assertEqual(SchedulingAction.STOP, action)
def test_multiple_any_config(self):
rule = SchedulingRule(MeetAllEventCondition().add_event(
event_key="key_1",
event_value="value_1",
event_type="*",
namespace="*",
sender="1-job-name"), action=JobAction.START)
configs = [rule]
config_str = json_utils.dumps(configs)
handler = AIFlowHandler(config=config_str)
event: BaseEvent = BaseEvent(key='key_1',
value='value_1',
namespace='default',
sender='1-job-name',
create_time=round(time.time() * 1000))
action, ts = handler.handle_event(event, None)
self.assertEqual(SchedulingAction.START, action)
event: BaseEvent = BaseEvent(key='key_1',
value='value_1',
namespace='default',
sender='aa',
create_time=round(time.time() * 1000))
action, ts = handler.handle_event(event, ts)
self.assertEqual(SchedulingAction.NONE, action)
event: BaseEvent = BaseEvent(key='key_1_1',
value='value_1',
namespace='default',
sender='1-job-name',
create_time=round(time.time() * 1000))
action, ts = handler.handle_event(event, ts)
self.assertEqual(SchedulingAction.NONE, action)
def test_control_edge(self):
condition = MeetAllEventCondition()
action = JobAction.START
rule = SchedulingRule(condition, action)
edge = ControlEdge('task', rule)
self.assertEqual('*', edge.source)
self.assertEqual('task', edge.destination)
self.assertEqual(rule, edge.scheduling_rule)
def test_sender_any(self):
rule = SchedulingRule(MeetAllEventCondition().add_event(
event_key="key_1",
event_value="value_1",
namespace="aa",
sender="*"), action=JobAction.START)
configs = [rule]
config_str = json_utils.dumps(configs)
handler = AIFlowHandler(config=config_str)
event: BaseEvent = BaseEvent(key='key_1',
value='value_1',
namespace='aa',
sender='aa',
create_time=int(time.time() * 1000))
action, ts = handler.handle_event(event, None)
self.assertEqual(SchedulingAction.START, action)
event: BaseEvent = BaseEvent(key='key_1',
value='value_1',
namespace='aa',
sender='bb',
create_time=int(time.time() * 1000 + 1))
action, ts = handler.handle_event(event, None)
self.assertEqual(SchedulingAction.START, action)
rule = SchedulingRule(MeetAllEventCondition().add_event(
event_key="key_1",
event_value="value_1",
namespace="aa",
sender="aa"), action=JobAction.START)
configs = [rule]
config_str = json_utils.dumps(configs)
handler = AIFlowHandler(config=config_str)
event: BaseEvent = BaseEvent(key='key_1',
value='value_1',
namespace='aa',
sender='bb',
create_time=int(time.time() * 1000))
action, ts = handler.handle_event(event, None)
self.assertEqual(SchedulingAction.NONE, action)
def test_workflow_serde(self):
workflow_config_file = os.path.join(os.path.dirname(__file__), 'workflow_1.yaml')
workflow_config = load_workflow_config(workflow_config_file)
workflow = Workflow()
workflow.workflow_config = workflow_config
jobs = []
for job_config in workflow_config.job_configs.values():
job = Job(job_config=job_config)
workflow.add_job(job)
jobs.append(job)
edge = ControlEdge(destination=jobs[0].job_name,
scheduling_rule=SchedulingRule(MeetAnyEventCondition().add_event('a', 'a'),
JobAction.START))
workflow.add_edge(jobs[0].job_name, edge)
edge = ControlEdge(destination=jobs[0].job_name,
scheduling_rule=SchedulingRule(MeetAnyEventCondition().add_event('b', 'b'),
JobAction.START))
workflow.add_edge(jobs[0].job_name, edge)
json_text = json_utils.dumps(workflow)
w: Workflow = json_utils.loads(json_text)
self.assertEqual(3, len(w.jobs))
self.assertEqual(2, len(w.edges.get(jobs[0].job_name)))
def action_on_events(job_name: Text, event_condition: EventCondition,
action: JobAction):
"""
Defines a rule which is the combination of a :class:`~ai_flow.workflow.control_edge.EventCondition` and
:class:`~ai_flow.workflow.control_edge.JobAction`, on the job with the given name. When the
:class:`~ai_flow.workflow.control_edge.EventCondition` of the rule met, the corresponding
:class:`~ai_flow.workflow.control_edge.JobAction` will be triggered on the job. User could call the method
multiple times on the same job to add multiples rules, the rules will be checked according to the order of the
the method call.
:param job_name: The name of the job to add the rule.
:param event_condition: The :class:`~ai_flow.workflow.control_edge.EventCondition` of the rule.
:param action: The :class:`~ai_flow.workflow.control_edge.JobAction` to take when the
:class:`~ai_flow.workflow.control_edge.EventCondition` is met.
"""
rule = SchedulingRule(event_condition, action)
control_edge = ControlEdge(destination=job_name, scheduling_rule=rule)
current_graph().add_edge(job_name, control_edge)
def test_unordered_event_life_repeated(self):
condition = MeetAnyEventCondition() \
.add_event(event_key="key_1", event_value="value_1", sender="1-job-name", life=EventLife.REPEATED)
rule1 = SchedulingRule(condition, action=JobAction.START)
configs = [rule1]
config_str = json_utils.dumps(configs)
handler = AIFlowHandler(config=config_str)
event1: BaseEvent = BaseEvent(key='key_1',
value='value_1',
namespace='default',
sender='1-job-name',
create_time=round(time.time() * 1000))
event1_in_past: BaseEvent = BaseEvent(key='key_1',
value='value_1',
namespace='default',
sender='1-job-name',
create_time=0)
action, ts = handler.handle_event(event1, None)
self.assertEqual(SchedulingAction.START, action)
action, ts = handler.handle_event(event1_in_past, ts)
self.assertEqual(SchedulingAction.START, action)
def test_scheduling_rule_serde(self):
condition = MeetAllEventCondition()
action = JobAction.START
rule = SchedulingRule(condition, action)
loaded_rule = json_utils.loads(json_utils.dumps(rule))
self.assertEqual(rule, loaded_rule)
def test_scheduling_rule(self):
condition = MeetAllEventCondition()
action = JobAction.START
rule = SchedulingRule(condition, action)
self.assertEqual(condition, rule.event_condition)
self.assertEqual(action, rule.action)
def add_control_edge(graph, job_name, upstream_job_name):
rule = SchedulingRule(MeetAnyEventCondition().add_event('', '', sender=upstream_job_name), JobAction.START)
graph.add_edge(job_name,
ControlEdge(destination=job_name,
scheduling_rule=rule))