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

mage_ai/data_preparation/models/block/dynamic/factory.py

@@ -0,0 +1,163 @@
+import time
+from functools import reduce
+from typing import Dict, List, Optional
+
+from mage_ai.data_preparation.models.block.dynamic.counter import DynamicItemCounter
+from mage_ai.data_preparation.models.block.dynamic.wrappers import (
+    DynamicBlockWrapperBase,
+)
+from mage_ai.data_preparation.models.block.settings.dynamic.constants import (
+    DEFAULT_STREAM_POLL_INTERVAL,
+    ModeType,
+)
+from mage_ai.orchestration.db.models.schedules import BlockRun
+
+
+class DynamicBlockFactory(DynamicBlockWrapperBase):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._counters = None
+        self._item_count = None
+
+    def is_complete(self) -> bool:
+        upstream_blocks_count = len(self.block.upstream_blocks)
+
+        completed_block_runs = len([
+            block_run
+            for block_run in self.__fetch_cloned_block_runs().values()
+            if block_run.status == BlockRun.BlockRunStatus.COMPLETED
+        ])
+
+        # If item count is 0 because the upstream blocks haven’t output anything yet
+        item_count_total = max(self.__calculate_item_count(), upstream_blocks_count)
+
+        upstream_block_runs = self.__upstream_block_runs()
+        upstream_block_runs_completed = len([
+            block_run
+            for block_run in upstream_block_runs
+            if block_run.status == BlockRun.BlockRunStatus.COMPLETED
+        ])
+
+        completed = (
+            upstream_block_runs_completed >= upstream_blocks_count
+            and completed_block_runs >= item_count_total
+        )
+
+        if not completed:
+            if self.block_run is not None:
+                upstream_block_uuids = ', '.join([b.uuid for b in self.block.upstream_blocks])
+                message = (
+                    f'Dynamic block factory for {self.block.uuid} '
+                    f'(block run ID {self.block_run_id}) is still working: '
+                    f'{completed_block_runs} of at least {item_count_total} '
+                    'dynamic child blocks completed. '
+                    f'Upstream blocks: {upstream_block_runs_completed} completed of '
+                    f'{len(upstream_block_runs)} block runs ({upstream_block_uuids})'
+                )
+                if self.logger:
+                    self.logger.info(message)
+                else:
+                    print(message)
+            self.__sleep()
+
+        return completed
+
+    def execute_sync(
+        self,
+        execution_partition: Optional[str] = None,
+        logging_tags: Optional[Dict] = None,
+        **kwargs,
+    ) -> List[Dict]:
+        cloned_block_runs = self.__fetch_cloned_block_runs()
+        runs_count = len(cloned_block_runs)
+        item_count = self.__calculate_item_count()
+
+        counts = ', '.join([
+            f'{block_uuid}={counter.item_count() if counter is not None else 1}'
+            for block_uuid, counter in self.__counters.items()
+        ])
+        print(
+            f'[DynamicBlockFactory.item_count] {self.block_run().block_uuid}: {item_count} '
+            f'items ({counts})'
+        )
+
+        block_runs = []
+        pipeline_run = self.pipeline_run()
+
+        for dynamic_block_index in range(runs_count, item_count):
+            block_run_uuid = ':'.join([self.block.uuid, str(dynamic_block_index)])
+            block_run = pipeline_run.create_block_run(
+                block_run_uuid,
+                metrics=dict(dynamic_block_index=dynamic_block_index),
+                skip_if_exists=True,
+            )
+            block_runs.append(block_run)
+
+        return block_runs
+
+    @property
+    def __counters(self) -> Dict[str, DynamicItemCounter]:
+        if self._counters is None:
+            self._counters = {}
+            for upstream_block in self.block.upstream_blocks:
+                counter = DynamicItemCounter.build_counter(
+                    upstream_block,
+                    downstream_block=self.block,
+                    partition=self.execution_partition,
+                )
+                if counter is not None:
+                    self._counters[upstream_block.uuid] = counter
+
+        return self._counters
+
+    def __calculate_item_count(self) -> int:
+        return reduce(
+            lambda a, b: a * b,
+            # Multiply by 0 so that no block runs are created until all
+            # upstream blocks have at least 1 output.
+            [
+                counter.item_count() if counter is not None else 1
+                for counter in self.__counters.values()
+            ],
+            1,
+        )
+
+    def __fetch_cloned_block_runs(self) -> Dict[str, BlockRun]:
+        cloned_block_runs = {}
+        for block_run in self.block_runs():
+            if block_run.block_uuid == self.block.uuid or (
+                self.block.uuid_replicated is not None
+                and block_run.block_uuid == self.block.uuid_replicated
+            ):
+                continue
+
+            block = self.pipeline.get_block(block_run.block_uuid)
+            if block and block.uuid == self.block.uuid:
+                block_run_previous = cloned_block_runs.get(block_run.block_uuid)
+                if (
+                    block_run_previous is None
+                    or block_run.started_at > block_run_previous.started_at
+                ):
+                    cloned_block_runs[block_run.block_uuid] = block_run
+        return cloned_block_runs
+
+    def __upstream_block_runs(self) -> List[BlockRun]:
+        upstream_block_uuids = [
+            upstream_block.uuid_replicated
+            if upstream_block.replicated_block
+            else upstream_block.uuid
+            for upstream_block in self.block.upstream_blocks
+        ]
+        return [
+            block_run
+            for block_run in self.block_runs()
+            if block_run.block_uuid in upstream_block_uuids
+        ]
+
+    def __sleep(self) -> None:
+        mode_settings = self.block.settings_for_mode(ModeType.STREAM)
+        poll_interval = mode_settings.poll_interval if mode_settings is not None else None
+        poll_interval = (
+            poll_interval if poll_interval is not None else DEFAULT_STREAM_POLL_INTERVAL
+        )
+        time.sleep(poll_interval)

mage_ai/data_preparation/models/block/dynamic/models.py

@@ -0,0 +1,19 @@
+from dataclasses import dataclass, field
+from typing import Any, List
+
+from botocore.tokens import Optional
+
+from mage_ai.shared.models import BaseDataClass
+
+
+@dataclass
+class DynamicRun(BaseDataClass):
+    block: Any
+    block_run: Optional[Any] = None
+    clones: List[Any] = field(default_factory=list)
+
+    def add(self, block_run):
+        if block_run.block_uuid in [self.block.uuid_replicated, self.block.uuid]:
+            self.block_run = block_run
+        else:
+            self.clones.append(block_run)

mage_ai/data_preparation/models/block/dynamic/shared.py

@@ -0,0 +1,92 @@
+from typing import Optional
+
+
+def is_dynamic_block(block) -> bool:
+    """
+    Checks if the given block is a dynamic block.
+
+    Args:
+        block: The block.
+
+    Returns:
+        bool: True if the block is a dynamic block, False otherwise.
+    """
+    if block is None or block.configuration is None:
+        return False
+
+    if block.is_dynamic_v2:
+        return block.is_dynamic_parent
+
+    dynamic = block.configuration.get('dynamic')
+    if isinstance(dynamic, bool):
+        return dynamic
+
+    return False
+
+
+def should_reduce_output(block) -> bool:
+    """
+    Checks if the given block should reduce its output.
+
+    Args:
+        block: The block.
+
+    Returns:
+        bool: True if the block should reduce its output, False otherwise.
+    """
+    if block is None or block.configuration is None:
+        return False
+
+    if block.is_dynamic_v2:
+        return block.should_reduce_output
+
+    reduce_output = block.configuration.get('reduce_output')
+    if isinstance(reduce_output, bool):
+        return reduce_output
+
+    return False
+
+
+def has_reduce_output_from_upstreams(block) -> bool:
+    return any([should_reduce_output(upstream_block) for upstream_block in block.upstream_blocks])
+
+
+def has_dynamic_block_upstream_parent(block) -> bool:
+    return block.upstream_blocks and any([is_dynamic_block(b) for b in block.upstream_blocks])
+
+
+def is_dynamic_block_child(block, include_reduce_output: Optional[bool] = None) -> bool:
+    """
+    Checks if the given block is a dynamic block child.
+
+    Args:
+        block: The block.
+
+    Returns:
+        bool: True if the block is a dynamic block child, False otherwise.
+    """
+    if not block:
+        return False
+
+    if block.is_dynamic_v2:
+        return block.is_dynamic_child
+
+    dynamic_or_child = []
+
+    for upstream_block in block.upstream_blocks:
+        if is_dynamic_block(upstream_block) or is_dynamic_block_child(
+            upstream_block, include_reduce_output=include_reduce_output
+        ):
+            dynamic_or_child.append(upstream_block)
+
+    if len(dynamic_or_child) == 0:
+        return False
+
+    if include_reduce_output:
+        return len(dynamic_or_child) >= 1
+
+    dynamic_or_child_with_reduce = list(
+        filter(lambda x: should_reduce_output(x), dynamic_or_child)
+    )
+
+    return len(dynamic_or_child) > len(dynamic_or_child_with_reduce)