mage-ai 0.9.44__py3-none-any.whl → 0.9.46__py3-none-any.whl

mage_ai/services/compute/models.py

@@ -0,0 +1,212 @@
+from abc import abstractmethod
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Dict, List
+
+from mage_ai.data_preparation.models.project import Project
+from mage_ai.services.compute.constants import (
+    ComputeConnectionActionUUID,
+    ComputeConnectionState,
+    ComputeManagementApplicationTab,
+)
+from mage_ai.services.spark.constants import ComputeServiceUUID
+from mage_ai.shared.hash import merge_dict
+from mage_ai.shared.models import BaseDataClass
+
+
+@dataclass
+class ErrorMessage(BaseDataClass):
+    message: str
+    variables: Dict = field(default_factory=dict)
+
+
+@dataclass
+class ConnectionCredential(BaseDataClass):
+    uuid: str
+    description: str = None
+    error: ErrorMessage = None
+    name: str = None
+    required: bool = False
+    valid: bool = False
+    value: str = None
+
+    def __post_init__(self):
+        if self.error and isinstance(self.error, dict):
+            self.error = ErrorMessage.load(**self.error)
+
+
+class SetupStepStatus(str, Enum):
+    COMPLETED = 'completed'
+    ERROR = 'error'
+    INCOMPLETE = 'incomplete'
+
+
+@dataclass
+class SetupStep(BaseDataClass):
+    name: str
+    uuid: str
+    description: str = None
+    error: ErrorMessage = None
+    group: bool = False
+    required: bool = True
+    status: SetupStepStatus = None
+    steps: List[SetupStepStatus] = field(default_factory=list)
+    tab: ComputeManagementApplicationTab = None
+
+    def __post_init__(self):
+        if self.error and isinstance(self.error, dict):
+            self.error = ErrorMessage.load(**self.error)
+
+        if self.status and isinstance(self.status, str):
+            self.status = SetupStepStatus(self.status)
+
+        if self.steps and isinstance(self.steps, list):
+            steps = []
+
+            for step in self.steps:
+                if isinstance(step, dict):
+                    steps.append(SetupStep.load(**step))
+                else:
+                    steps.append(step)
+
+            self.steps = steps
+
+        if self.tab and isinstance(self.tab, str):
+            self.tab = ComputeManagementApplicationTab(self.tab)
+
+    def status_calculated(self) -> SetupStepStatus:
+        if not hasattr(self, '_status_calculated'):
+            self._status_calculated = None
+
+        if self._status_calculated:
+            return self._status_calculated
+
+        if self.status:
+            self._status_calculated = self.status
+            return self._status_calculated
+
+        if self.steps:
+            if all([
+                not step.required or SetupStepStatus.COMPLETED == step.status_calculated()
+                for step in self.steps
+            ]):
+                self._status_calculated = SetupStepStatus.COMPLETED
+            elif any([
+                not step.required or SetupStepStatus.ERROR == step.status_calculated()
+                for step in self.steps
+            ]):
+                self._status_calculated = SetupStepStatus.ERROR
+            elif any([
+                not step.required or SetupStepStatus.INCOMPLETE == step.status_calculated()
+                for step in self.steps
+            ]):
+                self._status_calculated = SetupStepStatus.INCOMPLETE
+
+        return self._status_calculated
+
+    def to_dict(self, **kwargs) -> Dict:
+        return merge_dict(super().to_dict(), dict(
+            status_calculated=self.status_calculated(),
+        ))
+
+
+@dataclass
+class ComputeConnectionAction(BaseDataClass):
+    name: str
+    uuid: ComputeConnectionActionUUID
+    description: str = None
+
+    def __post_init__(self):
+        self.serialize_attribute_enum('uuid', ComputeConnectionActionUUID)
+
+
+@dataclass
+class ComputeConnection(SetupStep):
+    actions: List[ComputeConnectionAction] = field(default_factory=list)
+    attributes: Dict = field(default_factory=dict)
+    connection: Dict = field(default_factory=dict)
+    required: bool = False
+    state: ComputeConnectionState = None
+    steps: List[SetupStep] = field(default_factory=list)
+
+    def __post_init__(self):
+        super().__post_init__()
+        self.serialize_attribute_classes('actions', ComputeConnectionAction)
+        self.serialize_attribute_classes('steps', SetupStep)
+        self.serialize_attribute_enum('state', ComputeConnectionState)
+
+
+class ComputeService:
+    uuid = ComputeServiceUUID.STANDALONE_CLUSTER
+
+    def __init__(self, project: Project, with_clusters: bool = False):
+        self.project = project
+        self.with_clusters = with_clusters
+
+    @classmethod
+    def build(self, project: Project, with_clusters: bool = False):
+        service_class = self
+
+        if project and project.spark_config:
+            if project.emr_config:
+                from mage_ai.services.compute.aws.models import AWSEMRComputeService
+
+                service_class = AWSEMRComputeService
+
+        return service_class(project=project, with_clusters=with_clusters)
+
+    def to_dict(self, **kwargs) -> Dict:
+        result = dict(
+            connection_credentials=[m.to_dict() for m in self.connection_credentials()],
+            setup_steps=[m.to_dict() for m in self.setup_steps()],
+            uuid=self.uuid,
+        )
+
+        if self.with_clusters:
+            result['clusters'] = self.clusters_and_metadata()
+
+        return result
+
+    @property
+    def uuid(self) -> ComputeServiceUUID:
+        return self.__class__.uuid
+
+    @abstractmethod
+    def active_cluster(self, **kwargs) -> Dict:
+        pass
+
+    @abstractmethod
+    def activate_cluster(self, **kwargs) -> Dict:
+        pass
+
+    @abstractmethod
+    def create_cluster(self, **kwargs) -> Dict:
+        pass
+
+    @abstractmethod
+    def clusters_and_metadata(self, **kwargs) -> Dict:
+        pass
+
+    @abstractmethod
+    def get_cluster_details(self, **kwargs) -> Dict:
+        pass
+
+    @abstractmethod
+    def connection_credentials(self) -> List[ConnectionCredential]:
+        pass
+
+    @abstractmethod
+    def compute_connections(self) -> List[ComputeConnection]:
+        pass
+
+    @abstractmethod
+    def setup_steps(self) -> List[SetupStep]:
+        pass
+
+    @abstractmethod
+    def terminate_clusters(self, cluster_ids: List[str]) -> None:
+        pass
+
+    @abstractmethod
+    def update_cluster(self, cluster_id: str, payload: Dict) -> Dict:
+        pass

mage_ai/services/k8s/job_manager.py

@@ -173,6 +173,9 @@ class JobManager():
         mage_server_container_spec = self.get_mage_server_container()
         container_spec.env = container_spec.env + \
             [item for item in mage_server_container_spec.env if item not in container_spec.env]
+        container_spec.env_from = (container_spec.env_from or []) + \
+            [item for item in (mage_server_container_spec.env_from or [])
+             if item not in (container_spec.env_from or [])]
         container_spec.volume_mounts = container_spec.volume_mounts + \
             [item for item in mage_server_container_spec.volume_mounts
                 if item not in container_spec.volume_mounts]

mage_ai/services/spark/api/aws_emr.py

@@ -0,0 +1,38 @@
+from mage_ai.services.spark.api.constants import SPARK_UI_HOST, SPARK_UI_PORT_AWS_EMR
+from mage_ai.services.spark.api.local import LocalAPI
+from mage_ai.services.ssh.aws.emr.models import SSHTunnel
+
+
+class AwsEmrAPI(LocalAPI):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+        if not self.application_id:
+            applications = self.applications_sync()
+            if applications:
+                applications = sorted(applications, key=lambda x: x.id, reverse=True)
+                application = applications[0]
+                self.application_id = application.calculated_id()
+
+    @property
+    def spark_ui_url(self) -> str:
+        url = f'http://{SPARK_UI_HOST}:{SPARK_UI_PORT_AWS_EMR}'
+
+        tunnel = SSHTunnel()
+        if tunnel:
+            connection_details = tunnel.connection_details()
+            host = connection_details.get('host')
+            port = connection_details.get('port')
+            url = f'http://{host or SPARK_UI_HOST}:{port or SPARK_UI_PORT_AWS_EMR}'
+
+        if self.application_spark_ui_url:
+            url = self.application_spark_ui_url
+        elif self.spark_session:
+            url = self.spark_session.sparkContext.uiWebUrl
+
+        return url
+
+    def ready_for_requests(self, **kwargs) -> bool:
+        tunnel = SSHTunnel()
+        ready = tunnel and tunnel.is_active()
+        return True if ready else False

mage_ai/services/spark/api/base.py

@@ -60,7 +60,7 @@ class BaseAPI(ABC):
         if self._application:
             return self._application
 
-        self._application = Application(
+        self._application = Application.load(
             id=self.application_id,
             spark_ui_url=self.application_spark_ui_url,
         )
@@ -154,8 +154,11 @@ class BaseAPI(ABC):
     async def environment(self, application_id: str = None, **kwargs) -> Environment:
         pass
 
+    def ready_for_requests(self, **kwargs) -> bool:
+        return True if self.spark_session else False
+
     async def get(self, path: str, host: str = None, query: Dict = None):
-        if not self.spark_session:
+        if not self.ready_for_requests():
             return {}
 
         url = f'{self.endpoint(host=host)}{path}'
@@ -167,12 +170,12 @@ class BaseAPI(ABC):
         if response.status_code == 200:
             return response.json()
         else:
-            print(f'[WARNING] {self.__class__.__name__} {url}: {response}')
+            print(f'[WARNING] {self.__class__.__name__} async {url}: {response}')
 
         return {}
 
     def get_sync(self, path: str, host: str = None, query: Dict = None):
-        if not self.spark_session:
+        if not self.ready_for_requests():
             return {}
 
         url = f'{self.endpoint(host=host)}{path}'

mage_ai/services/spark/api/constants.py

@@ -0,0 +1,4 @@
+SPARK_UI_API_VERSION = 'v1'
+SPARK_UI_HOST = 'localhost'
+SPARK_UI_PORT = 4040
+SPARK_UI_PORT_AWS_EMR = 18080

mage_ai/services/spark/api/local.py

@@ -1,6 +1,11 @@
 from typing import Dict, List
 
 from mage_ai.services.spark.api.base import BaseAPI
+from mage_ai.services.spark.api.constants import (
+    SPARK_UI_API_VERSION,
+    SPARK_UI_HOST,
+    SPARK_UI_PORT,
+)
 from mage_ai.services.spark.models.applications import Application
 from mage_ai.services.spark.models.environments import Environment
 from mage_ai.services.spark.models.executors import Executor
@@ -15,10 +20,6 @@ from mage_ai.services.spark.models.stages import (
 from mage_ai.services.spark.models.threads import Thread
 from mage_ai.shared.hash import index_by
 
-API_VERSION = 'v1'
-SPARK_UI_HOST = 'localhost'
-SPARK_UI_PORT = '4040'
-
 
 class LocalAPI(BaseAPI):
     @property
@@ -33,7 +34,7 @@ class LocalAPI(BaseAPI):
         return url
 
     def endpoint(self, host: str = None, **kwargs) -> str:
-        return f'{host or self.spark_ui_url}/api/{API_VERSION}'
+        return f'{host or self.spark_ui_url}/api/{SPARK_UI_API_VERSION}'
 
     def applications_sync(self, **kwargs) -> List[Application]:
         models = self.get_sync('/applications')
@@ -46,10 +47,10 @@ class LocalAPI(BaseAPI):
         applications_cache = Application.get_applications_from_cache()
         if applications_cache:
             for application in applications_cache.values():
-                if application.id in mapping:
+                if application.calculated_id() in mapping:
                     continue
                 applications.append(application)
-                mapping[application.id] = application
+                mapping[application.calculated_id()] = application
 
         return applications
 
@@ -64,11 +65,11 @@ class LocalAPI(BaseAPI):
             applications = [self.application]
 
         for application in applications:
-            if application.id == application_id:
+            if application.calculated_id() == application_id:
                 arr.extend([Job.load(application=application, **model) for model in models])
             else:
                 jobs = self.get_sync(
-                    f'/applications/{application.id}/jobs',
+                    f'/applications/{application.calculated_id()}/jobs',
                     host=application.spark_ui_url,
                 )
                 arr.extend([Job.load(application=application, **model) for model in jobs])
@@ -89,7 +90,7 @@ class LocalAPI(BaseAPI):
             applications = [self.application]
 
         for application in applications:
-            if application.id == application_id:
+            if application.calculated_id() == application_id:
                 arr.extend(sorted(
                     [Sql.load(
                         application=application,
@@ -100,7 +101,7 @@ class LocalAPI(BaseAPI):
                 ))
             else:
                 models2 = self.get_sync(
-                    f'/applications/{application.id}/sql',
+                    f'/applications/{application.calculated_id()}/sql',
                     host=application.spark_ui_url,
                     query=query,
                 )
@@ -129,14 +130,14 @@ class LocalAPI(BaseAPI):
             applications = [self.application]
 
         for application in applications:
-            if application.id == application_id:
+            if application.calculated_id() == application_id:
                 arr.extend([model_class.load(
                     application=application,
                     **model,
                 ) for model in models])
             else:
                 models2 = self.get_sync(
-                    f'/applications/{application.id}/stages',
+                    f'/applications/{application.calculated_id()}/stages',
                     host=application.spark_ui_url,
                     query=query,
                 )
@@ -158,10 +159,10 @@ class LocalAPI(BaseAPI):
         applications_cache = Application.get_applications_from_cache()
         if applications_cache:
             for application in applications_cache.values():
-                if application.id in mapping:
+                if application.calculated_id() in mapping:
                     continue
                 applications.append(application)
-                mapping[application.id] = application
+                mapping[application.calculated_id()] = application
 
         return applications
 
@@ -176,11 +177,11 @@ class LocalAPI(BaseAPI):
             applications = [self.application]
 
         for application in applications:
-            if application.id == application_id:
+            if application.calculated_id() == application_id:
                 arr.extend([Job.load(application=application, **model) for model in models])
             else:
                 jobs = await self.get(
-                    f'/applications/{application.id}/jobs',
+                    f'/applications/{application.calculated_id()}/jobs',
                     host=application.spark_ui_url,
                 )
                 arr.extend([Job.load(application=application, **model) for model in jobs])
@@ -200,7 +201,7 @@ class LocalAPI(BaseAPI):
             host=application_spark_ui_url,
         )
         return Job.load(
-            application=Application(
+            application=Application.load(
                 id=application_id,
                 spark_ui_url=application_spark_ui_url,
             ),
@@ -221,14 +222,14 @@ class LocalAPI(BaseAPI):
             applications = [self.application]
 
         for application in applications:
-            if application.id == application_id:
+            if application.calculated_id() == application_id:
                 arr.extend([model_class.load(
                     application=application,
                     **model,
                 ) for model in models])
             else:
                 models2 = await self.get(
-                    f'/applications/{application.id}/stages',
+                    f'/applications/{application.calculated_id()}/stages',
                     host=application.spark_ui_url,
                     query=query,
                 )
@@ -254,7 +255,7 @@ class LocalAPI(BaseAPI):
             query=query,
         )
         return Stage.load(
-            application=Application(
+            application=Application.load(
                 id=application_id,
                 spark_ui_url=application_spark_ui_url,
             ),
@@ -344,7 +345,7 @@ class LocalAPI(BaseAPI):
             applications = [self.application]
 
         for application in applications:
-            if application.id == application_id:
+            if application.calculated_id() == application_id:
                 arr.extend(sorted(
                     [Sql.load(
                         application=application,
@@ -355,7 +356,7 @@ class LocalAPI(BaseAPI):
                 ))
             else:
                 models2 = await self.get(
-                    f'/applications/{application.id}/sql',
+                    f'/applications/{application.calculated_id()}/sql',
                     host=application.spark_ui_url,
                     query=query,
                 )
@@ -383,7 +384,7 @@ class LocalAPI(BaseAPI):
             host=application_spark_ui_url,
         )
         return Sql.load(
-            application=Application(
+            application=Application.load(
                 id=application_id,
                 spark_ui_url=application_spark_ui_url,
             ),

mage_ai/services/spark/api/service.py

@@ -1,5 +1,6 @@
+from mage_ai.services.spark.api.base import BaseAPI
 from mage_ai.services.spark.api.local import LocalAPI
-from mage_ai.services.spark.constants import ComputeService
+from mage_ai.services.spark.constants import ComputeServiceUUID
 from mage_ai.services.spark.utils import get_compute_service
 
 
@@ -7,18 +8,27 @@ class API:
     @classmethod
     def build(
         self,
-        all_applications: bool = False,
+        all_applications: bool = True,
         application_id: str = None,
         application_spark_ui_url: str = None,
         repo_config=None,
         spark_session=None,
-    ):
-        compute_service = get_compute_service(repo_config)
+    ) -> BaseAPI:
+        compute_service = get_compute_service(repo_config, ignore_active_kernel=True)
 
-        if ComputeService.STANDALONE_CLUSTER == compute_service:
+        if ComputeServiceUUID.STANDALONE_CLUSTER == compute_service:
             return LocalAPI(
                 all_applications=all_applications,
                 spark_session=spark_session,
                 application_id=application_id,
                 application_spark_ui_url=application_spark_ui_url,
             )
+        elif ComputeServiceUUID.AWS_EMR == compute_service:
+            from mage_ai.services.spark.api.aws_emr import AwsEmrAPI
+
+            return AwsEmrAPI(
+                all_applications=all_applications,
+                spark_session=spark_session,
+                application_id=application_id,
+                application_spark_ui_url=application_spark_ui_url,
+            )

mage_ai/services/spark/constants.py

@@ -3,7 +3,7 @@ from enum import Enum
 SPARK_DIRECTORY_NAME = '.spark'
 
 
-class ComputeService(str, Enum):
+class ComputeServiceUUID(str, Enum):
     AWS_EMR = 'AWS_EMR'
     STANDALONE_CLUSTER = 'STANDALONE_CLUSTER'
 

mage_ai/services/spark/models/applications.py

@@ -1,9 +1,13 @@
 import json
 import os
+import urllib.parse
 from dataclasses import dataclass
 from typing import Dict, List
 
+from mage_ai.services.spark.constants import ComputeServiceUUID
 from mage_ai.services.spark.models.base import BaseSparkModel
+from mage_ai.services.spark.utils import get_compute_service
+from mage_ai.shared.hash import merge_dict
 
 
 @dataclass
@@ -24,12 +28,32 @@ class ApplicationAttempt(BaseSparkModel):
 class Application(BaseSparkModel):
     id: str  # local-1697360611228
     attempts: List[ApplicationAttempt] = None
+    attempts_count: int = None
     name: str = None  # my spark app
     spark_ui_url: str = None
 
     def __post_init__(self):
         if self.attempts and isinstance(self.attempts, list):
-            self.attempts = [ApplicationAttempt.load(**m) for m in self.attempts]
+            arr = []
+            for m in self.attempts:
+                if m and isinstance(m, dict):
+                    arr.append(ApplicationAttempt.load(**m))
+                else:
+                    arr.append(m)
+            self.attempts = arr
+
+    @classmethod
+    def load(self, **kwargs):
+        payload = kwargs.copy() if kwargs else {}
+
+        if ComputeServiceUUID.AWS_EMR == get_compute_service(ignore_active_kernel=True):
+            if payload.get('id'):
+                parts = urllib.parse.unquote(payload.get('id')).split('/')
+                if len(parts) >= 2:
+                    payload['id'] = parts[0]
+                    payload['attempts_count'] = parts[1]
+
+        return super().load(**payload)
 
     @classmethod
     def __cache_file_path(self):
@@ -52,7 +76,7 @@ class Application(BaseSparkModel):
                     data.update(json.loads(content) or {})
 
         data.update({
-            application.id: application.to_dict(),
+            application.calculated_id(): application.to_dict(),
         })
 
         with open(self.__cache_file_path(), 'w') as f:
@@ -68,6 +92,24 @@ class Application(BaseSparkModel):
                 if content:
                     for application_dict in json.loads(content).values():
                         application = self.load(**application_dict)
-                        data[application.id] = application
+                        data[application.calculated_id()] = application
 
         return data
+
+    def calculated_id(self) -> str:
+        if ComputeServiceUUID.AWS_EMR == get_compute_service(ignore_active_kernel=True):
+            count = 1
+            if self.attempts:
+                count = len(self.attempts)
+            elif self.attempts_count is not None:
+                count = self.attempts_count
+
+            return f'{self.id}/{count}'
+
+        return self.id
+
+    def to_dict(self, **kwargs) -> Dict:
+        return merge_dict(
+            super().to_dict(**kwargs),
+            dict(calculated_id=self.calculated_id()),
+        )

mage_ai/services/spark/models/base.py

@@ -1,27 +1,14 @@
 import os
-from dataclasses import asdict, dataclass
-from typing import Dict
-
-import inflection
+from dataclasses import dataclass
 
 from mage_ai.data_preparation.repo_manager import get_repo_config
 from mage_ai.services.spark.constants import SPARK_DIRECTORY_NAME
 from mage_ai.settings.repo import get_repo_path
+from mage_ai.shared.models import BaseDataClass
 
 
 @dataclass
-class BaseSparkModel:
-    @classmethod
-    def load(self, **kwargs):
-        return self(**self.load_to_dict(**kwargs))
-
-    @classmethod
-    def load_to_dict(self, **kwargs) -> Dict:
-        data = {}
-        for key, value in kwargs.items():
-            data[inflection.underscore(key)] = value
-        return data
-
+class BaseSparkModel(BaseDataClass):
     @classmethod
     def cache_dir_path(self) -> str:
         repo_config = get_repo_config(repo_path=get_repo_path())
@@ -30,6 +17,3 @@ class BaseSparkModel:
             repo_config.variables_dir,
             SPARK_DIRECTORY_NAME,
         )
-
-    def to_dict(self) -> Dict:
-        return asdict(self)