SIENTIAPDE-1182

Implement prediction store policy handling in Gates activity

- Added a new method `get_prediction_store_policy` to validate and parse the prediction store policy.
- Updated `format_prediction` method to utilize the new policy handling, allowing for sorting of predictions based on the specified policy.
- Enhanced test coverage for the new policy handling, including various scenarios for valid and invalid policies.
- Removed the obsolete `coverage.sh` script.
This commit is contained in:
vitor-aignosi
2025-08-28 16:31:55 -03:00
parent 7f3ecc9add
commit 30c1d6746a
6 changed files with 201 additions and 8 deletions

View File

@@ -242,6 +242,28 @@ class Gates(BaseActivity):
self.info("Nothing was filtered by the mlflow content gate", metadata)
return None, 0, ""
def get_prediction_store_policy(self,
prediction_store_policy: str,
metadata: dict[str, Any]) -> tuple[str, int]:
policy_elements = prediction_store_policy.split(':')
if len(policy_elements) < 2:
self.error(
f"Invalid prediction store policy: {prediction_store_policy}, using default policy", metadata)
return 'lts', 1
policy_type = policy_elements[0]
policy_value = policy_elements[1]
# If the policy_type is not lts or erl, we use the default policy
# If the policty_value is not a number or 0, we use the default policy
if policy_type not in ['lts', 'erl'] or not policy_value.isdigit() or int(policy_value) == 0:
self.error(
f"Invalid prediction store policy: {prediction_store_policy}, using default policy", metadata)
return 'lts', 1
return policy_type, int(policy_value)
@activity.defn(name="format_prediction")
async def format_prediction(self, input_data: dict[str, Any]) -> dict[Any, Any]:
"""
@@ -252,21 +274,58 @@ class Gates(BaseActivity):
- timestamp (str): The timestamp of the data.
- model_id (str): The id of the model.
- prediction_confidence (float): The confidence of the prediction.
- prediction_store_policy (str): The policy to store the prediction.
Returns:
dict: The formatted data.
"""
metadata = input_data['metadata']
prediction_store_policy = input_data['prediction_store_policy']
self.info("Formatting prediction...", metadata)
data = DataFrame(input_data['data'])
data['timestamp'] = input_data['timestamp']
self.debug(
f"Prediction store policy: {prediction_store_policy}", metadata)
policy_type, policy_value = self.get_prediction_store_policy(
prediction_store_policy, metadata)
# If data has no timestamp, we use the default timestamp and not sort the data
if 'timestamp' not in data.columns:
self.warning(
"Data has no timestamp, using default timestamp", metadata)
data['timestamp'] = input_data['timestamp']
else:
self.debug(
f"Data has timestamp, sorting data by timestamp", metadata)
# If policy_type is lts, we need to sort the data by timestamp descending and take the first policy_value rows
if policy_type == 'lts':
self.debug(
f"Sorting data by timestamp descending", metadata)
data = data.sort_values(by='timestamp', ascending=False)
# If policy_type is erl, we need to sort the data by timestamp ascending and take the first policy_value rows
elif policy_type == 'erl':
self.debug(
f"Sorting data by timestamp ascending", metadata)
data = data.sort_values(by='timestamp', ascending=True)
else:
self.error(
f"Invalid policy type: {policy_type}, using default policy", metadata)
raise ValueError(
f"Invalid policy type: {policy_type}")
data = data.head(int(policy_value))
data['model_id'] = input_data['model_id']
data['prediction_confidence'] = input_data['prediction_confidence']
data['prediction_status'] = 'Good'
data['comments'] = ""
data = data.sort_values(by='timestamp')
data = data.sort_values(by='timestamp', ascending=False)
data = data.reset_index(drop=True)
self.info(f"Prediction formatted: {data.size} rows", metadata)
self.debug(f"Prediction data: {data.to_string()}", metadata)
return data.to_dict()