[dy] Add frontend support for managing instances (mage-ai#1165)

* [dy] Add frontend support for managing instances

* [dy] Update manage page

* [dy] Update

* [dy] Update variables

* [dy] Clean up PR

* [dy] Address comments

Author: dy46

.gitignore

@@ -132,6 +132,7 @@ dmypy.json
 # data files
 mage_ai/server/data/files/**
 default_repo
+instance_metadata.json
 
 # test notebook
 test.ipynb

docker-compose.yml

@@ -5,10 +5,13 @@ services:
     build:
       context: .
       dockerfile: ./dev.Dockerfile
-    command: "python mage_ai/server/server.py --host ${HOST} --port ${PORT} --project ${PROJECT}"
+    command: "python mage_ai/server/server.py --host ${HOST} --port ${PORT} --project ${PROJECT} --manage-instance ${MANAGE_INSTANCE}"
     environment:
       - AWS_ACCESS_KEY_ID=$AWS_ACCESS_KEY_ID
       - AWS_SECRET_ACCESS_KEY=$AWS_SECRET_ACCESS_KEY
+      - ECS_CLUSTER_NAME=$ECS_CLUSTER_NAME
+      - ECS_TASK_DEFINITION=$ECS_TASK_DEFINITION
+      - ECS_CONTAINER_NAME=$ECS_CONTAINER_NAME
     ports:
       - 6789:6789
     volumes:

mage_ai/cli/main.py

@@ -57,6 +57,7 @@ def main():
         parser.add_argument('repo_path', metavar='project_path', type=str)
         parser.add_argument('--host', nargs='?', type=str)
         parser.add_argument('--port', nargs='?', type=int)
+        parser.add_argument('--manage-instance', nargs='?', type=str)
 
         args = dict()
         if len(sys.argv) >= 3:
@@ -67,6 +68,7 @@ def main():
 
         start_server(
             host=args.get('host'),
+            manage=args.get('manage_instance') == '1',
             port=args.get('port'),
             project=repo_path,
         )

mage_ai/cluster_manager/aws/ecs_task_manager.py

@@ -1,54 +1,87 @@
 from botocore.config import Config
+from functools import reduce
 from mage_ai.services.aws.ecs.config import EcsConfig
-from mage_ai.services.aws.ecs.ecs import list_tasks, run_task
+from mage_ai.services.aws.ecs.ecs import list_tasks, run_task, stop_task
 from mage_ai.shared.array import find
 from mage_ai.shared.hash import dig
-from typing import List
+from typing import Dict, List
 
 import boto3
+import json
 import os
 
-
-CLUSTER_NAME = 'mage-data-prep-development-cluster'
-
 class EcsTaskManager:
-    def __init__(self, cluster_name=CLUSTER_NAME):
+    def __init__(self, cluster_name):
         self.cluster_name = cluster_name
 
+        self.metadata_file = os.path.join(
+            os.getcwd(),
+            'instance_metadata.json',
+        )
+
+        if not os.path.exists(self.metadata_file):
+            self.instance_metadata = {}
+
+    @property
+    def instance_metadata(self):
+        metadata = {}
+        with open(self.metadata_file, 'r', encoding='utf-8') as file:
+            metadata = json.load(file)
+        return metadata
+
+    @instance_metadata.setter
+    def instance_metadata(self, metadata):
+        with open(self.metadata_file, 'w', encoding='utf-8') as file:
+            json.dump(metadata, file)
+
     def list_tasks(self):
         region_name = os.getenv('AWS_REGION_NAME', 'us-west-2')
         config = Config(region_name=region_name)
         ec2_client = boto3.client('ec2', config=config)
-        response = list_tasks(self.cluster_name)['tasks']
 
+        response = list_tasks(self.cluster_name)['tasks']
         network_interfaces = self.__get_network_interfaces(response, ec2_client)
 
         tasks = []
-
-        for index, task in enumerate(response):
-            public_ip = dig(network_interfaces[index], 'Association.PublicIp')
+        for task in response:
+            public_ip = dig(network_interfaces.get(task['taskArn']), 'Association.PublicIp')
 
             tags = task['tags']
             name = find(lambda tag: tag.get('key') == 'name', tags)
 
             tasks.append(dict(
                 ip=public_ip,
-                group=task['group'],
                 name=name.get('value') if name is not None else None,
                 status=task['lastStatus'],
+                task_arn=task['taskArn'],
                 type=task['launchType'],
             ))
 
-        return tasks
+        running_instance_names = set(map(lambda x: x['name'], tasks))
+
+        stopped_instance_names = \
+            [name for name in list(self.instance_metadata.keys()) if name not in running_instance_names]
+        stopped_instances = \
+            list(
+                map(
+                    lambda name: { 'name': name, 'status': 'STOPPED' },
+                    stopped_instance_names
+                )
+            )
+
+        return tasks + stopped_instances
 
     def create_task(self, name: str, task_definition: str, container_name: str):
         region_name = os.getenv('AWS_REGION_NAME', 'us-west-2')
         config = Config(region_name=region_name)
         ec2_client = boto3.client('ec2', config=config)
 
         # create new task
-        task = list_tasks(self.cluster_name)['tasks'][0]
-        network_interface = self.__get_network_interfaces([task], ec2_client)[0]
+        task = find(
+            lambda task: task.get('lastStatus') == 'RUNNING',
+            list_tasks(self.cluster_name)['tasks'],
+        )
+        network_interface = self.__get_network_interfaces([task], ec2_client)[task['taskArn']]
 
         subnets = [network_interface['SubnetId']]
         security_groups = [g['GroupId'] for g in network_interface['Groups']]
@@ -67,19 +100,56 @@ def create_task(self, name: str, task_definition: str, container_name: str):
             ],
         )
 
+        self.instance_metadata = {
+            **self.instance_metadata,
+            name: dict()
+        }
+
         return run_task(f'mage start {name}', ecs_config=ecs_config)
 
-    def __get_network_interface_id(self, task: str):
+    def stop_task(self, task_arn: str):
+        return stop_task(task_arn, self.cluster_name)
+
+    def delete_task(self, name, task_arn: str = None):
+        if task_arn:
+            self.stop_task(task_arn)
+
+        updated_metadata = self.instance_metadata
+
+        if name in updated_metadata:
+            del updated_metadata[name]
+            self.instance_metadata = updated_metadata
+
+    def __get_network_interface_id(self, task):
+        if task.get('lastStatus') != 'RUNNING':
+            return None
+
         attachment = \
             find(lambda a: a['type'] == 'ElasticNetworkInterface', task.get('attachments', []))
         network_interface = \
             find(lambda d: d['name'] == 'networkInterfaceId', attachment.get('details', []))
         return network_interface.get('value', None)
 
+    def __get_network_interfaces(self, tasks: List, ec2_client) -> Dict:
+        task_mapping = dict()
+        for task in tasks:
+            nii = self.__get_network_interface_id(task)
+            if nii is not None:
+                task_mapping[task['taskArn']] = nii
 
-    def __get_network_interfaces(self, tasks: List, ec2_client):
-        network_interface_ids = [self.__get_network_interface_id(task) for task in tasks]
+        network_interface_ids = list(task_mapping.values())
 
-        return ec2_client.describe_network_interfaces(
+        network_interfaces = ec2_client.describe_network_interfaces(
             NetworkInterfaceIds=network_interface_ids
         )['NetworkInterfaces']
+
+        def aggregate(obj, task):
+            task_arn = task['taskArn']
+            if task_arn in task_mapping:
+                obj[task_arn] = find(
+                    lambda i: i['NetworkInterfaceId'] == task_mapping[task_arn],
+                    network_interfaces,
+                )
+            return obj
+
+        return reduce(aggregate, tasks, {})

mage_ai/data_preparation/shared/constants.py

@@ -1 +1,2 @@
-REPO_PATH_ENV_VAR = 'MAGE_REPO_PATH'
+REPO_PATH_ENV_VAR = 'MAGE_REPO_PATH'
+MANAGE_ENV_VAR = 'MAGE_IS_MANAGE_INSTANCE'

mage_ai/frontend/api/index.ts

@@ -35,6 +35,7 @@ export const FILE_CONTENTS: 'file_contents' = 'file_contents';
 export const FILES: 'files' = 'files';
 export const FEATURE_SETS: 'feature_sets' = 'feature_sets';
 export const FEATURE_SET_VERSIONS: 'feature_set_versions' = 'feature_set_versions';
+export const INSTANCES: 'instances' = 'instances';
 export const KERNELS: 'kernels' = 'kernels';
 export const KERNEL_ACTION_INTERRUPT: 'interrupt' = 'interrupt';
 export const KERNEL_ACTION_RESTART: 'restart' = 'restart';
@@ -69,6 +70,7 @@ const RESOURCES: any[][] = [
   [FEATURE_SETS],
   [FILES],
   [FILE_CONTENTS],
+  [INSTANCES, CLUSTERS],
   [KERNELS],
   [KERNEL_ACTION_INTERRUPT, KERNELS],
   [KERNEL_ACTION_RESTART, KERNELS],
@@ -155,8 +157,8 @@ RESOURCES.forEach(([resource, parentResource, grandchildResource, swrOptions]) =
     apis[resource][parentResource].useUpdate = (parentId: string, id: string, opts?: any) => async (body: any) =>
       fetchUpdateWithParent(resource, parentResource, parentId, id, body, opts);
 
-    apis[resource][parentResource].useDelete = (parentId: string, id: string) => async () => {
-      const response = await useDeleteWithParent(resource, parentResource, parentId, id);
+    apis[resource][parentResource].useDelete = (parentId: string, id: string, query?: object) => async () => {
+      const response = await useDeleteWithParent(resource, parentResource, parentId, id, query);
 
       return await handle(response);
     },
@@ -208,8 +210,8 @@ RESOURCES.forEach(([resource, parentResource, grandchildResource, swrOptions]) =
     apis[resource].useCreate = (opts?: any) =>
       async (body: any) => fetchCreate(resource, body, opts);
 
-    apis[resource].useDelete = (id: string) => async () => {
-      const response = await useDelete(resource, id);
+    apis[resource].useDelete = (id: string, query?: object) => async () => {
+      const response = await useDelete(resource, id, query);
 
       return await handle(response);
     },

mage_ai/frontend/api/utils/use.ts

@@ -115,16 +115,17 @@ export function useDetailWithParent(
   };
 }
 
-export function useDelete(resource: string, id: string) {
-  return buildFetch(buildUrl(resource, id), { method: DELETE });
+export function useDelete(resource: string, id: string, query: object = {}) {
+  return buildFetch(buildUrl(resource, id), { query, method: DELETE });
 }
 
 export function useDeleteWithParent(resource: string,
   parentResource: string,
   parentId: string,
   id: string,
+  query: object = {},
 ) {
-  return buildFetch(buildUrl(parentResource, parentId, resource, id), { method: DELETE });
+  return buildFetch(buildUrl(parentResource, parentId, resource, id), { query, method: DELETE });
 }
 
 export function useList(

mage_ai/frontend/components/Dashboard/index.tsx

@@ -127,9 +127,11 @@ function Dashboard({
       />
 
       <ContainerStyle>
-        <VerticalNavigationStyle>
-          <VerticalNavigation navigationItems={navigationItems} />
-        </VerticalNavigationStyle>
+        {navigationItems?.length !== 0 && (
+          <VerticalNavigationStyle>
+            <VerticalNavigation navigationItems={navigationItems} />
+          </VerticalNavigationStyle>
+        )}
 
         <Flex
           flex={1}

mage_ai/frontend/components/PipelineDetail/Runs/Table.style.tsx

@@ -30,4 +30,15 @@ export const PopupContainerStyle = styled.div<any>`
     background-color: ${(props.theme.interactive || dark.interactive).defaultBackground};
   `}
 
+  ${props => props.leftOffset && `
+    left: ${props.leftOffset}px;
+  `}
+
+  ${props => props.topOffset && `
+    top: ${props.topOffset}px;
+  `}
+
+  ${props => props.width && `
+    width: ${props.width}px;
+  `}
 `;

mage_ai/frontend/pages/index.tsx

@@ -1,19 +1,26 @@
 import { useEffect } from 'react';
 import { useRouter } from 'next/router';
 
+import api from '@api';
+
 const Home = () => {
   const router = useRouter();
   const completePath = router.asPath;
   const basePath = completePath.split('?')[0];
 
+  const { data: dataStatus } = api.status.list();
+  const manage = dataStatus?.status?.['is_instance_manager'];
+
   let pathname = completePath;
   if (basePath === '/') {
-    pathname = '/pipelines';
+    pathname = manage ? '/manage' : '/pipelines';
   }
 
   useEffect(() => {
-    router.replace(pathname);
-  }, []);
+    if (dataStatus) {
+      router.replace(pathname);
+    }
+  }, [dataStatus]);
 };
 
 export default Home;