mage-ai 0.9.44__py3-none-any.whl → 0.9.46__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
@@ -9,15 +10,9 @@ from datetime import datetime
 from time import sleep
 from typing import Optional, Union
 
-import prometheus_client
 import pytz
 import tornado.ioloop
 import tornado.web
-from opentelemetry import metrics
-from opentelemetry.exporter.prometheus import PrometheusMetricReader
-from opentelemetry.instrumentation.tornado import TornadoInstrumentor
-from opentelemetry.sdk.metrics import MeterProvider
-from opentelemetry.sdk.resources import SERVICE_NAME, Resource
 from tornado import autoreload
 from tornado.ioloop import PeriodicCallback
 from tornado.log import enable_pretty_logging
@@ -52,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,
@@ -85,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,
@@ -168,6 +164,8 @@ class ApiSchedulerHandler(BaseHandler):
 
 class PrometheusMetricsHandler(BaseHandler):
     def get(self):
+        import prometheus_client
+
         self.set_header('Content-Type', prometheus_client.CONTENT_TYPE_LATEST)
         self.write(prometheus_client.generate_latest(prometheus_client.REGISTRY))
 
@@ -296,6 +294,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
         (
@@ -331,6 +334,12 @@ def make_app(template_dir: str = None, update_routes: bool = False):
     ]
 
     if ENABLE_PROMETHEUS:
+        from opentelemetry import metrics
+        from opentelemetry.exporter.prometheus import PrometheusMetricReader
+        from opentelemetry.instrumentation.tornado import TornadoInstrumentor
+        from opentelemetry.sdk.metrics import MeterProvider
+        from opentelemetry.sdk.resources import SERVICE_NAME, Resource
+
         TornadoInstrumentor().instrument()
         # Service name is required for most backends
         resource = Resource(attributes={
@@ -353,6 +362,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,
@@ -403,7 +421,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)
@@ -498,6 +517,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,
@@ -555,13 +590,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:

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/ecs/config.py

@@ -17,11 +17,13 @@ class EcsConfig(BaseConfig):
     cluster: str
     security_groups: List[str]
     subnets: List[str]
-    tags: List = field(default_factory=list)
-    network_configuration: Dict = None
+    assign_public_ip: bool = True
     cpu: int = 512
-    memory: int = 1024
+    enable_execute_command: bool = False
     launch_type: str = 'FARGATE'
+    memory: int = 1024
+    network_configuration: Dict = None
+    tags: List = field(default_factory=list)
     wait_timeout: int = 600
 
     @classmethod
@@ -71,20 +73,21 @@ class EcsConfig(BaseConfig):
             network_configuration = {
                 'awsvpcConfiguration': {
                     'subnets': self.subnets,
-                    'assignPublicIp': 'ENABLED',
+                    'assignPublicIp': 'ENABLED' if self.assign_public_ip else 'DISABLED',
                     'securityGroups': self.security_groups,
                 }
             }
 
         task_config = dict(
-            taskDefinition=self.task_definition,
-            launchType=self.launch_type,
             cluster=self.cluster,
-            platformVersion='LATEST',
             count=1,
+            enableExecuteCommand=self.enable_execute_command,
+            launchType=self.launch_type,
             networkConfiguration=network_configuration,
-            tags=self.tags,
+            platformVersion='LATEST',
             propagateTags='TASK_DEFINITION',
+            tags=self.tags,
+            taskDefinition=self.task_definition,
         )
         if command is not None:
             task_config['overrides'] = {

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,
+]