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

mage_ai/services/spark/models/environments.py

@@ -19,11 +19,11 @@ class ExecutorResources(BaseSparkModel):
     off_heap: Resource = None
 
     def __post_init__(self):
-        if self.cores:
+        if self.cores and isinstance(self.cores, dict):
             self.cores = Resource.load(**self.cores)
-        if self.memory:
+        if self.memory and isinstance(self.memory, dict):
             self.memory = Resource.load(**self.memory)
-        if self.off_heap:
+        if self.off_heap and isinstance(self.off_heap, dict):
             self.off_heap = Resource.load(**self.off_heap)
 
 
@@ -32,7 +32,7 @@ class TaskResources(BaseSparkModel):
     cpus: Resource = None
 
     def __post_init__(self):
-        if self.cpus:
+        if self.cpus and isinstance(self.cpus, dict):
             self.cpus = Resource.load(**self.cpus)
 
 
@@ -43,9 +43,9 @@ class ResourceProfiles(BaseSparkModel):
     task_resources: TaskResources = None
 
     def __post_init__(self):
-        if self.executor_resources:
+        if self.executor_resources and isinstance(self.executor_resources, dict):
             self.executor_resources = ExecutorResources.load(**self.executor_resources)
-        if self.task_resources:
+        if self.task_resources and isinstance(self.task_resources, dict):
             self.task_resources = TaskResources.load(**self.task_resources)
 
 
@@ -117,9 +117,14 @@ class Environment(BaseSparkModel):
     system_properties: List[List[str]] = field(default_factory=list)
 
     def __post_init__(self):
-        if self.resource_profiles:
-            self.resource_profiles = \
-                [ResourceProfiles.load(**model) for model in self.resource_profiles]
-
-        if self.runtime:
+        if self.resource_profiles and isinstance(self.resource_profiles, list):
+            arr = []
+            for model in self.resource_profiles:
+                if model and isinstance(model, dict):
+                    arr.append(ResourceProfiles.load(**model))
+                else:
+                    arr.append(model)
+            self.resource_profiles = arr
+
+        if self.runtime and isinstance(self.runtime, dict):
             self.runtime = Runtime.load(**self.runtime)

mage_ai/services/spark/models/executors.py

@@ -38,8 +38,8 @@ class Executor(BaseSparkModel):
     total_tasks: int = None  # 28
 
     def __post_init__(self):
-        if self.memory_metrics:
+        if self.memory_metrics and isinstance(self.memory_metrics, dict):
             self.memory_metrics = MemoryMetrics.load(**self.memory_metrics)
 
-        if self.peak_memory_metrics:
+        if self.peak_memory_metrics and isinstance(self.peak_memory_metrics, dict):
             self.peak_memory_metrics = Metrics.load(**self.peak_memory_metrics)

mage_ai/services/spark/models/sqls.py

@@ -32,8 +32,14 @@ class Node(BaseSparkModel):
     whole_stage_codegen_id: int = None  # 2
 
     def __post_init__(self):
-        if self.metrics:
-            self.metrics = [Metric.load(**model) for model in self.metrics]
+        if self.metrics and isinstance(self.metrics, list):
+            arr = []
+            for model in self.metrics:
+                if model and isinstance(model, dict):
+                    arr.append(Metric.load(**model))
+                else:
+                    arr.append(model)
+            self.metrics = arr
 
     @property
     def id(self) -> int:
@@ -100,19 +106,44 @@ class Sql(BaseSparkModel):
         if self.application and isinstance(self.application, dict):
             self.application = Application.load(**self.application)
 
-        if self.edges:
-            self.edges = [Edge.load(**edge) for edge in self.edges]
-
-        if self.jobs:
-            self.jobs = [Job.load(**m) for m in self.jobs]
-
-        if self.nodes:
-            self.nodes = [Node.load(**node) for node in self.nodes]
-
-        if self.stages:
-            self.stages = [StageAttempt.load(**m) for m in self.stages]
-
-        if self.status:
+        if self.edges and isinstance(self.edges, list):
+            arr = []
+            for edge in self.edges:
+                if edge and isinstance(edge, dict):
+                    arr.append(Edge.load(**edge))
+                else:
+                    arr.append(edge)
+            self.edges = arr
+
+        if self.jobs and isinstance(self.jobs, list):
+            arr = []
+            for m in self.jobs:
+                if m and isinstance(m, dict):
+                    arr.append(Job.load(**m))
+                else:
+                    arr.append(m)
+            self.jobs = arr
+
+        if self.nodes and isinstance(self.nodes, list):
+            arr = []
+            for node in self.nodes:
+                if node and isinstance(node, dict):
+                    arr.append(Node.load(**node))
+                else:
+                    arr.append(node)
+            self.nodes = arr
+
+        if self.stages and isinstance(self.stages, list):
+            arr = []
+            for m in self.stages:
+                if m and isinstance(m, dict):
+                    arr.append(StageAttempt.load(**m))
+                else:
+                    arr.append(m)
+
+            self.stages = arr
+
+        if self.status and isinstance(self.status, str):
             try:
                 self.status = SqlStatus(self.status)
             except ValueError as err:

mage_ai/services/spark/models/stages.py

@@ -10,10 +10,12 @@ from mage_ai.services.spark.models.metrics import Metrics
 class Locality(str, Enum):
     NODE_LOCAL = 'NODE_LOCAL'
     PROCESS_LOCAL = 'PROCESS_LOCAL'
+    RACK_LOCAL = 'RACK_LOCAL'
 
 
 class StageStatus(str, Enum):
     COMPLETE = 'COMPLETE'
+    PENDING = 'PENDING'
     SKIPPED = 'SKIPPED'
 
 
@@ -42,7 +44,7 @@ class Driver(BaseSparkModel):
     task_time: int = None  # 11
 
     def __post_init__(self):
-        if self.peak_memory_metrics:
+        if self.peak_memory_metrics and isinstance(self.peak_memory_metrics, dict):
             self.peak_memory_metrics = Metrics.load(**self.peak_memory_metrics)
 
 
@@ -51,7 +53,7 @@ class ExecutorSummary(BaseSparkModel):
     driver: Driver = None
 
     def __post_init__(self):
-        if self.driver:
+        if self.driver and isinstance(self.driver, dict):
             self.driver = Driver.load(**self.driver)
 
 
@@ -100,7 +102,10 @@ class ShuffleReadMetricsQuantile(BaseSparkModel):
     total_blocks_fetched: List[int] = field(default_factory=list)
 
     def __post_init__(self):
-        if self.shuffle_push_read_metrics_dist:
+        if self.shuffle_push_read_metrics_dist and isinstance(
+            self.shuffle_push_read_metrics_dist,
+            dict,
+        ):
             self.shuffle_push_read_metrics_dist = ShufflePushReadMetricsDistQuantile.load(
                 **self.shuffle_push_read_metrics_dist,
             )
@@ -145,7 +150,7 @@ class ShuffleReadMetrics(BaseSparkModel):
     shuffle_push_read_metrics: PushReadMetrics = None
 
     def __post_init__(self):
-        if self.shuffle_push_read_metrics:
+        if self.shuffle_push_read_metrics and isinstance(self.shuffle_push_read_metrics, dict):
             self.shuffle_push_read_metrics = PushReadMetrics.load(**self.shuffle_push_read_metrics)
 
 
@@ -178,16 +183,16 @@ class TaskMetricsQuantiles(BaseSparkModel):
     shuffle_write_metrics: ShuffleWriteMetricsQuantile = None
 
     def __post_init__(self):
-        if self.input_metrics:
+        if self.input_metrics and isinstance(self.input_metrics, dict):
             self.input_metrics = InputMetricsQuantile.load(**self.input_metrics)
 
-        if self.output_metrics:
+        if self.output_metrics and isinstance(self.output_metrics, dict):
             self.output_metrics = OutputMetricsQuantile.load(**self.output_metrics)
 
-        if self.shuffle_read_metrics:
+        if self.shuffle_read_metrics and isinstance(self.shuffle_read_metrics, dict):
             self.shuffle_read_metrics = ShuffleReadMetricsQuantile.load(**self.shuffle_read_metrics)
 
-        if self.shuffle_write_metrics:
+        if self.shuffle_write_metrics and isinstance(self.shuffle_write_metrics, dict):
             self.shuffle_write_metrics = ShuffleWriteMetricsQuantile.load(
                 **self.shuffle_write_metrics,
             )
@@ -211,13 +216,13 @@ class TaskMetrics(BaseSparkModel):
     shuffle_write_metrics: ShuffleWriteMetrics = None
 
     def __post_init__(self):
-        if self.input_metrics:
+        if self.input_metrics and isinstance(self.input_metrics, dict):
             self.input_metrics = InputMetrics.load(**self.input_metrics)
-        if self.output_metrics:
+        if self.output_metrics and isinstance(self.output_metrics, dict):
             self.output_metrics = OutputMetrics.load(**self.output_metrics)
-        if self.shuffle_read_metrics:
+        if self.shuffle_read_metrics and isinstance(self.shuffle_read_metrics, dict):
             self.shuffle_read_metrics = ShuffleReadMetrics.load(**self.shuffle_read_metrics)
-        if self.shuffle_write_metrics:
+        if self.shuffle_write_metrics and isinstance(self.shuffle_write_metrics, dict):
             self.shuffle_write_metrics = ShuffleWriteMetrics.load(**self.shuffle_write_metrics)
 
 
@@ -241,8 +246,8 @@ class ExecutorMetricsDistributions(BaseSparkModel):
     task_time: List[int] = field(default_factory=list)
 
     def __post_init__(self):
-        if self.peak_memory_metrics:
-            self.peak_memory_metrics = Metrics(**self.peak_memory_metrics)
+        if self.peak_memory_metrics and isinstance(self.peak_memory_metrics, dict):
+            self.peak_memory_metrics = Metrics.load(**self.peak_memory_metrics)
 
 
 @dataclass
@@ -265,21 +270,21 @@ class Task(BaseSparkModel):
     task_metrics: TaskMetrics = None  # {}
 
     def __post_init__(self):
-        if self.status:
+        if self.status and isinstance(self.status, str):
             try:
                 self.status = TaskStatus(self.status)
             except ValueError as err:
                 print(f'[WARNING] Task: {err}')
                 self.status = self.status
 
-        if self.task_locality:
+        if self.task_locality and isinstance(self.task_locality, str):
             try:
                 self.task_locality = Locality(self.task_locality)
             except ValueError as err:
                 print(f'[WARNING] Task: {err}')
                 self.task_locality = self.task_locality
 
-        if self.task_metrics:
+        if self.task_metrics and isinstance(self.task_metrics, dict):
             self.task_metrics = TaskMetrics.load(**self.task_metrics)
 
     @property
@@ -370,16 +375,19 @@ class StageBase(BaseSparkModel):
         if self.application and isinstance(self.application, dict):
             self.application = Application.load(**self.application)
 
-        if self.executor_metrics_distributions:
+        if self.executor_metrics_distributions and isinstance(
+            self.executor_metrics_distributions,
+            dict,
+        ):
             if 'quantiles' not in self.executor_metrics_distributions:
                 self.executor_metrics_distributions = ExecutorMetricsDistributions.load(
                     **self.executor_metrics_distributions,
                 )
 
-        if self.peak_executor_metrics:
+        if self.peak_executor_metrics and isinstance(self.peak_executor_metrics, dict):
             self.peak_executor_metrics = Metrics.load(**self.peak_executor_metrics)
 
-        if self.task_metrics_distributions:
+        if self.task_metrics_distributions and isinstance(self.task_metrics_distributions, dict):
             if 'quantiles' in self.task_metrics_distributions:
                 self.task_metrics_distributions = TaskMetricsQuantiles.load(
                     **self.task_metrics_distributions,
@@ -389,7 +397,7 @@ class StageBase(BaseSparkModel):
                     **self.task_metrics_distributions,
                 )
 
-        if self.status:
+        if self.status and isinstance(self.status, str):
             try:
                 self.status = StageStatus(self.status)
             except ValueError as err:
@@ -481,13 +489,16 @@ class StageAttempt(StageBase):
     def __post_init__(self):
         super().__post_init__()
 
-        if self.executor_summary:
+        if self.executor_summary and isinstance(self.executor_summary, dict):
             self.executor_summary = ExecutorSummary.load(**self.executor_summary)
 
-        if self.tasks:
+        if self.tasks and isinstance(self.tasks, dict):
             tasks = {}
             for k, v in self.tasks.items():
-                tasks[k] = Task.load(**v)
+                if v and isinstance(v, dict):
+                    tasks[k] = Task.load(**v)
+                else:
+                    tasks[k] = v
             self.tasks = tasks
 
     @property
@@ -497,13 +508,22 @@ class StageAttempt(StageBase):
 
 @dataclass
 class Stage(StageBase):
-    stage_attempts: List[Dict] = field(default_factory=list)
+    stage_attempts: List[StageAttempt] = field(default_factory=list)
 
     def __post_init__(self):
         super().__post_init__()
 
-        if self.stage_attempts:
-            self.stage_attempts = [StageAttempt.load(**m) for m in self.stage_attempts]
+        if self.stage_attempts and isinstance(self.stage_attempts, list):
+            arr = []
+            for m in self.stage_attempts:
+                if m and isinstance(m, dict):
+                    m_updated = m.copy()
+                    if not m_updated.get('application'):
+                        m_updated['application'] = self.application
+                    arr.append(StageAttempt.load(**m_updated))
+                else:
+                    arr.append(m)
+            self.stage_attempts = arr
 
 
 @dataclass
@@ -557,7 +577,10 @@ class ShuffleReadMetricsSummary(BaseSparkModel):
     total_blocks_fetched: List[float] = field(default_factory=list)  # [ 0.0, 0.0, 0.0, 0.0, 0.0 ]
 
     def __post_init__(self):
-        if self.shuffle_push_read_metrics_dist:
+        if self.shuffle_push_read_metrics_dist and isinstance(
+            self.shuffle_push_read_metrics_dist,
+            dict,
+        ):
             self.shuffle_push_read_metrics_dist = PushReadMetricsDistSummary(
                 **self.shuffle_push_read_metrics_dist,
             )
@@ -597,13 +620,13 @@ class StageAttemptTaskSummary(BaseSparkModel):
     shuffle_write_metrics: ShuffleWriteMetricsSummary = None
 
     def __post_init__(self):
-        if self.input_metrics:
+        if self.input_metrics and isinstance(self.input_metrics, dict):
             self.input_metrics = InputMetricsSummary.load(**self.input_metrics)
-        if self.output_metrics:
+        if self.output_metrics and isinstance(self.output_metrics, dict):
             self.output_metrics = OutputMetricsSummary.load(**self.output_metrics)
-        if self.shuffle_read_metrics:
+        if self.shuffle_read_metrics and isinstance(self.shuffle_read_metrics, dict):
             self.shuffle_read_metrics = ShuffleReadMetricsSummary.load(**self.shuffle_read_metrics)
-        if self.shuffle_write_metrics:
+        if self.shuffle_write_metrics and isinstance(self.shuffle_write_metrics, dict):
             self.shuffle_write_metrics = ShuffleWriteMetricsSummary.load(
                 **self.shuffle_write_metrics,
             )

mage_ai/services/spark/models/threads.py

@@ -30,10 +30,10 @@ class Thread(BaseSparkModel):
     thread_state: ThreadState = None  # "TIMED_WAITING"
 
     def __post_init__(self):
-        if self.stack_trace:
+        if self.stack_trace and isinstance(self.stack_trace, dict):
             self.stack_trace = StackTrace.load(**self.stack_trace)
 
-        if self.thread_state:
+        if self.thread_state and isinstance(self.thread_state, str):
             try:
                 self.thread_state = ThreadState(self.thread_state)
             except ValueError as err:

mage_ai/services/spark/utils.py

@@ -1,21 +1,37 @@
+from typing import Dict
+
 from mage_ai.data_preparation.repo_manager import get_repo_config
-from mage_ai.server.active_kernel import get_active_kernel_name
 from mage_ai.server.kernels import KernelName
-from mage_ai.services.spark.constants import ComputeService, SparkMaster
+from mage_ai.services.spark.constants import ComputeServiceUUID, SparkMaster
 from mage_ai.shared.utils import is_spark_env
 
 
-def get_compute_service(repo_config=None) -> ComputeService:
+def get_compute_service(
+    emr_config: Dict = None,
+    repo_config=None,
+    ignore_active_kernel: bool = False,
+    kernel_name: KernelName = None,
+) -> ComputeServiceUUID:
     if not repo_config:
         repo_config = get_repo_config()
 
+    if repo_config and emr_config:
+        repo_config.emr_config = emr_config
+
     if not repo_config:
+        # This is probably happening because there is no repo config in Spark EMR
         return None
 
-    if repo_config.emr_config and get_active_kernel_name() == KernelName.PYSPARK:
-        return ComputeService.AWS_EMR
+    if not kernel_name:
+        from mage_ai.server.active_kernel import get_active_kernel_name
+
+        kernel_name = get_active_kernel_name()
+
+    if repo_config.emr_config and (KernelName.PYSPARK == kernel_name or ignore_active_kernel):
+        return ComputeServiceUUID.AWS_EMR
     elif is_spark_env() and repo_config.spark_config and \
             SparkMaster.LOCAL.value == repo_config.spark_config.get('spark_master'):
-        return ComputeService.STANDALONE_CLUSTER
+
+        return ComputeServiceUUID.STANDALONE_CLUSTER
 
     return None

mage_ai/services/ssh/aws/emr/constants.py

@@ -0,0 +1,10 @@
+from mage_ai.services.spark.api.constants import SPARK_UI_HOST, SPARK_UI_PORT_AWS_EMR
+
+SSH_USERNAME_DEFAULT = 'hadoop'
+SSH_DEFAULTS = dict(
+    spark_ui_host_local=SPARK_UI_HOST,
+    spark_ui_host_remote=SPARK_UI_HOST,
+    spark_ui_port_local=SPARK_UI_PORT_AWS_EMR,
+    spark_ui_port_remote=SPARK_UI_PORT_AWS_EMR,
+    ssh_username=SSH_USERNAME_DEFAULT,
+)