2a76f2e3c6
Enhance logging in Gates, MongoDB, and Redis activities by replacing debug statements with info level logs, improving observability of data processing steps. Update worker configuration to adjust concurrency settings and enable autoscaling for task polling.
289 lines
9.5 KiB
Python
289 lines
9.5 KiB
Python
from temporalio import workflow, activity
|
|
|
|
from scouter import metrics
|
|
|
|
with workflow.unsafe.imports_passed_through():
|
|
from sientia_do.notifications.models import NotificationLevel
|
|
from sientia_do.temporal.activities.base import BaseActivity
|
|
from sientia_do.notifications.handlers import NotificationHandler
|
|
from sientia_do.observability.logger import Logger
|
|
from scouter.utils.quality.filters import null_values_filter, out_of_bounds_filter
|
|
from typing import Any
|
|
import traceback
|
|
from pandas import DataFrame
|
|
|
|
quality_gate_filters = {
|
|
'NULL_VALUES_FILTER': null_values_filter,
|
|
'OUT_OF_BOUNDS_FILTER': out_of_bounds_filter
|
|
}
|
|
|
|
|
|
class Gates(BaseActivity):
|
|
|
|
def __init__(self, logger: Logger, notification_handler: NotificationHandler):
|
|
BaseActivity.__init__(
|
|
self, logger, notification_handler, set_error_counter=True)
|
|
|
|
def apply_aggregation(self, group: DataFrame, aggr_function: str,
|
|
metadata: dict[str, Any]) -> float | None | str:
|
|
"""
|
|
Apply aggregation function to a group of data.
|
|
|
|
Args:
|
|
group (DataFrame): The group of data to apply the aggregation function to.
|
|
aggr_function (str): The aggregation function to apply.
|
|
|
|
Returns:
|
|
float | None | str: The result of the aggregation function.
|
|
"""
|
|
if len(group) == 1:
|
|
return group['value'].item()
|
|
|
|
# Apply aggregation function to value
|
|
if aggr_function == 'lts':
|
|
return group['value'].iloc[-1]
|
|
else:
|
|
group.dropna(inplace=True, subset=['value'])
|
|
|
|
if group.empty:
|
|
return None
|
|
|
|
if aggr_function == 'avg':
|
|
return group['value'].mean()
|
|
elif aggr_function == 'mdn':
|
|
return group['value'].median()
|
|
elif aggr_function == 'max':
|
|
return group['value'].max()
|
|
elif aggr_function == 'min':
|
|
return group['value'].min()
|
|
else:
|
|
self.send_notification(
|
|
metadata=metadata,
|
|
notification_id="AGGREGATION_ISSUES",
|
|
message=f"Invalid aggregation function: {aggr_function}",
|
|
block="aggregate_data",
|
|
level=NotificationLevel.ERROR,
|
|
attachment_content=traceback.format_exc()
|
|
)
|
|
return 'continue'
|
|
|
|
@activity.defn(name="aggregate_data")
|
|
async def aggregate_data(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
|
"""
|
|
Aggregates time series data by tag and name, applying specified
|
|
aggregation functions and taking the latest timestamp.
|
|
|
|
Args:
|
|
input_data (dict[str, Any]): The data to aggregate. Contains:
|
|
data (dict[str, Any]): The time series data.
|
|
model_tags (dict[str, Any]): The tags configuration
|
|
containing aggregation functions.
|
|
|
|
Returns:
|
|
dict[str, Any]: The aggregated data.
|
|
"""
|
|
|
|
metadata = input_data['metadata']
|
|
|
|
try:
|
|
# Convert input data to DataFrame
|
|
df = DataFrame(input_data['data'])
|
|
|
|
self.info(
|
|
f"Aggregating time series data for {len(df)} rows",
|
|
metadata=metadata
|
|
)
|
|
|
|
# Initialize result dictionary
|
|
result = {}
|
|
|
|
# Group by tag and name
|
|
grouped = df.groupby(['tag', 'name'])
|
|
|
|
for (tag, name), group in grouped:
|
|
# Get the aggregation function from model_tags
|
|
aggr_function = input_data['model_tags'].get(
|
|
name, {}).get('aggr_func', 'lts')
|
|
|
|
group.sort_values(by='timestamp', inplace=True)
|
|
|
|
# Get the latest timestamp
|
|
latest_timestamp = group['timestamp'].max()
|
|
|
|
aggr_value = self.apply_aggregation(
|
|
group, aggr_function, metadata)
|
|
|
|
if aggr_value == 'continue':
|
|
continue
|
|
|
|
self.debug(
|
|
f"Aggregated data: {aggr_value}",
|
|
metadata=metadata
|
|
)
|
|
self.debug(
|
|
f"Latest timestamp: {latest_timestamp}",
|
|
metadata=metadata
|
|
)
|
|
self.debug(
|
|
f"Groups: {group.to_string()}",
|
|
metadata=metadata
|
|
)
|
|
self.debug(
|
|
f"group name: {name}",
|
|
metadata=metadata
|
|
)
|
|
self.debug(
|
|
f"group tag: {tag}",
|
|
metadata=metadata
|
|
)
|
|
|
|
# Store the result
|
|
result[f"{tag}_{name}"] = {
|
|
'tag': tag,
|
|
'name': name,
|
|
'value': aggr_value,
|
|
'timestamp': latest_timestamp,
|
|
'aggregation_function': aggr_function
|
|
}
|
|
|
|
result_df = DataFrame(list(result.values()))
|
|
self.info(
|
|
f"Aggregated data has {len(result_df)} rows",
|
|
metadata=metadata
|
|
)
|
|
|
|
self.debug(
|
|
f"Aggregated data: {result_df.to_string()}",
|
|
metadata=metadata
|
|
)
|
|
|
|
return result_df.to_dict()
|
|
|
|
except Exception as e:
|
|
trace = traceback.format_exc()
|
|
|
|
self.send_notification(
|
|
metadata=metadata,
|
|
notification_id="AGGREGATION_ISSUES",
|
|
message=f"Error aggregating data: {e}",
|
|
block="aggregate_data",
|
|
level=NotificationLevel.ERROR,
|
|
attachment_content=trace
|
|
)
|
|
|
|
self.error(trace, metadata=metadata)
|
|
raise e
|
|
|
|
@activity.defn(name="data_quality_gate")
|
|
async def data_quality_gate(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
|
"""
|
|
Data quality gate activity. for each selected filter,
|
|
extracts filtered data, discards or keeps filtered data
|
|
based on the filter.
|
|
|
|
Args:
|
|
input_data (dict[str, Any]): The data to validate. Contains:
|
|
filters (dict[str, str]): The filters to apply. In format:
|
|
{filter_name: policy}.
|
|
filter_name: The name of the filter.
|
|
policy: The policy to apply. Can be "DISCARD" or "KEEP".
|
|
data (dict[str, Any]): The data to validate.
|
|
model_tags (dict[str, Any]): The tags of the model.
|
|
And it's respective configuration.
|
|
|
|
Returns:
|
|
dict[str, Any]: The data validated.
|
|
"""
|
|
|
|
metadata = input_data['metadata']
|
|
|
|
filters = input_data['filters']
|
|
data = DataFrame(input_data['data'])
|
|
model_tags = input_data['model_tags']
|
|
|
|
self.info(
|
|
f"Applying quality gate to data to {len(data)} rows",
|
|
metadata=metadata
|
|
)
|
|
|
|
tags = list(model_tags.keys())
|
|
|
|
data = data[data['name'].isin(tags)]
|
|
|
|
for filter_name, config in filters.items():
|
|
policy = config['policy']
|
|
if filter_name not in quality_gate_filters:
|
|
self.warning(
|
|
f"Filter {filter_name} not found",
|
|
metadata=metadata
|
|
)
|
|
continue
|
|
|
|
try:
|
|
filtered_data = quality_gate_filters[filter_name](
|
|
data, model_tags)
|
|
|
|
except Exception as e:
|
|
trace = traceback.format_exc()
|
|
self.send_notification(
|
|
metadata=metadata,
|
|
notification_id="DATA_QUALITY_GATE_ISSUES",
|
|
message=f"Error applying filter {filter_name}: {e}",
|
|
block="data_quality_gate",
|
|
level=NotificationLevel.ERROR,
|
|
attachment_content=trace
|
|
)
|
|
|
|
self.error(trace, metadata=metadata)
|
|
|
|
else:
|
|
if filtered_data.empty:
|
|
continue
|
|
|
|
message = f"{len(filtered_data)} rows has quality issues: {filter_name}: {policy}"
|
|
attachment = filtered_data.to_string()
|
|
|
|
self.send_notification(
|
|
metadata=metadata,
|
|
notification_id=f"DATA_QUALITY_GATE_ISSUES__{filter_name}",
|
|
message=message,
|
|
block="data_quality_gate",
|
|
level=NotificationLevel.WARNING,
|
|
attachment_content=attachment
|
|
)
|
|
|
|
if policy == "DISCARD":
|
|
data = data[~data.index.isin(filtered_data.index)]
|
|
|
|
self.info(
|
|
f"Data quality gate applied, final data has {len(data)} rows",
|
|
metadata=metadata
|
|
)
|
|
|
|
return data.to_dict()
|
|
|
|
@activity.defn(name="write_metrics")
|
|
async def write_metrics(self, input_data: dict[str, Any]):
|
|
"""
|
|
Write metrics to the database.
|
|
input_data:
|
|
metadata: dict[str, Any]
|
|
"""
|
|
metadata = input_data['metadata']
|
|
|
|
self.info(
|
|
f"Writing metrics for {metadata['model_name']}",
|
|
metadata=metadata
|
|
)
|
|
|
|
metrics.LABORIOUS_DATA_WRITTEN_COUNT.labels(
|
|
pod_id=self.pod_id,
|
|
model_name=metadata['model_name'],
|
|
pipeline_name=metadata['workflow_name']
|
|
).inc()
|
|
|
|
self.info(
|
|
f"Metrics written for {metadata['model_name']}",
|
|
metadata=metadata
|
|
)
|