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

mage_ai/server/scheduler_manager.py

@@ -9,7 +9,6 @@ import sentry_sdk
 
 from mage_ai.orchestration.db.database_manager import database_manager
 from mage_ai.orchestration.db.process import create_process
-from mage_ai.orchestration.job_manager import job_manager
 from mage_ai.server.logger import Logger
 from mage_ai.services.newrelic import initialize_new_relic
 from mage_ai.settings import (
@@ -26,6 +25,7 @@ logger = Logger().new_server_logger(__name__)
 
 
 def run_scheduler():
+    from mage_ai.orchestration.job_manager import job_manager
     from mage_ai.orchestration.triggers.loop_time_trigger import LoopTimeTrigger
 
     sentry_dsn = SENTRY_DSN
@@ -75,6 +75,15 @@ class SchedulerManager:
     def is_alive(self):
         return self.scheduler_process is not None and self.scheduler_process.is_alive()
 
+    def get_scheduler_pid(self):
+        if not self.scheduler_process:
+            return None
+        try:
+            return self.scheduler_process.pid
+        except Exception:
+            traceback.print_exc()
+            return None
+
     def get_status(self, auto_restart: bool = False):
         if auto_restart and self.status == self.SchedulerStatus.RUNNING and not self.is_alive:
             logger.info('Restarting pipeline scheduler.')
@@ -98,6 +107,8 @@ class SchedulerManager:
                 time.sleep(SCHEDULER_AUTO_RESTART_INTERVAL / 1000)
 
     def stop_scheduler(self):
+        from mage_ai.orchestration.job_manager import job_manager
+
         logger.info('Stop scheduler.')
         if self.is_alive:
             self.scheduler_process.terminate()

mage_ai/server/server.py

@@ -13,6 +13,7 @@ from typing import Optional, Union
 import pytz
 import tornado.ioloop
 import tornado.web
+from sqlalchemy import or_
 from tornado import autoreload
 from tornado.ioloop import PeriodicCallback
 from tornado.log import enable_pretty_logging
@@ -42,7 +43,7 @@ from mage_ai.data_preparation.repo_manager import (
 )
 from mage_ai.data_preparation.shared.constants import MANAGE_ENV_VAR
 from mage_ai.orchestration.constants import Entity
-from mage_ai.orchestration.db import db_connection, set_db_schema
+from mage_ai.orchestration.db import db_connection, safe_db_query, set_db_schema
 from mage_ai.orchestration.db.database_manager import database_manager
 from mage_ai.orchestration.db.models.oauth import Oauth2Application, Role, User
 from mage_ai.orchestration.utils.distributed_lock import DistributedLock
@@ -55,6 +56,7 @@ from mage_ai.server.api.events import (
     ApiEventMatcherDetailHandler,
     ApiEventMatcherListHandler,
 )
+from mage_ai.server.api.runs import ApiRunHandler
 from mage_ai.server.api.triggers import ApiTriggerPipelineHandler
 from mage_ai.server.api.v1 import (
     ApiChildDetailHandler,
@@ -87,6 +89,7 @@ from mage_ai.services.ssh.aws.emr.utils import file_path as file_path_aws_emr
 from mage_ai.settings import (
     AUTHENTICATION_MODE,
     DISABLE_AUTO_BROWSER_OPEN,
+    DISABLE_AUTORELOAD,
     ENABLE_PROMETHEUS,
     OAUTH2_APPLICATION_CLIENT_ID,
     OTEL_EXPORTER_OTLP_ENDPOINT,
@@ -112,9 +115,10 @@ from mage_ai.settings.repo import (
     set_repo_path,
 )
 from mage_ai.shared.constants import ENV_VAR_INSTANCE_TYPE, InstanceType
-from mage_ai.shared.environments import is_debug
+from mage_ai.shared.environments import is_debug, is_dev
 from mage_ai.shared.io import chmod
 from mage_ai.shared.logger import LoggingLevel, set_logging_format
+from mage_ai.shared.singletons.memory import get_memory_manager_controller
 from mage_ai.shared.utils import is_port_in_use
 from mage_ai.usage_statistics.logger import UsageStatisticLogger
 
@@ -309,7 +313,6 @@ def make_app(
             r'/api/pipelines/(?P<pipeline_uuid>\w+)/blocks/(?P<block_uuid>[\w\-\%2f\.]+)/analyses',
             ApiPipelineBlockAnalysisHandler,
         ),
-
         # Trigger pipeline via API
         # Original route for backwards compatibility
         (
@@ -320,7 +323,15 @@ def make_app(
             r'/api/pipeline_schedules/(?P<pipeline_schedule_id>\w+)/api_trigger',
             ApiTriggerPipelineHandler,
         ),
-
+        # Run a single block and get a response immediately
+        (
+            r'/api/runs',
+            ApiRunHandler,
+        ),
+        (
+            r'/api/runs/(?P<token>\w+)',
+            ApiRunHandler,
+        ),
         # Download block output
         (
             r'/api/pipelines/(?P<pipeline_uuid>\w+)/block_outputs/'
@@ -328,12 +339,9 @@ def make_app(
             ApiDownloadHandler,
         ),
         # Download resource
+        (r'/api/downloads/(?P<token>[\w/%.-]+)', ApiResourceDownloadHandler),
         (
-            r'/api/downloads/(?P<token>[\w/%.-]+)',
-            ApiResourceDownloadHandler
-        ),
-        (
-            r'/api/(?P<resource>\w+)/(?P<pk>[\w\-\%2f\.]+)' \
+            r'/api/(?P<resource>\w+)/(?P<pk>[\w\-\%2f\.]+)'
             r'/(?P<child>\w+)/(?P<child_pk>[\w\-\%2f\.]+)',
             ApiChildDetailHandler,
         ),
@@ -364,6 +372,7 @@ def make_app(
 
     if ENABLE_PROMETHEUS or OTEL_EXPORTER_OTLP_ENDPOINT:
         from opentelemetry.instrumentation.tornado import TornadoInstrumentor
+
         TornadoInstrumentor().instrument()
         logger.info('OpenTelemetry instrumentation enabled.')
 
@@ -378,10 +387,12 @@ def make_app(
         from opentelemetry.sdk.trace import TracerProvider
         from opentelemetry.sdk.trace.export import BatchSpanProcessor
 
-        service_name = "mage-ai-server"
-        resource = Resource(attributes={
-            "service.name": service_name,
-        })
+        service_name = 'mage-ai-server'
+        resource = Resource(
+            attributes={
+                'service.name': service_name,
+            }
+        )
 
         # Set up a TracerProvider and attach an OTLP exporter to it
         trace.set_tracer_provider(TracerProvider(resource=resource))
@@ -392,7 +403,7 @@ def make_app(
             # Endpoint of your OpenTelemetry Collector
             endpoint=OTEL_EXPORTER_OTLP_ENDPOINT,
             # Use insecure channel if your collector does not support TLS
-            insecure=True
+            insecure=True,
         )
 
         # Attach the OTLP exporter to the TracerProvider
@@ -408,9 +419,7 @@ def make_app(
 
         TornadoInstrumentor().instrument()
         # Service name is required for most backends
-        resource = Resource(attributes={
-            SERVICE_NAME: 'mage'
-        })
+        resource = Resource(attributes={SERVICE_NAME: 'mage'})
 
         # Initialize PrometheusMetricReader which pulls metrics from the SDK
         # on-demand to respond to scrape requests
@@ -422,28 +431,31 @@ def make_app(
     if update_routes:
         updated_routes = []
         for route in routes:
-            updated_routes.append(
-                (route[0].replace('/', f'/{ROUTES_BASE_PATH}/', 1), *route[1:]))
+            updated_routes.append((route[0].replace('/', f'/{ROUTES_BASE_PATH}/', 1), *route[1:]))
     else:
         updated_routes = routes
 
-    autoreload.add_reload_hook(scheduler_manager.stop_scheduler)
-
     file_path = file_path_aws_emr()
     if not os.path.exists(file_path):
         os.makedirs(os.path.dirname(file_path), exist_ok=True)
         with open(file_path, 'w') as f:
             f.write(json.dumps({}))
 
-    autoreload.watch(file_path)
+    if is_dev() and not DISABLE_AUTORELOAD:
+        should_autoreload = True
+        autoreload.add_reload_hook(scheduler_manager.stop_scheduler)
+        autoreload.watch(file_path)
+    else:
+        should_autoreload = False
 
     return tornado.web.Application(
         updated_routes,
-        autoreload=True,
+        autoreload=should_autoreload,
         template_path=template_dir,
     )
 
 
+@safe_db_query
 def initialize_user_authentication(project_type: ProjectType) -> Oauth2Application:
     logger.info('User authentication is enabled.')
     # We need to sleep for a few seconds after creating all the tables or else there
@@ -471,19 +483,29 @@ def initialize_user_authentication(project_type: ProjectType) -> Oauth2Applicati
     if not legacy_owner_user and len(owner_users) == 0:
         logger.info('User with owner permission doesn’t exist, creating owner user.')
         if AUTHENTICATION_MODE.lower() == 'ldap':
-            user = User.create(
-                roles_new=[default_owner_role],
-                username=get_settings_value(LDAP_ADMIN_USERNAME, 'admin'),
-            )
+            username = get_settings_value(LDAP_ADMIN_USERNAME, 'admin')
+            user = User.query.filter(User.username == username).first()
+            if not user:
+                user = User.create(
+                    roles_new=[default_owner_role],
+                    username=get_settings_value(LDAP_ADMIN_USERNAME, 'admin'),
+                )
         else:
             password_salt = generate_salt()
-            user = User.create(
-                email='admin@admin.com',
-                password_hash=create_bcrypt_hash('admin', password_salt),
-                password_salt=password_salt,
-                roles_new=[default_owner_role],
-                username='admin',
-            )
+            user = User.query.filter(
+                or_(
+                    User.email == 'admin@admin.com',
+                    User.username == 'admin',
+                ),
+            ).first()
+            if not user:
+                user = User.create(
+                    email='admin@admin.com',
+                    password_hash=create_bcrypt_hash('admin', password_salt),
+                    password_salt=password_salt,
+                    roles_new=[default_owner_role],
+                    username='admin',
+                )
         owner_user = user
     else:
         if legacy_owner_user and not legacy_owner_user.roles_new:
@@ -494,8 +516,7 @@ def initialize_user_authentication(project_type: ProjectType) -> Oauth2Applicati
         Oauth2Application.client_id == OAUTH2_APPLICATION_CLIENT_ID,
     ).first()
     if not oauth_client:
-        logger.info(
-            'OAuth2 application doesn’t exist for frontend, creating OAuth2 application.')
+        logger.info('OAuth2 application doesn’t exist for frontend, creating OAuth2 application.')
         oauth_client = Oauth2Application.create(
             client_id=OAUTH2_APPLICATION_CLIENT_ID,
             client_type=Oauth2Application.ClientType.PUBLIC,
@@ -573,6 +594,7 @@ async def main(
         try:
             from mage_ai.data_preparation.sync import GitConfig
             from mage_ai.data_preparation.sync.git_sync import GitSync
+
             sync_config = GitConfig.load(config=preferences.sync_config)
             sync = GitSync(sync_config, setup_repo=True)
             if sync_config.remote_repo_link and sync_config.sync_on_start is True:
@@ -610,14 +632,14 @@ async def main(
         logger.info('Initializing tag cache.')
         await TagCache.initialize_cache(replace=True)
 
+        project_model = Project(root_project=True)
         logger.info('Initializing block action object cache.')
-        await BlockActionObjectCache.initialize_cache(replace=True)
+        await BlockActionObjectCache.initialize_cache(project=project_model, replace=True)
 
-        project_model = Project(root_project=True)
         if project_model:
-            if project_model.spark_config and \
-                    project_model.is_feature_enabled(FeatureUUID.COMPUTE_MANAGEMENT):
-
+            if project_model.spark_config and project_model.is_feature_enabled(
+                FeatureUUID.COMPUTE_MANAGEMENT
+            ):
                 Application.clear_cache()
 
             if project_model.is_feature_enabled(FeatureUUID.DBT_V2):
@@ -649,7 +671,10 @@ async def main(
         try:
             from mage_ai.services.ssh.aws.emr.models import create_tunnel
 
-            tunnel = create_tunnel(clean_up_on_failure=True)
+            tunnel = create_tunnel(
+                clean_up_on_failure=True,
+                project=project_model,
+            )
             if tunnel:
                 print(f'SSH tunnel active: {tunnel.is_active()}')
         except Exception as err:
@@ -670,6 +695,8 @@ async def main(
         )
         periodic_callback.start()
 
+    get_memory_manager_controller().events
+
     update_settings_on_metadata_change()
     observer = Observer()
     event_handler = MetadataEventHandler()

mage_ai/server/utils/custom_output.py

@@ -0,0 +1,284 @@
+INTERNAL_PREFIX = '__internal_output__'
+
+
+def __custom_output():
+    import warnings
+
+    import pandas as pd
+    import polars as pl
+    import simplejson
+
+    from mage_ai.ai.utils.xgboost import create_tree_visualization
+    from mage_ai.data_preparation.models.block.dynamic.utils import (
+        combine_transformed_output_for_multi_output,
+        transform_output_for_display,
+        transform_output_for_display_dynamic_child,
+        transform_output_for_display_reduce_output,
+    )
+    from mage_ai.data_preparation.models.block.outputs import format_output_data
+    from mage_ai.data_preparation.models.pipeline import Pipeline
+    from mage_ai.data_preparation.models.utils import infer_variable_type
+    from mage_ai.data_preparation.models.variable import VariableType
+    from mage_ai.presenters.utils import render_output_tags
+    from mage_ai.server.kernel_output_parser import DataType
+    from mage_ai.shared.parsers import (
+        convert_matrix_to_dataframe,
+        encode_complex,
+        polars_to_dict_split,
+        sample_output,
+    )
+
+    if pd.__version__ < '1.5.0':
+        from pandas.core.common import SettingWithCopyWarning
+    else:
+        from pandas.errors import SettingWithCopyWarning
+
+    warnings.simplefilter(action='ignore', category=SettingWithCopyWarning)
+
+    pipeline = Pipeline.get('{pipeline_uuid}', repo_path='{repo_path}')
+    block = pipeline.get_block(
+        '{block_uuid}', extension_uuid='{extension_uuid}', widget=bool('{widget}')
+    )
+
+    is_dynamic = bool('{is_dynamic}')
+    is_dynamic_child = bool('{is_dynamic_child}')
+
+    if block.executable:
+        if is_dynamic_child:
+            # The output of a dynamic child block is printed to the client from execute_custom_code
+            return
+
+        outputs = block.get_outputs()
+
+        if outputs is not None and isinstance(outputs, list):
+            outputs = outputs[: int('{DATAFRAME_SAMPLE_COUNT_PREVIEW}')]
+
+        if outputs is not None and len(outputs) >= 1:
+            _json_string = simplejson.dumps(
+                outputs,
+                default=encode_complex,
+                ignore_nan=True,
+            )
+
+            return print(render_output_tags(_json_string))
+
+    _internal_output_return = '{last_line}'
+    variable_type, basic_iterable = infer_variable_type(_internal_output_return)
+
+    if (
+        not is_dynamic
+        and not is_dynamic_child
+        and isinstance(_internal_output_return, list)
+        and len(_internal_output_return) >= 2
+        and any([infer_variable_type(i)[0] for i in _internal_output_return])
+    ):
+        output_transformed = []
+        for idx, item in enumerate(
+            _internal_output_return[: int('{DATAFRAME_SAMPLE_COUNT_PREVIEW}')]
+        ):
+            data_output, _ = format_output_data(
+                block,
+                item,
+                'output_' + str(idx),
+                automatic_sampling=True,
+                sample_count=int('{DATAFRAME_SAMPLE_COUNT_PREVIEW}'),
+            )
+            if isinstance(data_output, dict):
+                data_output['multi_output'] = True
+            output_transformed.append(data_output)
+
+        try:
+            _json_string = simplejson.dumps(
+                output_transformed,
+                default=encode_complex,
+                ignore_nan=True,
+            )
+
+            return print(f'[{INTERNAL_PREFIX}]{_json_string}')
+        except Exception as err:
+            print(f'[ERROR] Failed to serialize output: {err}')
+            raise err
+
+    if (
+        variable_type is None
+        and is_dynamic
+        and not is_dynamic_child
+        and isinstance(_internal_output_return, list)
+        and len(_internal_output_return) >= 1
+    ):
+        variable_type, basic_iterable = infer_variable_type(_internal_output_return[0])
+
+        if VariableType.LIST_COMPLEX == variable_type:
+            list_of_lists = []
+
+            for list_of_items in _internal_output_return:
+                list_of_lists.append([encode_complex(item) for item in list_of_items])
+            _internal_output_return = list_of_lists
+
+    if VariableType.MATRIX_SPARSE == variable_type:
+        if basic_iterable:
+            _internal_output_return = convert_matrix_to_dataframe(_internal_output_return[0])
+        else:
+            _internal_output_return = convert_matrix_to_dataframe(_internal_output_return)
+    elif VariableType.SERIES_PANDAS == variable_type:
+        if basic_iterable:
+            _internal_output_return = pd.DataFrame(_internal_output_return).T
+        else:
+            _internal_output_return = _internal_output_return.to_frame()
+    elif VariableType.CUSTOM_OBJECT == variable_type:
+        return _internal_output_return
+    elif VariableType.MODEL_SKLEARN == variable_type:
+        return _internal_output_return
+    elif VariableType.MODEL_XGBOOST == variable_type:
+        text_data, success = create_tree_visualization(_internal_output_return)
+
+        if not success:
+            print(text_data)
+            return
+
+        data = dict(
+            text_data=text_data,
+            type=DataType.IMAGE_PNG if success else DataType.TEXT,
+            variable_uuid=variable_type,
+        )
+        _json_string = simplejson.dumps(
+            data,
+            default=encode_complex,
+            ignore_nan=True,
+        )
+        return print(f'[{INTERNAL_PREFIX}]{_json_string}')
+    elif VariableType.LIST_COMPLEX == variable_type:
+        _internal_output_return = [encode_complex(item) for item in _internal_output_return]
+
+    # Dynamic block child logic always takes precedence over dynamic block logic
+    if is_dynamic_child:
+        output_transformed = []
+
+        if _internal_output_return and isinstance(_internal_output_return, list):
+            for output in _internal_output_return:
+                output_tf = transform_output_for_display_dynamic_child(
+                    output,
+                    is_dynamic=is_dynamic,
+                    is_dynamic_child=is_dynamic_child,
+                    single_item_only=True,
+                )
+                output_transformed.append(output_tf)
+
+        output_transformed = output_transformed[: int('{DATAFRAME_SAMPLE_COUNT_PREVIEW}')]
+
+        try:
+            _json_string = simplejson.dumps(
+                combine_transformed_output_for_multi_output(output_transformed),
+                default=encode_complex,
+                ignore_nan=True,
+            )
+
+            return print(f'[{INTERNAL_PREFIX}]{_json_string}')
+        except Exception as err:
+            print(type(_internal_output_return))
+            print(type(output_transformed))
+            raise err
+    elif is_dynamic:
+        _json_string = simplejson.dumps(
+            transform_output_for_display(
+                _internal_output_return,
+                is_dynamic=is_dynamic,
+                is_dynamic_child=is_dynamic_child,
+                sample_count=int('{DATAFRAME_ANALYSIS_MAX_COLUMNS}'),
+                sample_columns=int('{DATAFRAME_ANALYSIS_MAX_COLUMNS}'),
+            ),
+            default=encode_complex,
+            ignore_nan=True,
+        )
+        return print(f'[{INTERNAL_PREFIX}]{_json_string}')
+    elif bool('{has_reduce_output}') and is_dynamic_child:
+        _json_string = simplejson.dumps(
+            transform_output_for_display_reduce_output(
+                _internal_output_return,
+                is_dynamic=is_dynamic,
+                is_dynamic_child=is_dynamic_child,
+                sample_count=int('{DATAFRAME_SAMPLE_COUNT_PREVIEW}'),
+                sample_columns=int('{DATAFRAME_ANALYSIS_MAX_COLUMNS}'),
+            ),
+            default=encode_complex,
+            ignore_nan=True,
+        )
+        return print(f'[{INTERNAL_PREFIX}]{_json_string}')
+    elif isinstance(_internal_output_return, (pd.DataFrame, pl.DataFrame, dict, list)) and (
+        type(_internal_output_return).__module__ != 'geopandas.geodataframe'
+    ):
+        if isinstance(_internal_output_return, dict):
+            _internal_output_return = pd.DataFrame([_internal_output_return])
+        elif isinstance(_internal_output_return, list):
+            _internal_output_return = pd.DataFrame(_internal_output_return, columns=['column'])
+
+        _is_polars = isinstance(_internal_output_return, pl.DataFrame)
+
+        if _is_polars:
+            _sample = _internal_output_return[: int('{DATAFRAME_SAMPLE_COUNT_PREVIEW}')]
+            _columns = _sample.columns[: int('{DATAFRAME_ANALYSIS_MAX_COLUMNS}')]
+        else:
+            _sample = _internal_output_return.iloc[: int('{DATAFRAME_SAMPLE_COUNT_PREVIEW}')]
+            _columns = _sample.columns.tolist()[: int('{DATAFRAME_ANALYSIS_MAX_COLUMNS}')]
+
+        for col in _columns:
+            try:
+                _sample[col] = _sample[col].fillna('')
+            except Exception:
+                pass
+
+        if _is_polars:
+            _rows = polars_to_dict_split(_sample[_columns])['data']
+            _index = [i for i in range(len(_sample))]
+        else:
+            _rows = simplejson.loads(
+                _sample[_columns].to_json(
+                    date_format='iso',
+                    default_handler=str,
+                    orient='split',
+                )
+            )['data']
+            _index = _sample.index.tolist()
+
+        _shape = _internal_output_return.shape
+
+        _json_string = simplejson.dumps(
+            dict(
+                data=dict(
+                    columns=_columns,
+                    index=_index,
+                    rows=_rows,
+                    shape=_shape,
+                ),
+                type='table',
+            ),
+            default=encode_complex,
+            ignore_nan=True,
+        )
+        return print(f'[{INTERNAL_PREFIX}]{_json_string}')
+    elif type(_internal_output_return).__module__ == 'pyspark.sql.dataframe':
+        _sample = _internal_output_return.limit(int('{DATAFRAME_SAMPLE_COUNT_PREVIEW}')).toPandas()
+        _columns = _sample.columns.tolist()[:40]
+        _rows = _sample.to_numpy().tolist()
+        _shape = [_internal_output_return.count(), len(_sample.columns.tolist())]
+        _index = _sample.index.tolist()
+
+        _json_string = simplejson.dumps(
+            dict(
+                data=dict(
+                    columns=_columns,
+                    index=_index,
+                    rows=_rows,
+                    shape=_shape,
+                ),
+                type='table',
+            ),
+            default=encode_complex,
+            ignore_nan=True,
+        )
+        return print(f'[{INTERNAL_PREFIX}]{_json_string}')
+    elif not bool('{is_print_statement}'):
+        output, sampled = sample_output(encode_complex(_internal_output_return))
+        if sampled:
+            print('Sampled output is provided here for preview.')
+        return output