Support selected benchmarks

README.md

@@ -61,3 +61,42 @@ e.g.
 ### 2. GitHub Pull Request
 
 TODO
+
+## Command Arguments
+
+### 1. `--target` / `--baseline` / `--deps-list`
+
+See [Use cases - Single Package](#i-single-package) and [Use cases - Multiple Packages](#ii-multiple-packages).
+
+### 2. `--enable` / `--disable`
+
+Benchmarking always takes amount of time. In order to focus on the targets and reduce the time consumption of our benchmarking tool, the `--enable` and `--disable` options are used to specify **the parts to be included** and **the parts to be excluded** respectively.
+
+```shell
+> julia --project=benchmark benchmark/runbenchmarks-cli.jl  \
+>   --enable=<ENABLED_PARTS> \
+>   --disable=<DISABLED_PARTS> \
+>   --deps-list=<Dependencies List>
+```
+
+For specification, `Enabled Parts` and `Disabled Parts` have the same format, which is a single string that simulates an array, with each element separated by a comma.
+
+More precisely, the granularity of the element of `Enabled Parts` and `Disabled Parts` is currently at the file-level, which means now **our tool will recognize the name of each file in `benchmark/benchmark` before benchmarking** and **each element in `Enabled Parts` and `Disable Parts` should be exactly the name of those files**.
+
+`--enable` is used to specify the files that should be included, and by default (`--enable` not specified) all files in the `benchmark/benchmark` are included. `--disable` is used to specify the files that should be excluded, and the default value is an empty string.
+
+> I don't recommend using `--enable` and `--disable` at the same time. But if you do, `--disable` takes priority over `--enable`.
+> 
+> e.g. if `--enable` is set "flux,nnlib" while `--disable` is set "nnlib", only benchmarks in "benchmark/benchmark/flux.jl" will be executed.
+
+e.g.
+
+```shell
+> DEPS_LIST="https://github.com/FluxML/NNlib.jl#backports-0.8.21,https://github.com/skyleaworlder/NNlib.jl#dummy-benchmark-test;Flux,Flux@0.13.12"
+> # Only Flux and NNlib
+> julia --project=benchmark benchmark/runbenchmarks-cli.jl --enable="flux,nnlib" --deps-list=$DEPS_LIST
+> # All benchmarks except Flux and NNlib
+> julia --project=benchmark benchmark/runbenchmarks-cli.jl --disable="flux,nnlib" --deps-list=$DEPS_LIST
+> # Only Flux
+> julia --project=benchmark benchmark/runbenchmarks-cli.jl --enable="flux,nnlib" --disable="nnlib" --deps-list=$DEPS_LIST
+```

benchmark/benchmarks.jl

@@ -1,8 +1,12 @@
 using BenchmarkTools
 using Random
 
+foreach(println, ENV) # to check environment variables
+
 Random.seed!(1234567890)
 const SUITE = BenchmarkGroup()
 
-include("benchmark/nnlib.jl")
-include("benchmark/flux.jl")
+get(ENV, "FLUXML_BENCHMARK_NNLIB", "false") == "true" &&
+    include("benchmark/nnlib.jl")
+get(ENV, "FLUXML_BENCHMARK_FLUX", "false") == "true" &&
+    include("benchmark/flux.jl")

benchmark/runbenchmarks-cli.jl

@@ -4,9 +4,13 @@ using Pkg
 
 Pkg.develop(PackageSpec(path = ENV["PWD"]))
 using FluxMLBenchmarks
+
 parsed_args = parse_commandline()
 deps_list = parsed_args["deps-list"]
 baseline_fluxml_deps, target_fluxml_deps = parse_deps_list(deps_list)
+enable_arg = parsed_args["enable"]
+disable_arg = parsed_args["disable"]
+enabled_benchmarks = parse_enabled_benchmarks(enable_arg, disable_arg)
 
 setup_fluxml_env(baseline_fluxml_deps)
 
@@ -18,7 +22,10 @@ using PkgBenchmark
 group_baseline = benchmarkpkg(
     dirname(@__DIR__),
     BenchmarkConfig(
-        env = Dict("JULIA_NUM_THREADS" => get(ENV, "JULIA_NUM_THREADS", "1"))
+        env = merge(
+            Dict("JULIA_NUM_THREADS" => get(ENV, "JULIA_NUM_THREADS", "1")),
+            enabled_benchmarks
+        )
     ),
     resultfile = joinpath(@__DIR__, "result-baseline.json")
 )
@@ -40,7 +47,10 @@ using PkgBenchmark
 group_target = benchmarkpkg(
     dirname(@__DIR__),
     BenchmarkConfig(
-        env = Dict("JULIA_NUM_THREADS" => get(ENV, "JULIA_NUM_THREADS", "1"))
+        env = merge(
+            Dict("JULIA_NUM_THREADS" => get(ENV, "JULIA_NUM_THREADS", "1")),
+            enabled_benchmarks
+        )
     ),
     resultfile = joinpath(@__DIR__, "result-target.json"),
 )

benchmark/runbenchmarks-pr.jl

@@ -4,8 +4,13 @@ using Pkg
 
 Pkg.develop(PackageSpec(path = ENV["PWD"]))
 using FluxMLBenchmarks
+
 parsed_args = parse_commandline()
 
+enable_arg = parsed_args["enable"]
+disable_arg = parsed_args["disable"]
+enabled_benchmarks = parse_enabled_benchmarks(enable_arg, disable_arg)
+
 baseline_url = parsed_args["baseline"]
 setup_fluxml_env([baseline_url])
 
@@ -17,7 +22,10 @@ using PkgBenchmark
 group_baseline = benchmarkpkg(
     dirname(@__DIR__),
     BenchmarkConfig(
-        env = Dict("JULIA_NUM_THREADS" => get(ENV, "JULIA_NUM_THREADS", "1"))
+        env = merge(
+            Dict("JULIA_NUM_THREADS" => get(ENV, "JULIA_NUM_THREADS", "1")),
+            enabled_benchmarks
+        )
     ),
     resultfile = joinpath(@__DIR__, "result-baseline.json")
 )
@@ -40,7 +48,10 @@ using PkgBenchmark
 group_target = benchmarkpkg(
     dirname(@__DIR__),
     BenchmarkConfig(
-        env = Dict("JULIA_NUM_THREADS" => get(ENV, "JULIA_NUM_THREADS", "1"))
+        env = merge(
+            Dict("JULIA_NUM_THREADS" => get(ENV, "JULIA_NUM_THREADS", "1")),
+            enabled_benchmarks
+        )
     ),
     resultfile = joinpath(@__DIR__, "result-target.json"),
 )

src/FluxMLBenchmarks.jl

@@ -4,7 +4,9 @@ module FluxMLBenchmarks
 include("env_utils.jl")
 export Dependency, get_name, init_dependencies,
     parse_commandline, parse_deps_list,
-    setup_fluxml_env, teardown
+    parse_enabled_benchmarks,
+    setup_fluxml_env, teardown,
+    BENCHMARK_PKG_PATH, BENCHMARK_FILES_PATH, FLUXML_AVAILABLE_BENCHMARKS
 
 include("judge_utils.jl")
 export markdown_report, display_markdown_report

src/env_utils.jl

@@ -6,7 +6,7 @@ using URIParser
 BENCHMARK_PKG_PATH means the relative path of benchmark folder,
 which should be changed if the benchmark code is moved elsewhere.
 """
-const BENCHMARK_PKG_PATH = "./benchmark/"
+const BENCHMARK_PKG_PATH = "./benchmark"
 
 """
 BENCHMARK_BASIC_DEPS mean the dependencies required to be installed before
@@ -27,6 +27,21 @@ const FLUXML_PKGS = [
     "Functors", "ZygoteRules", "IRTools", "MacroTools"
 ]
 
+"""
+BENCHMARK_FILES_PATH means the folder containing benchmark files.
+FLUXML_AVAILABLE_BENCHMARKS is a vector, each element of it means
+an available benchmark file under BENCHMARK_FILES_PATH.
+"""
+const BENCHMARK_FILES_PATH = "$(BENCHMARK_PKG_PATH)/benchmark"
+const FLUXML_AVAILABLE_BENCHMARKS = filter(
+    !isnothing,
+    map(readdir(BENCHMARK_FILES_PATH)) do file_name
+        if (m = match(r"(.*?).jl$", file_name)) !== nothing
+            string(m.captures[1])
+        end
+    end
+)
+
 
 """
     Dependency
@@ -177,6 +192,18 @@ About url passed to Pkg.add, see https://pkgdocs.julialang.org/v1/managing-packa
 function parse_commandline()
     s = ArgParseSettings()
     @add_arg_table! s begin
+        "--enable"
+            help = "Specified benchmark sections to execute.
+                    e.g. flux,nnlib,optimisers
+                    By default, all benchmarks are enabled."
+            action = :store_arg
+            default = reduce((x,y) -> "$x,$y", FLUXML_AVAILABLE_BENCHMARKS)
+        "--disable"
+            help = "Specified benchmark sections not to execute,
+                    e.g. nnlib,flux
+                    no benchmarks are disabled by default."
+            action = :store_arg
+            default = ""
         "--target"
             help = "Repo URL to use as target. No default value.
                     e.g. https://github.com/FluxML/NNlib.jl#segfault"
@@ -284,3 +311,40 @@ function teardown()
     println("pwd: $pwd")
     cd(pwd)
 end
+
+
+"""
+    parse_enabled_benchmarks(enable_cmd_arg::String, disable_cmd_arg::String)
+
+is used to parses command-line arguments to determine the enabled benchmarks.
+Return a Dict as a part of environment variables, which will be used in BenchmarkConfig.
+
+* enable_cmd_arg: A string containing a comma-separated list of enabled benchmarks.
+* disable_cmd_arg: A string containing a comma-separated list of disabled benchmarks.
+
+If `enable_cmd_arg` is not included in FLUXML_AVAILABLE_BENCHMARKS, it will be reported
+and ignored, which is similarly when `disable_cmd_arg` is not included in `enable_cmd_arg`.
+"""
+function parse_enabled_benchmarks(
+                                    enable_cmd_arg::String,
+                                    disable_cmd_arg::String
+                                )::Dict{String, Bool}
+    function remove_invalid(input, baseline)
+        invalid_benchmarks = [e for e in input if !(e in baseline)]
+        valid_benchmarks = [e for e in input if e in baseline]
+        for e in invalid_benchmarks
+            @warn "$e is not a part of benchmarks, please check the files in $BENCHMARK_FILES_PATH"
+        end
+        return valid_benchmarks
+    end
+
+    cmd_enable = filter(!isempty, map(string, split(enable_cmd_arg, ",")))
+    cmd_disable = filter(!isempty, map(string, split(disable_cmd_arg, ",")))
+    valid_cmd_enable = remove_invalid(cmd_enable, FLUXML_AVAILABLE_BENCHMARKS)
+    valid_cmd_disable = remove_invalid(cmd_disable, union(cmd_enable, FLUXML_AVAILABLE_BENCHMARKS))
+    remain_benchmark_files_name = setdiff(valid_cmd_enable, valid_cmd_disable)
+    return Dict(
+        "FLUXML_BENCHMARK_$(uppercase(fn))" => true
+        for fn in remain_benchmark_files_name
+    )
+end

test/env_utils_test.jl

@@ -1,5 +1,16 @@
 @testset "env_utils_test" begin
 
+    @testset "check benchmark location" begin
+        @test BENCHMARK_PKG_PATH == "./benchmark"
+        @test BENCHMARK_FILES_PATH == "./benchmark/benchmark"
+    end
+
+    @testset "check existed benchmarks" begin
+        @test (length(FLUXML_AVAILABLE_BENCHMARKS) == 2 &&
+               "flux" in FLUXML_AVAILABLE_BENCHMARKS &&
+               "nnlib" in FLUXML_AVAILABLE_BENCHMARKS)
+    end
+
     @testset "Dependency" begin
         flux_dep = Dependency("Flux")
         @test flux_dep.name == "Flux"
@@ -46,4 +57,36 @@
                target_deps[2].name == "Flux" &&
                target_deps[2].version == "0.13.12")
     end
+
+    @testset "parse enabled benchmarks" begin
+        default_enable = reduce((x,y) -> "$x,$y", FLUXML_AVAILABLE_BENCHMARKS)
+        default_disable = ""
+
+        eb0 = parse_enabled_benchmarks(default_enable, default_disable)
+        @test length(eb0) == length(FLUXML_AVAILABLE_BENCHMARKS)
+
+        disable0 = "nnlib,flux"
+        eb1 = parse_enabled_benchmarks(default_enable, disable0)
+        @test length(eb1) == length(FLUXML_AVAILABLE_BENCHMARKS) - 2
+
+        enable0 = "flux,nnlib,zygote"
+        eb2 = parse_enabled_benchmarks(enable0, default_disable)
+        @test (length(eb2) == 2 &&
+               get(eb2, "FLUXML_BENCHMARK_FLUX", false) &&
+               get(eb2, "FLUXML_BENCHMARK_NNLIB", false) &&
+               !get(eb2, "FLUXML_BENCHMARK_ZYGOTE", false))
+
+        enable1 = "flux,nnlib,zygote"
+        disable1 = "zygote,flux"
+        eb3 = parse_enabled_benchmarks(enable1, disable1)
+        @test (length(eb3) == 1 &&
+               !get(eb3, "FLUXML_BENCHMARK_FLUX", false) &&
+               !get(eb3, "FLUXML_BENCHMARK_ZYGOTE", false) &&
+               get(eb3, "FLUXML_BENCHMARK_NNLIB", false))
+
+        disable2 = "flux,not_existed_package,unknown_package"
+        eb4 = parse_enabled_benchmarks(default_enable, disable2)
+        @test (length(eb4) == length(FLUXML_AVAILABLE_BENCHMARKS) - 1 &&
+               !get(eb3, "FLUXML_BENCHMARK_FLUX", false))
+    end
 end