Integrate checkin worker (#466)

Author: toondaey

cognite/examples/unstable/extractors/simple_extractor/config/connection_config.yaml

@@ -1,10 +1,10 @@
 project: ${COGNITE_PROJECT}
 base_url: ${COGNITE_BASE_URL}
 integration:
-  external_id: "utils-test-keyvault-remote"
+  external_id: ${COGNITE_INTEGRATION_ID}
 authentication:
   type: "client-credentials"
   client_id: ${COGNITE_CLIENT_ID}
   client_secret: ${COGNITE_CLIENT_SECRET}
   token_url: ${COGNITE_TOKEN_URL}
-  scopes: ${COGNITE_BASE_URL}/.default
+  scopes: ${COGNITE_SCOPES}

cognite/examples/unstable/extractors/simple_extractor/main.py

@@ -2,9 +2,10 @@
 An example extractor that logs messages at various levels.
 """
 
-from cognite.extractorutils.unstable.configuration.models import ExtractorConfig
+from cognite.extractorutils.unstable.configuration.models import ExtractorConfig, IntervalConfig, TimeIntervalConfig
 from cognite.extractorutils.unstable.core.base import Extractor, StartupTask, TaskContext
 from cognite.extractorutils.unstable.core.runtime import Runtime
+from cognite.extractorutils.unstable.core.tasks import ScheduledTask
 
 
 class SimpleConfig(ExtractorConfig):
@@ -32,6 +33,13 @@ def __init_tasks__(self) -> None:
         Initializes and adds tasks to the extractor.
         """
         self.add_task(StartupTask(name="main_task", target=self.run_my_task))
+        self.add_task(
+            ScheduledTask(
+                name="scheduled_task",
+                target=self.scheduled_task,
+                schedule=IntervalConfig(type="interval", expression=TimeIntervalConfig("3s")),
+            )
+        )
 
     # example task that logs messages at different levels
     def run_my_task(self, ctx: TaskContext) -> None:
@@ -46,6 +54,18 @@ def run_my_task(self, ctx: TaskContext) -> None:
         ctx.warning("This is a warning message.")
         ctx.info("Test finished.")
 
+    def scheduled_task(self, ctx: TaskContext) -> None:
+        """
+        An example scheduled task that logs a message.
+
+        Args:
+            ctx: The context for the task execution, used for logging.
+        """
+        ctx.info("This is a scheduled task running.")
+        ctx.warning("This is a warning from the scheduled task.")
+        ctx.debug("Debugging the scheduled task execution.")
+        ctx.error("This is an error message from the scheduled task.")
+
     # add more tasks as needed
 
 

cognite/extractorutils/unstable/configuration/models.py

@@ -433,6 +433,7 @@ class ExtractorConfig(ConfigModel):
     """
 
     log_handlers: list[LogHandlerConfig] = Field(default_factory=_log_handler_default)
+    retry_startup: bool = True
 
 
 ConfigType = TypeVar("ConfigType", bound=ExtractorConfig)

cognite/extractorutils/unstable/core/base.py

@@ -74,15 +74,12 @@ def my_task_function(self, task_context: TaskContext) -> None:
     ExtractorInfo,
     StartupRequest,
     TaskType,
-    TaskUpdate,
-)
-from cognite.extractorutils.unstable.core._dto import (
-    Error as DtoError,
 )
 from cognite.extractorutils.unstable.core._dto import (
     Task as DtoTask,
 )
 from cognite.extractorutils.unstable.core._messaging import RuntimeMessage
+from cognite.extractorutils.unstable.core.checkin_worker import CheckinWorker
 from cognite.extractorutils.unstable.core.errors import Error, ErrorLevel
 from cognite.extractorutils.unstable.core.logger import CogniteLogger
 from cognite.extractorutils.unstable.core.restart_policy import WHEN_CONTINUOUS_TASKS_CRASHES, RestartPolicy
@@ -142,8 +139,11 @@ class Extractor(Generic[ConfigType], CogniteLogger):
 
     RESTART_POLICY: RestartPolicy = WHEN_CONTINUOUS_TASKS_CRASHES
 
-    def __init__(self, config: FullConfig[ConfigType]) -> None:
+    cancellation_token: CancellationToken
+
+    def __init__(self, config: FullConfig[ConfigType], checkin_worker: CheckinWorker) -> None:
         self._logger = logging.getLogger(f"{self.EXTERNAL_ID}.main")
+        self._checkin_worker = checkin_worker
 
         self.cancellation_token = CancellationToken()
         self.cancellation_token.cancel_on_interrupt()
@@ -161,8 +161,6 @@ def __init__(self, config: FullConfig[ConfigType]) -> None:
         self._scheduler = TaskScheduler(self.cancellation_token.create_child_token())
 
         self._tasks: list[Task] = []
-        self._task_updates: list[TaskUpdate] = []
-        self._errors: dict[str, Error] = {}
         self._start_time: datetime
 
         self.__init_tasks__()
@@ -253,43 +251,6 @@ def _attach_runtime_controls(self, *, cancel_event: MpEvent, message_queue: Queu
         self._set_runtime_message_queue(message_queue)
         self._setup_cancellation_watcher(cancel_event)
 
-    def _checkin(self) -> None:
-        with self._checkin_lock:
-            task_updates = [t.model_dump(mode="json") for t in self._task_updates]
-            self._task_updates.clear()
-
-            error_updates = [
-                DtoError(
-                    external_id=e.external_id,
-                    level=e.level,
-                    description=e.description,
-                    details=e.details,
-                    start_time=e.start_time,
-                    end_time=e.end_time,
-                    task=e._task_name if e._task_name is not None else None,
-                ).model_dump(mode="json")
-                for e in self._errors.values()
-            ]
-            self._errors.clear()
-
-        res = self.cognite_client.post(
-            f"/api/v1/projects/{self.cognite_client.config.project}/integrations/checkin",
-            json={
-                "externalId": self.connection_config.integration.external_id,
-                "taskEvents": task_updates,
-                "errors": error_updates,
-            },
-            headers={"cdf-version": "alpha"},
-        )
-        new_config_revision = res.json().get("lastConfigRevision")
-
-        if (
-            new_config_revision
-            and self.current_config_revision != "local"
-            and new_config_revision > self.current_config_revision
-        ):
-            self.restart()
-
     def _get_startup_request(self) -> StartupRequest:
         return StartupRequest(
             external_id=self.connection_config.integration.external_id,
@@ -310,17 +271,10 @@ def _get_startup_request(self) -> StartupRequest:
         )
 
     def _run_checkin(self) -> None:
-        while not self.cancellation_token.is_cancelled:
-            try:
-                self._logger.debug("Running checkin")
-                self._checkin()
-            except Exception:
-                self._logger.exception("Error during checkin")
-            self.cancellation_token.wait(10)
+        self._checkin_worker.run_periodic_checkin(self.cancellation_token, self._get_startup_request())
 
     def _report_error(self, error: Error) -> None:
-        with self._checkin_lock:
-            self._errors[error.external_id] = error
+        self._checkin_worker.report_error(error)
 
     def _new_error(
         self,
@@ -348,8 +302,8 @@ def restart(self) -> None:
         self.cancellation_token.cancel()
 
     @classmethod
-    def _init_from_runtime(cls, config: FullConfig[ConfigType]) -> Self:
-        return cls(config)
+    def _init_from_runtime(cls, config: FullConfig[ConfigType], checkin_worker: CheckinWorker) -> Self:
+        return cls(config, checkin_worker)
 
     def add_task(self, task: Task) -> None:
         """
@@ -368,10 +322,7 @@ def run_task(task_context: TaskContext) -> None:
             A wrapped version of the task's target, with tracking and error handling.
             """
             # Record a task start
-            with self._checkin_lock:
-                self._task_updates.append(
-                    TaskUpdate(type="started", name=task.name, timestamp=now()),
-                )
+            self._checkin_worker.report_task_start(name=task.name, timestamp=now())
 
             try:
                 # Run task
@@ -390,10 +341,7 @@ def run_task(task_context: TaskContext) -> None:
 
             finally:
                 # Record task end
-                with self._checkin_lock:
-                    self._task_updates.append(
-                        TaskUpdate(type="ended", name=task.name, timestamp=now()),
-                    )
+                self._checkin_worker.report_task_end(name=task.name, timestamp=now())
 
         task.target = run_task
         self._tasks.append(task)
@@ -411,13 +359,6 @@ def run_task(task_context: TaskContext) -> None:
                     ),
                 )
 
-    def _report_extractor_info(self) -> None:
-        self.cognite_client.post(
-            f"/api/v1/projects/{self.cognite_client.config.project}/integrations/extractorinfo",
-            json=self._get_startup_request().model_dump(mode="json"),
-            headers={"cdf-version": "alpha"},
-        )
-
     def start(self) -> None:
         """
         Start the extractor.
@@ -427,7 +368,6 @@ def start(self) -> None:
         """
         self._setup_logging()
         self._start_time = datetime.now(tz=timezone.utc)
-        self._report_extractor_info()
         Thread(target=self._run_checkin, name="ExtractorCheckin", daemon=True).start()
 
     def stop(self) -> None:
@@ -456,10 +396,7 @@ def __exit__(
         Stop the extractor when exiting the context manager.
         """
         self.stop()
-        with self._checkin_lock:
-            self._checkin()
-
-        self._logger.info("Shutting down extractor")
+        self._checkin_worker.flush(self.cancellation_token)
         return exc_val is None
 
     def run(self) -> None:

cognite/extractorutils/unstable/core/checkin_worker.py

@@ -87,7 +87,7 @@ def __init__(
         self._retry_startup: bool = retry_startup
         self._has_reported_startup: bool = False
         self._active_revision: ConfigRevision = active_revision
-        self._errors: dict[str, DtoError] = {}
+        self._errors: dict[str, Error] = {}
         self._task_updates: list[TaskUpdate] = []
 
     @property
@@ -127,7 +127,7 @@ def run_periodic_checkin(
             self._logger.debug("Running periodic check-in with interval %.2f seconds", report_interval)
             self.flush(cancellation_token)
             self._logger.debug(f"Check-in worker finished check-in, sleeping for {report_interval:.2f} seconds")
-            sleep(report_interval)
+            cancellation_token.wait(report_interval)
 
     def _run_startup_report(
         self, cancellation_token: CancellationToken, startup_request: StartupRequest, interval: float | None = None
@@ -168,9 +168,7 @@ def _handle_checkin_response(self, response: JSONType) -> None:
                     "and configured to use remote config for the new config to take effect.",
                 )
             elif self._active_revision < checkin_response.last_config_revision:
-                self._logger.info(
-                    f"Remote config changed from {self._active_revision} to {checkin_response.last_config_revision}."
-                )
+                self._active_revision = checkin_response.last_config_revision
                 self._on_revision_change(checkin_response.last_config_revision)
 
     def flush(self, cancellation_token: CancellationToken) -> None:
@@ -197,7 +195,7 @@ def report_checkin(self, cancellation_token: CancellationToken) -> None:
         """
         with self._lock:
             if not self._has_reported_startup:
-                new_errors = [error for error in self._errors.values() if error.task is None]
+                new_errors = [error for error in self._errors.values() if error._task_name is None]
                 if len(new_errors) == 0:
                     self._logger.info("No startup request has been reported yet, skipping check-in.")
                     return
@@ -225,11 +223,8 @@ def report_checkin(self, cancellation_token: CancellationToken) -> None:
                 task_updates_to_write = task_updates
                 task_updates = []
                 self.try_write_checkin(
-                    CheckinRequest(
-                        external_id=self._integration,
-                        errors=errors_to_write if len(errors_to_write) > 0 else None,
-                        task_events=task_updates_to_write if len(task_updates_to_write) > 0 else None,
-                    ),
+                    errors_to_write,
+                    task_updates_to_write,
                 )
                 break
 
@@ -264,11 +259,8 @@ def report_checkin(self, cancellation_token: CancellationToken) -> None:
             if tasks_idx > 0:
                 task_updates = task_updates[tasks_idx:]
             self.try_write_checkin(
-                CheckinRequest(
-                    external_id=self._integration,
-                    errors=errors_to_write if len(errors_to_write) > 0 else None,
-                    task_events=task_updates_to_write if len(task_updates_to_write) > 0 else None,
-                )
+                errors_to_write,
+                task_updates_to_write,
             )
             if errs_idx == 0 and tasks_idx == 0:
                 self._logger.debug("Check-in worker finished writing check-in.")
@@ -278,15 +270,21 @@ def report_checkin(self, cancellation_token: CancellationToken) -> None:
             self._logger.debug("Extractor was stopped during check-in, requeuing remaining errors and task updates.")
             self._requeue_checkin(new_errors, task_updates)
 
-    def try_write_checkin(self, checkin_request: CheckinRequest) -> None:
+    def try_write_checkin(self, errors: list[Error], task_updates: list[TaskUpdate]) -> None:
         """
         We try to write a check-in.
 
         This will try to write a check in to integrations.
 
         Arguments:
-        checkin_request: The check-in request to write.
+        errors(list[Error]): The errors to write.
+        task_updates(list[TaskUpdate]): The task updates to write.
         """
+        checkin_request = CheckinRequest(
+            external_id=self._integration,
+            errors=list(map(DtoError.from_internal, errors)) if len(errors) > 0 else None,
+            task_events=task_updates if len(task_updates) > 0 else None,
+        )
         should_requeue = self._wrap_checkin_like_request(
             lambda: self._cognite_client.post(
                 f"/api/v1/projects/{self._cognite_client.config.project}/integrations/checkin",
@@ -296,7 +294,7 @@ def try_write_checkin(self, checkin_request: CheckinRequest) -> None:
         )
 
         if should_requeue:
-            self._requeue_checkin(checkin_request.errors, checkin_request.task_events)
+            self._requeue_checkin(errors, checkin_request.task_events)
 
     def report_error(self, error: Error) -> None:
         """
@@ -307,7 +305,7 @@ def report_error(self, error: Error) -> None:
         """
         with self._lock:
             if error.external_id not in self._errors:
-                self._errors[error.external_id] = DtoError.from_internal(error)
+                self._errors[error.external_id] = error
             else:
                 self._logger.warning(f"Error {error.external_id} already reported, skipping re-reporting.")
 
@@ -335,7 +333,7 @@ def report_task_end(self, name: str, message: MessageType | None = None, timesta
                 TaskUpdate(type="ended", name=name, timestamp=timestamp or (int(now() * 1000)), message=message)
             )
 
-    def _requeue_checkin(self, errors: list[DtoError] | None, task_updates: list[TaskUpdate] | None) -> None:
+    def _requeue_checkin(self, errors: list[Error] | None, task_updates: list[TaskUpdate] | None) -> None:
         with self._lock:
             for error in errors or []:
                 if error.external_id not in self._errors:
@@ -373,6 +371,7 @@ def _wrap_checkin_like_request(self, request: Callable[[], Response]) -> bool:
 
             return True
         except Exception as e:
+            self._logger.critical(f"Extractor could not connect to CDF {e!s}")
             self._on_fatal_error(e)
             return True