Various enhancements around Prefect

Author: dgarros

backend/infrahub/cli/__init__.py

@@ -8,6 +8,7 @@
 from infrahub.cli.events import app as events_app
 from infrahub.cli.git_agent import app as git_app
 from infrahub.cli.server import app as server_app
+from infrahub.cli.tasks import app as tasks_app
 from infrahub.core.initialization import initialization
 from infrahub.database import InfrahubDatabase, get_db
 
@@ -26,6 +27,7 @@ def common(ctx: typer.Context) -> None:
 app.add_typer(git_app, name="git-agent")
 app.add_typer(db_app, name="db")
 app.add_typer(events_app, name="events", help="Interact with the events system.")
+app.add_typer(tasks_app, name="tasks", hidden=True)
 
 
 async def _init_shell(config_file: str) -> None:

backend/infrahub/cli/tasks.py

@@ -0,0 +1,50 @@
+import logging
+
+import typer
+from infrahub_sdk.async_typer import AsyncTyper
+from prefect.client.orchestration import get_client
+
+from infrahub import config
+from infrahub.services.adapters.workflow.worker import WorkflowWorkerExecution
+from infrahub.tasks.dummy import DUMMY_FLOW, DummyInput
+from infrahub.workflows.initialization import setup_task_manager
+from infrahub.workflows.models import WorkerPoolDefinition
+
+app = AsyncTyper()
+
+# pylint: disable=unused-argument
+
+
+@app.command()
+async def init(
+    ctx: typer.Context,
+    debug: bool = typer.Option(False, help="Enable advanced logging and troubleshooting"),
+    config_file: str = typer.Argument("infrahub.toml", envvar="INFRAHUB_CONFIG"),
+) -> None:
+    """Initialize the task manager"""
+    logging.getLogger("prefect").setLevel(logging.ERROR)
+
+    config.load_and_exit(config_file_name=config_file)
+
+    await setup_task_manager()
+
+
+@app.command()
+async def execute(
+    ctx: typer.Context,
+    debug: bool = typer.Option(False, help="Enable advanced logging and troubleshooting"),
+    config_file: str = typer.Argument("infrahub.toml", envvar="INFRAHUB_CONFIG"),
+) -> None:
+    """Check the current format of the internal graph and apply the necessary migrations"""
+    logging.getLogger("infrahub").setLevel(logging.WARNING)
+    logging.getLogger("neo4j").setLevel(logging.ERROR)
+    logging.getLogger("prefect").setLevel(logging.ERROR)
+
+    config.load_and_exit(config_file_name=config_file)
+
+    async with get_client(sync_client=False) as client:
+        worker = WorkflowWorkerExecution()
+        await DUMMY_FLOW.save(client=client, work_pool=WorkerPoolDefinition(name="testing", worker_type="process"))
+
+        result = await worker.execute(workflow=DUMMY_FLOW, data=DummyInput(firstname="John", lastname="Doe"))  # type: ignore[var-annotated]
+        print(result)

backend/infrahub/config.py

@@ -269,8 +269,8 @@ class CacheSettings(BaseSettings):
     )
     database: int = Field(default=0, ge=0, le=15, description="Id of the database to use")
     driver: CacheDriver = CacheDriver.Redis
-    username: str = "infrahub"
-    password: str = "infrahub"
+    username: str = ""
+    password: str = ""
     tls_enabled: bool = Field(default=False, description="Indicates if TLS is enabled for the connection")
     tls_insecure: bool = Field(default=False, description="Indicates if TLS certificates are verified")
     tls_ca_file: Optional[str] = Field(default=None, description="File path to CA cert or bundle in PEM format")
@@ -290,6 +290,9 @@ class WorkflowSettings(BaseSettings):
     port: Optional[int] = Field(default=None, ge=1, le=65535, description="Specified if running on a non default port.")
     tls_enabled: bool = Field(default=False, description="Indicates if TLS is enabled for the connection")
     driver: WorkflowDriver = WorkflowDriver.WORKER
+    worker_polling_interval: int = Field(
+        default=2, ge=1, le=30, description="Specify how often the worker should poll the server for tasks (sec)"
+    )
 
     @property
     def api_endpoint(self) -> str:

backend/infrahub/graphql/app.py

@@ -38,7 +38,7 @@
 )
 from opentelemetry import trace
 from starlette.datastructures import UploadFile
-from starlette.requests import HTTPConnection, Request
+from starlette.requests import ClientDisconnect, HTTPConnection, Request
 from starlette.responses import JSONResponse, Response
 from starlette.websockets import WebSocket, WebSocketDisconnect, WebSocketState
 
@@ -191,6 +191,9 @@ async def _handle_http_request(
             operations = await _get_operation_from_request(request)
         except ValueError as exc:
             return JSONResponse({"errors": [exc.args[0]]}, status_code=400)
+        except ClientDisconnect as exc:
+            self.logger.error("Exception ClientDisconnect in _handle_http_request")
+            return JSONResponse({"errors": [str(exc)]}, status_code=400)
 
         if isinstance(operations, list):
             return JSONResponse({"errors": ["This server does not support batching"]}, status_code=400)

backend/infrahub/services/__init__.py

@@ -100,10 +100,10 @@ def git_report(
 
     async def initialize(self) -> None:
         """Initialize the Services"""
-        await self.component.initialize(service=self)
-        await self.http.initialize(service=self)
         await self.message_bus.initialize(service=self)
         await self.cache.initialize(service=self)
+        await self.http.initialize(service=self)
+        await self.component.initialize(service=self)
         await self.scheduler.initialize(service=self)
         await self.workflow.initialize(service=self)
         await self.event.initialize(service=self)

backend/infrahub/services/adapters/workflow/worker.py

@@ -2,12 +2,10 @@
 
 from typing import TYPE_CHECKING, Any, Awaitable, Callable
 
-from prefect.client.orchestration import get_client
-from prefect.client.schemas.actions import WorkPoolCreate
+from prefect.client.schemas import StateType
 from prefect.deployments import run_deployment
-from prefect.exceptions import ObjectAlreadyExists
 
-from infrahub.workflows.catalogue import worker_pools, workflows
+from infrahub.workflows.initialization import setup_task_manager
 
 from . import InfrahubWorkflow, Return
 
@@ -22,23 +20,8 @@ class WorkflowWorkerExecution(InfrahubWorkflow):
     async def initialize(self, service: InfrahubServices) -> None:
         """Initialize the Workflow engine"""
 
-        async with get_client(sync_client=False) as client:
-            for worker in worker_pools:
-                wp = WorkPoolCreate(
-                    name=worker.name,
-                    type=worker.worker_type,
-                    description=worker.description,
-                )
-                try:
-                    await client.create_work_pool(work_pool=wp)
-                    service.log.info(f"work pool {worker} created successfully ... ")
-                except ObjectAlreadyExists:
-                    service.log.info(f"work pool {worker} already present ")
-
-            # Create deployment
-            for workflow in workflows:
-                flow_id = await client.create_flow_from_name(workflow.name)
-                await client.create_deployment(flow_id=flow_id, **workflow.to_deployment())
+        if await service.component.is_primary_api():
+            await setup_task_manager()
 
     async def execute(
         self,
@@ -50,7 +33,11 @@ async def execute(
             response: FlowRun = await run_deployment(name=workflow.full_name, parameters=kwargs or {})  # type: ignore[return-value, misc]
             if not response.state:
                 raise RuntimeError("Unable to read state from the response")
-            return response.state.result(raise_on_failure=True)
+
+            if response.state.type == StateType.CRASHED:
+                raise RuntimeError(response.state.message)
+
+            return await response.state.result(raise_on_failure=True, fetch=True)  # type: ignore[call-overload]
 
         if function:
             return await function(**kwargs)

backend/infrahub/services/component.py

@@ -42,6 +42,8 @@ async def initialize(self, service: InfrahubServices) -> None:
         """Initialize the Message bus"""
         self._service = service
 
+        await self.refresh_heartbeat()
+
     async def is_primary_api(self) -> bool:
         primary_identity = await self.service.cache.get(PRIMARY_API_SERVER)
         return primary_identity == WORKER_IDENTITY

backend/infrahub/tasks/dummy.py

@@ -1,6 +1,20 @@
 from prefect import flow, task
 from pydantic import BaseModel
 
+from infrahub.workflows.models import WorkflowDefinition
+
+DUMMY_FLOW = WorkflowDefinition(
+    name="dummy_flow",
+    module="infrahub.tasks.dummy",
+    function="dummy_flow",
+)
+
+DUMMY_FLOW_BROKEN = WorkflowDefinition(
+    name="dummy_flow_broken",
+    module="infrahub.tasks.dummy",
+    function="dummy_flow_broken",
+)
+
 
 class DummyInput(BaseModel):
     firstname: str
@@ -19,3 +33,9 @@ async def aggregate_name(firstname: str, lastname: str) -> str:
 @flow(persist_result=True)
 async def dummy_flow(data: DummyInput) -> DummyOutput:
     return DummyOutput(full_name=await aggregate_name(firstname=data.firstname, lastname=data.lastname))
+
+
+@flow(persist_result=True)
+async def dummy_flow_broken(data: DummyInput) -> DummyOutput:
+    response = await aggregate_name(firstname=data.firstname, lastname=data.lastname)
+    return DummyOutput(not_valid=response)  # type: ignore[call-arg]

backend/infrahub/workers/infrahub_async.py

@@ -7,6 +7,7 @@
 from anyio.abc import TaskStatus
 from infrahub_sdk import Config, InfrahubClient
 from infrahub_sdk.exceptions import Error as SdkError
+from prefect import settings as prefect_settings
 from prefect.client.schemas.objects import FlowRun
 from prefect.flow_engine import run_flow_async
 from prefect.workers.base import BaseJobConfiguration, BaseVariables, BaseWorker, BaseWorkerResult
@@ -26,6 +27,7 @@
 from infrahub.services.adapters.message_bus.rabbitmq import RabbitMQMessageBus
 from infrahub.services.adapters.workflow.local import WorkflowLocalExecution
 from infrahub.services.adapters.workflow.worker import WorkflowWorkerExecution
+from infrahub.workflows.models import TASK_RESULT_STORAGE_NAME
 
 
 class InfrahubWorkerAsyncConfiguration(BaseJobConfiguration):
@@ -49,8 +51,6 @@ class InfrahubWorkerAsync(BaseWorker):
     _description = "Infrahub worker designed to run the flow in the main async loop."
 
     async def setup(self, **kwargs: dict[str, Any]) -> None:
-        await super().setup(**kwargs)
-
         logging.getLogger("websockets").setLevel(logging.ERROR)
         logging.getLogger("httpx").setLevel(logging.ERROR)
         logging.getLogger("httpcore").setLevel(logging.ERROR)
@@ -67,7 +67,19 @@ async def setup(self, **kwargs: dict[str, Any]) -> None:
         self._logger.info(f"Starting metric endpoint on port {metric_port}")
         start_http_server(metric_port)
 
-        self._logger.info(f"Using Infrahub API at {config.SETTINGS.main.internal_address}")
+        self._exit_stack.enter_context(
+            prefect_settings.temporary_settings(
+                updates={  # type: ignore[arg-type]
+                    prefect_settings.PREFECT_WORKER_QUERY_SECONDS: config.SETTINGS.workflow.worker_polling_interval,
+                    prefect_settings.PREFECT_RESULTS_PERSIST_BY_DEFAULT: True,
+                    prefect_settings.PREFECT_DEFAULT_RESULT_STORAGE_BLOCK: f"redisstoragecontainer/{TASK_RESULT_STORAGE_NAME}",
+                }
+            )
+        )
+
+        await super().setup(**kwargs)
+
+        self._logger.debug(f"Using Infrahub API at {config.SETTINGS.main.internal_address}")
         client = InfrahubClient(
             config=Config(address=config.SETTINGS.main.internal_address, retry_on_failure=True, log=self._logger)
         )

backend/infrahub/workflows/catalogue.py

@@ -1,59 +1,52 @@
+from .constants import WorkflowType
 from .models import WorkerPoolDefinition, WorkflowDefinition
 
-WORKER_POOL = WorkerPoolDefinition(
-    name="infrahub-internal", worker_type="infrahubasync", description="Pool for internal tasks"
+INFRAHUB_WORKER_POOL = WorkerPoolDefinition(
+    name="infrahub-worker", worker_type="infrahubasync", description="Default Pool for internal tasks"
 )
 
 WEBHOOK_SEND = WorkflowDefinition(
     name="webhook_send",
-    work_pool=WORKER_POOL,
+    type=WorkflowType.USER,
     module="infrahub.message_bus.operations.send.webhook",
     function="send_webhook",
 )
 
 TRANSFORM_JINJA2_RENDER = WorkflowDefinition(
     name="transform_render_jinja2_template",
-    work_pool=WORKER_POOL,
+    type=WorkflowType.USER,
     module="infrahub.message_bus.operations.transform.jinja",
     function="transform_render_jinja2_template",
 )
 
 ANONYMOUS_TELEMETRY_SEND = WorkflowDefinition(
     name="anonymous_telemetry_send",
-    work_pool=WORKER_POOL,
+    type=WorkflowType.INTERNAL,
     cron="0 2 * * *",
     module="infrahub.message_bus.operations.send.telemetry",
     function="send_telemetry_push",
 )
 
-DUMMY_FLOW = WorkflowDefinition(
-    name="dummy_flow",
-    work_pool=WORKER_POOL,
-    module="infrahub.tasks.dummy",
-    function="dummy_flow",
-)
-
 SCHEMA_APPLY_MIGRATION = WorkflowDefinition(
     name="schema_apply_migrations",
-    work_pool=WORKER_POOL,
+    type=WorkflowType.INTERNAL,
     module="infrahub.core.migrations.schema.tasks",
     function="schema_apply_migrations",
 )
 
 SCHEMA_VALIDATE_MIGRATION = WorkflowDefinition(
     name="schema_validate_migrations",
-    work_pool=WORKER_POOL,
+    type=WorkflowType.INTERNAL,
     module="infrahub.core.validators.tasks",
     function="schema_validate_migrations",
 )
 
-worker_pools = [WORKER_POOL]
+worker_pools = [INFRAHUB_WORKER_POOL]
 
 workflows = [
     WEBHOOK_SEND,
     TRANSFORM_JINJA2_RENDER,
     ANONYMOUS_TELEMETRY_SEND,
-    DUMMY_FLOW,
     SCHEMA_APPLY_MIGRATION,
     SCHEMA_VALIDATE_MIGRATION,
 ]