forked from ci4s/aim
97 lines
3.7 KiB
Python
97 lines
3.7 KiB
Python
import os
|
|
import click
|
|
|
|
from aim.sdk.repo import Repo
|
|
from aim.sdk.utils import clean_repo_path
|
|
from aim.cli.utils import (
|
|
ShellCommandException,
|
|
build_uvicorn_command,
|
|
exec_cmd,
|
|
get_free_port_num,
|
|
get_repo_instance,
|
|
set_log_level
|
|
)
|
|
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('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('--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,
|
|
readable=True))
|
|
@click.option('--ssl-certfile', required=False, type=click.Path(exists=True,
|
|
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,
|
|
repo, ssl_keyfile, ssl_certfile,
|
|
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_inst = get_repo_instance(repo_path, yes)
|
|
|
|
if not repo_inst:
|
|
return
|
|
|
|
os.environ[AIM_SERVER_MOUNTED_REPO_PATH] = repo_inst.path
|
|
|
|
click.secho('Running Aim Server on repo `{}`'.format(repo), fg='yellow')
|
|
click.echo('Server is mounted on aim://{}:{}'.format(host, port), err=True)
|
|
click.echo('Press Ctrl+C to exit')
|
|
|
|
try:
|
|
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 ShellCommandException:
|
|
click.echo('Failed to run Aim Tracking Server. '
|
|
'Please see the logs above for details.')
|
|
exit(1)
|