diff --git a/.ci/lumen_cli/README.md b/.ci/lumen_cli/README.md new file mode 100644 index 000000000000..a0bb8b19a000 --- /dev/null +++ b/.ci/lumen_cli/README.md @@ -0,0 +1,31 @@ +# 🔧 Lumen_cli +A Python CLI tool for building and testing PyTorch-based components, using a YAML configuration file for structured, repeatable workflows. + + +## Features +- **Build** + - external projects (e.g. vLLM) + +## 📦 Installation +at the root of the pytorch repo +```bash +pip install -e .ci/lumen_cli +``` + +## Run the cli tool +The cli tool must be used at root of pytorch repo, as example to run build external vllm: +```bash +python -m cli.run build external vllm +``` +this will run the build steps with default behaviour for vllm project. + +to see help messages, run +```bash +python3 -m cli.run --help +``` + +## Add customized external build logics +To add a new external build, for instance, add a new external build logics: +1. create the build function in cli/lib folder +2. register your target and the main build function at EXTERNAL_BUILD_TARGET_DISPATCH in `cli/build_cli/register_build.py` +3. [optional] create your ci config file in .github/ci_configs/${EXTERNAL_PACKAGE_NAME}.yaml diff --git a/.ci/lumen_cli/cli/build_cli/__init__.py b/.ci/lumen_cli/cli/build_cli/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/.ci/lumen_cli/cli/build_cli/register_build.py b/.ci/lumen_cli/cli/build_cli/register_build.py new file mode 100644 index 000000000000..eb12753933e0 --- /dev/null +++ b/.ci/lumen_cli/cli/build_cli/register_build.py @@ -0,0 +1,41 @@ +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) diff --git a/.ci/lumen_cli/cli/lib/__init__.py b/.ci/lumen_cli/cli/lib/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/.ci/lumen_cli/cli/lib/common/cli_helper.py b/.ci/lumen_cli/cli/lib/common/cli_helper.py new file mode 100644 index 000000000000..82e46a2f0e61 --- /dev/null +++ b/.ci/lumen_cli/cli/lib/common/cli_helper.py @@ -0,0 +1,92 @@ +""" +Cli Argparser Utility helpers for CLI tasks. + +""" + +import argparse +from abc import ABC, abstractmethod + + +try: + from typing import Any, Callable, Required, TypedDict # Python 3.11+ +except ImportError: + from typing import Any, Callable, TypedDict + + from typing_extensions import Required # Fallback for Python <3.11 + + +class BaseRunner(ABC): + def __init__(self, args: Any) -> None: + self.args = args + + @abstractmethod + def run(self) -> None: + """runs main logics, required""" + + +# Pretty help: keep newlines + show defaults +class RichHelp( + argparse.ArgumentDefaultsHelpFormatter, argparse.RawDescriptionHelpFormatter +): + pass + + +class TargetSpec(TypedDict, total=False): + # Registry entry for a CLI subcommand. + # + # Keys: + # runner – class implementing BaseRunner; will be constructed + # with the parsed argparse Namespace. + # help – short help text for the subparser (shows in parent help). + # description – long help text (falls back to runner.__doc__ if missing). + # add_arguments – optional function to add subparser-specific CLI args. + + runner: Required[type[BaseRunner]] + help: str + description: str + add_arguments: Callable[[argparse.ArgumentParser], None] + + +def register_target_commands_and_runner( + parser: argparse.ArgumentParser, + target_specs: dict[str, TargetSpec], + common_args: Callable[[argparse.ArgumentParser], None] = lambda _: None, + placeholder_name: str = "target", +) -> None: + """ + Given an argparse parser and a mapping of target names → TargetSpec, + register each target as a subcommand and wire it to its runner class. + - Creates a subparser for each target name. + - description: defaults to TargetSpec['description'] or runner.__doc__. + - description has higher priority than runner.__doc__. + - add_arguments[Optional]: add target-specific CLI args, if 'add_arguments' exists + - common_args[Optional]: add shared CLI args to each target parser. + - Sets parser defaults: + func: lambda that constructs the runner with parsed args + and calls its .run(). + _runner_class: stored runner class for introspection/testing. + """ + targets = parser.add_subparsers( + dest=placeholder_name, + required=True, + metavar="{" + ",".join(target_specs.keys()) + "}", + ) + + for name, spec in target_specs.items(): + desc = (spec.get("description") or (spec["runner"].__doc__ or "")).strip() + + p = targets.add_parser( + name, + help=spec.get("help", ""), + description=desc, + formatter_class=RichHelp, + ) + p.set_defaults( + func=lambda args, _cls=spec["runner"]: _cls(args).run(), + _runner_class=spec["runner"], + ) + + if "add_arguments" in spec and callable(spec["add_arguments"]): + spec["add_arguments"](p) + if common_args: + common_args(p) diff --git a/.ci/lumen_cli/cli/lib/common/logger.py b/.ci/lumen_cli/cli/lib/common/logger.py new file mode 100644 index 000000000000..7a638206d931 --- /dev/null +++ b/.ci/lumen_cli/cli/lib/common/logger.py @@ -0,0 +1,14 @@ +""" +Logger Utility helpers for CLI tasks. +""" + +import logging +import sys + + +def setup_logging(level: int = logging.INFO): + logging.basicConfig( + level=level, + format="%(asctime)s [%(levelname)s] %(name)s: %(message)s", + stream=sys.stdout, + ) diff --git a/.ci/lumen_cli/cli/lib/common/path_helper.py b/.ci/lumen_cli/cli/lib/common/path_helper.py new file mode 100644 index 000000000000..f7c07bb3d389 --- /dev/null +++ b/.ci/lumen_cli/cli/lib/common/path_helper.py @@ -0,0 +1,9 @@ +""" +File And Path Utility helpers for CLI tasks. +""" + +import logging + + +logger = logging.getLogger(__name__) +# TODO(elainewy): Add path_helper utils diff --git a/.ci/lumen_cli/cli/lib/common/utils.py b/.ci/lumen_cli/cli/lib/common/utils.py new file mode 100644 index 000000000000..407de0af22e8 --- /dev/null +++ b/.ci/lumen_cli/cli/lib/common/utils.py @@ -0,0 +1,9 @@ +""" +General Utility helpers for CLI tasks. +""" + +import logging + + +logger = logging.getLogger(__name__) +# TODO(elainewy): Add common utils diff --git a/.ci/lumen_cli/cli/lib/core/vllm.py b/.ci/lumen_cli/cli/lib/core/vllm.py new file mode 100644 index 000000000000..3a11e2db2d15 --- /dev/null +++ b/.ci/lumen_cli/cli/lib/core/vllm.py @@ -0,0 +1,15 @@ +import logging + +from cli.lib.common.cli_helper import BaseRunner + + +logger = logging.getLogger(__name__) + + +class VllmBuildRunner(BaseRunner): + """ + Build vllm whels in ci + """ + + def run(self): + logger.info("Running vllm build") diff --git a/.ci/lumen_cli/cli/run.py b/.ci/lumen_cli/cli/run.py new file mode 100644 index 000000000000..7b91858c91ce --- /dev/null +++ b/.ci/lumen_cli/cli/run.py @@ -0,0 +1,42 @@ +# main.py + +import argparse +import logging + +from cli.build_cli.register_build import register_build_commands +from cli.lib.common.logger import setup_logging + + +logger = logging.getLogger(__name__) + + +def main(): + # Define top-level parser + parser = argparse.ArgumentParser(description="Lumos CLI") + subparsers = parser.add_subparsers(dest="command", required=True) + # Add top-level args + parser.add_argument( + "--config", required=False, help="Path to config file for build and test" + ) + parser.add_argument( + "--log-level", default="INFO", help="Log level (DEBUG, INFO, WARNING, ERROR)" + ) + + # registers second-level subcommands + register_build_commands(subparsers) + + # parse args after all options are registered + args = parser.parse_args() + + # setup global logging + setup_logging(getattr(logging, args.log_level.upper(), logging.INFO)) + logger.debug("Parsed args: %s", args) + + if hasattr(args, "func"): + args.func(args) + else: + parser.print_help() + + +if __name__ == "__main__": + main() diff --git a/.ci/lumen_cli/pyproject.toml b/.ci/lumen_cli/pyproject.toml new file mode 100644 index 000000000000..299b6cea2f60 --- /dev/null +++ b/.ci/lumen_cli/pyproject.toml @@ -0,0 +1,19 @@ +[project] +name = "lumen-ci" +version = "0.1.0" +dependencies = [ + "pytest==7.3.2" +] + +[tool.setuptools] +packages = ["cli"] + +[tool.setuptools.package-dir] +cli = "cli" + +[tool.ruff.lint] +# Enable preview mode for linting +preview = true + +# Now you can select your preview rules, like RUF048 +extend-select = ["RUF048"] diff --git a/.ci/lumen_cli/tests/test_app.py b/.ci/lumen_cli/tests/test_app.py new file mode 100644 index 000000000000..612bb7afe5d5 --- /dev/null +++ b/.ci/lumen_cli/tests/test_app.py @@ -0,0 +1,45 @@ +# tests/test_cli.py +import io +import sys +import unittest +from contextlib import redirect_stderr, redirect_stdout +from unittest.mock import patch + +from cli.run import main + + +class TestArgparseCLI(unittest.TestCase): + def test_cli_run_build_external(self): + test_args = ["cli.run", "build", "external", "vllm"] + + with patch.object(sys, "argv", test_args): + with self.assertLogs(level="INFO") as caplog: + # if argparse could exit on error, wrap in try/except SystemExit if needed + main() + + # stdout print from your CLI plumbing + # logs emitted inside your code (info/debug/error etc.) + logs_text = "\n".join(caplog.output) + self.assertIn("Running vllm build", logs_text) + + def test_build_help(self): + test_args = ["cli.run", "build", "--help"] + + with patch.object(sys, "argv", test_args): + stdout = io.StringIO() + stderr = io.StringIO() + + # --help always raises SystemExit(0) + with self.assertRaises(SystemExit) as cm: + with redirect_stdout(stdout), redirect_stderr(stderr): + main() + + self.assertEqual(cm.exception.code, 0) + + output = stdout.getvalue() + self.assertIn("usage", output) + self.assertIn("external", output) + + +if __name__ == "__main__": + unittest.main() diff --git a/.ci/lumen_cli/tests/test_cli_helper.py b/.ci/lumen_cli/tests/test_cli_helper.py new file mode 100644 index 000000000000..984054fdef86 --- /dev/null +++ b/.ci/lumen_cli/tests/test_cli_helper.py @@ -0,0 +1,120 @@ +import argparse +import io +import unittest +from contextlib import redirect_stderr +from unittest.mock import patch + +from cli.lib.common.cli_helper import ( + BaseRunner, + register_target_commands_and_runner, + RichHelp, + TargetSpec, +) + + +# ---- Dummy runners for unittests---- +class FooRunner(BaseRunner): + """Foo description from docstring.""" + + def run(self) -> None: # replaced by mock + pass + + +class BarRunner(BaseRunner): + def run(self) -> None: # replaced by mock + pass + + +def add_foo_args(p: argparse.ArgumentParser) -> None: + p.add_argument("--x", type=int, required=True, help="x value") + + +def common_args(p: argparse.ArgumentParser) -> None: + p.add_argument("--verbose", action="store_true", help="verbose flag") + + +def build_parser(specs: dict[str, TargetSpec]) -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(prog="app", formatter_class=RichHelp) + register_target_commands_and_runner( + parser=parser, + target_specs=specs, + common_args=common_args, + ) + return parser + + +def get_subparser( + parser: argparse.ArgumentParser, name: str +) -> argparse.ArgumentParser: + subparsers_action = next( + a + for a in parser._subparsers._group_actions # type: ignore[attr-defined] + if isinstance(a, argparse._SubParsersAction) + ) + return subparsers_action.choices[name] + + +class TestRegisterTargets(unittest.TestCase): + def test_metavar_lists_targets(self): + specs: dict[str, TargetSpec] = { + "foo": {"runner": FooRunner, "add_arguments": add_foo_args}, + "bar": {"runner": BarRunner}, + } + parser = build_parser(specs) + subparsers_action = next( + a + for a in parser._subparsers._group_actions # type: ignore[attr-defined] + if isinstance(a, argparse._SubParsersAction) + ) + self.assertEqual(subparsers_action.metavar, "{foo,bar}") + + def test_add_arguments_and_common_args_present(self): + specs: dict[str, TargetSpec] = { + "foo": {"runner": FooRunner, "add_arguments": add_foo_args}, + } + parser = build_parser(specs) + foo = get_subparser(parser, "foo") + help_text = foo.format_help() + self.assertIn("--x", help_text) + self.assertIn("--verbose", help_text) + + def test_runner_constructed_with_ns_and_run_called(self): + specs: dict[str, TargetSpec] = { + "foo": {"runner": FooRunner, "add_arguments": add_foo_args}, + } + parser = build_parser(specs) + + with ( + patch.object(FooRunner, "__init__", return_value=None) as mock_init, + patch.object(FooRunner, "run", return_value=None) as mock_run, + ): + ns = parser.parse_args(["foo", "--x", "3", "--verbose"]) + ns.func(ns) # set by register_target_commands_and_runner + # __init__ received the Namespace + self.assertEqual(mock_init.call_count, 1) + (called_ns,), _ = mock_init.call_args + self.assertIsInstance(called_ns, argparse.Namespace) + # run() called with no args + mock_run.assert_called_once_with() + + def test_runner_docstring_used_as_description_when_missing(self): + specs: dict[str, TargetSpec] = { + "foo": {"runner": FooRunner, "add_arguments": add_foo_args}, + } + parser = build_parser(specs) + foo = get_subparser(parser, "foo") + help_text = foo.format_help() + self.assertIn("Foo description from docstring.", help_text) + + def test_missing_target_raises_systemexit_with_usage(self): + specs: dict[str, TargetSpec] = {"foo": {"runner": FooRunner}} + parser = build_parser(specs) + buf = io.StringIO() + with self.assertRaises(SystemExit), redirect_stderr(buf): + parser.parse_args([]) + err = buf.getvalue() + self.assertIn("usage:", err) + + +if __name__ == "__main__": + unittest.main() diff --git a/.github/workflows/tools-unit-tests.yml b/.github/workflows/tools-unit-tests.yml new file mode 100644 index 000000000000..7b80d2b37c62 --- /dev/null +++ b/.github/workflows/tools-unit-tests.yml @@ -0,0 +1,41 @@ +name: test-scripts-and-ci-tools + +on: + push: + branches: + - main + paths: + - scripts/lumen_cli/** + - .github/workflows/tools-unit-tests.yml + pull_request: + paths: + - scripts/lumen_cli/** + - .github/workflows/tools-unit-tests.yml + +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ github.event_name == 'workflow_dispatch' }} + cancel-in-progress: true + +jobs: + torch-cli-unit-tests: + permissions: + contents: read + pull-requests: write + if: ${{ github.repository_owner == 'pytorch' }} + runs-on: ubuntu-latest + steps: + - name: Checkout pytorch + uses: pytorch/pytorch/.github/actions/checkout-pytorch@main + with: + submodules: true + fetch-depth: 0 + + - name: Run tests + continue-on-error: true + run: | + set -ex + python3 -m venv /tmp/venv + source /tmp/venv/bin/activate + python -m pip install --upgrade pip + pip install -e .ci/lumen_cli/ + pytest -v -s .ci/lumen_cli/tests/*