diff --git a/orchestrator/activities/activities.py b/orchestrator/activities/activities.py index 5049e3f..02437c2 100644 --- a/orchestrator/activities/activities.py +++ b/orchestrator/activities/activities.py @@ -17,7 +17,7 @@ class Activities( # Couchbase, TemporalManager, SlotManager, Formatters, MongoDB): def __init__(self, - temporal_client: Client, + temporal_config: dict[str, Any], # couchbase_config: dict[str, Any], redis_config: dict[str, Any], mongodb_config: dict[str, Any], @@ -32,7 +32,9 @@ class Activities( # Couchbase, # notification_handler=notification_handler) TemporalManager.__init__(self, - temporal_client=temporal_client, + host=temporal_config['temporal_host'], + scouter_namespace=temporal_config['temporal_scouter_namespace'], + laborious_namespace=temporal_config['temporal_laborious_namespace'], logger=logger, notification_handler=notification_handler) @@ -45,6 +47,8 @@ class Activities( # Couchbase, notification_handler=notification_handler) Formatters.__init__(self, + scouter_namespace=temporal_config['temporal_scouter_namespace'], + laborious_namespace=temporal_config['temporal_laborious_namespace'], logger=logger, notification_handler=notification_handler) @@ -54,9 +58,5 @@ class Activities( # Couchbase, 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) - def shutdown(self): MongoDB.shutdown(self) diff --git a/orchestrator/activities/formatters.py b/orchestrator/activities/formatters.py index 8a06865..376ec44 100644 --- a/orchestrator/activities/formatters.py +++ b/orchestrator/activities/formatters.py @@ -14,7 +14,12 @@ with workflow.unsafe.imports_passed_through(): class Formatters(BaseActivity): - def __init__(self, logger: Logger, notification_handler: NotificationHandler): + def __init__(self, + scouter_namespace: str, + laborious_namespace: str, + logger: Logger, notification_handler: NotificationHandler): + self.scouter_namespace = scouter_namespace + self.laborious_namespace = laborious_namespace BaseActivity.__init__(self, logger=logger, notification_handler=notification_handler) @@ -37,14 +42,18 @@ class Formatters(BaseActivity): pipelines = input_data['pipelines'] - schedule_config = {} + schedule_config = { + self.scouter_namespace: {}, + self.laborious_namespace: {} + } for pipeline in pipelines: if pipeline['workflow_type'] == 'scouter': - schedule_config[pipeline['schedule_name']] = scouter(pipeline) + schedule_config[self.scouter_namespace][pipeline['schedule_name']] = scouter( + pipeline) elif pipeline['workflow_type'] == 'predictions_batch': - schedule_config[pipeline['schedule_name'] - ] = predictions_batch(pipeline) + schedule_config[self.laborious_namespace][pipeline['schedule_name'] + ] = predictions_batch(pipeline) self.logger.info("Processed schedules") @@ -110,6 +119,31 @@ class Formatters(BaseActivity): return slot_config + def compare_configs(self, + schedules: dict[str, Any], current_schedules: dict[str, Any], + to_update: dict[str, Any], to_create: dict[str, Any], + namespace: str): + """ + Compares the schedules and current schedules, and updates the + to_update and to_create dictionaries. + """ + + for schedule_name, schedule in schedules.items(): + if schedule_name in current_schedules: + old_config = current_schedules[schedule_name]['data'] + + self.logger.debug(f"Comparing {schedule_name}:") + self.logger.debug(json.dumps( + old_config, indent=4, sort_keys=True)) + self.logger.debug(json.dumps( + schedule, indent=4, sort_keys=True)) + + if schedule != old_config: + to_update[namespace][schedule_name] = schedule + + elif schedule_name not in current_schedules: + to_create[namespace][schedule_name] = schedule + @activity.defn(name="create_schedule_config") async def create_schedule_config(self, input_data: dict[str, Any]) -> dict[str, Any]: @@ -134,30 +168,28 @@ class Formatters(BaseActivity): current_schedule_config = input_data['current_schedule_config'] schedule_config = input_data['schedule_config'] - to_update = {} - to_create = {} - to_delete = [] + to_update = { + self.scouter_namespace: {}, + self.laborious_namespace: {} + } + to_create = { + self.scouter_namespace: {}, + self.laborious_namespace: {} + } + to_delete = { + self.scouter_namespace: [], + self.laborious_namespace: [] + } - for schedule_name, schedule in schedule_config.items(): - if schedule_name in current_schedule_config: - self.logger.debug(f"{current_schedule_config[schedule_name]}") - old_config = current_schedule_config[schedule_name]['data'] + for namespace, schedules in schedule_config.items(): + current_schedules = current_schedule_config.get(namespace, {}) + self.compare_configs( + schedules, current_schedules, to_update, to_create, namespace) - self.logger.debug(f"Comparing {schedule_name}:") - self.logger.debug(json.dumps( - old_config, indent=4, sort_keys=True)) - self.logger.debug(json.dumps( - schedule, indent=4, sort_keys=True)) - - if schedule != old_config: - 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) + for namespace, schedules in current_schedule_config.items(): + for schedule_name in schedules: + if schedule_name not in schedule_config[namespace]: + to_delete[namespace].append(schedule_name) output = { "to_update": to_update, diff --git a/orchestrator/activities/temporal_manager.py b/orchestrator/activities/temporal_manager.py index fc4c5a1..c07b621 100644 --- a/orchestrator/activities/temporal_manager.py +++ b/orchestrator/activities/temporal_manager.py @@ -1,4 +1,3 @@ -from asyncio import sleep from temporalio import activity, workflow from temporalio.client import ( Client, Schedule, ScheduleActionStartWorkflow, ScheduleIntervalSpec, ScheduleSpec, ScheduleUpdate, ScheduleUpdateInput) @@ -13,16 +12,23 @@ with workflow.unsafe.imports_passed_through(): import base64 from datetime import timedelta import json + from asyncio import sleep from orchestrator.utils.converters import parse_frequency class TemporalManager(BaseActivity): - def __init__(self, temporal_client: Client, logger: Logger, - notification_handler: NotificationHandler): + def __init__(self, host: str, scouter_namespace: str, laborious_namespace: str, + logger: Logger, notification_handler: NotificationHandler): - self.temporal_client = temporal_client + self.temporal_host = host + self.scouter_namespace = scouter_namespace + self.laborious_namespace = laborious_namespace + self.temporal_clients = {} - self.schedule_handles = {} + self.schedule_handles = { + self.scouter_namespace: {}, + self.laborious_namespace: {} + } self.model_id_id_key = SearchAttributeKey.for_keyword("model_id") self.model_name_id_key = SearchAttributeKey.for_keyword("model_name") @@ -33,6 +39,28 @@ class TemporalManager(BaseActivity): logger=logger, notification_handler=notification_handler) + async def connect_to_temporal(self): + self.logger.info( + f"Connecting to Temporal side namespaces at {self.temporal_host}") + self.logger.info(f"Scouter namespace: {self.scouter_namespace}") + + scouter_client = await Client.connect( + target_host=self.temporal_host, + namespace=self.scouter_namespace + ) + + self.logger.info(f"Laborious namespace: {self.laborious_namespace}") + + laborious_client = await Client.connect( + target_host=self.temporal_host, + namespace=self.laborious_namespace + ) + + self.temporal_clients = { + self.scouter_namespace: scouter_client, + self.laborious_namespace: laborious_client + } + @activity.defn(name="load_schedule") async def load_schedule(self) -> dict[str, Any]: """ @@ -48,46 +76,57 @@ class TemporalManager(BaseActivity): orchestrated_schedules = {} - async for schedule in await self.temporal_client.list_schedules(): - search_attrs = getattr(schedule, "search_attributes", {}) - if search_attrs.get("orchestrated", ["false"]) == ["true"]: - schedule_id = schedule.id + for namespace, client in self.temporal_clients.items(): + orchestrated_schedules[namespace] = {} - self.logger.debug("Schedule id: %s", schedule_id) + self.logger.info( + "Getting orchestrated schedules for %s", namespace) - handle = self.temporal_client.get_schedule_handle(schedule_id) + async for schedule in await client.list_schedules(): + search_attrs = getattr(schedule, "search_attributes", {}) + if search_attrs.get("orchestrated", ["false"]) == ["true"]: + schedule_id = schedule.id - self.logger.debug("Handle acquired") + self.logger.debug("Schedule id: %s", schedule_id) - self.schedule_handles[schedule_id] = handle + handle = client.get_schedule_handle( + schedule_id) - self.logger.debug("Describing schedule...") + self.logger.debug("Handle acquired") - desc = await handle.describe( - rpc_timeout=timedelta(seconds=60) - ) + self.schedule_handles[namespace][schedule_id] = handle - self.logger.debug("Parsing args...") + self.logger.debug("Describing schedule...") - for arg in desc.schedule.action.args: - data = MessageToDict(arg)['data'] - data = base64.b64decode(data).decode('utf-8') + desc = await handle.describe( + rpc_timeout=timedelta(seconds=60) + ) - frequency = desc.schedule.spec.intervals[0].every.seconds + self.logger.debug("Parsing args...") - orchestrated_schedules[schedule_id] = { - 'frequency': frequency, - 'data': json.loads(data), - } + for arg in desc.schedule.action.args: + data = MessageToDict(arg)['data'] + data = base64.b64decode(data).decode('utf-8') - await sleep(0.1) + frequency = desc.schedule.spec.intervals[0].every.seconds + + orchestrated_schedules[namespace][schedule_id] = { + 'frequency': frequency, + 'data': json.loads(data), + } + + await sleep(0.1) self.logger.info("Found %d orchestrated schedules", - len(orchestrated_schedules)) + len(orchestrated_schedules[self.scouter_namespace]) + + len(orchestrated_schedules[self.laborious_namespace])) self.logger.debug("Orchestrated schedules: %s", orchestrated_schedules) + self.logger.debug("Schedule handles: %s", + self.schedule_handles) + return orchestrated_schedules @activity.defn(name="create_schedules") @@ -104,66 +143,74 @@ class TemporalManager(BaseActivity): - dict[str, Any]: A report of the created schedules. """ - schedules = input_data['schedules'] + schedules_to_create = 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'] + for namespace, schedules in schedules_to_create.items(): + client = self.temporal_clients.get(namespace) - try: - self.logger.debug(f"Creating schedule {schedule_name}:") - self.logger.debug(json.dumps( - schedule, indent=4, sort_keys=True)) - await self.temporal_client.create_schedule( - schedule_name, - Schedule( - action=ScheduleActionStartWorkflow( - workflow_type, - schedule, - id=schedule_name, - task_queue=f"{workflow_type}-queue", - execution_timeout=timedelta(minutes=2) - ), - spec=ScheduleSpec( - intervals=[ - ScheduleIntervalSpec( - every=timedelta(seconds=parse_frequency( - schedule.get('frequency', '1m'))) - ) - ] - ) + if not client: + raise ValueError( + f"Temporal client for {namespace} not found, clients: {self.temporal_clients}") + + for schedule_name, schedule in schedules.items(): + search_attributes = TypedSearchAttributes([ + SearchAttributePair( + key=self.model_id_id_key, + value=schedule['model_id'] ), - search_attributes=search_attributes - ) + SearchAttributePair( + key=self.model_name_id_key, + value=schedule['model_name'] + ), + SearchAttributePair( + key=self.orchestrated_id_key, + value="true" + ) + ]) + workflow_type = schedule['workflow_type'] - 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) - } + try: + self.logger.debug(f"Creating schedule {schedule_name}:") + self.logger.debug(json.dumps( + schedule, indent=4, sort_keys=True)) - self.logger.info(f"Processed {len(schedules)} schedules") + await client.create_schedule( + schedule_name, + Schedule( + action=ScheduleActionStartWorkflow( + workflow_type, + schedule, + id=schedule_name, + task_queue=f"{workflow_type}-queue", + execution_timeout=timedelta(minutes=2) + ), + 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_to_create)} schedules") self.logger.debug("\n %s", json.dumps(report, indent=4, sort_keys=True)) @@ -184,55 +231,64 @@ class TemporalManager(BaseActivity): - dict[str, Any]: A report of the updated schedules. """ - schedules = input_data['schedules'] + schedules_to_update = input_data['schedules'] report = {} - for schedule_name, schedule in schedules.items(): - try: - handler = self.schedule_handles.get(schedule_name) + for namespace, schedules in schedules_to_update.items(): + schedule_handles = self.schedule_handles.get(namespace, None) - if not handler: - raise ValueError(f"Schedule {schedule_name} not found") + if schedule_handles is None: + raise ValueError( + f"Schedule handles for {namespace} not found, handles: {self.schedule_handles}") - async def update_schedule(input_data: ScheduleUpdateInput) -> ScheduleUpdate: - schedule_action = input_data.description.schedule.action + for schedule_name, schedule in schedules.items(): + try: + handler = schedule_handles.get(schedule_name) - self.logger.debug("Updating schedule:") + if not handler: + raise ValueError(f"Schedule {schedule_name} not found") - if hasattr(schedule_action, "args"): - self.logger.debug("New schedule:") - self.logger.debug(json.dumps( - schedule, indent=4, sort_keys=True)) + # fmt: off + async def update_schedule(input_data: ScheduleUpdateInput) -> ScheduleUpdate: # NOSONAR + schedule_action = input_data.description.schedule.action - schedule_action.args = [schedule] + self.logger.debug("Updating schedule:") - input_data.description.schedule.spec.intervals = [ - ScheduleIntervalSpec( - every=timedelta(seconds=parse_frequency( - schedule.get('frequency', '1m'))) - ) - ] + if hasattr(schedule_action, "args"): + self.logger.debug("New schedule:") + self.logger.debug(json.dumps( + schedule, indent=4, sort_keys=True)) # NOSONAR - return ScheduleUpdate(schedule=input_data.description.schedule) + schedule_action.args = [schedule] - await handler.update(update_schedule) + input_data.description.schedule.spec.intervals = [ + ScheduleIntervalSpec( + every=timedelta(seconds=parse_frequency( + schedule.get('frequency', '1m'))) + ) + ] - del update_schedule + return ScheduleUpdate(schedule=input_data.description.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) - } + # fmt: on + await handler.update(update_schedule) - self.logger.info(f"Processed {len(schedules)} schedules") + 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_to_update)} schedules") self.logger.debug("\n %s", json.dumps(report, indent=4, sort_keys=True)) @@ -253,34 +309,41 @@ class TemporalManager(BaseActivity): - dict[str, Any]: A report of the deleted schedules. """ - schedules = input_data['schedules'] + schedules_to_delete = input_data['schedules'] report = {} - for schedule_name in schedules: - try: - handler = self.schedule_handles.get(schedule_name) + for namespace, schedules in schedules_to_delete.items(): + schedule_handles = self.schedule_handles.get(namespace, None) - if not handler: - raise ValueError(f"Schedule {schedule_name} not found") + if schedule_handles is None: + raise ValueError( + f"Schedule handles for {namespace} not found, handles: {self.schedule_handles}") - await handler.delete() + for schedule_name in schedules: + try: + handler = schedule_handles.get(schedule_name) - del self.schedule_handles[schedule_name] + if not handler: + raise ValueError(f"Schedule {schedule_name} not found") - 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) - } + await handler.delete() - self.logger.info(f"Processed {len(schedules)} schedules") + del self.schedule_handles[namespace][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_to_delete)} schedules") self.logger.debug("\n %s", json.dumps(report, indent=4, sort_keys=True)) diff --git a/orchestrator/utils/connectors_config.py b/orchestrator/utils/connectors_config.py index 3808f72..3b7551a 100644 --- a/orchestrator/utils/connectors_config.py +++ b/orchestrator/utils/connectors_config.py @@ -28,3 +28,12 @@ def build_couchbase_config(): 'username': getenv('COUCHBASE_USERNAME', 'sientia'), 'password': getenv('COUCHBASE_PASSWORD', 'sientia') } + + +def build_temporal_config(): + return { + 'temporal_host': getenv('TEMPORAL_HOST', 'localhost:7233'), + 'temporal_namespace': getenv('TEMPORAL_NAMESPACE', 'default'), + 'temporal_scouter_namespace': getenv('TEMPORAL_SCOUTER_NAMESPACE', 'scouter'), + 'temporal_laborious_namespace': getenv('TEMPORAL_LABORIOUS_NAMESPACE', 'laborious') + } diff --git a/orchestrator/worker/worker.py b/orchestrator/worker/worker.py index 8aa019f..76c88f1 100644 --- a/orchestrator/worker/worker.py +++ b/orchestrator/worker/worker.py @@ -1,6 +1,8 @@ from temporalio import workflow, client from temporalio.worker import Worker +from orchestrator.utils.connectors_config import build_temporal_config + with workflow.unsafe.imports_passed_through(): import os import sys @@ -40,13 +42,15 @@ async def main(): logger.info('Starting Activities...') activities = Activities( - temporal_client=temporal_client, + temporal_config=build_temporal_config(), redis_config=build_redis_config(), mongodb_config=build_mongodb_config(), logger=logger, notification_handler=notification_handler ) + await activities.connect_to_temporal() + logger.info('Starting Workers...') workers = [ @@ -55,8 +59,6 @@ async def main(): task_queue='orchestrator-queue', workflows=[Orchestrator], activities=[ - # Base - activities.prepare_activity, # Redis activities.load_active_ingestors, activities.load_opc_slots, diff --git a/test.ipynb b/test.ipynb index 7a97f4b..1fcad1d 100644 --- a/test.ipynb +++ b/test.ipynb @@ -136,7 +136,7 @@ }, { "cell_type": "code", - "execution_count": 4, + "execution_count": null, "id": "bb750ae6", "metadata": {}, "outputs": [ @@ -163,9 +163,9 @@ ")\n", "from temporalio.common import TypedSearchAttributes, SearchAttributeKey, SearchAttributePair\n", "\n", - "# temporal operator search-attribute create --namespace default --name model_id --type Text\n", - "# temporal operator search-attribute create --namespace default --name orchestrated --type Text\n", - "# temporal operator search-attribute create --namespace default --name model_name --type Text\n", + "# temporal operator search-attribute create --namespace scouter --name model_id --type Text && temporal operator search-attribute create --namespace scouter --name orchestrated --type Text && temporal operator search-attribute create --namespace scouter --name model_name --type Text && temporal operator search-attribute create --namespace laborious --name model_id --type Text && temporal operator search-attribute create --namespace laborious --name orchestrated --type Text && temporal operator search-attribute create --namespace laborious --name model_name --type Text\n", + "\n", + "\n", "\n", "await temporal_client.create_schedule(\n", " \"orchestrator\",\n", diff --git a/tests/orchestrator/activities/test_activities.py b/tests/orchestrator/activities/test_activities.py index 557ec8f..353382a 100644 --- a/tests/orchestrator/activities/test_activities.py +++ b/tests/orchestrator/activities/test_activities.py @@ -28,12 +28,17 @@ def test___init__(mock_formatters_init, mock_slot_manager_init, 'password': 'password' } - temporal_client = MagicMock() + temporal_config = { + 'temporal_host': 'localhost', + 'temporal_scouter_namespace': 'scouter', + 'temporal_laborious_namespace': 'laborious' + } + logger = MagicMock() notification_handler = MagicMock() activities = Activities( - temporal_client=temporal_client, + temporal_config=temporal_config, redis_config=redis_config, mongodb_config=mongo_db_config, logger=logger, @@ -66,13 +71,17 @@ def test___init__(mock_formatters_init, mock_slot_manager_init, mock_temporal_manager_init.assert_called_once_with( ANY, - temporal_client=temporal_client, + host='localhost', + scouter_namespace='scouter', + laborious_namespace='laborious', logger=logger, notification_handler=notification_handler ) mock_formatters_init.assert_called_once_with( ANY, + scouter_namespace='scouter', + laborious_namespace='laborious', logger=logger, notification_handler=notification_handler ) diff --git a/tests/orchestrator/activities/test_formatters.py b/tests/orchestrator/activities/test_formatters.py index 73f29a7..db63ba5 100644 --- a/tests/orchestrator/activities/test_formatters.py +++ b/tests/orchestrator/activities/test_formatters.py @@ -9,6 +9,8 @@ from orchestrator.utils.orchestrator_functions import build_tag_config @fixture def formatters(): return Formatters( + scouter_namespace="scouter", + laborious_namespace="laborious", logger=MagicMock(), notification_handler=MagicMock() ) @@ -40,12 +42,18 @@ async def test_process_schedules(mock_predictions_batch, mock_scouter, formatter result = await formatters.process_schedules(input_data) assert result == { - "test_schedule_name": "test_scouter", - "test_schedule_name2": "test_predictions_batch" + "scouter": { + "test_schedule_name": "test_scouter" + }, + "laborious": { + "test_schedule_name2": "test_predictions_batch" + } } - mock_scouter.assert_called_once_with(input_data['pipelines'][0]) - mock_predictions_batch.assert_called_once_with(input_data['pipelines'][1]) + mock_scouter.assert_called_once_with( + input_data['pipelines'][0]) + mock_predictions_batch.assert_called_once_with( + input_data['pipelines'][1]) @mark.asyncio @@ -243,23 +251,29 @@ async def test_process_slots(mock_build_tag_config, mock_gather_read_tags, forma async def test_create_schedule_config(formatters): input_data = { "current_schedule_config": { - "test_schedule_name_to_delete": { - "frequency": 60, - "data": {"test": "test"} - }, - "test_schedule_name_to_update": { - "frequency": 60, - "data": {"test": "test"} + "scouter": { + "test_schedule_name_to_delete": { + "frequency": 60, + "data": {"test": "test"} + }, + "test_schedule_name_to_update": { + "frequency": 60, + "data": {"test": "test"} + } } }, "schedule_config": { - "test_schedule_name_to_create": { - "frequency": 60, - "data": {"test": "test"} + "laborious": { + "test_schedule_name_to_create": { + "frequency": 60, + "data": {"test": "test"} + } }, - "test_schedule_name_to_update": { - "frequency": 60, - "data": {"test": "test2"} + "scouter": { + "test_schedule_name_to_update": { + "frequency": 60, + "data": {"test": "test2"} + } } } } @@ -268,20 +282,29 @@ async def test_create_schedule_config(formatters): assert result == { "to_create": { - "test_schedule_name_to_create": { - "frequency": 60, - "data": {"test": "test"} - } + "laborious": { + "test_schedule_name_to_create": { + "frequency": 60, + "data": {"test": "test"} + } + }, + "scouter": {} }, "to_update": { - "test_schedule_name_to_update": { - "frequency": 60, - "data": {"test": "test2"} - } + "scouter": { + "test_schedule_name_to_update": { + "frequency": 60, + "data": {"test": "test2"} + } + }, + "laborious": {} }, - "to_delete": [ - "test_schedule_name_to_delete" - ] + "to_delete": { + "scouter": [ + "test_schedule_name_to_delete" + ], + "laborious": [] + } } diff --git a/tests/orchestrator/activities/test_temporal_manager.py b/tests/orchestrator/activities/test_temporal_manager.py index 7e504ef..5410052 100644 --- a/tests/orchestrator/activities/test_temporal_manager.py +++ b/tests/orchestrator/activities/test_temporal_manager.py @@ -3,18 +3,43 @@ from datetime import timedelta import base64 import json from pytest import fixture, mark +import pytest_asyncio from orchestrator.activities.temporal_manager import TemporalManager from orchestrator.utils.converters import parse_frequency @fixture -def temporal_manager(): - return TemporalManager( - temporal_client=MagicMock(), +@patch("orchestrator.activities.temporal_manager.Client.connect") +def temporal_manager(connect_mock): + temporal_manager = TemporalManager( + host='localhost:7233', + scouter_namespace='scouter', + laborious_namespace='laborious', logger=MagicMock(), notification_handler=MagicMock() ) + temporal_manager.temporal_clients['scouter'] = MagicMock() + temporal_manager.temporal_clients['laborious'] = MagicMock() + + return temporal_manager + + +@mark.asyncio +@patch("orchestrator.activities.temporal_manager.Client.connect", new_callable=AsyncMock) +async def test_connect_to_temporal(connect_mock, temporal_manager): + await temporal_manager.connect_to_temporal() + connect_mock.assert_has_calls([ + call( + target_host='localhost:7233', + namespace='scouter' + ), + call( + target_host='localhost:7233', + namespace='laborious' + ) + ]) + @mark.asyncio @patch("orchestrator.activities.temporal_manager.MessageToDict", @@ -41,9 +66,7 @@ async def test_load_schedule(_mock_message_to_dict, temporal_manager): } ) - temporal_manager.temporal_client.list_schedules = AsyncMock( - return_value=async_iter()) - temporal_manager.temporal_client.get_schedule_handle.return_value = MagicMock( + handle = MagicMock( describe=AsyncMock( return_value=MagicMock( schedule=MagicMock( @@ -59,24 +82,54 @@ async def test_load_schedule(_mock_message_to_dict, temporal_manager): ) ) ) - temporal_manager.temporal_client.get_schedule_handle.return_value.describe \ - .return_value.schedule.spec = MagicMock( - intervals=[ - MagicMock( - every=MagicMock( - seconds=60 - ) + + temporal_manager.temporal_clients['scouter'].list_schedules = AsyncMock( + return_value=async_iter() + ) + temporal_manager.temporal_clients['laborious'].list_schedules = AsyncMock( + return_value=async_iter() + ) + + temporal_manager.temporal_clients['scouter'].get_schedule_handle = MagicMock( + return_value=handle + ) + temporal_manager.temporal_clients['laborious'].get_schedule_handle = MagicMock( + return_value=handle + ) + + describe_mock = MagicMock( + intervals=[ + MagicMock( + every=MagicMock( + seconds=60 ) - ] - ) + ) + ] + ) + temporal_manager.temporal_clients['scouter'].get_schedule_handle.return_value.describe \ + .return_value.schedule.spec = describe_mock + temporal_manager.temporal_clients['laborious'].get_schedule_handle.return_value.describe \ + .return_value.schedule.spec = describe_mock response = await temporal_manager.load_schedule() - temporal_manager.temporal_client.list_schedules.assert_called_once() + temporal_manager.temporal_clients['scouter'].list_schedules.assert_awaited_once( + ) + temporal_manager.temporal_clients['laborious'].list_schedules.assert_awaited_once( + ) + assert response == { - "test-schedule-id": { - "frequency": 60, - "data": {"test": "test"} + "scouter": { + "test-schedule-id": { + "frequency": 60, + "data": {"test": "test"} + } + }, + "laborious": { + "test-schedule-id": { + "frequency": 60, + "data": {"test": "test"} + } } } @@ -102,68 +155,112 @@ async def test_create_schedule( input_data = { "schedules": { - "test-schedule": { - "model_id": 1, - "model_name": "test-model-name", - "workflow_type": "test-workflow", - "frequency": "1m", - "data": {"test": "test"} + "scouter": { + "test-schedule": { + "model_id": 1, + "model_name": "test-model-name", + "workflow_type": "test-workflow", + "frequency": "1m", + "data": {"test": "test"} + }, + "test-schedule-invalid-frequency": { + "model_id": 2, + "model_name": "test-model-name", + "workflow_type": "test-workflow", + "frequency": "10y", + "data": {"test": "test"} + } }, - "test-schedule-invalid-frequency": { - "model_id": 2, - "model_name": "test-model-name", - "workflow_type": "test-workflow", - "frequency": "10y", - "data": {"test": "test"} + "laborious": { + "test-schedule-laborious": { + "model_id": 1, + "model_name": "test-model-name", + "workflow_type": "test-workflow", + "frequency": "2m", + "data": {"test": "test"} + } } } } - temporal_manager.temporal_client.create_schedule = AsyncMock() + temporal_manager.temporal_clients['scouter'].create_schedule = AsyncMock() + temporal_manager.temporal_clients['laborious'].create_schedule = AsyncMock( + ) report = await temporal_manager.create_schedules(input_data) - temporal_manager.temporal_client.create_schedule.assert_called_once_with( + temporal_manager.temporal_clients['scouter'].create_schedule.assert_called_once_with( "test-schedule", mock_schedule.return_value, search_attributes=mock_typed_search_attributes.return_value ) - - mock_schedule.assert_called_once_with( - action=mock_schedule_action_start_workflow.return_value, - spec=mock_schedule_spec.return_value + temporal_manager.temporal_clients['laborious'].create_schedule.assert_called_once_with( + "test-schedule-laborious", + mock_schedule.return_value, + search_attributes=mock_typed_search_attributes.return_value ) + mock_schedule.assert_has_calls([ + call( + action=mock_schedule_action_start_workflow.return_value, + spec=mock_schedule_spec.return_value + ), + call( + action=mock_schedule_action_start_workflow.return_value, + spec=mock_schedule_spec.return_value + ) + ]) + mock_schedule_action_start_workflow.assert_has_calls([ call( "test-workflow", - input_data['schedules']['test-schedule'], + input_data['schedules']['scouter']['test-schedule'], id="test-schedule", task_queue="test-workflow-queue", execution_timeout=ANY ), call( "test-workflow", - input_data['schedules']['test-schedule-invalid-frequency'], + input_data['schedules']['scouter']['test-schedule-invalid-frequency'], id="test-schedule-invalid-frequency", task_queue="test-workflow-queue", execution_timeout=ANY + ), + call( + "test-workflow", + input_data['schedules']['laborious']['test-schedule-laborious'], + id="test-schedule-laborious", + task_queue="test-workflow-queue", + execution_timeout=ANY ) ]) - mock_schedule_spec.assert_called_once_with( - intervals=[ - mock_schedule_interval_spec.return_value - ] - ) + mock_schedule_spec.assert_has_calls([ + call( + intervals=[ + mock_schedule_interval_spec.return_value + ] + ), + call( + intervals=[ + mock_schedule_interval_spec.return_value + ] + ) + ]) - mock_schedule_interval_spec.assert_called_once_with( - every=timedelta(seconds=60) - ) + mock_schedule_interval_spec.assert_has_calls([ + call( + every=timedelta(seconds=60) + ), + call( + every=timedelta(seconds=120) + ) + ]) mock_parse_frequency.assert_has_calls([ call("1m"), - call("10y") + call("10y"), + call("2m") ]) mock_typed_search_attributes.assert_has_calls([ @@ -172,6 +269,11 @@ async def test_create_schedule( mock_search_attribute_pair.return_value, mock_search_attribute_pair.return_value ]), + call([ + mock_search_attribute_pair.return_value, + mock_search_attribute_pair.return_value, + mock_search_attribute_pair.return_value + ]), call([ mock_search_attribute_pair.return_value, mock_search_attribute_pair.return_value, @@ -188,6 +290,18 @@ async def test_create_schedule( key=temporal_manager.model_name_id_key, value="test-model-name" ), + call( + key=temporal_manager.orchestrated_id_key, + value="true" + ), + call( + key=temporal_manager.model_id_id_key, + value=2 + ), + call( + key=temporal_manager.model_name_id_key, + value="test-model-name" + ), call( key=temporal_manager.orchestrated_id_key, value="true" @@ -202,10 +316,35 @@ async def test_create_schedule( "test-schedule-invalid-frequency": { "success": False, "message": "Invalid frequency" + }, + "test-schedule-laborious": { + "success": True, + "message": "Schedule created successfully" } } +@mark.asyncio +async def test_create_schedules_with_no_client(temporal_manager): + temporal_manager.temporal_clients = {} + input_data = { + "schedules": { + "abc": { + "test-schedule": { + "frequency": "1m", + "data": {"test": "test"} + } + } + } + } + + try: + await temporal_manager.create_schedules(input_data) + except Exception as e: + assert str( + e) == f"Temporal client for abc not found, clients: {temporal_manager.temporal_clients}" + + @mark.asyncio @patch("orchestrator.activities.temporal_manager.parse_frequency", side_effect=parse_frequency) @@ -218,30 +357,50 @@ async def test_update_schedules( args=MagicMock() ) temporal_manager.schedule_handles = { - "test-schedule": MagicMock( - update=AsyncMock( + "scouter": { + "test-schedule": MagicMock( update=AsyncMock( - side_effect=lambda f: f(input_mock) + update=AsyncMock( + side_effect=lambda f: f(input_mock) + ) ) ) - ) + }, + "laborious": { + "test-schedule-laborious": MagicMock( + update=AsyncMock( + update=AsyncMock( + side_effect=lambda f: f(input_mock) + ) + ) + ) + } } input_data = { "schedules": { - "test-schedule": { - "frequency": "1m", - "data": {"test": "test"} + "scouter": { + "test-schedule": { + "frequency": "1m", + "data": {"test": "test"} + }, + "test-schedule_no_handler": { + "frequency": "1m", + "data": {"test": "test"} + } }, - "test-schedule_no_handler": { - "frequency": "1m", - "data": {"test": "test"} + "laborious": { + "test-schedule-laborious": { + "frequency": "2m", + "data": {"test": "test"} + } } } } report = await temporal_manager.update_schedules(input_data) - temporal_manager.schedule_handles['test-schedule'].update.assert_called_once() + temporal_manager.schedule_handles['scouter']['test-schedule'].update.assert_called_once() + temporal_manager.schedule_handles['laborious']['test-schedule-laborious'].update.assert_called_once() assert report == { "test-schedule": { @@ -251,21 +410,59 @@ async def test_update_schedules( "test-schedule_no_handler": { "success": False, "message": "Schedule test-schedule_no_handler not found" + }, + "test-schedule-laborious": { + "success": True, + "message": "Schedule updated successfully" } } +@mark.asyncio +async def test_update_schedules_with_no_handle(temporal_manager): + temporal_manager.temporal_clients = {} + input_data = { + "schedules": { + "abc": { + "test-schedule": { + "frequency": "1m", + "data": {"test": "test"} + } + } + } + } + + try: + await temporal_manager.update_schedules(input_data) + except Exception as e: + assert str( + e) == f"Schedule handles for abc not found, handles: {temporal_manager.schedule_handles}" + + @mark.asyncio async def test_delete_schedules(temporal_manager): temporal_manager.schedule_handles = { - "test-schedule": MagicMock( - delete=AsyncMock() - ) + "scouter": { + "test-schedule": MagicMock( + delete=AsyncMock() + ) + }, + "laborious": { + "test-schedule-laborious": MagicMock( + delete=AsyncMock() + ) + } } + input_data = { - "schedules": [ - "test-schedule", "test-schedule_no_handler" - ] + "schedules": { + "scouter": [ + "test-schedule", "test-schedule_no_handler" + ], + "laborious": [ + "test-schedule-laborious" + ] + } } report = await temporal_manager.delete_schedules(input_data) @@ -278,5 +475,27 @@ async def test_delete_schedules(temporal_manager): "test-schedule_no_handler": { "success": False, "message": "Schedule test-schedule_no_handler not found" + }, + "test-schedule-laborious": { + "success": True, + "message": "Schedule deleted successfully" } } + + +@mark.asyncio +async def test_delete_schedules_with_no_handle(temporal_manager): + temporal_manager.schedule_handles = {} + input_data = { + "schedules": { + "abc": [ + "test-schedule" + ] + } + } + + try: + await temporal_manager.delete_schedules(input_data) + except Exception as e: + assert str( + e) == f"Schedule handles for abc not found, handles: {temporal_manager.schedule_handles}" diff --git a/tests/orchestrator/utils/test_connectors_config.py b/tests/orchestrator/utils/test_connectors_config.py index 06b1606..287bc17 100644 --- a/tests/orchestrator/utils/test_connectors_config.py +++ b/tests/orchestrator/utils/test_connectors_config.py @@ -1,7 +1,7 @@ from os import environ from orchestrator.utils.connectors_config import (build_redis_config, build_couchbase_config, - build_mongodb_config) + build_mongodb_config, build_temporal_config) def test_build_redis_config_with_env_vars(): @@ -74,3 +74,27 @@ def test_build_mongo_db_config_with_defaults(): 'connection_string': 'mongodb://sientia:sientia@localhost:27017', 'database_name': 'sientia' } + + +def test_build_temporal_config_with_env_vars(): + environ['TEMPORAL_HOST'] = 'localhost:7233' + environ['TEMPORAL_SCOUTER_NAMESPACE'] = 'scouter' + environ['TEMPORAL_LABORIOUS_NAMESPACE'] = 'laborious' + assert build_temporal_config() == { + 'temporal_host': 'localhost:7233', + 'temporal_namespace': 'default', + 'temporal_scouter_namespace': 'scouter', + 'temporal_laborious_namespace': 'laborious' + } + + +def test_build_temporal_config_with_defaults(): + environ.pop('TEMPORAL_HOST', None) + environ.pop('TEMPORAL_SCOUTER_NAMESPACE', None) + environ.pop('TEMPORAL_LABORIOUS_NAMESPACE', None) + assert build_temporal_config() == { + 'temporal_host': 'localhost:7233', + 'temporal_namespace': 'default', + 'temporal_scouter_namespace': 'scouter', + 'temporal_laborious_namespace': 'laborious' + } diff --git a/values.yaml b/values.yaml index 7bb0447..067b888 100644 --- a/values.yaml +++ b/values.yaml @@ -123,7 +123,7 @@ env: - name: GITHUB_REPO_URL value: "git@github.com:Aignosi/sientia-dataops-orchestrator_temporal.git" - name: GITHUB_BRANCH - value: "SIENTIAPDE-1141-criar-testes-de-carga" + value: "SIENTIAPDE-1148-separar-scouter-laborious-e-orchestrator-por-namespaces" - name: PYTHON_APP value: "orchestrator.worker.worker" @@ -171,6 +171,10 @@ env: value: "temporal-frontend.temporal.svc.cluster.local:7233" - name: TEMPORAL_NAMESPACE value: "default" + - name: TEMPORAL_SCOUTER_NAMESPACE + value: "scouter" + - name: TEMPORAL_LABORIOUS_NAMESPACE + value: "laborious" ssh: enabled: true