Skip to content

Commit

Permalink
merge branch 35
Browse files Browse the repository at this point in the history
  • Loading branch information
MarkTNO committed Jun 21, 2024
2 parents 37d1a11 + ca4256a commit 335702b
Showing 1 changed file with 58 additions and 7 deletions.
65 changes: 58 additions & 7 deletions src/omotes_orchestrator/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,8 +22,16 @@
JobProgressUpdate,
JobCancel,
)
from omotes_sdk_protocol.work_flow_pb2 import RequestAvailableWorkflows

Check failure on line 25 in src/omotes_orchestrator/main.py

View workflow job for this annotation

GitHub Actions / Typecheck (3.11)

error: Cannot find implementation or library stub for module named "omotes_sdk_protocol.work_flow_pb2" [import]
from omotes_sdk.job import Job
from omotes_sdk.workflow_type import WorkflowTypeManager, WorkflowType
from omotes_sdk.workflow_type import (

Check failure on line 27 in src/omotes_orchestrator/main.py

View workflow job for this annotation

GitHub Actions / Typecheck (3.11)

error: Module "omotes_sdk.workflow_type" has no attribute "WorkflowParameter" [attr-defined]

Check failure on line 27 in src/omotes_orchestrator/main.py

View workflow job for this annotation

GitHub Actions / Typecheck (3.11)

error: Module "omotes_sdk.workflow_type" has no attribute "ParameterType" [attr-defined]

Check failure on line 27 in src/omotes_orchestrator/main.py

View workflow job for this annotation

GitHub Actions / Typecheck (3.11)

error: Module "omotes_sdk.workflow_type" has no attribute "ParameterSchema" [attr-defined]

Check failure on line 27 in src/omotes_orchestrator/main.py

View workflow job for this annotation

GitHub Actions / Typecheck (3.11)

error: Module "omotes_sdk.workflow_type" has no attribute "ParameterStringFormat" [attr-defined]
WorkflowTypeManager,
WorkflowType,
WorkflowParameter,
ParameterType,
ParameterSchema,
ParameterStringFormat,
)
from google.protobuf import json_format

from omotes_orchestrator.celery_interface import CeleryInterface
Expand Down Expand Up @@ -187,11 +195,7 @@ def start(self) -> None:
self._resume_init_barriers(self.postgresql_if.get_all_jobs())

self.celery_if.start()
self.omotes_if.start()
self.omotes_if.connect_to_job_submissions(
callback_on_new_job=self.new_job_submitted_handler
)
self.omotes_if.connect_to_job_cancellations(self.job_cancellation_handler)

self.jobs_broker_if.start()
self.jobs_broker_if.add_queue_subscription(
"omotes_task_result_events", self.task_result_received
Expand All @@ -200,6 +204,19 @@ def start(self) -> None:
"omotes_task_progress_events", self.task_progress_update
)

self.omotes_if.start()
self.omotes_if.connect_to_request_available_workflows(

Check failure on line 208 in src/omotes_orchestrator/main.py

View workflow job for this annotation

GitHub Actions / Typecheck (3.11)

error: "OrchestratorInterface" has no attribute "connect_to_request_available_workflows" [attr-defined]
callback_on_request_workflows=self.request_workflows_handler
)
self.omotes_if.connect_to_job_submissions(
callback_on_new_job=self.new_job_submitted_handler
)
self.omotes_if.connect_to_job_cancellations(
callback_on_job_cancel=self.job_cancellation_handler
)

self.omotes_if.send_available_workflows(self.workflow_manager)

Check failure on line 218 in src/omotes_orchestrator/main.py

View workflow job for this annotation

GitHub Actions / Typecheck (3.11)

error: "OrchestratorInterface" has no attribute "send_available_workflows" [attr-defined]

def stop(self) -> None:
"""Stop the orchestrator."""
self.omotes_if.stop()
Expand Down Expand Up @@ -325,6 +342,14 @@ def job_cancellation_handler(self, job_cancellation: JobCancel) -> None:
)
self._cleanup_job(job_id)

def request_workflows_handler(self, request_workflows: RequestAvailableWorkflows) -> None:
"""When a available work flows request is received from the SDK.
:param request_workflows: Request available work flows.
"""
logger.info("Received an available workflows request")
self.omotes_if.send_available_workflows(self.workflow_manager)

Check failure on line 351 in src/omotes_orchestrator/main.py

View workflow job for this annotation

GitHub Actions / Typecheck (3.11)

error: "OrchestratorInterface" has no attribute "send_available_workflows" [attr-defined]

def _cleanup_job(self, job_id: uuid.UUID) -> None:
"""Cleanup any references to job with id `job_id`.
Expand Down Expand Up @@ -526,7 +551,8 @@ def main() -> None:
workflow_type_description_name="Grow Optimizer default workflow",
),
WorkflowType(
workflow_type_name="grow_simulator", workflow_type_description_name="Grow Simulator"
workflow_type_name="grow_simulator",
workflow_type_description_name="Grow Simulator",
),
WorkflowType(
workflow_type_name="grow_optimizer_no_heat_losses",
Expand All @@ -539,6 +565,31 @@ def main() -> None:
WorkflowType(

Check failure on line 565 in src/omotes_orchestrator/main.py

View workflow job for this annotation

GitHub Actions / Typecheck (3.11)

error: Unexpected keyword argument "workflow_parameters" for "WorkflowType" [call-arg]
workflow_type_name="simulator",
workflow_type_description_name="High fidelity simulator",
workflow_parameters=[
WorkflowParameter(
key_name="start_time",
schema=ParameterSchema(
type=ParameterType.STRING,
format=ParameterStringFormat.DATETIME,
),
),
WorkflowParameter(
key_name="step_size_in_seconds",
schema=ParameterSchema(
type=ParameterType.INTEGER,
default="60",
minimum=0,
),
),
WorkflowParameter(
key_name="number_of_steps",
schema=ParameterSchema(
type=ParameterType.INTEGER,
minimum=0,
maximum=1,
),
),
],
),
WorkflowType(
workflow_type_name="test_worker",
Expand Down

0 comments on commit 335702b

Please sign in to comment.