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_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))