FlowerPower 0.9.13.1__py3-none-any.whl → 1.0.0b2__py3-none-any.whl

flowerpower/pipeline/job_queue.py

@@ -0,0 +1,505 @@
+# -*- coding: utf-8 -*-
+# pylint: disable=logging-fstring-interpolation
+# flake8: noqa: E501
+"""Pipeline Job Queue."""
+
+import datetime as dt
+from typing import Any, Callable
+from uuid import UUID
+
+from loguru import logger
+from rich import print as rprint
+
+# Import necessary config types
+from ..cfg import PipelineConfig, ProjectConfig
+from ..fs import AbstractFileSystem
+from ..utils.logging import setup_logging
+from ..job_queue import JobQueue
+from .registry import PipelineRegistry
+from .. import settings
+
+setup_logging()
+
+
+class PipelineJobQueue:
+    """Handles scheduling of pipeline runs via a configured worker backend."""
+
+    def __init__(
+        self,
+        project_cfg: ProjectConfig,
+        fs: AbstractFileSystem,
+        cfg_dir: str,
+        pipelines_dir: str,
+        job_queue_type: str | None = None,
+    ):
+        """Initialize PipelineJobQueue.
+
+        Args:
+            project_cfg: The project configuration object.
+            fs: The file system to use for file operations.
+            cfg_dir: The directory for configuration files.
+            pipelines_dir: The directory for pipeline files.
+            job_queue_type: The type of worker to use (e.g., 'rq', 'apscheduler'). If None, defaults to the project config.
+        """
+        self.project_cfg = project_cfg
+        self._fs = fs
+        self._cfg_dir = cfg_dir
+        self._pipelines_dir = pipelines_dir
+        self._job_queue_type = job_queue_type or project_cfg.job_queue.type
+        if not self._job_queue_type:
+            # Fallback or default if not specified in project config
+            self._job_queue_type = settings.DEFAULT_JOB_QUEUE
+            logger.warning(
+                f"Job queue type not specified in project config, defaulting to '{self._job_queue_type}'"
+            )
+
+    @property
+    def job_queue(self):
+        """
+        Lazily instantiate and cache a Job queue instance.
+        """
+        # Lazily instantiate worker using project_cfg attributes
+        logger.debug(
+            f"Instantiating worker of type: {self._job_queue_type} for project '{self.project_cfg.name}'"
+        )
+        # Pass the necessary parts of project_cfg to the Job queue
+        return JobQueue(
+            type=self._job_queue_type,
+            fs=self._fs,
+        )
+
+    def _get_schedule_ids(self) -> list[Any]:
+        """Get all schedules from the worker backend."""
+
+        with self.job_queue as worker:
+            logger.debug("Fetching schedules ids from worker")
+            return worker.schedule_ids
+
+    def run_job(
+        self,
+        run_func: Callable,
+        name: str,  # name: str,
+        inputs: dict | None = None,
+        final_vars: list | None = None,
+        config: dict | None = None,
+        cache: bool | dict = False,
+        executor_cfg: str | dict | Any | None = None,
+        with_adapter_cfg: dict | Any | None = None,
+        pipeline_adapter_cfg: dict | Any | None = None,
+        project_adapter_cfg: dict | Any | None = None,
+        adapter: dict[str, Any] | None = None,
+        reload: bool = False,
+        log_level: str | None = None,
+        **kwargs,
+    ) -> dict[str, Any]:
+        """
+        Add a job to run the pipeline immediately via the worker queue.
+
+        Args:
+            run_func (Callable): The function to execute in the worker (e.g., a configured PipelineRunner.run).
+            name (str): The name of the pipeline (used for logging).
+            inputs (dict | None): Inputs for the pipeline run.
+            final_vars (list | None): Final variables for the pipeline run.
+            config (dict | None): Hamilton driver config.
+            cache (bool | dict): Cache configuration.
+            executor_cfg (str | dict | ExecutorConfig | None): Executor configuration.
+            with_adapter_cfg (dict | WithAdapterConfig | None): Adapter configuration.
+            pipeline_adapter_cfg (dict | PipelineAdapterConfig | None): Pipeline adapter configuration.
+            project_adapter_cfg (dict | ProjectAdapterConfig | None): Project adapter configuration.
+            adapter (dict[str, Any] | None): Additional adapter configuration.
+            reload (bool): Whether to reload the pipeline module.
+            log_level (str | None): Log level for the run.
+            **kwargs: Additional keyword arguments passed directly to the worker's add_job method.
+
+        Returns:
+            dict[str, Any]: The result of the job execution.
+        """
+        logger.debug(f"Adding immediate job for pipeline: {name}")
+
+       
+        pipeline_run_args = {
+            # 'name' is not passed to run_func, it's part of the context already in PipelineRunner
+            "inputs": inputs,
+            "final_vars": final_vars,
+            "config": config,
+            "cache": cache,
+            "executor_cfg": executor_cfg,
+            "with_adapter_cfg": with_adapter_cfg,
+            "pipeline_adapter_cfg": pipeline_adapter_cfg,
+            "project_adapter_cfg": project_adapter_cfg,
+            "adapter": adapter,
+            "reload": reload,
+            "log_level": log_level,
+        }
+        pipeline_run_args = {
+            k: v for k, v in pipeline_run_args.items() if v is not None
+        }
+        logger.debug(
+            f"Resolved arguments for target run_func for job '{name}': {pipeline_run_args}"
+        )
+
+        with self.job_queue as worker:
+            res = worker.run_job(
+                func=run_func,
+                func_kwargs=pipeline_run_args,
+                **kwargs,
+            )
+
+        return res
+
+    def add_job(
+        self,
+        run_func: Callable,  # The actual function to run (e.g., PipelineRunner(...).run)
+        name: str,
+        inputs: dict | None = None,
+        final_vars: list | None = None,
+        config: dict | None = None,
+        cache: bool | dict = False,
+        executor_cfg: str | dict | Any | None = None,
+        with_adapter_cfg: dict | Any | None = None,
+        pipeline_adapter_cfg: dict | Any | None = None,
+        project_adapter_cfg: dict | Any | None = None,
+        adapter: dict[str, Any] | None = None,
+        result_ttl: int | dt.timedelta = 120,
+        run_at: dt.datetime | None = None,
+        run_in: float | dt.timedelta | None = None,
+        reload: bool = False,
+        log_level: str | None = None,
+        max_retries: int | None = None,
+        retry_delay: float | None = None,
+        jitter_factor: float | None = None,
+        retry_exceptions: tuple | list | None = None,
+        **kwargs,  # Allow other worker-specific args if needed
+    ) -> str | UUID:
+        """
+        Add a job to run the pipeline immediately via the worker queue, storing the result.
+
+        Executes the job immediately and returns the job id (UUID). The job result will be stored
+        by the worker backend for the given `result_ttl` and can be fetched using the job id.
+
+        Args:
+            run_func (Callable): The function to execute in the worker (e.g., a configured PipelineRunner.run).
+            name (str): The name of the pipeline (used for logging).
+            inputs (dict | None): Inputs for the pipeline run.
+            final_vars (list | None): Final variables for the pipeline run.
+            config (dict | None): Hamilton driver config.
+            cache (bool | dict): Cache configuration.
+            executor_cfg (str | dict | ExecutorConfig | None): Executor configuration.
+            with_adapter_cfg (dict | WithAdapterConfig | None): Adapter configuration.
+            pipeline_adapter_cfg (dict | PipelineAdapterConfig | None): Pipeline adapter configuration.
+            project_adapter_cfg (dict | ProjectAdapterConfig | None): Project adapter configuration.
+            adapter (dict[str, Any] | None): Additional adapter configuration.
+            reload (bool): Whether to reload the pipeline module.
+            log_level (str | None): Log level for the run.
+            result_ttl (int | dt.timedelta): How long the job result should be stored. Defaults to 0 (don't store).
+            run_at (dt.datetime | None): Optional datetime to run the job at.
+            run_in (float | dt.timedelta | None): Optional delay before running the job.
+            max_retries (int): Maximum number of retries for the job.
+            retry_delay (float): Delay between retries.
+            jitter_factor (float): Jitter factor for retry delay.
+            retry_exceptions (tuple): Exceptions that should trigger a retry.
+            **kwargs: Additional keyword arguments passed directly to the worker's add_job method.
+
+        Returns:
+            str | UUID: The ID of the added job.
+        """
+        logger.debug(f"Adding immediate job with result TTL for pipeline: {name}")
+
+        pipeline_run_args = {
+            "inputs": inputs,
+            "final_vars": final_vars,
+            "config": config,
+            "cache": cache,
+            "executor_cfg": executor_cfg,
+            "with_adapter_cfg": with_adapter_cfg,
+            "pipeline_adapter_cfg": pipeline_adapter_cfg,
+            "project_adapter_cfg": project_adapter_cfg,
+            "adapter": adapter,
+            "reload": reload,
+            "log_level": log_level,
+            "max_retries": max_retries,
+            "retry_delay": retry_delay,
+            "jitter_factor": jitter_factor,
+            "retry_exceptions": retry_exceptions,
+        }
+        pipeline_run_args = {
+            k: v for k, v in pipeline_run_args.items() if v is not None
+        }
+        logger.debug(
+            f"Resolved arguments for target run_func for job (TTL) '{name}': {pipeline_run_args}"
+        )
+
+        with self.job_queue as worker:
+            job_id = worker.add_job(
+                func=run_func,
+                func_kwargs=pipeline_run_args,
+                result_ttl=result_ttl,
+                run_at=run_at,
+                run_in=run_in,
+                **kwargs,
+            )
+        rprint(
+            f"✅ Successfully added job for "
+            f"[blue]{self.project_cfg.name}.{name}[/blue] with ID [green]{job_id}[/green]"
+            f" and result TTL of {result_ttl} seconds."
+        )
+        return job_id
+
+    # --- End Moved from PipelineManager ---
+
+    def schedule(
+        self,
+        run_func: Callable,
+        pipeline_cfg: PipelineConfig,
+        # --- Run Parameters (passed to run_func) ---
+        inputs: dict | None = None,
+        final_vars: list | None = None,
+        config: dict | None = None,  # Driver config
+        cache: bool | dict = False,
+        executor_cfg: str | dict | Any | None = None,
+        with_adapter_cfg: dict | Any | None = None,
+        pipeline_adapter_cfg: dict | Any | None = None,
+        project_adapter_cfg: dict | Any | None = None,
+        adapter: dict[str, Any] | None = None,
+        reload: bool = False,
+        log_level: str | None = None,
+        max_retries: int | None = None,
+        retry_delay: float | None = None,
+        jitter_factor: float | None = None,
+        retry_exceptions: tuple = (Exception,),
+        # --- Schedule Parameters (passed to worker.add_schedule) ---
+        cron: str | dict[str, str | int] | None = None,
+        interval: int | str | dict[str, str | int] | None = None,
+        date: dt.datetime | None = None,
+        overwrite: bool = False,
+        schedule_id: str | None = None,
+        **kwargs,
+    ) -> str | UUID:
+        """
+        Schedule a pipeline for execution using the configured worker.
+
+        Args:
+            run_func (Callable): The function to execute in the worker.
+            pipeline_cfg (PipelineConfig): The pipeline configuration object.
+            inputs (dict | None): Inputs for the pipeline run (overrides config).
+            final_vars (list | None): Final variables for the pipeline run (overrides config).
+            config (dict | None): Hamilton driver config (overrides config).
+            cache (bool | dict): Cache configuration (overrides config).
+            executor_cfg (str | dict | ExecutorConfig | None): Executor configuration (overrides config).
+            with_adapter_cfg (dict | WithAdapterConfig | None): Adapter configuration (overrides config).
+            pipeline_adapter_cfg (dict | PipelineAdapterConfig | None): Pipeline adapter configuration (overrides config).
+            project_adapter_cfg (dict | ProjectAdapterConfig | None): Project adapter configuration (overrides config).
+            adapter (dict | None): Additional Hamilton adapters (overrides config).
+            reload (bool): Whether to reload module (overrides config).
+            log_level (str | None): Log level for the run (overrides config).
+            max_retries (int): Maximum number of retries for the job.
+            retry_delay (float): Delay between retries.
+            jitter_factor (float): Jitter factor for retry delay.
+            retry_exceptions (tuple): Exceptions that should trigger a retry.
+            cron (str | dict | None): Cron expression or dict for cron trigger.
+            interval (int | str | dict | None): Interval in seconds or dict for interval trigger.
+            date (dt.datetime | None): Date for date trigger.
+            overwrite (bool): If True and id_ is None, generates ID '{name}-1', potentially overwriting.
+            schedule_id (str | None): Optional ID for the schedule. If None, generates a new ID.
+            **kwargs: Additional keyword arguments passed to the worker's add_schedule method,
+                For RQ this includes:
+                    - repeat: Repeat count (int or dict)
+                    - result_ttl: Time to live for the job result (float or timedelta)
+                    - ttl: Time to live for the job (float or timedelta)
+                    - use_local_time_zone: Whether to use local time zone for scheduling (bool)
+                For APScheduler, this includes:
+                    - misfire_grace_time: Grace time for misfires (timedelta)
+                    - coalesce: Whether to coalesce jobs (bool)
+                    - max_running_jobs: Maximum instances of the job (int)
+                    - max_jitter: Maximum jitter for scheduling (int)
+                    - conflict_policy: Policy for conflicting jobs (str)
+                    - paused: Whether to pause the job (bool)
+
+
+        Returns:
+            str | UUID: The ID of the scheduled pipeline.
+
+        Raises:
+            ValueError: If trigger_type is invalid or required args are missing.
+            Exception: Can raise exceptions from the worker backend.
+        """
+
+        project_name = self.project_cfg.name
+        name = pipeline_cfg.name
+        logger.debug(
+            f"Attempting to schedule pipeline: {project_name}.{name} with id: {schedule_id}"
+        )
+
+        # --- Resolve Parameters using pipeline_cfg for defaults ---
+        schedule_cfg = pipeline_cfg.schedule
+        #run_cfg = pipeline_cfg.run
+
+        pipeline_run_args = {
+            "inputs": inputs,
+            "final_vars": final_vars,
+            "config": config,
+            "cache": cache,
+            "executor_cfg": executor_cfg,
+            "with_adapter_cfg": with_adapter_cfg,
+            "pipeline_adapter_cfg": pipeline_adapter_cfg,
+            "project_adapter_cfg": project_adapter_cfg,
+            "adapter": adapter,
+            "reload": reload,
+            "log_level": log_level,
+            "max_retries": max_retries,
+            "retry_delay": retry_delay,
+            "jitter_factor": jitter_factor,
+            "retry_exceptions": retry_exceptions,
+        }
+        pipeline_run_args = {
+            k: v for k, v in pipeline_run_args.items() if v is not None
+        }
+        logger.debug(f"Resolved run_kwargs for '{name}': {pipeline_run_args}")
+
+        cron = cron if cron is not None else schedule_cfg.cron
+        interval = interval if interval is not None else schedule_cfg.interval
+        date = date if date is not None else schedule_cfg.date
+        logger.debug(
+            f"Resolved schedule parameters for '{name}': cron={cron}, interval={interval}, date={date}"
+        )
+
+        # --- Generate ID if not provided ---
+        # (Keep _generate_id function as is, it uses self._get_schedules())
+        def _generate_id(
+            pipeline_name: str, explicit_id: str | None, force_overwrite_base: bool
+        ) -> str:
+            if explicit_id:
+                logger.debug(f"Using explicit schedule ID: {explicit_id}")
+                return explicit_id
+
+            base_id = f"{pipeline_name}-1"
+
+            if force_overwrite_base:
+                logger.debug(f"Overwrite specified, using base ID: {base_id}")
+                return base_id
+
+            try:
+                existing_ids = self._get_schedule_ids()
+                logger.debug(f"Existing schedule IDs: {existing_ids}")
+
+                if not any(
+                    id_val.startswith(f"{pipeline_name}-") for id_val in existing_ids
+                ):
+                    logger.debug(
+                        f"No existing schedules found for '{pipeline_name}', using base ID: {base_id}"
+                    )
+                    return base_id
+
+                # Find highest existing number for this pipeline name
+                max_num = 0
+                for id_val in existing_ids:
+                    if id_val.startswith(f"{pipeline_name}-"):
+                        try:
+                            num_part = id_val.split("-")[-1]
+                            num = int(num_part)
+                            if num > max_num:
+                                max_num = num
+                        except (ValueError, IndexError):
+                            logger.warning(
+                                f"Could not parse number from existing schedule ID: {id_val}"
+                            )
+                            continue  # Skip malformed IDs
+
+                new_id = f"{pipeline_name}-{max_num + 1}"
+                logger.debug(f"Generated new schedule ID: {new_id}")
+                return new_id
+
+            except Exception as e:
+                logger.error(
+                    f"Error getting existing schedules to generate ID: {e}. Falling back to base ID: {base_id}"
+                )
+                # Fallback in case of error fetching schedules
+                return base_id
+
+        schedule_id = _generate_id(name, schedule_id, overwrite)
+
+        # --- Add Schedule via Job queue ---
+        try:
+            with self.job_queue as worker:
+                # Job queue is now responsible for creating the trigger object
+                # Pass trigger type and kwargs directly
+                added_id = worker.add_schedule(
+                    func=run_func,
+                    func_kwargs=pipeline_run_args,  # Pass resolved run parameters
+                    cron=cron,
+                    interval=interval,
+                    date=date,
+                    schedule_id=schedule_id,
+                    **kwargs,  # Pass resolved schedule run parameters
+                )
+                logger.info(
+                    f"✅ Successfully scheduled job for "
+                    f"[blue]{project_name}.{name}[/blue] with ID [green]{added_id}[/green]"
+                )
+                return added_id
+        except Exception as e:
+            logger.error(
+                f"Failed to add schedule '{schedule_id}' for pipeline '{name}': {e}"
+            )
+            raise
+
+    # --- schedule_all method removed ---
+    # PipelineManager will be responsible for iterating and calling schedule()
+
+    def schedule_all(self, registry: PipelineRegistry, **kwargs):
+        """
+        Schedule all pipelines found by the registry.
+
+        Args:
+            **kwargs: Arguments passed directly to the `schedule` method for each pipeline.
+                      Note: Pipeline-specific configurations will still take precedence for
+                      defaults if not overridden by kwargs.
+        """
+        try:
+            registry = self._get_registry_func()
+            names = registry._get_names()  # Use registry to find pipelines
+            if not names:
+                logger.info("[yellow]No pipelines found to schedule.[/yellow]")
+                return
+
+            logger.info(f"Attempting to schedule {len(names)} pipelines...")
+            scheduled_ids = []
+            errors = []
+            for name in names:
+                try:
+                    # Load config specifically for this pipeline to get defaults
+                    # Note: schedule() will load it again, potential optimization later
+                    cfg = self._load_config_func(name=name)
+                    if (
+                        not cfg
+                        or not cfg.pipeline
+                        or not cfg.pipeline.schedule
+                        or not cfg.pipeline.schedule.enabled
+                    ):
+                        logger.info(
+                            f"🟡 Skipping schedule for [cyan]{name}[/cyan]: Not configured or disabled in config."
+                        )
+                        continue
+
+                    logger.info(f"Scheduling [cyan]{name}[/cyan]...")
+                    # Pass kwargs, allowing overrides of config defaults
+                    schedule_id = self.schedule(name=name, **kwargs)
+                    scheduled_ids.append(schedule_id)
+                except Exception as e:
+                    logger.error(f"Failed to schedule pipeline '{name}': {e}")
+                    errors.append(name)
+                    logger.error(f"❌ Error scheduling [cyan]{name}[/cyan]: {e}")
+
+            if errors:
+                logger.error(
+                    f"\n[bold red]Finished scheduling with errors for: {', '.join(errors)}[/bold red]"
+                )
+            else:
+                logger.success(
+                    f"\n[bold green]Successfully scheduled {len(scheduled_ids)} pipelines.[/bold green]"
+                )
+
+        except Exception as e:
+            logger.error(
+                f"[bold red]An unexpected error occurred during schedule_all: {e}[/bold red]"
+            )