cognite-neat 0.104.0__py3-none-any.whl → 0.105.0__py3-none-any.whl

cognite/neat/_workflows/steps_registry.py

@@ -1,218 +0,0 @@
-import importlib
-import inspect
-import logging
-import os
-import sys
-import types
-from pathlib import Path
-from typing import Any
-
-from pydantic import BaseModel
-
-# steps
-import cognite.neat._workflows.steps.lib.current
-import cognite.neat._workflows.steps.lib.io
-from cognite.neat._app.monitoring.metrics import NeatMetricsCollector
-from cognite.neat._config import Config
-from cognite.neat._issues.errors import WorkflowConfigurationNotSetError, WorkFlowMissingDataError
-from cognite.neat._workflows.model import FlowMessage, WorkflowConfigs
-from cognite.neat._workflows.steps.step_model import Configurable, DataContract, Step
-
-
-class StepMetadata(BaseModel):
-    name: str
-    description: str = ""
-    category: str = "default"  # defines the category the step belongs to (e.g. "data", "model", "test")
-    scope: str = "global"  # defines the scope of the step (e.g. "global", local to a specific workflow)
-    input: list[str]
-    output: list[str]
-    configurables: list[Configurable] = []
-    version: str = "1.0.0"  # version of the step. All alpha versions considered as experimental
-    docs_url: str = ""  # url to the extended documentation of the step
-    source: str = (
-        "cognite"  # source of the step , can be source identifier or url , for instance github url for instance.
-    )
-
-
-class StepsRegistry:
-    def __init__(self, config: Config):
-        self.config = config
-        self._step_classes: list[type[Step]] = []
-        self.user_steps_path: Path = config.data_store_path / "steps"
-        self.data_store_path: str = str(config.data_store_path)
-        self.categorized_steps: dict[str, set] = {
-            "legacy": set(),
-            "current": set(),
-            "io": set(),
-        }
-
-    def load_step_classes(self):
-        if self._step_classes:
-            # classes already loaded - no need to reload
-            return
-        for name, step_cls in inspect.getmembers(cognite.neat._workflows.steps.lib.current):
-            if inspect.isclass(step_cls):
-                logging.info(f"Loading NEAT step {name}")
-                self._step_classes.append(step_cls)
-                self.categorized_steps["current"].add(name)
-        for name, step_cls in inspect.getmembers(cognite.neat._workflows.steps.lib.io):
-            if inspect.isclass(step_cls):
-                logging.info(f"Loading NEAT step {name}")
-                self._step_classes.append(step_cls)
-                self.categorized_steps["io"].add(name)
-        sys.path.append(str(Path(self.data_store_path) / "workflows"))
-        try:
-            if self.user_steps_path:
-                sys.path.append(str(self.user_steps_path))
-                self.load_custom_step_classes(self.user_steps_path, scope="user_global")
-        except Exception as e:
-            logging.info(f"No user defined modules provided in {self.user_steps_path}. Error: {e}")
-
-    def load_workflow_step_classes(self, workflow_name: str):
-        if not self.data_store_path:
-            raise WorkflowConfigurationNotSetError("data_store_path")
-        workflow_steps_path = Path(self.data_store_path) / "workflows" / workflow_name
-        if workflow_steps_path.exists():
-            self.load_custom_step_classes(workflow_steps_path, scope="workflow")
-
-    def load_custom_step_classes(self, custom_steps_path: Path, scope: str = "global"):
-        for step_module_name in os.listdir(custom_steps_path):
-            step_module_path = custom_steps_path / Path(step_module_name)
-            if step_module_name.startswith("__") or (
-                step_module_path.is_file() and not step_module_name.endswith(".py")
-            ):
-                continue
-            if scope == "user_global":
-                full_module_name = step_module_name.replace(".py", "")
-            else:
-                full_module_name = custom_steps_path.name + "." + step_module_name.replace(".py", "")
-            logging.info(f"Loading user defined step module {full_module_name}")
-            if full_module_name in sys.modules:
-                logging.info(f"Reloading existing workflow module {full_module_name}")
-                steps_module = importlib.reload(sys.modules[full_module_name])
-            else:
-                steps_module = importlib.import_module(full_module_name)
-                logging.info(f"Loading user defined step from {steps_module}")
-            for name, step_cls in inspect.getmembers(steps_module):
-                base_class = getattr(step_cls, "__bases__", None)
-                if (
-                    name.startswith("__")
-                    or base_class is None
-                    or len(base_class) == 0
-                    or base_class[0].__name__ != "Step"
-                ):
-                    continue
-                logging.info(f"Loading user defined step {name} from {steps_module}")
-                if inspect.isclass(step_cls):
-                    step_cls.scope = scope
-                    is_new_class = True
-                    for i, step in enumerate(self._step_classes):
-                        if step.__name__ == step_cls.__name__:
-                            logging.info(f"Reloading user defined step {name} class {step_cls.__name__}")
-                            self._step_classes[i] = step_cls
-                            is_new_class = False
-                    if is_new_class:
-                        logging.info(f"Loading NEW user defined step {name} class {step_cls.__name__}")
-                        self._step_classes.append(step_cls)
-
-    def run_step(
-        self,
-        step_name: str,
-        flow_context: dict[str, DataContract],
-        step_configs: dict[str, Any],
-        metrics: NeatMetricsCollector | None = None,
-        workflow_configs: WorkflowConfigs | None = None,
-        workflow_id: str = "",
-        workflow_run_id: str = "",
-        step_complex_configs: dict[str, Any] | None = None,
-    ) -> DataContract | tuple[FlowMessage, DataContract] | FlowMessage:
-        if step_complex_configs is None:
-            step_complex_configs = {}
-        for step_cls in self._step_classes:
-            if step_cls.__name__ == step_name:
-                step_obj: Step = step_cls(config=self.config)
-                step_obj.configure(step_configs, step_complex_configs)
-                step_obj.set_flow_context(flow_context)
-                step_obj.set_metrics(metrics)
-                step_obj.set_workflow_configs(workflow_configs)
-                step_obj.set_workflow_metadata(workflow_id, workflow_run_id)
-                signature = inspect.signature(step_obj.run)
-                parameters = signature.parameters
-                is_valid = True
-                input_data = []
-                missing_data = []
-                for parameter in parameters.values():
-                    try:
-                        if parameter.annotation is FlowMessage:
-                            input_data.append(flow_context["FlowMessage"])
-                        else:
-                            if isinstance(parameter.annotation, types.UnionType):
-                                for param in parameter.annotation.__args__:
-                                    if (
-                                        param.__name__ != "_empty"
-                                        and param.__name__ in flow_context
-                                        and flow_context[param.__name__] is not None
-                                    ):
-                                        input_data.append(flow_context[param.__name__])
-                                        break  # Only one variable can be used as input
-                            else:
-                                input_data.append(flow_context[parameter.annotation.__name__])
-                    except KeyError:
-                        is_valid = False
-                        logging.error(f"Missing data for step {step_name} parameter {parameter.name}")
-                        missing_data.append(parameter.annotation.__name__)
-                        continue
-                if not is_valid:
-                    raise WorkFlowMissingDataError(step_name, frozenset(missing_data))
-                return step_obj.run(*input_data)
-        raise WorkFlowMissingDataError(step_name, frozenset({}))
-
-    def get_list_of_steps(self) -> list[StepMetadata]:
-        steps: list[StepMetadata] = []
-        for step_cls in self._step_classes:
-            try:
-                signature = inspect.signature(step_cls.run)
-                parameters = signature.parameters
-                input_data = []
-                output_data = []
-                for parameter in parameters.values():
-                    if isinstance(parameter.annotation, types.UnionType):
-                        for param in parameter.annotation.__args__:
-                            if param.__name__ != "_empty":
-                                input_data.append(param.__name__)
-                    elif parameter.annotation.__name__ != "_empty":
-                        input_data.append(parameter.annotation.__name__)
-                return_annotation = signature.return_annotation
-                if return_annotation:
-                    if isinstance(return_annotation, types.UnionType):
-                        for annotation in return_annotation.__args__:
-                            output_data.append(annotation.__name__)
-                    elif isinstance(return_annotation, tuple):
-                        for annotation in return_annotation:
-                            if isinstance(annotation, types.UnionType):
-                                for annotation_l2 in annotation.__args__:
-                                    output_data.append(annotation_l2.__name__)
-                            else:
-                                output_data.append(annotation.__name__)
-                    else:
-                        output_data.append(return_annotation.__name__)
-                steps.append(
-                    StepMetadata(
-                        name=step_cls.__name__,
-                        scope=step_cls.scope,
-                        input=input_data,
-                        description=step_cls.description,
-                        category=step_cls.category,
-                        output=output_data,
-                        configurables=step_cls.configurables,
-                        version=step_cls.version,
-                        source=step_cls.source,
-                        docs_url=step_cls.docs_url,
-                    )
-                )
-            except AttributeError as e:
-                logging.error(
-                    f"Step {type(step_cls).__name__} does not have a run method or types can't be infered.Error: {e}"
-                )
-
-        return steps

cognite/neat/_workflows/tasks.py

@@ -1,18 +0,0 @@
-from cognite.client import CogniteClient
-
-from cognite.neat._workflows.model import FlowMessage
-
-
-class WorkflowTaskBuilder:
-    """Collection of all base tasks for workflows.All tasks must run in the context of a workflow including threads."""
-
-    def __init__(self, cdf_client: CogniteClient | None, workflow_manager):
-        # TODO : figure out circular import and set type to WorkflowManager
-        self.cdf_client = cdf_client
-        self.workflow_manager = workflow_manager
-
-    def start_workflow_task(self, workflow_name: str, sync: bool, flow_message: FlowMessage | None):
-        """Call a workflow task from another workflow"""
-        return self.workflow_manager.start_workflow_instance(
-            workflow_name=workflow_name, flow_msg=flow_message, sync=sync
-        )

cognite/neat/_workflows/triggers.py

@@ -1,169 +0,0 @@
-import json
-import logging
-import threading
-import time
-
-import schedule
-from fastapi import Request
-
-from cognite.neat._workflows.manager import WorkflowManager
-from cognite.neat._workflows.model import FlowMessage, StepType, WorkflowState
-
-
-class TriggerManager:
-    """Triggers are the way to start a workflow. They are defined in the workflow definition file."""
-
-    def __init__(self, workflow_manager: WorkflowManager):
-        self.workflow_manager = workflow_manager
-        self.is_running = False
-        self.thread = None
-
-    def start_workflow_from_http_request(self, workflow_name: str, step_id: str, request: Request, body: bytes):
-        logging.info(f"New HTTP trigger request for workflow {workflow_name} step {step_id}")
-        headers = dict(request.headers)
-        logging.debug(f"Request headers = {headers}")
-        json_payload = None
-        try:
-            json_payload = json.loads(body)
-        except Exception as e:
-            logging.info(f"Error parsing json body {e}")
-        logging.debug(f"Request object = {json_payload}")
-
-        flow_msg = FlowMessage(payload=json_payload, headers=dict(headers))
-        start_status = self.workflow_manager.start_workflow_instance(
-            workflow_name=workflow_name, step_id=step_id, flow_msg=flow_msg
-        )
-        if start_status.is_success and start_status.workflow_instance:
-            return start_status.workflow_instance.flow_message
-
-    def resume_workflow_from_http_request(
-        self, workflow_name: str, step_id: str, instance_id: str, request: Request, body: bytes
-    ):
-        if instance_id != "default":
-            workflow = self.workflow_manager.get_workflow_instance(instance_id)
-        else:
-            returned = self.workflow_manager.get_workflow(workflow_name)
-            if returned is None:
-                return {"result": "Workflow instance not found"}
-            workflow = returned
-
-        json_payload = None
-        try:
-            json_payload = json.loads(body)
-        except ValueError as e:
-            logging.info(f"Error parsing json body {e}")
-        flow_msg = FlowMessage(payload=json_payload)
-        if workflow.state == WorkflowState.RUNNING_WAITING:
-            workflow.resume_workflow(flow_message=flow_msg, step_id=step_id)
-            return {"result": "Workflow instance resumed"}
-
-        return {"result": "Workflow instance not in RUNNING_WAITING state"}
-
-    def _start_scheduler_main_loop(self):
-        """Starts a scheduler main loop for the workflows
-
-        Parameters
-        ----------
-
-        """
-        logging.info("Starting scheduler main loop")
-        self.is_running = True
-
-        def main_loop():
-            while self.is_running:
-                schedule.run_pending()
-                time.sleep(1)
-            logging.info("Scheduler main loop stopped")
-
-        self.thread = threading.Thread(target=main_loop)
-        self.thread.start()
-
-    def stop_scheduler_main_loop(self):
-        """Stops a scheduler main loop for the workflows
-
-        Parameters
-        ----------
-
-        """
-        logging.info("Stopping scheduler main loop")
-        self.is_running = False
-        if self.thread:
-            self.thread.join()
-        logging.info("Scheduler main loop stopped")
-
-    def every_weekday_schedule(self, weekday_str: str):
-        return getattr(schedule.every(), weekday_str)
-
-    def start_time_schedulers(self):
-        """Starts a time scheduler for the workflows
-
-        Parameters
-        ----------
-
-        """
-        logging.info("Starting time trigger scheduler")
-        weekdays = ["monday", "tuesday", "wednesday", "thursday", "friday", "saturday", "sunday"]
-
-        trigger_steps_count = 0
-        for workflow in self.workflow_manager.workflow_registry.values():
-            for step in workflow.workflow_steps:
-                if step.trigger and step.stype == StepType.TIME_TRIGGER and step.enabled:
-                    # params
-                    every_str = step.params["interval"]
-                    logging.info(f"Starting time trigger scheduler for {workflow.name} {step.id} {every_str}")
-                    every_str = every_str.replace("every", "").strip()
-
-                    if "at" in every_str:  # "day at 10:30:00"
-                        every = every_str.split("at ")
-                        if len(every) == 2:
-                            interval_unit = every[0]
-                            interval_value = every[1]
-                            if "day" in interval_unit:
-                                if interval_unit.lower() in weekdays:
-                                    trigger_steps_count += 1
-                                    self.every_weekday_schedule(interval_unit.lower()).at(interval_value).do(
-                                        workflow.start, start_step_id=step.id
-                                    )
-                                else:
-                                    trigger_steps_count += 1
-                                    schedule.every().day.at(interval_value).do(workflow.start, start_step_id=step.id)
-                        else:
-                            logging.error(f"Invalid time trigger interval {every_str}")
-
-                    else:  # "5 minutes"
-                        every = every_str.split(" ")
-                        if len(every) == 2:
-                            interval_unit = every[1]
-                            interval_value = every[0]
-                            trigger_steps_count += 1
-                            if "minutes" in interval_unit:  # "day at 10:30:00" , "5 minutes"
-                                schedule.every(int(interval_value)).minutes.do(workflow.start, start_step_id=step.id)
-                            elif "hours" in interval_unit:
-                                schedule.every(int(interval_value)).hours.do(workflow.start, start_step_id=step.id)
-                            elif "days" in interval_unit:
-                                schedule.every(int(interval_value)).days.do(workflow.start, start_step_id=step.id)
-                            elif "seconds" in interval_unit:
-                                schedule.every(int(interval_value)).seconds.do(workflow.start, start_step_id=step.id)
-                            else:
-                                logging.error(f"Invalid time trigger interval {every_str}")
-                                trigger_steps_count -= 1
-                        else:
-                            logging.error(f"Invalid time trigger interval {every_str}")
-
-        if trigger_steps_count > 0 and not self.is_running:
-            self._start_scheduler_main_loop()
-            logging.info("ALl Time trigger scheduler started")
-        else:
-            logging.info("No Time trigger scheduler started")
-
-    def reload_all_triggers(self):
-        """Reloads all triggers
-
-        Parameters
-        ----------
-        """
-        logging.info("Reloading all triggers")
-        schedule.clear()
-        self.is_running = False
-        time.sleep(1)
-        self.start_time_schedulers()

cognite/neat/_workflows/utils.py

@@ -1,19 +0,0 @@
-import hashlib
-import uuid
-from datetime import datetime
-from pathlib import Path
-
-
-def get_iso8601_timestamp_now_unaware():
-    return datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")
-
-
-def get_file_hash(file_path: Path) -> str:
-    if isinstance(file_path, str):
-        file_path = Path(file_path)
-    return hashlib.md5(file_path.read_bytes()).hexdigest()
-
-
-def generate_run_id() -> str:
-    # Generate a random run guid
-    return str(uuid.uuid4())