class DBBasedMockRawDataVisitor(RawDataVisitor): def __init__(self, kvstore, raw_data_options, mock_data_source_name, raw_data_sub_dir, partition_id): mock_data_source = common_pb.DataSource( state=common_pb.DataSourceState.Processing, raw_data_sub_dir=raw_data_sub_dir, data_source_meta=common_pb.DataSourceMeta( name=mock_data_source_name, partition_num=partition_id + 1)) self._mock_rd_manifest_manager = RawDataManifestManager( kvstore, mock_data_source, False) self._partition_id = partition_id super(DBBasedMockRawDataVisitor, self).__init__(kvstore, mock_data_source, partition_id, raw_data_options) def active_visitor(self): self._mock_rd_manifest_manager.sub_new_raw_data(self._partition_id) if self.is_visitor_stale(): self._finished = False def is_input_data_finish(self): manager = self._mock_rd_manifest_manager return manager.get_manifest(self._partition_id).finished def cleanup_meta_data(self): self._mock_rd_manifest_manager.cleanup_meta_data()
def __init__(self, etcd, raw_data_options, mock_data_source_name, raw_data_sub_dir): mock_data_source = common_pb.DataSource( state=common_pb.DataSourceState.Processing, data_source_meta=common_pb.DataSourceMeta( name=mock_data_source_name, partition_num=1, raw_data_sub_dir=raw_data_sub_dir)) self._mock_rd_manifest_manager = RawDataManifestManager( etcd, mock_data_source) super(EtcdBasedMockRawDataVisitor, self).__init__(etcd, mock_data_source, 0, raw_data_options)
def __init__(self, kvstore, raw_data_options, mock_data_source_name, raw_data_sub_dir, partition_id): mock_data_source = common_pb.DataSource( state=common_pb.DataSourceState.Processing, raw_data_sub_dir=raw_data_sub_dir, data_source_meta=common_pb.DataSourceMeta( name=mock_data_source_name, partition_num=partition_id + 1)) self._mock_rd_manifest_manager = RawDataManifestManager( kvstore, mock_data_source, False) self._partition_id = partition_id super(DBBasedMockRawDataVisitor, self).__init__(kvstore, mock_data_source, partition_id, raw_data_options)
class FileBasedMockRawDataVisitor(RawDataVisitor): def __init__(self, etcd, raw_data_options, mock_data_source_name, input_fpaths): mock_data_source = common_pb.DataSource( state=common_pb.DataSourceState.Processing, data_source_meta=common_pb.DataSourceMeta( name=mock_data_source_name, partition_num=1 ) ) self._mock_rd_manifest_manager = RawDataManifestManager( etcd, mock_data_source ) manifest = self._mock_rd_manifest_manager.get_manifest(0) if not manifest.finished: metas = [] for fpath in input_fpaths: metas.append(dj_pb.RawDataMeta(file_path=fpath, start_index=-1)) self._mock_rd_manifest_manager.add_raw_data(0, metas, True) self._mock_rd_manifest_manager.finish_raw_data(0) super(FileBasedMockRawDataVisitor, self).__init__( etcd, mock_data_source, 0, raw_data_options ) def cleanup_meta_data(self): self._mock_rd_manifest_manager.cleanup_meta_data()
def __init__(self, peer_client, data_source_name, kvstore, batch_mode): self._lock = threading.Lock() self._peer_client = peer_client self._data_source_name = data_source_name self._kvstore = kvstore self._batch_mode = batch_mode self._init_fsm_action() self._data_source = None self._sync_data_source() self._reset_batch_mode() self._raw_data_manifest_manager = RawDataManifestManager( kvstore, self._data_source, batch_mode) self._data_source_meta = self._data_source.data_source_meta if self._data_source.role == common_pb.FLRole.Leader: self._role_repr = "leader" else: self._role_repr = "follower" self._fsm_worker = None self._started = False
def __init__(self, peer_client, data_source_name, etcd): self._lock = threading.Lock() self._peer_client = peer_client self._data_source_name = data_source_name self._etcd = etcd self._init_fsm_action() self._data_source = None self._sync_data_source() assert self._data_source is not None, \ "data source must not None if sync data source success" self._raw_data_manifest_manager = RawDataManifestManager( etcd, self._data_source ) self._data_source_meta = self._data_source.data_source_meta if self._data_source.role == common_pb.FLRole.Leader: self._role_repr = "leader" else: self._role_repr = "follower" self._fsm_worker = None self._started = False
class EtcdBasedMockRawDataVisitor(RawDataVisitor): def __init__(self, etcd, raw_data_options, mock_data_source_name, raw_data_sub_dir): mock_data_source = common_pb.DataSource( state=common_pb.DataSourceState.Processing, data_source_meta=common_pb.DataSourceMeta( name=mock_data_source_name, partition_num=1, raw_data_sub_dir=raw_data_sub_dir)) self._mock_rd_manifest_manager = RawDataManifestManager( etcd, mock_data_source) super(EtcdBasedMockRawDataVisitor, self).__init__(etcd, mock_data_source, 0, raw_data_options) def active_visitor(self): self._mock_rd_manifest_manager.sub_new_raw_data() if self.is_visitor_stale(): self._finished = False def is_input_data_finish(self): return self._mock_rd_manifest_manager.get_manifest(0).finished
class MasterFSM(object): INVALID_PEER_FSM_STATE = {} INVALID_PEER_FSM_STATE[common_pb.DataSourceState.Init] = set( [common_pb.DataSourceState.Failed, common_pb.DataSourceState.Ready, common_pb.DataSourceState.Finished] ) INVALID_PEER_FSM_STATE[common_pb.DataSourceState.Processing] = set( [common_pb.DataSourceState.Failed, common_pb.DataSourceState.Finished] ) INVALID_PEER_FSM_STATE[common_pb.DataSourceState.Ready] = set( [common_pb.DataSourceState.Failed, common_pb.DataSourceState.Init] ) INVALID_PEER_FSM_STATE[common_pb.DataSourceState.Finished] = set( [common_pb.DataSourceState.Failed, common_pb.DataSourceState.Init, common_pb.DataSourceState.Processing] ) def __init__(self, peer_client, data_source_name, etcd): self._lock = threading.Lock() self._peer_client = peer_client self._data_source_name = data_source_name self._etcd = etcd self._init_fsm_action() self._data_source = None self._sync_data_source() assert self._data_source is not None, \ "data source must not None if sync data source success" self._raw_data_manifest_manager = RawDataManifestManager( etcd, self._data_source ) self._data_source_meta = self._data_source.data_source_meta if self._data_source.role == common_pb.FLRole.Leader: self._role_repr = "leader" else: self._role_repr = "follower" self._fsm_worker = None self._started = False def get_mainifest_manager(self): return self._raw_data_manifest_manager def get_data_source(self): with self._lock: self._sync_data_source() return self._data_source def set_failed(self): return self.set_state(common_pb.DataSourceState.Failed, None) def set_state(self, new_state, origin_state=None): with self._lock: try: self._sync_data_source() if self._data_source.state == new_state: return True if (origin_state is None or self._data_source.state == origin_state): self._data_source.state = new_state self._update_data_source(self._data_source) return True logging.warning("DataSource: %s failed to set to state: " "%d, origin state mismatch(%d != %d)", self._data_source_name, new_state, origin_state, self._data_source.state) return False except Exception as e: # pylint: disable=broad-except logging.warning("Faile to set state to %d with exception %s", new_state, e) return False return True def start_fsm_worker(self): with self._lock: if not self._started: assert self._fsm_worker is None, \ "fsm_woker must be None if FSM is not started" self._started = True self._fsm_worker = RoutineWorker( '{}_fsm_worker'.format(self._data_source_name), self._fsm_routine_fn, self._fsm_routine_cond, 5 ) self._fsm_worker.start_routine() def stop_fsm_worker(self): tmp_worker = None with self._lock: if self._fsm_worker is not None: tmp_worker = self._fsm_worker self._fsm_worker = None if tmp_worker is not None: tmp_worker.stop_routine() def _fsm_routine_fn(self): peer_info = self._get_peer_data_source_status() with self._lock: self._sync_data_source() state = self._data_source.state if self._fallback_failed_state(peer_info): logging.warning("%s at state %d, Peer at state %d "\ "state invalid! abort data source %s", self._role_repr, state, peer_info.state, self._data_source_name) elif state not in self._fsm_driven_handle: logging.error("%s at error state %d for data_source %s", self._role_repr, state, self._data_source_name) else: state_changed = self._fsm_driven_handle[state](peer_info) if state_changed: self._sync_data_source() new_state = self._data_source.state logging.warning("%s state changed from %d to %d", self._role_repr, state, new_state) state = new_state if state in (common_pb.DataSourceState.Init, common_pb.DataSourceState.Processing): self._raw_data_manifest_manager.sub_new_raw_data() def _fsm_routine_cond(self): return True def _sync_data_source(self): if self._data_source is None: self._data_source = \ retrieve_data_source(self._etcd, self._data_source_name) def _init_fsm_action(self): self._fsm_driven_handle = { common_pb.DataSourceState.UnKnown: self._get_fsm_action('unknown'), common_pb.DataSourceState.Init: self._get_fsm_action('init'), common_pb.DataSourceState.Processing: self._get_fsm_action('processing'), common_pb.DataSourceState.Ready: self._get_fsm_action('ready'), common_pb.DataSourceState.Finished: self._get_fsm_action('finished'), common_pb.DataSourceState.Failed: self._get_fsm_action('failed') } def _get_fsm_action(self, action): def _not_implement(useless): raise NotImplementedError('state is not NotImplemented') name = '_fsm_{}_action'.format(action) return getattr(self, name, _not_implement) def _fsm_init_action(self, peer_info): state_changed = False if self._data_source.role == common_pb.FLRole.Leader: if peer_info.state == common_pb.DataSourceState.Init: state_changed = True elif peer_info.state == common_pb.DataSourceState.Processing: state_changed = True if state_changed: self._data_source.state = common_pb.DataSourceState.Processing self._update_data_source(self._data_source) return True return False def _fsm_processing_action(self, peer_info): if self._all_partition_finished(): state_changed = False if self._data_source.role == common_pb.FLRole.Leader: if peer_info.state == common_pb.DataSourceState.Processing: state_changed = True elif peer_info.state == common_pb.DataSourceState.Ready: state_changed = True if state_changed: self._data_source.state = common_pb.DataSourceState.Ready self._update_data_source(self._data_source) return True return False def _fsm_ready_action(self, peer_info): state_changed = False if self._data_source.role == common_pb.FLRole.Leader: if peer_info.state == common_pb.DataSourceState.Ready: state_changed = True elif peer_info.state == common_pb.DataSourceState.Finished: state_changed = True if state_changed: self._data_source.state = common_pb.DataSourceState.Finished self._update_data_source(self._data_source) return True return False def _fsm_finished_action(self, peer_info): return False def _fsm_failed_action(self, peer_info): if peer_info.state != common_pb.DataSourceState.Failed: request = dj_pb.DataSourceRequest( data_source_meta=self._data_source_meta ) self._peer_client.AbortDataSource(request) return False def _fallback_failed_state(self, peer_info): state = self._data_source.state if (state in self.INVALID_PEER_FSM_STATE and peer_info.state in self.INVALID_PEER_FSM_STATE[state]): self._data_source.state = common_pb.DataSourceState.Failed self._update_data_source(self._data_source) return True return False 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 _get_peer_data_source_status(self): request = dj_pb.DataSourceRequest( data_source_meta=self._data_source_meta ) return self._peer_client.GetDataSourceStatus(request) def _all_partition_finished(self): all_manifest = self._raw_data_manifest_manager.list_all_manifest() assert len(all_manifest) == \ self._data_source.data_source_meta.partition_num, \ "manifest number should same with partition number" for manifest in all_manifest.values(): if manifest.sync_example_id_rep.state != \ dj_pb.SyncExampleIdState.Synced or \ manifest.join_example_rep.state != \ dj_pb.JoinExampleState.Joined: return False return True
class MasterFSM(object): INVALID_PEER_FSM_STATE = {} INVALID_PEER_FSM_STATE[common_pb.DataSourceState.Init] = set([ common_pb.DataSourceState.Failed, common_pb.DataSourceState.Ready, common_pb.DataSourceState.Finished ]) INVALID_PEER_FSM_STATE[common_pb.DataSourceState.Processing] = set( [common_pb.DataSourceState.Failed, common_pb.DataSourceState.Finished]) INVALID_PEER_FSM_STATE[common_pb.DataSourceState.Ready] = set( [common_pb.DataSourceState.Failed, common_pb.DataSourceState.Init]) INVALID_PEER_FSM_STATE[common_pb.DataSourceState.Finished] = set([ common_pb.DataSourceState.Failed, common_pb.DataSourceState.Init, common_pb.DataSourceState.Processing ]) def __init__(self, peer_client, data_source_name, etcd): self._lock = threading.Lock() self._peer_client = peer_client self._data_source_name = data_source_name self._master_etcd_key = os.path.join(data_source_name, 'master') self._etcd = etcd self._init_fsm_action() self._data_source = None self._sync_data_source() assert self._data_source is not None self._raw_data_manifest_manager = RawDataManifestManager( etcd, self._data_source) assert self._data_source is not None self._data_source_meta = self._data_source.data_source_meta if self._data_source.role == common_pb.FLRole.Leader: self._role_repr = "leader" else: self._role_repr = "follower" self._fsm_worker = None self._started = False def get_mainifest_manager(self): return self._raw_data_manifest_manager def get_data_source(self): with self._lock: self._sync_data_source() assert self._data_source is not None return self._data_source def set_failed(self): return self.set_state(common_pb.DataSourceState.Failed, None) def set_state(self, new_state, origin_state=None): with self._lock: try: self._sync_data_source() assert self._data_source is not None if self._data_source.state == new_state: return True if (origin_state is None or self._data_source.state == origin_state): self._data_source.state = new_state self._update_data_source(self._data_source) return True logging.warning( "DataSource: %s failed to set to state: " "%d, origin state mismatch(%d != %d)", self._data_source_name, new_state, origin_state, self._data_source.state) return False except Exception as e: # pylint: disable=broad-except logging.warning("Faile to set state to %d with exception %s", new_state, e) return False return True def start_fsm_worker(self): with self._lock: if not self._started: assert self._fsm_worker is None self._started = True self._fsm_worker = RoutineWorker( '{}_fsm_worker'.format(self._data_source_name), self._fsm_routine_fn, self._fsm_routine_cond, 5) self._fsm_worker.start_routine() def stop_fsm_worker(self): tmp_worker = None with self._lock: if self._fsm_worker is not None: tmp_worker = self._fsm_worker self._fsm_worker = None if tmp_worker is not None: tmp_worker.stop_routine() def _fsm_routine_fn(self): peer_info = self._get_peer_data_source_state() if peer_info.status.code != 0: logging.error("Failed to get peer state, %s", peer_info.status.error_message) return with self._lock: self._sync_data_source() assert self._data_source is not None state = self._data_source.state if self._fallback_failed_state(peer_info): logging.warning( "Self(%s) at state %d, Peer at state %d "\ "state invalid! abort data source %s", self._role_repr, state, peer_info.state, self._data_source_name ) elif state not in self._fsm_driven_handle: logging.error("Self(%s) at error state %d for data_source %s", self._role_repr, state, self._data_source_name) else: state_changed = self._fsm_driven_handle[state](peer_info) if state_changed: self._sync_data_source() assert self._data_source is not None new_state = self._data_source.state logging.warning("Self(%s) state changed from %d to %d", self._role_repr, state, new_state) def _fsm_routine_cond(self): return True def _sync_data_source(self): if self._data_source is None: raw_data = self._etcd.get_data(self._master_etcd_key) if raw_data is None: raise ValueError("etcd master key is None for {}".format( self._data_source_name)) self._data_source = text_format.Parse(raw_data, common_pb.DataSource()) def _init_fsm_action(self): self._fsm_driven_handle = { common_pb.DataSourceState.UnKnown: self._get_fsm_action('unknown'), common_pb.DataSourceState.Init: self._get_fsm_action('init'), common_pb.DataSourceState.Processing: self._get_fsm_action('processing'), common_pb.DataSourceState.Ready: self._get_fsm_action('ready'), common_pb.DataSourceState.Finished: self._get_fsm_action('finished'), common_pb.DataSourceState.Failed: self._get_fsm_action('failed') } def _get_fsm_action(self, action): def _not_implement(useless): raise NotImplementedError('state is not NotImplemented') name = '_fsm_{}_action'.format(action) return getattr(self, name, _not_implement) def _fsm_init_action(self, peer_info): state_changed = False if self._data_source.role == common_pb.FLRole.Leader: if peer_info.state == common_pb.DataSourceState.Init: state_changed = True elif peer_info.state == common_pb.DataSourceState.Processing: state_changed = True if state_changed: self._data_source.state = common_pb.DataSourceState.Processing self._update_data_source(self._data_source) return True return False def _fsm_processing_action(self, peer_info): if self._all_partition_finished(): state_changed = False if self._data_source.role == common_pb.FLRole.Leader: if peer_info.state == common_pb.DataSourceState.Processing: state_changed = True elif peer_info.state == common_pb.DataSourceState.Ready: state_changed = True if state_changed: self._data_source.state = common_pb.DataSourceState.Ready self._update_data_source(self._data_source) return True return False def _fsm_ready_action(self, peer_info): state_changed = False if self._data_source.role == common_pb.FLRole.Leader: if peer_info.state == common_pb.DataSourceState.Ready: state_changed = True elif peer_info.state == common_pb.DataSourceState.Finished: state_changed = True if state_changed: self._data_source.state = common_pb.DataSourceState.Finished self._update_data_source(self._data_source) return True return False def _fsm_finished_action(self, peer_info): return False def _fsm_failed_action(self, peer_info): if peer_info.state != common_pb.DataSourceState.Failed: self._peer_client.AbortDataSource(self._data_source_meta) return False def _fallback_failed_state(self, peer_info): state = self._data_source.state if (state in self.INVALID_PEER_FSM_STATE and peer_info.state in self.INVALID_PEER_FSM_STATE[state]): self._data_source.state = common_pb.DataSourceState.Failed self._update_data_source(self._data_source) return True return False def _update_data_source(self, data_source): self._data_source = None try: self._etcd.set_data(self._master_etcd_key, text_format.MessageToString(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 _get_peer_data_source_state(self): return self._peer_client.GetDataSourceState(self._data_source_meta) def _all_partition_finished(self): all_manifest = self._raw_data_manifest_manager.list_all_manifest() assert (len(all_manifest) == self._data_source.data_source_meta.partition_num) for manifest in all_manifest.values(): if manifest.state != dj_pb.RawDataState.Done: return False return True