def __init__(self):
super().__init__()
self._yaml_config = None
with open(
os.path.dirname(os.path.abspath(__file__)) + '/source.yaml',
'r') as yaml_file:
self._yaml_config = yaml.load(yaml_file)
self._aiflow_client = AIFlowClient(
server_uri=self._yaml_config.get('master_uri'))
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 setUpClass(cls) -> None:
print("TestAIFlowClientMySQL setUpClass")
db_server_url = get_mysql_server_url()
cls.db_name = 'test_aiflow_client'
cls.engine = sqlalchemy.create_engine(db_server_url)
cls.engine.execute('DROP DATABASE IF EXISTS %s' % cls.db_name)
cls.engine.execute('CREATE DATABASE IF NOT EXISTS %s' % cls.db_name)
cls.store_uri = '%s/%s' % (db_server_url, cls.db_name)
cls.server = AIFlowServer(store_uri=cls.store_uri, port=_PORT)
cls.server.run()
test_client.client = AIFlowClient(server_uri='localhost:' + _PORT)
test_client.client1 = AIFlowClient(server_uri='localhost:' + _PORT)
test_client.client2 = AIFlowClient(server_uri='localhost:' + _PORT)
def setUpClass(cls) -> None:
if os.path.exists(_SQLITE_DB_FILE):
os.remove(_SQLITE_DB_FILE)
cls.server = AIFlowServer(store_uri=_SQLITE_DB_URI, port=_PORT)
cls.server.run()
global client
client = AIFlowClient(server_uri='localhost:' + _PORT)
def ensure_project_registered():
""" Ensure the project configured in project.yaml has been registered. """
client = AIFlowClient(server_uri=_default_project_config.get_master_uri())
project_meta = client.get_project_by_name(
_default_project_config.get_project_name())
pp = {}
for k, v in _default_project_config.items():
pp[k] = str(v)
if project_meta is None:
project_meta = client.register_project(
name=_default_project_config.get_project_name(), properties=pp)
else:
project_meta = client.update_project(
project_name=_default_project_config.get_project_name(),
properties=pp)
_default_project_config.set_project_uuid(str(project_meta.uuid))
def setUpClass(cls) -> None:
if os.path.exists(_SQLITE_DB_FILE):
os.remove(_SQLITE_DB_FILE)
cls.server = AIFlowServer(store_uri=_SQLITE_DB_URI, port=_PORT)
cls.sc_manager = cls.server.deploy_service.scheduler_manager
cls.server.run()
global client
client = AIFlowClient(server_uri='localhost:' + _PORT)
cls.ls_manager = cls.sc_manager.listener_manager
def get_client(self) -> Optional[AIFlowClient]:
"""
Get the ai flow client.
:return: AIFlowClient.
"""
if self.client is None:
self.client = AIFlowClient(
server_uri=project_config().get_master_uri())
return self.client
def start(self):
self.client = AIFlowClient(server_uri=self.server_uri)
self.listener_manager = ListenerManager(client=self.client)
self.listener_manager.start_dispatcher()
self.listener_manager.start_job_status_listener()
self.scheduler_pool = SchedulerPool(
client=self.client, listener_manager=self.listener_manager)
self.scheduler_pool.setDaemon(True)
self.scheduler_pool.setName("scheduler-pool")
self.scheduler_pool.start()
class TestTensorFlowIrisModel(unittest.TestCase):
def setUp(self) -> None:
if os.path.exists(_SQLITE_DB_FILE):
os.remove(_SQLITE_DB_FILE)
self.server = AIFlowServer(store_uri=_SQLITE_DB_URI, port=_PORT)
self.server.run()
self.client = AIFlowClient(server_uri='localhost:' + _PORT)
def tearDown(self) -> None:
self.client.stop_listen_event()
self.server.stop()
os.remove(_SQLITE_DB_FILE)
def test_save_and_load_model(self):
iris_model = fit_and_save_model()
tf_graph = tf.Graph()
registered_model = self.client.create_registered_model(
model_name='iris_model',
model_type=ModelType.SAVED_MODEL,
model_desc='iris model')
self.client.create_model_version(
model_name=registered_model.model_name,
model_path=iris_model.path,
model_metric='http://metric',
model_flavor=
'{"meta_graph_tags":["serve"],"signature_def_map_key":"predict"}',
version_desc='iris model')
class IrisWatcher(EventWatcher):
def process(self, notifications):
for notification in notifications:
model_path = json.loads(
notification.value).get('_model_path')
model_flavor = json.loads(
notification.value).get('_model_flavor')
signature_def = tensorflow.load_model(
model_uri=model_path,
meta_graph_tags=json.loads(model_flavor).get(
'meta_graph_tags'),
signature_def_map_key=json.loads(model_flavor).get(
'signature_def_map_key'),
tf_session=tf.Session(graph=tf_graph))
for _, input_signature in signature_def.inputs.items():
t_input = tf_graph.get_tensor_by_name(
input_signature.name)
assert t_input is not None
for _, output_signature in signature_def.outputs.items():
t_output = tf_graph.get_tensor_by_name(
output_signature.name)
assert t_output is not None
self.client.start_listen_event(key=registered_model.model_name,
watcher=IrisWatcher())
def get_ai_flow_client():
""" Get AI flow Client. """
global _default_ai_flow_client, _default_master_uri
ensure_project_registered()
if _default_ai_flow_client is None:
current_uri = _default_project_config.get_master_uri()
if current_uri is None:
return None
else:
_default_master_uri = current_uri
_default_ai_flow_client \
= AIFlowClient(server_uri=_default_master_uri,
notification_service_uri=_default_project_config.get_notification_service_uri())
return _default_ai_flow_client
else:
current_uri = _default_project_config.get_master_uri()
if current_uri != _default_master_uri:
_default_master_uri = current_uri
_default_ai_flow_client \
= AIFlowClient(server_uri=_default_master_uri,
notification_service_uri=_default_project_config.get_notification_service_uri())
return _default_ai_flow_client
class TestSklearnModel(unittest.TestCase):
def setUp(self) -> None:
if os.path.exists(_SQLITE_DB_FILE):
os.remove(_SQLITE_DB_FILE)
self.server = AIFlowServer(store_uri=_SQLITE_DB_URI, port=_PORT)
self.server.run()
self.client = AIFlowClient(server_uri='localhost:' + _PORT)
def tearDown(self) -> None:
self.client.stop_listen_event()
self.server.stop()
os.remove(_SQLITE_DB_FILE)
def test_save_and_load_model(self):
knn_model = fit_model()
model_path = tempfile.mkdtemp()
if not os.path.exists(model_path):
os.makedirs(model_path)
model_path = os.path.join(model_path, 'model.pkl')
save_model(sk_model=knn_model.model,
output_path=model_path,
serialization_format=SERIALIZATION_FORMAT_PICKLE)
registered_model = self.client.create_registered_model(
model_name='knn_model',
model_type=ModelType.SAVED_MODEL,
model_desc='knn model')
self.client.create_model_version(
model_name=registered_model.model_name,
model_path=model_path,
model_metric='http://metric',
model_flavor='sklearn',
version_desc='knn model')
class KnnWatcher(EventWatcher):
def process(self, notifications):
for notification in notifications:
load_path = json.loads(
notification.value).get('_model_path')
reloaded_knn_model = sklearn.load_model(
model_uri=load_path)
numpy.testing.assert_array_equal(
knn_model.model.predict(knn_model.inference_data),
reloaded_knn_model.predict(knn_model.inference_data))
os.remove(load_path)
self.client.start_listen_event(key=registered_model.model_name,
watcher=KnnWatcher())
def get_client(self):
self.client = AIFlowClient(
server_uri=project_config().get_master_uri())
return self.client
class Source(object):
"""
Generate inference online read example messages when listening to the source notification.
"""
def __init__(self):
super().__init__()
self._yaml_config = None
with open(
os.path.dirname(os.path.abspath(__file__)) + '/source.yaml',
'r') as yaml_file:
self._yaml_config = yaml.load(yaml_file)
self._aiflow_client = AIFlowClient(
server_uri=self._yaml_config.get('master_uri'))
def listen_notification(self):
class SourceWatcher(Watcher):
def __init__(self, yaml_config: Dict):
super().__init__()
self._yaml_config = yaml_config
def process(self, listener_name, notifications):
self.process_notification()
def process_notification(self):
bootstrap_servers = self._yaml_config.get('bootstrap_servers')
read_example_topic = self._yaml_config.get(
'read_example_topic')
write_example_topic = self._yaml_config.get(
'write_example_topic')
admin_client = KafkaAdminClient(
bootstrap_servers=bootstrap_servers)
topics = admin_client.list_topics()
if read_example_topic in topics:
process = Popen(args=[
'kafka-topics.sh',
'--bootstrap-server',
bootstrap_servers,
'--delete',
'--topic',
read_example_topic,
],
shell=False)
print('Delete kafka topic {} status: {}'.format(
read_example_topic, process.wait()))
if write_example_topic in topics:
process = Popen(args=[
'kafka-topics.sh',
'--bootstrap-server',
bootstrap_servers,
'--delete',
'--topic',
write_example_topic,
],
shell=False)
print('Delete kafka topic {} status: {}'.format(
write_example_topic, process.wait()))
# Create inference online read example topic.
admin_client.create_topics(new_topics=[
NewTopic(name=read_example_topic,
num_partitions=1,
replication_factor=1)
])
# Create inference vector write example topic.
admin_client.create_topics(new_topics=[
NewTopic(name=write_example_topic,
num_partitions=1,
replication_factor=1)
])
self.generate_read_example()
def generate_read_example(self):
"""
Generate inference online read example messages.
"""
bootstrap_servers = self._yaml_config.get('bootstrap_servers')
read_example_topic = self._yaml_config.get(
'read_example_topic')
# Read inference online read example.
df = pd.read_csv(
filepath_or_buffer=self._yaml_config.get('dataset_uri'),
delimiter=';',
header=None)
producer = KafkaProducer(bootstrap_servers=[bootstrap_servers])
for index, row in df.iterrows():
print(
'Send inference online read example message: topic=%s, key=%s, value=%s'
% (read_example_topic, row.get(1), '%s,%s,%s' %
(row.get(1), row.get(2), row.get(3))))
# Send inference online read example messages.
producer.send(
read_example_topic,
key=bytes(row.get(1), encoding='utf8'),
value=bytes('%s,%s,%s' %
(row.get(1), row.get(2), row.get(3)),
encoding='utf8'))
time.sleep(self._yaml_config.get('time_interval') / 1000)
self._aiflow_client.start_listen_notification(
listener_name='source_listener',
key=self._yaml_config.get('notification_key'),
watcher=SourceWatcher(self._yaml_config))
def setUp(self) -> None:
if os.path.exists(_SQLITE_DB_FILE):
os.remove(_SQLITE_DB_FILE)
self.server = AIFlowServer(store_uri=_SQLITE_DB_URI, port=_PORT)
self.server.run()
self.client = AIFlowClient(server_uri='localhost:' + _PORT)
class TestHighAvailableAIFlowServer(unittest.TestCase):
@staticmethod
def start_aiflow_server(host, port):
port = str(port)
server_uri = host + ":" + port
server = HighAvailableAIFlowServer(
store_uri=_SQLITE_DB_URI, port=port,
server_uri=server_uri)
server.run()
return server
def wait_for_new_members_detected(self, new_member_uri):
for i in range(100):
living_member = self.client.living_aiflow_members
if new_member_uri in living_member:
break
else:
time.sleep(0.1)
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 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 test_server_change(self) -> None:
self.client.register_project("test_project")
projects = self.client.list_project(10, 0)
self.assertEqual(self.client.current_aiflow_uri, "localhost:50051")
self.assertEqual(projects[0].name, "test_project")
self.server2 = self.start_aiflow_server("localhost", 50052)
self.wait_for_new_members_detected("localhost:50052")
self.server1.stop()
projects = self.client.list_project(10, 0)
self.assertEqual(self.client.current_aiflow_uri, "localhost:50052")
self.assertEqual(projects[0].name, "test_project")
self.server3 = self.start_aiflow_server("localhost", 50053)
self.wait_for_new_members_detected("localhost:50053")
self.server2.stop()
projects = self.client.list_project(10, 0)
self.assertEqual(self.client.current_aiflow_uri, "localhost:50053")
self.assertEqual(projects[0].name, "test_project")