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

mage_ai/cluster_manager/aws/emr_cluster_manager.py

@@ -8,7 +8,6 @@ from mage_ai.data_preparation.repo_manager import get_repo_config
 from mage_ai.services.aws.emr.config import EmrConfig
 from mage_ai.services.aws.emr.emr import describe_cluster, list_clusters
 from mage_ai.services.aws.emr.launcher import create_cluster
-from mage_ai.settings.repo import get_repo_path
 from mage_ai.shared.array import find
 from mage_ai.shared.hash import merge_dict
 
@@ -37,12 +36,12 @@ class EmrClusterManager(ClusterManager):
             is_active=c['Id'] == self.active_cluster_id,
         ) for c in valid_clusters]
 
-    def create_cluster(self, emr_config: Dict = None):
+    def create_cluster(self, repo_path: str, emr_config: Dict = None):
         emr_config = EmrConfig.load(
             config=merge_dict(get_repo_config().emr_config or dict(), emr_config or dict()))
 
         return create_cluster(
-            get_repo_path(),
+            repo_path,
             done_status=None,
             emr_config=emr_config,
             tags=dict(name=CLUSTER_NAME),
@@ -51,6 +50,7 @@ class EmrClusterManager(ClusterManager):
 
     def set_active_cluster(
         self,
+        repo_path: str = None,
         auto_creation: bool = True,
         auto_selection: bool = False,
         cluster_id=None,
@@ -61,13 +61,17 @@ class EmrClusterManager(ClusterManager):
             if len(clusters) > 0:
                 cluster_id = clusters[0]['id']
             elif auto_creation:
-                self.create_cluster(emr_config=emr_config)
+                self.create_cluster(repo_path, emr_config=emr_config)
         if cluster_id is None:
             return
 
         self.active_cluster_id = cluster_id
         emr_config = EmrConfig.load(
-            config=merge_dict(get_repo_config().emr_config or dict(), emr_config or dict()))
+            config=merge_dict(
+                get_repo_config(repo_path=repo_path).emr_config or dict(),
+                emr_config or dict(),
+            )
+        )
 
         # Fetch cluster master instance public DNS
         cluster_info = describe_cluster(cluster_id)

mage_ai/cluster_manager/config.py

@@ -65,7 +65,7 @@ class PostStart(BaseConfig):
         hook_path = config.get('hook_path')
         if hook_path and not os.path.isabs(hook_path):
             config['hook_path'] = os.path.join(
-                get_repo_path(), hook_path
+                get_repo_path(root_project=True), hook_path
             )
 
         return config
@@ -82,7 +82,7 @@ class LifecycleConfig(BaseConfig):
         pre_start_script_path = config.get('pre_start_script_path')
         if pre_start_script_path and not os.path.isabs(pre_start_script_path):
             config['pre_start_script_path'] = os.path.join(
-                get_repo_path(), pre_start_script_path
+                get_repo_path(root_project=True), pre_start_script_path
             )
 
         return config

mage_ai/cluster_manager/kubernetes/workload_manager.py

@@ -477,6 +477,57 @@ class WorkloadManager:
 
         return k8s_service
 
+    def patch_workload(
+        self,
+        name: str,
+        workspace_config: KubernetesWorkspaceConfig,
+        update_workspace_settings: bool = False,
+    ) -> None:
+        """
+        Update workload for k8s. Currently the only fields that can be updated are
+        container_config and workspace settings.
+
+        Args:
+            name (str): name of the workload/workspace
+            workspace_config (KubernetesWorkspaceConfig): new workspace config
+            update_workspace_settings (bool): whether to update workspace settings with
+                the current pod's environment variables.
+        """
+        container_config_yaml = workspace_config.container_config
+        container_config = dict()
+        if isinstance(container_config_yaml, str):
+            container_config = yaml.full_load(container_config_yaml)
+
+        if update_workspace_settings:
+            env_vars = self.__populate_env_vars(
+                name,
+                container_config=container_config,
+            )
+            container_config['env'] = env_vars
+
+        mage_container_config = {
+            'name': f'{name}-container',
+            **container_config,
+        }
+
+        stateful_set_template_spec = {
+            'containers': [mage_container_config],
+        }
+
+        updated_stateful_set = {
+            'spec': {
+                'template': {
+                    'spec': stateful_set_template_spec,
+                },
+            },
+        }
+
+        self.apps_client.patch_namespaced_stateful_set(
+            name,
+            namespace=self.namespace,
+            body=updated_stateful_set,
+        )
+
     def add_service_to_ingress_paths(
         self,
         ingress_name: str,
@@ -712,7 +763,7 @@ class WorkloadManager:
     def __populate_env_vars(
         self,
         name,
-        project_type: str = 'standalone',
+        project_type: str = None,
         project_uuid: str = None,
         container_config: Dict = None,
         initial_metadata: Dict = None,

mage_ai/cluster_manager/manage.py

@@ -67,7 +67,7 @@ def get_workspaces(cluster_type: ClusterType, **kwargs) -> List[Workspace]:
     """
     is_main_project = get_project_type() == ProjectType.MAIN
 
-    repo_path = get_repo_path()
+    repo_path = get_repo_path(root_project=True)
     projects_folder = os.path.join(repo_path, 'projects')
     if is_main_project:
         # sort by file last modified time

mage_ai/cluster_manager/workspace/base.py

@@ -33,7 +33,7 @@ class Workspace(abc.ABC):
 
     @classproperty
     def project_folder(self) -> str:
-        return os.path.join(get_repo_path(), 'projects')
+        return os.path.join(get_repo_path(root_project=True), 'projects')
 
     @property
     def config_path(self) -> str:
@@ -148,6 +148,12 @@ class Workspace(abc.ABC):
         """
         raise NotImplementedError('Initialize method not implemented')
 
+    def update(self, payload: Dict, **kwargs):
+        """
+        Update the workspace configuration.
+        """
+        raise NotImplementedError('Update method not implemented')
+
     @abc.abstractmethod
     @safe_db_query
     def delete(self, **kwargs):

mage_ai/cluster_manager/workspace/kubernetes.py

@@ -1,4 +1,5 @@
 import os
+from typing import Dict
 
 import yaml
 
@@ -7,7 +8,7 @@ from mage_ai.cluster_manager.constants import KUBE_NAMESPACE, ClusterType
 from mage_ai.cluster_manager.kubernetes.workload_manager import WorkloadManager
 from mage_ai.cluster_manager.workspace.base import Workspace
 from mage_ai.data_preparation.repo_manager import ProjectType, get_project_type
-from mage_ai.shared.hash import merge_dict
+from mage_ai.shared.hash import extract, merge_dict
 
 
 class KubernetesWorkspace(Workspace):
@@ -82,6 +83,26 @@ class KubernetesWorkspace(Workspace):
 
         return cls(name)
 
+    def update(self, payload: Dict, **kwargs):
+        update_workspace_settings = payload.pop('update_workspace_settings', False)
+        extracted_payload = extract(payload, [
+            'container_config',
+        ])
+        updated_config = merge_dict(
+            self.config.to_dict(),
+            extracted_payload,
+        )
+        workspace_config = KubernetesWorkspaceConfig.load(
+            config=updated_config
+        )
+        self.workload_manager.patch_workload(
+            self.name,
+            workspace_config,
+            update_workspace_settings=update_workspace_settings,
+        )
+        with open(self.config_path, 'w', encoding='utf-8') as fp:
+            yaml.dump(workspace_config.to_dict(), fp)
+
     def delete(self, **kwargs):
         try:
             self.workload_manager.delete_workload(

mage_ai/command_center/applications/factory.py

@@ -21,7 +21,7 @@ class ApplicationFactory(BaseFactory):
     async def fetch_items(self, **kwargs) -> List[Dict]:
         items = []
 
-        for item_dict in (self.page_history or []):
+        for item_dict in self.page_history or []:
             item_scored = self.filter_score(item_dict)
             if item_scored:
                 items.append(item_scored)
@@ -70,12 +70,15 @@ class ApplicationFactory(BaseFactory):
                 if app:
                     status = (app.get('state') or {}).get('status')
                     if ApplicationExpansionStatus.MINIMIZED == status:
-                        pass
-
-                    self.filter_score_mutate_accumulator(
-                        await build_close_application(uuid),
-                        items,
-                    )
+                        self.filter_score_mutate_accumulator(
+                            await build_open_application(uuid),
+                            items,
+                        )
+                    else:
+                        self.filter_score_mutate_accumulator(
+                            await build_close_application(uuid),
+                            items,
+                        )
 
             if not app:
                 self.filter_score_mutate_accumulator(

mage_ai/command_center/applications/utils.py

@@ -17,7 +17,7 @@ from mage_ai.version_control.models import File
 
 """
 Icons:
-Arcane Library (file browser): ChurnV3
+Text Editor (file browser): ChurnV3
 Portal Terminal (terminal):    RankingV3
 Version Control (file diffs):  ForecastV3
 
@@ -65,7 +65,7 @@ async def build_application_arcane_library() -> Dict:
     return build_generic(model_class=File, item_dict=dict(
         item_type=ItemType.DETAIL,
         object_type=ObjectType.FILE,
-        title='Arcane Library',
+        title='Text editor',
         description='browse and edit files across all projects',
         subtitle='Browser / Editor',
         applications=[

mage_ai/command_center/files/factory.py

@@ -23,13 +23,13 @@ class FileFactory(BaseFactory):
         files = []
 
         async def build_and_score(line: str, files=files, factory=self):
-            full_path, file_size, modified_timestamp = line.split(',')
+            full_path, file_size, modified_timestamp = line.split(",")
             filename = os.path.basename(full_path)
             modified_at = datetime.fromtimestamp(float(modified_timestamp)).isoformat()
 
             path_dict = shorten_directory(full_path)
-            extension = path_dict.get('extension')
-            directory = path_dict.get('directory')
+            extension = path_dict.get("extension")
+            directory = path_dict.get("directory")
 
             item_dict = dict(
                 item_type=ItemType.DETAIL,
@@ -55,7 +55,6 @@ class FileFactory(BaseFactory):
                     ),
                 ),
             )
-
             scored = factory.filter_score(item_dict)
             if scored:
                 files.append(scored)
@@ -64,23 +63,26 @@ class FileFactory(BaseFactory):
             now = datetime.utcnow().timestamp()
             cache = FileCache.initialize_cache_with_settings()
             lines = await cache.load() or []
-            print(f'[FileFactory] Load files: {len(lines)} - {datetime.utcnow().timestamp() - now}')
+            print(
+                f"[FileFactory] Load files: {len(lines)} - {datetime.utcnow().timestamp() - now}"
+            )
 
             now = datetime.utcnow().timestamp()
-            await asyncio.gather(
-                *[build_and_score(item_dict) for item_dict in lines]
-            )
+            await asyncio.gather(*[build_and_score(item_dict) for item_dict in lines])
             print(
-                f'[FileFactory] Search {self.search}: '
-                f'{len(files)} - {datetime.utcnow().timestamp() - now}',
+                f"[FileFactory] Search {self.search}: "
+                f"{len(files)} - {datetime.utcnow().timestamp() - now}",
             )
 
             now = datetime.utcnow().timestamp()
             files = await self.rank_items(files)
-            files = [merge_dict(
-                item_dict,
-                add_application_actions(item_dict),
-            ) for item_dict in files]
-            print(f'[FileFactory] Rank items: {datetime.utcnow().timestamp() - now}')
+            files = [
+                merge_dict(
+                    item_dict,
+                    add_application_actions(item_dict),
+                )
+                for item_dict in files
+            ]
+            print(f"[FileFactory] Rank items: {datetime.utcnow().timestamp() - now}")
 
         return files[:10] + items

mage_ai/command_center/presenters/text.py

@@ -5,5 +5,5 @@ def application_title(uuid: ApplicationExpansionUUID) -> str:
     if ApplicationExpansionUUID.VersionControlFileDiffs == uuid:
         return 'version control file diffs'
     if ApplicationExpansionUUID.ArcaneLibrary == uuid:
-        return 'the Arcane Library'
+        return 'the Text Editor'
     return f'the {uuid} application'

mage_ai/command_center/utils.py

@@ -4,27 +4,39 @@ from typing import Dict
 
 
 def shorten_directory(full_path: str) -> Dict:
-    parts = Path(os.path.dirname(full_path)).parts
+    # Use Path object for simplicity and cross-platform compatibility
+    path_obj = Path(full_path)
+    parts = path_obj.parts
     parts_count = len(parts)
-    dir_name = ''
-    if parts_count >= 1:
-        os.path.join(*parts[max([0, parts_count - 3]):])
 
+    # Initialize dir_name as an empty string
+    dir_name = ""
+    # Handle directory part; ignore the filename in parts
+    if parts_count > 1:
+        # Join all parts except the last one (filename)
+        dir_name = os.path.join(*parts[: parts_count - 1])
+
+    # Simplifying the directory to show only relevant parts (3 levels up max)
     if parts_count >= 4:
-        if parts_count >= 5:
-            dir_name = os.path.join('.', dir_name)
+        if parts_count > 4:
+            # For deeply nested paths, show as relative path starting with ..
+            relevant_parts = parts[-3:]
+            dir_name = os.path.join("..", *relevant_parts)
         else:
-            dir_name = os.path.join('..', dir_name)
+            # If exactly 4 levels deep, show entire path
+            dir_name = os.path.join(*parts[:-1])
 
-        for _i in range(min(1, parts_count - 4)):
-            dir_name = os.path.join('..', dir_name)
+    # Extracting extension (without the dot)
+    extension = (
+        path_obj.suffix[1:] if path_obj.suffix.startswith(".") else path_obj.suffix
+    )
 
-    extension = Path(full_path or '').suffix
-    if extension:
-        extension = extension.replace('.', '')
+    # Include the filename in the output
+    filename = path_obj.name
 
     return dict(
         directory=dir_name,
         extension=extension,
-        parts=parts,
+        parts=parts[:-1],  # Exclude the filename from 'parts'
+        filename=filename,
     )

mage_ai/data/constants.py

@@ -0,0 +1,45 @@
+from typing import AsyncGenerator, Generator, Iterable, List, Union
+
+import polars as pl
+
+from mage_ai.data.models.generator import DataGenerator
+from mage_ai.data.models.pyarrow.record_batch import RecordBatch, TaggedRecordBatch
+from mage_ai.data_preparation.models.variables.constants import VariableType
+from mage_ai.shared.models import BaseEnum
+
+SUPPORTED_VARIABLE_TYPES = [
+    VariableType.DATAFRAME,
+    VariableType.POLARS_DATAFRAME,
+    VariableType.SERIES_PANDAS,
+    VariableType.SERIES_POLARS,
+]
+
+ScanBatchDatasetResult = Union[RecordBatch, TaggedRecordBatch, pl.DataFrame]
+OutputData = Union[
+    DataGenerator,
+    List[pl.DataFrame],
+    Iterable[ScanBatchDatasetResult],
+    List[ScanBatchDatasetResult],
+]
+
+RecordBatchGenerator = Union[Generator[ScanBatchDatasetResult, None, None], DataGenerator]
+
+AsyncRecordBatchGenerator = AsyncGenerator[ScanBatchDatasetResult, None]
+
+
+class InputDataType(BaseEnum):
+    # Batch settings will be used to fetch the data in batches,
+    # and execute the decorated function inside a yield loop.
+    BATCH = 'batch'
+    # @chunker, @chunking, @chunk
+    # CHUNKS = 'chunks' # This will be handled in the block code
+    # Original data will be passed as is.
+    DEFAULT = 'default'
+    # A generator to iterate over
+    GENERATOR = 'generator'
+    # Input variable will be a Callable function that loads the data; e.g. Variable object
+    READER = 'reader'
+
+
+class ReadModeType(BaseEnum):
+    MEMORY = 'memory'

mage_ai/data/models/base.py

@@ -0,0 +1,119 @@
+import os
+from pathlib import Path
+from typing import Any, List, Optional, Union
+
+from mage_ai.data.constants import InputDataType
+from mage_ai.data.models.constants import CHUNKS_DIRECTORY_NAME
+from mage_ai.data.models.utils import variable_type_supported
+from mage_ai.data.tabular.models import BatchSettings
+from mage_ai.data_preparation.models.variables.constants import VariableType
+from mage_ai.data_preparation.storage.base_storage import BaseStorage
+from mage_ai.settings.repo import get_variables_dir
+from mage_ai.system.models import ResourceUsage
+from mage_ai.system.storage.utils import size_of_path
+
+
+class BaseData:
+    def __init__(
+        self,
+        storage: BaseStorage,
+        variable_dir_path: str,
+        variable_path: str,
+        batch_settings: Optional[BatchSettings] = None,
+        chunks: Optional[List[Any]] = None,
+        input_data_types: Optional[List[InputDataType]] = None,
+        poll_interval: Optional[int] = None,
+        uuid: Optional[str] = None,
+        variable_type: Optional[VariableType] = None,
+        variable_types: Optional[List[VariableType]] = None,
+        variables_dir: Optional[str] = None,
+    ):
+        self.storage = storage
+        """
+        ~/.mage_data/[project_name]/pipelines/[pipeline_uuid]/.variables
+            /[pipeline_run.id]/[execution_partition]/[block.uuid]
+        ~/.mage_data/unit_3_observability/pipelines/sklearn_training/.variables
+            /9/20240518T144726_954384/load_models
+        """
+        self.variable_dir_path = variable_dir_path
+        """
+        ~/.mage_data/[project_name]/pipelines/[pipeline_uuid]/.variables
+            /[pipeline_run.id]/[execution_partition]/[block.uuid]
+            /[variable_name]
+
+        ~/.mage_data/unit_3_observability/pipelines/sklearn_training/.variables
+            /9/20240518T144726_954384/load_models
+            /output_0
+        """
+        self.batch_settings = batch_settings
+        self.chunks = chunks
+        self.poll_interval = poll_interval
+        self.variable_path = variable_path
+        self.variable_type = variable_type
+        self.variable_types = variable_types or []
+        self.variables_dir = variables_dir or get_variables_dir(root_project=False)
+        self.uuid = uuid or str(Path(self.variable_dir_path).relative_to(Path(self.variables_dir)))
+
+    @property
+    def data_source(self) -> Union[str, List[str]]:
+        if self.variable_types and len(self.variable_types) >= 1:
+            return self.__data_source_paths
+        return self.data_source_directory_path
+
+    @property
+    def number_of_outputs(self) -> int:
+        return (
+            len(self.variable_types)
+            if self.variable_types and len(self.variable_types) >= 1
+            else 1
+        )
+
+    @property
+    def data_source_directory_path(self) -> str:
+        return self.__build_data_source_path()
+
+    @property
+    def __data_source_paths(self) -> List[str]:
+        if self.variable_types:
+            return [
+                self.__build_data_source_path(str(idx)) for idx in range(len(self.variable_types))
+            ]
+        return [self.data_source_directory_path]
+
+    def __build_data_source_path(self, subdirectory: Optional[str] = None) -> str:
+        return os.path.join(self.variable_path, subdirectory or '', CHUNKS_DIRECTORY_NAME)
+
+    def is_dataframe(self) -> bool:
+        return all(
+            variable_type
+            in [
+                VariableType.DATAFRAME,
+                VariableType.POLARS_DATAFRAME,
+                VariableType.SERIES_PANDAS,
+                VariableType.SERIES_POLARS,
+            ]
+            for variable_type in (self.variable_types or [self.variable_type])
+        )
+
+    @property
+    def resource_usage(self) -> ResourceUsage:
+        return ResourceUsage.combine(self.resource_usages)
+
+    @property
+    def resource_usages(self) -> List[ResourceUsage]:
+        return [
+            ResourceUsage.load(
+                directory=path,
+                size=size_of_path(path, file_extension='parquet'),
+            )
+            for path in self.__data_source_paths
+        ]
+
+    def supported(self, data: Optional[Any] = None) -> bool:
+        if self.variable_types and len(self.variable_types) >= 1:
+            return all(
+                variable_type_supported(variable_type, data)
+                for variable_type in self.variable_types
+            )
+
+        return variable_type_supported(self.variable_type, data)

mage_ai/data/models/constants.py

@@ -0,0 +1 @@
+CHUNKS_DIRECTORY_NAME = 'chunks'

mage_ai/data/models/generator.py

@@ -0,0 +1,115 @@
+from collections.abc import Generator
+from typing import Any, Callable, Dict, List, Optional
+
+
+class DataGenerator:
+    def __init__(
+        self,
+        data: Optional[Any] = None,
+        load_data: Optional[Callable[[int], Any]] = None,
+        measure_data: Optional[Callable[[int], int]] = None,
+    ):
+        self.data = data
+        self.index = 0
+        self.load_data = load_data
+        self.measure_data = measure_data
+        self._length = None
+
+    def __len__(self):
+        # Check if the total length has been determined already
+        if self._length is None:
+            # Logic to determine the length from the datasource
+            # This could involve querying the database,
+            # reading file info, etc., without loading all data
+            self._length = self._data_length(self.index)
+        return self._length
+
+    def __iter__(self):
+        if self.data is not None and hasattr(self.data, '__iter__'):
+            for item in self.data:
+                yield item
+        elif self.data is not None and hasattr(self.data, '__getitem__'):
+            for item in self.__sequence_iterator():
+                yield item
+        else:
+            while self.index < self._data_length(self.index):
+                item = None
+                if self.load_data:
+                    item = self.load_data(self.index)
+                elif self.data is not None:
+                    item = self.data[self.index]
+                self.index += 1
+                yield item
+        self.data = None  # Explicitly set data to None after iteration is finished
+
+    def __next__(self):
+        if self.index < self._data_length(self.index):
+            # Load the next item
+            item = None
+            if self.load_data:
+                item = self.load_data(self.index)
+            elif self.data is not None:
+                item = self.data[self.index]
+            self.index += 1
+            return item
+        else:
+            self.data = None  # Set data to None when the generator is exhausted
+            # No more data, stop iteration
+            raise StopIteration
+
+    def iterable(self) -> bool:
+        return hasattr(self.data, '__iter__') or hasattr(self.data, '__getitem__')
+
+    def _data_length(self, index: Optional[int] = None) -> int:
+        if self.measure_data:
+            return self.measure_data(index or 0)
+        return len(self.data) if self.data is not None else 0
+
+    def __sequence_iterator(self) -> Any:
+        # Custom iterator for objects supporting the sequence protocol without __iter__
+        class SeqIterator:
+            def __init__(self, seq):
+                self.seq = seq
+                self.index = 0
+
+            def __next__(self):
+                try:
+                    result = self.seq[self.index]
+                except IndexError:
+                    raise StopIteration
+                self.index += 1
+                return result
+
+        return SeqIterator(self.data)
+
+
+class GeneratorWithMetadata:
+    def __init__(self, generator: Generator, metadata: Optional[List[Dict[str, Any]]]):
+        """
+        Initializes the custom generator wrapper with a generator and associated metadata.
+
+        :param generator: The generator to wrap.
+        :param metadata: A dictionary containing metadata related to the generator.
+        """
+        self._generator = generator
+        self.metadata = metadata
+
+    def __iter__(self):
+        """
+        Returns the iterator object itself.
+        """
+        return self
+
+    def __next__(self):
+        """
+        Returns the next item from the generator.
+        """
+        return next(self._generator)
+
+    def get_metadata(self) -> Dict[str, Any]:
+        """
+        Returns the stored metadata.
+
+        :return: A dictionary containing metadata.
+        """
+        return self.metadata