Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[feat] Ability to change UI active package after aim ui has been started #2935

Merged
merged 9 commits into from
Aug 4, 2023
9 changes: 9 additions & 0 deletions src/aimcore/cli/package/commands.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
26 changes: 20 additions & 6 deletions src/aimcore/cli/package/watcher.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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
33 changes: 16 additions & 17 deletions src/aimcore/cli/server/commands.py
Original file line number Diff line number Diff line change
@@ -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 (
Expand All @@ -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,
Expand All @@ -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:
Expand Down Expand Up @@ -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')
Expand All @@ -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)
48 changes: 29 additions & 19 deletions src/aimcore/cli/ui/commands.py
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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,
Expand All @@ -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
Expand All @@ -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:
Expand Down Expand Up @@ -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)
Expand Down
34 changes: 0 additions & 34 deletions src/aimcore/cli/ui/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
5 changes: 5 additions & 0 deletions src/aimcore/cli/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
1 change: 0 additions & 1 deletion src/aimcore/transport/handlers.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down
8 changes: 0 additions & 8 deletions src/aimcore/transport/server.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
6 changes: 6 additions & 0 deletions src/aimcore/transport/tracking.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import uuid
import base64
import logging

from typing import Dict, Union
from fastapi import WebSocket, Request, APIRouter
Expand All @@ -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())
Expand Down Expand Up @@ -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)
Expand All @@ -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)
Expand Down Expand Up @@ -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)
7 changes: 1 addition & 6 deletions src/aimcore/web/api/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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')
Expand Down
8 changes: 7 additions & 1 deletion src/aimcore/web/api/boards/views.py
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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)
Expand Down
Loading
Loading