gitea_admin/sientia-dataops-orchestrator_temporal
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #6 from Aignosi/SIENTIAPDE-1148-separar-scouter-laborious-e-orchestrator-por-namespaces

Browse Source 
Sientiapde 1148 separar scouter laborious e orchestrator por namespaces
This commit is contained in:
Matheus Demoner
2025-07-14 10:35:47 -03:00
committed by GitHub
parent d84b721798 9821b052ae
commit e3a88eb1d5
11 changed files with 659 additions and 274 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
orchestrator/activities/activities.py
View File

@@ -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)

86
orchestrator/activities/formatters.py
View File

@@ -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,

335
orchestrator/activities/temporal_manager.py
View File

@@ -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))

9
orchestrator/utils/connectors_config.py
View File

@@ -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')
}

8
orchestrator/worker/worker.py
View File

@@ -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,

8
test.ipynb
View File

@@ -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",

15
tests/orchestrator/activities/test_activities.py
View File

@@ -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
)

79
tests/orchestrator/activities/test_formatters.py
View File

@@ -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": []
}
}

349
tests/orchestrator/activities/test_temporal_manager.py
View File

@@ -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}"

26
tests/orchestrator/utils/test_connectors_config.py
View File

@@ -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'
}

6
values.yaml
View File

@@ -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