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

mage_ai/api/resources/KernelResource.py

@@ -1,4 +1,10 @@
+from datetime import datetime
+
 from mage_ai.api.resources.GenericResource import GenericResource
+from mage_ai.data_preparation.models.project import Project
+from mage_ai.data_preparation.models.project.constants import FeatureUUID
+from mage_ai.kernels.models import KernelProcess, KernelWrapper
+from mage_ai.kernels.utils import find_ipykernel_launchers_info_async
 from mage_ai.orchestration.db import safe_db_query
 from mage_ai.server.active_kernel import (
     interrupt_kernel,
@@ -8,18 +14,31 @@ from mage_ai.server.active_kernel import (
 )
 from mage_ai.server.kernels import DEFAULT_KERNEL_NAME, KernelName, kernel_managers
 from mage_ai.services.ssh.aws.emr.utils import tunnel
+from mage_ai.shared.singletons.memory import get_memory_manager_controller
 
 
 class KernelResource(GenericResource):
     @classmethod
     @safe_db_query
-    def collection(self, query, meta, user, **kwargs):
+    async def collection(self, query, meta, user, **kwargs):
         kernels = []
 
+        if Project().is_feature_enabled(FeatureUUID.AUTOMATIC_KERNEL_CLEANUP):
+            # Only do this every minute
+            if int(datetime.utcnow().timestamp()) % 60 == 0:
+                kill_count, memory_freed = KernelProcess.terminate_inactive(
+                    await find_ipykernel_launchers_info_async(),
+                )
+                if kill_count >= 1:
+                    print(
+                        f'[KernelResource] Automatic kernel cleanup: {kill_count} '
+                        f'kernels terminated, {(memory_freed / 1024**3):.3f} GBs if memory freed.'
+                    )
+
         for kernel_name in KernelName:
             kernel = kernel_managers[kernel_name]
             if kernel.has_kernel:
-                kernels.append(kernel)
+                kernels.append(KernelWrapper(kernel))
 
         return self.build_result_set(
             kernels,
@@ -47,7 +66,7 @@ class KernelResource(GenericResource):
         if not kernel:
             kernel = kernel_managers[DEFAULT_KERNEL_NAME]
 
-        return self(kernel, user, **kwargs)
+        return self(KernelWrapper(kernel), user, **kwargs)
 
     @safe_db_query
     def update(self, payload, **kwargs):
@@ -56,13 +75,16 @@ class KernelResource(GenericResource):
         switch_active_kernel(self.model.kernel_name)
 
         if 'interrupt' == action_type:
+            get_memory_manager_controller().stop_all_events()
             interrupt_kernel()
         elif 'restart' == action_type:
             try:
+                get_memory_manager_controller().stop_all_events()
                 restart_kernel()
             except RuntimeError as e:
                 # RuntimeError: Cannot restart the kernel. No previous call to 'start_kernel'.
                 if 'start_kernel' in str(e):
+                    get_memory_manager_controller().stop_all_events()
                     start_kernel()
 
         def _callback(*args, **kwargs):

mage_ai/api/resources/OauthResource.py

@@ -20,7 +20,7 @@ from mage_ai.authentication.oauth.utils import (
     refresh_token_for_client,
 )
 from mage_ai.authentication.providers.constants import NAME_TO_PROVIDER
-from mage_ai.data_preparation.git.api import get_oauth_client_id
+from mage_ai.data_preparation.git.utils import get_oauth_client_id
 from mage_ai.orchestration.db import safe_db_query
 from mage_ai.orchestration.db.models.oauth import Oauth2AccessToken, Oauth2Application
 

mage_ai/api/resources/OutputResource.py

@@ -9,20 +9,35 @@ from mage_ai.data_preparation.models.block.data_integration.utils import (
     read_data_from_cache,
 )
 from mage_ai.data_preparation.models.block.utils import serialize_output
+from mage_ai.data_preparation.models.constants import DATAFRAME_SAMPLE_COUNT
 from mage_ai.data_preparation.models.pipeline import Pipeline
 from mage_ai.orchestration.db import safe_db_query
 from mage_ai.orchestration.db.models.schedules import BlockRun
+from mage_ai.settings.repo import get_repo_path
 
 
 class OutputResource(GenericResource):
     @classmethod
     @safe_db_query
-    def collection(self, query, meta, user, **kwargs):
+    def collection(self, query_arg, meta, user, **kwargs):
+        query = {}
+        for key in [
+            'sample_count',
+        ]:
+            value = query_arg.get(key, [None])
+            if value:
+                value = value[0]
+                if value is not None:
+                    query[key] = value
+
         parent_model = kwargs['parent_model']
 
         outputs = []
         if type(parent_model) is BlockRun:
-            outputs = parent_model.get_outputs()
+            outputs = parent_model.get_outputs(
+                sample_count=query.get('sample_count', DATAFRAME_SAMPLE_COUNT),
+                exclude_blank_variable_uuids=True,
+            )
 
         return self.build_result_set(
             outputs,
@@ -45,7 +60,8 @@ class OutputResource(GenericResource):
         model = dict(outputs=[])
 
         if block_uuid and pipeline_uuid:
-            pipeline = Pipeline.get(pipeline_uuid)
+            repo_path = get_repo_path(user=user)
+            pipeline = Pipeline.get(pipeline_uuid, repo_path=repo_path)
             block = pipeline.get_block(block_uuid)
 
             if block.is_data_integration():
@@ -98,10 +114,12 @@ class OutputResource(GenericResource):
                             else:
                                 output_serialized = serialize_output(block, output)
 
-                        model['outputs'].append(dict(
-                            data=output_serialized,
-                            uuid=stream_id,
-                        ))
+                        model['outputs'].append(
+                            dict(
+                                data=output_serialized,
+                                uuid=stream_id,
+                            )
+                        )
 
                         if refresh and persist:
                             persist_data_for_stream(
@@ -134,6 +152,8 @@ class OutputResource(GenericResource):
         sample_count = payload.get('sample_count') or None
         if sample_count:
             sample_count = int(sample_count)
+        else:
+            sample_count = DATAFRAME_SAMPLE_COUNT
 
         model = dict(outputs=[])
 
@@ -163,9 +183,11 @@ class OutputResource(GenericResource):
                         if not output_serialized:
                             output_serialized = serialize_output(block, output)
 
-                        model['outputs'].append(dict(
-                            data=output_serialized,
-                            uuid=stream_id,
-                        ))
+                        model['outputs'].append(
+                            dict(
+                                data=output_serialized,
+                                uuid=stream_id,
+                            )
+                        )
 
         return self(model, user, **kwargs)

mage_ai/api/resources/PageBlockLayoutResource.py

@@ -25,27 +25,33 @@ class PageBlockLayoutResource(GenericResource):
         return self(model, user, **kwargs)
 
     async def update(self, payload, **kwargs):
-        blocks = payload.get('blocks') or {}
-        layout = payload.get('layout') or []
+        blocks = payload.get("blocks") or {}
+        layout = payload.get("layout") or []
 
         blocks_with_content = {}
         blocks_with_new_names = {}
 
         for block_uuid, block_config in blocks.items():
-            block_name = block_config.get('name')
-            block_name_new = block_config.get('name_new')
-            block_type = block_config.get('type')
+            block_name = block_config.get("name")
+            block_name_new = block_config.get("name_new")
+            block_type = block_config.get("type")
 
             if block_name_new and block_name_new != block_name:
-                blocks_with_new_names[block_uuid] = merge_dict(ignore_keys(block_config, [
-                    'name_new',
-                ]), dict(
-                    name=block_name_new,
-                    uuid=clean_name(block_name_new),
-                ))
-
-            if 'content' in block_config:
-                file_path = block_config.get('file_path')
+                blocks_with_new_names[block_uuid] = merge_dict(
+                    ignore_keys(
+                        block_config,
+                        [
+                            "name_new",
+                        ],
+                    ),
+                    dict(
+                        name=block_name_new,
+                        uuid=clean_name(block_name_new),
+                    ),
+                )
+
+            if "content" in block_config:
+                file_path = block_config.get("file_path")
                 block_uuid_use = block_uuid
                 if file_path:
                     parts = file_path.split(os.path.sep)
@@ -56,35 +62,40 @@ class PageBlockLayoutResource(GenericResource):
                     block_uuid_use,
                     block_uuid_use,
                     block_type,
-                    **extract(block_config, [
-                        'language',
-                    ]),
+                    **extract(
+                        block_config,
+                        [
+                            "language",
+                        ],
+                    ),
                 )
                 blocks_with_content[block_uuid] = block
 
         for block_uuid, block_config in blocks_with_new_names.items():
             blocks.pop(block_uuid, None)
-            block_name_new = block_config.get('name')
-            block_uuid_new = block_config.get('uuid')
+            block_name_new = block_config.get("name")
+            block_uuid_new = block_config.get("uuid")
             blocks[block_uuid_new] = block_config
 
             if block_uuid in blocks_with_content:
                 block = blocks_with_content[block_uuid]
                 block.update(dict(name=block_name_new))
+                blocks_with_content.pop(block_uuid, None)
+                blocks_with_content[block_uuid_new] = block
 
         for block_uuid, block in blocks_with_content.items():
             await block.update_content_async(
-                blocks.get(block_uuid, {}).get('content'),
+                blocks.get(block_uuid, {}).get("content"),
                 error_if_file_missing=False,
             )
-            blocks[block_uuid].pop('content', None)
+            blocks[block_uuid].pop("content", None)
 
         for idx1, row in enumerate(layout):
             for idx2, column in enumerate(row):
-                block_uuid = column.get('block_uuid')
+                block_uuid = column.get("block_uuid")
                 block = blocks_with_new_names.get(block_uuid)
                 if block:
-                    layout[idx1][idx2]['block_uuid'] = block.get('uuid')
+                    layout[idx1][idx2]["block_uuid"] = block.get("uuid")
 
         self.model.blocks = blocks
         self.model.layout = layout

mage_ai/api/resources/PipelineInteractionResource.py

@@ -5,6 +5,7 @@ from mage_ai.api.resources.InteractionResource import InteractionResource
 from mage_ai.data_preparation.models.pipeline import Pipeline
 from mage_ai.data_preparation.models.pipelines.interactions import PipelineInteractions
 from mage_ai.orchestration.db import safe_db_query
+from mage_ai.settings.repo import get_repo_path
 from mage_ai.shared.hash import extract, merge_dict
 
 
@@ -13,13 +14,17 @@ class PipelineInteractionResource(GenericResource):
     @safe_db_query
     async def get_model(self, pk, **kwargs):
         uuid = urllib.parse.unquote(pk)
-        pipeline = await Pipeline.get_async(uuid)
+        user = kwargs.get('user')
+        repo_path = get_repo_path(user=user, root_project=True)
+        pipeline = await Pipeline.get_async(
+            uuid, all_projects=True, repo_path=repo_path
+        )
         return PipelineInteractions(pipeline)
 
     @classmethod
     @safe_db_query
     async def member(self, pk, user, **kwargs):
-        model = await self.get_model(pk)
+        model = await self.get_model(pk, user=user)
 
         query = kwargs.get('query', {})
         filter_for_permissions = query.get('filter_for_permissions', [False])
@@ -36,11 +41,14 @@ class PipelineInteractionResource(GenericResource):
         if 'content' in payload:
             payload_update['content'] = payload.get('content')
         else:
-            payload_update['content_parsed'] = extract(payload, [
-                'blocks',
-                'layout',
-                'permissions',
-            ])
+            payload_update['content_parsed'] = extract(
+                payload,
+                [
+                    'blocks',
+                    'layout',
+                    'permissions',
+                ],
+            )
 
         await self.model.update(**payload_update)
 
@@ -50,13 +58,21 @@ class PipelineInteractionResource(GenericResource):
                 resource = InteractionResource.member(
                     interaction_uuid,
                     self.current_user,
-                    **merge_dict(kwargs, dict(
-                        parent_model=self.model.pipeline,
-                    )),
+                    **merge_dict(
+                        kwargs,
+                        dict(
+                            parent_model=self.model.pipeline,
+                        ),
+                    ),
                 )
 
-                await resource.update(extract(interaction, [
-                    'inputs',
-                    'layout',
-                    'variables',
-                ]))
+                await resource.update(
+                    extract(
+                        interaction,
+                        [
+                            'inputs',
+                            'layout',
+                            'variables',
+                        ],
+                    )
+                )