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

mage_ai/server/websocket_server.py

@@ -1,4 +1,5 @@
 import asyncio
+import base64
 import json
 import multiprocessing
 import os
@@ -9,6 +10,7 @@ from datetime import datetime, timedelta
 from distutils.file_util import copy_file
 from typing import Dict, List
 
+import simplejson
 import tornado.websocket
 from jupyter_client import KernelClient
 
@@ -127,8 +129,9 @@ def run_pipeline(
         )
     except Exception:
         trace = traceback.format_exc().splitlines()
-        add_pipeline_message(f'Pipeline {pipeline.uuid} execution failed with error:',
-                             metadata=metadata)
+        add_pipeline_message(
+            f'Pipeline {pipeline.uuid} execution failed with error:', metadata=metadata
+        )
         add_pipeline_message(trace, execution_state='idle', metadata=metadata)
 
     if pipeline.type == PipelineType.PYTHON:
@@ -186,45 +189,54 @@ class WebSocketServer(tornado.websocket.WebSocketHandler):
         token = message.get('token')
 
         pipeline_uuid = message.get('pipeline_uuid')
+        oauth_client, oauth_token, user = None, None, None
+
+        if api_key and token:
+            oauth_client = Oauth2Application.query.filter(
+                Oauth2Application.client_id == api_key,
+            ).first()
+            if oauth_client:
+                oauth_token, _ = authenticate_client_and_token(oauth_client.id, token)
+                if oauth_token:
+                    user = oauth_token.user
+
+        repo_path = get_repo_path(user=user)
         pipeline = None
         if pipeline_uuid:
             pipeline = Pipeline.get(
                 pipeline_uuid,
-                repo_path=get_repo_path(),
                 all_projects=project_platform_activated(),
+                repo_path=repo_path,
+                use_repo_path=True,
             )
 
         if REQUIRE_USER_AUTHENTICATION or DISABLE_NOTEBOOK_EDIT_ACCESS:
             valid = not REQUIRE_USER_AUTHENTICATION
 
-            if api_key and token:
-                oauth_client = Oauth2Application.query.filter(
-                    Oauth2Application.client_id == api_key,
-                ).first()
-                if oauth_client:
-                    oauth_token, valid = authenticate_client_and_token(oauth_client.id, token)
-                    if valid and oauth_token and oauth_token.user:
-                        try:
-                            if pipeline_uuid:
-                                await PipelineResource.policy_class()(
-                                    PipelineResource(
-                                        pipeline,
-                                        oauth_token.user,
-                                    ),
-                                    oauth_token.user,
-                                ).authorize_action(action=OperationType.UPDATE)
-                            else:
-                                await ProjectResource.policy_class()(
-                                    ProjectResource(
-                                        {},
-                                        oauth_token.user,
-                                    ),
-                                    oauth_token.user,
-                                ).authorize_action(action=OperationType.UPDATE)
-                            valid = True
-                        except ApiError as err:
-                            print(f'[WARNING] WebSocketServer.on_message: {err}.')
-                            valid = False
+            if oauth_client:
+                valid = (oauth_token is not None) and oauth_token.is_valid()
+                if valid and oauth_token and user:
+                    try:
+                        if pipeline_uuid:
+                            await PipelineResource.policy_class()(
+                                PipelineResource(
+                                    pipeline,
+                                    user,
+                                ),
+                                user,
+                            ).authorize_action(action=OperationType.UPDATE)
+                        else:
+                            await ProjectResource.policy_class()(
+                                ProjectResource(
+                                    {},
+                                    user,
+                                ),
+                                user,
+                            ).authorize_action(action=OperationType.UPDATE)
+                        valid = True
+                    except ApiError as err:
+                        print(f'[WARNING] WebSocketServer.on_message: {err}.')
+                        valid = False
 
             if not valid or DISABLE_NOTEBOOK_EDIT_ACCESS == 1:
                 return self.send_message(
@@ -300,14 +312,19 @@ class WebSocketServer(tornado.websocket.WebSocketHandler):
     @classmethod
     def send_message(cls, message: dict) -> None:
         def should_filter_message(message):
-            if message.get('data') is None and message.get('error') is None \
-                    and message.get('execution_state') is None and message.get('type') is None:
+            if (
+                message.get('data') is None
+                and message.get('error') is None
+                and message.get('execution_state') is None
+                and message.get('type') is None
+            ):
                 return True
 
             try:
                 # Filter out messages meant for jupyter widgets that we can't render
-                if message.get('msg_type') == 'display_data' and \
-                        message.get('data')[0].startswith('FloatProgress'):
+                if message.get('msg_type') == 'display_data' and message.get('data')[0].startswith(
+                    'FloatProgress'
+                ):
                     return True
             except IndexError:
                 pass
@@ -327,8 +344,12 @@ class WebSocketServer(tornado.websocket.WebSocketHandler):
         msg_id = message.get('msg_id')
         if msg_id is None:
             return
-        if message.get('data') is None and message.get('error') is None \
-           and message.get('execution_state') is None and message.get('type') is None:
+        if (
+            message.get('data') is None
+            and message.get('error') is None
+            and message.get('execution_state') is None
+            and message.get('type') is None
+        ):
             return
 
         if should_filter_message(message):
@@ -337,8 +358,11 @@ class WebSocketServer(tornado.websocket.WebSocketHandler):
         message = filter_out_sensitive_data(message)
 
         execution_metadata = message.get('execution_metadata')
-        msg_id_value = execution_metadata if execution_metadata is not None \
+        msg_id_value = (
+            execution_metadata
+            if execution_metadata is not None
             else WebSocketServer.running_executions_mapping.get(msg_id, dict())
+        )
         block_type = msg_id_value.get('block_type')
         block_uuid = msg_id_value.get('block_uuid')
         replicated_block = msg_id_value.get('replicated_block')
@@ -410,12 +434,12 @@ class WebSocketServer(tornado.websocket.WebSocketHandler):
         if is_disable_pipeline_edit_access():
             custom_code = block.content
 
-        reload_all_repo_modules(custom_code)
-
         code = custom_code
 
         client = self.init_kernel_client(kernel_name)
 
+        reload_all_repo_modules(custom_code, client)
+
         value = dict(
             block_type=block_type or block.type,
             block_uuid=block_uuid,
@@ -441,26 +465,37 @@ class WebSocketServer(tornado.websocket.WebSocketHandler):
 
                 execution_uuid = None
                 # Need to cache everything here
-                if block.should_track_spark() and \
-                        ComputeServiceUUID.AWS_EMR == block.compute_service_uuid:
-
+                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)
 
+                pipeline_config = pipeline.get_config_from_yaml()
+                repo_config = pipeline.repo_config.to_dict(remote=remote_execution)
+
                 code = add_execution_code(
                     pipeline_uuid,
                     block_uuid,
                     custom_code,
                     global_vars,
+                    pipeline.repo_path,
                     block_type=block_type,
                     execution_uuid=execution_uuid,
                     extension_uuid=extension_uuid,
                     kernel_name=kernel_name,
                     output_messages_to_logs=output_messages_to_logs,
-                    pipeline_config=pipeline.get_config_from_yaml(),
-                    repo_config=pipeline.repo_config.to_dict(remote=remote_execution),
+                    pipeline_config=pipeline_config,
+                    pipeline_config_json_encoded=base64.b64encode(
+                        simplejson.dumps(pipeline_config).encode()
+                    ).decode(),
+                    repo_config=repo_config,
+                    repo_config_json_encoded=base64.b64encode(
+                        simplejson.dumps(repo_config).encode()
+                    ).decode(),
                     # repo_config=get_repo_config().to_dict(remote=remote_execution),
                     run_incomplete_upstream=run_incomplete_upstream,
                     # The UI can execute a block and send run_settings to control the behavior
@@ -486,13 +521,21 @@ db_connection.start_session()
 """
                 client.execute(initialize_db_connection)
 
-            msg_id = client.execute(add_internal_output_info(block, code))
+            msg_id = client.execute(
+                add_internal_output_info(
+                    block,
+                    code,
+                    extension_uuid=extension_uuid,
+                    widget=widget,
+                )
+            )
 
             WebSocketServer.running_executions_mapping[msg_id] = value
 
             block_output_process_code = get_block_output_process_code(
                 pipeline_uuid,
                 block_uuid,
+                pipeline.repo_path,
                 block_type=block_type,
                 kernel_name=kernel_name,
             )
@@ -505,14 +548,18 @@ db_connection.start_session()
                     if BlockType.CHART != downstream_block.type:
                         continue
 
-                    await self.on_message(json.dumps(dict(
-                        api_key=message.get('api_key'),
-                        code=downstream_block.file.content(),
-                        pipeline_uuid=pipeline_uuid,
-                        token=message.get('token'),
-                        type=downstream_block.type,
-                        uuid=downstream_block.uuid,
-                    )))
+                    await self.on_message(
+                        json.dumps(
+                            dict(
+                                api_key=message.get('api_key'),
+                                code=downstream_block.file.content(),
+                                pipeline_uuid=pipeline_uuid,
+                                token=message.get('token'),
+                                type=downstream_block.type,
+                                uuid=downstream_block.uuid,
+                            )
+                        )
+                    )
 
     def __execute_pipeline(
         self,
@@ -525,6 +572,7 @@ db_connection.start_session()
         if kernel_name == KernelName.PYSPARK:
             code = get_pipeline_execution_code(
                 pipeline_uuid,
+                pipeline.repo_path,
                 global_vars=global_vars,
                 kernel_name=kernel_name,
                 pipeline_config=pipeline.to_dict(include_content=True),
@@ -534,9 +582,7 @@ db_connection.start_session()
             client = self.init_kernel_client(kernel_name)
             msg_id = client.execute(code)
 
-            WebSocketServer.running_executions_mapping[msg_id] = dict(
-                pipeline_uuid=pipeline_uuid
-            )
+            WebSocketServer.running_executions_mapping[msg_id] = dict(pipeline_uuid=pipeline_uuid)
         else:
             # TODO: save config for other kernel types.
             def save_pipeline_config() -> str:
@@ -571,8 +617,7 @@ db_connection.start_session()
 
             queue = multiprocessing.Queue()
             proc = multiprocessing.Process(
-                target=run_pipeline,
-                args=(pipeline, config_copy_path, global_vars, queue)
+                target=run_pipeline, args=(pipeline, config_copy_path, global_vars, queue)
             )
             proc.start()
             set_current_pipeline_process(proc)
@@ -590,8 +635,7 @@ db_connection.start_session()
                             metadata=metadata,
                             msg_type=msg.get('msg_type'),
                         )
-                        if execution_state == 'idle' and \
-                                metadata.get('block_uuid') is None:
+                        if execution_state == 'idle' and metadata.get('block_uuid') is None:
                             loop = False
                             break
                     await asyncio.sleep(0.5)
@@ -658,9 +702,9 @@ db_connection.start_session()
 
         try:
             if initial_idx and end_idx:
-                return error[:initial_idx - 1] + error[end_idx:]
+                return error[: initial_idx - 1] + error[end_idx:]
             elif initial_idx and custom_block_end_idx:
-                return error[:initial_idx - 1] + error[custom_block_end_idx:]
+                return error[: initial_idx - 1] + error[custom_block_end_idx:]
         except Exception:
             pass
 
@@ -672,6 +716,7 @@ class StreamBlockOutputToQueue(object):
     Fake file-like stream object that redirects block output to a queue
     to be streamed to the websocket.
     """
+
     def __init__(
         self,
         queue,
@@ -690,9 +735,11 @@ class StreamBlockOutputToQueue(object):
 
     def write(self, buf):
         for line in buf.rstrip().splitlines():
-            self.queue.put(dict(
-                message=f'[{self.block_uuid}] {line.rstrip()}',
-                execution_state=self.execution_state,
-                metadata=self.metadata,
-                msg_type=self.msg_type,
-            ))
+            self.queue.put(
+                dict(
+                    message=f'[{self.block_uuid}] {line.rstrip()}',
+                    execution_state=self.execution_state,
+                    metadata=self.metadata,
+                    msg_type=self.msg_type,
+                )
+            )

mage_ai/services/aws/ecs/ecs.py

@@ -16,6 +16,7 @@ def run_task(
     response = client.run_task(**ecs_config.get_task_config(command=command))
 
     print(json.dumps(response, indent=4, default=str))
+    wait_for_completion = False if ecs_config.wait_timeout == -1 else wait_for_completion
 
     if wait_for_completion:
         arn = response['tasks'][0]['taskArn']

mage_ai/services/k8s/config.py

@@ -2,12 +2,13 @@ from dataclasses import dataclass
 from typing import Dict
 
 from kubernetes.client import (
-    V1Affinity,
     V1Container,
     V1EnvVar,
+    V1JobSpec,
     V1LocalObjectReference,
     V1ObjectMeta,
     V1PodSpec,
+    V1PodTemplateSpec,
     V1ResourceRequirements,
     V1Toleration,
     V1Volume,
@@ -15,8 +16,9 @@ from kubernetes.client import (
 )
 
 from mage_ai.services.k8s.constants import CONFIG_FILE, DEFAULT_NAMESPACE
+from mage_ai.services.k8s.utils import parse_affinity_config
 from mage_ai.shared.config import BaseConfig
-from mage_ai.shared.hash import camel_case_keys_to_snake_case, get_safe_value
+from mage_ai.shared.hash import get_safe_value
 
 # import traceback
 
@@ -43,8 +45,10 @@ class K8sExecutorConfig(BaseConfig):
     metadata: Dict = None
     container: Dict = None
     pod: Dict = None
+    job: Dict = None
     # parsed k8s objects
     pod_config: V1PodSpec = None
+    job_config: V1JobSpec = None
     meta: V1ObjectMeta = None
 
     @classmethod
@@ -61,8 +65,8 @@ class K8sExecutorConfig(BaseConfig):
                 executor_config.pod.get('service_account_name') or DEFAULT_SERVICE_ACCOUNT_NAME
             )
             if executor_config.pod.get('affinity'):
-                affinity = V1Affinity(
-                    **camel_case_keys_to_snake_case(executor_config.pod['affinity']))
+                # Convert the affinity to a V1Affinity object
+                affinity = parse_affinity_config(executor_config.pod['affinity'])
 
             if executor_config.pod.get('tolerations'):
                 tolerations += [V1Toleration(**e) for e in executor_config.pod['tolerations']]
@@ -119,6 +123,25 @@ class K8sExecutorConfig(BaseConfig):
             tolerations=tolerations,
             volumes=volumes,
         )
+
+        active_deadline_seconds = None
+        backoff_limit = 0
+        ttl_seconds_after_finished = None
+
+        if executor_config.job:
+            active_deadline_seconds = executor_config.job.get('active_deadline_seconds')
+            backoff_limit = executor_config.job.get('backoff_limit')
+            ttl_seconds_after_finished = executor_config.job.get('ttl_seconds_after_finished')
+
+        executor_config.job_config = V1JobSpec(
+            active_deadline_seconds=active_deadline_seconds,
+            backoff_limit=backoff_limit,
+            ttl_seconds_after_finished=ttl_seconds_after_finished,
+            template=V1PodTemplateSpec(
+                spec=executor_config.pod_config
+            )
+        )
+
         return executor_config
 
     @classmethod

mage_ai/services/k8s/job_manager.py

@@ -137,7 +137,12 @@ class JobManager():
         )
 
         # Create the specification of deployment
-        spec = client.V1JobSpec(template=template, backoff_limit=0)
+        spec = client.V1JobSpec(
+            template=template,
+            active_deadline_seconds=k8s_config.job_config.active_deadline_seconds,
+            backoff_limit=k8s_config.job_config.backoff_limit,
+            ttl_seconds_after_finished=k8s_config.job_config.ttl_seconds_after_finished)
+
         # Instantiate the job object
         job = client.V1Job(
             api_version=self.api_version,

mage_ai/services/k8s/utils.py

@@ -0,0 +1,97 @@
+from kubernetes import client
+
+
+def parse_affinity_config(affinity_config):
+    if not affinity_config:
+        return None
+
+    affinity = client.V1Affinity(
+        node_affinity=parse_node_affinity(affinity_config.get('nodeAffinity')),
+        pod_affinity=parse_pod_affinity(affinity_config.get('podAffinity')),
+        pod_anti_affinity=parse_pod_anti_affinity(affinity_config.get('podAntiAffinity'))
+    )
+    return affinity
+
+
+def parse_node_affinity(node_affinity_config):
+    if not node_affinity_config:
+        return None
+
+    required_terms = node_affinity_config.get(
+        'requiredDuringSchedulingIgnoredDuringExecution', {}).get('nodeSelectorTerms', [])
+    node_selector_terms = []
+    for term in required_terms:
+        match_expressions = [client.V1NodeSelectorRequirement(
+            key=expr.get('key'),
+            operator=expr.get('operator'),
+            values=expr.get('values')
+        ) for expr in term.get('matchExpressions', [])]
+
+        node_selector_terms.append(client.V1NodeSelectorTerm(match_expressions=match_expressions))
+
+    return client.V1NodeAffinity(
+        required_during_scheduling_ignored_during_execution=client.V1NodeSelector(
+            node_selector_terms=node_selector_terms)
+    )
+
+
+def parse_pod_affinity(pod_affinity_config):
+    if not pod_affinity_config:
+        return None
+
+    required_terms = pod_affinity_config.get('requiredDuringSchedulingIgnoredDuringExecution', [])
+    preferred_terms = pod_affinity_config.get('preferredDuringSchedulingIgnoredDuringExecution', [])
+
+    required_affinity = parse_affinity_term_list(required_terms)
+    preferred_affinity = parse_affinity_term_list(preferred_terms)
+
+    return client.V1PodAffinity(
+        required_during_scheduling_ignored_during_execution=required_affinity,
+        preferred_during_scheduling_ignored_during_execution=preferred_affinity
+    )
+
+
+def parse_pod_anti_affinity(pod_anti_affinity_config):
+    if not pod_anti_affinity_config:
+        return None
+
+    required_terms = pod_anti_affinity_config.get(
+        'requiredDuringSchedulingIgnoredDuringExecution', [])
+    preferred_terms = pod_anti_affinity_config.get(
+        'preferredDuringSchedulingIgnoredDuringExecution', [])
+
+    required_affinity = parse_affinity_term_list(required_terms)
+    preferred_affinity = parse_affinity_term_list(preferred_terms)
+
+    return client.V1PodAntiAffinity(
+        required_during_scheduling_ignored_during_execution=required_affinity,
+        preferred_during_scheduling_ignored_during_execution=preferred_affinity
+    )
+
+
+def parse_affinity_term_list(term_list):
+    if not term_list:
+        return None
+
+    affinity_term_list = []
+    for term in term_list:
+        label_selector = client.V1LabelSelector(
+            match_labels=term.get('labelSelector', {}).get('matchLabels', {}),
+            match_expressions=[client.V1LabelSelectorRequirement(
+                key=expr.get('key'),
+                operator=expr.get('operator'),
+                values=expr.get('values')
+            ) for expr in term.get('labelSelector', {}).get('matchExpressions', [])]
+        )
+
+        topology_key = term.get('topologyKey')
+
+        affinity_term_list.append(client.V1WeightedPodAffinityTerm(
+            weight=term.get('weight'),
+            pod_affinity_term=client.V1PodAffinityTerm(
+                label_selector=label_selector,
+                topology_key=topology_key
+            )
+        ))
+
+    return affinity_term_list

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

@@ -32,6 +32,14 @@ def should_tunnel(
     if not project:
         project = Project()
 
+    if not project.emr_config:
+        error_message = 'Project has incomplete configuration for AWS EMR.'
+
+        if raise_error:
+            raise Exception(error_message)
+        else:
+            return False
+
     if not project.is_feature_enabled(FeatureUUID.COMPUTE_MANAGEMENT):
         error_message = 'Compute management feature isn’t enabled. '\
             'Turn the feature on in the project’s settings.'
@@ -53,14 +61,6 @@ def should_tunnel(
         else:
             return False
 
-    if not project.emr_config:
-        error_message = 'Project has incomplete configuration for AWS EMR.'
-
-        if raise_error:
-            raise Exception(error_message)
-        else:
-            return False
-
     ec2_key_name = project.emr_config.get('ec2_key_name')
     ec2_key_path = project.emr_config.get('ec2_key_path')
 

mage_ai/settings/keys/auth.py

@@ -76,3 +76,4 @@ AZURE_DEVOPS_ORGANIZATION = 'AZURE_DEVOPS_ORGANIZATION'
 OIDC_CLIENT_ID = 'OIDC_CLIENT_ID'
 OIDC_CLIENT_SECRET = 'OIDC_CLIENT_SECRET'
 OIDC_DISCOVERY_URL = 'OIDC_DISCOVERY_URL'
+OIDC_ROLES_MAPPING = 'OIDC_ROLES_MAPPING'