orchestrator-core 3.2.2__py3-none-any.whl → 4.0.0__py3-none-any.whl

orchestrator/__init__.py

@@ -13,7 +13,7 @@
 
 """This is the orchestrator workflow engine."""
 
-__version__ = "3.2.2"
+__version__ = "4.0.0"
 
 from orchestrator.app import OrchestratorCore
 from orchestrator.settings import app_settings

orchestrator/api/api_v1/endpoints/settings.py

@@ -22,9 +22,10 @@ from sqlalchemy.exc import SQLAlchemyError
 from oauth2_lib.fastapi import OIDCUserModel
 from orchestrator.api.error_handling import raise_status
 from orchestrator.db import EngineSettingsTable
-from orchestrator.schemas import EngineSettingsBaseSchema, EngineSettingsSchema, GlobalStatusEnum, WorkerStatus
+from orchestrator.schemas import EngineSettingsBaseSchema, EngineSettingsSchema, WorkerStatus
 from orchestrator.security import authenticate
 from orchestrator.services import processes, settings
+from orchestrator.services.settings import generate_engine_global_status
 from orchestrator.settings import ExecutorType, app_settings
 from orchestrator.utils.json import json_dumps
 from orchestrator.utils.redis import delete_keys_matching_pattern
@@ -165,17 +166,6 @@ def generate_engine_status_response(
         Engine StatusEnum
 
     """
-
-    if engine_settings.global_lock and engine_settings.running_processes > 0:
-        result = EngineSettingsSchema.model_validate(engine_settings)
-        result.global_status = GlobalStatusEnum.PAUSING
-        return result
-
-    if engine_settings.global_lock and engine_settings.running_processes == 0:
-        result = EngineSettingsSchema.model_validate(engine_settings)
-        result.global_status = GlobalStatusEnum.PAUSED
-        return result
-
     result = EngineSettingsSchema.model_validate(engine_settings)
-    result.global_status = GlobalStatusEnum.RUNNING
+    result.global_status = generate_engine_global_status(engine_settings)
     return result

orchestrator/api/api_v1/endpoints/subscription_customer_descriptions.py

@@ -28,7 +28,6 @@ from orchestrator.domain.customer_description import (
 from orchestrator.schemas import SubscriptionDescriptionBaseSchema, SubscriptionDescriptionSchema
 from orchestrator.schemas.subscription_descriptions import UpdateSubscriptionDescriptionSchema
 from orchestrator.utils.errors import StaleDataError
-from orchestrator.utils.redis import delete_from_redis
 
 router = APIRouter()
 
@@ -55,7 +54,6 @@ def delete_subscription_customer_descriptions(_id: UUID) -> None:
     description = db.session.get(SubscriptionCustomerDescriptionTable, _id)
     if description:
         delete(SubscriptionCustomerDescriptionTable, _id)
-        delete_from_redis(description.subscription_id)
 
 
 @router.get("/{_id}", response_model=SubscriptionDescriptionSchema)

orchestrator/api/api_v1/endpoints/subscriptions.py

@@ -163,6 +163,7 @@ def subscriptions_search(
 @router.get(
     "/workflows/{subscription_id}",
     response_model=SubscriptionWorkflowListsSchema,
+    response_model_by_alias=True,
     response_model_exclude_none=True,
     deprecated=True,
     description="This endpoint is deprecated and will be removed in a future release. Please use the GraphQL query",

orchestrator/app.py

@@ -25,6 +25,7 @@ import structlog
 import typer
 from fastapi.applications import FastAPI
 from fastapi_etag.dependency import add_exception_handler
+from prometheus_client import make_asgi_app
 from sentry_sdk.integrations import Integration
 from sentry_sdk.integrations.asyncio import AsyncioIntegration
 from sentry_sdk.integrations.fastapi import FastApiIntegration
@@ -54,6 +55,7 @@ from orchestrator.graphql.schema import ContextGetterFactory
 from orchestrator.graphql.schemas.subscription import SubscriptionInterface
 from orchestrator.graphql.types import ScalarOverrideType, StrawberryModelType
 from orchestrator.log_config import LOGGER_OVERRIDES
+from orchestrator.metrics import ORCHESTRATOR_METRICS_REGISTRY, initialize_default_metrics
 from orchestrator.services.process_broadcast_thread import ProcessDataBroadcastThread
 from orchestrator.settings import AppSettings, ExecutorType, app_settings
 from orchestrator.version import GIT_COMMIT_HASH
@@ -143,9 +145,14 @@ class OrchestratorCore(FastAPI):
         self.add_exception_handler(ProblemDetailException, problem_detail_handler)  # type: ignore[arg-type]
         add_exception_handler(self)
 
+        if base_settings.ENABLE_PROMETHEUS_METRICS_ENDPOINT:
+            initialize_default_metrics()
+            metrics_app = make_asgi_app(registry=ORCHESTRATOR_METRICS_REGISTRY)
+            self.mount("/api/metrics", metrics_app)
+
         @self.router.get("/", response_model=str, response_class=JSONResponse, include_in_schema=False)
         def _index() -> str:
-            return "Orchestrator orchestrator"
+            return "Orchestrator Core"
 
     def add_sentry(
         self,
@@ -154,6 +161,7 @@ class OrchestratorCore(FastAPI):
         server_name: str,
         environment: str,
         release: str | None = GIT_COMMIT_HASH,
+        **sentry_kwargs: Any,
     ) -> None:
         logger.info("Adding Sentry middleware to app", app=self.title)
         if self.base_settings.EXECUTOR == ExecutorType.WORKER:
@@ -173,6 +181,7 @@ class OrchestratorCore(FastAPI):
             integrations=sentry_integrations,
             propagate_traces=True,
             profiles_sample_rate=trace_sample_rate,
+            **sentry_kwargs,
         )
 
     @staticmethod

orchestrator/cli/generator/templates/new_product_migration.j2

@@ -63,24 +63,28 @@ new_workflows = [
     {
         "name": "create_{{ product.variable }}",
         "target": Target.CREATE,
+        "is_task": False,
         "description": "Create {{ product.name }}",
         "product_type": "{{ product.type }}",
     },
     {
         "name": "modify_{{ product.variable }}",
         "target": Target.MODIFY,
+        "is_task": False,
         "description": "Modify {{ product.name }}",
         "product_type": "{{ product.type }}",
     },
     {
         "name": "terminate_{{ product.variable }}",
         "target": Target.TERMINATE,
+        "is_task": False,
         "description": "Terminate {{ product.name }}",
         "product_type": "{{ product.type }}",
     },
     {
         "name": "validate_{{ product.variable }}",
-        "target": Target.SYSTEM,
+        "target": Target.VALIDATE,
+        "is_task": True,
         "description": "Validate {{ product.name }}",
         "product_type": "{{ product.type }}",
     },

orchestrator/cli/migrate_tasks.py

@@ -150,22 +150,22 @@ def create_tasks_migration_wizard() -> tuple[list[dict], list[dict]]:
         - list of task items to add in the migration
         - list of task items to delete in the migration
     """
-    database_tasks = {
-        task.name: task for task in list(db.session.scalars(select(WorkflowTable))) if task.target == Target.SYSTEM
-    }
+    database_tasks = {task.name: task for task in list(db.session.scalars(select(WorkflowTable))) if task.is_task}
     registered_wf_instances = {
         task: cast(Workflow, get_workflow(task)) for task in orchestrator.workflows.ALL_WORKFLOWS.keys()
     }
 
+    is_task = [Target.SYSTEM, Target.VALIDATE]
+
     registered_tasks = dict(
         filter(
-            lambda task: task[1].target == Target.SYSTEM and task[0] in database_tasks.keys(),
+            lambda task: task[1].target in is_task and task[0] in database_tasks.keys(),
             registered_wf_instances.items(),
         )
     )
     available_tasks = dict(
         filter(
-            lambda task: task[1].target == Target.SYSTEM and task[0] not in database_tasks.keys(),
+            lambda task: task[1].target in is_task and task[0] not in database_tasks.keys(),
             registered_wf_instances.items(),
         )
     )

orchestrator/cli/migrate_workflows.py

@@ -11,7 +11,6 @@ import orchestrator.workflows
 from orchestrator.cli.helpers.input_helpers import _enumerate_menu_keys, _prompt_user_menu, get_user_input
 from orchestrator.cli.helpers.print_helpers import COLOR, noqa_print, print_fmt, str_fmt
 from orchestrator.db import ProductTable, WorkflowTable, db
-from orchestrator.targets import Target
 from orchestrator.workflows import LazyWorkflowInstance, get_workflow
 
 # Workflows are registered via migrations with product type. For every product with the given product_type, there will be an entry in products_workflows.
@@ -183,7 +182,7 @@ def create_workflows_migration_wizard() -> tuple[list[dict], list[dict]]:
     """
     database_workflows = list(db.session.scalars(select(WorkflowTable)))
     registered_workflows = orchestrator.workflows.ALL_WORKFLOWS
-    system_workflow_names = {wf.name for wf in database_workflows if wf.target == Target.SYSTEM}
+    system_workflow_names = {wf.name for wf in database_workflows if wf.is_task}
     registered_non_system_workflows = {k: v for k, v in registered_workflows.items() if k not in system_workflow_names}
 
     unregistered_workflows = [wf for wf in database_workflows if wf.name not in registered_workflows.keys()]

orchestrator/db/models.py

@@ -400,7 +400,7 @@ class WorkflowTable(BaseModel):
     workflow_id = mapped_column(UUIDType, server_default=text("uuid_generate_v4()"), primary_key=True)
     name = mapped_column(String(), nullable=False, unique=True)
     target = mapped_column(String(), nullable=False)
-    description = mapped_column(Text(), nullable=True)
+    description = mapped_column(Text(), nullable=False)
     created_at = mapped_column(UtcTimestamp, nullable=False, server_default=text("current_timestamp()"))
     deleted_at = mapped_column(UtcTimestamp, deferred=True)
 
@@ -412,6 +412,8 @@ class WorkflowTable(BaseModel):
     )
     processes = relationship("ProcessTable", cascade="all, delete-orphan", back_populates="workflow")
 
+    is_task = mapped_column(Boolean, nullable=False, server_default=text("false"))
+
     @staticmethod
     def select() -> Select:
         return (

orchestrator/domain/base.py

@@ -1294,27 +1294,13 @@ class SubscriptionModel(DomainModel):
     # Some common functions shared by from_other_product and from_subscription
     @classmethod
     def _get_subscription(cls: type[S], subscription_id: UUID | UUIDstr) -> SubscriptionTable | None:
-        from orchestrator.settings import app_settings
 
         if not isinstance(subscription_id, UUID | UUIDstr):
             raise TypeError(f"subscription_id is of type {type(subscription_id)} instead of UUID | UUIDstr")
 
-        if app_settings.ENABLE_SUBSCRIPTION_MODEL_OPTIMIZATIONS:
-            # TODO #900 remove toggle and make this path the default
-            loaders = [
-                joinedload(SubscriptionTable.product).selectinload(ProductTable.fixed_inputs),
-            ]
-
-        else:
-            loaders = [
-                selectinload(SubscriptionTable.instances)
-                .joinedload(SubscriptionInstanceTable.product_block)
-                .selectinload(ProductBlockTable.resource_types),
-                selectinload(SubscriptionTable.instances).selectinload(
-                    SubscriptionInstanceTable.in_use_by_block_relations
-                ),
-                selectinload(SubscriptionTable.instances).selectinload(SubscriptionInstanceTable.values),
-            ]
+        loaders = [
+            joinedload(SubscriptionTable.product).selectinload(ProductTable.fixed_inputs),
+        ]
 
         return db.session.get(SubscriptionTable, subscription_id, options=loaders)
 
@@ -1394,7 +1380,6 @@ class SubscriptionModel(DomainModel):
     def from_subscription(cls: type[S], subscription_id: UUID | UUIDstr) -> S:
         """Use a subscription_id to return required fields of an existing subscription."""
         from orchestrator.domain.context_cache import get_from_cache, store_in_cache
-        from orchestrator.settings import app_settings
 
         if cached_model := get_from_cache(subscription_id):
             return cast(S, cached_model)
@@ -1421,12 +1406,7 @@ class SubscriptionModel(DomainModel):
 
         fixed_inputs = {fi.name: fi.value for fi in subscription.product.fixed_inputs}
 
-        instances: dict[str, Any]
-        if app_settings.ENABLE_SUBSCRIPTION_MODEL_OPTIMIZATIONS:
-            # TODO #900 remove toggle and make this path the default
-            instances = cls._load_root_instances(subscription_id)
-        else:
-            instances = cls._load_instances(subscription.instances, status, match_domain_attr=False)
+        instances = cls._load_root_instances(subscription_id)
 
         try:
             model = cls(

orchestrator/domain/customer_description.py

@@ -21,7 +21,6 @@ from sqlalchemy import select
 from orchestrator.api.models import delete
 from orchestrator.db import SubscriptionCustomerDescriptionTable, db
 from orchestrator.utils.errors import StaleDataError
-from orchestrator.utils.redis import delete_subscription_from_redis
 from orchestrator.utils.validate_data_version import validate_data_version
 from orchestrator.websocket import invalidate_subscription_cache
 
@@ -38,7 +37,6 @@ def get_customer_description_by_customer_subscription(
     return db.session.scalars(stmt).one_or_none()
 
 
-@delete_subscription_from_redis()
 async def create_subscription_customer_description(
     customer_id: str, subscription_id: UUID, description: str
 ) -> SubscriptionCustomerDescriptionTable:
@@ -53,7 +51,6 @@ async def create_subscription_customer_description(
     return customer_description
 
 
-@delete_subscription_from_redis()
 async def update_subscription_customer_description(
     customer_description: SubscriptionCustomerDescriptionTable,
     description: str,
@@ -70,7 +67,6 @@ async def update_subscription_customer_description(
     return customer_description
 
 
-@delete_subscription_from_redis()
 async def delete_subscription_customer_description_by_customer_subscription(
     customer_id: str, subscription_id: UUID
 ) -> SubscriptionCustomerDescriptionTable | None:

orchestrator/graphql/mutations/customer_description.py

@@ -60,7 +60,7 @@ async def resolve_remove_customer_description(
     )
     if not description:
         return NotFoundError(message="Customer description not found")
-    return CustomerDescription.from_pydantic(description)
+    return CustomerDescription.from_pydantic(description)  # type: ignore
 
 
 @strawberry.type(description="Customer subscription description mutations")

orchestrator/metrics/__init__.py

@@ -0,0 +1,3 @@
+from orchestrator.metrics.init import ORCHESTRATOR_METRICS_REGISTRY, initialize_default_metrics
+
+__all__ = ["initialize_default_metrics", "ORCHESTRATOR_METRICS_REGISTRY"]

orchestrator/metrics/engine.py

@@ -0,0 +1,49 @@
+from typing import Iterable
+
+from prometheus_client import Metric
+from prometheus_client.metrics_core import GaugeMetricFamily, StateSetMetricFamily
+from prometheus_client.registry import Collector
+
+from orchestrator.schemas.engine_settings import GlobalStatusEnum
+from orchestrator.services import settings
+
+
+def _get_engine_status() -> tuple[GlobalStatusEnum, int]:
+    """Query for getting the current status of the workflow engine.
+
+    This includes the engine status, and the amount of currently running processes.
+    """
+    engine_settings = settings.get_engine_settings()
+    engine_status = settings.generate_engine_global_status(engine_settings)
+
+    return engine_status, int(engine_settings.running_processes)
+
+
+class WorkflowEngineCollector(Collector):
+    """Initialize a Prometheus enum and a gauge.
+
+    The enum of the current workflow engine status takes three values:
+        - RUNNING
+        - PAUSING
+        - PAUSED
+
+    This metric also exports the amount of currently running processes.
+    """
+
+    def collect(self) -> Iterable[Metric]:
+        current_engine_status, running_process_count = _get_engine_status()
+
+        engine_status = StateSetMetricFamily(
+            "wfo_engine_status",
+            documentation="Current workflow engine status.",
+            value={status: status == current_engine_status for status in GlobalStatusEnum.values()},
+        )
+
+        engine_process_count = GaugeMetricFamily(
+            "wfo_active_process_count",
+            unit="count",
+            value=running_process_count,
+            documentation="Number of currently running processes in the workflow engine.",
+        )
+
+        return [engine_status, engine_process_count]

orchestrator/metrics/init.py

@@ -0,0 +1,14 @@
+from prometheus_client import CollectorRegistry
+
+from orchestrator.metrics.engine import WorkflowEngineCollector
+from orchestrator.metrics.processes import ProcessCollector
+from orchestrator.metrics.subscriptions import SubscriptionCollector
+
+ORCHESTRATOR_METRICS_REGISTRY = CollectorRegistry(auto_describe=True)
+
+
+def initialize_default_metrics() -> None:
+    """Register default Prometheus collectors."""
+    ORCHESTRATOR_METRICS_REGISTRY.register(SubscriptionCollector())
+    ORCHESTRATOR_METRICS_REGISTRY.register(ProcessCollector())
+    ORCHESTRATOR_METRICS_REGISTRY.register(WorkflowEngineCollector())

orchestrator/metrics/processes.py

@@ -0,0 +1,147 @@
+from typing import Iterable
+
+from prometheus_client.metrics_core import GaugeMetricFamily, Metric
+from prometheus_client.registry import Collector
+from pydantic import BaseModel
+from sqlalchemy import desc, func
+
+from orchestrator.db import ProcessTable, ProductTable, SubscriptionTable, WorkflowTable, db
+from orchestrator.db.models import ProcessSubscriptionTable
+from orchestrator.targets import Target
+from orchestrator.workflow import ProcessStatus
+from pydantic_forms.types import UUIDstr
+
+
+class ProcessTableQueryResult(BaseModel):
+    workflow_name: str
+    customer_id: UUIDstr
+    workflow_target: Target
+    last_status: ProcessStatus
+    created_by: str
+    is_task: bool
+    product_name: str
+    total_runtime: float
+    process_count: int
+
+
+def _get_processes() -> list[ProcessTableQueryResult]:
+    """Query for getting all processes.
+
+    Equivalent to the following SQL statement:
+    ```sql
+    SELECT
+      workflows."name"                                                                          AS workflow_name
+      , subscriptions.customer_id
+      , workflows.target                                                                        AS workflow_target
+      , processes.last_status
+      , processes.created_by
+      , processes.is_task
+      , products."name"                                                                         AS product_name
+      , Coalesce(Sum(Extract(EPOCH FROM processes.last_modified_at - processes.started_at)), 0) AS total_runtime
+      , Count(workflows."name")                                                                 AS process_count
+    FROM
+      processes
+    JOIN workflows
+      ON processes.workflow_id = workflows.workflow_id
+    JOIN processes_subscriptions
+      ON processes.pid = processes_subscriptions.pid
+    JOIN subscriptions
+      ON processes_subscriptions.subscription_id = subscriptions.subscription_id
+    JOIN products
+      ON subscriptions.product_id = products.product_id
+    GROUP BY
+      workflows."name"
+      , subscriptions.customer_id
+      , workflows.target
+      , processes.last_status
+      , processes.created_by
+      , processes.is_task
+      , products."name"
+    ;
+    ```
+    """
+    process_count = func.count(WorkflowTable.name).label("process_count")
+    total_process_time = func.coalesce(
+        func.sum(func.extract("epoch", (ProcessTable.last_modified_at - ProcessTable.started_at))), 0
+    ).label("total_runtime")
+    return (
+        db.session.query(
+            ProcessTable.last_status,
+            ProcessTable.created_by,
+            ProcessTable.is_task,
+            ProductTable.name.label("product_name"),
+            WorkflowTable.name.label("workflow_name"),
+            SubscriptionTable.customer_id,
+            WorkflowTable.target.label("workflow_target"),
+            process_count,
+            total_process_time,
+        )
+        .join(WorkflowTable, WorkflowTable.workflow_id == ProcessTable.workflow_id)
+        .join(ProcessSubscriptionTable, ProcessSubscriptionTable.process_id == ProcessTable.process_id)
+        .join(SubscriptionTable, SubscriptionTable.subscription_id == ProcessSubscriptionTable.subscription_id)
+        .join(ProductTable, ProductTable.product_id == SubscriptionTable.product_id)
+        .group_by(
+            ProcessTable.last_status,
+            ProcessTable.created_by,
+            ProcessTable.is_task,
+            ProductTable.name,
+            WorkflowTable.name,
+            SubscriptionTable.customer_id,
+            WorkflowTable.target,
+        )
+        .order_by(desc(process_count))
+    ).all()
+
+
+class ProcessCollector(Collector):
+    """Collector that contains two Prometheus gauges with process counts and total runtime.
+
+    These gauges contain the amount of processes, and the total runtime in seconds, per every combination of the labels
+    that are defined:
+        - Process last status
+        - Process created by
+        - Process is task
+        - Product name
+        - Workflow name
+        - Customer ID
+        - Workflow target
+    """
+
+    def collect(self) -> Iterable[Metric]:
+        label_names = [
+            "last_status",
+            "created_by",
+            "is_task",
+            "product_name",
+            "workflow_name",
+            "customer_id",
+            "workflow_target",
+        ]
+        process_counts = GaugeMetricFamily(
+            "wfo_process_count",
+            labels=label_names,
+            unit="count",
+            documentation="Number of processes per status, creator, task, product, workflow, customer, and target.",
+        )
+        process_seconds_total = GaugeMetricFamily(
+            "wfo_process_seconds_total",
+            labels=label_names,
+            unit="count",
+            documentation="Total time spent on processes in seconds.",
+        )
+
+        for row in _get_processes():
+            label_values = [
+                row.last_status,
+                str(row.created_by),
+                str(row.is_task),
+                row.product_name,
+                row.workflow_name,
+                row.customer_id,
+                row.workflow_target,
+            ]
+
+            process_counts.add_metric(label_values, row.process_count)
+            process_seconds_total.add_metric(label_values, row.total_runtime)
+
+        return [process_counts, process_seconds_total]

orchestrator/metrics/subscriptions.py

@@ -0,0 +1,93 @@
+from typing import Iterable
+
+from prometheus_client import Metric
+from prometheus_client.metrics_core import GaugeMetricFamily
+from prometheus_client.registry import Collector
+from pydantic import BaseModel
+from sqlalchemy import desc, func
+
+from orchestrator.db import ProductTable, SubscriptionTable, db
+from orchestrator.types import SubscriptionLifecycle
+from pydantic_forms.types import UUIDstr
+
+
+class SubscriptionTableQueryResult(BaseModel):
+    lifecycle_state: SubscriptionLifecycle
+    customer_id: UUIDstr
+    insync: bool
+    product_name: str
+    subscription_count: int
+
+
+def _get_subscriptions() -> list[SubscriptionTableQueryResult]:
+    """Query for getting all subscriptions.
+
+    Equivalent to the following SQL statement:
+    ```sql
+    SELECT
+        subscriptions.status                   AS lifecycle_state
+        , subscriptions.customer_id
+        , subscriptions.insync
+        , products."name"                      AS product_name
+        , Count(subscriptions.subscription_id) AS subscription_count
+    FROM
+      subscriptions
+      JOIN products
+        ON subscriptions.product_id = products.product_id
+    GROUP BY  subscriptions.status
+              , subscriptions.customer_id
+              , insync
+              , products."name"
+    ;
+    ```
+    """
+    subscription_count = func.count(SubscriptionTable.subscription_id).label("subscription_count")
+    return (
+        db.session.query(
+            SubscriptionTable.status.label("lifecycle_state"),
+            SubscriptionTable.customer_id,
+            SubscriptionTable.insync,
+            ProductTable.name.label("product_name"),
+            subscription_count,
+        )
+        .outerjoin(ProductTable, ProductTable.product_id == SubscriptionTable.product_id)
+        .group_by(
+            SubscriptionTable.status,
+            SubscriptionTable.customer_id,
+            SubscriptionTable.insync,
+            ProductTable.name,
+        )
+        .order_by(desc(subscription_count))
+    ).all()
+
+
+class SubscriptionCollector(Collector):
+    """Collector for Subscriptions stored in the subscription database.
+
+    This collector contains one gauge that contains the amount of subscriptions, per every combination of the labels
+    that are defined:
+        - Product name
+        - Subscription lifecycle
+        - Customer ID
+        - `insync` state
+    """
+
+    def collect(self) -> Iterable[Metric]:
+        subscriptions = GaugeMetricFamily(
+            name="wfo_subscriptions_count",
+            labels=[
+                "product_name",
+                "lifecycle_state",
+                "customer_id",
+                "insync",
+            ],
+            unit="count",
+            documentation="Number of subscriptions per product, lifecycle state, customer, and in sync state.",
+        )
+
+        for row in _get_subscriptions():
+            subscriptions.add_metric(
+                [row.product_name, row.lifecycle_state, row.customer_id, str(row.insync)], row.subscription_count
+            )
+
+        return [subscriptions]