mage-ai 0.9.74__py3-none-any.whl → 0.9.79__py3-none-any.whl

mage_ai/orchestration/job_manager.py

@@ -8,6 +8,7 @@ class JobType(StrEnum):
     BLOCK_RUN = 'block_run'
     PIPELINE_RUN = 'pipeline_run'
     INTEGRATION_STREAM = 'integration_stream'
+    GENERIC_JOB = 'generic_job'
 
 
 class JobManager:
@@ -62,6 +63,19 @@ class JobManager:
             logging_tags=logging_tags,
         )
 
+    def has_generic_job(
+        self,
+        generic_job_id: int,
+        logger=None,
+        logging_tags: Dict = None,
+    ) -> bool:
+        job_id = self.__job_id(JobType.GENERIC_JOB, generic_job_id)
+        return self.queue.has_job(
+            job_id,
+            logger=logger,
+            logging_tags=logging_tags,
+        )
+
     def kill_block_run_job(self, block_run_id):
         print(f'Kill block run id: {block_run_id}')
         job_id = self.__job_id(JobType.BLOCK_RUN, block_run_id)
@@ -78,6 +92,11 @@ class JobManager:
         job_id = self.__job_id(JobType.INTEGRATION_STREAM, id)
         return self.queue.kill_job(job_id)
 
+    def kill_generic_job(self, generic_job_id):
+        print(f'Kill generic job id: {generic_job_id}')
+        job_id = self.__job_id(JobType.GENERIC_JOB, generic_job_id)
+        return self.queue.kill_job(job_id)
+
     def start(self):
         self.queue.start()
 

mage_ai/orchestration/metrics/pipeline_run.py

@@ -220,7 +220,7 @@ def __calculate_metrics(pipeline_run: PipelineRun, pipeline) -> Dict:
 
     pipeline_metrics_by_stream = {}
     pipeline_logs_by_stream = {}
-    pipeline_logs = pipeline_run.logs['content'].split('\n')
+    pipeline_logs = pipeline_run.logs[0]['content'].split('\n')
     for pipeline_log in pipeline_logs:
         tags = parse_line(pipeline_log)
         stream = tags.get('stream')

mage_ai/orchestration/notification/sender.py

@@ -69,7 +69,7 @@ class NotificationSender:
 
         if self.config.teams_config is not None and self.config.teams_config.is_valid:
             try:
-                send_teams_message(self.config.teams_config, summary)
+                send_teams_message(self.config.teams_config, summary, title)
             except Exception:
                 traceback.print_exc()
 
@@ -87,7 +87,7 @@ class NotificationSender:
 
         if self.config.google_chat_config is not None and self.config.google_chat_config.is_valid:
             try:
-                send_google_chat_message(self.config.google_chat_config, summary)
+                send_google_chat_message(self.config.google_chat_config, summary, title)
             except Exception:
                 traceback.print_exc()
 

mage_ai/orchestration/pipeline_scheduler_original.py

@@ -1,4 +1,3 @@
-import asyncio
 import collections
 import os
 import traceback
@@ -33,6 +32,7 @@ from mage_ai.orchestration.db.models.schedules import (
     Backfill,
     BlockRun,
     EventMatcher,
+    GenericJob,
     PipelineRun,
     PipelineSchedule,
 )
@@ -260,7 +260,7 @@ class PipelineScheduler:
                         pipeline=self.pipeline,
                         pipeline_run=self.pipeline_run,
                     )
-                    asyncio.run(UsageStatisticLogger().pipeline_run_ended(self.pipeline_run))
+                    UsageStatisticLogger().pipeline_run_ended_sync(self.pipeline_run)
 
                 self.logger_manager.output_logs_to_destination()
 
@@ -339,7 +339,7 @@ class PipelineScheduler:
         error_msg: str,
         status=PipelineRun.PipelineRunStatus.FAILED,
     ) -> None:
-        asyncio.run(UsageStatisticLogger().pipeline_run_ended(self.pipeline_run))
+        UsageStatisticLogger().pipeline_run_ended_sync(self.pipeline_run)
 
         if status == PipelineRun.PipelineRunStatus.FAILED:
             # Only send notification when pipeline run status is FAILED
@@ -929,6 +929,81 @@ class PipelineScheduler:
             self.memory_usage_failure(tags=tags)
 
 
+def on_pipeline_run_cancelled(
+    job_id: int,
+    pipeline_run_id: int,
+    cancelled_block_run_ids: List,
+):
+    job = GenericJob.query.get(job_id)
+    if job.status not in [
+        GenericJob.JobStatus.INITIAL,
+        GenericJob.JobStatus.QUEUED,
+        GenericJob.JobStatus.RUNNING,
+    ]:
+        return
+    job.mark_running()
+
+    if not pipeline_run_id:
+        job.mark_failed(
+            metadata=dict(
+                error_message='Pipeline run id is not provided in the job.',
+            ),
+        )
+        return
+    if not cancelled_block_run_ids:
+        job.mark_failed(
+            metadata=dict(
+                error_message='Cancelled block run ids are not provided in the job.',
+            ),
+        )
+        return
+
+    try:
+        # Run callback blocks for cancelled block runs
+        pipeline_run = PipelineRun.query.get(pipeline_run_id)
+        if not pipeline_run:
+            raise Exception(f'Fail to retrieve pipeline run with id {pipeline_run_id}')
+        pipeline_schedule = pipeline_run.pipeline_schedule
+        pipeline = get_pipeline_from_platform(
+            pipeline_run.pipeline_uuid,
+            repo_path=pipeline_schedule.repo_path if pipeline_schedule else None,
+        )
+        if pipeline_run.status != PipelineRun.PipelineRunStatus.CANCELLED:
+            job.mark_failed(
+                metadata=dict(
+                    error_message=f'Pipeline run {pipeline_run_id} is not in cancelled status.',
+                ),
+            )
+            return
+        for block_run_id in cancelled_block_run_ids:
+            block_run = BlockRun.query.get(block_run_id)
+            if not block_run:
+                continue
+            if block_run.status != BlockRun.BlockRunStatus.CANCELLED:
+                continue
+            ExecutorFactory.get_block_executor(
+                pipeline,
+                block_run.block_uuid,
+                block_run_id=block_run.id,
+                execution_partition=pipeline_run.execution_partition,
+                executor_type=ExecutorType.LOCAL_PYTHON,
+            ).execute_callback(
+                callback='on_cancelled',
+                global_vars=pipeline_run.get_variables(),
+                logging_tags=None,
+                pipeline_run=pipeline_run,
+                block_run_id=block_run_id,
+            )
+    except Exception as e:
+        job.mark_failed(
+            metadata=dict(
+                error_message=f'Failed to run pipeline run cancellation hook: {e}.',
+            ),
+        )
+        return
+    job.mark_completed()
+
+
 def run_integration_streams(
     streams: List[Dict],
     *args,
@@ -1372,7 +1447,7 @@ def stop_pipeline_run(
     # Update pipeline run status to cancelled
     pipeline_run.update(status=status)
 
-    asyncio.run(UsageStatisticLogger().pipeline_run_ended(pipeline_run))
+    UsageStatisticLogger().pipeline_run_ended_sync(pipeline_run)
 
     # Cancel all the block runs
     cancel_block_runs_and_jobs(pipeline_run, pipeline)
@@ -1407,8 +1482,9 @@ def cancel_block_runs_and_jobs(
             block_runs_to_cancel.append(b)
         if b.status == BlockRun.BlockRunStatus.RUNNING:
             running_blocks.append(b)
+    cancelled_block_run_ids = [b.id for b in block_runs_to_cancel]
     BlockRun.batch_update_status(
-        [b.id for b in block_runs_to_cancel],
+        cancelled_block_run_ids,
         BlockRun.BlockRunStatus.CANCELLED,
     )
 
@@ -1436,6 +1512,18 @@ def cancel_block_runs_and_jobs(
         for b in running_blocks:
             job_manager.kill_block_run_job(b.id)
 
+    GenericJob.enqueue_cancel_pipeline_run(pipeline_run.id, cancelled_block_run_ids)
+
+
+def kill_cancelled_runs(pipeline_run_id: int, cancelled_block_run_ids: List):
+    job_manager = get_job_manager()
+    if pipeline_run_id and job_manager.has_pipeline_run_job(pipeline_run_id):
+        job_manager.kill_pipeline_run_job(pipeline_run_id)
+    if cancelled_block_run_ids:
+        for block_run_id in cancelled_block_run_ids:
+            if job_manager.has_block_run_job(block_run_id):
+                job_manager.kill_block_run_job(block_run_id)
+
 
 def check_sla():
     repo_pipelines = set(
@@ -1584,7 +1672,7 @@ def schedule_all():
         pipeline_runs_excluded_by_limit = []
         for pipeline_schedule in active_schedules:
             lock_key = f'pipeline_schedule_{pipeline_schedule.id}'
-            if not lock.try_acquire_lock(lock_key):
+            if not lock.try_acquire_lock(lock_key, timeout=30):
                 continue
 
             trigger_pipeline_run_limit = pipeline_schedule.get_settings().pipeline_run_limit
@@ -1733,6 +1821,13 @@ def schedule_all():
             logger.exception(f'Failed to schedule {r}')
             traceback.print_exc()
             continue
+
+    try:
+        schedule_generic_jobs()
+    except Exception as e:
+        logger.exception(f'Failed to schedule generic jobs {e}')
+        traceback.print_exc()
+
     job_manager = get_job_manager()
     if job_manager is None:
         logger.info('Job manager is None.')
@@ -1913,3 +2008,52 @@ def sync_schedules(pipeline_uuids: List[str]):
             trigger_configs.append(pipeline_trigger)
 
     PipelineSchedule.create_or_update_batch(trigger_configs)
+
+
+def schedule_generic_jobs():
+    """
+    Schedule generic jobs that are in INITIAL status.
+
+    This method schedules jobs like cancel_pipeline_run that have been enqueued
+    by the server process and are waiting to be executed by the job manager.
+    """
+    try:
+        # Get all jobs with INITIAL status across all repositories
+        jobs = GenericJob.get_jobs_with_initial_status()
+        job_manager = get_job_manager()
+
+        if not jobs:
+            return
+
+        logger.info(f'Scheduling {len(jobs)} generic jobs')
+
+        # Schedule each job with the job manager
+        for job in jobs:
+            try:
+                if job.job_type == GenericJob.JobType.CANCEL_PIPELINE_RUN:
+                    pipeline_run_id = job.payload.get('pipeline_run_id')
+                    cancelled_block_run_ids = job.payload.get('cancelled_block_run_ids')
+                    try:
+                        kill_cancelled_runs(pipeline_run_id, cancelled_block_run_ids)
+                    except Exception:
+                        pass
+                    # Add job to job manager for execution
+                    job_manager.add_job(
+                        JobType.GENERIC_JOB,
+                        job.id,
+                        on_pipeline_run_cancelled,
+                        # args
+                        job.id,
+                        pipeline_run_id,
+                        cancelled_block_run_ids,
+                    )
+                    job.mark_queued()
+                    logger.info(f'Scheduled generic job {job.id} ({job.job_type})')
+                else:
+                    raise ValueError(f"Unknown job type: {job.job_type}")
+            except Exception as e:
+                error_message = f"Failed to schedule generic job: {str(e)}"
+                logger.exception(f'Failed to schedule generic job {job.id}: {error_message}')
+                job.mark_failed({'error': error_message})
+    except Exception as e:
+        logger.exception(f'Failed to schedule generic jobs: {e}')

mage_ai/orchestration/pipeline_scheduler_project_platform.py

@@ -1,4 +1,3 @@
-import asyncio
 import collections
 import os
 import traceback
@@ -262,7 +261,7 @@ class PipelineScheduler:
                         pipeline_run=self.pipeline_run,
                     )
 
-                asyncio.run(UsageStatisticLogger().pipeline_run_ended(self.pipeline_run))
+                UsageStatisticLogger().pipeline_run_ended_sync(self.pipeline_run)
 
                 self.logger_manager.output_logs_to_destination()
 
@@ -338,7 +337,7 @@ class PipelineScheduler:
 
     @safe_db_query
     def on_pipeline_run_failure(self, error: str) -> None:
-        asyncio.run(UsageStatisticLogger().pipeline_run_ended(self.pipeline_run))
+        UsageStatisticLogger().pipeline_run_ended_sync(self.pipeline_run)
         self.notification_sender.send_pipeline_run_failure_message(
             pipeline=self.pipeline,
             pipeline_run=self.pipeline_run,
@@ -1317,7 +1316,7 @@ def stop_pipeline_run(
     # Update pipeline run status to cancelled
     pipeline_run.update(status=status)
 
-    asyncio.run(UsageStatisticLogger().pipeline_run_ended(pipeline_run))
+    UsageStatisticLogger().pipeline_run_ended_sync(pipeline_run)
 
     # Cancel all the block runs
     cancel_block_runs_and_jobs(pipeline_run, pipeline)
@@ -1568,7 +1567,7 @@ def schedule_all():
                 concurrency_config = ConcurrencyConfig.load(config=pipeline.concurrency_config)
 
                 lock_key = f'pipeline_schedule_{pipeline_schedule.id}'
-                if not lock.try_acquire_lock(lock_key):
+                if not lock.try_acquire_lock(lock_key, timeout=30):
                     continue
 
                 try:

mage_ai/orchestration/queue/config.py

@@ -1,8 +1,11 @@
+import traceback
 from dataclasses import dataclass
 
 from mage_ai.shared.config import BaseConfig
 from mage_ai.shared.enum import StrEnum
 
+DEFAULT_CONCURRENCY = 20
+
 
 class QueueType(StrEnum):
     CELERY = 'celery'
@@ -17,5 +20,12 @@ class ProcessQueueConfig(BaseConfig):
 @dataclass
 class QueueConfig(BaseConfig):
     queue_type: QueueType = QueueType.PROCESS
-    concurrency: int = 20
+    concurrency: int = DEFAULT_CONCURRENCY
     process_queue_config: ProcessQueueConfig = None
+
+    def __post_init__(self):
+        try:
+            self.concurrency = int(self.concurrency)
+        except Exception:
+            traceback.print_exc()
+            self.concurrency = DEFAULT_CONCURRENCY

mage_ai/orchestration/queue/process_queue.py

@@ -19,6 +19,7 @@ from mage_ai.settings import (
     HOSTNAME,
     REDIS_URL,
     SENTRY_DSN,
+    SENTRY_SERVER_NAME,
     SENTRY_TRACES_SAMPLE_RATE,
     SERVER_LOGGING_FORMAT,
     SERVER_VERBOSITY,
@@ -287,7 +288,10 @@ class Worker(mp.Process):
             sentry_sdk.init(
                 self.dsn,
                 traces_sample_rate=SENTRY_TRACES_SAMPLE_RATE,
+                server_name=SENTRY_SERVER_NAME,
             )
+            import atexit
+            atexit.register(lambda: sentry_sdk.flush(timeout=5))
         initialize_new_relic()
 
         set_logging_format(

mage_ai/orchestration/utils/distributed_lock.py

@@ -1,12 +1,19 @@
+import os
+
 from mage_ai.services.redis.redis import init_redis_client
 from mage_ai.settings import REDIS_URL
 
+try:
+    REDIS_LOCK_DEFAULT_TIMEOUT = int(os.getenv('REDIS_LOCK_DEFAULT_TIMEOUT', 30) or 30)
+except Exception:
+    REDIS_LOCK_DEFAULT_TIMEOUT = 30
+
 
 class DistributedLock:
     def __init__(
         self,
         lock_key_prefix='LOCK_KEY',
-        lock_timeout=30,
+        lock_timeout=REDIS_LOCK_DEFAULT_TIMEOUT,
     ):
         self.lock_key_prefix = lock_key_prefix
         self.lock_timeout = lock_timeout

mage_ai/orchestration/utils/resources.py

@@ -4,6 +4,9 @@ from typing import Tuple
 
 import psutil
 
+# Threshold value for cgroup memory limit indicating no limit is set
+NO_LIMIT_THRESHOLD = 9e18
+
 
 def get_compute() -> Tuple[float, float, float, float]:
     # Getting loadover15 minutes
@@ -14,14 +17,65 @@ def get_compute() -> Tuple[float, float, float, float]:
 
 
 def get_memory() -> Tuple[float, float, float]:
+    """
+    Get memory usage information.
+
+    For containerized environments (e.g., Google Cloud Run), this reads from cgroup files
+    which provide accurate container-specific memory limits and usage.
+
+    For non-containerized environments, falls back to the 'free' command.
+
+    Returns:
+        Tuple[float, float, float]: (free_memory, used_memory, total_memory) in MB
+    """
     free_memory = None
     total_memory = None
     used_memory = None
 
+    # Skip checking the memory in Windows
+    if os.name == 'nt':
+        return free_memory, used_memory, total_memory
+
+    # Try to read from cgroup files first (for containerized environments)
     try:
-        # Skip check the memory in Windows
-        if os.name == 'nt':
+        # Define MB conversion factor (1024 * 1024)
+        mb_factor = 1024 * 1024
+
+        # Try cgroup v1 paths first
+        usage_path = "/sys/fs/cgroup/memory/memory.usage_in_bytes"
+        limit_path = "/sys/fs/cgroup/memory/memory.limit_in_bytes"
+
+        # Check if cgroup v2 is being used (single unified hierarchy)
+        if not os.path.exists(usage_path):
+            usage_path = "/sys/fs/cgroup/memory.current"
+            limit_path = "/sys/fs/cgroup/memory.max"
+
+        # Read memory usage
+        if os.path.exists(usage_path):
+            with open(usage_path, "r") as f:
+                usage_bytes = int(f.read().strip())
+                used_memory = usage_bytes / mb_factor
+
+        # Read memory limit
+        if os.path.exists(limit_path):
+            with open(limit_path, "r") as f:
+                limit_value = f.read().strip()
+                # Handle cgroup v2 "max" value or v1 very large limit (no limit set)
+                if limit_value != "max":
+                    limit_bytes = int(limit_value)
+                    # A value of 9223372036854771712 or greater indicates no limit set
+                    if limit_bytes < NO_LIMIT_THRESHOLD:
+                        total_memory = limit_bytes / mb_factor
+        # If we successfully read both values, calculate free memory
+        if used_memory is not None and total_memory is not None:
+            free_memory = total_memory - used_memory
             return free_memory, used_memory, total_memory
+    except (FileNotFoundError, ValueError, PermissionError):
+        # Cgroup files not available or unreadable, will fall back to 'free' command
+        pass
+
+    # Fallback to 'free' command for non-containerized environments
+    try:
         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)