def __init__(self):
self.name = os.getenv('PROJECT_NAME') or self.DEFAULT_PROJECT_NAME
self.path = os.getenv('PROJECT_PATH') or self.DEFAULT_PROJECT_PATH
self.description = ''
self.tf_enabled = TFSummaryAdapter.exists()
self.repo = AimRepo(repo_full_path=self.REPO_PATH,
mode=AimRepo.READING_MODE)
def init():
"""
Initializes new repository in the current working directory:
- Creates .aim directory
- Adds .aim/config.json file with initial configuration
"""
repo = AimRepo(os.getcwd())
re_init = False
# Check whether repo already exists
if repo.exists():
re_init = click.confirm('Aim repository is already initialized. ' +
'Do you want to re-initialize it?')
if not re_init:
return
# Clear old repo
repo.rm()
# Init repo
new_repo = AimRepo(os.getcwd())
if new_repo.init():
if re_init:
click.echo(
'Re-initialized empty Aim repository at {}'.format(new_repo))
else:
click.echo(('Initialized a new ' +
'Aim repository at {}').format(new_repo))
def _save_record(repo: AimRepo,
artifact: Artifact,
record: Record,
dir_path: str = None):
if record.binary_type is Artifact.IMAGE:
# Get image name and abs path
img_name_time = math.floor(time.time() * 1000)
img_name_random = random_str(10)
img_name = '{time}__{random}.jpg'.format(time=img_name_time,
random=img_name_random)
res = repo.store_image(img_name, record.cat)
# Save image at specified path
artifact.save_blobs(res['path'], res['abs_path'])
elif record.binary_type == Artifact.MODEL:
# Get model name, directory and zip archive paths
file_res = repo.store_model_file(record.name, record.cat)
# Save model at specified path
model_save_res = artifact.save_blobs(file_res)
res = repo.store_model(record.name, record.data['model_name'],
record.data['epoch'], record.data['meta'],
model_save_res, record.cat)
# Archive model directory
archive_dir(res['zip_path'], res['dir_path'])
elif record.binary_type == Artifact.PROTOBUF:
writer_type = RecordWriter.AIMRECORDS_WRITER
write_mode = 'w' if record.is_singular else 'a'
writer = RecordWriter.get_writer(writer_type, repo.records_storage)
writer.write(artifact.get_inst_unique_name(), write_mode,
record.content, record.context)
res = repo.store_artifact(record.name, record.cat, record.data,
writer_type, record.binary_type,
record.context)
else:
if record.binary_type == Artifact.JSON:
writer = RecordWriter.JSON_WRITER
else:
writer = RecordWriter.JSON_LOG_WRITER
file_name = '{}.log'.format(record.name)
res = repo.store_file(file_name, record.name, record.cat,
record.data, dir_path)
write_mode = 'w' if record.is_singular else 'a'
writer = RecordWriter.get_writer(writer)
writer.write(res['abs_path'], write_mode, record.content)
return res
def select_runs(expression: Optional[str] = None,
repo_path: Optional[str] = None):
if repo_path is not None:
repo = AimRepo(repo_full_path=repo_path)
else:
repo = AimRepo.get_working_repo(mode=AimRepo.READING_MODE)
if not repo:
return None
if expression and 'run.archived' not in expression:
default_expression = 'run.archived is not True'
else:
default_expression = None
return repo.select_runs(expression, default_expression)
def _parse_env_vars(self, env_vars):
env_vars = env_vars or ''
env_vars_arr = env_vars.split(' ')
filtered_env_vars = []
automated = False
automated_branch = None
automated_commit = None
for e in env_vars_arr:
if AIM_AUTOMATED_EXEC_ENV_VAR in e:
automated = True
elif AIM_BRANCH_ENV_VAR in e:
_, _, automated_branch = e.rpartition('=')
else:
filtered_env_vars.append(e)
if automated:
if not automated_branch:
automated_branch = AIM_DEFAULT_BRANCH_NAME
automated_commit = AimRepo.generate_commit_hash()
filtered_env_vars.append('{}={}'.format(AIM_BRANCH_ENV_VAR,
automated_branch))
filtered_env_vars.append('{}={}'.format(AIM_COMMIT_ENV_VAR,
automated_commit))
filtered_env_vars.append('{}={}'.format(AIM_PROCESS_ENV_VAR,
self.process_uuid))
filtered_env_vars.append('{}=1'.format(AIM_AUTOMATED_EXEC_ENV_VAR))
return {
'env_vars': ' '.join(filtered_env_vars),
'automated': automated,
'automated_branch': automated_branch,
'automated_commit': automated_commit,
}
def down(repo_inst, repo):
check_docker_dependency()
repo_path = clean_repo_path(repo)
if repo_path:
repo_inst = AimRepo(repo_path)
if repo_inst is None or not repo_inst.exists():
repo_init_alert()
return
# Kill all identical running containers
cont = AimContainer(repo_inst)
if cont.get_container() is not None:
click.echo('Shutting down Aim UI at `{}`'.format(repo_inst.path))
cont.kill()
click.echo('Done')
else:
click.echo('No running Aim UI at `{}`'.format(repo_inst.path))
def de_entry_point(ctx):
repo = ctx.obj or AimRepo.get_working_repo()
if repo is None:
repo_init_alert()
exit()
ctx.obj = repo
if not AimContainer.is_docker_installed():
docker_requirement_alert()
exit()
def select_metrics(search_statement: str, repo_path: Optional[str] = None):
if repo_path is not None:
repo = AimRepo(repo_full_path=repo_path)
else:
repo = AimRepo.get_working_repo(mode=AimRepo.READING_MODE)
if not repo:
return None
parser = Statement()
parsed_stmt = parser.parse(search_statement.strip())
statement_select = parsed_stmt.node['select']
statement_expr = parsed_stmt.node['expression']
if 'run.archived' not in search_statement:
default_expression = 'run.archived is not True'
else:
default_expression = None
return repo.select_metrics(statement_select, statement_expr,
default_expression)
class Project:
DEFAULT_PROJECT_NAME = 'Project'
DEFAULT_PROJECT_PATH = '/project'
REPO_PATH = '/store'
def __init__(self):
self.name = os.getenv('PROJECT_NAME') or self.DEFAULT_PROJECT_NAME
self.path = os.getenv('PROJECT_PATH') or self.DEFAULT_PROJECT_PATH
self.description = ''
self.tf_enabled = TFSummaryAdapter.exists()
self.repo = AimRepo(repo_full_path=self.REPO_PATH,
mode=AimRepo.READING_MODE)
def exists(self):
return self.repo and self.repo.exists()
def init():
repo = AimRepo(os.getcwd())
re_init = False
# Check whether repo already exists
if repo.exists():
re_init = click.confirm('Aim repository is already initialized. ' +
'Do you want to re-initialize it?')
if not re_init:
return
# Clear old repo
repo.rm()
# Init repo
new_repo = AimRepo(os.getcwd())
if new_repo.init():
if re_init:
click.echo(
'Re-initialized empty Aim repository in {}'.format(new_repo))
else:
click.echo(('Initialized empty ' +
'Aim repository in {}').format(new_repo))
def save(self, repo: AimRepo, artifact: Artifact) -> bool:
"""
Stores serialized instance into .aim repo
"""
item = artifact.serialize()
res = None
if isinstance(item, Record):
res = ArtifactWriter._save_record(repo, artifact, item)
elif isinstance(item, RecordCollection):
dir_path, dir_rel_path = repo.store_dir(item.name, item.cat,
item.data)
res = []
for record in item.records:
# Store dir files
res.append(
ArtifactWriter._save_record(repo, artifact, record,
dir_rel_path))
# Save dict
return res
def get_repo(path: Optional[str] = None,
experiment_name: Optional[str] = None) -> AimRepo:
# Auto commit
if os.getenv(AIM_AUTOMATED_EXEC_ENV_VAR):
# Get Aim environment variables
branch_name = os.getenv(AIM_BRANCH_ENV_VAR)
commit_hash = os.getenv(AIM_COMMIT_ENV_VAR)
else:
commit_hash = AimRepo.generate_commit_hash()
if experiment_name is not None:
branch_name = experiment_name
else:
# FIXME: Get active experiment name from given repo
# if path is specified. Currently active experiment name of
# the highest repo in the hierarchy will be returned.
branch_name = AimRepo.get_active_branch_if_exists() \
or AIM_DEFAULT_BRANCH_NAME
if path is not None:
repo = AimRepo(path)
if not repo.exists():
if not repo.init():
raise ValueError('can not create repo `{}`'.format(path))
repo = AimRepo(path, branch_name, commit_hash)
else:
if AimRepo.get_working_repo() is None:
path = os.getcwd()
repo = AimRepo(path)
if not repo.init():
raise ValueError('can not create repo `{}`'.format(path))
repo = AimRepo(path, branch_name, commit_hash)
else:
repo = AimRepo.get_working_repo(branch_name, commit_hash)
return repo
def cli_entry_point(ctx, verbose):
if verbose:
click.echo('Verbose mode is on')
# Init repo instance
ctx.obj = AimRepo.get_working_repo()
def up(repo_inst, dev, host, port, version, repo, tf_logs, detach):
check_docker_dependency()
repo_path = clean_repo_path(repo)
if repo_path:
repo_inst = AimRepo(repo_path)
if repo_inst is None or not repo_inst.exists():
repo_init_alert()
return
# Dev param
# 0 => pull and run official image of container
# 1 => run production build from local environment
# 2 => run in development mode
cont = AimContainer(repo_inst, dev=bool(dev))
if tf_logs:
cont.mount_volume(tf_logs, AIM_TF_LOGS_PATH)
click.echo(
click.style('Running Aim UI on repo `{}`'.format(repo_inst),
fg='yellow'))
# Check if image exist
if dev == 0 and not cont.image_exist(version):
click.echo('Pulling Aim UI docker image, please wait...')
if not cont.pull(version):
docker_image_pull_fail_alert()
return
else:
click.echo('Successfully pulled Aim UI image')
if cont.get_container(running_only=True):
kill = click.confirm('Aim UI is already running. ' +
'Do you want to restart it?')
if not kill:
return
# Kill all identical running containers
cont.kill()
# if dev < 2
# cont.bind(port + 1, AIM_CONTAINER_CMD_PORT)
cont_id = cont.up(port, host, version)
if not cont_id:
click.echo('Failed to run Aim UI. Please try again.')
return
# cont_cmd = AimContainerCommandManager((port + 1) if dev < 2
# else AIM_CONTAINER_CMD_PORT)
# cont_cmd.listen()
# Kill container after keyboard interruption
def graceful_shutdown():
# cont_cmd.kill()
click.echo('')
click.echo('Shutting down...')
try:
cont.kill()
except Exception:
pass
click.echo('Done')
sys.exit(0)
# Add keyboard signal interruption listener
# signal.signal(signal.SIGINT, graceful_shutdown)
# signal.pause()
try:
sleep(4)
click.echo('Open http://{}:{}'.format(host, port))
if not detach:
click.echo('Press Ctrl+C to exit')
while True:
if cont.get_container() is not None:
sleep(0.3)
else:
click.echo('Exited')
return
except KeyboardInterrupt:
graceful_shutdown()
sys.exit(0)
def up(repo_inst, dev, host, port, version, repo, tf_logs, detach):
check_docker_dependency()
repo_path = clean_repo_path(repo)
if repo_path:
repo_inst = AimRepo(repo_path)
if repo_inst is None or not repo_inst.exists():
repo_init_alert()
return
# Dev param
# 0 => pull and run official image of container
# 1 => run production build from local environment
# 2 => run in development mode
cont = AimContainer(repo_inst, dev=bool(dev))
if tf_logs:
cont.mount_volume(tf_logs, AIM_TF_LOGS_PATH)
if os.getenv(AIM_CONTAINER_TELEMETRY_FLAG) is not None \
and os.getenv(AIM_CONTAINER_TELEMETRY_FLAG) == '0':
cont.turn_telemetry_off()
else:
cont.turn_telemetry_on()
alert_msg = 'Aim UI collects anonymous usage analytics.'
opt_out_msg = 'Read how to opt-out here: '
opt_out_url = 'https://github.com/aimhubio/aim#anonymized-telemetry'
line_width = max(len(opt_out_msg), len(alert_msg)) + 16
click.echo('┌{}┐'.format('-' * (line_width - 2)))
click.echo('{}{}{}'.format(' ' * ((line_width - len(alert_msg)) // 2),
alert_msg,
' ' * ((line_width - len(alert_msg)) // 2)))
click.echo('{}{}{}'.format(
' ' * ((line_width - len(opt_out_msg)) // 2), opt_out_msg,
' ' * ((line_width - len(opt_out_msg)) // 2)))
click.echo('{}{}{}'.format(
' ' * ((line_width - len(opt_out_url)) // 2), opt_out_url,
' ' * ((line_width - len(opt_out_url)) // 2)))
click.echo('└{}┘'.format('-' * (line_width - 2)))
click.echo(
click.style('Running Aim UI on repo `{}`'.format(repo_inst),
fg='yellow'))
# Check if image exist
if dev == 0 and not cont.image_exist(version):
click.echo('Pulling Aim UI docker image, please wait...')
if not cont.pull(version):
docker_image_pull_fail_alert()
return
else:
click.echo('Successfully pulled Aim UI image')
if cont.get_container(running_only=True):
kill = click.confirm('Aim UI is already running. ' +
'Do you want to restart it?')
if not kill:
return
# Kill all identical running containers
cont.kill()
# if dev < 2
# cont.bind(port + 1, AIM_CONTAINER_CMD_PORT)
cont_id = cont.up(port, host, version)
if not cont_id:
click.echo('Failed to run Aim UI. Please try again.')
return
# cont_cmd = AimContainerCommandManager((port + 1) if dev < 2
# else AIM_CONTAINER_CMD_PORT)
# cont_cmd.listen()
# Kill container after keyboard interruption
def graceful_shutdown():
# cont_cmd.kill()
click.echo('')
click.echo('Shutting down...')
try:
cont.kill()
except Exception:
pass
click.echo('Done')
sys.exit(0)
# Add keyboard signal interruption listener
# signal.signal(signal.SIGINT, graceful_shutdown)
# signal.pause()
try:
sleep(4)
click.echo('Open http://{}:{}'.format(host, port))
if not detach:
click.echo('Press Ctrl+C to exit')
while True:
if cont.get_container() is not None:
sleep(0.3)
else:
click.echo('Exited')
return
except KeyboardInterrupt:
graceful_shutdown()
sys.exit(0)
class Project:
DEFAULT_PROJECT_NAME = 'Project'
DEFAULT_PROJECT_PATH = '/project'
REPO_PATH = '/store'
def __init__(self):
self.name = os.getenv('PROJECT_NAME') or self.DEFAULT_PROJECT_NAME
self.path = os.getenv('PROJECT_PATH') or self.DEFAULT_PROJECT_PATH
self.description = ''
self.tf_enabled = TFSummaryAdapter.exists()
self.repo = AimRepo(repo_full_path=self.REPO_PATH,
mode=AimRepo.READING_MODE)
def exists(self):
return self.repo and self.repo.exists()
def get_run_config(self, experiment_name, run_hash):
config_file_path = os.path.join(self.repo.path,
experiment_name,
run_hash)
config_file_path = os.path.join(config_file_path,
AIM_COMMIT_CONFIG_FILE_NAME)
if not os.path.isfile(config_file_path):
return None
with open(config_file_path, 'r+') as config_file:
try:
configs = json.loads(config_file.read())
except:
configs = None
return configs
def get_modified_runs(self, last_modified=0):
updated_runs = {}
if not self.repo:
return updated_runs
experiments = self.repo.list_branches()
for experiment in experiments:
runs = self.repo.list_branch_commits(experiment)
for run in runs:
run_path = os.path.join(self.repo.path, experiment, run)
run_config_path = os.path.join(run_path,
AIM_COMMIT_CONFIG_FILE_NAME)
if not os.path.exists(run_config_path):
continue
run_config_path = Path(run_config_path)
if not run_config_path or not run_config_path.exists():
continue
run_stats = run_config_path.stat()
if not run_stats:
continue
modified_time = run_config_path.stat().st_mtime
if not modified_time or modified_time > last_modified:
updated_runs.setdefault(experiment, [])
updated_runs[experiment].append((run, modified_time))
return updated_runs
def commit(repo, message, code):
if repo is None:
click.echo('Repository does not exist')
return
commit_hash = AimRepo.generate_commit_hash()
message = message.strip() or int(time.time())
# Check if there is anything to commit
if repo.is_index_empty():
click.echo('Nothing to commit')
return
branch_name = branch_hash = None
if code:
try:
from aim.version_control.factory import Factory
except Exception:
click.echo('Git executable not found')
return
vc = Factory.create(Factory.GIT)
# Check if version control repo exists
if vc is None or vc.get_repo() is None:
click.echo('No git repository (or any parent up to mount ' +
'point /) found. Initialize git repository ' +
'by running `git init`')
return
# Check untracked files
if len(vc.get_untracked_files()):
click.echo('You have untracked files, please add them to the ' +
'index before committing your changes by running ' +
'`git add -A`')
return
# Check if HEAD exists
if not vc.get_head_hash():
click.echo('Needed a single revision. ' +
'You do not have the git initial commit yet')
return
# Commit changes to a new created branch and return branch name
branch_name, branch_hash = vc.commit_changes_to_branch(
message, commit_hash)
# Get the latest branch
latest_branch = repo.get_latest_vc_branch() or {}
latest_branch = latest_branch.get('branch')
if latest_branch is None:
latest_branch = vc.get_head_hash()
# Get diff between current commit and latest commit(or HEAD)
diff = vc.get_diff_text(latest_branch, branch_name)
repo.save_diff(diff)
# Commit
commit_res = repo.commit(commit_hash, message, branch_name, branch_hash)
click.echo(
click.style('[{b}/{c} commit]'.format(b=commit_res['branch'],
c=commit_res['commit']),
fg='yellow'))
click.echo(message)
def test_run_archivation(self):
"""
Test aim repo can be removed correctly.
"""
prefix = os.getcwd() + '/'
with tempfile.TemporaryDirectory(prefix=prefix) as tmp_dir_path:
AimRepo(tmp_dir_path).init()
experiment = 'test'
run_hash = AimRepo.generate_commit_hash()
repo = AimRepo(tmp_dir_path, experiment, run_hash)
repo.commit_init()
repo.commit_finish()
self.assertFalse(repo.is_archived(experiment, run_hash))
repo.archive(experiment, run_hash)
self.assertTrue(repo.is_archived(experiment, run_hash))
repo.unarchive(experiment, run_hash)
self.assertFalse(repo.is_archived(experiment, run_hash))