load("//bazel:python.bzl", "py_test_run_all_notebooks")

filegroup(
    name = "tune_examples",
    srcs = glob(["*.ipynb"]),
    visibility = ["//doc:__subpackages__"],
)

# --------------------------------------------------------------------
# Test all doc/source/tune/examples notebooks.
# --------------------------------------------------------------------

# pbt_ppo_example.ipynb is not tested right now due to large resource
# requirements

py_test_run_all_notebooks(
    size = "medium",
    include = ["*.ipynb"],
    data = ["//doc/source/tune/examples:tune_examples"],
    exclude = [
        "pbt_ppo_example.ipynb",
        "tune-xgboost.ipynb",
        "pbt_transformers.ipynb",  # Transformers uses legacy Tune APIs.
        "tune-aim.ipynb",  # CI does not have `aim`
        "bohb_example.ipynb",  # CI does not have bohb requirements
   ],
    tags = [
        "exclusive",
        "team:ml",
    ],
    env = {"RAY_TRAIN_V2_ENABLED": "1"},
)

# GPU tests
py_test_run_all_notebooks(
    size = "large",
    include = ["tune-xgboost.ipynb"],
    data = ["//doc/source/tune/examples:tune_examples"],
    exclude = [],
    tags = [
        "exclusive",
        "gpu",
        "team:ml",
    ],
    env = {"RAY_TRAIN_V2_ENABLED": "1"},
)

filegroup(
    name = "tune_examples_ci_configs",
    srcs = glob([
        "**/ci/aws.yaml",
        "**/ci/gce.yaml",
    ]),
    visibility = ["//doc:__pkg__"],
)