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_list_all_events_with_id_range(self): event1 = self.client.send_event(BaseEvent(key="key", value="value1")) self.client.send_event(BaseEvent(key="key", value="value2")) event3 = self.client.send_event(BaseEvent(key="key", value="value3")) events = self.client.list_all_events(start_version=event1.version, end_version=event3.version) self.assertEqual(2, len(events))
def test_list_events(self): self.client._default_namespace = "a" self.client._sender = 's' event1 = self.client.send_event(BaseEvent(key="key", value="value1")) self.client._default_namespace = "b" self.client.send_event( BaseEvent(key="key", value="value2", event_type="a")) self.client.send_event(BaseEvent(key="key", value="value3")) self.client.send_event(BaseEvent(key="key2", value="value3")) events = self.client.list_events(["key", "key2"], version=event1.version) self.assertEqual(3, len(events)) self.assertEqual('s', events[0].sender) self.client._default_namespace = "a" events = self.client.list_events("key") self.assertEqual(1, len(events)) self.client._default_namespace = "b" events = self.client.list_events("key") self.assertEqual(2, len(events)) events = self.client.list_events("key", event_type="a") self.assertEqual(1, len(events)) events = self.client.list_events("key", sender='s') self.assertEqual(2, len(events)) events = self.client.list_events("key", sender='p') self.assertEqual(0, len(events))
def test_all_listen_events(self): self.client.send_event(BaseEvent(key="key", value="value1")) event2 = self.client.send_event(BaseEvent(key="key", value="value2")) start_time = event2.create_time self.client.send_event(BaseEvent(key="key", value="value3")) events = self.client.list_all_events(start_time) self.assertEqual(2, len(events))
def run_ai_flow_function(self): client = NotificationClient(server_uri="localhost:{}".format( self.port), default_namespace="default", sender='1-job-name') while True: with create_session() as session: tes = session.query(TaskExecution).filter( TaskExecution.dag_id == 'workflow_1', TaskExecution.task_id == '1-job-name').all() if len(tes) > 0: time.sleep(5) client.send_event( BaseEvent(key='key_1', value='value_1', event_type='UNDEFINED')) client.send_event( BaseEvent(key='key_2', value='value_2', event_type='UNDEFINED')) while True: with create_session() as session_2: tes_2 = session_2.query(TaskExecution).filter( TaskExecution.dag_id == 'workflow_1').all() if len(tes_2) == 3: break else: time.sleep(1) break else: time.sleep(1) time.sleep(3) client.send_event(StopSchedulerEvent(job_id=0).to_event())
def test_inner_event(self): self.assertTrue(SchedulerInnerEventUtil.is_inner_event(BaseEvent('', '', 'STOP_SCHEDULER'))) self.assertFalse(SchedulerInnerEventUtil.is_inner_event(BaseEvent('', '', 'TEST'))) self.assertEqual(SchedulerInnerEventUtil.event_type(BaseEvent('', '', 'STOP_SCHEDULER')), SchedulerInnerEventType.STOP_SCHEDULER) event = SchedulerInnerEventUtil.to_inner_event(BaseEvent('1', '', 'STOP_SCHEDULER')) self.assertTrue(isinstance(event, StopSchedulerEvent))
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_listen_events(self): event_list = [] class TestWatch(EventWatcher): def __init__(self, event_list) -> None: super().__init__() self.event_list = event_list def process(self, events: List[BaseEvent]): self.event_list.extend(events) self.client._default_namespace = "a" event1 = self.client.send_event(BaseEvent(key="key", value="value1")) self.client.start_listen_event(key="key", watcher=TestWatch(event_list), version=event1.version) self.client.send_event(BaseEvent(key="key", value="value2")) self.client.send_event(BaseEvent(key="key", value="value3")) self.client._default_namespace = None self.client.send_event(BaseEvent(key="key", value="value4")) self.client._default_namespace = "a" self.client.stop_listen_event("key") events = self.client.list_events("key", version=event1.version) self.assertEqual(2, len(events)) self.assertEqual(2, len(event_list))
def test_list_events(self): event = self.client.send_event(BaseEvent(key="key", value="value1")) event = self.client.send_event(BaseEvent(key="key", value="value2")) event = self.client.send_event(BaseEvent(key="key", value="value3")) events = self.client.list_events("key", version=1) self.assertEqual(2, len(events)) events = self.client.list_events("key") self.assertEqual(3, len(events))
def test_all_listen_events(self): event = self.client.send_event(BaseEvent(key="key", value="value1")) event = self.client.send_event(BaseEvent(key="key", value="value2")) start_time = event.create_time print("#####start time: {}".format(start_time)) event = self.client.send_event(BaseEvent(key="key", value="value3")) events = self.client.list_all_events(start_time) self.assertEqual(2, len(events))
def test_list_events(self): event = self.client.send_event(BaseEvent(key="key", value="value1")) first_version = event.version print("######first version: {}".format(first_version)) event = self.client.send_event(BaseEvent(key="key", value="value2")) event = self.client.send_event(BaseEvent(key="key", value="value3")) events = self.client.list_events("key", version=first_version) self.assertEqual(2, len(events)) events = self.client.list_events("key") self.assertEqual(3, len(events))
def add_event(self, event: BaseEvent): self.max_id += 1 event.id = self.max_id event.create_time = time.time_ns() if event.key not in self.kv: self.kv[event.key] = 0 self.kv[event.key] += 1 event.version = self.kv[event.key] self.store.append(event) return event
def test_event_type_any(self): met_config_1 = {"action": "START", "condition": "NECESSARY", "event_key": "key_1", "event_type": "*", "event_value": "value_1", "life": "ONCE", "namespace": "aa", "sender": "1-job-name", "value_condition": "EQUAL"} configs = [met_config_1] config_str = json.dumps(configs) handler = AIFlowHandler(config=config_str) event: BaseEvent = BaseEvent(key='key_1', value='value_1', event_type='aa', namespace='aa', sender='1-job-name', create_time=int(time.time()*1000)) ts = AiFlowTs() action, ts = handler.handle_event(event, ts) self.assertEqual(SchedulingAction.START, action) event: BaseEvent = BaseEvent(key='key_1', value='value_1', namespace='aa', event_type='bb', sender='1-job-name', create_time=int(time.time() * 1000+1)) action, ts = handler.handle_event(event, ts) self.assertEqual(SchedulingAction.START, action) met_config_1 = {"action": "START", "condition": "NECESSARY", "event_key": "key_1", "event_type": "aa", "event_value": "value_1", "life": "ONCE", "namespace": "aa", "sender": "1-job-name", "value_condition": "EQUAL"} configs = [met_config_1] config_str = json.dumps(configs) handler = AIFlowHandler(config=config_str) event: BaseEvent = BaseEvent(key='key_1', value='value_1', event_type='bb', namespace='aa', sender='1-job-name', create_time=int(time.time() * 1000)) ts = AiFlowTs() action, ts = handler.handle_event(event, ts) self.assertEqual(SchedulingAction.NONE, action)
def add_event(self, event: BaseEvent): kwargs = { "server_id": self.server_id, "create_time": time.time_ns(), "event_type": event.event_type, "key": event.key, "value": event.value, } mongo_event = MongoEvent(**kwargs) mongo_event.save() mongo_event.reload() event.create_time = mongo_event.create_time event.version = mongo_event.version event.id = mongo_event.auto_increase_id return event
def run_one_task_function(self): self.wait_for_running() self.client.send_event(BaseEvent(key='a', value='a')) time.sleep(5) self.client.send_event(BaseEvent(key='a', value='a')) while True: with create_session() as session: tes = session.query(TaskExecution).filter( TaskExecution.dag_id == 'single', TaskExecution.task_id == 'task_1').all() if len(tes) >= 2: break else: time.sleep(1) self.client.send_event(StopSchedulerEvent(job_id=0).to_event())
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 process(self, execution_context: ExecutionContext, input_list: List) -> List: model_meta: af.ModelMeta = execution_context.config.get('model_info') model_name = model_meta.name validated_model = af.get_latest_validated_model_version(model_name) # Deprecate deployed model deployed_model_version = af.get_deployed_model_version(model_name) if deployed_model_version is not None: af.update_model_version( model_name=model_name, model_version=deployed_model_version.version, current_stage=ModelVersionStage.DEPRECATED) af.update_model_version(model_name=model_name, model_version=validated_model.version, current_stage=ModelVersionStage.DEPLOYED) af.get_ai_flow_client().send_event( BaseEvent(key='START_PREDICTION', value=validated_model.version)) print(validated_model.version) # Copy deployed model to deploy_model_dir deployed_model_dir = af.get_artifact_by_name(self.artifact).uri if not os.path.exists(deployed_model_dir): os.makedirs(deployed_model_dir) for file in os.listdir(deployed_model_dir): file_path = os.path.join(deployed_model_dir, file) if os.path.isfile(file_path): os.remove(file_path) elif os.path.isdir(file_path): shutil.rmtree(file_path, True) deployed_model_version = af.get_deployed_model_version(model_name) shutil.copy(deployed_model_version.model_path, deployed_model_dir) return []
def run_event_task_function(self): client = NotificationClient(server_uri="localhost:{}".format( self.port), default_namespace="") while True: with create_session() as session: tes = session.query(TaskExecution).filter( TaskExecution.dag_id == 'event_dag', TaskExecution.task_id == 'task_1').all() if len(tes) > 0: time.sleep(5) client.send_event( BaseEvent(key='start', value='', event_type='', namespace='')) while True: with create_session() as session_2: tes_2 = session_2.query(TaskExecution).filter( TaskExecution.dag_id == 'event_dag', TaskExecution.task_id == 'task_2').all() if len(tes_2) > 0: break else: time.sleep(1) break else: time.sleep(1) client.send_event(StopSchedulerEvent(job_id=0).to_event())
def test_dag_run_event_manager(self): mailbox = Mailbox() event_manager = DagRunEventManager(mailbox) event_manager.start() self.create_task_state(dag_run=self._dag_run, task_id='operator_toggle_handler') event = BaseEvent("test_event", "test_event", namespace="default") event_manager.handle_event( DagRunId(self._dag_run.dag_id, self._dag_run.run_id), event) event_manager.handle_event( DagRunId(self._dag_run.dag_id, self._dag_run.run_id), event) handle_event = mailbox.get_message() message = EventHandleResult.from_event(handle_event) assert message == EventHandleResult( DagRunId(self._dag_run.dag_id, self._dag_run.run_id), "operator_toggle_handler", SchedulingAction.START) handle_event = mailbox.get_message() message = EventHandleResult.from_event(handle_event) assert message == EventHandleResult( DagRunId(self._dag_run.dag_id, self._dag_run.run_id), "operator_toggle_handler", SchedulingAction.STOP) time.sleep(2) event_manager.handle_event( DagRunId(self._dag_run.dag_id, self._dag_run.run_id), event) handle_event = mailbox.get_message() message = EventHandleResult.from_event(handle_event) assert message == EventHandleResult( DagRunId(self._dag_run.dag_id, self._dag_run.run_id), "operator_toggle_handler", SchedulingAction.START) event_manager.end()
def deleteModelVersion(self, request, context): model_meta_param = ModelVersion.from_proto(request) self.model_repo_store.delete_model_version(model_meta_param) self.notification_client.send_event(BaseEvent(model_meta_param.model_name, json.dumps(model_meta_param.__dict__), ModelVersionEventType.MODEL_DELETED)) return _wrap_response(request.model_meta)
def test_two_config_2(self): met_config_1 = {"action": "START", "condition": "SUFFICIENT", "event_key": "key_1", "event_type": "UNDEFINED", "event_value": "value_1", "life": "ONCE", "namespace": "default", "sender": "1-job-name", "value_condition": "EQUAL"} met_config_2 = {"action": "START", "condition": "NECESSARY", "event_key": "key_2", "event_type": "UNDEFINED", "event_value": "value_2", "life": "ONCE", "namespace": "default", "sender": "1-job-name", "value_condition": "EQUAL"} configs = [met_config_1, met_config_2] config_str = json.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)) ts = AiFlowTs() action, ts = handler.handle_event(event, ts) self.assertEqual(SchedulingAction.START, action)
def test_dag_run_event_manager_release_runner(self): dag_run1 = self._dag_run _, dag_run2 = self.init_dag_and_dag_run( '../../dags/test_task_event_handler_dag.py', 'test_event_handler', timezone.datetime(2017, 1, 2)) self.create_task_state(dag_run1, 'operator_toggle_handler') self.create_task_state(dag_run2, 'operator_toggle_handler') event = BaseEvent("test_event", "test_event", namespace="default") mailbox = Mailbox() event_manager = DagRunEventManager(mailbox=mailbox) event_manager.handle_event(DagRunId(dag_run1.dag_id, dag_run1.run_id), event) time.sleep(5) event_manager.handle_event(DagRunId(dag_run1.dag_id, dag_run2.run_id), event) assert (DagRunId( dag_run2.dag_id, dag_run2.run_id)) in event_manager._event_executor_runners assert (DagRunId( dag_run1.dag_id, dag_run1.run_id)) not in event_manager._event_executor_runners event_manager.end()
def test_task_event_executor_runner(self): event = BaseEvent("test_event", "test_event", namespace="default") self.create_task_state(dag_run=self._dag_run, task_id='operator_toggle_handler') mailbox = Mailbox() executor_runner = DagRunEventExecutorRunner( mailbox, DagRunId(self._dag_run.dag_id, self._dag_run.run_id), 10) executor_runner.put_event(event) executor_runner.put_event(event) executor_runner.run() handle_event = mailbox.get_message() message = EventHandleResult.from_event( SchedulerInnerEventUtil.to_inner_event(handle_event)) assert message == EventHandleResult( DagRunId(self._dag_run.dag_id, self._dag_run.run_id), "operator_toggle_handler", SchedulingAction.START) handle_event = mailbox.get_message() message = EventHandleResult.from_event( SchedulerInnerEventUtil.to_inner_event(handle_event)) assert message == EventHandleResult( DagRunId(self._dag_run.dag_id, self._dag_run.run_id), "operator_toggle_handler", SchedulingAction.STOP)
def run_task_function(client: NotificationClient): with af.global_config_file(workflow_config_file()): with af.config('task_2'): executor_1 = af.user_define_operation( af.PythonObjectExecutor(SimpleExecutor())) with af.config('task_5'): executor_2 = af.user_define_operation( af.PythonObjectExecutor(SimpleExecutor())) af.user_define_control_dependency(src=executor_2, dependency=executor_1, namespace='test', event_key='key_1', event_value='value_1', sender='*') workflow_info = af.workflow_operation.submit_workflow( workflow_name) af.workflow_operation.start_new_workflow_execution(workflow_name) flag = True while True: with create_session() as session: tes = session.query(TaskExecution).filter( TaskExecution.dag_id == 'test_project.test_workflow', TaskExecution.task_id == 'task_2').all() if 1 == len(tes) and flag: client.send_event( BaseEvent(key='key_1', value='value_1')) flag = False dag_run = session.query(DagRun).filter( DagRun.dag_id == 'test_project.test_workflow').first() if dag_run is not None and dag_run.state in State.finished: break else: time.sleep(1)
def handle_event(self, event: BaseEvent, task_state: object) -> Tuple[SchedulingAction, object]: rules = self._parse_configs(self.config) if SchedulerInnerEventType.TASK_STATUS_CHANGED.value == event.event_type: event = copy.deepcopy(event) event.value = _airflow_task_state_to_aiflow_status(event.value) if task_state is None: task_state = AiFlowTs(len(rules)) af_ts = copy.deepcopy(task_state) af_ts.add_event(event) actions = self._check_rules(rules, af_ts) # pick the first triggered action for action in actions: if action is not None: scheduling_action = job_action_to_scheduling_action(action) logging.debug( "AIFlowHandler {} handle event {} triggered action: {}". format(self.config, event, scheduling_action)) return scheduling_action, af_ts logging.debug( "AIFlowHandler {} handle event {} no action triggered return action {}" .format(self.config, event, SchedulingAction.NONE)) return SchedulingAction.NONE, af_ts
async def _send_event(self, request): event_proto = request.event event = BaseEvent( key=event_proto.key, value=event_proto.value, event_type=None if event_proto.event_type == "" else event_proto.event_type, context=None if event_proto.context == "" else event_proto.context, namespace=None if event_proto.namespace == "" else event_proto.namespace) uuid = request.uuid key = event.key namespace = event.namespace # Lock conditions dict for get/check/update of key await self.lock.acquire() if self.notification_conditions.get((key, namespace)) is None: self.notification_conditions.update({((key, namespace), asyncio.Condition())}) # Release lock after check/update key of notification conditions dict self.lock.release() async with self.notification_conditions.get( (key, namespace)), self.write_condition: event: BaseEvent = self.storage.add_event(event, uuid) self.notification_conditions.get((key, namespace)).notify_all() self.write_condition.notify_all() result_event_proto = event_to_proto(event) return notification_service_pb2.SendEventsResponse( event=result_event_proto, return_code=notification_service_pb2.ReturnStatus.SUCCESS, return_msg='')
def event_proto_to_event(event_proto): return BaseEvent(id=event_proto.id, key=event_proto.key, value=event_proto.value, event_type=event_proto.event_type, version=event_proto.version, create_time=event_proto.create_time)
def schedule_task_function(self): stopped = False while not stopped: with create_session() as session: ti_sleep_1000_secs = self._get_task_instance( EVENT_BASED_SCHEDULER_DAG, 'sleep_1000_secs', session) ti_python_sleep = self._get_task_instance( EVENT_BASED_SCHEDULER_DAG, 'python_sleep', session) if ti_sleep_1000_secs and ti_sleep_1000_secs.state == State.SCHEDULED and \ ti_python_sleep and ti_python_sleep.state == State.SCHEDULED: self.client.send_event( BaseEvent(key='start', value='', event_type='', namespace='test_namespace')) while not stopped: ti_sleep_1000_secs.refresh_from_db() ti_python_sleep.refresh_from_db() if ti_sleep_1000_secs and ti_sleep_1000_secs.state == State.RUNNING and \ ti_python_sleep and ti_python_sleep.state == State.RUNNING: time.sleep(10) break else: time.sleep(1) self.client.send_event( BaseEvent(key='stop', value='', event_type=UNDEFINED_EVENT_TYPE, namespace='test_namespace')) self.client.send_event( BaseEvent(key='restart', value='', event_type=UNDEFINED_EVENT_TYPE, namespace='test_namespace')) while not stopped: ti_sleep_1000_secs.refresh_from_db() ti_python_sleep.refresh_from_db() if ti_sleep_1000_secs and ti_sleep_1000_secs.state == State.KILLED and \ ti_python_sleep and ti_python_sleep.state == State.RUNNING: stopped = True else: time.sleep(1) else: time.sleep(1) self.client.send_event(StopSchedulerEvent(job_id=0).to_event())
def add_event(self, event: BaseEvent, uuid: str): kwargs = { "server_id": self.server_id, "create_time": int(time.time() * 1000), "event_type": event.event_type, "key": event.key, "value": event.value, "context": event.context, "namespace": event.namespace, "uuid": uuid } mongo_event = MongoEvent(**kwargs) mongo_event.save() mongo_event.reload() event.create_time = mongo_event.create_time event.version = mongo_event.version return event
def event_proto_to_event(event_proto): return BaseEvent(key=event_proto.key, value=event_proto.value, event_type=event_proto.event_type, version=event_proto.version, create_time=event_proto.create_time, context=event_proto.context, namespace=event_proto.namespace)