# flake8: noqa
import os
import os.path as osp
import subprocess
from typing import Tuple, Union
import click
from mim.click import CustomCommand, param2lowercase
from mim.utils import (
echo_success,
echo_warning,
exit_with_error,
get_installed_path,
highlighted_error,
is_installed,
recursively_find,
)
@click.command(
'run',
context_settings=dict(ignore_unknown_options=True),
cls=CustomCommand)
@click.argument('package', type=str, callback=param2lowercase)
@click.argument('command', type=str)
@click.argument('other_args', nargs=-1, type=click.UNPROCESSED)
def cli(package: str, command: str, other_args: tuple = ()) -> None:
"""Run arbitrary command of a codebase.
Example:
\b
# Get the Flops of a model
> mim run mmcls get_flops resnet101_b16x8_cifar10.py
# Publish a model
> mim run mmcls publish_model input.pth output.pth
# Train models on a slurm HPC with one GPU
> srun -p partition --gres=gpu:1 mim run mmcls train \
resnet101_b16x8_cifar10.py --work-dir tmp
# Test models on a slurm HPC with one GPU, report accuracy
> srun -p partition --gres=gpu:1 mim run mmcls test \
resnet101_b16x8_cifar10.py tmp/epoch_3.pth --metrics accuracy
# Print help messages of sub-command run
> mim run -h
# Print help messages of sub-command run, list all available scripts in
# codebase mmcls
> mim run mmcls -h
# Print help messages of sub-command run, print the help message of
# training script in mmcls
> mim run mmcls train -h
"""
is_success, msg = run(
package=package, command=command, other_args=other_args)
if is_success:
echo_success(msg) # type: ignore
else:
exit_with_error(msg)
[docs]def run(
package: str, command: str,
other_args: tuple = ()) -> Tuple[bool, Union[str, Exception]]:
"""Run arbitrary command of a codebase.
This command assumes the command scripts have been put into the
``package/tools`` directory.
Args:
package (str): The codebase name.
other_args (tuple, optional): Other arguments, will be passed to the
codebase's script. Defaults to ().
"""
if not is_installed(package):
msg = (f'The codebase {package} is not installed, '
'do you want to install it? ')
if click.confirm(msg):
click.echo(f'Installing {package}')
cmd = ['mim', 'install', package]
ret = subprocess.check_call(cmd)
if ret != 0:
msg = f'{package} is not successfully installed'
raise RuntimeError(highlighted_error(msg))
else:
click.echo(f'{package} is successfully installed')
else:
msg = (f"You can't run commands in {package} without "
f'{package} installed.')
return False, msg
pkg_root = get_installed_path(package)
prefix = osp.join(pkg_root, 'tools/')
command_domain = ''
if ':' in command:
split_command = command.split(':')
command_domain = '/'.join(split_command[:-1])
command = split_command[-1]
files = recursively_find(prefix, command + '.py')
if command_domain == '':
suffix = f'/{command}.py'
else:
suffix = f'/{command_domain}/{command}.py'
files = [f for f in files if f.endswith(suffix)]
if len(files) == 0:
msg = f'No script in codebase {package} has suffix {suffix}.'
raise ValueError(highlighted_error(msg))
elif len(files) > 1:
echo_warning(
f'Multiple scripts in codebase {package} have suffix {suffix}: ')
for f in files:
echo_warning(f)
# Use the shortest path
files.sort(key=lambda x: len(x.split('/')))
echo_warning(f'We are using the script {files[0]}. ')
echo_warning('To use other scripts, you need to use these commands: ')
for f in files[1:]:
cmd = f.split(prefix)[1].split('.')[0].replace('/', ':')
echo_warning(f'Command for {f}: {cmd}')
script = files[0]
click.echo(f'Use the script {script} for command {command}.')
cmd = ['python', script] + list(other_args)
cmd_text = ' '.join(cmd)
click.echo(f'The command to call is {cmd_text}. ')
ret = subprocess.check_call(cmd, env=dict(os.environ))
if ret == 0:
return True, 'The script finished successfully.'
else:
return False, 'The script not finished successfully.'