mage-ai 0.8.97__py3-none-any.whl → 0.8.99__py3-none-any.whl

mage_ai/cli/main.py

@@ -9,6 +9,8 @@ from rich import print
 from typer.core import TyperGroup
 
 from mage_ai.cli.utils import parse_runtime_variables
+from mage_ai.data_preparation.repo_manager import ProjectType
+from mage_ai.services.newrelic import initialize_new_relic
 from mage_ai.shared.constants import InstanceType
 
 
@@ -26,7 +28,17 @@ app = typer.Typer(
 
 INIT_PROJECT_PATH_DEFAULT = typer.Argument(..., help='path of the Mage project to be created.')
 INIT_PROJECT_TYPE_DEFAULT = typer.Option(
-    'standalone', help='type of project to create, options are main, sub, or standalone')
+    ProjectType.STANDALONE.value,
+    help='type of project to create, options are main, sub, or standalone'
+)
+INIT_CLUSTER_TYPE_DEFAULT = typer.Option(
+    None,
+    help='type of instance to create for workspace management',
+)
+INIT_PROJECT_UUID_DEFAULT = typer.Option(
+    None,
+    help='project uuid for the new project',
+)
 
 START_PROJECT_PATH_DEFAULT = typer.Argument(
     os.getcwd(), help='path of the Mage project to be loaded.')
@@ -36,6 +48,18 @@ START_MANAGE_INSTANCE_DEFAULT = typer.Option('0', help='')
 START_DBT_DOCS_INSTANCE_DEFAULT = typer.Option('0', help='')
 START_INSTANCE_TYPE_DEFAULT = typer.Option(
     InstanceType.SERVER_AND_SCHEDULER.value, help='specify the instance type.')
+START_PROJECT_TYPE_DEFAULT = typer.Option(
+    ProjectType.STANDALONE.value,
+    help='create project of this type if does not exist, options are main, sub, or standalone',
+)
+START_CLUSTER_TYPE_DEFAULT = typer.Option(
+    None,
+    help='type of instance to create for workspace management',
+)
+START_PROJECT_UUID_DEFAULT = typer.Option(
+    None,
+    help='set project uuid if it has not been set for the project already',
+)
 
 RUN_PROJECT_PATH_DEFAULT = typer.Argument(
     ..., help='path of the Mage project that contains the pipeline.'
@@ -90,6 +114,8 @@ CREATE_SPARK_CLUSTER_PROJECT_PATH_DEFAULT = typer.Argument(
 def init(
     project_path: str = INIT_PROJECT_PATH_DEFAULT,
     project_type: Union[str, None] = INIT_PROJECT_TYPE_DEFAULT,
+    cluster_type: str = INIT_CLUSTER_TYPE_DEFAULT,
+    project_uuid: str = INIT_PROJECT_UUID_DEFAULT,
 ):
     """
     Initialize Mage project.
@@ -97,7 +123,12 @@ def init(
     from mage_ai.data_preparation.repo_manager import init_repo
 
     repo_path = os.path.join(os.getcwd(), project_path)
-    init_repo(repo_path, project_type=project_type)
+    init_repo(
+        repo_path,
+        project_type=project_type,
+        cluster_type=cluster_type,
+        project_uuid=project_uuid,
+    )
     print(f'Initialized Mage project at {repo_path}')
 
 
@@ -109,6 +140,9 @@ def start(
     manage_instance: str = START_MANAGE_INSTANCE_DEFAULT,
     dbt_docs_instance: str = START_DBT_DOCS_INSTANCE_DEFAULT,
     instance_type: str = START_INSTANCE_TYPE_DEFAULT,
+    project_type: str = START_PROJECT_TYPE_DEFAULT,
+    cluster_type: str = START_CLUSTER_TYPE_DEFAULT,
+    project_uuid: str = START_PROJECT_UUID_DEFAULT,
 ):
     """
     Start Mage server and UI.
@@ -128,6 +162,9 @@ def start(
         manage=manage_instance == "1",
         dbt_docs=dbt_docs_instance == "1",
         instance_type=instance_type,
+        project_type=project_type,
+        cluster_type=cluster_type,
+        project_uuid=project_uuid,
     )
 
 
@@ -155,6 +192,9 @@ def run(
     project_path = os.path.abspath(project_path)
     set_repo_path(project_path)
 
+    from contextlib import nullcontext
+
+    import newrelic.agent
     import sentry_sdk
 
     from mage_ai.data_preparation.executors.executor_factory import ExecutorFactory
@@ -171,55 +211,58 @@ def run(
             sentry_dsn,
             traces_sample_rate=SENTRY_TRACES_SAMPLE_RATE,
         )
-
-    runtime_variables = dict()
-    if runtime_vars is not None:
-        runtime_variables = parse_runtime_variables(runtime_vars)
-
-    sys.path.append(os.path.dirname(project_path))
-    pipeline = Pipeline.get(pipeline_uuid, repo_path=project_path)
-
-    db_connection.start_session()
-
-    if pipeline_run_id is None:
-        default_variables = get_global_variables(pipeline_uuid)
-        global_vars = merge_dict(default_variables, runtime_variables)
-    else:
-        pipeline_run = PipelineRun.query.get(pipeline_run_id)
-        global_vars = pipeline_run.get_variables(extra_variables=runtime_variables)
-
-    if template_runtime_configuration is not None:
-        template_runtime_configuration = json.loads(template_runtime_configuration)
-
-    if block_uuid is None:
-        ExecutorFactory.get_pipeline_executor(
-            pipeline,
-            execution_partition=execution_partition,
-            executor_type=executor_type,
-        ).execute(
-            analyze_outputs=False,
-            global_vars=global_vars,
-            pipeline_run_id=pipeline_run_id,
-            run_sensors=not skip_sensors,
-            run_tests=test,
-            update_status=False,
-        )
-    else:
-        ExecutorFactory.get_block_executor(
-            pipeline,
-            block_uuid,
-            execution_partition=execution_partition,
-            executor_type=executor_type,
-        ).execute(
-            analyze_outputs=False,
-            block_run_id=block_run_id,
-            callback_url=callback_url,
-            global_vars=global_vars,
-            pipeline_run_id=pipeline_run_id,
-            template_runtime_configuration=template_runtime_configuration,
-            update_status=False,
-        )
-    print('Pipeline run completed.')
+    (enable_new_relic, application) = initialize_new_relic()
+
+    with newrelic.agent.BackgroundTask(application, name="mage-run", group='Task') \
+         if enable_new_relic else nullcontext():
+        runtime_variables = dict()
+        if runtime_vars is not None:
+            runtime_variables = parse_runtime_variables(runtime_vars)
+
+        sys.path.append(os.path.dirname(project_path))
+        pipeline = Pipeline.get(pipeline_uuid, repo_path=project_path)
+
+        db_connection.start_session()
+
+        if pipeline_run_id is None:
+            default_variables = get_global_variables(pipeline_uuid)
+            global_vars = merge_dict(default_variables, runtime_variables)
+        else:
+            pipeline_run = PipelineRun.query.get(pipeline_run_id)
+            global_vars = pipeline_run.get_variables(extra_variables=runtime_variables)
+
+        if template_runtime_configuration is not None:
+            template_runtime_configuration = json.loads(template_runtime_configuration)
+
+        if block_uuid is None:
+            ExecutorFactory.get_pipeline_executor(
+                pipeline,
+                execution_partition=execution_partition,
+                executor_type=executor_type,
+            ).execute(
+                analyze_outputs=False,
+                global_vars=global_vars,
+                pipeline_run_id=pipeline_run_id,
+                run_sensors=not skip_sensors,
+                run_tests=test,
+                update_status=False,
+            )
+        else:
+            ExecutorFactory.get_block_executor(
+                pipeline,
+                block_uuid,
+                execution_partition=execution_partition,
+                executor_type=executor_type,
+            ).execute(
+                analyze_outputs=False,
+                block_run_id=block_run_id,
+                callback_url=callback_url,
+                global_vars=global_vars,
+                pipeline_run_id=pipeline_run_id,
+                template_runtime_configuration=template_runtime_configuration,
+                update_status=False,
+            )
+        print('Pipeline run completed.')
 
 
 @app.command()

mage_ai/cluster_manager/kubernetes/workload_manager.py

@@ -1,6 +1,7 @@
 import os
 from typing import Dict, List
 
+import yaml
 from kubernetes import client, config
 
 from mage_ai.cluster_manager.constants import (
@@ -8,18 +9,24 @@ from mage_ai.cluster_manager.constants import (
     CONNECTION_URL_SECRETS_NAME,
     DB_SECRETS_NAME,
     GCP_BACKEND_CONFIG_ANNOTATION,
+    KUBE_NAMESPACE,
+    KUBE_SERVICE_ACCOUNT_NAME,
     KUBE_SERVICE_GCP_BACKEND_CONFIG,
     KUBE_SERVICE_TYPE,
+    KUBE_STORAGE_CLASS_NAME,
     NODE_PORT_SERVICE_TYPE,
     SERVICE_ACCOUNT_CREDENTIAL_FILE_PATH,
     SERVICE_ACCOUNT_SECRETS_NAME,
 )
+from mage_ai.data_preparation.repo_manager import ProjectType
 from mage_ai.orchestration.constants import (
     DATABASE_CONNECTION_URL_ENV_VAR,
     DB_NAME,
     DB_PASS,
     DB_USER,
 )
+from mage_ai.settings import MAGE_SETTINGS_ENVIRONMENT_VARIABLES
+from mage_ai.shared.array import find
 
 
 class WorkloadManager:
@@ -47,46 +54,93 @@ class WorkloadManager:
 
         return False
 
-    def list_services(self):
+    def list_workloads(self):
         services = self.core_client.list_namespaced_service(self.namespace).items
-        services_list = []
+        workloads_list = []
+
+        pods = self.core_client.list_namespaced_pod(self.namespace).items
+        pod_map = dict()
+        for pod in pods:
+            try:
+                name = pod.metadata.labels.get('app')
+                pod_map[name] = pod
+            except Exception:
+                pass
         for service in services:
             try:
                 labels = service.metadata.labels
                 if not labels.get('dev-instance'):
                     continue
-                ip_address = service.status.load_balancer.ingress[0].ip
-                conditions = service.status.conditions or list()
-                services_list.append(dict(
-                    ip=ip_address,
+                name = labels.get('app')
+                pod = pod_map[name]
+                service_type = service.spec.type
+                workload = dict(
                     name=labels.get('app'),
-                    status='RUNNING' if len(conditions) == 0 else conditions[0].status,
-                    type='kubernetes',
-                ))
+                    type=service_type,
+                )
+                if pod:
+                    status = pod.status.phase
+                    workload['status'] = status.upper()
+
+                    node_name = pod.spec.node_name
+                    ip = None
+                    if service_type == 'NodePort':
+                        try:
+                            if node_name:
+                                items = self.core_client.list_node(
+                                    field_selector=f'metadata.name={node_name}').items
+                                node = items[0]
+                                ip = find(
+                                    lambda a: a.type == 'ExternalIP',
+                                    node.status.addresses
+                                ).address
+                                if ip:
+                                    node_port = service.spec.ports[0].node_port
+                                    workload['ip'] = f'{ip}:{node_port}'
+                        except Exception:
+                            pass
+
+                workloads_list.append(workload)
             except Exception:
                 pass
 
-        return services_list
+        return workloads_list
 
-    def create_deployment(
+    def create_workload(
         self,
-        deployment_name,
-        container_config: Dict = None,
-        service_account_name: str = None,
-        storage_class_name: str = None,
-        volume_host_path: str = None,
+        name: str,
+        project_type: str = ProjectType.STANDALONE,
+        project_uuid: str = None,
+        **kwargs,
     ):
-        if container_config is None:
-            container_config = dict()
+        container_config_yaml = kwargs.get('container_config')
+        container_config = dict()
+        if container_config_yaml:
+            container_config = yaml.full_load(container_config_yaml)
 
-        env_vars = self.__populate_env_vars(container_config)
+        service_account_name = kwargs.get(
+            'service_account_name',
+            os.getenv(KUBE_SERVICE_ACCOUNT_NAME),
+        )
+        storage_class_name = kwargs.get(
+            'storage_class_name',
+            os.getenv(KUBE_STORAGE_CLASS_NAME, 'default'),
+        )
+        storage_access_mode = kwargs.get('storage_access_mode', 'ReadWriteOnce')
+        storage_request_size = kwargs.get('storage_request_size', 2)
+
+        env_vars = self.__populate_env_vars(
+            name,
+            project_type=project_type,
+            project_uuid=project_uuid,
+            container_config=container_config,
+        )
         container_config['env'] = env_vars
 
         containers = [
             {
-                'name': f'{deployment_name}-container',
+                'name': f'{name}-container',
                 'image': 'mageai/mageai:latest',
-                'command': ['mage', 'start', deployment_name],
                 'ports': [
                     {
                         'containerPort': 6789,
@@ -146,32 +200,23 @@ class WorkloadManager:
                 }
             )
 
-        if volume_host_path:
-            volumes.append(
-                {
-                    'name': 'mage-data',
-                    'hostPath': {
-                        'path': volume_host_path
-                    },
-                }
-            )
-        deployment_template_spec = dict()
+        stateful_set_template_spec = dict()
         if service_account_name:
-            deployment_template_spec['serviceAccountName'] = service_account_name
+            stateful_set_template_spec['serviceAccountName'] = service_account_name
 
-        deployment = {
+        stateful_set = {
             'apiVersion': 'apps/v1',
-            'kind': 'Deployment',
+            'kind': 'StatefulSet',
             'metadata': {
-                'name': deployment_name,
+                'name': name,
                 'labels': {
-                    'app': deployment_name
+                    'app': name
                 }
             },
             'spec': {
                 'selector': {
                     'matchLabels': {
-                        'app': deployment_name
+                        'app': name
                     }
                 },
                 'replicas': 1,
@@ -179,22 +224,38 @@ class WorkloadManager:
                 'template': {
                     'metadata': {
                         'labels': {
-                            'app': deployment_name
+                            'app': name
                         }
                     },
                     'spec': {
                         'terminationGracePeriodSeconds': 10,
                         'containers': containers,
                         'volumes': volumes,
-                        **deployment_template_spec
+                        **stateful_set_template_spec
                     }
                 },
+                'volumeClaimTemplates': [
+                    {
+                        'metadata': {
+                            'name': 'mage-data'
+                        },
+                        'spec': {
+                            'accessModes': [storage_access_mode],
+                            'storageClassName': storage_class_name,
+                            'resources': {
+                                'requests': {
+                                    'storage': f'{storage_request_size}Gi'
+                                }
+                            }
+                        }
+                    }
+                ]
             }
         }
 
-        self.apps_client.create_namespaced_deployment(self.namespace, deployment)
+        self.apps_client.create_namespaced_stateful_set(self.namespace, stateful_set)
 
-        service_name = f'{deployment_name}-service'
+        service_name = f'{name}-service'
 
         annotations = {}
         if os.getenv(KUBE_SERVICE_GCP_BACKEND_CONFIG):
@@ -207,7 +268,7 @@ class WorkloadManager:
             'metadata': {
                 'name': service_name,
                 'labels': {
-                    'app': deployment_name,
+                    'app': name,
                     'dev-instance': '1',
                 },
                 'annotations': annotations
@@ -220,7 +281,7 @@ class WorkloadManager:
                     }
                 ],
                 'selector': {
-                    'app': deployment_name
+                    'app': name
                 },
                 'type': os.getenv(KUBE_SERVICE_TYPE, NODE_PORT_SERVICE_TYPE)
             }
@@ -228,8 +289,33 @@ class WorkloadManager:
 
         return self.core_client.create_namespaced_service(self.namespace, service)
 
-    def __populate_env_vars(self, container_config) -> List:
-        env_vars = []
+    def delete_workload(self, name: str):
+        self.apps_client.delete_namespaced_stateful_set(name, self.namespace)
+        self.core_client.delete_namespaced_service(f'{name}-service', self.namespace)
+
+    def __populate_env_vars(
+        self,
+        name,
+        project_type: str = 'standalone',
+        project_uuid: str = None,
+        container_config: Dict = None
+    ) -> List:
+        env_vars = [
+            {
+                'name': 'USER_CODE_PATH',
+                'value': name,
+            }
+        ]
+        if project_type:
+            env_vars.append({
+                'name': 'PROJECT_TYPE',
+                'value': project_type,
+            })
+        if project_uuid:
+            env_vars.append({
+                'name': 'PROJECT_UUID',
+                'value': project_uuid,
+            })
 
         connection_url_secrets_name = os.getenv(CONNECTION_URL_SECRETS_NAME)
         if connection_url_secrets_name:
@@ -245,6 +331,16 @@ class WorkloadManager:
                 }
             )
 
+        for var in MAGE_SETTINGS_ENVIRONMENT_VARIABLES + [
+            DATABASE_CONNECTION_URL_ENV_VAR,
+            KUBE_NAMESPACE,
+        ]:
+            if os.getenv(var) is not None:
+                env_vars.append({
+                    'name': var,
+                    'value': str(os.getenv(var)),
+                })
+
         # For connecting to CloudSQL PostgreSQL database.
         db_secrets_name = os.getenv(DB_SECRETS_NAME)
         if db_secrets_name:

mage_ai/data_preparation/git/__init__.py

@@ -102,7 +102,7 @@ class Git:
         self.repo.create_remote(name, url)
 
     def remove_remote(self, name: str) -> None:
-        self.repo.remotes[name].remove(self.repo, name)
+        self.repo.git.remote('remove', name)
 
     def staged_files(self) -> List[str]:
         files_string = self.repo.git.diff('--name-only', '--cached')
@@ -308,31 +308,101 @@ class Git:
         if message:
             self.repo.index.commit(message)
 
-    def remotes(self, limit: int = 40) -> List[Dict]:
+    def remotes(self, limit: int = 40, user: User = None) -> List[Dict]:
         arr = []
 
         for idx, remote in enumerate(self.repo.remotes):
+            from git.exc import GitCommandError
+
             if idx >= limit:
                 break
 
-            refs = []
-            for ref in remote.refs:
-                refs.append(dict(
-                    name=ref.name,
-                    commit=dict(
-                        author=dict(
-                            email=ref.commit.author.email,
-                            name=ref.commit.author.name,
+            error = None
+            remote_url = None
+            remote_exists = False
+            try:
+                remote_url = [url for url in remote.urls][0]
+                remote_exists = True
+            except GitCommandError as err:
+                print('WARNING (mage_ai.data_preparation.git.remotes):')
+                print(err)
+
+            try:
+                remote_refs = remote.refs
+                if len(remote_refs) == 0 and user:
+                    from mage_ai.data_preparation.git import api
+
+                    access_token = api.get_access_token_for_user(user)
+                    if access_token:
+
+                        if remote_exists:
+                            token = access_token.token
+                            username = api.get_username(token)
+                            url = api.build_authenticated_remote_url(
+                                remote_url,
+                                username,
+                                token,
+                            )
+
+                            remote.set_url(url)
+
+                            authenticated = False
+                            try:
+                                api.check_connection(self.repo, url)
+                                authenticated = True
+                            except Exception as err:
+                                print('WARNING (mage_ai.data_preparation.git.remotes):')
+                                print(err)
+
+                            if authenticated:
+                                try:
+                                    remote.fetch()
+                                    remote_refs = remote.refs
+                                except Exception as err:
+                                    print('WARNING (mage_ai.data_preparation.git.remotes):')
+                                    print(err)
+
+                refs = []
+                for ref in remote_refs:
+                    refs.append(dict(
+                        name=ref.name,
+                        commit=dict(
+                            author=dict(
+                                email=ref.commit.author.email,
+                                name=ref.commit.author.name,
+                            ),
+                            date=datetime.fromtimestamp(ref.commit.authored_date).isoformat(),
+                            message=ref.commit.message,
                         ),
-                        date=datetime.fromtimestamp(ref.commit.authored_date).isoformat(),
-                        message=ref.commit.message,
-                    ),
-                ))
+                    ))
+            except Exception as err:
+                error = err
+
+            try:
+                remote.set_url(remote_url)
+            except GitCommandError as err:
+                print('WARNING (mage_ai.data_preparation.git.remotes):')
+                print(err)
+
+            if error:
+                raise error
+
+            repository_names = []
+            urls = []
+            try:
+                for url in remote.urls:
+                    if url.lower().startswith('https'):
+                        repository_names.append('/'.join(url.split('/')[-2:]).replace('.git', ''))
+                    urls.append(url)
+            except GitCommandError as err:
+                print('WARNING (mage_ai.data_preparation.git.remotes):')
+                print(err)
 
             arr.append(dict(
                 name=remote.name,
                 refs=refs,
-                urls=[url for url in remote.urls],
+                repository_names=repository_names,
+                urls=urls,
             ))
 
         return arr
@@ -416,7 +486,7 @@ class Git:
                     self._run_command(cmd)
                 print(f'Installing {requirements_file} completed successfully.')
             except Exception as err:
-                print(f'Skip installing {requirements_file} due to error: {err}')
+                print(f'Skip installing {requirements_file} due to error: {str(err)}')
                 pass
 
     def __create_ssh_keys(self) -> str: