diff --git a/orchestrator/activities/activities.py b/orchestrator/activities/activities.py index 2bf0388..ff36b17 100644 --- a/orchestrator/activities/activities.py +++ b/orchestrator/activities/activities.py @@ -82,4 +82,9 @@ class Activities( # Couchbase, notification_handler=notification_handler) def shutdown(self): + """ + Shutdown the MongoDB connection and clean up resources. + """ MongoDB.shutdown(self) + Postgres.close(self) + Email.shutdown(self) diff --git a/orchestrator/activities/email.py b/orchestrator/activities/email.py index e47ec1e..1401876 100644 --- a/orchestrator/activities/email.py +++ b/orchestrator/activities/email.py @@ -42,6 +42,12 @@ class Email(BaseActivity): logger=logger, notification_handler=notification_handler) + def shutdown(self): + """ + Shutdown the Email connection and clean up resources. + """ + self.server.quit() + @activity.defn(name="build_email_html") async def build_email_html(self, input_data: dict[str, Any]) -> dict[str, Any]: """ @@ -105,9 +111,15 @@ class Email(BaseActivity): def try_send_email(self, msg: MIMEMultipart, receivers: str): """ - Sends an email to the receivers. - """ + Sends an email to the receivers with automatic reconnection handling. + Args: + msg (MIMEMultipart): The email message to send. + receivers (str): Comma-separated list of email addresses to send to. + + Raises: + Exception: If email sending fails after reconnection attempts. + """ try: self.server.sendmail( self.sender_email, receivers, msg.as_string()) diff --git a/orchestrator/activities/mongo_db.py b/orchestrator/activities/mongo_db.py index 80acd0c..be0d40f 100644 --- a/orchestrator/activities/mongo_db.py +++ b/orchestrator/activities/mongo_db.py @@ -64,7 +64,7 @@ class MongoDB(BaseActivity): def shutdown(self): """ - Close the MongoDB client connection. + Shutdown the MongoDB connection and clean up resources. """ try: if self.client: @@ -81,6 +81,16 @@ class MongoDB(BaseActivity): self.shutdown() def find(self, collection_name: str, filters: dict[str, Any]) -> list[dict[str, Any]]: + """ + Find documents in a MongoDB collection based on the provided filters. + + Args: + collection_name (str): The name of the collection to search in. + filters (dict[str, Any]): The query filters to apply. + + Returns: + list[dict[str, Any]]: List of documents matching the filters, with _id fields removed. + """ collection = self.database[collection_name] documents = list(collection.find(filters, {"_id": 0})) diff --git a/orchestrator/activities/slot_manager.py b/orchestrator/activities/slot_manager.py index 18e96d0..14aa70c 100644 --- a/orchestrator/activities/slot_manager.py +++ b/orchestrator/activities/slot_manager.py @@ -202,6 +202,14 @@ class SlotManager(Redis): async def get_last_data_timestamp(self, input_data: dict[str, Any]) -> str | None: """ Gets the last data timestamp from redis. + + Args: + input_data (dict[str, Any]): The input data containing: + - metadata (dict): Metadata for logging purposes. + - mail_type (str): The type of mail to get timestamp for. + + Returns: + str | None: The last data timestamp as a string, or None if no timestamp exists. """ metadata = input_data['metadata'] key = f"notification_last_timestamp:{input_data['mail_type']}" @@ -233,6 +241,15 @@ class SlotManager(Redis): async def put_last_data_timestamp(self, input_data: dict[str, Any]): """ Puts the last data timestamp into redis. + + Args: + input_data (dict[str, Any]): The input data containing: + - metadata (dict): Metadata for logging purposes. + - data (list[dict]): The data to extract timestamp from. + - mail_type (str): The type of mail to store timestamp for. + + Returns: + str | None: The last data timestamp that was stored, or None if no data exists. """ metadata = input_data['metadata'] key = f"notification_last_timestamp:{input_data['mail_type']}" @@ -270,7 +287,18 @@ class SlotManager(Redis): @activity.defn(name="filter_notification_alerts") async def filter_notification_alerts(self, input_data: dict[str, Any]) -> dict[str, Any]: """ - Filter notification alerts + Filter notification alerts based on sending configurations and notification package. + + Args: + input_data (dict[str, Any]): The input data containing: + - metadata (dict): Metadata for logging purposes. + - notification_package (list): The package of notifications to filter. + - sending_configs (list): The configurations for sending notifications. + Each config should have 'group_name', 'contents', and optionally 'ignore' fields. + - notification_ttl (int): Time to live for notifications in seconds. + + Returns: + dict[str, Any]: The filtered receiver groups with their notifications. """ metadata = input_data['metadata'] notification_package = input_data['notification_package'] @@ -330,7 +358,13 @@ class SlotManager(Redis): @activity.defn(name="store_notification_cache") async def store_notification_cache(self, input_data: dict[str, Any]) -> None: """ - Store notification cache + Store notification cache in Redis to track recently sent notifications. + + Args: + input_data (dict[str, Any]): The input data containing: + - metadata (dict): Metadata for logging purposes. + - log_report (list[dict]): The log report containing notification statuses. + - sent_ttl (int): Time to live for sent notification cache in seconds. """ metadata = input_data['metadata'] log_report = DataFrame(input_data['log_report']) diff --git a/orchestrator/activities/temporal_manager.py b/orchestrator/activities/temporal_manager.py index eb7d361..d6f64bf 100644 --- a/orchestrator/activities/temporal_manager.py +++ b/orchestrator/activities/temporal_manager.py @@ -42,6 +42,10 @@ class TemporalManager(BaseActivity): notification_handler=notification_handler) async def connect_to_temporal(self): + """ + Connect to Temporal server namespaces for scouter and laborious workflows. + Creates client connections to both namespaces and stores them for later use. + """ self.logger.info( f"Connecting to Temporal side namespaces at {self.temporal_host}") self.logger.info(f"Scouter namespace: {self.scouter_namespace}") diff --git a/orchestrator/utils/connectors_config.py b/orchestrator/utils/connectors_config.py index 7237a05..8ab0c76 100644 --- a/orchestrator/utils/connectors_config.py +++ b/orchestrator/utils/connectors_config.py @@ -2,6 +2,12 @@ from os import getenv def build_redis_config(): + """ + Build Redis configuration from environment variables. + + Returns: + dict: Redis configuration with host, port, username, and password. + """ return { 'host': getenv('REDIS_HOST', 'localhost'), 'port': int(getenv('REDIS_PORT', '6379')), @@ -11,6 +17,12 @@ def build_redis_config(): def build_mongodb_config(): + """ + Build MongoDB configuration from environment variables. + + Returns: + dict: MongoDB configuration with connection string, database name, and TTL index seconds. + """ username = getenv('MONGODB_USERNAME', 'root') password = getenv('MONGODB_PASSWORD', 'wKZDbMNU1c') uri = getenv('MONGODB_URL', 'localhost:27018') @@ -24,6 +36,12 @@ def build_mongodb_config(): def build_couchbase_config(): + """ + Build Couchbase configuration from environment variables. + + Returns: + dict: Couchbase configuration with connection string, username, and password. + """ return { 'connection_string': getenv('COUCHBASE_CONNECTION_STRING', 'couchbase://localhost'), 'username': getenv('COUCHBASE_USERNAME', 'sientia'), @@ -32,6 +50,12 @@ def build_couchbase_config(): def build_temporal_config(): + """ + Build Temporal configuration from environment variables. + + Returns: + dict: Temporal configuration with host and namespace settings. + """ return { 'temporal_host': getenv('TEMPORAL_HOST', 'localhost:7233'), 'temporal_namespace': getenv('TEMPORAL_NAMESPACE', 'default'), @@ -41,6 +65,12 @@ def build_temporal_config(): def build_postgres_config(): + """ + Build PostgreSQL configuration from environment variables. + + Returns: + dict: PostgreSQL configuration with connection details and connection pool settings. + """ return { 'host': getenv('POSTGRES_HOST', 'localhost'), 'port': int(getenv('POSTGRES_PORT', '5432')), @@ -53,6 +83,12 @@ def build_postgres_config(): def build_email_config(): + """ + Build email configuration from environment variables. + + Returns: + dict: Email configuration with SMTP server settings and sender credentials. + """ return { 'sender_email': getenv('EMAIL_SENDER', 'sientia-alerts@aignosi.com'), 'sender_password': getenv('EMAIL_SENDER_PASSWORD', 'sientia'), diff --git a/orchestrator/utils/email_builder.py b/orchestrator/utils/email_builder.py index 88a765f..64b5963 100644 --- a/orchestrator/utils/email_builder.py +++ b/orchestrator/utils/email_builder.py @@ -18,12 +18,33 @@ class EmailBuilder: self.general_template = file.read() def replace_parameters(self, template: str, parameters: dict) -> str: + """ + Replace parameters in a Jinja2 template with provided values. + + Args: + template (str): The Jinja2 template string. + parameters (dict): Dictionary of parameters to replace in the template. + + Returns: + str: The rendered template with parameters replaced. + """ # Criar um template Jinja2 template = Template(template) return template.render(parameters) def parameters(self, general_events: dict, mail_type: str) -> dict: + """ + Build parameters dictionary for email templates based on general events and mail type. + + Args: + general_events (dict): Dictionary containing events categorized by level (ERROR, WARNING, INFO). + Each level contains a 'models' key with model-specific event data. + mail_type (str): The type of email being sent. + + Returns: + dict: Dictionary with mail_type and rendered event sections for each notification level. + """ error_models = general_events.get('ERROR', {}).get('models', []) warning_models = general_events.get('WARNING', {}).get('models', []) info_models = general_events.get('INFO', {}).get('models', []) @@ -43,7 +64,16 @@ class EmailBuilder: def build_email(self, report_data: list[dict], mail_type: str) -> str: """ - Builds the email html. + Builds the email HTML by organizing report data by notification level and model. + + Args: + report_data (list[dict]): List of notification reports, each containing: + - level (str): Notification level (ERROR, WARNING, INFO) + - model_name (str): Name of the model + - Additional notification details + + Returns: + str: Complete HTML email content ready for sending. """ general_events = {} diff --git a/orchestrator/utils/orchestrator_functions.py b/orchestrator/utils/orchestrator_functions.py index 7b6289f..38bde51 100644 --- a/orchestrator/utils/orchestrator_functions.py +++ b/orchestrator/utils/orchestrator_functions.py @@ -2,6 +2,21 @@ from typing import Any def common_config(config: dict[str, Any]): + """ + Extract common configuration parameters from a pipeline configuration. + + Args: + config (dict[str, Any]): Pipeline configuration containing: + - workflow_type (str): Type of workflow (e.g., 'scouter', 'predictions_batch') + - schedule_name (str): Name of the schedule + - frequency (str, optional): Frequency of execution (default: '1m') + - max_retry_policy (int, optional): Maximum retry attempts (default: 1) + - model_id (str): ID of the model + - models (dict): Model configuration containing 'name' field + + Returns: + dict[str, Any]: Common configuration dictionary with extracted parameters. + """ return { "workflow_type": config['workflow_type'], "schedule_name": config['schedule_name'], @@ -14,6 +29,18 @@ def common_config(config: dict[str, Any]): def minimal_retrain(config: dict[str, Any]): + """ + Build minimal retrain configuration from pipeline config. + + Args: + config (dict[str, Any]): Pipeline configuration containing: + - schedule_name (str): Name of the schedule + - query (str): SQL query for retraining + - Additional fields from common_config + + Returns: + dict[str, Any]: Minimal retrain configuration with workflow type set to 'minimal_retrain'. + """ return { **common_config(config), "workflow_type": "minimal_retrain", @@ -25,6 +52,25 @@ def minimal_retrain(config: dict[str, Any]): def scouter(config: dict[str, Any]): + """ + Build scouter configuration from pipeline config. + + Args: + config (dict[str, Any]): Pipeline configuration containing: + - filters (list[dict], optional): List of filter configurations + - read_tags (list[dict]): List of tag configurations with: + - filter_name (str): Name of the filter + - policy (str): Filter policy + - tag_name (str): Name of the tag + - aggr_func (str, optional): Aggregation function (default: 'lts') + - data_range (list[int], optional): Data range limits (default: [-100, 100]) + - tag_retention_minutes (int, optional): Tag retention time in minutes (default: 60) + - debug_data_package (bool, optional): Enable debug data package (default: False) + - Additional fields from common_config + + Returns: + dict[str, Any]: Scouter configuration with topic, filters, tags, and retention settings. + """ filters = {} for f in config.get('filters', []): filters[f['filter_name']] = { @@ -53,6 +99,19 @@ def scouter(config: dict[str, Any]): def overlap_filter_config(base_filter_config: dict[str, Any], config: list[dict[str, Any]]): + """ + Overlap filter configuration with base filter config. + + Args: + base_filter_config (dict[str, Any]): Base filter configuration to extend. + config (list[dict[str, Any]]): List of filter configurations to add, each containing: + - filter_name (str): Name of the filter + - policy (str): Filter policy + - config (dict, optional): Additional filter configuration + + Returns: + dict[str, Any]: Extended filter configuration with new filters added. + """ for fil in config: base_filter_config[fil['filter_name']] = { "policy": fil['policy'], @@ -63,6 +122,15 @@ def overlap_filter_config(base_filter_config: dict[str, Any], config: list[dict[ def process_path_priority(path_priority: list[str]): + """ + Process and normalize path priority list to ensure it contains the required priorities. + + Args: + path_priority (list[str]): List of path priorities to process. + + Returns: + list[str]: Normalized path priority list with exactly 3 elements: ["STOP", "CONTINUE", "REPEAT"]. + """ for priority in path_priority[:]: if priority not in ["STOP", "CONTINUE", "REPEAT"]: path_priority.remove(priority) @@ -75,6 +143,26 @@ def process_path_priority(path_priority: list[str]): def predictions_batch(config: dict[str, Any]): + """ + Build predictions batch configuration from pipeline config. + + Args: + config (dict[str, Any]): Pipeline configuration containing: + - write_tags (list[dict]): List of tag configurations with: + - server_id (str): ID of the OPC server + - type (str): Tag type ('prediction' or 'confidence') + - addr (str): Tag address + - data_type (str, optional): Data type (default: 'float') + - path_priority (list[str], optional): List of path priorities (default: ["STOP", "CONTINUE", "REPEAT"]) + - input_filters (list[dict], optional): List of input filter configurations + - mlflow_transform_filters (list[dict], optional): List of MLflow transform filter configurations + - mlflow_predict_filters (list[dict], optional): List of MLflow predict filter configurations + - model_retention_minutes (int, optional): Model retention time in minutes (default: 60) + - Additional fields from common_config + + Returns: + dict[str, Any]: Predictions batch configuration with OPC output config, filters, and path priority. + """ tags = {} for tag in config.get('write_tags', []): if tag['server_id'] not in tags: @@ -156,6 +244,23 @@ def gather_read_tags(pipelines: list[dict[str, Any]]) -> dict[str, Any]: def build_tag_config(tag: dict[str, Any], slot_config: dict[str, Any], opc_servers: dict[str, Any], i: int): + """ + Build tag configuration for a specific slot and OPC server. + + Args: + tag (dict[str, Any]): Tag configuration containing: + - server_id (str): ID of the OPC server + - tag_address (str): Address of the tag + slot_config (dict[str, Any]): Current slot configuration to update. + opc_servers (dict[str, Any]): Dictionary of OPC server configurations. + i (int): Slot number to configure. + + Returns: + dict[str, Any]: Updated slot configuration with the new tag. + + Raises: + ValueError: If the specified server_id is not found in opc_servers. + """ server_id = tag['server_id'] if server_id not in opc_servers: diff --git a/orchestrator/worker/worker.py b/orchestrator/worker/worker.py index 89959e3..49fc31e 100644 --- a/orchestrator/worker/worker.py +++ b/orchestrator/worker/worker.py @@ -29,6 +29,13 @@ POD_ID = os.getenv("POD_ID") async def main(): + """ + Main function to initialize and run the Temporal worker. + + Sets up MongoDB connection, notification handler, Temporal client, and starts + multiple workers for different task queues (orchestrator, alerts, reports). + Handles graceful shutdown and error handling. + """ host = os.getenv('TEMPORAL_HOST', 'localhost:7233') namespace = os.getenv('TEMPORAL_NAMESPACE', 'default') logger = get_logger(__name__) @@ -176,6 +183,12 @@ async def main(): def start_prometheus_server(): + """ + Start the Prometheus metrics server on the configured port. + + Sets up HTTP server for metrics collection and marks the application as UP. + Exits the application if the server fails to start. + """ try: port = int(os.getenv("HTTP_METRICS_PORT", 9090)) start_http_server(port)