diff --git a/pyproject.toml b/pyproject.toml index d9fb4502..c0cec29d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -6,6 +6,9 @@ build-backend = "scikit_build_core.build" name = "ark" version = "0.5.0" +[project.scripts] +arkprof = "ark.profiler:main" + [tool.scikit-build] cmake.version = ">=3.25" cmake.args = [] diff --git a/python/ark/profiler.py b/python/ark/profiler.py index da346cb7..23925245 100644 --- a/python/ark/profiler.py +++ b/python/ark/profiler.py @@ -1,6 +1,7 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. +import argparse import sys import time from typing import Optional, List @@ -69,9 +70,8 @@ def run( ) -if __name__ == "__main__": - import argparse - +def main(argv=None): + """CLI entry point for the ARK profiler.""" parser = argparse.ArgumentParser(description="ARK Profiler") parser.add_argument( "--iter", @@ -97,13 +97,15 @@ def run( parser.add_argument( "--plan", type=str, help="Path to the plan file", required=True ) - args = parser.parse_args() + args = parser.parse_args(argv) target_processor_groups = None if args.target_processor_groups is not None: - target_processor_groups = list( - map(int, args.target_processor_groups.split(",")) - ) + target_processor_groups = [ + int(s.strip()) + for s in args.target_processor_groups.split(",") + if s.strip() + ] plan = Plan.from_file(args.plan) profiler = Profiler(plan) @@ -113,3 +115,7 @@ def run( profile_processor_groups=args.profile_processor_groups, target_processor_groups=target_processor_groups, ) + + +if __name__ == "__main__": # pragma: no cover + main() diff --git a/python/unittest/test_profiler.py b/python/unittest/test_profiler.py index 17b836d2..271cadc1 100644 --- a/python/unittest/test_profiler.py +++ b/python/unittest/test_profiler.py @@ -1,9 +1,131 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. +import subprocess +import sys +from unittest.mock import patch, MagicMock + from common import ark, pytest_ark import pytest +try: + from ark import core as _ark_core # noqa: F401 + + _has_ark_core = True +except ImportError: + _has_ark_core = False + + +@pytest.mark.skipif( + not _has_ark_core, reason="native _ark_core extension not available" +) +def test_profiler_cli_help(): + """Test that `python -m ark.profiler --help` exits 0 and shows usage.""" + result = subprocess.run( + [sys.executable, "-m", "ark.profiler", "--help"], + capture_output=True, + text=True, + timeout=30, + ) + assert result.returncode == 0 + assert "ARK Profiler" in result.stdout + + +@pytest.mark.skipif( + not _has_ark_core, reason="native _ark_core extension not available" +) +def test_profiler_cli_missing_plan(): + """Test that `python -m ark.profiler` without --plan exits with code 2. + Validates that --plan is configured as a required argument.""" + result = subprocess.run( + [sys.executable, "-m", "ark.profiler"], + capture_output=True, + text=True, + timeout=30, + ) + assert result.returncode == 2 + assert "--plan" in result.stderr + + +def test_profiler_main_arg_parsing(): + """Test main() parses args and delegates to Profiler.run.""" + mock_plan = MagicMock() + mock_profiler = MagicMock() + + with patch("ark.profiler.Plan") as MockPlan, patch( + "ark.profiler.Profiler" + ) as MockProfiler: + MockPlan.from_file.return_value = mock_plan + MockProfiler.return_value = mock_profiler + + ark.profiler.main( + [ + "--plan", + "test.json", + "--iter", + "5", + "--loop_mode", + "--profile_processor_groups", + "--target_processor_groups", + "0,1", + ] + ) + + MockPlan.from_file.assert_called_once_with("test.json") + MockProfiler.assert_called_once_with(mock_plan) + mock_profiler.run.assert_called_once_with( + iter=5, + loop_mode=True, + profile_processor_groups=True, + target_processor_groups=[0, 1], + ) + + +def test_profiler_main_defaults(): + """Test main() uses default args when only --plan is given.""" + mock_plan = MagicMock() + mock_profiler = MagicMock() + + with patch("ark.profiler.Plan") as MockPlan, patch( + "ark.profiler.Profiler" + ) as MockProfiler: + MockPlan.from_file.return_value = mock_plan + MockProfiler.return_value = mock_profiler + + ark.profiler.main(["--plan", "plan.json"]) + + MockPlan.from_file.assert_called_once_with("plan.json") + MockProfiler.assert_called_once_with(mock_plan) + mock_profiler.run.assert_called_once_with( + iter=1000, + loop_mode=False, + profile_processor_groups=False, + target_processor_groups=None, + ) + + +def test_profiler_main_target_groups_whitespace(): + """Test that target_processor_groups handles whitespace and empty segments.""" + mock_plan = MagicMock() + mock_profiler = MagicMock() + + with patch("ark.profiler.Plan") as MockPlan, patch( + "ark.profiler.Profiler" + ) as MockProfiler: + MockPlan.from_file.return_value = mock_plan + MockProfiler.return_value = mock_profiler + + ark.profiler.main( + ["--plan", "t.json", "--target_processor_groups", "0, 1"] + ) + + mock_profiler.run.assert_called_once_with( + iter=1000, + loop_mode=False, + profile_processor_groups=False, + target_processor_groups=[0, 1], + ) + @pytest_ark() def test_profiler_non_loop_mode():