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:
@@ -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()
|
||||
|
||||
|
||||
@@ -52,6 +52,8 @@ class FormatAndExportPrediction():
|
||||
'timestamp': input_data['timestamp'],
|
||||
'model_id': input_data['model_id'],
|
||||
'prediction_confidence': prediction_confidence,
|
||||
'prediction_store_policy': input_data.get(
|
||||
'prediction_store_policy', 'lts:1')
|
||||
},
|
||||
retry_policy=retry_policy,
|
||||
start_to_close_timeout=timedelta(seconds=60)
|
||||
|
||||
Reference in New Issue
Block a user