mage-ai 0.9.68__py3-none-any.whl → 0.9.70__py3-none-any.whl

mage_ai/orchestration/db/models/schedules.py

@@ -48,6 +48,7 @@ from mage_ai.data_preparation.models.constants import (
     PipelineType,
 )
 from mage_ai.data_preparation.models.pipeline import Pipeline
+from mage_ai.data_preparation.models.project import Project
 from mage_ai.data_preparation.models.triggers import (
     ScheduleInterval,
     ScheduleStatus,
@@ -934,7 +935,10 @@ class PipelineRun(PipelineRunProjectPlatformMixin, BaseModel):
 
     @property
     def pipeline_tags(self):
-        pipeline_config = Pipeline.get_config(self.pipeline_uuid)
+        try:
+            pipeline_config = Pipeline.get_config(self.pipeline_uuid)
+        except Exception:
+            pipeline_config = dict()
 
         return pipeline_config.get('tags') if pipeline_config is not None else []
 
@@ -989,6 +993,9 @@ class PipelineRun(PipelineRunProjectPlatformMixin, BaseModel):
         block_runs_all = []
 
         data_integration_block_uuids_mapping = {}
+
+        pipeline_project = Project(pipeline.repo_config)
+
         for block_run in self.block_runs:
             block_runs_all.append(block_run)
 
@@ -1021,7 +1028,7 @@ class PipelineRun(PipelineRunProjectPlatformMixin, BaseModel):
                 "original": 1,
             }
             """
-            if metrics and block and block.is_data_integration():
+            if metrics and block and block.is_data_integration(pipeline_project=pipeline_project):
                 original_block_uuid = metrics.get('original_block_uuid')
 
                 if original_block_uuid and metrics.get('child'):
@@ -1906,3 +1913,24 @@ class Backfill(BaseModel):
                 Backfill.pipeline_schedule_id.in_(pipeline_schedule_ids),
             )
         return []
+
+    @property
+    def pipeline_run_status_counts(self) -> Dict:
+        status_counts = dict()
+        execution_dates_counted = set()
+
+        # Sort the pipeline runs by id in reverse order so the first pipeline run
+        # checked is the latest pipeline run created for a given execution date.
+        pipeline_runs_sorted = sorted(
+            self.pipeline_runs,
+            key=lambda pr: (pr.execution_date, pr.id),
+            reverse=True,
+        )
+
+        for pr in pipeline_runs_sorted:
+            # Only count a pipeline run once per execution date
+            if pr.execution_date not in execution_dates_counted:
+                status_counts[pr.status] = (status_counts.get(pr.status) or 0) + 1
+                execution_dates_counted.add(pr.execution_date)
+
+        return status_counts

mage_ai/orchestration/job_manager.py

@@ -86,6 +86,12 @@ class JobManager:
         job_id = self.__job_id(JobType.INTEGRATION_STREAM, id)
         return self.queue.kill_job(job_id)
 
+    def start(self):
+        self.queue.start()
+
+    def stop(self):
+        self.queue.stop()
+
     def __job_id(self, job_type: JobType, uid: Union[str, int]):
         return f'{job_type}_{uid}'
 

mage_ai/orchestration/notification/sender.py

@@ -1,4 +1,5 @@
 import os
+import traceback
 from typing import Dict
 
 from mage_ai.orchestration.notification.config import (
@@ -61,33 +62,54 @@ class NotificationSender:
         if summary is None:
             return
         if self.config.slack_config is not None and self.config.slack_config.is_valid:
-            send_slack_message(self.config.slack_config, details or summary, title)
+            try:
+                send_slack_message(self.config.slack_config, details or summary, title)
+            except Exception:
+                traceback.print_exc()
 
         if self.config.teams_config is not None and self.config.teams_config.is_valid:
-            send_teams_message(self.config.teams_config, summary)
+            try:
+                send_teams_message(self.config.teams_config, summary)
+            except Exception:
+                traceback.print_exc()
 
         if self.config.discord_config is not None and self.config.discord_config.is_valid:
-            send_discord_message(self.config.discord_config, summary, title)
+            try:
+                send_discord_message(self.config.discord_config, summary, title)
+            except Exception:
+                traceback.print_exc()
 
         if self.config.telegram_config is not None and self.config.telegram_config.is_valid:
-            send_telegram_message(self.config.telegram_config, summary, title)
+            try:
+                send_telegram_message(self.config.telegram_config, summary, title)
+            except Exception:
+                traceback.print_exc()
 
         if self.config.google_chat_config is not None and self.config.google_chat_config.is_valid:
-            send_google_chat_message(self.config.google_chat_config, summary)
+            try:
+                send_google_chat_message(self.config.google_chat_config, summary)
+            except Exception:
+                traceback.print_exc()
 
         if self.config.email_config is not None and title is not None:
-            send_email(
-                self.config.email_config,
-                subject=title,
-                message=details or summary,
-            )
+            try:
+                send_email(
+                    self.config.email_config,
+                    subject=title,
+                    message=details or summary,
+                )
+            except Exception:
+                traceback.print_exc()
 
         if self.config.opsgenie_config is not None and self.config.opsgenie_config.is_valid:
-            send_opsgenie_alert(
-                self.config.opsgenie_config,
-                message=title,
-                description=details or summary,
-            )
+            try:
+                send_opsgenie_alert(
+                    self.config.opsgenie_config,
+                    message=title,
+                    description=details or summary,
+                )
+            except Exception:
+                traceback.print_exc()
 
     def send_pipeline_run_success_message(self, pipeline, pipeline_run) -> None:
         if AlertOn.PIPELINE_RUN_SUCCESS in self.config.alert_on:

mage_ai/orchestration/pipeline_scheduler_original.py

@@ -284,8 +284,7 @@ class PipelineScheduler:
                     or PipelineRun.PipelineRunStatus.FAILED
                 )
                 self.pipeline_run.update(status=status)
-
-                self.on_pipeline_run_failure('Pipeline run timed out.')
+                self.on_pipeline_run_failure('Pipeline run timed out.', status=status)
             elif self.pipeline_run.any_blocks_failed() and not self.allow_blocks_to_fail:
                 self.pipeline_run.update(
                     status=PipelineRun.PipelineRunStatus.FAILED)
@@ -318,30 +317,38 @@ class PipelineScheduler:
                     self.__schedule_blocks(block_runs)
 
     @safe_db_query
-    def on_pipeline_run_failure(self, error_msg: str) -> None:
-        failed_block_runs = self.pipeline_run.failed_block_runs
-        stacktrace = None
-        for br in failed_block_runs:
-            if br.metrics:
-                message = br.metrics.get('error', {}).get('message')
-                if message:
-                    message_split = message.split('\n')
-                    # Truncate the error message if it has too many lines, set max
-                    # lines at 50
-                    if len(message_split) > 50:
-                        message_split = message_split[-50:]
-                        message_split.insert(0, '... (error truncated)')
-                    message = '\n'.join(message_split)
-                    stacktrace = f'Error for block {br.block_uuid}:\n{message}'
-                    break
-
+    def on_pipeline_run_failure(
+        self,
+        error_msg: str,
+        status=PipelineRun.PipelineRunStatus.FAILED,
+    ) -> None:
         asyncio.run(UsageStatisticLogger().pipeline_run_ended(self.pipeline_run))
-        self.notification_sender.send_pipeline_run_failure_message(
-            pipeline=self.pipeline,
-            pipeline_run=self.pipeline_run,
-            error=error_msg,
-            stacktrace=stacktrace,
-        )
+
+        if status == PipelineRun.PipelineRunStatus.FAILED:
+            # Only send notification when pipeline run status is FAILED
+            failed_block_runs = self.pipeline_run.failed_block_runs
+            stacktrace = None
+            for br in failed_block_runs:
+                if br.metrics:
+                    message = br.metrics.get('error', {}).get('message')
+                    if message:
+                        message_split = message.split('\n')
+                        # Truncate the error message if it has too many lines, set max
+                        # lines at 50
+                        if len(message_split) > 50:
+                            message_split = message_split[-50:]
+                            message_split.insert(0, '... (error truncated)')
+                        message = '\n'.join(message_split)
+                        stacktrace = f'Error for block {br.block_uuid}:\n{message}'
+                        break
+
+            self.notification_sender.send_pipeline_run_failure_message(
+                pipeline=self.pipeline,
+                pipeline_run=self.pipeline_run,
+                error=error_msg,
+                stacktrace=stacktrace,
+            )
+
         # Cancel block runs that are still in progress for the pipeline run.
         cancel_block_runs_and_jobs(self.pipeline_run, self.pipeline)
 
@@ -830,6 +837,8 @@ class PipelineScheduler:
         Returns:
             List[BlockRun]: A list of crashed block runs.
         """
+        for b in self.pipeline_run.block_runs:
+            b.refresh()
         running_or_queued_block_runs = [b for b in self.pipeline_run.block_runs if b.status in [
             BlockRun.BlockRunStatus.RUNNING,
             BlockRun.BlockRunStatus.QUEUED,
@@ -1679,7 +1688,11 @@ def gen_pipeline_with_schedules_single_project(
     # Iterate through pipeline schedules by pipeline to handle pipeline run limits for
     # each pipeline.
     for pipeline_uuid, active_schedules in pipeline_schedules_by_pipeline.items():
-        pipeline = Pipeline.get(pipeline_uuid)
+        try:
+            pipeline = Pipeline.get(pipeline_uuid)
+        except Exception as e:
+            print(f'Error fetching pipeline {pipeline_uuid}: {e}')
+            continue
         yield pipeline_uuid, pipeline, active_schedules
 
 
@@ -1740,11 +1753,15 @@ def gen_pipeline_with_schedules_project_platform(
     for pair in pipeline_schedules_by_pipeline_by_repo_path.items():
         repo_path, pipeline_schedules_by_pipeline = pair
         for pipeline_uuid, active_schedules in pipeline_schedules_by_pipeline.items():
-            pipeline = get_pipeline_from_platform(
-                pipeline_uuid,
-                repo_path=repo_path,
-                mapping=pipeline_schedule_repo_paths_to_repo_path_mapping,
-            )
+            try:
+                pipeline = get_pipeline_from_platform(
+                    pipeline_uuid,
+                    repo_path=repo_path,
+                    mapping=pipeline_schedule_repo_paths_to_repo_path_mapping,
+                )
+            except Exception as e:
+                print(f'Error fetching pipeline {pipeline_uuid}: {e}')
+                continue
             yield pipeline_uuid, pipeline, active_schedules
 
 

mage_ai/orchestration/queue/celery_queue.py

@@ -1,6 +1,7 @@
+from typing import Callable
+
 from mage_ai.orchestration.queue.config import QueueConfig
 from mage_ai.orchestration.queue.queue import Queue
-from typing import Callable
 
 
 class CeleryQueue(Queue):
@@ -27,3 +28,9 @@ class CeleryQueue(Queue):
 
     def kill_job(self, job_id: str):
         pass
+
+    def start(self):
+        pass
+
+    def stop(self):
+        pass

mage_ai/orchestration/queue/process_queue.py

@@ -36,6 +36,11 @@ class JobStatus(str, Enum):
     CANCELLED = 'cancelled'
 
 
+class QueueStatus(str, Enum):
+    ACTIVE = 'active'
+    INACTIVE = 'inactive'
+
+
 class ProcessQueue(Queue):
     def __init__(self, queue_config: QueueConfig):
         """
@@ -53,6 +58,7 @@ class ProcessQueue(Queue):
                 jobs.
 
         """
+        self.status = QueueStatus.INACTIVE
         self.queue_config = queue_config
         self.process_queue_config = self.queue_config.process_queue_config
         self.queue = mp.Queue()
@@ -76,12 +82,16 @@ class ProcessQueue(Queue):
 
     def clean_up_jobs(self):
         """
-        Cleans up completed jobs from the job dictionary.
+        1. Cleans up completed jobs from the job dictionary.
+        2. Check whether there're jobs need to be killed.
         """
         job_ids = self.job_dict.keys()
         for job_id in job_ids:
-            if job_id in self.job_dict and not self.has_job(job_id):
-                del self.job_dict[job_id]
+            if job_id in self.job_dict:
+                if not self.has_job(job_id):
+                    del self.job_dict[job_id]
+                elif self.__should_kill_job(job_id):
+                    self.kill_job(job_id)
 
     def enqueue(self, job_id: str, target: Callable, *args, **kwargs):
         """
@@ -94,6 +104,9 @@ class ProcessQueue(Queue):
             **kwargs: Keyword arguments for the target function.
 
         """
+        if self.status != QueueStatus.ACTIVE:
+            self._print('Cannot enqueue a job to an inactive queue.')
+            return
         if self.has_job(job_id):
             self._print(f'Job {job_id} exists. Skip enqueue.')
             return
@@ -155,6 +168,7 @@ class ProcessQueue(Queue):
         print(f'Kill job {job_id}, job_dict {self.job_dict}')
         job = self.job_dict.get(job_id)
         if not job:
+            self.__set_kill_job(job_id)
             return
         if isinstance(job, int):
             if job == os.getpid():
@@ -165,6 +179,7 @@ class ProcessQueue(Queue):
             except Exception as err:
                 print(err)
         self.job_dict[job_id] = JobStatus.CANCELLED
+        self.__unset_kill_job(job_id)
 
     def start_worker_pool(self):
         """
@@ -182,6 +197,28 @@ class ProcessQueue(Queue):
         )
         self.worker_pool_proc.start()
 
+    def start(self):
+        self.status = QueueStatus.ACTIVE
+
+    def stop(self):
+        """
+        1. Stop enqueueing new jobs
+        2. Clear the queue
+        3. Kill all the running jobs
+        """
+        self.status = QueueStatus.INACTIVE
+        while not self.queue.empty():
+            try:
+                self.queue.get_nowait()
+            except self.queue.Empty:
+                break
+        job_ids = self.job_dict.keys()
+        for job_id in job_ids:
+            if job_id in self.job_dict:
+                if isinstance(self.job_dict.get(job_id), int):
+                    self.kill_job(job_id)
+                del self.job_dict[job_id]
+
     def is_worker_pool_alive(self) -> bool:
         """
         Checks if the worker pool process is alive.
@@ -197,6 +234,31 @@ class ProcessQueue(Queue):
     def __is_process_alive(self, pid: int) -> bool:
         return psutil.pid_exists(pid)
 
+    def __redis_key_kill_job(self, job_id):
+        return f'kill_job_{job_id}'
+
+    def __set_kill_job(self, job_id):
+        if not self.redis_client:
+            return
+        return self.redis_client.set(
+            self.__redis_key_kill_job(job_id),
+            '1',
+            ex=LIVENESS_TIMEOUT_SECONDS,
+        )
+
+    def __unset_kill_job(self, job_id):
+        if not self.redis_client:
+            return
+        key = self.__redis_key_kill_job(job_id)
+        if self.redis_client.get(key):
+            self.redis_client.delete(key)
+
+    def __should_kill_job(self, job_id):
+        if not self.redis_client:
+            return False
+        value = self.redis_client.get(self.__redis_key_kill_job(job_id))
+        return value is not None
+
 
 class Worker(mp.Process):
     def __init__(
@@ -275,10 +337,11 @@ def poll_job_and_execute(
         job_dict: The shared job dictionary.
 
     """
+    pid = os.getpid()
     workers = []
     while True:
         workers = [w for w in workers if w.is_alive()]
-        print(f'Worker pool size: {len(workers)}')
+        print(f'[Process {pid}] Worker pool size: {len(workers)}')
         if not workers and queue.empty():
             break
         while not queue.empty():

mage_ai/orchestration/queue/queue.py

@@ -19,5 +19,13 @@ class Queue(ABC):
     def kill_job(self, job_id: str):
         pass
 
+    @abstractmethod
+    def start(self):
+        pass
+
+    @abstractmethod
+    def stop(self):
+        pass
+
     def _print(self, msg):
         print(f'[{self.__class__.__name__}] {msg}')

mage_ai/orchestration/triggers/api.py

@@ -1,7 +1,8 @@
 from datetime import datetime
-from typing import Dict, Optional
+from typing import Dict, List, Optional, Union
 
 from mage_ai.api.resources.PipelineScheduleResource import PipelineScheduleResource
+from mage_ai.data_preparation.models.block.remote.models import RemoteBlock
 from mage_ai.data_preparation.models.pipeline import Pipeline
 from mage_ai.data_preparation.models.triggers import ScheduleStatus, ScheduleType
 from mage_ai.orchestration.db.models.schedules import PipelineRun, PipelineSchedule
@@ -25,11 +26,21 @@ def trigger_pipeline(
     poll_timeout: Optional[float] = None,
     schedule_name: str = None,
     verbose: bool = True,
+    remote_blocks: List[Union[Dict, RemoteBlock]] = None,
+    return_remote_blocks: bool = False,
     _should_schedule: bool = False,  # For internal use only (e.g. running hooks from notebook).
 ) -> PipelineRun:
     if variables is None:
         variables = {}
 
+    if remote_blocks:
+        arr = []
+        for remote_block in remote_blocks:
+            if isinstance(remote_block, dict):
+                remote_block = RemoteBlock.load(**remote_block)
+            arr.append(remote_block.to_dict())
+        variables['remote_blocks'] = arr
+
     pipeline = Pipeline.get(pipeline_uuid, all_projects=project_platform_activated())
 
     pipeline_schedule = __fetch_or_create_pipeline_schedule(pipeline, schedule_name=schedule_name)
@@ -50,6 +61,23 @@ def trigger_pipeline(
             verbose=verbose,
         )
 
+    if return_remote_blocks and pipeline_run and pipeline_run.pipeline:
+        pipeline = pipeline_run.pipeline
+
+        return [
+            dict(
+                remote_blocks=[
+                    RemoteBlock.load(
+                        block_uuid=block.uuid,
+                        execution_partition=pipeline_run.execution_partition,
+                        pipeline_uuid=pipeline.uuid,
+                        repo_path=pipeline.repo_path,
+                    )
+                    for block in pipeline.blocks_by_uuid.values()
+                ],
+            ),
+        ]
+
     return pipeline_run
 
 

mage_ai/orchestration/triggers/global_data_product.py

@@ -3,9 +3,10 @@ import json
 from datetime import datetime, timedelta, timezone
 from logging import Logger
 from time import sleep
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Union
 
 from mage_ai.data_preparation.logging.logger import DictLogger
+from mage_ai.data_preparation.models.block.remote.models import RemoteBlock
 from mage_ai.data_preparation.models.global_data_product import GlobalDataProduct
 from mage_ai.data_preparation.models.triggers import ScheduleStatus, ScheduleType
 from mage_ai.orchestration.db import safe_db_query
@@ -145,6 +146,7 @@ def trigger_and_check_status(
     round_number: int = 0,
     verbose: bool = True,
     should_schedule: bool = False,
+    remote_blocks: List[Union[Dict, RemoteBlock]] = None,
 ):
     tags = merge_dict(logging_tags, dict(
         block_uuid=block.uuid if block else None,
@@ -163,6 +165,7 @@ def trigger_and_check_status(
             print(log_message)
             print(json.dumps(tags, indent=2))
 
+    pipeline_run = None
     pipeline_run_created = None
     tries = 0
 
@@ -305,6 +308,7 @@ def trigger_and_check_status(
                         global_data_product.pipeline,
                         pipeline_schedule,
                         dict(variables=variables),
+                        remote_blocks=remote_blocks,
                         should_schedule=should_schedule,
                     )
                     if pipeline_run_created:
@@ -312,11 +316,9 @@ def trigger_and_check_status(
                             f'Created pipeline run {pipeline_run_created.id} for '
                             f'global data product {global_data_product.uuid}.'
                         )
+                        pipeline_run = pipeline_run_created
 
                     lock.release_lock(__lock_key_for_creating_pipeline_run(global_data_product))
-
-                    if pipeline_run_created:
-                        break
             else:
                 __log(
                     f'No pipeline run for global data product {global_data_product.uuid} '
@@ -339,6 +341,7 @@ def trigger_and_check_status(
                     poll_timeout=poll_timeout,
                     verbose=verbose,
                     should_schedule=should_schedule,
+                    remote_blocks=remote_blocks,
                 )
                 break
         else:
@@ -359,6 +362,8 @@ def trigger_and_check_status(
         else:
             break
 
+    return pipeline_run or pipeline_run_created
+
 
 def fetch_or_create_pipeline_schedule(global_data_product: GlobalDataProduct) -> PipelineSchedule:
     pipeline_uuid = global_data_product.object_uuid

mage_ai/orchestration/triggers/utils.py

@@ -1,7 +1,8 @@
 from datetime import datetime, timedelta
 from time import sleep
-from typing import Dict, Optional
+from typing import Dict, List, Optional, Union
 
+from mage_ai.data_preparation.models.block.remote.models import RemoteBlock
 from mage_ai.data_preparation.models.pipeline import Pipeline
 from mage_ai.data_preparation.models.triggers import ScheduleStatus
 from mage_ai.orchestration.db import db_connection, safe_db_query
@@ -84,6 +85,7 @@ def create_and_start_pipeline_run(
     pipeline_schedule: PipelineSchedule,
     payload: Dict = None,
     should_schedule: bool = False,
+    remote_blocks: List[Union[Dict, RemoteBlock]] = None,
 ) -> PipelineRun:
     if payload is None:
         payload = {}
@@ -94,6 +96,13 @@ def create_and_start_pipeline_run(
         payload,
     )
 
+    if remote_blocks:
+        variables = configured_payload.get('variables', {})
+        if variables.get('remote_blocks'):
+            remote_blocks = variables.get('remote_blocks', []) + remote_blocks
+        variables['remote_blocks'] = remote_blocks
+        configured_payload['variables'] = variables
+
     pipeline_run = PipelineRun.create(**configured_payload)
 
     # Do not start the pipeline run immediately due to concurrency control

mage_ai/orchestration/utils/resources.py

@@ -19,6 +19,9 @@ def get_memory() -> Tuple[float, float, float]:
     used_memory = None
 
     try:
+        # Skip check the memory in Windows
+        if os.name == 'nt':
+            return free_memory, used_memory, total_memory
         output = subprocess.check_output('free -t -m', shell=True).decode('utf-8')
         values = output.splitlines()[-1].split()[1:]
         total_memory, used_memory, free_memory = map(float, values)

mage_ai/server/api/downloads.py

@@ -114,7 +114,10 @@ class ApiResourceDownloadHandler(BaseHandler):
     # file pointer points to either a singular file or a temporary zip
     def get_file_pointer(self, file_list, relative_file_list):
         if len(file_list) == 1:
-            return open(file_list[0])
+            if file_list[0].endswith('.xlsx'):  # Check if it's an XLSX file
+                return open(file_list[0], 'rb')  # Open in binary mode for XLSX
+            else:
+                return open(file_list[0])
         return self.zip_files(file_list, relative_file_list)
 
     # creates a temporary zip and returns the (open) file pointer