Skip to content

Commit

Permalink
[feat] Change transport from grpc to http(fastapi) for tracking server (
Browse files Browse the repository at this point in the history
  • Loading branch information
mihran113 committed Feb 27, 2024
1 parent 2cda7e9 commit 3373cbb
Show file tree
Hide file tree
Showing 48 changed files with 757 additions and 3,808 deletions.
5 changes: 5 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
@@ -1,5 +1,10 @@
# Changelog

## 3.19.0

### Enhancements:
- Replace grpc with http/ws as transport for aim tracking server (mihran113)

## 3.18.1 Feb 7, 2024

### Enhancements:
Expand Down
73 changes: 53 additions & 20 deletions aim/cli/server/commands.py
Original file line number Diff line number Diff line change
@@ -1,21 +1,25 @@
import os
import click

from aim.cli.utils import set_log_level
from aim.sdk.repo import Repo, RepoStatus
from aim.sdk.utils import clean_repo_path
from aim.ext.transport.config import AIM_SERVER_DEFAULT_HOST, AIM_SERVER_DEFAULT_PORT, AIM_SERVER_MOUNTED_REPO_PATH
from aim.ext.transport.server import run_router
from aim.cli.utils import set_log_level, build_uvicorn_command, get_free_port_num, exec_cmd
from aim.ext.transport.config import (
AIM_SERVER_DEFAULT_PORT,
AIM_SERVER_DEFAULT_HOST,
AIM_SERVER_MOUNTED_REPO_PATH,
AIM_SERVER_BASE_PATH
)
from aim.web.configs import AIM_ENV_MODE_KEY


@click.command()
@click.command('server')
@click.option('-h', '--host', default=AIM_SERVER_DEFAULT_HOST, type=str)
@click.option('-p', '--port', default=AIM_SERVER_DEFAULT_PORT, type=int)
@click.option('-w', '--workers', default=1, type=int)
@click.option('--repo', required=False, type=click.Path(exists=True,
file_okay=False,
dir_okay=True,
writable=True))
@click.option('--repo', required=False, default=os.getcwd(), type=click.Path(exists=True,
file_okay=False,
dir_okay=True,
writable=True))
@click.option('--ssl-keyfile', required=False, type=click.Path(exists=True,
file_okay=True,
dir_okay=False,
Expand All @@ -24,23 +28,51 @@
file_okay=True,
dir_okay=False,
readable=True))
@click.option('--base-path', required=False, default='', type=str)
@click.option('--log-level', required=False, default='', type=str)
@click.option('--dev', is_flag=True, default=False)
@click.option('-y', '--yes', is_flag=True, help='Automatically confirm prompt')
def server(host, port, workers,
def server(host, port,
repo, ssl_keyfile, ssl_certfile,
log_level, yes):
# TODO [MV, AT] remove code duplication with aim up cmd implementation
base_path, log_level, dev, yes):
"""
Starts the Aim remote tracking server for real-time logging.
The Aim tracking server facilitates real-time logging of experiments
from remote locations. This command launches the server with specified
configurations, including host, port, and associated repository.
Like the UI, the server can also run in production or development mode.
"""
if dev:
os.environ[AIM_ENV_MODE_KEY] = 'dev'
log_level = log_level or 'debug'
else:
os.environ[AIM_ENV_MODE_KEY] = 'prod'

if log_level:
set_log_level(log_level)

if base_path:
if base_path.endswith('/'):
base_path = base_path[:-1]
if not base_path.startswith('/'):
base_path = f'/{base_path}'
os.environ[AIM_SERVER_BASE_PATH] = base_path

if port == 0:
try:
port = get_free_port_num()
except Exception:
pass

repo_path = clean_repo_path(repo) or Repo.default_repo_path()
repo_status = Repo.check_repo_status(repo_path)
if repo_status == RepoStatus.MISSING:
if yes:
init_repo = True
else:
init_repo = click.confirm(f'\'{repo_path}\' is not a valid Aim repository. Do you want to initialize it?')

if not init_repo:
click.echo('To initialize repo please run the following command:')
click.secho('aim init', fg='yellow')
Expand All @@ -63,15 +95,16 @@ def server(host, port, workers,

os.environ[AIM_SERVER_MOUNTED_REPO_PATH] = repo_inst.path

click.echo(
click.style('Running Aim Server on repo `{}`'.format(repo_inst),
fg='yellow'))
click.echo('Server is mounted on {}:{}'.format(host, port), err=True)
click.secho('Running Aim Server on repo `{}`'.format(repo), fg='yellow')
click.echo('Server is mounted on aim:https://{}:{}'.format(host, port), err=True)
click.echo('Press Ctrl+C to exit')

try:
run_router(host, port, workers, ssl_keyfile, ssl_certfile)
cmd = build_uvicorn_command('aim.ext.transport.run:app',
host=host, port=port,
ssl_keyfile=ssl_keyfile, ssl_certfile=ssl_certfile, log_level=log_level)
exec_cmd(cmd, stream_output=True)
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)
25 changes: 17 additions & 8 deletions aim/cli/up/commands.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,13 @@
import click

from aim.cli.utils import set_log_level
from aim.cli.up.utils import build_db_upgrade_command, build_uvicorn_command, get_free_port_num
from aim.cli.utils import (
build_db_upgrade_command,
build_uvicorn_command,
get_free_port_num,
exec_cmd,
ShellCommandException
)
from aim.web.configs import (
AIM_ENV_MODE_KEY,
AIM_TF_LOGS_PATH_KEY,
Expand All @@ -15,8 +21,6 @@
)
from aim.sdk.repo import Repo, RepoStatus
from aim.sdk.utils import clean_repo_path
from aim.web.utils import exec_cmd
from aim.web.utils import ShellCommandException


@click.command()
Expand All @@ -42,14 +46,14 @@
dir_okay=False,
readable=True))
@click.option('--base-path', required=False, default='', type=str)
@click.option('--force-init', is_flag=True, default=False)
@click.option('--profiler', is_flag=True, default=False)
@click.option('--log-level', required=False, default='', type=str)
@click.option('-y', '--yes', is_flag=True, help='Automatically confirm prompt')
def up(dev, host, port, workers, uds,
repo, tf_logs,
ssl_keyfile, ssl_certfile,
base_path, force_init,
profiler, log_level):
base_path, profiler,
log_level, yes):
if dev:
os.environ[AIM_ENV_MODE_KEY] = 'dev'
log_level = log_level or 'debug'
Expand All @@ -71,7 +75,7 @@ def up(dev, host, port, workers, uds,
repo_status = Repo.check_repo_status(repo_path)
if repo_status == RepoStatus.MISSING:
init_repo = None
if not force_init:
if not yes:
init_repo = click.confirm(f'\'{repo_path}\' is not a valid Aim repository. Do you want to initialize it?')
else:
init_repo = True
Expand Down Expand Up @@ -129,7 +133,12 @@ def up(dev, host, port, workers, uds,
os.environ[AIM_PROFILER_KEY] = '1'

try:
server_cmd = build_uvicorn_command(host, port, workers, uds, ssl_keyfile, ssl_certfile, log_level)
server_cmd = build_uvicorn_command(
'aim.web.run:app',
host=host, port=port,
workers=workers, uds=uds,
ssl_keyfile=ssl_keyfile, ssl_certfile=ssl_certfile,
log_level=log_level)
exec_cmd(server_cmd, stream_output=True)
except ShellCommandException:
click.echo('Failed to run Aim UI. Please see the logs above for details.')
Expand Down
45 changes: 0 additions & 45 deletions aim/cli/up/utils.py

This file was deleted.

109 changes: 109 additions & 0 deletions aim/cli/utils.py
Original file line number Diff line number Diff line change
@@ -1,11 +1,120 @@
import os
import sys
import logging
import subprocess

from aim.web.configs import AIM_ENV_MODE_KEY
from aim.web.configs import AIM_LOG_LEVEL_KEY


class ShellCommandException(Exception):
pass


def exec_cmd(
cmd, throw_on_error=True, env=None, stream_output=False, cwd=None, cmd_stdin=None, **kwargs
):
"""
Runs a command as a child process.
A convenience wrapper for running a command from a Python script.
Keyword arguments:
cmd -- the command to run, as a list of strings
throw_on_error -- if true, raises an Exception if the exit code of the program is nonzero
env -- additional environment variables to be defined when running the child process
cwd -- working directory for child process
stream_output -- if true, does not capture standard output and error; if false, captures these
streams and returns them
cmd_stdin -- if specified, passes the specified string as stdin to the child process.
Note on the return value: If stream_output is true, then only the exit code is returned. If
stream_output is false, then a tuple of the exit code, standard output and standard error is
returned.
"""
cmd_env = os.environ.copy()
if env:
cmd_env.update(env)
if stream_output:
child = subprocess.Popen(
cmd, env=cmd_env, cwd=cwd, universal_newlines=True, stdin=subprocess.PIPE, **kwargs
)
child.communicate(cmd_stdin)
exit_code = child.wait()
if throw_on_error and exit_code != 0:
raise ShellCommandException("Non-zero exitcode: %s" % (exit_code))
return exit_code
else:
child = subprocess.Popen(
cmd,
env=cmd_env,
stdout=subprocess.PIPE,
stdin=subprocess.PIPE,
stderr=subprocess.PIPE,
cwd=cwd,
universal_newlines=True,
**kwargs
)
(stdout, stderr) = child.communicate(cmd_stdin)
exit_code = child.wait()
if throw_on_error and exit_code != 0:
raise ShellCommandException(
"Non-zero exit code: %s\n\nSTDOUT:\n%s\n\nSTDERR:%s" % (exit_code, stdout, stderr)
)
return exit_code, stdout, stderr


def set_log_level(log_level):
numeric_level = getattr(logging, log_level.upper(), None)
if not isinstance(numeric_level, int):
raise ValueError('Invalid log level: %s' % log_level)
os.environ[AIM_LOG_LEVEL_KEY] = str(numeric_level)
logging.basicConfig(level=numeric_level)


def build_db_upgrade_command():
from aim import web
web_dir = os.path.dirname(web.__file__)
migrations_dir = os.path.join(web_dir, 'migrations')
if os.getenv(AIM_ENV_MODE_KEY, 'prod') == 'prod':
ini_file = os.path.join(migrations_dir, 'alembic.ini')
else:
ini_file = os.path.join(migrations_dir, 'alembic_dev.ini')
return [sys.executable, '-m', 'alembic', '-c', ini_file, 'upgrade', 'head']


def build_uvicorn_command(app,
host='0.0.0.0',
port=0,
workers=1,
uds=None,
ssl_keyfile=None,
ssl_certfile=None,
log_level='warning',
):
cmd = [sys.executable, '-m', 'uvicorn',
'--host', host, '--port', f'{port}',
'--workers', f'{workers}']

if os.getenv(AIM_ENV_MODE_KEY, 'prod') == 'prod':
log_level = log_level or 'error'
else:
import aim
cmd += ['--reload', '--reload-dir', os.path.dirname(aim.__file__)]
log_level = log_level or 'debug'

if uds:
cmd += ['--uds', uds]
if ssl_keyfile:
cmd += ['--ssl-keyfile', ssl_keyfile]
if ssl_certfile:
cmd += ['--ssl-certfile', ssl_certfile]
cmd += ['--log-level', log_level.lower()]
cmd += [app]
return cmd


def get_free_port_num():
import socket
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
s.bind(('', 0))
port_num = s.getsockname()[1]
s.close()
return port_num
1 change: 1 addition & 0 deletions aim/ext/transport/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
from .client import Client
Loading

0 comments on commit 3373cbb

Please sign in to comment.