def __init__(self,
store_uri=None,
port=_PORT,
start_default_notification: bool = True,
notification_uri=None,
ha_manager=None,
server_uri=None,
ha_storage=None,
ttl_ms: int = 10000):
super(HighAvailableAIFlowServer,
self).__init__(store_uri, port, start_default_notification,
notification_uri)
if ha_manager is None:
ha_manager = SimpleAIFlowServerHaManager()
if server_uri is None:
raise ValueError("server_uri is required!")
if ha_storage is None:
db_engine = extract_db_engine_from_uri(store_uri)
if DBType.value_of(db_engine) == DBType.MONGODB:
username, password, host, port, db = parse_mongo_uri(store_uri)
ha_storage = MongoStore(host=host,
port=int(port),
username=username,
password=password,
db=db)
else:
ha_storage = SqlAlchemyStore(store_uri)
self.ha_service = HighAvailableService(ha_manager, server_uri,
ha_storage, ttl_ms)
add_HighAvailabilityManagerServicer_to_server(self.ha_service,
self.server)
def __init__(self, backend_store_uri):
db_engine = extract_db_engine_from_uri(backend_store_uri)
if DBType.value_of(db_engine) == DBType.MONGODB:
username, password, host, port, db = parse_mongo_uri(
backend_store_uri)
super().__init__(storage=MongoEventStorage(host=host,
port=int(port),
username=username,
password=password,
db=db))
else:
super().__init__(storage=DbEventStorage(backend_store_uri))
def __init__(self, db_uri, server_uri):
db_engine = extract_db_engine_from_uri(db_uri)
if DBType.value_of(db_engine) == DBType.MONGODB:
username, password, host, port, db = parse_mongo_uri(db_uri)
self.store = MongoStore(host=host,
port=int(port),
username=username,
password=password,
db=db)
else:
self.store = SqlAlchemyStore(db_uri)
self.model_center_client = ModelCenterClient(server_uri)
def __init__(self, db_uri):
self.db_uri = db_uri
db_engine = extract_db_engine_from_uri(self.db_uri)
if DBType.value_of(db_engine) == DBType.MONGODB:
username, password, host, port, db = parse_mongo_uri(self.db_uri)
self.store = MongoStore(host=host,
port=int(port),
username=username,
password=password,
db=db)
else:
self.store = SqlAlchemyStore(self.db_uri)
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 __init__(self, store_uri, notification_uri=None):
db_engine = extract_db_engine_from_uri(store_uri)
if DBType.value_of(db_engine) == DBType.MONGODB:
username, password, host, port, db = parse_mongo_uri(store_uri)
self.model_repo_store = MongoStore(host=host,
port=int(port),
username=username,
password=password,
db=db)
else:
self.model_repo_store = SqlAlchemyStore(store_uri)
self.notification_client = None
if notification_uri is not None:
self.notification_client = NotificationClient(notification_uri, default_namespace=DEFAULT_NAMESPACE)
def _add_ha_service(self, ha_manager, ha_server_uri, ha_storage, store_uri,
ttl_ms):
if ha_manager is None:
ha_manager = SimpleAIFlowServerHaManager()
if ha_server_uri is None:
raise ValueError("ha_server_uri is required with ha enabled!")
if ha_storage is None:
db_engine = extract_db_engine_from_uri(store_uri)
if DBType.value_of(db_engine) == DBType.MONGODB:
username, password, host, port, db = parse_mongo_uri(store_uri)
ha_storage = MongoStore(host=host,
port=int(port),
username=username,
password=password,
db=db)
else:
ha_storage = SqlAlchemyStore(store_uri)
self.ha_service = HighAvailableService(ha_manager, ha_server_uri,
ha_storage, ttl_ms)
add_HighAvailabilityManagerServicer_to_server(self.ha_service,
self.server)
def create_default_sever_config(root_dir_path, db_uri, airflow_deploy_path):
"""
Generate default server config which use Apache Airflow as scheduler.
"""
from ai_flow.store.db.db_util import extract_db_engine_from_uri
db_type = extract_db_engine_from_uri(db_uri)
content = textwrap.dedent(f"""\
# Config of master server
# endpoint of AI Flow Server
server_ip: localhost
server_port: 50051
# uri of database backend of AIFlow server
db_uri: {db_uri}
# type of database backend in master
db_type: {db_type}
# whether to start the scheduler service
start_scheduler_service: True
# Whether to start default notification service, if not a custom URI should be set
start_default_notification: False
notification_uri: localhost:50052
# scheduler config
scheduler:
scheduler_class_name: ai_flow_plugins.scheduler_plugins.airflow.airflow_scheduler.AirFlowScheduler
scheduler_config:
airflow_deploy_path: {airflow_deploy_path}
notification_service_uri: localhost:50052
""")
master_yaml_path = root_dir_path + "/master.yaml"
with open(master_yaml_path, "w") as f:
f.write(content)
return master_yaml_path
