1c78cfb558
Add write_metrics activity to Gates for database metric logging and import metrics module
273 lines
9.2 KiB
Python
273 lines
9.2 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.temporal.utils.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.debug(
|
|
f"Aggregating time series data: {df.to_string()}",
|
|
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.debug(
|
|
f"Aggregated data:\n {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.debug(
|
|
f"Applying quality gate to data: {data.to_string()}",
|
|
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.debug(
|
|
"Data quality gate applied",
|
|
metadata=metadata
|
|
)
|
|
|
|
return data.to_dict()
|
|
|
|
@activity.defn(name="write_metrics")
|
|
async def write_metrics(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
|
"""
|
|
Write metrics to the database.
|
|
input_data:
|
|
metadata: dict[str, Any]
|
|
"""
|
|
metadata = input_data['metadata']
|
|
|
|
metrics.LABORIOUS_DATA_WRITTEN_COUNT(
|
|
pod_id=metadata['pod_id'],
|
|
model_name=metadata['model_name'],
|
|
pipeline_name=metadata['pipeline_name']
|
|
).inc()
|