src/starfish/data_factory/job_manager.py

import asyncio
import copy
import datetime
import hashlib
import json
import uuid
from asyncio import Queue
from typing import Any, Callable, Dict, List
import traceback

from starfish.common.logger import get_logger
from starfish.data_factory.config import PROGRESS_LOG_INTERVAL
from starfish.data_factory.constants import (
    IDX,
    RECORD_STATUS,
    RUN_MODE_DRY_RUN,
    STATUS_COMPLETED,
    STATUS_DUPLICATE,
    STATUS_FAILED,
    STATUS_FILTERED,
)
from starfish.data_factory.event_loop import run_in_event_loop
from starfish.data_factory.storage.base import Storage
from starfish.data_factory.storage.models import GenerationJob, Record
from starfish.data_factory.task_runner import TaskRunner
from starfish.data_factory.utils.data_class import FactoryJobConfig, FactoryMasterConfig
from starfish.data_factory.utils.errors import TimeoutErrorAsyncio
from starfish.data_factory.utils.state import MutableSharedState

logger = get_logger(__name__)


class JobManager:
    """Manages the execution of data generation jobs.

    This class handles the orchestration of tasks, including task execution,
    progress tracking, and result storage. It manages concurrency and provides
    mechanisms for job monitoring and control.

    Attributes:
        master_job_id (str): ID of the master job
        job_config (FactoryJobConfig): Configuration for the job
        storage (Storage): Storage instance for persisting results
        state (dict): Current state of the job
        semaphore (asyncio.Semaphore): Semaphore for concurrency control
        lock (asyncio.Lock): Lock for thread-safe operations
        task_runner (TaskRunner): Runner for executing tasks
        job_input_queue (Queue): Queue for input data
        job_output (Queue): Queue for output data
        completed_count (int): Count of completed tasks
        duplicate_count (int): Count of duplicate tasks
        filtered_count (int): Count of filtered tasks
        failed_count (int): Count of failed tasks
        total_count (int): Total number of tasks attempted
        job_run_stop_threshold (int): Threshold for stopping job
        active_operations (set): Set of active operations
        _progress_ticker_task (asyncio.Task): Task for progress logging
        dead_queue (Queue): Queue for failed tasks
        task_failure_count (dict): Track failure count per task
    """

    # ====================
    # Initialization
    # ====================
    def __init__(
        self, master_job_config: FactoryMasterConfig, state: MutableSharedState, storage: Storage, user_func: Callable, input_data_queue: Queue = None
    ):
        """Initialize the JobManager with job configuration and storage.

        Args:
            master_job_config (FactoryMasterConfig): Configuration for the master job
            state (MutableSharedState): Shared state object for job state management
            storage (Storage): Storage instance for persisting job results and metadata
            user_func (Callable): Function to execute for each task
            input_data_queue (Queue, optional): Queue for input data. Defaults to None.
        """
        self.master_job_id = master_job_config.master_job_id
        self.job_config = FactoryJobConfig(
            batch_size=master_job_config.batch_size,
            target_count=master_job_config.target_count,
            dead_queue_threshold=master_job_config.dead_queue_threshold,
            show_progress=master_job_config.show_progress,
            user_func=user_func,
            run_mode=master_job_config.run_mode,
            max_concurrency=master_job_config.max_concurrency,
            on_record_complete=master_job_config.on_record_complete,
            on_record_error=master_job_config.on_record_error,
            job_run_stop_threshold=master_job_config.job_run_stop_threshold,
        )
        self.storage = storage
        self.state = state
        self.task_runner = TaskRunner(timeout=master_job_config.task_runner_timeout)
        self.job_input_queue = input_data_queue if input_data_queue else Queue()
        self.job_output = Queue()
        self.prev_job = master_job_config.prev_job
        # Initialize counters
        self._initialize_counters()
        self.active_operations = set()
        self._progress_ticker_task = None
        self.execution_time = 0
        self.err_type_counter = {}
        self.dead_queue = Queue()  # Add dead queue for failed tasks
        self.task_failure_count = {}  # Track failure count per task

    def _initialize_counters(self):
        """Initialize all job counters."""
        self.completed_count = 0
        self.duplicate_count = 0
        self.filtered_count = 0
        self.failed_count = 0
        self.total_count = 0
        self.dead_queue_count = 0

    async def setup_input_output_queue(self):
        pass

    # ====================
    # Job Execution
    # ====================
    def run_orchestration(self):
        """Start the job orchestration process.

        This method initiates the main event loop for processing tasks, managing concurrency,
        and handling task completion. It runs until either the target count is reached or
        the stop threshold is triggered.
        """
        start_time = datetime.datetime.now(datetime.timezone.utc)
        run_in_event_loop(self._async_run_orchestration())
        self.execution_time = int((datetime.datetime.now(datetime.timezone.utc) - start_time).total_seconds())

    async def _async_run_orchestration(self):
        """Main asynchronous orchestration loop for the job.

        This method manages the core job execution loop, including:
        - Starting the progress ticker
        - Processing tasks from the input queue
        - Managing concurrency with semaphores
        - Handling task completion and cleanup
        """
        self._initialize_concurrency_controls()

        if self.job_config.show_progress:
            self._progress_ticker_task = asyncio.create_task(self._progress_ticker())

        self.running_tasks = set()

        try:
            if not self.job_input_queue.empty():
                await self._process_tasks()
        finally:
            await self._cleanup()

    def _initialize_concurrency_controls(self):
        """Initialize semaphore and lock for concurrency control."""
        if hasattr(self, "semaphore"):
            del self.semaphore
        if hasattr(self, "lock"):
            del self.lock
        self.semaphore = asyncio.Semaphore(self.job_config.max_concurrency)
        self.lock = asyncio.Lock()

    async def _process_tasks(self):
        """Process tasks from the input queue until stop condition is met."""
        while not self._is_job_to_stop():
            if not self.job_input_queue.empty():
                await self.semaphore.acquire()
                input_data = await self.job_input_queue.get()
                task = self._create_single_task(input_data)
                self.running_tasks.add(task)
                task.add_done_callback(self.running_tasks.discard)
            else:
                await asyncio.sleep(1)

    # ====================
    # Task Management
    # ====================
    def _create_single_task(self, input_data) -> asyncio.Task:
        """Create and manage a single task."""
        task = asyncio.create_task(self._run_single_task(input_data))
        asyncio.create_task(self._handle_task_completion(task))
        return task

    async def _run_single_task(self, input_data) -> List[Dict[str, Any]]:
        """Execute a single task with error handling."""
        output = []
        output_ref = []
        task_status = STATUS_COMPLETED
        err_output = {}
        input_data_idx = input_data.get(IDX, None)
        if input_data_idx == None:
            logger.warning(f"found an input_data without index ")

        try:
            output = await self.task_runner.run_task(self.job_config.user_func, input_data, input_data_idx)
            task_status = self._evaluate_task_output(output)
            output_ref = await self._save_record_data(copy.deepcopy(output), task_status, input_data)
        except (Exception, TimeoutErrorAsyncio) as e:
            task_status, err_output = self._handle_task_error(e)

        if task_status != STATUS_COMPLETED:
            await self._requeue_task(input_data, input_data_idx)

        return self._create_task_result(input_data_idx, task_status, output_ref, output, err_output)

    def _evaluate_task_output(self, output):
        """Evaluate task output and determine status."""
        hooks_output = [hook(output, self.state) for hook in self.job_config.on_record_complete]
        if STATUS_DUPLICATE in hooks_output:
            return STATUS_DUPLICATE
        if STATUS_FILTERED in hooks_output:
            return STATUS_FILTERED
        return STATUS_COMPLETED

    def _handle_task_error(self, error):
        """Handle task errors and update state."""
        err_str = str(error)
        # [-1]
        err_trace = traceback.format_exc().splitlines()
        logger.error(f"Error running task: {err_str}")

        for hook in self.job_config.on_record_error:
            hook(err_str, self.state)

        return STATUS_FAILED, {"err_str": err_str, "err_trace": err_trace}

    async def _requeue_task(self, input_data, input_data_idx):
        """Requeue a task that needs to be retried or move to dead queue if failed too many times."""
        task_key = str(input_data_idx)
        async with self.lock:
            self.task_failure_count[task_key] = self.task_failure_count.get(task_key, 0) + 1
            if self.task_failure_count[task_key] >= self.job_config.dead_queue_threshold:
                await self.dead_queue.put(input_data)
                self.dead_queue_count += 1
                logger.warning(f"Task {task_key} failed 3 times, moving to dead queue")
            else:
                await self.job_input_queue.put(input_data)
                logger.debug(f"Requeuing task {task_key} (failure count: {self.task_failure_count[task_key]})")

    def _create_task_result(self, input_data_idx, task_status, output_ref, output, err_output):
        """Create a standardized task result dictionary."""
        return {IDX: input_data_idx, RECORD_STATUS: task_status, "output_ref": output_ref, "output": output, "err": [err_output]}

    async def _handle_task_completion(self, task):
        """Handle task completion and update counters.

        Args:
            task (asyncio.Task): The completed task

        Updates:
            - Job counters (completed, failed, etc.)
            - Output queue
            - Semaphore
        """
        result = await task
        async with self.lock:
            await self.job_output.put(result)
            self.total_count += 1
            task_status = result.get(RECORD_STATUS)
            # Update counters based on task status
            if task_status == STATUS_COMPLETED:
                self.completed_count += 1
            elif task_status == STATUS_DUPLICATE:
                self.duplicate_count += 1
            elif task_status == STATUS_FILTERED:
                self.filtered_count += 1
            elif task_status == STATUS_FAILED:
                self.failed_count += 1
                # Safely extract error information with better type checking
                err_output = result.get("err", [{}])[0]  # Default to empty dict if no error
                err_str = err_output.get("err_str", "Unknown error").strip()  # Clean up whitespace

                # Normalize error string for consistent counting
                normalized_err = err_str.lower().strip()
                self.err_type_counter[normalized_err] = self.err_type_counter.get(normalized_err, 0) + 1
                # Optionally log the error count for this type
                logger.debug(f"Error type '{normalized_err}' count: {self.err_type_counter[normalized_err]}")
            # await self._update_progress(task_status, STATUS_MOJO_MAP[task_status])
            self.semaphore.release()

    # ====================
    # Cleanup & Utilities
    # ====================
    async def _cleanup(self):
        """Clean up resources after job completion."""
        await self._del_progress_ticker()
        await self._del_running_tasks()
        await self._cancel_operations()

    async def _del_running_tasks(self):
        """Cancel all running tasks."""
        for task in self.running_tasks:
            task.cancel()
        await asyncio.gather(*self.running_tasks, return_exceptions=True)

    async def _cancel_operations(self):
        """Cancel all active operations."""
        for task in self.active_operations:
            task.cancel()
        await asyncio.gather(*self.active_operations, return_exceptions=True)

    async def _create_execution_job(self, job_uuid: str, input_data: Dict[str, Any]):
        """Create and log a new execution job in storage.

        Args:
            job_uuid: Unique identifier for the execution job
            input_data: Dictionary containing input data for the job
        """
        logger.debug("\n3. Creating execution job...")
        input_data_str = json.dumps(input_data, sort_keys=True) if isinstance(input_data, dict) else str(input_data)
        # input_data_str = json.dumps(input_data)
        self.job = GenerationJob(
            job_id=job_uuid,
            master_job_id=self.master_job_id,
            status="pending",
            worker_id="test-worker-1",
            run_config=input_data_str,
            run_config_hash=hashlib.sha256(input_data_str.encode()).hexdigest(),
        )
        await self.storage.log_execution_job_start(self.job)
        logger.debug(f"  - Created execution job: {job_uuid}")

    async def _save_record_data(self, *args, **kwargs):
        task = asyncio.create_task(self._job_save_record_data(*args, **kwargs))
        self.active_operations.add(task)
        try:
            return await task
        finally:
            self.active_operations.discard(task)

    async def _job_save_record_data(self, records, task_status: str, input_data: Dict[str, Any]) -> List[str]:
        output_ref_list = []
        if self.job_config.run_mode == RUN_MODE_DRY_RUN:
            return output_ref_list
        logger.debug("\n5. Saving record data...")
        storage_class_name = self.storage.__class__.__name__
        if storage_class_name == "LocalStorage":
            job_uuid = str(uuid.uuid4())
            await self._create_execution_job(job_uuid, input_data)
            for i, record in enumerate(records):
                record_uid = str(uuid.uuid4())
                output_ref = await self.storage.save_record_data(record_uid, self.master_job_id, job_uuid, record)
                record["job_id"] = job_uuid
                record["master_job_id"] = self.master_job_id
                record["status"] = task_status
                record["output_ref"] = output_ref
                record["end_time"] = datetime.datetime.now(datetime.timezone.utc)
                record_model = Record(**record)
                await self.storage.log_record_metadata(record_model)
                logger.debug(f"  - Saved data for record {i}: {output_ref}")
                output_ref_list.append(output_ref)
            await self._complete_execution_job(job_uuid, status=task_status, num_records=len(records))
        return output_ref_list

    async def _complete_execution_job(self, job_uuid: str, status: str, num_records: int):
        logger.debug("\n6. Completing execution job...")
        now = datetime.datetime.now(datetime.timezone.utc)
        completed_num_records = 0
        fitered_num_records = 0
        duplicated_num_records = 0
        if status == STATUS_COMPLETED:
            completed_num_records = num_records
        elif status == STATUS_FILTERED:
            fitered_num_records = num_records
        elif status == STATUS_DUPLICATE:
            duplicated_num_records = num_records

        counts = {
            STATUS_COMPLETED: completed_num_records,
            STATUS_FILTERED: fitered_num_records,
            STATUS_DUPLICATE: duplicated_num_records,
        }
        await self.storage.log_execution_job_end(job_uuid, status, counts, now, now)
        logger.debug("  - Marked execution job as completed")

    def _is_job_to_stop(self) -> bool:
        # Immutable Snapshot: Creates a snapshot of the queue at the time of the list() call, so modifications to the queue during iteration won't affect the loop.
        job_output_list = list(self.job_output._queue)
        queue_size = len(job_output_list)
        if queue_size == 0:
            return False

        items = []
        for _ in range(min(self.job_config.job_run_stop_threshold, queue_size)):
            items.append(job_output_list[-1])

        consecutive_not_completed = len(items) == self.job_config.job_run_stop_threshold and all(item[RECORD_STATUS] != STATUS_COMPLETED for item in items)

        if consecutive_not_completed:
            logger.error(
                f"consecutive_not_completed: in {self.job_config.job_run_stop_threshold} times, "
                f"stopping this job; please adjust factory config and input data then "
                f"resume_from_checkpoint({self.master_job_id})"
            )
        target_not_reach_count = self.job_config.target_count - self.completed_count
        completed_tasks_reach_target = target_not_reach_count <= 0
        if target_not_reach_count > 0 and target_not_reach_count == self.dead_queue_count:
            logger.warning(f"there are {target_not_reach_count} input data not able to process, pls remove them")
            completed_tasks_reach_target = True

        return consecutive_not_completed or completed_tasks_reach_target

    # ====================
    # Progress Tracking
    # ====================
    async def _progress_ticker(self):
        """Log job progress at regular intervals."""
        while not self._is_job_to_stop():
            """Format and log the current job progress."""
            logger.info(
                f"[JOB PROGRESS] "
                f"\033[32mCompleted: {self.completed_count}/{self.job_config.target_count}\033[0m | "
                f"\033[33mRunning: {self.job_config.max_concurrency - self.semaphore._value}\033[0m | "
                f"\033[36mAttempted: {self.total_count}\033[0m"
                f"    (\033[32mCompleted: {self.completed_count}\033[0m, "
                f"\033[31mFailed: {self.failed_count}\033[0m, "
                f"\033[35mFiltered: {self.filtered_count}\033[0m, "
                f"\033[34mDuplicate: {self.duplicate_count}\033[0m, "
                f"\033[1;31mInDeadQueue: {self.dead_queue_count}\033[0m)"
            )
            await asyncio.sleep(PROGRESS_LOG_INTERVAL)

    async def _del_progress_ticker(self):
        """Safely stop the progress ticker."""
        if self._progress_ticker_task:
            self._progress_ticker_task.cancel()
            try:
                await self._progress_ticker_task
            except asyncio.CancelledError:
                pass

    # reserve this function for reference
    # async def _async_run_orchestration_3_11(self):
    #     """Main asynchronous orchestration loop for the job.

    #     This method manages the core job execution loop, including:
    #     - Starting the progress ticker
    #     - Processing tasks from the input queue
    #     - Managing concurrency with semaphores
    #     - Handling task completion and cleanup
    #     """
    #     # Start the ticker task
    #     if self.job_config.show_progress:
    #         self._progress_ticker_task = asyncio.create_task(self._progress_ticker())

    #     try:
    #         async with asyncio.TaskGroup() as tg:
    #             while not self._is_job_to_stop():
    #                 logger.debug("Job is not to stop, checking job input queue")
    #                 if not self.job_input_queue.empty():
    #                     logger.debug("Job input queue is not empty, acquiring semaphore")
    #                     await self.semaphore.acquire()
    #                     logger.debug("Semaphore acquired, waiting for task to complete")
    #                     input_data = self.job_input_queue.get()
    #                     task = tg.create_task(self._run_single_task(input_data))
    #                     asyncio.create_task(self._handle_task_completion(task))
    #     finally:
    #         await self._del_progress_ticker()
    #         await self._cancel_operations()