|
|
|
|
|
import importlib.util |
|
|
import os |
|
|
import subprocess |
|
|
import sys |
|
|
from typing import Dict, List, Optional |
|
|
|
|
|
from swift.utils import get_logger |
|
|
|
|
|
logger = get_logger() |
|
|
|
|
|
ROUTE_MAPPING: Dict[str, str] = { |
|
|
'pt': 'swift.cli.pt', |
|
|
'sft': 'swift.cli.sft', |
|
|
'infer': 'swift.cli.infer', |
|
|
'merge-lora': 'swift.cli.merge_lora', |
|
|
'web-ui': 'swift.cli.web_ui', |
|
|
'deploy': 'swift.cli.deploy', |
|
|
'rollout': 'swift.cli.rollout', |
|
|
'rlhf': 'swift.cli.rlhf', |
|
|
'sample': 'swift.cli.sample', |
|
|
'export': 'swift.cli.export', |
|
|
'eval': 'swift.cli.eval', |
|
|
'app': 'swift.cli.app', |
|
|
} |
|
|
|
|
|
|
|
|
def use_torchrun() -> bool: |
|
|
nproc_per_node = os.getenv('NPROC_PER_NODE') |
|
|
nnodes = os.getenv('NNODES') |
|
|
if nproc_per_node is None and nnodes is None: |
|
|
return False |
|
|
return True |
|
|
|
|
|
|
|
|
def get_torchrun_args() -> Optional[List[str]]: |
|
|
if not use_torchrun(): |
|
|
return |
|
|
torchrun_args = [] |
|
|
for env_key in ['NPROC_PER_NODE', 'MASTER_PORT', 'NNODES', 'NODE_RANK', 'MASTER_ADDR']: |
|
|
env_val = os.getenv(env_key) |
|
|
if env_val is None: |
|
|
continue |
|
|
torchrun_args += [f'--{env_key.lower()}', env_val] |
|
|
return torchrun_args |
|
|
|
|
|
|
|
|
def _compat_web_ui(argv): |
|
|
|
|
|
method_name = argv[0] |
|
|
if method_name in {'web-ui', 'web_ui'} and ('--model' in argv or '--adapters' in argv or '--ckpt_dir' in argv): |
|
|
argv[0] = 'app' |
|
|
logger.warning('Please use `swift app`.') |
|
|
|
|
|
|
|
|
def cli_main(route_mapping: Optional[Dict[str, str]] = None) -> None: |
|
|
route_mapping = route_mapping or ROUTE_MAPPING |
|
|
argv = sys.argv[1:] |
|
|
_compat_web_ui(argv) |
|
|
method_name = argv[0].replace('_', '-') |
|
|
argv = argv[1:] |
|
|
file_path = importlib.util.find_spec(route_mapping[method_name]).origin |
|
|
torchrun_args = get_torchrun_args() |
|
|
python_cmd = sys.executable |
|
|
if torchrun_args is None or method_name not in {'pt', 'sft', 'rlhf', 'infer'}: |
|
|
args = [python_cmd, file_path, *argv] |
|
|
else: |
|
|
args = [python_cmd, '-m', 'torch.distributed.run', *torchrun_args, file_path, *argv] |
|
|
print(f"run sh: `{' '.join(args)}`", flush=True) |
|
|
result = subprocess.run(args) |
|
|
if result.returncode != 0: |
|
|
sys.exit(result.returncode) |
|
|
|
|
|
|
|
|
if __name__ == '__main__': |
|
|
cli_main() |
|
|
|
