mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
00d7d6f123
# Description set up torch_cli using argparses ## Details: - add vllm placeholer in the cli - add unittest for cli command see Readme.md to see how to run the cli Pull Request resolved: https://github.com/pytorch/pytorch/pull/160043 Approved by: https://github.com/huydhn
42 lines
1.2 KiB
Python
42 lines
1.2 KiB
Python
import argparse
|
|
import logging
|
|
|
|
from cli.lib.common.cli_helper import (
|
|
register_target_commands_and_runner,
|
|
RichHelp,
|
|
TargetSpec,
|
|
)
|
|
from cli.lib.core.vllm import VllmBuildRunner
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Maps targets to their argparse configuration and runner
|
|
# it adds new target to path python -m cli.run build external {target} with buildrunner
|
|
_TARGETS: dict[str, TargetSpec] = {
|
|
"vllm": {
|
|
"runner": VllmBuildRunner,
|
|
"help": "Build vLLM using docker buildx.",
|
|
}
|
|
# add yours ...
|
|
}
|
|
|
|
|
|
def register_build_commands(subparsers: argparse._SubParsersAction) -> None:
|
|
build_parser = subparsers.add_parser(
|
|
"build",
|
|
help="Build related commands",
|
|
formatter_class=RichHelp,
|
|
)
|
|
build_subparsers = build_parser.add_subparsers(dest="build_command", required=True)
|
|
overview = "\n".join(
|
|
f" {name:12} {spec.get('help', '')}" for name, spec in _TARGETS.items()
|
|
)
|
|
external_parser = build_subparsers.add_parser(
|
|
"external",
|
|
help="Build external targets",
|
|
description="Build third-party targets.\n\nAvailable targets:\n" + overview,
|
|
formatter_class=RichHelp,
|
|
)
|
|
register_target_commands_and_runner(external_parser, _TARGETS)
|