mage-ai 0.9.45__py3-none-any.whl → 0.9.47__py3-none-any.whl

mage_ai/server/server.py

@@ -1,5 +1,6 @@
 import argparse
 import asyncio
+import json
 import os
 import shutil
 import stat
@@ -46,7 +47,7 @@ from mage_ai.orchestration.utils.distributed_lock import DistributedLock
 from mage_ai.server.active_kernel import switch_active_kernel
 from mage_ai.server.api.base import BaseHandler
 from mage_ai.server.api.blocks import ApiPipelineBlockAnalysisHandler
-from mage_ai.server.api.downloads import ApiDownloadHandler
+from mage_ai.server.api.downloads import ApiDownloadHandler, ApiResourceDownloadHandler
 from mage_ai.server.api.events import (
     ApiEventHandler,
     ApiEventMatcherDetailHandler,
@@ -79,6 +80,7 @@ from mage_ai.server.terminal_server import (
 from mage_ai.server.websocket_server import WebSocketServer
 from mage_ai.services.redis.redis import init_redis_client
 from mage_ai.services.spark.models.applications import Application
+from mage_ai.services.ssh.aws.emr.utils import file_path as file_path_aws_emr
 from mage_ai.settings import (
     AUTHENTICATION_MODE,
     ENABLE_PROMETHEUS,
@@ -93,7 +95,13 @@ from mage_ai.settings import (
     SHELL_COMMAND,
     USE_UNIQUE_TERMINAL,
 )
-from mage_ai.settings.repo import DEFAULT_MAGE_DATA_DIR, get_repo_name, set_repo_path
+from mage_ai.settings.repo import (
+    DEFAULT_MAGE_DATA_DIR,
+    MAGE_CLUSTER_TYPE_ENV_VAR,
+    MAGE_PROJECT_TYPE_ENV_VAR,
+    get_repo_name,
+    set_repo_path,
+)
 from mage_ai.shared.constants import ENV_VAR_INSTANCE_TYPE, InstanceType
 from mage_ai.shared.io import chmod
 from mage_ai.shared.logger import LoggingLevel
@@ -292,6 +300,11 @@ def make_app(template_dir: str = None, update_routes: bool = False):
             r'(?P<block_uuid>[\w\-\%2f\.(/.*)?]+)/downloads',
             ApiDownloadHandler,
         ),
+        # Download resource
+        (
+            r'/api/downloads/(?P<token>[\w/%.-]+)',
+            ApiResourceDownloadHandler
+        ),
 
         # API v1 routes
         (
@@ -355,6 +368,15 @@ def make_app(template_dir: str = None, update_routes: bool = False):
         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)
+
     return tornado.web.Application(
         updated_routes,
         autoreload=True,
@@ -405,7 +427,8 @@ async def main(
         address=host if host != 'localhost' else None,
     )
 
-    url = f'http://{host or "localhost"}:{port}'
+    host = host or 'localhost'
+    url = f'http://{host}:{port}'
     if update_routes:
         url = f'{url}/{ROUTES_BASE_PATH}'
     webbrowser.open_new_tab(url)
@@ -500,6 +523,22 @@ async def main(
     logger.info('Initializing block action object cache.')
     await BlockActionObjectCache.initialize_cache(replace=True)
 
+    project_model = Project()
+    if project_model and \
+            project_model.spark_config and \
+            project_model.is_feature_enabled(FeatureUUID.COMPUTE_MANAGEMENT):
+
+        Application.clear_cache()
+
+    try:
+        from mage_ai.services.ssh.aws.emr.models import create_tunnel
+
+        tunnel = create_tunnel(clean_up_on_failure=True)
+        if tunnel:
+            print(f'SSH tunnel active: {tunnel.is_active()}')
+    except Exception as err:
+        print(f'[WARNING] SSH tunnel failed to create and connect: {err}')
+
     # Check scheduler status periodically
     periodic_callback = PeriodicCallback(
         check_scheduler_status,
@@ -557,13 +596,6 @@ def start_server(
 
     asyncio.run(UsageStatisticLogger().project_impression())
 
-    project_model = Project()
-    if project_model and \
-            project_model.spark_config and \
-            project_model.is_feature_enabled(FeatureUUID.COMPUTE_MANAGEMENT):
-
-        Application.clear_cache()
-
     if dbt_docs:
         run_docs_server()
     else:
@@ -622,8 +654,8 @@ if __name__ == '__main__':
     manage = args.manage_instance == '1'
     dbt_docs = args.dbt_docs_instance == '1'
     instance_type = os.getenv(ENV_VAR_INSTANCE_TYPE, args.instance_type)
-    project_type = os.getenv('PROJECT_TYPE', ProjectType.STANDALONE)
-    cluster_type = os.getenv('CLUSTER_TYPE')
+    project_type = os.getenv(MAGE_PROJECT_TYPE_ENV_VAR, ProjectType.STANDALONE)
+    cluster_type = os.getenv(MAGE_CLUSTER_TYPE_ENV_VAR)
 
     start_server(
         host=host,

mage_ai/server/utils/output_display.py

@@ -221,6 +221,7 @@ def add_execution_code(
     code: str,
     global_vars,
     block_type: BlockType = None,
+    execution_uuid: str = None,
     extension_uuid: str = None,
     kernel_name: str = None,
     output_messages_to_logs: bool = False,
@@ -237,6 +238,9 @@ def add_execution_code(
 ) -> str:
     escaped_code = code.replace("'''", "\"\"\"")
 
+    if execution_uuid:
+        execution_uuid = f"'{execution_uuid}'"
+
     if extension_uuid:
         extension_uuid = f"'{extension_uuid}'"
     if upstream_blocks:
@@ -290,6 +294,7 @@ def execute_custom_code():
         config={pipeline_config},
         repo_config={repo_config},
     )
+
     block = pipeline.get_block(block_uuid, extension_uuid={extension_uuid}, widget={widget})
 
     upstream_blocks = {upstream_blocks}
@@ -328,6 +333,7 @@ def execute_custom_code():
         global_vars['logger'] = logger
     block_output = block.execute_with_callback(
         custom_code=code,
+        execution_uuid={execution_uuid},
         from_notebook=True,
         global_vars=global_vars,
         logger=logger,

mage_ai/server/websocket_server.py

@@ -49,6 +49,7 @@ from mage_ai.server.utils.output_display import (
     get_block_output_process_code,
     get_pipeline_execution_code,
 )
+from mage_ai.services.spark.constants import ComputeServiceUUID
 from mage_ai.settings import (
     DISABLE_NOTEBOOK_EDIT_ACCESS,
     HIDE_ENV_VAR_VALUES,
@@ -418,12 +419,24 @@ class WebSocketServer(tornado.websocket.WebSocketHandler):
                     remote_execution = True
                 else:
                     remote_execution = False
+
+                execution_uuid = None
+                # Need to cache everything here
+                if block.should_track_spark() and \
+                        ComputeServiceUUID.AWS_EMR == block.compute_service_uuid:
+
+                    execution_uuid = str(uuid.uuid4()).split('-')[0]
+                    block.clear_spark_jobs_cache()
+                    block.cache_spark_application()
+                    block.set_spark_job_execution_start(execution_uuid=execution_uuid)
+
                 code = add_execution_code(
                     pipeline_uuid,
                     block_uuid,
                     custom_code,
                     global_vars,
                     block_type=block_type,
+                    execution_uuid=execution_uuid,
                     extension_uuid=extension_uuid,
                     kernel_name=kernel_name,
                     output_messages_to_logs=output_messages_to_logs,
@@ -452,6 +465,7 @@ from mage_ai.orchestration.db import db_connection
 db_connection.start_session()
 """
                 client.execute(initialize_db_connection)
+
             msg_id = client.execute(add_internal_output_info(code))
 
             WebSocketServer.running_executions_mapping[msg_id] = value

mage_ai/services/aws/emr/config.py

@@ -26,6 +26,7 @@ class EmrScalingPocliy(BaseConfig):
 class EmrConfig(BaseConfig):
     bootstrap_script_path: str = None
     ec2_key_name: str = None
+    ec2_key_path: str = None
     master_security_group: str = None
     slave_security_group: str = None
     master_instance_type: str = DEFAULT_INSTANCE_TYPE

mage_ai/services/aws/emr/constants.py

@@ -0,0 +1 @@
+SECURITY_GROUP_NAME_MASTER_DEFAULT = 'ElasticMapReduce-master'

mage_ai/services/aws/emr/emr.py

@@ -4,12 +4,14 @@ import random
 import sys
 import time
 from datetime import datetime
+from typing import List
 
 from botocore.exceptions import ClientError
 
 from mage_ai.services.aws import get_aws_boto3_client
 from mage_ai.services.aws.emr import emr_basics
 from mage_ai.services.aws.emr.config import EmrConfig
+from mage_ai.services.compute.aws.constants import ClusterStatusState
 
 MAX_STEPS_IN_CLUSTER = 255 - 55
 MAX_RUNNING_OR_PENDING_STEPS = 15
@@ -123,20 +125,31 @@ def describe_cluster(cluster_id):
     return emr_basics.describe_cluster(cluster_id, emr_client)
 
 
-def list_clusters():
+def list_clusters(cluster_states: List[ClusterStatusState] = None):
     emr_client = get_aws_boto3_client('emr')
 
+    if cluster_states is None:
+        cluster_states = [
+            ClusterStatusState.BOOTSTRAPPING,
+            ClusterStatusState.RUNNING,
+            ClusterStatusState.STARTING,
+            ClusterStatusState.WAITING,
+        ]
+
     clusters = emr_client.list_clusters(
-        ClusterStates=[
-            'BOOTSTRAPPING',
-            'RUNNING',
-            'STARTING',
-            'WAITING',
-        ],
+        ClusterStates=cluster_states,
     )
     return clusters
 
 
+def terminate_clusters(cluster_ids: List[str]) -> None:
+    emr_client = get_aws_boto3_client('emr')
+
+    return emr_client.terminate_job_flows(
+        JobFlowIds=cluster_ids,
+    )
+
+
 def submit_spark_job(
     cluster_name,
     steps,

mage_ai/services/compute/aws/constants.py

@@ -0,0 +1,21 @@
+from enum import Enum
+
+CONNECTION_CREDENTIAL_AWS_ACCESS_KEY_ID = 'AWS_ACCESS_KEY_ID'
+CONNECTION_CREDENTIAL_AWS_SECRET_ACCESS_KEY = 'AWS_SECRET_ACCESS_KEY'
+
+
+class ClusterStatusState(str, Enum):
+    BOOTSTRAPPING = 'BOOTSTRAPPING'
+    RUNNING = 'RUNNING'
+    STARTING = 'STARTING'
+    TERMINATED = 'TERMINATED'
+    TERMINATED_WITH_ERRORS = 'TERMINATED_WITH_ERRORS'
+    TERMINATING = 'TERMINATING'
+    WAITING = 'WAITING'
+
+
+INVALID_STATES = [
+    ClusterStatusState.TERMINATED,
+    ClusterStatusState.TERMINATED_WITH_ERRORS,
+    ClusterStatusState.TERMINATING,
+]