[feat] Ability to change UI active package after aim ui has been star…

…ted (#2935)

src/aimcore/cli/package/commands.py

@@ -88,3 +88,12 @@ def sync_package(name, repo):
         package_watcher = PackageSourceWatcher(repo_inst, name, src_path)
         package_watcher.initialize()
         package_watcher.start()
+
+
+@package.command('set-active')
+@click.option('--name', '-n', required=True, type=str)
+@click.option('--repo', default='', type=str)
+def set_active_package(name, repo):
+    repo_inst = Repo.from_path(repo) if repo else Repo.default()
+    click.echo(f'Setting \'{name}\' as active package for Repo \'{repo_inst}\'')
+    repo_inst.set_active_package(pkg_name=name)

src/aimcore/cli/package/watcher.py

@@ -51,26 +51,32 @@ async def watch_events(self):
                 fs_events.append(self.queue.get())
             with self.repo.storage_engine.write_batch(0):
                 for fs_event in fs_events:
+                    file = pathlib.Path(fs_event.src_path)
+                    file_path = str(file.relative_to(self.src_dir))
+                    if self.is_black_listed(file_path):
+                        continue
                     try:
-                        file = pathlib.Path(fs_event.src_path)
-                        file_path = str(file.relative_to(self.src_dir))
-                        click.echo(f'Detected change in file \'{file_path}\'. Syncing.')
+                        click.echo(f'Detected change in: \'{file_path}\', type: \'{fs_event.event_type}\'. Syncing...')
 
                         if fs_event.event_type in (events.EVENT_TYPE_CREATED, events.EVENT_TYPE_MODIFIED):
                             with file.open('r') as fh:
                                 contents = fh.read()
                                 self.package.sync(str(file_path), contents)
+                                click.echo(f'Updated file {file_path}.')
                         elif fs_event.event_type == events.EVENT_TYPE_DELETED:
                             self.package.remove(str(file_path))
-                        elif fs_event == events.EVENT_TYPE_MOVED:
-                            dest_path = pathlib.Path(fs_event.dest_path).relative_to(self.src_dir)
+                            click.echo(f'Removed file {file_path}.')
+                        elif fs_event.event_type == events.EVENT_TYPE_MOVED:
+                            dest_path = str(pathlib.Path(fs_event.dest_path).relative_to(self.src_dir))
                             self.package.move(file_path, dest_path)
+                            click.echo(f'Moved file {file_path} to {dest_path}.')
                     except Exception:
                         pass
 
             await asyncio.sleep(5)
 
     def start(self):
+        click.echo(f'Watching for change in package \'{self.package_name}\' sources...')
         self.observer = Observer()
         event_hanlder = SourceFileChangeHandler(self.src_dir, self.queue)
         self.observer.schedule(event_hanlder, self.src_dir, recursive=True)
@@ -88,7 +94,15 @@ def initialize(self):
             click.echo(f'Initializing package \'{self.package_name}\'.')
             for file_path in self.src_dir.glob('**/*'):
                 file_name = file_path.relative_to(self.src_dir)
-                if file_path.is_file():
+                if file_path.is_file() and not self.is_black_listed(str(file_name)):
                     with file_path.open('r') as fh:
                         self.package.sync(str(file_name), fh.read())
             self.package.install()
+            click.echo(f'Package \'{self.package_name}\' initialized.')
+
+    def is_black_listed(self, file_path: str) -> bool:
+        if '__pycache__' in file_path:
+            return True
+        if file_path.endswith('.pyc'):
+            return True
+        return False

src/aimcore/cli/server/commands.py

@@ -1,7 +1,7 @@
 import os
 import click
 
-from aimcore.cli.utils import set_log_level
+from aimcore.cli.utils import set_log_level, start_uvicorn_app
 from aim._sdk.repo import Repo
 from aim._sdk.package_utils import Package
 from aimcore.transport.config import (
@@ -21,7 +21,6 @@
                                                                              file_okay=False,
                                                                              dir_okay=True,
                                                                              writable=True))
-@click.option('--package', '--pkg', required=False, default='asp', type=str)
 @click.option('--ssl-keyfile', required=False, type=click.Path(exists=True,
                                                                file_okay=True,
                                                                dir_okay=False,
@@ -35,7 +34,7 @@
 @click.option('--dev', is_flag=True, default=False)
 @click.option('-y', '--yes', is_flag=True, help='Automatically confirm prompt')
 def server(host, port,
-           repo, package, ssl_keyfile, ssl_certfile,
+           repo, ssl_keyfile, ssl_certfile,
            base_path, log_level, dev, yes):
     # TODO [MV, AT] remove code duplication with aim up cmd implementation
     if not log_level:
@@ -67,9 +66,6 @@ def server(host, port,
         return
     os.environ[AIM_SERVER_MOUNTED_REPO_PATH] = repo
 
-    if package not in Package.pool:
-        Package.load_package(package)
-
     click.secho('Running Aim Server on repo `{}`'.format(repo_inst), fg='yellow')
     click.echo('Server is mounted on {}:{}'.format(host, port), err=True)
     click.echo('Press Ctrl+C to exit')
@@ -79,17 +75,20 @@ def server(host, port,
     # delete the repo as it needs to be opened in a child process in dev mode
     del repo_inst
 
-    try:
-        from aimcore.transport.server import start_server
+    if dev:
+        import aim
+        import aimcore
+
+        reload_dirs = [os.path.dirname(aim.__file__), os.path.dirname(aimcore.__file__), dev_package_dir]
+    else:
+        reload_dirs = []
 
-        if dev:
-            import aim
-            import aimcore
-            reload_dirs = (os.path.dirname(aim.__file__), os.path.dirname(aim.__file__), dev_package_dir)
-            start_server(host, port, ssl_keyfile, ssl_certfile, log_level=log_level, reload=dev, reload_dirs=reload_dirs)
-        else:
-            start_server(host, port, ssl_keyfile, ssl_certfile, log_level=log_level)
+    try:
+        start_uvicorn_app('aimcore.transport.server:app',
+                          host=host, port=port,
+                          ssl_keyfile=ssl_keyfile, ssl_certfile=ssl_certfile, log_level=log_level,
+                          reload=dev, reload_dirs=reload_dirs)
     except Exception:
         click.echo('Failed to run Aim Tracking Server. '
-                   'Please see the logs for details.')
-        return
+                   'Please see the logs above for details.')
+        exit(1)

src/aimcore/cli/ui/commands.py

@@ -1,14 +1,13 @@
 import os
 import click
 
-from aimcore.cli.utils import set_log_level
-from aimcore.cli.ui.utils import build_db_upgrade_command, build_uvicorn_command, get_free_port_num
+from aimcore.cli.utils import set_log_level, start_uvicorn_app
+from aimcore.cli.ui.utils import build_db_upgrade_command, get_free_port_num
 from aimcore.web.configs import (
     AIM_UI_BASE_PATH,
     AIM_UI_DEFAULT_HOST,
     AIM_UI_DEFAULT_PORT,
     AIM_UI_MOUNTED_REPO_PATH,
-    AIM_UI_PACKAGE_NAME,
     AIM_UI_TELEMETRY_KEY,
     AIM_PROXY_URL,
     AIM_PROFILER_KEY
@@ -34,7 +33,7 @@
                                                                              file_okay=False,
                                                                              dir_okay=True,
                                                                              writable=True))
-@click.option('--package', '--pkg', required=False, default='asp', type=str)
+@click.option('--package', '--pkg', required=False, default='', show_default='asp', type=str)
 @click.option('--dev', is_flag=True, default=False)
 @click.option('--ssl-keyfile', required=False, type=click.Path(exists=True,
                                                                file_okay=True,
@@ -57,14 +56,11 @@ def ui(dev, host, port, workers, uds,
     """
     Start Aim UI with the --repo repository.
     """
-    if dev:
-        os.environ[AIM_ENV_MODE_KEY] = 'dev'
-        log_level = log_level or 'debug'
-    else:
-        os.environ[AIM_ENV_MODE_KEY] = 'prod'
+    if not log_level:
+        log_level = 'debug' if dev else 'warning'
+    set_log_level(log_level)
 
-    if log_level:
-        set_log_level(log_level)
+    os.environ[AIM_ENV_MODE_KEY] = 'dev' if dev else 'prod'
 
     if base_path:
         # process `base_path` as ui requires leading slash
@@ -82,17 +78,19 @@ def ui(dev, host, port, workers, uds,
             return
         Repo.init(repo)
     repo_inst = Repo.from_path(repo, read_only=True)
-
     os.environ[AIM_UI_MOUNTED_REPO_PATH] = repo
-    os.environ[AIM_UI_PACKAGE_NAME] = package
+
+    dev_package_dir = repo_inst.dev_package_dir
+    if package:
+        repo_inst.set_active_package(pkg_name=package)
 
     try:
         db_cmd = build_db_upgrade_command()
         exec_cmd(db_cmd, stream_output=True)
     except ShellCommandException:
         click.echo('Failed to initialize Aim DB. '
                    'Please see the logs above for details.')
-        return
+        exit(1)
 
     if port == 0:
         try:
@@ -123,12 +121,24 @@ def ui(dev, host, port, workers, uds,
     if profiler:
         os.environ[AIM_PROFILER_KEY] = '1'
 
+    if dev:
+        import aim
+        import aimstack
+        import aimcore
+
+        reload_dirs = [os.path.dirname(aim.__file__), os.path.dirname(aimcore.__file__), os.path.dirname(aimstack.__file__), dev_package_dir]
+    else:
+        reload_dirs = []
+
     try:
-        server_cmd = build_uvicorn_command(host, port, workers, uds, ssl_keyfile, ssl_certfile, log_level, package)
-        exec_cmd(server_cmd, stream_output=True)
-    except ShellCommandException:
-        click.echo('Failed to run Aim UI. Please see the logs above for details.')
-        return
+        start_uvicorn_app('aimcore.web.run:app',
+                          host=host, port=port, workers=workers, uds=uds,
+                          ssl_keyfile=ssl_keyfile, ssl_certfile=ssl_certfile, log_level=log_level,
+                          reload=dev, reload_dirs=reload_dirs)
+    except Exception:
+        click.echo('Failed to run Aim UI. '
+                   'Please see the logs above for details.')
+        exit(1)
 
 
 @click.command('up', context_settings={'ignore_unknown_options': True, 'allow_extra_args': True}, hidden=True)

src/aimcore/cli/ui/utils.py

@@ -15,40 +15,6 @@ def build_db_upgrade_command():
     return [sys.executable, '-m', 'alembic', '-c', ini_file, 'upgrade', 'head']
 
 
-def build_uvicorn_command(host, port, num_workers, uds_path, ssl_keyfile, ssl_certfile, log_level, pkg_name):
-    cmd = [sys.executable, '-m', 'uvicorn',
-           '--host', host, '--port', f'{port}',
-           '--workers', f'{num_workers}']
-    if os.getenv(AIM_ENV_MODE_KEY, 'prod') == 'prod':
-        log_level = log_level or 'error'
-    else:
-        import aim
-        import aimstack
-        from aimcore import web as aim_web
-
-        cmd += ['--reload']
-        cmd += ['--reload-dir', os.path.dirname(aim.__file__)]
-        cmd += ['--reload-dir', os.path.dirname(aim_web.__file__)]
-        cmd += ['--reload-dir', os.path.dirname(aimstack.__file__)]
-
-        from aim._sdk.package_utils import Package
-        if pkg_name not in Package.pool:
-            Package.load_package(pkg_name)
-            pkg = Package.pool[pkg_name]
-            cmd += ['--reload-dir', os.path.dirname(pkg._path)]
-
-        log_level = log_level or 'debug'
-    if uds_path:
-        cmd += ['--uds', uds_path]
-    if ssl_keyfile:
-        cmd += ['--ssl-keyfile', ssl_keyfile]
-    if ssl_certfile:
-        cmd += ['--ssl-certfile', ssl_certfile]
-    cmd += ['--log-level', log_level.lower()]
-    cmd += ['aimcore.web.run:app']
-    return cmd
-
-
 def get_free_port_num():
     import socket
     s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

src/aimcore/cli/utils.py

@@ -9,3 +9,8 @@ def set_log_level(log_level):
         raise ValueError('Invalid log level: %s' % log_level)
     os.environ[AIM_LOG_LEVEL_KEY] = str(numeric_level)
     logging.basicConfig(level=numeric_level)
+
+
+def start_uvicorn_app(app: str, **uvicorn_args):
+    import uvicorn
+    uvicorn.run(app, **uvicorn_args)

src/aimcore/transport/handlers.py

@@ -5,7 +5,6 @@
 
 from aim import Repo
 from aim._sdk.local_storage import LocalFileManager
-from aim._sdk.dev_package import DevPackage
 from aimcore.cleanup import AutoClean
 
 

src/aimcore/transport/server.py

@@ -137,11 +137,3 @@ async def run_write_instructions(websocket: WebSocket, client_uri: str):
 
 
 app = create_app()
-
-
-def start_server(host, port, ssl_keyfile=None, ssl_certfile=None, *, log_level='info', reload=False, reload_dirs=()):
-    import uvicorn
-    uvicorn.run('aimcore.transport.server:app', host=host, port=port,
-                ssl_keyfile=ssl_keyfile, ssl_certfile=ssl_certfile,
-                log_level=log_level,
-                reload=reload, reload_dirs=reload_dirs)

src/aimcore/transport/tracking.py

@@ -1,5 +1,6 @@
 import uuid
 import base64
+import logging
 
 from typing import Dict, Union
 from fastapi import WebSocket, Request, APIRouter
@@ -8,6 +9,8 @@
 import aimcore.transport.message_utils as utils
 from aim._core.storage.treeutils import encode_tree, decode_tree
 
+logger = logging.getLogger(__name__)
+
 
 def get_handler():
     return str(uuid.uuid4())
@@ -109,6 +112,7 @@ async def get_resource(self,
             except KeyError:
                 pass
 
+            logger.debug(f'Caught exception {e}. Sending response 400.')
             return JSONResponse({
                 'exception': utils.build_exception(e),
             }, status_code=400)
@@ -118,6 +122,7 @@ async def release_resource(self, client_uri, resource_handler):
             self._verify_resource_handler(resource_handler, client_uri)
             del self.resource_pool[resource_handler]
         except Exception as e:
+            logger.debug(f'Caught exception {e}. Sending response 400.')
             return JSONResponse({
                 'exception': utils.build_exception(e),
             }, status_code=400)
@@ -157,6 +162,7 @@ async def run_instruction(self, client_uri: str,
 
             return StreamingResponse(utils.pack_stream(encode_tree(result)))
         except Exception as e:
+            logger.debug(f'Caught exception {e}. Sending response 400.')
             return JSONResponse({
                 'exception': utils.build_exception(e),
             }, status_code=400)

src/aimcore/web/api/__init__.py

@@ -7,9 +7,8 @@
 from fastapi.responses import JSONResponse
 
 from aim._sdk.configs import get_aim_repo_name
-from aim._sdk.package_utils import Package
 
-from aimcore.web.configs import AIM_PROFILER_KEY, AIM_UI_PACKAGE_NAME
+from aimcore.web.configs import AIM_PROFILER_KEY
 from aimcore.web.middlewares.profiler import PyInstrumentProfilerMiddleware
 from aimcore.web.utils import get_root_path
 
@@ -74,10 +73,6 @@ def create_app():
         api_app.add_middleware(PyInstrumentProfilerMiddleware,
                                repo_path=os.path.join(get_root_path(), get_aim_repo_name()))
 
-    ui_pkg_name = os.environ.get(AIM_UI_PACKAGE_NAME)
-    if ui_pkg_name not in Package.pool:
-        Package.load_package(ui_pkg_name)
-
     api_app.include_router(dashboard_apps_router, prefix='/apps')
     api_app.include_router(dashboards_router, prefix='/dashboards')
     api_app.include_router(boards_router, prefix='/boards')

src/aimcore/web/api/boards/views.py

@@ -1,6 +1,6 @@
 from fastapi import Depends, HTTPException
 from fastapi.responses import JSONResponse
-from aimcore.web.utils import get_root_package
+from aimcore.web.utils import get_root_package, get_root_package_name
 from aimcore.web.api.utils import APIRouter  # wrapper for fastapi.APIRouter
 
 from aimcore.web.api.boards.pydantic_models import BoardOut, BoardListOut
@@ -13,12 +13,18 @@
 
 @boards_router.get('/', response_model=BoardOut)
 async def board_list_api(package=Depends(get_root_package)):
+    if package is None:
+        raise HTTPException(status_code=400, detail=f'Failed to load current active '
+                                                    f'package \'{get_root_package_name()}\'.')
     result = [board.as_posix() for board in package.boards]
     return JSONResponse(result)
 
 
 @boards_router.get('/{board_path:path}', response_model=BoardListOut)
 async def board_get_api(board_path: str, package=Depends(get_root_package)):
+    if package is None:
+        raise HTTPException(status_code=400, detail=f'Failed to load current active '
+                                                    f'package \'{get_root_package_name()}\'.')
     board: pathlib.Path = package.boards_directory / board_path
     if not board.exists():
         raise HTTPException(status_code=404)