def setUp(self):
data_source = common_pb.DataSource()
data_source.data_source_meta.name = "milestone-x"
data_source.data_source_meta.partition_num = 4
data_source.data_source_meta.start_time = 0
data_source.data_source_meta.end_time = 10000
data_source.output_base_dir = "./ds_output"
data_source.role = common_pb.FLRole.Follower
self.data_source = data_source
self.db_database = 'test_cluster'
self.db_addr = 'localhost:2379'
self.db_base_dir = 'fedlearner'
self.db_username = '******'
self.db_password = '******'
self.kvstore = mysql_client.DBClient(self.db_database, self.db_addr,
self.db_username,
self.db_password,
self.db_base_dir, True)
common.commit_data_source(self.kvstore, self.data_source)
if gfile.Exists(data_source.output_base_dir):
gfile.DeleteRecursively(data_source.output_base_dir)
self.data_block_matas = []
self.manifest_manager = raw_data_manifest_manager.RawDataManifestManager(
self.kvstore, self.data_source)
partition_num = self.data_source.data_source_meta.partition_num
for i in range(partition_num):
self._create_data_block(i)
def _setUpDataSource(self):
self._data_source_name = 'test_data_source'
self._etcd_l.delete_prefix(self._data_source_name)
self._etcd_f.delete_prefix(self._data_source_name)
self._data_source_l = common_pb.DataSource()
self._data_source_l.role = common_pb.FLRole.Leader
self._data_source_l.state = common_pb.DataSourceState.Init
self._data_source_l.data_block_dir = "./data_block_l"
self._data_source_l.raw_data_dir = "./raw_data_l"
self._data_source_l.example_dumped_dir = "./example_dumped_l"
self._data_source_l.raw_data_sub_dir = "./raw_data_sub_dir_l"
self._data_source_f = common_pb.DataSource()
self._data_source_f.role = common_pb.FLRole.Follower
self._data_source_f.state = common_pb.DataSourceState.Init
self._data_source_f.data_block_dir = "./data_block_f"
self._data_source_f.raw_data_dir = "./raw_data_f"
self._data_source_f.example_dumped_dir = "./example_dumped_f"
self._data_source_f.raw_data_sub_dir = "./raw_data_sub_dir_f"
data_source_meta = common_pb.DataSourceMeta()
data_source_meta.name = self._data_source_name
data_source_meta.partition_num = 4
data_source_meta.start_time = 0
data_source_meta.end_time = 100000000
self._data_source_l.data_source_meta.MergeFrom(data_source_meta)
self._data_source_f.data_source_meta.MergeFrom(data_source_meta)
common.commit_data_source(self._etcd_l, self._data_source_l)
common.commit_data_source(self._etcd_f, self._data_source_f)
def _setUpDataSource(self):
self._data_source_name = 'test_data_source'
self._kvstore_l.delete_prefix(
common.data_source_kvstore_base_dir(self._data_source_name))
self._kvstore_f.delete_prefix(
common.data_source_kvstore_base_dir(self._data_source_name))
self._data_source_l = common_pb.DataSource()
self._data_source_l.role = common_pb.FLRole.Leader
self._data_source_l.state = common_pb.DataSourceState.Init
self._data_source_l.output_base_dir = "./ds_output_l"
self._raw_data_dir_l = "./raw_data_l"
self._data_source_l.raw_data_sub_dir = "./raw_data_sub_dir_l"
self._data_source_f = common_pb.DataSource()
self._data_source_f.role = common_pb.FLRole.Follower
self._data_source_f.state = common_pb.DataSourceState.Init
self._data_source_f.output_base_dir = "./ds_output_f"
self._raw_data_dir_f = "./raw_data_f"
self._data_source_f.raw_data_sub_dir = "./raw_data_sub_dir_f"
data_source_meta = common_pb.DataSourceMeta()
data_source_meta.name = self._data_source_name
data_source_meta.partition_num = 4
data_source_meta.start_time = 0
data_source_meta.end_time = 100000000
self._data_source_l.data_source_meta.MergeFrom(data_source_meta)
self._data_source_f.data_source_meta.MergeFrom(data_source_meta)
common.commit_data_source(self._kvstore_l, self._data_source_l)
common.commit_data_source(self._kvstore_f, self._data_source_f)
def _update_data_source(self, data_source):
self._data_source = None
try:
commit_data_source(self._etcd, data_source)
except Exception as e:
logging.error("Failed to update data source: %s since "\
"exception: %s", self._data_source_name, e)
raise
self._data_source = data_source
logging.debug("Success update to update data source: %s.",
self._data_source_name)
def setUp(self):
self.sche = _TaskScheduler(30)
self.kv_store = [None, None]
self.app_id = "test_trainer_v1"
db_database, db_addr, db_username, db_password, db_base_dir = \
get_kvstore_config("etcd")
data_source = [
self._gen_ds_meta(common_pb.FLRole.Leader),
self._gen_ds_meta(common_pb.FLRole.Follower)
]
for role in range(2):
self.kv_store[role] = mysql_client.DBClient(
data_source[role].data_source_meta.name, db_addr, db_username,
db_password, db_base_dir, True)
self.data_source = data_source
(x, y) = (None, None)
if debug_mode:
(x, y), _ = tf.keras.datasets.mnist.load_data(local_mnist_path)
else:
(x, y), _ = tf.keras.datasets.mnist.load_data()
x = x[:200, ]
x = x.reshape(x.shape[0], -1).astype(np.float32) / 255.0
y = y.astype(np.int64)
xl = x[:, :x.shape[1] // 2]
xf = x[:, x.shape[1] // 2:]
self._create_local_data(xl, xf, y)
x = [xl, xf]
for role in range(2):
common.commit_data_source(self.kv_store[role], data_source[role])
if gfile.Exists(data_source[role].output_base_dir):
gfile.DeleteRecursively(data_source[role].output_base_dir)
manifest_manager = raw_data_manifest_manager.RawDataManifestManager(
self.kv_store[role], data_source[role])
partition_num = data_source[role].data_source_meta.partition_num
for i in range(partition_num):
self._create_data_block(data_source[role], i, x[role], y)
#x[role], y if role == 0 else None)
manifest_manager._finish_partition(
'join_example_rep', dj_pb.JoinExampleState.UnJoined,
dj_pb.JoinExampleState.Joined, -1, i)
def __init__(self,
base_path,
name,
role,
partition_num=1,
start_time=0,
end_time=100000):
if role == 'leader':
role = 0
elif role == 'follower':
role = 1
else:
raise ValueError("Unknown role %s" % role)
data_source = common_pb.DataSource()
data_source.data_source_meta.name = name
data_source.data_source_meta.partition_num = partition_num
data_source.data_source_meta.start_time = start_time
data_source.data_source_meta.end_time = end_time
data_source.output_base_dir = "{}/{}_{}/data_source/".format(
base_path, data_source.data_source_meta.name, role)
data_source.role = role
if gfile.Exists(data_source.output_base_dir):
gfile.DeleteRecursively(data_source.output_base_dir)
self._data_source = data_source
db_database, db_addr, db_username, db_password, db_base_dir = \
get_kvstore_config("etcd")
self._kv_store = mysql_client.DBClient(db_database, db_addr,
db_username, db_password,
db_base_dir, True)
common.commit_data_source(self._kv_store, self._data_source)
self._dbms = []
for i in range(partition_num):
manifest_manager = raw_data_manifest_manager.RawDataManifestManager(
self._kv_store, self._data_source)
manifest_manager._finish_partition('join_example_rep',
dj_pb.JoinExampleState.UnJoined,
dj_pb.JoinExampleState.Joined,
-1, i)
self._dbms.append(
data_block_manager.DataBlockManager(self._data_source, i))
def setUp(self):
data_source = common_pb.DataSource()
data_source.data_source_meta.name = "milestone-x"
data_source.data_source_meta.partition_num = 4
data_source.data_source_meta.start_time = 0
data_source.data_source_meta.end_time = 10000
data_source.output_base_dir = "./ds_output"
data_source.role = common_pb.FLRole.Follower
self.data_source = data_source
self.kvstore = db_client.DBClient('etcd', True)
common.commit_data_source(self.kvstore, self.data_source)
if gfile.Exists(data_source.output_base_dir):
gfile.DeleteRecursively(data_source.output_base_dir)
self.data_block_matas = []
self.manifest_manager = raw_data_manifest_manager.RawDataManifestManager(
self.kvstore, self.data_source)
partition_num = self.data_source.data_source_meta.partition_num
for i in range(partition_num):
self._create_data_block(i)
def setUp(self):
data_source = common_pb.DataSource()
data_source.data_source_meta.name = "milestone-x"
data_source.data_source_meta.partition_num = 4
data_source.data_source_meta.start_time = 0
data_source.data_source_meta.end_time = 10000
data_source.data_block_dir = "./data_block"
data_source.role = common_pb.FLRole.Follower
self.data_source = data_source
self.etcd_name = 'test_cluster'
self.etcd_addrs = 'localhost:2379'
self.etcd_base_dir = 'fedlearner'
self.etcd = etcd_client.EtcdClient(self.etcd_name, self.etcd_addrs,
self.etcd_base_dir, True)
common.commit_data_source(self.etcd, self.data_source)
if gfile.Exists(data_source.data_block_dir):
gfile.DeleteRecursively(data_source.data_block_dir)
self.data_block_matas = []
self.manifest_manager = raw_data_manifest_manager.RawDataManifestManager(
self.etcd, self.data_source)
partition_num = self.data_source.data_source_meta.partition_num
for i in range(partition_num):
self._create_data_block(i)
def setUp(self):
etcd_name = 'test_etcd'
etcd_addrs = 'localhost:2379'
etcd_base_dir_l = 'byefl_l'
etcd_base_dir_f = 'byefl_f'
data_source_name = 'test_data_source'
etcd_l = EtcdClient(etcd_name, etcd_addrs, etcd_base_dir_l, True)
etcd_f = EtcdClient(etcd_name, etcd_addrs, etcd_base_dir_f, True)
etcd_l.delete_prefix(
common.data_source_etcd_base_dir(data_source_name))
etcd_f.delete_prefix(
common.data_source_etcd_base_dir(data_source_name))
data_source_l = common_pb.DataSource()
self.raw_data_pub_dir_l = './raw_data_pub_dir_l'
data_source_l.raw_data_sub_dir = self.raw_data_pub_dir_l
data_source_l.role = common_pb.FLRole.Leader
data_source_l.state = common_pb.DataSourceState.Init
data_source_l.data_block_dir = "./data_block_l"
data_source_l.raw_data_dir = "./raw_data_l"
data_source_l.example_dumped_dir = "./example_dumped_l"
data_source_f = common_pb.DataSource()
self.raw_data_pub_dir_f = './raw_data_pub_dir_f'
data_source_f.role = common_pb.FLRole.Follower
data_source_f.raw_data_sub_dir = self.raw_data_pub_dir_f
data_source_f.state = common_pb.DataSourceState.Init
data_source_f.data_block_dir = "./data_block_f"
data_source_f.raw_data_dir = "./raw_data_f"
data_source_f.example_dumped_dir = "./example_dumped_f"
data_source_meta = common_pb.DataSourceMeta()
data_source_meta.name = data_source_name
data_source_meta.partition_num = 2
data_source_meta.start_time = 0
data_source_meta.end_time = 100000000
data_source_l.data_source_meta.MergeFrom(data_source_meta)
common.commit_data_source(etcd_l, data_source_l)
data_source_f.data_source_meta.MergeFrom(data_source_meta)
common.commit_data_source(etcd_f, data_source_f)
master_options = dj_pb.DataJoinMasterOptions(use_mock_etcd=True)
master_addr_l = 'localhost:4061'
master_addr_f = 'localhost:4062'
master_l = data_join_master.DataJoinMasterService(
int(master_addr_l.split(':')[1]),
master_addr_f,
data_source_name,
etcd_name,
etcd_base_dir_l,
etcd_addrs,
master_options,
)
master_l.start()
master_f = data_join_master.DataJoinMasterService(
int(master_addr_f.split(':')[1]), master_addr_l, data_source_name,
etcd_name, etcd_base_dir_f, etcd_addrs, master_options)
master_f.start()
channel_l = make_insecure_channel(master_addr_l, ChannelType.INTERNAL)
master_client_l = dj_grpc.DataJoinMasterServiceStub(channel_l)
channel_f = make_insecure_channel(master_addr_f, ChannelType.INTERNAL)
master_client_f = dj_grpc.DataJoinMasterServiceStub(channel_f)
while True:
req_l = dj_pb.DataSourceRequest(
data_source_meta=data_source_l.data_source_meta)
req_f = dj_pb.DataSourceRequest(
data_source_meta=data_source_f.data_source_meta)
dss_l = master_client_l.GetDataSourceStatus(req_l)
dss_f = master_client_f.GetDataSourceStatus(req_f)
self.assertEqual(dss_l.role, common_pb.FLRole.Leader)
self.assertEqual(dss_f.role, common_pb.FLRole.Follower)
if dss_l.state == common_pb.DataSourceState.Processing and \
dss_f.state == common_pb.DataSourceState.Processing:
break
else:
time.sleep(2)
self.master_client_l = master_client_l
self.master_client_f = master_client_f
self.master_addr_l = master_addr_l
self.master_addr_f = master_addr_f
self.etcd_l = etcd_l
self.etcd_f = etcd_f
self.data_source_l = data_source_l
self.data_source_f = data_source_f
self.master_l = master_l
self.master_f = master_f
self.data_source_name = data_source_name,
self.etcd_name = etcd_name
self.etcd_addrs = etcd_addrs
self.etcd_base_dir_l = etcd_base_dir_l
self.etcd_base_dir_f = etcd_base_dir_f
self.raw_data_publisher_l = raw_data_publisher.RawDataPublisher(
self.etcd_l, self.raw_data_pub_dir_l)
self.raw_data_publisher_f = raw_data_publisher.RawDataPublisher(
self.etcd_f, self.raw_data_pub_dir_f)
if gfile.Exists(data_source_l.data_block_dir):
gfile.DeleteRecursively(data_source_l.data_block_dir)
if gfile.Exists(data_source_l.example_dumped_dir):
gfile.DeleteRecursively(data_source_l.example_dumped_dir)
if gfile.Exists(data_source_l.raw_data_dir):
gfile.DeleteRecursively(data_source_l.raw_data_dir)
if gfile.Exists(data_source_f.data_block_dir):
gfile.DeleteRecursively(data_source_f.data_block_dir)
if gfile.Exists(data_source_f.example_dumped_dir):
gfile.DeleteRecursively(data_source_f.example_dumped_dir)
if gfile.Exists(data_source_f.raw_data_dir):
gfile.DeleteRecursively(data_source_f.raw_data_dir)
self.worker_options = dj_pb.DataJoinWorkerOptions(
use_mock_etcd=True,
raw_data_options=dj_pb.RawDataOptions(raw_data_iter='TF_RECORD',
compressed_type=''),
example_id_dump_options=dj_pb.ExampleIdDumpOptions(
example_id_dump_interval=1, example_id_dump_threshold=1024),
example_joiner_options=dj_pb.ExampleJoinerOptions(
example_joiner='STREAM_JOINER',
min_matching_window=64,
max_matching_window=256,
data_block_dump_interval=30,
data_block_dump_threshold=1000),
batch_processor_options=dj_pb.BatchProcessorOptions(
batch_size=512, max_flying_item=2048),
data_block_builder_options=dj_pb.WriterOptions(
output_writer='TF_RECORD'))
self.total_index = 1 << 13
def setUp(self):
etcd_name = 'test_etcd'
etcd_addrs = 'localhost:2379'
etcd_base_dir_l = 'byefl_l'
etcd_base_dir_f= 'byefl_f'
data_source_name = 'test_data_source'
etcd_l = EtcdClient(etcd_name, etcd_addrs, etcd_base_dir_l, True)
etcd_f = EtcdClient(etcd_name, etcd_addrs, etcd_base_dir_f, True)
etcd_l.delete_prefix(common.data_source_etcd_base_dir(data_source_name))
etcd_f.delete_prefix(common.data_source_etcd_base_dir(data_source_name))
data_source_l = common_pb.DataSource()
self.raw_data_pub_dir_l = './raw_data_pub_dir_l'
data_source_l.raw_data_sub_dir = self.raw_data_pub_dir_l
data_source_l.role = common_pb.FLRole.Leader
data_source_l.state = common_pb.DataSourceState.Init
data_source_l.output_base_dir = "./ds_output_l"
self.raw_data_dir_l = "./raw_data_l"
data_source_f = common_pb.DataSource()
self.raw_data_pub_dir_f = './raw_data_pub_dir_f'
data_source_f.role = common_pb.FLRole.Follower
data_source_f.raw_data_sub_dir = self.raw_data_pub_dir_f
data_source_f.state = common_pb.DataSourceState.Init
data_source_f.output_base_dir = "./ds_output_f"
self.raw_data_dir_f = "./raw_data_f"
data_source_meta = common_pb.DataSourceMeta()
data_source_meta.name = data_source_name
data_source_meta.partition_num = 2
data_source_meta.start_time = 0
data_source_meta.end_time = 100000000
data_source_l.data_source_meta.MergeFrom(data_source_meta)
common.commit_data_source(etcd_l, data_source_l)
data_source_f.data_source_meta.MergeFrom(data_source_meta)
common.commit_data_source(etcd_f, data_source_f)
self.etcd_l = etcd_l
self.etcd_f = etcd_f
self.data_source_l = data_source_l
self.data_source_f = data_source_f
self.data_source_name = data_source_name
self.etcd_name = etcd_name
self.etcd_addrs = etcd_addrs
self.etcd_base_dir_l = etcd_base_dir_l
self.etcd_base_dir_f = etcd_base_dir_f
self.raw_data_publisher_l = raw_data_publisher.RawDataPublisher(
self.etcd_l, self.raw_data_pub_dir_l
)
self.raw_data_publisher_f = raw_data_publisher.RawDataPublisher(
self.etcd_f, self.raw_data_pub_dir_f
)
if gfile.Exists(data_source_l.output_base_dir):
gfile.DeleteRecursively(data_source_l.output_base_dir)
if gfile.Exists(self.raw_data_dir_l):
gfile.DeleteRecursively(self.raw_data_dir_l)
if gfile.Exists(data_source_f.output_base_dir):
gfile.DeleteRecursively(data_source_f.output_base_dir)
if gfile.Exists(self.raw_data_dir_f):
gfile.DeleteRecursively(self.raw_data_dir_f)
self.worker_options = dj_pb.DataJoinWorkerOptions(
use_mock_etcd=True,
raw_data_options=dj_pb.RawDataOptions(
raw_data_iter='TF_RECORD',
read_ahead_size=1<<20,
read_batch_size=128
),
example_id_dump_options=dj_pb.ExampleIdDumpOptions(
example_id_dump_interval=1,
example_id_dump_threshold=1024
),
example_joiner_options=dj_pb.ExampleJoinerOptions(
example_joiner='STREAM_JOINER',
min_matching_window=64,
max_matching_window=256,
data_block_dump_interval=30,
data_block_dump_threshold=1000
),
batch_processor_options=dj_pb.BatchProcessorOptions(
batch_size=512,
max_flying_item=2048
),
data_block_builder_options=dj_pb.WriterOptions(
output_writer='TF_RECORD'
)
)
self.total_index = 1 << 12
def test_api(self):
logging.getLogger().setLevel(logging.DEBUG)
etcd_name = 'test_etcd'
etcd_addrs = 'localhost:2379'
etcd_base_dir_l = 'byefl_l'
etcd_base_dir_f = 'byefl_f'
data_source_name = 'test_data_source'
etcd_l = EtcdClient(etcd_name, etcd_addrs, etcd_base_dir_l, True)
etcd_f = EtcdClient(etcd_name, etcd_addrs, etcd_base_dir_f, True)
etcd_l.delete_prefix(
common.data_source_etcd_base_dir(data_source_name))
etcd_f.delete_prefix(
common.data_source_etcd_base_dir(data_source_name))
data_source_l = common_pb.DataSource()
data_source_l.role = common_pb.FLRole.Leader
data_source_l.state = common_pb.DataSourceState.Init
data_source_l.output_base_dir = "./ds_output_l"
data_source_f = common_pb.DataSource()
data_source_f.role = common_pb.FLRole.Follower
data_source_f.state = common_pb.DataSourceState.Init
data_source_f.output_base_dir = "./ds_output_f"
data_source_meta = common_pb.DataSourceMeta()
data_source_meta.name = data_source_name
data_source_meta.partition_num = 1
data_source_meta.start_time = 0
data_source_meta.end_time = 100000000
data_source_l.data_source_meta.MergeFrom(data_source_meta)
common.commit_data_source(etcd_l, data_source_l)
data_source_f.data_source_meta.MergeFrom(data_source_meta)
common.commit_data_source(etcd_f, data_source_f)
master_addr_l = 'localhost:4061'
master_addr_f = 'localhost:4062'
options = dj_pb.DataJoinMasterOptions(use_mock_etcd=True)
master_l = data_join_master.DataJoinMasterService(
int(master_addr_l.split(':')[1]), master_addr_f, data_source_name,
etcd_name, etcd_base_dir_l, etcd_addrs, options)
master_l.start()
master_f = data_join_master.DataJoinMasterService(
int(master_addr_f.split(':')[1]), master_addr_l, data_source_name,
etcd_name, etcd_base_dir_f, etcd_addrs, options)
master_f.start()
channel_l = make_insecure_channel(master_addr_l, ChannelType.INTERNAL)
client_l = dj_grpc.DataJoinMasterServiceStub(channel_l)
channel_f = make_insecure_channel(master_addr_f, ChannelType.INTERNAL)
client_f = dj_grpc.DataJoinMasterServiceStub(channel_f)
while True:
req_l = dj_pb.DataSourceRequest(
data_source_meta=data_source_l.data_source_meta)
req_f = dj_pb.DataSourceRequest(
data_source_meta=data_source_f.data_source_meta)
dss_l = client_l.GetDataSourceStatus(req_l)
dss_f = client_f.GetDataSourceStatus(req_f)
self.assertEqual(dss_l.role, common_pb.FLRole.Leader)
self.assertEqual(dss_f.role, common_pb.FLRole.Follower)
if dss_l.state == common_pb.DataSourceState.Processing and \
dss_f.state == common_pb.DataSourceState.Processing:
break
else:
time.sleep(2)
rdreq = dj_pb.RawDataRequest(
data_source_meta=data_source_f.data_source_meta,
rank_id=0,
partition_id=-1,
join_example=empty_pb2.Empty())
rdrsp = client_f.RequestJoinPartition(rdreq)
self.assertTrue(rdrsp.status.code == 0)
self.assertTrue(rdrsp.HasField('manifest'))
self.assertEqual(rdrsp.manifest.join_example_rep.rank_id, 0)
self.assertEqual(rdrsp.manifest.join_example_rep.state,
dj_pb.JoinExampleState.Joining)
self.assertEqual(rdrsp.manifest.partition_id, 0)
#check idempotent
rdrsp = client_f.RequestJoinPartition(rdreq)
self.assertTrue(rdrsp.status.code == 0)
self.assertTrue(rdrsp.HasField('manifest'))
self.assertEqual(rdrsp.manifest.join_example_rep.rank_id, 0)
self.assertEqual(rdrsp.manifest.join_example_rep.state,
dj_pb.JoinExampleState.Joining)
self.assertEqual(rdrsp.manifest.partition_id, 0)
rdreq = dj_pb.RawDataRequest(
data_source_meta=data_source_l.data_source_meta,
rank_id=0,
partition_id=0,
join_example=empty_pb2.Empty())
rdrsp = client_l.RequestJoinPartition(rdreq)
self.assertTrue(rdrsp.status.code == 0)
self.assertTrue(rdrsp.HasField('manifest'))
self.assertEqual(rdrsp.manifest.join_example_rep.rank_id, 0)
self.assertEqual(rdrsp.manifest.join_example_rep.state,
dj_pb.JoinExampleState.Joining)
self.assertEqual(rdrsp.manifest.partition_id, 0)
#check idempotent
rdrsp = client_l.RequestJoinPartition(rdreq)
self.assertTrue(rdrsp.status.code == 0)
self.assertTrue(rdrsp.HasField('manifest'))
self.assertEqual(rdrsp.manifest.join_example_rep.rank_id, 0)
self.assertEqual(rdrsp.manifest.join_example_rep.state,
dj_pb.JoinExampleState.Joining)
self.assertEqual(rdrsp.manifest.partition_id, 0)
rdreq = dj_pb.RawDataRequest(
data_source_meta=data_source_l.data_source_meta,
rank_id=1,
partition_id=-1,
sync_example_id=empty_pb2.Empty())
rdrsp = client_l.RequestJoinPartition(rdreq)
self.assertEqual(rdrsp.status.code, 0)
self.assertTrue(rdrsp.HasField('manifest'))
self.assertEqual(rdrsp.manifest.sync_example_id_rep.rank_id, 1)
self.assertEqual(rdrsp.manifest.sync_example_id_rep.state,
dj_pb.SyncExampleIdState.Syncing)
self.assertEqual(rdrsp.manifest.partition_id, 0)
#check idempotent
rdrsp = client_l.RequestJoinPartition(rdreq)
self.assertEqual(rdrsp.status.code, 0)
self.assertTrue(rdrsp.HasField('manifest'))
self.assertEqual(rdrsp.manifest.sync_example_id_rep.rank_id, 1)
self.assertEqual(rdrsp.manifest.sync_example_id_rep.state,
dj_pb.SyncExampleIdState.Syncing)
self.assertEqual(rdrsp.manifest.partition_id, 0)
rdreq = dj_pb.RawDataRequest(
data_source_meta=data_source_f.data_source_meta,
rank_id=1,
partition_id=0,
sync_example_id=empty_pb2.Empty())
rdrsp = client_f.RequestJoinPartition(rdreq)
self.assertEqual(rdrsp.status.code, 0)
self.assertTrue(rdrsp.HasField('manifest'))
self.assertEqual(rdrsp.manifest.sync_example_id_rep.rank_id, 1)
self.assertEqual(rdrsp.manifest.sync_example_id_rep.state,
dj_pb.SyncExampleIdState.Syncing)
self.assertEqual(rdrsp.manifest.partition_id, 0)
#check idempotent
rdrsp = client_f.RequestJoinPartition(rdreq)
self.assertEqual(rdrsp.status.code, 0)
self.assertTrue(rdrsp.HasField('manifest'))
self.assertEqual(rdrsp.manifest.sync_example_id_rep.rank_id, 1)
self.assertEqual(rdrsp.manifest.sync_example_id_rep.state,
dj_pb.SyncExampleIdState.Syncing)
self.assertEqual(rdrsp.manifest.partition_id, 0)
rdreq1 = dj_pb.RawDataRequest(
data_source_meta=data_source_l.data_source_meta,
rank_id=1,
partition_id=0,
sync_example_id=empty_pb2.Empty())
try:
rsp = client_l.FinishJoinPartition(rdreq1)
except Exception as e:
self.assertTrue(True)
else:
self.assertTrue(False)
rdreq2 = dj_pb.RawDataRequest(
data_source_meta=data_source_l.data_source_meta,
rank_id=0,
partition_id=0,
join_example=empty_pb2.Empty())
try:
rsp = client_l.FinishJoinPartition(rdreq2)
except Exception as e:
self.assertTrue(True)
else:
self.assertTrue(False)
rdreq = dj_pb.RawDataRequest(
data_source_meta=data_source_l.data_source_meta,
rank_id=0,
partition_id=0,
)
manifest_l = client_l.QueryRawDataManifest(rdreq)
self.assertTrue(manifest_l is not None)
self.assertFalse(manifest_l.finished)
self.assertEqual(manifest_l.next_process_index, 0)
rdreq = dj_pb.RawDataRequest(
data_source_meta=data_source_l.data_source_meta,
rank_id=0,
partition_id=0,
)
rsp = client_l.GetRawDataLatestTimeStamp(rdreq)
self.assertEqual(rsp.status.code, 0)
self.assertTrue(rsp.HasField('timestamp'))
self.assertEqual(rsp.timestamp.seconds, 0)
self.assertEqual(rsp.timestamp.nanos, 0)
rdreq = dj_pb.RawDataRequest(
data_source_meta=data_source_l.data_source_meta,
rank_id=0,
partition_id=0,
added_raw_data_metas=dj_pb.AddedRawDataMetas(
dedup=False,
raw_data_metas=[
dj_pb.RawDataMeta(
file_path='a',
timestamp=timestamp_pb2.Timestamp(seconds=3))
]))
rsp = client_l.AddRawData(rdreq)
self.assertEqual(rsp.code, 0)
manifest_l = client_l.QueryRawDataManifest(rdreq)
self.assertTrue(manifest_l is not None)
self.assertFalse(manifest_l.finished)
self.assertEqual(manifest_l.next_process_index, 1)
rdreq = dj_pb.RawDataRequest(
data_source_meta=data_source_l.data_source_meta,
rank_id=0,
partition_id=0,
)
rsp = client_l.GetRawDataLatestTimeStamp(rdreq)
self.assertEqual(rsp.status.code, 0)
self.assertTrue(rsp.HasField('timestamp'))
self.assertEqual(rsp.timestamp.seconds, 3)
self.assertEqual(rsp.timestamp.nanos, 0)
rdreq = dj_pb.RawDataRequest(
data_source_meta=data_source_l.data_source_meta,
rank_id=0,
partition_id=0,
added_raw_data_metas=dj_pb.AddedRawDataMetas(
dedup=False,
raw_data_metas=[
dj_pb.RawDataMeta(
file_path='b',
timestamp=timestamp_pb2.Timestamp(seconds=5))
]))
rsp = client_l.AddRawData(rdreq)
self.assertEqual(rsp.code, 0)
manifest_l = client_l.QueryRawDataManifest(rdreq)
self.assertTrue(manifest_l is not None)
self.assertFalse(manifest_l.finished)
self.assertEqual(manifest_l.next_process_index, 2)
rdreq = dj_pb.RawDataRequest(
data_source_meta=data_source_l.data_source_meta,
rank_id=0,
partition_id=0,
)
rsp = client_l.GetRawDataLatestTimeStamp(rdreq)
self.assertEqual(rsp.status.code, 0)
self.assertTrue(rsp.HasField('timestamp'))
self.assertEqual(rsp.timestamp.seconds, 5)
self.assertEqual(rsp.timestamp.nanos, 0)
rdreq = dj_pb.RawDataRequest(
data_source_meta=data_source_l.data_source_meta,
rank_id=0,
partition_id=0,
added_raw_data_metas=dj_pb.AddedRawDataMetas(
dedup=True,
raw_data_metas=[
dj_pb.RawDataMeta(
file_path='b',
timestamp=timestamp_pb2.Timestamp(seconds=5)),
dj_pb.RawDataMeta(
file_path='a',
timestamp=timestamp_pb2.Timestamp(seconds=3))
]))
rsp = client_l.AddRawData(rdreq)
self.assertEqual(rsp.code, 0)
manifest_l = client_l.QueryRawDataManifest(rdreq)
self.assertTrue(manifest_l is not None)
self.assertFalse(manifest_l.finished)
self.assertEqual(manifest_l.next_process_index, 2)
rsp = client_l.GetRawDataLatestTimeStamp(rdreq)
self.assertEqual(rsp.status.code, 0)
self.assertTrue(rsp.HasField('timestamp'))
self.assertEqual(rsp.timestamp.seconds, 5)
self.assertEqual(rsp.timestamp.nanos, 0)
rdreq = dj_pb.RawDataRequest(
data_source_meta=data_source_f.data_source_meta,
rank_id=0,
partition_id=0,
)
manifest_f = client_f.QueryRawDataManifest(rdreq)
self.assertTrue(manifest_f is not None)
self.assertFalse(manifest_f.finished)
self.assertEqual(manifest_f.next_process_index, 0)
rdreq = dj_pb.RawDataRequest(
data_source_meta=data_source_l.data_source_meta,
rank_id=0,
partition_id=0,
)
rsp = client_f.GetRawDataLatestTimeStamp(rdreq)
self.assertEqual(rsp.status.code, 0)
self.assertTrue(rsp.HasField('timestamp'))
self.assertEqual(rsp.timestamp.seconds, 0)
self.assertEqual(rsp.timestamp.nanos, 0)
rdreq = dj_pb.RawDataRequest(
data_source_meta=data_source_f.data_source_meta,
rank_id=0,
partition_id=0,
added_raw_data_metas=dj_pb.AddedRawDataMetas(
dedup=False,
raw_data_metas=[
dj_pb.RawDataMeta(
file_path='a',
timestamp=timestamp_pb2.Timestamp(seconds=1))
]))
rsp = client_f.AddRawData(rdreq)
self.assertEqual(rsp.code, 0)
manifest_f = client_f.QueryRawDataManifest(rdreq)
self.assertTrue(manifest_f is not None)
self.assertFalse(manifest_f.finished)
self.assertEqual(manifest_f.next_process_index, 1)
rdreq = dj_pb.RawDataRequest(
data_source_meta=data_source_l.data_source_meta,
rank_id=0,
partition_id=0,
)
rsp = client_f.GetRawDataLatestTimeStamp(rdreq)
self.assertEqual(rsp.status.code, 0)
self.assertTrue(rsp.HasField('timestamp'))
self.assertEqual(rsp.timestamp.seconds, 1)
self.assertEqual(rsp.timestamp.nanos, 0)
rdreq = dj_pb.RawDataRequest(
data_source_meta=data_source_f.data_source_meta,
rank_id=0,
partition_id=0,
added_raw_data_metas=dj_pb.AddedRawDataMetas(
dedup=False,
raw_data_metas=[
dj_pb.RawDataMeta(
file_path='b',
timestamp=timestamp_pb2.Timestamp(seconds=2))
]))
rsp = client_f.AddRawData(rdreq)
self.assertEqual(rsp.code, 0)
manifest_f = client_f.QueryRawDataManifest(rdreq)
self.assertTrue(manifest_f is not None)
self.assertFalse(manifest_f.finished)
self.assertEqual(manifest_f.next_process_index, 2)
rdreq = dj_pb.RawDataRequest(
data_source_meta=data_source_l.data_source_meta,
rank_id=0,
partition_id=0,
)
rsp = client_f.GetRawDataLatestTimeStamp(rdreq)
self.assertEqual(rsp.status.code, 0)
self.assertTrue(rsp.HasField('timestamp'))
self.assertEqual(rsp.timestamp.seconds, 2)
self.assertEqual(rsp.timestamp.nanos, 0)
rdreq = dj_pb.RawDataRequest(
data_source_meta=data_source_f.data_source_meta,
rank_id=0,
partition_id=0,
added_raw_data_metas=dj_pb.AddedRawDataMetas(
dedup=True,
raw_data_metas=[
dj_pb.RawDataMeta(
file_path='a',
timestamp=timestamp_pb2.Timestamp(seconds=1)),
dj_pb.RawDataMeta(
file_path='b',
timestamp=timestamp_pb2.Timestamp(seconds=2))
]))
rsp = client_f.AddRawData(rdreq)
self.assertEqual(rsp.code, 0)
manifest_f = client_f.QueryRawDataManifest(rdreq)
self.assertTrue(manifest_f is not None)
self.assertFalse(manifest_f.finished)
self.assertEqual(manifest_f.next_process_index, 2)
rdreq = dj_pb.RawDataRequest(
data_source_meta=data_source_l.data_source_meta,
rank_id=0,
partition_id=0,
)
rsp = client_f.GetRawDataLatestTimeStamp(rdreq)
self.assertEqual(rsp.status.code, 0)
self.assertTrue(rsp.HasField('timestamp'))
self.assertEqual(rsp.timestamp.seconds, 2)
self.assertEqual(rsp.timestamp.nanos, 0)
rdreq = dj_pb.RawDataRequest(
data_source_meta=data_source_l.data_source_meta,
rank_id=0,
partition_id=0,
finish_raw_data=empty_pb2.Empty())
rsp = client_l.FinishRawData(rdreq)
self.assertEqual(rsp.code, 0)
#check idempotent
rsp = client_l.FinishRawData(rdreq)
self.assertEqual(rsp.code, 0)
manifest_l = client_l.QueryRawDataManifest(rdreq)
self.assertTrue(manifest_l is not None)
self.assertTrue(manifest_l.finished)
self.assertEqual(manifest_l.next_process_index, 2)
rdreq = dj_pb.RawDataRequest(
data_source_meta=data_source_l.data_source_meta,
rank_id=0,
partition_id=0,
added_raw_data_metas=dj_pb.AddedRawDataMetas(
dedup=False,
raw_data_metas=[
dj_pb.RawDataMeta(
file_path='x',
timestamp=timestamp_pb2.Timestamp(seconds=4))
]))
try:
rsp = client_l.AddRawData(rdreq)
except Exception as e:
self.assertTrue(True)
else:
self.assertTrue(False)
try:
rsp = client_f.FinishJoinPartition(rdreq2)
except Exception as e:
self.assertTrue(True)
else:
self.assertTrue(False)
rsp = client_l.FinishJoinPartition(rdreq1)
self.assertEqual(rsp.code, 0)
#check idempotent
rsp = client_l.FinishJoinPartition(rdreq1)
self.assertEqual(rsp.code, 0)
rsp = client_f.FinishJoinPartition(rdreq1)
self.assertEqual(rsp.code, 0)
#check idempotent
rsp = client_f.FinishJoinPartition(rdreq1)
self.assertEqual(rsp.code, 0)
try:
rsp = client_f.FinishJoinPartition(rdreq2)
except Exception as e:
self.assertTrue(True)
else:
self.assertTrue(False)
rdreq = dj_pb.RawDataRequest(
data_source_meta=data_source_f.data_source_meta,
rank_id=0,
partition_id=0,
finish_raw_data=empty_pb2.Empty())
rsp = client_f.FinishRawData(rdreq)
self.assertEqual(rsp.code, 0)
#check idempotent
rsp = client_f.FinishRawData(rdreq)
self.assertEqual(rsp.code, 0)
manifest_f = client_f.QueryRawDataManifest(rdreq)
self.assertTrue(manifest_f is not None)
self.assertTrue(manifest_f.finished)
self.assertEqual(manifest_f.next_process_index, 2)
rdreq = dj_pb.RawDataRequest(
data_source_meta=data_source_f.data_source_meta,
rank_id=0,
partition_id=0,
added_raw_data_metas=dj_pb.AddedRawDataMetas(
dedup=True,
raw_data_metas=[
dj_pb.RawDataMeta(
file_path='x',
timestamp=timestamp_pb2.Timestamp(seconds=3))
]))
try:
rsp = client_f.AddRawData(rdreq)
except Exception as e:
self.assertTrue(True)
else:
self.assertTrue(False)
rsp = client_f.FinishJoinPartition(rdreq2)
self.assertEqual(rsp.code, 0)
#check idempotent
rsp = client_f.FinishJoinPartition(rdreq2)
self.assertEqual(rsp.code, 0)
rsp = client_l.FinishJoinPartition(rdreq2)
self.assertEqual(rsp.code, 0)
#check idempotent
rsp = client_l.FinishJoinPartition(rdreq2)
self.assertEqual(rsp.code, 0)
while True:
req_l = dj_pb.DataSourceRequest(
data_source_meta=data_source_l.data_source_meta)
req_f = dj_pb.DataSourceRequest(
data_source_meta=data_source_f.data_source_meta)
dss_l = client_l.GetDataSourceStatus(req_l)
dss_f = client_f.GetDataSourceStatus(req_f)
self.assertEqual(dss_l.role, common_pb.FLRole.Leader)
self.assertEqual(dss_f.role, common_pb.FLRole.Follower)
if dss_l.state == common_pb.DataSourceState.Finished and \
dss_f.state == common_pb.DataSourceState.Finished:
break
else:
time.sleep(2)
master_l.stop()
master_f.stop()
parser.add_argument('--raw_data_sub_dir',
type=str,
required=True,
help='the mysql base dir to subscribe new raw data')
args = parser.parse_args()
data_source = common_pb.DataSource()
data_source.data_source_meta.name = args.data_source_name
data_source.data_source_meta.partition_num = args.partition_num
data_source.data_source_meta.start_time = args.start_time
data_source.data_source_meta.end_time = args.end_time
data_source.data_source_meta.negative_sampling_rate = \
args.negative_sampling_rate
if args.role.upper() == 'LEADER':
data_source.role = common_pb.FLRole.Leader
else:
assert args.role.upper() == 'FOLLOWER'
data_source.role = common_pb.FLRole.Follower
data_source.output_base_dir = args.output_base_dir
data_source.raw_data_sub_dir = args.raw_data_sub_dir
data_source.state = common_pb.DataSourceState.Init
kvstore = DBClient(args.kvstore_type)
master_kvstore_key = common.data_source_kvstore_base_dir(
data_source.data_source_meta.name)
raw_data = kvstore.get_data(master_kvstore_key)
if raw_data is None:
logging.info("data source %s is not existed", args.data_source_name)
common.commit_data_source(kvstore, data_source)
logging.info("apply new data source %s", args.data_source_name)
else:
logging.info("data source %s has been existed", args.data_source_name)
required=True,
help='the etcd base dir to subscribe new raw data')
args = parser.parse_args()
data_source = common_pb.DataSource()
data_source.data_source_meta.name = args.data_source_name
data_source.data_source_meta.partition_num = args.partition_num
data_source.data_source_meta.start_time = args.start_time
data_source.data_source_meta.end_time = args.end_time
data_source.data_source_meta.negative_sampling_rate = \
args.negative_sampling_rate
if args.role == 'leader':
data_source.role = common_pb.FLRole.Leader
else:
assert args.role == 'follower'
data_source.role = common_pb.FLRole.Follower
data_source.example_dumped_dir = args.example_dump_dir
data_source.data_block_dir = args.data_block_dir
data_source.raw_data_sub_dir = args.raw_data_sub_dir
data_source.state = common_pb.DataSourceState.Init
etcd = EtcdClient(args.etcd_name, args.etcd_addrs, args.etcd_base_dir)
master_etcd_key = common.data_source_etcd_base_dir(
data_source.data_source_meta.name)
raw_data = etcd.get_data(master_etcd_key)
if raw_data is None:
logging.info("data source %s is not existed", args.data_source_name)
common.commit_data_source(etcd, data_source)
logging.info("apply new data source %s", args.data_source_name)
else:
logging.info("data source %s has been existed", args.data_source_name)
etcd.destroy_client_pool()
