def start_notification_service(port: int = 50052,
db_conn: str = None,
enable_ha: bool = False,
server_uris: str = None,
create_table_if_not_exists: bool = True):
if db_conn:
storage = DbEventStorage(db_conn, create_table_if_not_exists)
else:
raise Exception(
'Failed to start notification service without database connection info.'
)
if enable_ha:
if not server_uris:
raise Exception("When HA enabled, server_uris must be set.")
ha_storage = DbHighAvailabilityStorage(db_conn=db_conn)
ha_manager = SimpleNotificationServerHaManager()
service = HighAvailableNotificationService(storage, ha_manager,
server_uris, ha_storage,
5000)
master = NotificationMaster(service=service, port=int(port))
else:
master = NotificationMaster(service=NotificationService(storage),
port=port)
master.run(is_block=True)
def setUpClass(cls):
kwargs = {
"host": "127.0.0.1",
"port": 27017,
"db": "test"
}
cls.storage = MongoEventStorage(**kwargs)
cls.master = NotificationMaster(NotificationService(cls.storage))
cls.master.run()
cls.client = NotificationClient(server_uri="localhost:50051")
def _execute(self):
"""
1. Init the DagRun route.
2. Start the executor.
3. Option of start the notification master.
4. Create the notification client.
5. Start the DagTrigger.
6. Run the scheduler event loop.
:return:
"""
notification_client = None
try:
self._init_route()
self.executor.set_use_nf(True)
self.executor.start()
self.dag_trigger = DagTrigger(
subdir=self.subdir,
mailbox=self.mail_box,
run_duration=self.run_duration,
using_sqlite=self.using_sqlite,
num_runs=self.num_runs,
processor_poll_interval=self._processor_poll_interval)
if self.use_local_nf:
self.notification_master \
= NotificationMaster(service=NotificationService(EventModelStorage()), port=self.nf_port)
self.notification_master.run()
self.log.info("start notification service {0}".format(
self.nf_port))
notification_client = NotificationClient(
server_uri="localhost:{0}".format(self.nf_port))
else:
notification_client \
= NotificationClient(server_uri="{0}:{1}".format(self.nf_host, self.nf_port))
notification_client.start_listen_events(
watcher=SCEventWatcher(self.mail_box))
self.dag_trigger.start()
self._start_executor_heartbeat()
self._run_event_loop()
except Exception as e:
self.log.exception("Exception when executing _execute {0}".format(
str(e)))
finally:
self.running = False
self._stop_executor_heartheat()
if self.dag_trigger is not None:
self.dag_trigger.stop()
if notification_client is not None:
notification_client.stop_listen_events()
if self.notification_master is not None:
self.notification_master.stop()
self.executor.end()
self.log.info("Exited execute event scheduler")
def __init__(self,
store_uri=None,
port=_PORT,
start_default_notification: bool = True,
notification_uri=None,
start_meta_service: bool = True,
start_model_center_service: bool = True,
start_metric_service: bool = True,
start_scheduler_service: bool = True,
scheduler_service_config: Dict = None,
enabled_ha: bool = False,
ha_manager=None,
ha_server_uri=None,
ha_storage=None,
ttl_ms: int = 10000):
self.store_uri = store_uri
self.db_type = DBType.value_of(extract_db_engine_from_uri(store_uri))
self.executor = Executor(futures.ThreadPoolExecutor(max_workers=10))
self.server = grpc.server(self.executor)
self.start_default_notification = start_default_notification
self.enabled_ha = enabled_ha
server_uri = 'localhost:{}'.format(port)
if start_default_notification:
logging.info("start default notification service.")
notification_service_pb2_grpc.add_NotificationServiceServicer_to_server(
NotificationService.from_storage_uri(store_uri), self.server)
if start_model_center_service:
logging.info("start model center service.")
model_center_service_pb2_grpc.add_ModelCenterServiceServicer_to_server(
ModelCenterService(
store_uri=store_uri,
notification_uri=server_uri if start_default_notification
and notification_uri is None else notification_uri),
self.server)
if start_meta_service:
logging.info("start meta service.")
metadata_service_pb2_grpc.add_MetadataServiceServicer_to_server(
MetadataService(db_uri=store_uri, server_uri=server_uri),
self.server)
if start_metric_service:
logging.info("start metric service.")
metric_service_pb2_grpc.add_MetricServiceServicer_to_server(
MetricService(db_uri=store_uri), self.server)
if start_scheduler_service:
self._add_scheduler_service(scheduler_service_config)
if enabled_ha:
self._add_ha_service(ha_manager, ha_server_uri, ha_storage,
store_uri, ttl_ms)
self.server.add_insecure_port('[::]:' + str(port))
def test_user_trigger_parse_dag(self):
port = 50101
service_uri = 'localhost:{}'.format(port)
storage = MemoryEventStorage()
master = NotificationMaster(NotificationService(storage), port)
master.run()
mailbox = Mailbox()
dag_trigger = DagTrigger("../../dags/test_scheduler_dags.py", -1, [], False, mailbox, 5, service_uri)
dag_trigger.start()
message = mailbox.get_message()
message = SchedulerInnerEventUtil.to_inner_event(message)
# only one dag is executable
assert "test_task_start_date_scheduling" == message.dag_id
sc = EventSchedulerClient(server_uri=service_uri, namespace='a')
sc.trigger_parse_dag()
dag_trigger.end()
master.stop()
def setUp(self):
db.clear_db_jobs()
db.clear_db_dags()
db.clear_db_serialized_dags()
db.clear_db_runs()
db.clear_db_task_execution()
db.clear_db_message()
self.scheduler = None
self.port = 50102
self.storage = MemoryEventStorage()
self.master = NotificationMaster(NotificationService(self.storage),
self.port)
self.master.run()
self.client = NotificationClient(server_uri="localhost:{}".format(
self.port),
default_namespace="test_namespace")
time.sleep(1)
def test_trigger_parse_dag(self):
import os
port = 50102
server_uri = "localhost:{}".format(port)
storage = MemoryEventStorage()
master = NotificationMaster(NotificationService(storage), port)
master.run()
dag_folder = os.path.abspath(os.path.dirname(__file__)) + "/../../dags"
mailbox = Mailbox()
dag_trigger = DagTrigger(dag_folder, -1, [], False, mailbox, notification_service_uri=server_uri)
dag_trigger.start()
to_be_triggered = [dag_folder + "/test_event_based_scheduler.py",
dag_folder + "/test_event_task_dag.py",
dag_folder + "/test_event_based_executor.py",
dag_folder + "/test_scheduler_dags.py",
]
for file in to_be_triggered:
self._send_request_and_receive_response(server_uri, file)
dag_trigger.end()
def setUp(self) -> None:
SqlAlchemyStore(_SQLITE_DB_URI)
self.notification = NotificationMaster(
service=NotificationService(storage=MemoryEventStorage()),
port=30031)
self.notification.run()
self.server1 = AIFlowServer(store_uri=_SQLITE_DB_URI,
port=50051,
enabled_ha=True,
start_scheduler_service=False,
ha_server_uri='localhost:50051',
notification_uri='localhost:30031',
start_default_notification=False)
self.server1.run()
self.server2 = None
self.server3 = None
self.config = ProjectConfig()
self.config.set_enable_ha(True)
self.config.set_notification_service_uri('localhost:30031')
self.client = AIFlowClient(
server_uri='localhost:50052,localhost:50051',
project_config=self.config)
def start_notification_service(port: int = 50052):
storage = MemoryEventStorage()
notification_master \
= NotificationMaster(service=NotificationService(storage), port=port)
notification_master.run(is_block=True)
def set_up_class(cls):
cls.storage = MemoryEventStorage()
cls.master = NotificationMaster(NotificationService(cls.storage))
cls.master.run()
def test_from_storage_uri(self):
service = NotificationService.from_storage_uri("sqlite:///aiflow.db")
self.assertTrue(isinstance(service.storage, DbEventStorage))
service = NotificationService.from_storage_uri(
"mongodb://*****:*****@localhost:27017/test")
self.assertTrue(isinstance(service.storage, MongoEventStorage))
def setUpClass(cls):
cls.master = NotificationMaster(
NotificationService(EventModelStorage()))
cls.master.run()
cls.client = NotificationClient(server_uri="localhost:50051")
def setUpClass(cls):
cls.storage = EventModelStorage()
cls.master = NotificationMaster(NotificationService(cls.storage))
cls.master.run()
def run_server():
storage = MemoryEventStorage()
master = NotificationMaster(service=NotificationService(storage),
port=50051)
master.run(is_block=True)
def setUpClass(cls):
cls.storage = MemoryEventStorage()
cls.master = NotificationMaster(NotificationService(cls.storage))
cls.master.run()
cls.client = NotificationClient(server_uri="localhost:50051")
def setUp(self):
clear_db_event_model()
self.master = NotificationMaster(service=NotificationService(EventModelStorage()))
self.master.run()
self.client = NotificationClient(server_uri="localhost:50051")
