class FrontendDedicatedLoggingOp(StreamingOperator):
def __init__(self):
super().__init__(operator_name='FRONTEND_DEDICATED_LOGGING_OP')
self._pslx_dedicated_logging_storage = ProtoTableStorage(logger=pslx_frontend_logger)
self._pslx_dedicated_logging_storage.initialize_from_file(
file_name=pslx_dedicated_logging_storage_path
)
self._cached_logging = {}
self._logging_storage_capacity = max(int(EnvUtil.get_pslx_env_variable(var='PSLX_INTERNAL_CACHE')), 700)
def pubsub_msg_parser(self, exchange_name, topic_name, message):
self._cached_logging[str(TimezoneUtil.cur_time_in_pst())] = message
if len(self._cached_logging) >= int(EnvUtil.get_pslx_env_variable('PSLX_RPC_FLUSH_RATE')):
self._pslx_dedicated_logging_storage.write(
data=self._cached_logging
)
self._cached_logging.clear()
total_entries = self._pslx_dedicated_logging_storage.get_num_entries()
if total_entries > self._logging_storage_capacity:
try:
all_data = self._pslx_dedicated_logging_storage.read_all()
all_sorted_keys = sorted(list(all_data.keys()))
for i in range(total_entries - self._logging_storage_capacity):
key_to_delete = all_sorted_keys[i]
self._pslx_dedicated_logging_storage.delete(key=key_to_delete)
except Exception as _:
pass
def execute_impl(self):
pslx_dedicated_subscriber.bind_to_op(self)
pslx_dedicated_subscriber.start()
def test_write_1(self):
proto_table_storage = ProtoTableStorage()
proto_table_storage.initialize_from_file(file_name=self.TEST_DATA_3)
proto_table_storage.write(data={'test': self.EXAMPLE_PROTO_1}, )
result_proto = proto_table_storage.read(params={
'key': 'test',
'message_type': NodeSnapshot
})
self.assertEqual(result_proto, self.EXAMPLE_PROTO_1)
gclient_ext.cp_file(self.TEST_DATA_1, self.TEST_DATA_3)
def get_response_impl(backend_folder, request, lru_cache=None):
storage_value = ContainerBackendValue()
storage_value.container_name = request.container_name
storage_value.container_status = request.status
for operator_name, operator_snapshot in dict(
request.operator_snapshot_map).items():
operator_info = ContainerBackendValue.OperatorInfo()
operator_info.status = operator_snapshot.status
for parent in operator_snapshot.node_snapshot.parents_names:
operator_info.parents.append(parent)
operator_info.start_time = operator_snapshot.start_time
operator_info.end_time = operator_snapshot.end_time
operator_info.log_file = operator_snapshot.log_file
storage_value.operator_info_map[operator_name].CopyFrom(
operator_info)
storage_value.mode = request.mode
storage_value.data_model = request.data_model
storage_value.updated_time = str(TimezoneUtil.cur_time_in_pst())
storage_value.start_time = request.start_time
storage_value.end_time = request.end_time
storage_value.log_file = request.log_file
storage_value.run_cell = request.run_cell
storage_value.snapshot_cell = request.snapshot_cell
for key in request.counters:
storage_value.counters[key] = request.counters[key]
storage_value.ttl = int(
EnvUtil.get_pslx_env_variable('PSLX_BACKEND_CONTAINER_TTL'))
storage = lru_cache.get(
key=storage_value.container_name) if lru_cache else None
if not storage:
storage = ProtoTableStorage()
storage.initialize_from_file(file_name=FileUtil.join_paths_to_file(
root_dir=backend_folder,
base_name=storage_value.container_name + '.pb'))
if lru_cache:
lru_cache.set(key=backend_folder, value=storage)
all_data = storage.read_all()
if len(all_data) >= int(
EnvUtil.get_pslx_env_variable('PSLX_INTERNAL_CACHE')) > 0:
key_to_delete = sorted(all_data.keys())[0]
storage.delete(key=key_to_delete)
storage.write(data={storage_value.start_time: storage_value})
def test_write_2(self):
proto_table_storage = ProtoTableStorage()
proto_table_storage.initialize_from_file(file_name=self.TEST_DATA_4)
proto_table_storage.write(data={'test': self.EXAMPLE_PROTO_1})
proto_table_storage.write(data={'test_1': self.EXAMPLE_PROTO_2})
result_proto = proto_table_storage.read(
params={
'key': 'test_1',
'message_type': OperatorSnapshot
})
self.assertEqual(result_proto, self.EXAMPLE_PROTO_2)
proto_table_storage.write(data={'test_1': self.EXAMPLE_PROTO_3})
result_proto = proto_table_storage.read(
params={
'key': 'test_1',
'message_type': OperatorSnapshot
})
self.assertEqual(result_proto, self.EXAMPLE_PROTO_3)
result = proto_table_storage.read_all()
self.assertDictEqual(
result, {
'test': ProtoUtil.message_to_any(self.EXAMPLE_PROTO_1),
'test_1': ProtoUtil.message_to_any(self.EXAMPLE_PROTO_3),
})
self.assertEqual(proto_table_storage.get_num_entries(), 2)
gclient_ext.cp_file(self.TEST_DATA_2, self.TEST_DATA_4)
def write(self, data, params=None):
if not params:
params = {}
if 'overwrite' not in params:
params['overwrite'] = True
assert isinstance(data, dict)
exising_shard_to_data_map = defaultdict(dict)
new_data = {}
for key, val in data.items():
if key in self._index_map.index_map:
exising_shard_to_data_map[self._index_map.index_map[key]][key] = val
else:
new_data[key] = val
try:
for shard, existing_data in exising_shard_to_data_map.items():
related_proto_file = self._shard_to_file(shard=shard)
proto_table = ProtoTableStorage()
proto_table.initialize_from_file(file_name=related_proto_file)
proto_table.write(data=existing_data, params=params)
self.increment_rpc_count_by(n=proto_table.get_rpc_call_count_and_reset())
if new_data:
all_new_keys = list(new_data.keys())
latest_shard = self.get_latest_shard()
latest_proto_file = self._shard_to_file(shard=latest_shard)
proto_table = ProtoTableStorage()
proto_table.initialize_from_file(file_name=latest_proto_file)
latest_proto_table_size = proto_table.get_num_entries()
proto_table.write(
data={key: new_data[key] for key in all_new_keys[:self._size_per_shard - latest_proto_table_size]},
params=params
)
self.increment_rpc_count_by(n=proto_table.get_rpc_call_count_and_reset())
for key in all_new_keys[:self._size_per_shard - latest_proto_table_size]:
self._index_map.index_map[key] = latest_shard
if len(all_new_keys) > self._size_per_shard - latest_proto_table_size:
remaining_new_keys = all_new_keys[self._size_per_shard - latest_proto_table_size:]
start_index = 0
while start_index < len(remaining_new_keys):
data_for_new_shard = remaining_new_keys[start_index:start_index + self._size_per_shard]
latest_shard += 1
self._index_map.cur_shard += 1
proto_file = self._shard_to_file(shard=latest_shard)
self._logger.info("Write to new file with name [" + proto_file + '] and shard [' +
str(latest_shard) + '].')
proto_table = ProtoTableStorage()
proto_table.initialize_from_file(file_name=proto_file)
proto_table.write(
data={key: new_data[key] for key in data_for_new_shard},
params=params
)
self.increment_rpc_count_by(n=proto_table.get_rpc_call_count_and_reset())
for key in data_for_new_shard:
self._index_map.index_map[key] = latest_shard
start_index += self._size_per_shard
self._logger.info("Writing the index map to [" + self._index_map_file + '].')
self.increment_rpc_count_by(n=1)
FileUtil.write_proto_to_file(
proto=self._index_map,
file_name=self._index_map_file
)
except Exception as err:
self._SYS_LOGGER.error("Write to dir [" + self.get_dir_name() + "] got exception: " + str(err) + '.')
self._logger.error("Write to dir [" + self.get_dir_name() + "] got exception: " + str(err) + '.')
raise StorageWriteException("Write to dir [" + self.get_dir_name() + "] got exception: " + str(err) + '.')