mage-ai 0.9.46__py3-none-any.whl → 0.9.48__py3-none-any.whl

mage_ai/io/chroma.py

@@ -0,0 +1,162 @@
+import logging
+from typing import Dict
+
+import chromadb
+from pandas import DataFrame
+
+from mage_ai.io.base import BaseIO
+from mage_ai.io.config import BaseConfigLoader, ConfigKey
+
+logger = logging.getLogger(__name__)
+
+
+class Chroma(BaseIO):
+    def __init__(
+            self,
+            collection: str,
+            path: str = None,
+            verbose: bool = True,
+            **kwargs,) -> None:
+        """
+        Initializes settings to connect to chroma db.
+        """
+        super().__init__(verbose=verbose)
+        self.collection = collection
+        if path is None:
+            self.client = chromadb.Client()
+        else:
+            self.client = chromadb.PersistentClient(path=path)
+
+    @classmethod
+    def with_config(cls, config: BaseConfigLoader) -> 'Chroma':
+        return cls(
+            collection=config[ConfigKey.CHROMA_COLLECTION],
+            path=config[ConfigKey.CHROMA_PATH],
+        )
+
+    def __convert_to_multiple_rows(self, data: Dict, column_name: str, item_length: int):
+        """
+        Convert the data from chroma query into multiple rows.
+
+        Args:
+            data (Dict): Data returned from chroma query.
+            Example format:
+            {
+                'ids': [['id1', '0', 'id2']],
+                'distances': [[0.6036068181428754, 0.6036068181428754, 1.267588382105355]],
+                'metadatas': [[
+                                {'chapter': '3', 'verse': '16'},
+                                None,
+                                {'chapter': '3', 'verse': '5'}
+                              ]],
+                'embeddings': None,
+                'documents': [['abc', 'abc', 'def']],
+                'uris': None,
+                'data': None
+            }
+            column_name (str): Name of the field in data to convert.
+            item_length (int): Number of expected items in the list.
+
+        Returns:
+            DataFrame: Data frame object loaded from the chroma query fuction.
+            Chroma query function requries dictionary contains number of results matched.
+        """
+        converted_array = []
+        if data[column_name] is None:
+            for _ in range(item_length):
+                converted_array.append('')
+        else:
+            for list_items in data[column_name]:
+                for item in list_items:
+                    if item is None:
+                        converted_array.append('')
+                    else:
+                        converted_array.append(item)
+        return converted_array
+
+    def load(
+        self,
+        n_results: int,
+        collection: str = None,
+        query_embeddings: str = None,
+        query_texts: str = None,
+        **kwargs,
+    ) -> DataFrame:
+        """
+        Loads the data from Chroma with embeddings or query_texts.
+
+        Args:
+            n_results (int): Number of results to return.
+            query_embeddings (str): Embeddings to query.
+            query_texts (str): Texts to query.
+
+        Returns:
+            DataFrame: Data frame object loaded from the chroma query fuction.
+            Chroma query function requries dictionary contains number of results matched.
+        """
+        collection_client = self.client.get_or_create_collection(
+                name=self.collection if collection is None else collection)
+        # Chroma supports query by embeddings or query texts
+        if query_embeddings is not None:
+            data = collection_client.query(
+                query_embeddings=query_embeddings,
+                n_results=n_results
+            )
+        elif query_texts is not None:
+            data = collection_client.query(
+                query_texts=query_texts,
+                n_results=n_results
+            )
+        else:
+            raise ValueError("query_embeddings or query_texts must be specified")
+        # Flatten the query results into multiple rows
+        flatten_data = {}
+        item_length = len(data['ids'][0])
+        flatten_data['ids'] = self.__convert_to_multiple_rows(
+            data, 'ids', item_length)
+        flatten_data['distances'] = self.__convert_to_multiple_rows(
+            data, 'distances', item_length)
+        flatten_data['metadatas'] = self.__convert_to_multiple_rows(
+            data, 'metadatas', item_length)
+        flatten_data['embeddings'] = self.__convert_to_multiple_rows(
+            data, 'embeddings', item_length)
+        flatten_data['documents'] = self.__convert_to_multiple_rows(
+            data, 'documents', item_length)
+        flatten_data['data'] = self.__convert_to_multiple_rows(
+            data, 'data', item_length)
+        return DataFrame.from_dict(flatten_data)
+
+    def export(
+        self,
+        df: DataFrame,
+        document_column: str,
+        collection: str = None,
+        id_column: str = None,
+        metadata_column: str = None,
+        **kwargs,
+    ) -> None:
+        """
+        Exports the input dataframe to the chroma collection.
+        The column required is document_column.
+        It contains the contextual prompt for query and search.
+
+        Args:
+            df (DataFrame): Data frame to export.
+            document_column (str): name of the document.
+        """
+        docs = df[document_column].tolist()
+        if id_column is None:
+            ids = [str(x) for x in df.index.tolist()]
+        else:
+            ids = df[id_column].apply(str).tolist()
+        collection_client = self.client.get_or_create_collection(
+                name=self.collection if collection is None else collection)
+        insert_kwargs = dict(
+            documents=docs,
+            ids=ids,
+        )
+        if metadata_column:
+            insert_kwargs['metadatas'] = df[metadata_column].tolist()
+        return collection_client.add(
+            **insert_kwargs
+        )

mage_ai/io/config.py

@@ -27,6 +27,9 @@ class ConfigKey(str, Enum):
     AZURE_STORAGE_ACCOUNT_NAME = 'AZURE_STORAGE_ACCOUNT_NAME'
     AZURE_TENANT_ID = 'AZURE_TENANT_ID'
 
+    CHROMA_COLLECTION = 'CHROMA_COLLECTION'
+    CHROMA_PATH = 'CHROMA_PATH'
+
     CLICKHOUSE_DATABASE = 'CLICKHOUSE_DATABASE'
     CLICKHOUSE_HOST = 'CLICKHOUSE_HOST'
     CLICKHOUSE_INTERFACE = 'CLICKHOUSE_INTERFACE'
@@ -323,6 +326,7 @@ class VerboseConfigKey(str, Enum):
 
     AWS = 'AWS'
     BIGQUERY = 'BigQuery'
+    CHROMA = 'Chroma'
     CLICKHOUSE = 'ClickHouse'
     DRUID = 'Druid'
     DUCKDB = 'Duck DB'
@@ -370,6 +374,10 @@ class ConfigFileLoader(BaseConfigLoader):
             VerboseConfigKey.REDSHIFT,
             'user',
         ),
+        ConfigKey.CHROMA_COLLECTION: (
+            VerboseConfigKey.CHROMA, 'collection'),
+        ConfigKey.CHROMA_PATH: (
+            VerboseConfigKey.CHROMA, 'path'),
         ConfigKey.CLICKHOUSE_DATABASE: (
             VerboseConfigKey.CLICKHOUSE, 'database'),
         ConfigKey.CLICKHOUSE_HOST: (

mage_ai/orchestration/db/models/schedules.py

@@ -1,12 +1,12 @@
 import asyncio
+import collections
 import enum
 import traceback
 import uuid
 from datetime import datetime, timedelta, timezone
-from itertools import groupby
 from math import ceil
 from statistics import stdev
-from typing import Dict, List
+from typing import DefaultDict, Dict, List
 
 import dateutil.parser
 import pytz
@@ -119,6 +119,18 @@ class PipelineSchedule(BaseModel):
         query = query.filter(PipelineRun.pipeline_schedule_id.in_(ids))
         return query.all()
 
+    @classmethod
+    def fetch_latest_pipeline_runs_without_retries(self, ids: List[int]) -> List:
+        query = PipelineRun.query
+        query.cache = True
+        query = (
+            query.
+            filter(PipelineRun.pipeline_schedule_id.in_(ids)).
+            group_by(PipelineRun.execution_date).
+            order_by(func.max(PipelineRun.started_at))
+        )
+        return query.all()
+
     def get_settings(self) -> 'SettingsConfig':
         settings = self.settings if self.settings else dict()
         return SettingsConfig.load(config=settings)
@@ -491,46 +503,61 @@ class PipelineSchedule(BaseModel):
 
         return (self.settings or {}).get('landing_time_enabled', False)
 
-    def runtime_history(
+    def recently_completed_pipeline_runs(
         self,
         pipeline_run=None,
         sample_size: int = None,
-    ) -> List[float]:
-        sample_size_to_use = sample_size if sample_size else 7
-        previous_runtimes = []
+    ):
+        pipeline_runs = (
+            PipelineRun.
+            query.
+            filter(
+                PipelineRun.pipeline_schedule_id == self.id,
+                PipelineRun.status == PipelineRun.PipelineRunStatus.COMPLETED,
+            )
+        )
 
         if pipeline_run:
-            previous_runtimes += (pipeline_run.metrics or {}).get('previous_runtimes', [])
-
-        if len(previous_runtimes) < sample_size_to_use - 1 if pipeline_run else sample_size_to_use:
             pipeline_runs = (
-                PipelineRun.
-                query.
+                pipeline_runs.
                 filter(
-                    PipelineRun.pipeline_schedule_id == self.id,
-                    PipelineRun.status == PipelineRun.PipelineRunStatus.COMPLETED,
+                    PipelineRun.id != pipeline_run.id,
                 )
             )
 
-            if pipeline_run:
-                pipeline_runs = (
-                    pipeline_runs.
-                    filter(
-                        PipelineRun.id != pipeline_run.id,
-                    )
-                )
+        pipeline_runs = (
+            pipeline_runs.
+            order_by(PipelineRun.execution_date.desc())
+        )
 
-            pipeline_runs = (
-                pipeline_runs.
-                order_by(PipelineRun.execution_date.desc()).
-                limit(sample_size_to_use).
-                all()
-            )
+        if sample_size:
+            pipeline_runs = pipeline_runs.limit(sample_size)
+
+        pipeline_runs = pipeline_runs.all()
+
+        pipeline_runs = sorted(
+            pipeline_runs,
+            key=lambda pr: pr.execution_date,
+            reverse=True,
+        )
+
+        return pipeline_runs
+
+    def runtime_history(
+        self,
+        pipeline_run=None,
+        sample_size: int = None,
+    ) -> List[float]:
+        sample_size_to_use = sample_size if sample_size else 7
+        previous_runtimes = []
+
+        if pipeline_run:
+            previous_runtimes += (pipeline_run.metrics or {}).get('previous_runtimes', [])
 
-            pipeline_runs = sorted(
-                pipeline_runs,
-                key=lambda pr: pr.execution_date,
-                reverse=True,
+        if len(previous_runtimes) < sample_size_to_use - 1 if pipeline_run else sample_size_to_use:
+            pipeline_runs = self.recently_completed_pipeline_runs(
+                pipeline_run=pipeline_run,
+                sample_size=sample_size_to_use,
             )
 
             for pr in pipeline_runs:
@@ -667,6 +694,43 @@ class PipelineRun(BaseModel):
             repo_config=self.pipeline.repo_config,
         ).get_logs()
 
+    @classmethod
+    def recently_completed_pipeline_runs(
+        self,
+        pipeline_uuid: str,
+        pipeline_run_id: int = None,
+        pipeline_schedule_id: int = None,
+        sample_size: int = None,
+    ):
+        pipeline_runs = (
+            self.
+            query.
+            filter(
+                self.pipeline_uuid == pipeline_uuid,
+                self.status == self.PipelineRunStatus.COMPLETED,
+            )
+        )
+
+        if pipeline_run_id is not None:
+            pipeline_runs = pipeline_runs.filter(self.id != pipeline_run_id)
+
+        if pipeline_schedule_id is not None:
+            pipeline_runs = pipeline_runs.filter(self.pipeline_schedule_id == pipeline_schedule_id)
+
+        pipeline_runs = pipeline_runs.order_by(PipelineRun.execution_date.desc())
+
+        if sample_size:
+            pipeline_runs = pipeline_runs.limit(sample_size)
+
+        pipeline_runs = pipeline_runs.all()
+        pipeline_runs = sorted(
+            pipeline_runs,
+            key=lambda pr: pr.execution_date,
+            reverse=True,
+        )
+
+        return pipeline_runs
+
     async def logs_async(self):
         return await LoggerManagerFactory.get_logger_manager(
             pipeline_uuid=self.pipeline_uuid,
@@ -838,6 +902,77 @@ class PipelineRun(BaseModel):
 
         return executable_block_runs
 
+    def update_block_run_statuses(self, block_runs: List['BlockRun']) -> None:
+        """Update the statuses of the block runs to CONDITION_FAILED or UPSTREAM_FAILED.
+
+        This method updates the statuses of the block runs based on the pipeline run's block runs.
+        It retrieves the block UUIDs for failed block runs and conditionally failed block runs.
+        It maps the block run statuses to their corresponding block UUIDs.
+
+        The method iterates overthe provided block runs and checks if their dynamic upstream block
+        UUIDs or upstream block UUIDs match the failed or conditionally failed block UUIDs.
+        * If there is a match, the block run's status is updated accordingly.
+        * If no updates are made for a block run, it is added to the list of not updated block runs.
+
+        The method refreshes the pipeline run and continues iterating through block runs until no
+        more updates can be made.
+
+        Args:
+            block_runs (List[BlockRun]): A list of block runs to update.
+
+        Returns:
+            None
+        """
+        pipeline = self.pipeline
+
+        failed_block_uuids = set(
+            b.block_uuid for b in self.block_runs
+            if b.status in [
+                BlockRun.BlockRunStatus.UPSTREAM_FAILED,
+                BlockRun.BlockRunStatus.FAILED,
+            ]
+        )
+        condition_failed_block_uuids = set(
+            b.block_uuid for b in self.block_runs
+            if b.status in [
+                BlockRun.BlockRunStatus.CONDITION_FAILED,
+            ]
+        )
+
+        statuses = {
+            BlockRun.BlockRunStatus.CONDITION_FAILED: condition_failed_block_uuids,
+            BlockRun.BlockRunStatus.UPSTREAM_FAILED: failed_block_uuids,
+        }
+        not_updated_block_runs = []
+        for block_run in block_runs:
+            updated_status = False
+            dynamic_upstream_block_uuids = block_run.metrics and block_run.metrics.get(
+                'dynamic_upstream_block_uuids',
+            )
+
+            for status, block_uuids in statuses.items():
+                upstream_block_uuids = []
+                if dynamic_upstream_block_uuids:
+                    upstream_block_uuids = dynamic_upstream_block_uuids
+                else:
+                    block = pipeline.get_block(block_run.block_uuid)
+                    if block:
+                        upstream_block_uuids = block.upstream_block_uuids
+                if any(
+                    b in block_uuids
+                    for b in upstream_block_uuids
+                ):
+                    block_run.update(status=status)
+                    updated_status = True
+
+            if not updated_status:
+                not_updated_block_runs.append(block_run)
+
+        self.refresh()
+        # keep iterating through block runs until no more updates can be made
+        if len(block_runs) != len(not_updated_block_runs):
+            self.update_block_run_statuses(not_updated_block_runs)
+
     @classmethod
     @safe_db_query
     def active_runs_for_pipelines(
@@ -865,7 +1000,7 @@ class PipelineRun(BaseModel):
         self,
         pipeline_uuids: List[str],
         include_block_runs: bool = False,
-    ) -> Dict[str, List['PipelineRun']]:
+    ) -> DefaultDict[str, List['PipelineRun']]:
         """
         Get a dictionary of active pipeline runs grouped by pipeline uuid.
         """
@@ -874,9 +1009,9 @@ class PipelineRun(BaseModel):
             pipeline_uuids,
             include_block_runs=include_block_runs,
         )
-        grouped = {}
-        for key, runs in groupby(active_runs, key=lambda x: x.pipeline_uuid):
-            grouped[key] = list(runs)
+        grouped = collections.defaultdict(list)
+        for run in active_runs:
+            grouped[run.pipeline_uuid].append(run)
         return grouped
 
     @classmethod

mage_ai/orchestration/pipeline_scheduler.py

@@ -1,8 +1,8 @@
 import asyncio
+import collections
 import os
 import traceback
 from datetime import datetime, timedelta
-from itertools import groupby
 from typing import Any, Dict, List, Set, Tuple
 
 import pytz
@@ -75,11 +75,15 @@ class PipelineScheduler:
         # Get the list of integration stream if the pipeline is data integration pipeline
         self.streams = []
         if self.pipeline.type == PipelineType.INTEGRATION:
-            self.streams = self.pipeline.streams(
-                self.pipeline_run.get_variables(
-                    extra_variables=get_extra_variables(self.pipeline)
+            try:
+                self.streams = self.pipeline.streams(
+                    self.pipeline_run.get_variables(
+                        extra_variables=get_extra_variables(self.pipeline)
+                    )
                 )
-            )
+            except Exception:
+                logger.exception(f'Fail to get streams for {pipeline_run}')
+                traceback.print_exc()
 
         # Initialize the logger
         self.logger_manager = LoggerManagerFactory.get_logger_manager(
@@ -192,6 +196,19 @@ class PipelineScheduler:
         if PipelineType.STREAMING == self.pipeline.type:
             self.__schedule_pipeline()
         else:
+            schedule = PipelineSchedule.get(
+                self.pipeline_run.pipeline_schedule_id,
+            )
+            backfills = schedule.backfills if schedule else []
+            backfill = backfills[0] if len(backfills) >= 1 else None
+
+            if backfill is not None and \
+                backfill.status == Backfill.Status.INITIAL and \
+                    self.pipeline_run.status == PipelineRun.PipelineRunStatus.RUNNING:
+                backfill.update(
+                    status=Backfill.Status.RUNNING,
+                )
+
             if self.pipeline_run.all_blocks_completed(self.allow_blocks_to_fail):
                 if PipelineType.INTEGRATION == self.pipeline.type:
                     tags = self.build_tags()
@@ -221,17 +238,20 @@ class PipelineScheduler:
 
                 self.logger_manager.output_logs_to_destination()
 
-                schedule = PipelineSchedule.get(
-                    self.pipeline_run.pipeline_schedule_id,
-                )
-
                 if schedule:
-                    backfills = schedule.backfills
-                    # When all pipeline runs that are associated with backfill is done
-                    if len(backfills) >= 1:
-                        backfill = backfills[0]
+                    if backfill is not None:
+                        """
+                        Exclude old pipeline run retries associated with the backfill
+                        (if a backfill's runs had failed and the backfill was retried, those
+                        previous runs are no longer relevant) and check if the backfill's
+                        latest pipeline runs with different execution dates were successfull.
+                        """
+                        latest_pipeline_runs = \
+                            PipelineSchedule.fetch_latest_pipeline_runs_without_retries(
+                                [backfill.pipeline_schedule_id]
+                            )
                         if all([PipelineRun.PipelineRunStatus.COMPLETED == pr.status
-                                for pr in backfill.pipeline_runs]):
+                                for pr in latest_pipeline_runs]):
                             backfill.update(
                                 completed_at=datetime.now(tz=pytz.UTC),
                                 status=Backfill.Status.COMPLETED,
@@ -251,6 +271,20 @@ class PipelineScheduler:
                 self.pipeline_run.update(
                     status=PipelineRun.PipelineRunStatus.FAILED)
 
+                # Backfill status updated to "failed" if at least 1 of its pipeline runs failed
+                if backfill is not None:
+                    latest_pipeline_runs = \
+                        PipelineSchedule.fetch_latest_pipeline_runs_without_retries(
+                            [backfill.pipeline_schedule_id]
+                        )
+                    if any(
+                        [PipelineRun.PipelineRunStatus.FAILED == pr.status
+                            for pr in latest_pipeline_runs]
+                    ):
+                        backfill.update(
+                            status=Backfill.Status.FAILED,
+                        )
+
                 asyncio.run(UsageStatisticLogger().pipeline_run_ended(self.pipeline_run))
 
                 failed_block_runs = self.pipeline_run.failed_block_runs
@@ -472,75 +506,6 @@ class PipelineScheduler:
                 pass
         return any_block_run_timed_out
 
-    def __update_block_run_statuses(self, block_runs: List[BlockRun]) -> None:
-        """Update the statuses of the block runs to CONDITION_FAILED or UPSTREAM_FAILED.
-
-        This method updates the statuses of the block runs based on the pipeline run's block runs.
-        It retrieves the block UUIDs for failed block runs and conditionally failed block runs.
-        It maps the block run statuses to their corresponding block UUIDs.
-
-        The method iterates overthe provided block runs and checks if their dynamic upstream block
-        UUIDs or upstream block UUIDs match the failed or conditionally failed block UUIDs.
-        * If there is a match, the block run's status is updated accordingly.
-        * If no updates are made for a block run, it is added to the list of not updated block runs.
-
-        The method refreshes the pipeline run and continues iterating through block runs until no
-        more updates can be made.
-
-        Args:
-            block_runs (List[BlockRun]): A list of block runs to update.
-
-        Returns:
-            None
-        """
-        failed_block_uuids = set(
-            b.block_uuid for b in self.pipeline_run.block_runs
-            if b.status in [
-                BlockRun.BlockRunStatus.UPSTREAM_FAILED,
-                BlockRun.BlockRunStatus.FAILED,
-            ]
-        )
-        condition_failed_block_uuids = set(
-            b.block_uuid for b in self.pipeline_run.block_runs
-            if b.status in [
-                BlockRun.BlockRunStatus.CONDITION_FAILED,
-            ]
-        )
-
-        statuses = {
-            BlockRun.BlockRunStatus.CONDITION_FAILED: condition_failed_block_uuids,
-            BlockRun.BlockRunStatus.UPSTREAM_FAILED: failed_block_uuids,
-        }
-        not_updated_block_runs = []
-        for block_run in block_runs:
-            updated_status = False
-            dynamic_upstream_block_uuids = block_run.metrics and block_run.metrics.get(
-                'dynamic_upstream_block_uuids',
-            )
-
-            for status, block_uuids in statuses.items():
-                upstream_block_uuids = []
-                if dynamic_upstream_block_uuids:
-                    upstream_block_uuids = dynamic_upstream_block_uuids
-                else:
-                    block = self.pipeline.get_block(block_run.block_uuid)
-                    if block:
-                        upstream_block_uuids = block.upstream_block_uuids
-                if any(
-                    b in block_uuids
-                    for b in upstream_block_uuids
-                ):
-                    block_run.update(status=status)
-                    updated_status = True
-
-            if not updated_status:
-                not_updated_block_runs.append(block_run)
-
-        self.pipeline_run.refresh()
-        # keep iterating through block runs until no more updates can be made
-        if len(block_runs) != len(not_updated_block_runs):
-            self.__update_block_run_statuses(not_updated_block_runs)
-
     def __schedule_blocks(self, block_runs: List[BlockRun] = None) -> None:
         """Schedule the block runs for execution.
 
@@ -557,7 +522,7 @@ class PipelineScheduler:
         Returns:
             None
         """
-        self.__update_block_run_statuses(self.pipeline_run.initial_block_runs)
+        self.pipeline_run.update_block_run_statuses(self.pipeline_run.initial_block_runs)
         if block_runs is None:
             block_runs_to_schedule = self.pipeline_run.executable_block_runs(
                 allow_blocks_to_fail=self.allow_blocks_to_fail,
@@ -1417,10 +1382,9 @@ def schedule_all():
     active_pipeline_uuids = list(set([s.pipeline_uuid for s in active_pipeline_schedules]))
     pipeline_runs_by_pipeline = PipelineRun.active_runs_for_pipelines_grouped(active_pipeline_uuids)
 
-    pipeline_schedules_by_pipeline = {
-        key: list(runs)
-        for key, runs in groupby(active_pipeline_schedules, key=lambda x: x.pipeline_uuid)
-    }
+    pipeline_schedules_by_pipeline = collections.defaultdict(list)
+    for schedule in active_pipeline_schedules:
+        pipeline_schedules_by_pipeline[schedule.pipeline_uuid].append(schedule)
 
     # Iterate through pipeline schedules by pipeline to handle pipeline run limits for
     # each pipeline.
@@ -1546,7 +1510,13 @@ def schedule_all():
             )
 
         for r in quota_filtered_runs:
-            PipelineScheduler(r).start()
+            try:
+                PipelineScheduler(r).start()
+            except Exception:
+                logger.exception(f'Failed to start {r}')
+                traceback.print_exc()
+                r.update(status=PipelineRun.PipelineRunStatus.FAILED)
+                continue
 
         # If on_pipeline_run_limit_reached is set as SKIP, cancel the pipeline runs that
         # were not scheduled due to pipeline run limits.

mage_ai/server/constants.py

@@ -12,4 +12,4 @@ DATAFRAME_OUTPUT_SAMPLE_COUNT = 10
 # Dockerfile depends on it because it runs ./scripts/install_mage.sh and uses
 # the last line to determine the version to install.
 VERSION = \
-'0.9.46'
+'0.9.48'