def submit_workflow(workflow_name: Text = None) -> WorkflowInfo:
"""
Submits the user-defined workflow to the scheduler with the given name of workflow. Before the submission of
workflow in Scheduler Service, the metadata of workflow will be registered in Metadata Service.
The submission of workflow translates the current :class:`~ai_flow.ai_graph.ai_graph.AIGraph` , uploads the
project package, registers the metadata of the specified workflow and submits the workflow by Scheduler Service
which delegates the submission to the :class:`~ai_flow.plugin_interface.scheduler_interface.Scheduler`.
:param workflow_name: The name of the workflow.
:return: The :class:`~ai_flow.plugin_interface.scheduler_interface.WorkflowInfo` which contains the information
about the submitted workflow.
"""
if current_graph().is_empty():
raise EmptyGraphException("Cannot submit empty graph")
entry_module_path = current_project_context().get_workflow_entry_module(workflow_name=workflow_name)
namespace = current_project_config().get_project_name()
translator = get_translator()
workflow = translator.translate(graph=current_graph(), project_context=current_project_context())
_apply_full_info_to_workflow(workflow, entry_module_path)
current_graph().clear_graph()
workflow_meta = get_ai_flow_client().get_workflow_by_name(project_name=current_project_config().get_project_name(),
workflow_name=workflow_name)
if workflow_meta is None:
get_ai_flow_client().register_workflow(name=workflow_name,
project_id=int(current_project_config().get_project_uuid()))
return proto_to_workflow(get_ai_flow_client()
.submit_workflow_to_scheduler(namespace=namespace,
workflow_json=json_utils.dumps(workflow),
workflow_name=workflow_name,
args={}))
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_event_condition_serde(self):
condition = EventCondition([], ConditionType.MEET_ANY)
condition.add_event(event_key='k1', event_value='v1')
condition.add_event(event_key='k2', event_value='v2')
loaded_condition = json_utils.loads(json_utils.dumps(condition))
self.assertEqual(condition, loaded_condition)
loaded_condition.add_event(event_key='k1', event_value='v1')
loaded_condition.add_event(event_key='k2', event_value='v2')
self.assertEqual(condition, loaded_condition)
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_graph_serde(self):
graph = Graph()
nodes = []
for i in range(3):
node = Node(name=str(i))
graph.add_node(node)
nodes.append(node)
edge = Edge(destination=nodes[0].node_id, source=nodes[1].node_id)
graph.add_edge(nodes[0].node_id, edge)
edge = Edge(destination=nodes[0].node_id, source=nodes[2].node_id)
graph.add_edge(nodes[0].node_id, edge)
json_text = json_utils.dumps(graph)
g: Graph = json_utils.loads(json_text)
self.assertEqual(3, len(g.nodes))
self.assertEqual(2, len(graph.edges.get(nodes[0].node_id)))
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 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 __str__(self) -> str:
return json_utils.dumps(self)
def test_edge_serde(self):
edge = Edge("a", 'b')
json_text = json_utils.dumps(edge)
c2: Edge = json_utils.loads(json_text)
self.assertEqual(edge.source, c2.source)
def test_node_serde(self):
node = Node(name="a")
json_text = json_utils.dumps(node)
n2: Node = json_utils.loads(json_text)
self.assertEqual(node.name, n2.name)
