SIENTIAPDE-1030
Add comprehensive tests and utility functions for orchestrator activities - Introduced tests for the new Formatters class, covering methods for processing schedules and slots. - Enhanced SlotManager tests with update and delete slot functionalities. - Added TemporalManager tests for creating, updating, and deleting schedules, including frequency parsing. - Implemented utility functions for orchestrator operations, including frequency parsing and tag configuration building. - Created tests for utility functions to ensure correct behavior and integration with orchestrator activities. - Established a new converters module for parsing frequency strings into seconds.
This commit is contained in:
vitor-aignosi
@@ -5,12 +5,13 @@ with workflow.unsafe.imports_passed_through():
|
||||
from orchestrator.activities.couchbase import Couchbase
|
||||
from orchestrator.activities.temporal_manager import TemporalManager
|
||||
from orchestrator.activities.slot_manager import SlotManager
|
||||
from orchestrator.activities.formatters import Formatters
|
||||
from typing import Any
|
||||
from logging import Logger
|
||||
from sientia_do.notifications.handlers import NotificationHandler
|
||||
|
||||
|
||||
class Activities(Couchbase, TemporalManager, SlotManager):
|
||||
class Activities(Couchbase, TemporalManager, SlotManager, Formatters):
|
||||
|
||||
def __init__(self,
|
||||
temporal_client: Client,
|
||||
@@ -39,6 +40,10 @@ class Activities(Couchbase, TemporalManager, SlotManager):
|
||||
logger=logger,
|
||||
notification_handler=notification_handler)
|
||||
|
||||
Formatters.__init__(self,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler)
|
||||
|
||||
@activity.defn(name="prepare_activity")
|
||||
async def prepare_activity(self, input_data: dict[str, Any]):
|
||||
await super().prepare_activity(input_data)
|
||||
|
||||
@@ -1,10 +1,16 @@
|
||||
from sientia_do.notifications.models import NotificationLevel
|
||||
from temporalio import activity, workflow
|
||||
|
||||
with workflow.unsafe.imports_passed_through():
|
||||
from sientia_do.temporal.activities.base import BaseActivity
|
||||
import json
|
||||
from typing import Any
|
||||
from logging import Logger
|
||||
from sientia_do.notifications.handlers import NotificationHandler
|
||||
from sientia_do.temporal.activities.base import BaseActivity
|
||||
from orchestrator.utils.orchestrator_functions import (
|
||||
scouter, predictions_batch, gather_read_tags, build_tag_config
|
||||
)
|
||||
from math import ceil
|
||||
|
||||
|
||||
class Formatters(BaseActivity):
|
||||
@@ -13,7 +19,22 @@ class Formatters(BaseActivity):
|
||||
notification_handler=notification_handler)
|
||||
|
||||
@activity.defn(name="process_schedules")
|
||||
async def process_schedules(self, input_data: dict[str, Any]):
|
||||
async def process_schedules(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
||||
"""
|
||||
Process schedules. Generates a schedule config dictionary
|
||||
based on the input data workflow type.
|
||||
|
||||
Args:
|
||||
- input_data (dict[str, Any]): The input data containing
|
||||
the schedules to process.
|
||||
- pipelines (list[dict[str, Any]]): The schedules to process.
|
||||
|
||||
Returns:
|
||||
- dict[str, Any]: The schedule config dictionary
|
||||
"""
|
||||
|
||||
self.logger.info("Processing schedules...")
|
||||
|
||||
pipelines = input_data['pipelines']
|
||||
|
||||
schedule_config = {}
|
||||
@@ -21,115 +42,291 @@ class Formatters(BaseActivity):
|
||||
for pipeline in pipelines:
|
||||
if pipeline['workflow_type'] == 'scouter':
|
||||
schedule_config[pipeline['schedule_name']] = scouter(pipeline)
|
||||
elif pipeline['workflow_type'] == 'predictions_batch':
|
||||
schedule_config[pipeline['schedule_name']
|
||||
] = predictions_batch(pipeline)
|
||||
|
||||
return schedule_config
|
||||
|
||||
@activity.defn(name="process_slots")
|
||||
async def process_slots(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
||||
"""
|
||||
Extracts all read tags from input pipelines, divides them into slots and
|
||||
returns a slot config dictionary. If no ingestor is available, only one slot
|
||||
is created.
|
||||
|
||||
def common_config(config: dict[str, Any]):
|
||||
return {
|
||||
"workflow_type": "scouter",
|
||||
"schedule_name": config['schedule_name'],
|
||||
"frequency": config.get('frequency', '1m'),
|
||||
"max_retry_policy": config.get('max_retry_policy', 1),
|
||||
Args:
|
||||
- input_data (dict[str, Any]): The input data containing
|
||||
the schedules to process.
|
||||
- pipelines (list[dict[str, Any]]): The schedules to process.
|
||||
- opc_servers (list[str]): The OPC servers to create ingestor config.
|
||||
- active_ingestors (list[str]): The active ingestors to divide into slots.
|
||||
|
||||
"model_id": config['model_id'],
|
||||
"model_name": config['model_name'],
|
||||
}
|
||||
Returns:
|
||||
- dict[str, Any]: The slot config dictionary
|
||||
"""
|
||||
|
||||
self.logger.info("Processing slots...")
|
||||
|
||||
def scouter(config: dict[str, Any]):
|
||||
filters = {}
|
||||
for f in config['filters']:
|
||||
filters[f['filter_name']] = {
|
||||
"policy": f['policy']
|
||||
pipelines = input_data['pipelines']
|
||||
opc_servers_list = input_data['opc_servers']
|
||||
active_ingestors = input_data['active_ingestors']
|
||||
|
||||
opc_servers = {}
|
||||
for server in opc_servers_list:
|
||||
opc_servers[server['server_name']] = {
|
||||
**server,
|
||||
}
|
||||
|
||||
tags = list(gather_read_tags(pipelines).values())
|
||||
|
||||
number_of_tags = len(tags)
|
||||
number_of_slots = len(active_ingestors) if active_ingestors else 1
|
||||
tags_per_slot = ceil(number_of_tags / number_of_slots)
|
||||
|
||||
slot_config = {}
|
||||
last_index = 0
|
||||
|
||||
for i in range(1, number_of_slots):
|
||||
slot_config[f"{i}"] = {}
|
||||
for tag in tags[last_index:last_index + tags_per_slot]:
|
||||
slot_config = build_tag_config(
|
||||
tag, slot_config.copy(), opc_servers, i)
|
||||
last_index += tags_per_slot
|
||||
|
||||
slot_config[f"{number_of_slots}"] = {}
|
||||
for tag in tags[last_index:]:
|
||||
slot_config = build_tag_config(
|
||||
tag, slot_config.copy(), opc_servers, number_of_slots)
|
||||
|
||||
return slot_config
|
||||
|
||||
@activity.defn(name="create_schedule_config")
|
||||
async def create_schedule_config(self,
|
||||
input_data: dict[str, Any]) -> dict[str, Any]:
|
||||
"""
|
||||
Creates a schedule config dictionary based on the input data.
|
||||
Checks the existing schedule config and updates it with the new schedule config,
|
||||
deleting unnecessary schedules, creating new schedules and updating existing schedules.
|
||||
|
||||
Args:
|
||||
- input_data (dict[str, Any]): The input data containing
|
||||
the schedules to process.
|
||||
- current_schedule_config (dict[str, Any]): The current schedule
|
||||
config in Temporal server.
|
||||
- schedule_config (dict[str, Any]): The schedule config to process.
|
||||
|
||||
Returns:
|
||||
- dict[str, Any]: The schedule config dictionary
|
||||
"""
|
||||
|
||||
self.logger.info("Creating schedule config...")
|
||||
|
||||
current_schedule_config = input_data['current_schedule_config']
|
||||
schedule_config = input_data['schedule_config']
|
||||
|
||||
to_update = {}
|
||||
to_create = {}
|
||||
to_delete = []
|
||||
|
||||
for schedule_name, schedule in schedule_config.items():
|
||||
if schedule_name in current_schedule_config:
|
||||
if schedule != current_schedule_config[schedule_name]['data']:
|
||||
to_update[schedule_name] = schedule
|
||||
|
||||
elif schedule_name not in current_schedule_config:
|
||||
to_create[schedule_name] = schedule
|
||||
|
||||
for schedule_name in current_schedule_config:
|
||||
if schedule_name not in schedule_config:
|
||||
to_delete.append(schedule_name)
|
||||
|
||||
return {
|
||||
"to_update": to_update,
|
||||
"to_create": to_create,
|
||||
"to_delete": to_delete
|
||||
}
|
||||
|
||||
tags = {}
|
||||
for tag in config['read_tags']:
|
||||
tags[tag['tag_name']] = {
|
||||
"aggr_func": tag.get('aggr_func', 'lts'),
|
||||
"data_range": tag.get('data_range', [-100, 100])
|
||||
@activity.defn(name="create_slot_config")
|
||||
async def create_slot_config(self,
|
||||
input_data: dict[str, Any]) -> dict[str, Any]:
|
||||
"""
|
||||
Creates a slot config dictionary based on the input data.
|
||||
Checks the existing slot config and updates it with the new slot config,
|
||||
deleting unnecessary slots.
|
||||
|
||||
Args:
|
||||
- input_data (dict[str, Any]): The input data containing
|
||||
the slots to process.
|
||||
- current_slot_config (dict[str, Any]): The current slot
|
||||
config in Temporal server.
|
||||
- slot_config (dict[str, Any]): The slot config to process.
|
||||
|
||||
Returns:
|
||||
- dict[str, Any]: The slot config dictionary
|
||||
"""
|
||||
|
||||
self.logger.info("Creating slot config...")
|
||||
|
||||
current_slot_config = input_data['current_slot_config']
|
||||
slot_config = input_data['slot_config']
|
||||
to_delete = []
|
||||
|
||||
number_of_current_slots = len(current_slot_config)
|
||||
number_of_slots = len(slot_config)
|
||||
|
||||
if number_of_current_slots > number_of_slots:
|
||||
to_delete = [str(i) for i in range(
|
||||
number_of_slots + 1, number_of_current_slots + 1)]
|
||||
|
||||
return {
|
||||
"to_delete": to_delete,
|
||||
"to_insert": slot_config
|
||||
}
|
||||
|
||||
return {
|
||||
**common_config(config),
|
||||
def send_success_report(self, message: str, notification_id: str) -> None:
|
||||
self.notification_handler.build_and_send_notification(
|
||||
notification_id,
|
||||
message,
|
||||
"report_orchestration",
|
||||
NotificationLevel.INFO
|
||||
)
|
||||
|
||||
"topic": f"raw_{config['schedule_name']}",
|
||||
"trigger_laborious": False,
|
||||
"filters": filters,
|
||||
"schema": "sientia_data",
|
||||
"table_name": "laborious_data",
|
||||
"retention_time": config.get('tag_retention_minutes', 60) * 60,
|
||||
"model_tags": tags
|
||||
}
|
||||
def send_error_report(self, message: str, notification_id: str,
|
||||
attachment: dict[str, Any]) -> None:
|
||||
self.notification_handler.build_and_send_notification(
|
||||
notification_id,
|
||||
message,
|
||||
"report_orchestration",
|
||||
NotificationLevel.ERROR,
|
||||
attachment_content=json.dumps(attachment, indent=4, sort_keys=True)
|
||||
)
|
||||
|
||||
def parse_report(self, input_data: dict[str, Any]) -> tuple[list[str], list[str]]:
|
||||
success_keys = [key for key, value
|
||||
in input_data.items() if value['success']]
|
||||
|
||||
def overlap_filter_config(base_filter_config: dict[str, Any], config: dict[str, Any]):
|
||||
for fil in config['filters']:
|
||||
base_filter_config[fil['filter_name']] = {
|
||||
"policy": fil['policy'],
|
||||
"config": fil.get('config', {})
|
||||
}
|
||||
error_keys = [key for key, value
|
||||
in input_data.items() if not value['success']]
|
||||
|
||||
return base_filter_config
|
||||
return success_keys, error_keys
|
||||
|
||||
@activity.defn(name="report_schedule_orchestration")
|
||||
async def report_schedule_orchestration(self,
|
||||
input_data: dict[str, Any]) -> None:
|
||||
"""
|
||||
Reports the orchestration result to the notification handler.
|
||||
|
||||
def process_path_priority(path_priority: list[str]):
|
||||
for priority in path_priority[:]:
|
||||
if priority not in ["STOP", "CONTINUE", "REPEAT"]:
|
||||
path_priority.remove(priority)
|
||||
Args:
|
||||
- input_data (dict[str, Any]): The input data containing
|
||||
the orchestration result.
|
||||
- created_schedules (dict[str, Any]): The created schedules.
|
||||
- updated_schedules (dict[str, Any]): The updated schedules.
|
||||
- deleted_schedules (list[str]): The deleted schedules.
|
||||
"""
|
||||
|
||||
if len(path_priority) != 3:
|
||||
for priority in ["STOP", "CONTINUE", "REPEAT"]:
|
||||
if priority not in path_priority:
|
||||
path_priority.append(priority)
|
||||
self.logger.info("Reporting orchestration...")
|
||||
|
||||
return path_priority
|
||||
created_schedules = input_data['created_schedules']
|
||||
updated_schedules = input_data['updated_schedules']
|
||||
deleted_schedules = input_data['deleted_schedules']
|
||||
|
||||
# Send report for created schedules
|
||||
if len(created_schedules) > 0:
|
||||
success_keys, error_keys = self.parse_report(created_schedules)
|
||||
|
||||
def predictions_batch(config: dict[str, Any]):
|
||||
tags = {}
|
||||
for tag in config['write_tags']:
|
||||
if tag['server_name'] not in tags:
|
||||
tags[tag['server_name']] = {}
|
||||
if len(success_keys) > 0:
|
||||
self.send_success_report(
|
||||
f"Created schedules: \n {', '.join(success_keys)}",
|
||||
"REPORT_ORCHESTRATION_CREATED_SCHEDULES"
|
||||
)
|
||||
|
||||
tag_type = tag['type']
|
||||
if len(error_keys) > 0:
|
||||
self.send_error_report(
|
||||
f"Failed to create schedules: \n {', '.join(error_keys)}",
|
||||
"REPORT_ORCHESTRATION_CREATED_SCHEDULES",
|
||||
created_schedules
|
||||
)
|
||||
|
||||
if tag_type == 'prediction' or tag_type == 'confidence':
|
||||
tag_type_str = f"{tag_type}_tags"
|
||||
# Send report for updated schedules
|
||||
if len(updated_schedules) > 0:
|
||||
success_keys, error_keys = self.parse_report(updated_schedules)
|
||||
|
||||
if tag_type_str not in tags[tag['server_name']]:
|
||||
tags[tag['server_name']][tag_type_str] = {}
|
||||
if len(success_keys) > 0:
|
||||
self.send_success_report(
|
||||
f"Updated schedules: \n {', '.join(success_keys)}",
|
||||
"REPORT_ORCHESTRATION_UPDATED_SCHEDULES"
|
||||
)
|
||||
|
||||
tags[tag['server_name']][tag_type_str][tag['addr']] = {
|
||||
"data_type": tag.get('data_type', 'float'),
|
||||
}
|
||||
if len(error_keys) > 0:
|
||||
self.send_error_report(
|
||||
f"Failed to update schedules: \n {', '.join(error_keys)}",
|
||||
"REPORT_ORCHESTRATION_UPDATED_SCHEDULES",
|
||||
updated_schedules
|
||||
)
|
||||
|
||||
path_priority = process_path_priority(config.get(
|
||||
'path_priority', ["STOP", "CONTINUE", "REPEAT"]))
|
||||
if len(deleted_schedules) > 0:
|
||||
success_keys, error_keys = self.parse_report(deleted_schedules)
|
||||
|
||||
return {
|
||||
**common_config(config),
|
||||
if len(success_keys) > 0:
|
||||
self.send_success_report(
|
||||
f"Deleted schedules: \n {', '.join(success_keys)}",
|
||||
"REPORT_ORCHESTRATION_DELETED_SCHEDULES"
|
||||
)
|
||||
|
||||
"query": config['query'],
|
||||
"schema": "sientia_data",
|
||||
"table_name": "predictions",
|
||||
"retention_time": config.get('model_retention_minutes', 60) * 60,
|
||||
"opc_output_config": tags,
|
||||
"input_filters": overlap_filter_config({
|
||||
"EMPTY_DATA": {
|
||||
"policy": "STOP"
|
||||
}
|
||||
}, config['input_filters']),
|
||||
"mlflow_transform_filters": overlap_filter_config({
|
||||
"API_ERROR": {
|
||||
"policy": "STOP"
|
||||
}
|
||||
}, config['mlflow_transform_filters']),
|
||||
"mlflow_predict_filters": overlap_filter_config({
|
||||
"API_ERROR": {
|
||||
"policy": "STOP"
|
||||
}
|
||||
}, config['mlflow_predict_filters']),
|
||||
"path_priority": path_priority
|
||||
}
|
||||
if len(error_keys) > 0:
|
||||
self.send_error_report(
|
||||
f"Failed to delete schedules: \n {', '.join(error_keys)}",
|
||||
"REPORT_ORCHESTRATION_DELETED_SCHEDULES",
|
||||
deleted_schedules
|
||||
)
|
||||
|
||||
@activity.defn(name="report_slot_orchestration")
|
||||
async def report_slot_orchestration(self,
|
||||
input_data: dict[str, Any]) -> None:
|
||||
"""
|
||||
Reports the orchestration result to the notification handler.
|
||||
|
||||
Args:
|
||||
- input_data (dict[str, Any]): The input data containing
|
||||
the orchestration result.
|
||||
- inserted_slots (dict[str, Any]): The inserted slots.
|
||||
- deleted_slots (list[str]): The deleted slots.
|
||||
"""
|
||||
|
||||
self.logger.info("Reporting orchestration...")
|
||||
|
||||
inserted_slots = input_data['inserted_slots']
|
||||
deleted_slots = input_data['deleted_slots']
|
||||
|
||||
if len(inserted_slots) > 0:
|
||||
success_keys, error_keys = self.parse_report(inserted_slots)
|
||||
|
||||
if len(success_keys) > 0:
|
||||
self.send_success_report(
|
||||
f"Inserted slots: \n {', '.join(success_keys)}",
|
||||
"REPORT_ORCHESTRATION_INSERTED_SLOTS"
|
||||
)
|
||||
|
||||
if len(error_keys) > 0:
|
||||
self.send_error_report(
|
||||
f"Failed to insert slots: \n {', '.join(error_keys)}",
|
||||
"REPORT_ORCHESTRATION_INSERTED_SLOTS",
|
||||
inserted_slots
|
||||
)
|
||||
|
||||
if len(deleted_slots) > 0:
|
||||
success_keys, error_keys = self.parse_report(deleted_slots)
|
||||
|
||||
if len(success_keys) > 0:
|
||||
self.send_success_report(
|
||||
f"Deleted slots: \n {', '.join(success_keys)}",
|
||||
"REPORT_ORCHESTRATION_DELETED_SLOTS"
|
||||
)
|
||||
|
||||
if len(error_keys) > 0:
|
||||
self.send_error_report(
|
||||
f"Failed to delete slots: \n {', '.join(error_keys)}",
|
||||
"REPORT_ORCHESTRATION_DELETED_SLOTS",
|
||||
deleted_slots
|
||||
)
|
||||
|
||||
@@ -32,23 +32,20 @@ class SlotManager(Redis):
|
||||
|
||||
slot_keys = self.redis_client.keys("slot:opc_tags:*")
|
||||
|
||||
if slot_keys:
|
||||
decoded_keys = [key.decode('utf-8') for key in slot_keys]
|
||||
values = self.redis_client.mget(decoded_keys)
|
||||
self.logger.debug("Slot keys: %s", slot_keys)
|
||||
|
||||
for i, key in enumerate(decoded_keys):
|
||||
value = values[i]
|
||||
if value is not None:
|
||||
try:
|
||||
opc_slots[key] = value.decode('utf-8')
|
||||
except (UnicodeDecodeError, AttributeError):
|
||||
opc_slots[key] = value
|
||||
else:
|
||||
opc_slots[key] = None
|
||||
if slot_keys:
|
||||
if isinstance(slot_keys[0], bytes):
|
||||
decoded_keys = [key.decode('utf-8') for key in slot_keys]
|
||||
else:
|
||||
decoded_keys = slot_keys
|
||||
|
||||
for key in decoded_keys:
|
||||
opc_slots[key] = self.get(key)
|
||||
|
||||
self.logger.info(f"Loaded {len(opc_slots)} OPC slots")
|
||||
|
||||
self.logger.debug("OPC slots: \n %s",
|
||||
self.logger.debug("Loaded: \n %s",
|
||||
json.dumps(opc_slots, indent=4, sort_keys=True))
|
||||
|
||||
return opc_slots
|
||||
@@ -71,3 +68,86 @@ class SlotManager(Redis):
|
||||
self.logger.debug("Active ingestors: \n %s", active_ingestors)
|
||||
|
||||
return [ingestor.decode('utf-8') for ingestor in active_ingestors]
|
||||
|
||||
@activity.defn(name="update_slots")
|
||||
async def update_slots(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
||||
"""
|
||||
Update OPC slots in Redis
|
||||
|
||||
Args:
|
||||
- input_data (dict[str, Any]): The input data containing
|
||||
the slots to update.
|
||||
- to_insert (dict[str, Any]): The slots to insert.
|
||||
|
||||
Returns:
|
||||
- report (dict[str, Any]): A report of the updated slots.
|
||||
"""
|
||||
|
||||
to_insert = input_data['to_insert']
|
||||
self.logger.info("Updating OPC slots...")
|
||||
|
||||
report = {}
|
||||
|
||||
for slot in to_insert:
|
||||
try:
|
||||
self.set(f"slot:opc_tags:{slot}",
|
||||
to_insert[slot], ttl=None)
|
||||
report[slot] = {
|
||||
"success": True,
|
||||
"message": "Slot updated successfully"
|
||||
}
|
||||
except Exception as e:
|
||||
self.logger.error("Failed to update slot %s: %s",
|
||||
slot, str(e))
|
||||
report[slot] = {
|
||||
"success": False,
|
||||
"message": str(e)
|
||||
}
|
||||
|
||||
self.logger.info(f"Updated {len(to_insert)} OPC slots")
|
||||
|
||||
self.logger.debug("Report: \n %s",
|
||||
json.dumps(report, indent=4, sort_keys=True))
|
||||
|
||||
return report
|
||||
|
||||
@activity.defn(name="delete_slots")
|
||||
async def delete_slots(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
||||
"""
|
||||
Delete OPC slots from Redis
|
||||
|
||||
Args:
|
||||
- input_data (dict[str, Any]): The input data containing
|
||||
the slots to delete.
|
||||
- to_delete (list[str]): The slots to delete.
|
||||
|
||||
Returns:
|
||||
- report (dict[str, Any]): A report of the deleted slots.
|
||||
"""
|
||||
|
||||
to_delete = input_data['to_delete']
|
||||
self.logger.info("Deleting OPC slots...")
|
||||
|
||||
report = {}
|
||||
|
||||
for slot in to_delete:
|
||||
try:
|
||||
self.redis_client.delete(f"slot:opc_tags:{slot}")
|
||||
report[slot] = {
|
||||
"success": True,
|
||||
"message": "Slot deleted successfully"
|
||||
}
|
||||
except Exception as e:
|
||||
self.logger.error("Failed to delete slot %s: %s",
|
||||
slot, str(e))
|
||||
report[slot] = {
|
||||
"success": False,
|
||||
"message": str(e)
|
||||
}
|
||||
|
||||
self.logger.info(f"Deleted {len(to_delete)} OPC slots")
|
||||
|
||||
self.logger.debug("Report: \n %s",
|
||||
json.dumps(report, indent=4, sort_keys=True))
|
||||
|
||||
return report
|
||||
|
||||
@@ -1,5 +1,7 @@
|
||||
from temporalio import activity, workflow
|
||||
from temporalio.client import Client
|
||||
from temporalio.client import (
|
||||
Client, Schedule, ScheduleActionStartWorkflow, ScheduleIntervalSpec, ScheduleSpec, ScheduleUpdate, ScheduleUpdateInput)
|
||||
from temporalio.common import SearchAttributeKey, SearchAttributePair, TypedSearchAttributes
|
||||
|
||||
with workflow.unsafe.imports_passed_through():
|
||||
from typing import Any
|
||||
@@ -8,7 +10,9 @@ with workflow.unsafe.imports_passed_through():
|
||||
from sientia_do.temporal.activities.base import BaseActivity
|
||||
from google.protobuf.json_format import MessageToDict
|
||||
import base64
|
||||
from datetime import timedelta
|
||||
import json
|
||||
from orchestrator.utils.converters import parse_frequency
|
||||
|
||||
|
||||
class TemporalManager(BaseActivity):
|
||||
@@ -17,6 +21,13 @@ class TemporalManager(BaseActivity):
|
||||
|
||||
self.temporal_client = temporal_client
|
||||
|
||||
self.schedule_handles = {}
|
||||
|
||||
self.model_id_id_key = SearchAttributeKey.for_keyword("model_id")
|
||||
self.model_name_id_key = SearchAttributeKey.for_keyword("model_name")
|
||||
self.orchestrated_id_key = SearchAttributeKey.for_keyword(
|
||||
"orchestrated")
|
||||
|
||||
BaseActivity.__init__(self,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler)
|
||||
@@ -41,7 +52,9 @@ class TemporalManager(BaseActivity):
|
||||
if search_attrs.get("Orchestrated", ["false"]) == ["true"]:
|
||||
schedule_id = schedule.id
|
||||
|
||||
handle = self.temporal_client.get_schedule(schedule_id)
|
||||
handle = self.temporal_client.get_schedule_handle(schedule_id)
|
||||
|
||||
self.schedule_handles[schedule_id] = handle
|
||||
|
||||
desc = await handle.describe()
|
||||
|
||||
@@ -54,7 +67,6 @@ class TemporalManager(BaseActivity):
|
||||
orchestrated_schedules[schedule_id] = {
|
||||
'frequency': frequency,
|
||||
'data': json.loads(data),
|
||||
'handle': handle
|
||||
}
|
||||
|
||||
self.logger.info("Found %d orchestrated schedules",
|
||||
@@ -64,3 +76,190 @@ class TemporalManager(BaseActivity):
|
||||
orchestrated_schedules)
|
||||
|
||||
return orchestrated_schedules
|
||||
|
||||
@activity.defn(name="create_schedules")
|
||||
async def create_schedules(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
||||
"""
|
||||
Create schedules in Temporal
|
||||
|
||||
Args:
|
||||
- input_data (dict[str, Any]): The input data containing
|
||||
the schedules to create.
|
||||
- schedules (dict[str, Any]): The schedules to create.
|
||||
|
||||
Returns:
|
||||
- dict[str, Any]: A report of the created schedules.
|
||||
"""
|
||||
|
||||
schedules = input_data['schedules']
|
||||
|
||||
report = {}
|
||||
|
||||
for schedule_name, schedule in schedules.items():
|
||||
search_attributes = TypedSearchAttributes([
|
||||
SearchAttributePair(
|
||||
key=self.model_id_id_key,
|
||||
value=schedule['model_id']
|
||||
),
|
||||
SearchAttributePair(
|
||||
key=self.model_name_id_key,
|
||||
value=schedule['model_name']
|
||||
),
|
||||
SearchAttributePair(
|
||||
key=self.orchestrated_id_key,
|
||||
value="true"
|
||||
)
|
||||
])
|
||||
workflow_type = schedule['workflow_type']
|
||||
|
||||
try:
|
||||
await self.temporal_client.create_schedule(
|
||||
schedule_name,
|
||||
Schedule(
|
||||
action=ScheduleActionStartWorkflow(
|
||||
workflow=workflow_type,
|
||||
args=schedule,
|
||||
id=schedule_name,
|
||||
task_queue=f"{workflow_type}-queue"
|
||||
),
|
||||
spec=ScheduleSpec(
|
||||
intervals=[
|
||||
ScheduleIntervalSpec(
|
||||
every=timedelta(seconds=parse_frequency(
|
||||
schedule.get('frequency', '1m')))
|
||||
)
|
||||
]
|
||||
)
|
||||
),
|
||||
search_attributes=search_attributes
|
||||
)
|
||||
|
||||
report[schedule_name] = {
|
||||
"success": True,
|
||||
"message": "Schedule created successfully"
|
||||
}
|
||||
except Exception as e:
|
||||
self.logger.error("Failed to create schedule %s: %s",
|
||||
schedule_name, str(e))
|
||||
report[schedule_name] = {
|
||||
"success": False,
|
||||
"message": str(e)
|
||||
}
|
||||
|
||||
self.logger.info(f"Processed {len(schedules)} schedules")
|
||||
|
||||
self.logger.debug("\n %s",
|
||||
json.dumps(report, indent=4, sort_keys=True))
|
||||
|
||||
return report
|
||||
|
||||
@activity.defn(name="update_schedules")
|
||||
async def update_schedules(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
||||
"""
|
||||
Update schedules in Temporal
|
||||
|
||||
Args:
|
||||
- input_data (dict[str, Any]): The input data containing
|
||||
the schedules to update.
|
||||
- schedules (dict[str, Any]): The schedules to update.
|
||||
|
||||
Returns:
|
||||
- dict[str, Any]: A report of the updated schedules.
|
||||
"""
|
||||
|
||||
schedules = input_data['schedules']
|
||||
|
||||
report = {}
|
||||
|
||||
for schedule_name, schedule in schedules.items():
|
||||
try:
|
||||
handler = self.schedule_handles.get(schedule_name)
|
||||
|
||||
if not handler:
|
||||
raise ValueError(f"Schedule {schedule_name} not found")
|
||||
|
||||
async def update_schedule(input_data: ScheduleUpdateInput) -> ScheduleUpdate:
|
||||
schedule_action = input_data.description.schedule.action
|
||||
|
||||
if hasattr(schedule_action, "args"):
|
||||
schedule_action.args = schedule
|
||||
|
||||
input_data.description.schedule.spec.intervals = [
|
||||
ScheduleIntervalSpec(
|
||||
every=timedelta(seconds=parse_frequency(
|
||||
schedule.get('frequency', '1m')))
|
||||
)
|
||||
]
|
||||
|
||||
return ScheduleUpdate(schedule=input_data.description.schedule)
|
||||
|
||||
await handler.update(update_schedule)
|
||||
|
||||
del update_schedule
|
||||
|
||||
report[schedule_name] = {
|
||||
"success": True,
|
||||
"message": "Schedule updated successfully"
|
||||
}
|
||||
except Exception as e:
|
||||
self.logger.error("Failed to update schedule %s: %s",
|
||||
schedule_name, str(e))
|
||||
report[schedule_name] = {
|
||||
"success": False,
|
||||
"message": str(e)
|
||||
}
|
||||
|
||||
self.logger.info(f"Processed {len(schedules)} schedules")
|
||||
|
||||
self.logger.debug("\n %s",
|
||||
json.dumps(report, indent=4, sort_keys=True))
|
||||
|
||||
return report
|
||||
|
||||
@activity.defn(name="delete_schedules")
|
||||
async def delete_schedules(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
||||
"""
|
||||
Delete schedules in Temporal
|
||||
|
||||
Args:
|
||||
- input_data (dict[str, Any]): The input data containing
|
||||
the schedules to delete.
|
||||
- schedules (list[str]): The schedules to delete.
|
||||
|
||||
Returns:
|
||||
- dict[str, Any]: A report of the deleted schedules.
|
||||
"""
|
||||
|
||||
schedules = input_data['schedules']
|
||||
|
||||
report = {}
|
||||
|
||||
for schedule_name in schedules:
|
||||
try:
|
||||
handler = self.schedule_handles.get(schedule_name)
|
||||
|
||||
if not handler:
|
||||
raise ValueError(f"Schedule {schedule_name} not found")
|
||||
|
||||
await handler.delete()
|
||||
|
||||
del self.schedule_handles[schedule_name]
|
||||
|
||||
report[schedule_name] = {
|
||||
"success": True,
|
||||
"message": "Schedule deleted successfully"
|
||||
}
|
||||
except Exception as e:
|
||||
self.logger.error("Failed to delete schedule %s: %s",
|
||||
schedule_name, str(e))
|
||||
report[schedule_name] = {
|
||||
"success": False,
|
||||
"message": str(e)
|
||||
}
|
||||
|
||||
self.logger.info(f"Processed {len(schedules)} schedules")
|
||||
|
||||
self.logger.debug("\n %s",
|
||||
json.dumps(report, indent=4, sort_keys=True))
|
||||
|
||||
return report
|
||||
|
||||
Reference in New Issue
Block a user