def test_data_block_set(self):
db1 = DataBlock('1', 'data_path1', 'meta_path1')
db2 = DataBlock('2', 'data_path2', 'meta_path2')
db_set = DataBlockSet()
db_set.add(db1)
db_set.add(db2)
self.assertIsNone(db_set.get('3'))
self.assertEqual(db_set.get('1'), db1)
self.assertIsNone(db_set.get('1'))
self.assertEqual(db_set.get('2'), db2)
self.assertIsNone(db_set.get('2'))
class FollowerTrainerMaster(TrainerMaster):
def __init__(self, application_id, data_source, start_time, end_time,
online_training):
super(FollowerTrainerMaster, self).__init__(application_id, None,
online_training)
self._data_block_set = DataBlockSet()
self._data_block_visitor = DataBlockVisitor(data_source, ETCD_NAME,
ETCD_BASE_DIR, ETCD_ADDR)
self._start_time = start_time
self._end_time = end_time
def _load_data(self):
checkpoint = self._get_checkpoint()
# pylint: disable=line-too-long
for block_id, block_item in self._data_block_visitor.LoadDataBlockRepByTimeFrame(
self._start_time, self._end_time).items():
if block_id not in checkpoint:
logging.debug('load data block id %s path %s', block_id,
block_item.data_block_fpath)
self._data_block_set.add(block_item)
logging.debug("FollowerTrainerMaster: get all block %s",
self._data_block_set)
def _alloc_data_block(self, block_id=None):
logging.debug("FollowerTrainerMaster is getting block %s", block_id)
if not block_id:
raise Exception('follower tm need block_id to alloc.')
return self._data_block_set.get(block_id)
class FollowerTrainerMaster(TrainerMaster):
def __init__(self, application_id, data_source,
start_time, end_time, online_training):
super(FollowerTrainerMaster, self).__init__(application_id,
None, online_training)
self._data_block_set = DataBlockSet()
kvstore_use_mock = os.environ.get('KVSTORE_USE_MOCK', "off") == "on"
self._data_block_visitor = DataBlockVisitor(
data_source, db_database, db_base_dir, db_addr,
db_username, db_password, kvstore_use_mock)
self._start_time = start_time
self._end_time = end_time
def _load_data(self):
checkpoint = self._get_checkpoint()
# pylint: disable=line-too-long
for block_id, block_item in self._data_block_visitor.LoadDataBlockRepByTimeFrame(
self._start_time, self._end_time).items():
if block_id not in checkpoint:
logging.debug('load data block id %s path %s',
block_id, block_item.data_block_fpath)
self._data_block_set.add(block_item)
logging.debug("FollowerTrainerMaster: get all block %s",
self._data_block_set)
def _alloc_data_block(self, block_id=None):
logging.debug("FollowerTrainerMaster is getting block %s", block_id)
if not block_id:
raise Exception('follower tm need block_id to alloc.')
return self._data_block_set.get(block_id)