def _clear_db(self):
if self.master_config.get_db_type() == DBType.SQLITE:
store = SqlAlchemyStore(self.master_config.get_db_uri())
base.metadata.drop_all(store.db_engine)
base.metadata.create_all(store.db_engine)
elif self.master_config.get_db_type() == DBType.MONGODB:
MongoStoreConnManager().drop_all()
def _clear_db(self):
if self.db_type == DBType.SQLITE:
store = SqlAlchemyStore(self.store_uri)
base.metadata.drop_all(store.db_engine)
base.metadata.create_all(store.db_engine)
elif self.db_type == DBType.MONGODB:
MongoStoreConnManager().drop_all()
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 tearDown(self) -> None:
print('tearDown')
store = SqlAlchemyStore(_SQLITE_DB_URI)
base.metadata.drop_all(store.db_engine)
base.metadata.create_all(store.db_engine)
af.default_graph().clear_graph()
res = client.list_job(page_size=10, offset=0)
self.assertIsNone(res)
def tearDown(self) -> None:
self.client.stop_listen_event()
self.client.disable_high_availability()
if self.server1 is not None:
self.server1.stop()
if self.server2 is not None:
self.server2.stop()
if self.server3 is not None:
self.server3.stop()
store = SqlAlchemyStore(_SQLITE_DB_URI)
base.metadata.drop_all(store.db_engine)
def stop(self, clear_sql_lite_db_file=True) -> None:
"""
Stop the AI flow master.
:param clear_sql_lite_db_file: If True, the sqlite database files will be deleted When the server stops working.
"""
self.server.stop()
if self.master_config.get_db_type() == DBType.SQLITE and clear_sql_lite_db_file:
store = SqlAlchemyStore(self.master_config.get_db_uri())
base.metadata.drop_all(store.db_engine)
os.remove(self.master_config.get_sql_lite_db_file())
def setUp(self) -> None:
SqlAlchemyStore(_SQLITE_DB_URI)
self.server1 = HighAvailableAIFlowServer(
store_uri=_SQLITE_DB_URI, port=50051,
server_uri='localhost:50051')
self.server1.run()
self.server2 = None
self.server3 = None
self.config = ProjectConfig()
self.config.set_enable_ha(True)
self.client = AIFlowClient(server_uri='localhost:50052,localhost:50051', project_config=self.config)
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, 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 stop(self, clear_sql_lite_db_file=False):
self.executor.shutdown()
self.server.stop(0)
if self.enabled_ha:
self.ha_service.stop()
if self.db_type == DBType.SQLITE and clear_sql_lite_db_file:
store = SqlAlchemyStore(self.store_uri)
base.metadata.drop_all(store.db_engine)
os.remove(self.store_uri[10:])
elif self.db_type == DBType.MONGODB:
MongoStoreConnManager().disconnect_all()
logging.info('AIFlow server stopped.')
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 __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:
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 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 tearDownClass(cls) -> None:
client.stop_listen_event()
store = SqlAlchemyStore(_SQLITE_DB_URI)
base.metadata.drop_all(store.db_engine)
os.remove(_SQLITE_DB_FILE)
def _clear_db(self):
if self.master_config.get_db_type() == DBType.SQLITE:
store = SqlAlchemyStore(self.master_config.get_db_uri())
base.metadata.drop_all(store.db_engine)
base.metadata.create_all(store.db_engine)
def __init__(self, db_uri, server_uri):
db_uri = db_uri
self.store = SqlAlchemyStore(db_uri)
self.model_center_client = ModelCenterClient(server_uri)
def __init__(self, db_uri):
self.db_uri = db_uri
self.store = SqlAlchemyStore(db_uri)
def _get_store(db_uri=''):
return SqlAlchemyStore(db_uri)
def __init__(self, store_uri, server_uri, notification_uri=None):
self.model_repo_store = SqlAlchemyStore(store_uri)
if notification_uri is None:
self.notification_client = NotificationClient(server_uri)
else:
self.notification_client = NotificationClient(notification_uri)