wehub-resource-sync
1012 lines
22 KiB
Python
1012 lines
22 KiB
Python
load("@rules_python//python:defs.bzl", "py_library", "py_test")
|
|
load("//bazel:python.bzl", "doctest")
|
|
|
|
doctest(
|
|
name = "py_doctest[train]",
|
|
size = "large",
|
|
env = {
|
|
"RAY_TRAIN_V2_ENABLED": "1",
|
|
"TF_USE_LEGACY_KERAS": "1",
|
|
},
|
|
files = glob(
|
|
["**/*.py"],
|
|
exclude = [
|
|
"examples/**",
|
|
"tests/**",
|
|
"horovod/**", # CI do not have horovod installed
|
|
"mosaic/**", # CI do not have mosaicml installed
|
|
# GPU tests
|
|
"tensorflow/tensorflow_trainer.py",
|
|
"_internal/session.py",
|
|
"context.py",
|
|
],
|
|
),
|
|
tags = ["team:ml"],
|
|
)
|
|
|
|
doctest(
|
|
name = "py_doctest[train-gpu]",
|
|
size = "large",
|
|
# TODO: [V2] Migrate
|
|
env = {"RAY_TRAIN_V2_ENABLED": "0"},
|
|
files = [
|
|
"_internal/session.py",
|
|
"context.py",
|
|
"tensorflow/tensorflow_trainer.py",
|
|
],
|
|
gpu = True,
|
|
tags = ["team:ml"],
|
|
)
|
|
|
|
# --------------------------------------------------------------------
|
|
# Tests from the python/ray/train/examples directory.
|
|
# Please keep these sorted alphabetically.
|
|
# --------------------------------------------------------------------
|
|
py_library(
|
|
name = "conftest",
|
|
srcs = ["tests/conftest.py"],
|
|
)
|
|
|
|
############ Experiment tracking examples start ############
|
|
|
|
# no credentials needed
|
|
py_test(
|
|
name = "lightning_exp_tracking_mlflow",
|
|
size = "small",
|
|
srcs = [
|
|
"examples/experiment_tracking/lightning_exp_tracking_mlflow.py",
|
|
"examples/experiment_tracking/lightning_exp_tracking_model_dl.py",
|
|
],
|
|
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
|
main = "examples/experiment_tracking/lightning_exp_tracking_mlflow.py",
|
|
tags = [
|
|
"exclusive",
|
|
"no_main",
|
|
"team:ml",
|
|
"train_v2",
|
|
],
|
|
deps = [":train_lib"],
|
|
)
|
|
|
|
py_test(
|
|
name = "lightning_exp_tracking_tensorboard",
|
|
size = "small",
|
|
srcs = [
|
|
"examples/experiment_tracking/lightning_exp_tracking_model_dl.py",
|
|
"examples/experiment_tracking/lightning_exp_tracking_tensorboard.py",
|
|
],
|
|
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
|
main = "examples/experiment_tracking/lightning_exp_tracking_tensorboard.py",
|
|
tags = [
|
|
"exclusive",
|
|
"no_main",
|
|
"team:ml",
|
|
"train_v2",
|
|
],
|
|
deps = [":train_lib"],
|
|
)
|
|
|
|
py_test(
|
|
name = "torch_exp_tracking_mlflow",
|
|
size = "medium",
|
|
srcs = ["examples/experiment_tracking/torch_exp_tracking_mlflow.py"],
|
|
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
|
main = "examples/experiment_tracking/torch_exp_tracking_mlflow.py",
|
|
tags = [
|
|
"exclusive",
|
|
"no_main",
|
|
"team:ml",
|
|
"train_v2",
|
|
],
|
|
deps = [":train_lib"],
|
|
)
|
|
|
|
# credentials needed
|
|
py_test(
|
|
name = "lightning_exp_tracking_wandb",
|
|
size = "medium",
|
|
srcs = [
|
|
"examples/experiment_tracking/lightning_exp_tracking_model_dl.py",
|
|
"examples/experiment_tracking/lightning_exp_tracking_wandb.py",
|
|
],
|
|
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
|
main = "examples/experiment_tracking/lightning_exp_tracking_wandb.py",
|
|
tags = [
|
|
"exclusive",
|
|
"needs_credentials",
|
|
"no_main",
|
|
"team:ml",
|
|
"train_v2",
|
|
],
|
|
deps = [":train_lib"],
|
|
)
|
|
|
|
py_test(
|
|
name = "lightning_exp_tracking_comet",
|
|
size = "medium",
|
|
srcs = [
|
|
"examples/experiment_tracking/lightning_exp_tracking_comet.py",
|
|
"examples/experiment_tracking/lightning_exp_tracking_model_dl.py",
|
|
],
|
|
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
|
main = "examples/experiment_tracking/lightning_exp_tracking_comet.py",
|
|
tags = [
|
|
"exclusive",
|
|
"needs_credentials",
|
|
"no_main",
|
|
"team:ml",
|
|
"train_v2",
|
|
],
|
|
deps = [":train_lib"],
|
|
)
|
|
|
|
py_test(
|
|
name = "torch_exp_tracking_wandb",
|
|
size = "medium",
|
|
srcs = ["examples/experiment_tracking/torch_exp_tracking_wandb.py"],
|
|
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
|
main = "examples/experiment_tracking/torch_exp_tracking_wandb.py",
|
|
tags = [
|
|
"exclusive",
|
|
"needs_credentials",
|
|
"no_main",
|
|
"team:ml",
|
|
"train_v2",
|
|
],
|
|
deps = [":train_lib"],
|
|
)
|
|
|
|
############ Experiment tracking examples end ############
|
|
|
|
py_test(
|
|
name = "tensorflow_quick_start",
|
|
size = "medium",
|
|
srcs = ["examples/tf/tensorflow_quick_start.py"],
|
|
env = {
|
|
"RAY_TRAIN_V2_ENABLED": "1",
|
|
"TF_USE_LEGACY_KERAS": "1",
|
|
},
|
|
main = "examples/tf/tensorflow_quick_start.py",
|
|
tags = [
|
|
"exclusive",
|
|
"team:ml",
|
|
"train_v2",
|
|
],
|
|
deps = [":train_lib"],
|
|
)
|
|
|
|
py_test(
|
|
name = "torch_quick_start",
|
|
size = "medium",
|
|
srcs = ["examples/pytorch/torch_quick_start.py"],
|
|
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
|
main = "examples/pytorch/torch_quick_start.py",
|
|
tags = [
|
|
"exclusive",
|
|
"team:ml",
|
|
"train_v2",
|
|
],
|
|
deps = [":train_lib"],
|
|
)
|
|
|
|
# Formerly AIR examples
|
|
|
|
py_test(
|
|
name = "distributed_sage_example",
|
|
size = "small",
|
|
srcs = ["examples/pytorch_geometric/distributed_sage_example.py"],
|
|
args = [
|
|
"--use-gpu",
|
|
"--num-workers=2",
|
|
"--epochs=1",
|
|
"--dataset=fake",
|
|
],
|
|
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
|
main = "examples/pytorch_geometric/distributed_sage_example.py",
|
|
tags = [
|
|
"exclusive",
|
|
"team:ml",
|
|
"train_v2_gpu",
|
|
],
|
|
deps = [":train_lib"],
|
|
)
|
|
|
|
py_test(
|
|
name = "horovod_pytorch_example",
|
|
size = "small",
|
|
srcs = ["examples/horovod/horovod_pytorch_example.py"],
|
|
args = ["--num-epochs=1"],
|
|
env = {"RAY_TRAIN_V2_ENABLED": "0"},
|
|
tags = [
|
|
"exclusive",
|
|
"manual",
|
|
"team:ml",
|
|
],
|
|
deps = [":train_lib"],
|
|
)
|
|
|
|
py_test(
|
|
name = "horovod_tune_example",
|
|
size = "small",
|
|
srcs = ["examples/horovod/horovod_tune_example.py"],
|
|
args = ["--smoke-test"],
|
|
env = {"RAY_TRAIN_V2_ENABLED": "0"},
|
|
tags = [
|
|
"exclusive",
|
|
"manual",
|
|
"team:ml",
|
|
],
|
|
deps = [":train_lib"],
|
|
)
|
|
|
|
py_test(
|
|
name = "tensorflow_regression_example",
|
|
size = "medium",
|
|
srcs = ["examples/tf/tensorflow_regression_example.py"],
|
|
args = ["--smoke-test"],
|
|
env = {
|
|
"RAY_TRAIN_V2_ENABLED": "1",
|
|
"TF_USE_LEGACY_KERAS": "1",
|
|
},
|
|
main = "examples/tf/tensorflow_regression_example.py",
|
|
tags = [
|
|
"exclusive",
|
|
"team:ml",
|
|
],
|
|
deps = [":train_lib"],
|
|
)
|
|
|
|
# This is tested in test_examples!
|
|
# py_test(
|
|
# name = "tensorflow_mnist_example",
|
|
# size = "medium",
|
|
# main = "examples/tf/tensorflow_mnist_example.py",
|
|
# srcs = ["examples/tf/tensorflow_mnist_example.py"],
|
|
# tags = ["team:ml", "exclusive"],
|
|
# deps = [":train_lib"],
|
|
# args = ["--smoke-test"]
|
|
# )
|
|
|
|
# This is tested in test_examples!
|
|
# py_test(
|
|
# name = "torch_fashion_mnist_example",
|
|
# size = "medium",
|
|
# main = "examples/pytorch/torch_fashion_mnist_example.py",
|
|
# srcs = ["examples/pytorch/torch_fashion_mnist_example.py"],
|
|
# tags = ["team:ml", "exclusive"],
|
|
# deps = [":train_lib"],
|
|
# args = ["--smoke-test"]
|
|
# )
|
|
|
|
# This is tested in test_gpu_examples!
|
|
# py_test(
|
|
# name = "torch_fashion_mnist_example_gpu",
|
|
# size = "medium",
|
|
# main = "examples/pytorch/torch_fashion_mnist_example.py",
|
|
# srcs = ["examples/pytorch/torch_fashion_mnist_example.py"],
|
|
# tags = ["team:ml", "exclusive", "gpu"],
|
|
# deps = [":train_lib"],
|
|
# args = ["--use-gpu"]
|
|
# )
|
|
|
|
py_test(
|
|
name = "torch_regression_example",
|
|
size = "medium",
|
|
srcs = ["examples/pytorch/torch_regression_example.py"],
|
|
args = ["--smoke-test"],
|
|
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
|
main = "examples/pytorch/torch_regression_example.py",
|
|
tags = [
|
|
"exclusive",
|
|
"team:ml",
|
|
"train_v2",
|
|
],
|
|
deps = [":train_lib"],
|
|
)
|
|
|
|
# This is tested in test_examples!
|
|
# py_test(
|
|
# name = "torch_linear_example",
|
|
# size = "small",
|
|
# main = "examples/pytorch/torch_linear_example.py",
|
|
# srcs = ["examples/pytorch/torch_linear_example.py"],
|
|
# tags = ["team:ml", "exclusive"],
|
|
# deps = [":train_lib"],
|
|
# args = ["--smoke-test"]
|
|
# )
|
|
|
|
# --------------------------------------------------------------------
|
|
# Tests from the python/ray/train/tests directory.
|
|
# Please keep these sorted alphabetically.
|
|
# --------------------------------------------------------------------
|
|
|
|
py_test(
|
|
name = "test_torch_accelerate",
|
|
size = "large",
|
|
srcs = ["tests/test_torch_accelerate.py"],
|
|
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
|
tags = [
|
|
"exclusive",
|
|
"team:ml",
|
|
"train_v2_gpu",
|
|
],
|
|
deps = [
|
|
":conftest",
|
|
":train_lib",
|
|
],
|
|
)
|
|
|
|
py_test(
|
|
name = "test_api_migrations",
|
|
size = "small",
|
|
srcs = ["tests/test_api_migrations.py"],
|
|
# NOTE: This test explicitly tests V1 -> V2 migration.
|
|
env = {"RAY_TRAIN_V2_ENABLED": "0"},
|
|
tags = [
|
|
"exclusive",
|
|
"team:ml",
|
|
],
|
|
deps = [
|
|
":conftest",
|
|
":train_lib",
|
|
],
|
|
)
|
|
|
|
py_test(
|
|
name = "test_backend",
|
|
size = "large",
|
|
srcs = ["tests/test_backend.py"],
|
|
# NOTE: Relevant tests have been migrated to
|
|
# test_torch_trainer.py and test_worker_group.py
|
|
# TODO: [V2] There are still some accelerator integration tests left.
|
|
env = {"RAY_TRAIN_V2_ENABLED": "0"},
|
|
tags = [
|
|
"exclusive",
|
|
"team:ml",
|
|
],
|
|
deps = [
|
|
":conftest",
|
|
":train_lib",
|
|
],
|
|
)
|
|
|
|
py_test(
|
|
name = "test_base_trainer",
|
|
size = "medium",
|
|
srcs = ["tests/test_base_trainer.py"],
|
|
env = {"RAY_TRAIN_V2_ENABLED": "0"},
|
|
tags = [
|
|
"exclusive",
|
|
"team:ml",
|
|
],
|
|
deps = [
|
|
":conftest",
|
|
":train_lib",
|
|
],
|
|
)
|
|
|
|
py_test(
|
|
name = "test_base_worker_group",
|
|
size = "small",
|
|
srcs = ["tests/test_base_worker_group.py"],
|
|
tags = ["team:ml"],
|
|
deps = [":train_lib"],
|
|
)
|
|
|
|
py_test(
|
|
name = "test_checkpoint",
|
|
size = "small",
|
|
srcs = ["tests/test_checkpoint.py"],
|
|
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
|
tags = [
|
|
"exclusive",
|
|
"team:ml",
|
|
"train_v2",
|
|
],
|
|
deps = [":train_lib"],
|
|
)
|
|
|
|
py_test(
|
|
name = "test_checkpoint_manager",
|
|
size = "small",
|
|
srcs = ["tests/test_checkpoint_manager.py"],
|
|
# NOTE: This already has a V2 copy.
|
|
env = {"RAY_TRAIN_V2_ENABLED": "0"},
|
|
tags = [
|
|
"exclusive",
|
|
"team:ml",
|
|
],
|
|
deps = [":train_lib"],
|
|
)
|
|
|
|
py_test(
|
|
name = "test_data_parallel_trainer",
|
|
size = "medium",
|
|
srcs = ["tests/test_data_parallel_trainer.py"],
|
|
env = {"RAY_TRAIN_V2_ENABLED": "0"},
|
|
tags = [
|
|
"exclusive",
|
|
"team:ml",
|
|
],
|
|
deps = [":train_lib"],
|
|
)
|
|
|
|
py_test(
|
|
name = "test_data_parallel_trainer_checkpointing",
|
|
size = "medium",
|
|
srcs = ["tests/test_data_parallel_trainer_checkpointing.py"],
|
|
env = {"RAY_TRAIN_V2_ENABLED": "0"},
|
|
tags = [
|
|
"exclusive",
|
|
"team:ml",
|
|
],
|
|
deps = [":train_lib"],
|
|
)
|
|
|
|
py_test(
|
|
name = "test_examples",
|
|
size = "large",
|
|
srcs = ["tests/test_examples.py"],
|
|
env = {
|
|
"RAY_TRAIN_V2_ENABLED": "1",
|
|
"TF_USE_LEGACY_KERAS": "1",
|
|
},
|
|
tags = [
|
|
"exclusive",
|
|
"team:ml",
|
|
],
|
|
deps = [
|
|
":conftest",
|
|
":train_lib",
|
|
],
|
|
)
|
|
|
|
py_test(
|
|
name = "test_gpu",
|
|
size = "large",
|
|
srcs = ["tests/test_gpu.py"],
|
|
# NOTE: Migrated relevant tests to v2/tests/test_torch_gpu.py
|
|
env = {"RAY_TRAIN_V2_ENABLED": "0"},
|
|
tags = [
|
|
"exclusive",
|
|
"gpu",
|
|
"team:ml",
|
|
],
|
|
deps = [
|
|
":conftest",
|
|
":train_lib",
|
|
],
|
|
)
|
|
|
|
py_test(
|
|
name = "test_gpu_2",
|
|
size = "medium",
|
|
srcs = ["tests/test_gpu_2.py"],
|
|
# NOTE: Already covered by test_iter_torch_batches_gpu
|
|
env = {"RAY_TRAIN_V2_ENABLED": "0"},
|
|
tags = [
|
|
"exclusive",
|
|
"gpu",
|
|
"team:ml",
|
|
],
|
|
deps = [
|
|
":conftest",
|
|
":train_lib",
|
|
],
|
|
)
|
|
|
|
py_test(
|
|
name = "test_iter_torch_batches_gpu",
|
|
size = "medium",
|
|
srcs = ["tests/test_iter_torch_batches_gpu.py"],
|
|
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
|
tags = [
|
|
"data_integration",
|
|
"exclusive",
|
|
"team:ml",
|
|
"train_v2_gpu",
|
|
],
|
|
deps = [
|
|
":conftest",
|
|
":train_lib",
|
|
],
|
|
)
|
|
|
|
py_test(
|
|
name = "test_gpu_auto_transfer",
|
|
size = "medium",
|
|
srcs = ["tests/test_gpu_auto_transfer.py"],
|
|
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
|
tags = [
|
|
"exclusive",
|
|
"team:ml",
|
|
"train_v2_gpu",
|
|
],
|
|
deps = [
|
|
":conftest",
|
|
":train_lib",
|
|
],
|
|
)
|
|
|
|
py_test(
|
|
name = "test_gpu_examples",
|
|
size = "large",
|
|
srcs = ["tests/test_gpu_examples.py"],
|
|
env = {
|
|
"RAY_TRAIN_V2_ENABLED": "1",
|
|
"TF_USE_LEGACY_KERAS": "1",
|
|
},
|
|
tags = [
|
|
"exclusive",
|
|
"team:ml",
|
|
"train_v2_gpu",
|
|
],
|
|
deps = [
|
|
":conftest",
|
|
":train_lib",
|
|
],
|
|
)
|
|
|
|
py_test(
|
|
name = "test_torch_fsdp",
|
|
size = "small",
|
|
srcs = ["tests/test_torch_fsdp.py"],
|
|
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
|
tags = [
|
|
"exclusive",
|
|
"team:ml",
|
|
"train_v2_gpu",
|
|
],
|
|
deps = [":train_lib"],
|
|
)
|
|
|
|
py_test(
|
|
name = "test_horovod_trainer",
|
|
size = "large",
|
|
srcs = ["tests/test_horovod_trainer.py"],
|
|
env = {"RAY_TRAIN_V2_ENABLED": "0"},
|
|
tags = [
|
|
"exclusive",
|
|
"manual",
|
|
"team:ml",
|
|
],
|
|
deps = [":train_lib"],
|
|
)
|
|
|
|
py_test(
|
|
name = "test_lightgbm_trainer",
|
|
size = "medium",
|
|
srcs = ["tests/test_lightgbm_trainer.py"],
|
|
env = {"RAY_TRAIN_V2_ENABLED": "0"},
|
|
tags = [
|
|
"exclusive",
|
|
"team:ml",
|
|
],
|
|
deps = [":train_lib"],
|
|
)
|
|
|
|
py_test(
|
|
name = "test_torch_lightning_train",
|
|
size = "large",
|
|
srcs = ["tests/test_torch_lightning_train.py"],
|
|
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
|
tags = [
|
|
"exclusive",
|
|
"ptl_v1",
|
|
"team:ml",
|
|
"train_v2_gpu",
|
|
],
|
|
deps = [":train_lib"],
|
|
)
|
|
|
|
py_test(
|
|
name = "test_torch_transformers_train",
|
|
size = "large",
|
|
srcs = ["tests/test_torch_transformers_train.py"],
|
|
# NOTE: There's already a copy in V2.
|
|
env = {"RAY_TRAIN_V2_ENABLED": "0"},
|
|
tags = [
|
|
"exclusive",
|
|
"gpu",
|
|
"team:ml",
|
|
],
|
|
deps = [":train_lib"],
|
|
)
|
|
|
|
py_test(
|
|
name = "accelerate_torch_trainer",
|
|
size = "large",
|
|
srcs = ["examples/accelerate/accelerate_torch_trainer.py"],
|
|
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
|
tags = [
|
|
"exclusive",
|
|
"team:ml",
|
|
"train_v2_gpu",
|
|
],
|
|
deps = [":train_lib"],
|
|
)
|
|
|
|
py_test(
|
|
name = "accelerate_torch_trainer_no_raydata",
|
|
size = "large",
|
|
srcs = ["examples/accelerate/accelerate_torch_trainer_no_raydata.py"],
|
|
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
|
tags = [
|
|
"exclusive",
|
|
"team:ml",
|
|
"train_v2_gpu",
|
|
],
|
|
deps = [":train_lib"],
|
|
)
|
|
|
|
py_test(
|
|
name = "deepspeed_torch_trainer",
|
|
size = "large",
|
|
srcs = ["examples/deepspeed/deepspeed_torch_trainer.py"],
|
|
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
|
tags = [
|
|
"exclusive",
|
|
"team:ml",
|
|
"train_v2_gpu",
|
|
],
|
|
deps = [":train_lib"],
|
|
)
|
|
|
|
py_test(
|
|
name = "deepspeed_torch_trainer_no_raydata",
|
|
size = "large",
|
|
srcs = ["examples/deepspeed/deepspeed_torch_trainer_no_raydata.py"],
|
|
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
|
tags = [
|
|
"exclusive",
|
|
"team:ml",
|
|
"train_v2_gpu",
|
|
],
|
|
deps = [":train_lib"],
|
|
)
|
|
|
|
py_test(
|
|
name = "test_minimal",
|
|
size = "small",
|
|
srcs = ["tests/test_minimal.py"],
|
|
# TODO: [V2] The minimal test needs to install pydantic,
|
|
# which is a Train V2 dependency.
|
|
env = {"RAY_TRAIN_V2_ENABLED": "0"},
|
|
tags = [
|
|
"exclusive",
|
|
"minimal",
|
|
"team:ml",
|
|
],
|
|
deps = [":train_lib"],
|
|
)
|
|
|
|
py_test(
|
|
name = "test_new_persistence",
|
|
size = "large",
|
|
srcs = ["tests/test_new_persistence.py"],
|
|
# NOTE: There's already a copy in V2.
|
|
env = {"RAY_TRAIN_V2_ENABLED": "0"},
|
|
tags = [
|
|
"exclusive",
|
|
"team:ml",
|
|
],
|
|
deps = [
|
|
":conftest",
|
|
":train_lib",
|
|
],
|
|
)
|
|
|
|
py_test(
|
|
name = "test_result",
|
|
size = "medium",
|
|
srcs = ["tests/test_result.py"],
|
|
# NOTE: There's already a copy in V2.
|
|
env = {"RAY_TRAIN_V2_ENABLED": "0"},
|
|
tags = [
|
|
"exclusive",
|
|
"team:ml",
|
|
],
|
|
deps = [
|
|
":conftest",
|
|
":train_lib",
|
|
],
|
|
)
|
|
|
|
py_test(
|
|
name = "test_session",
|
|
size = "small",
|
|
srcs = ["tests/test_session.py"],
|
|
env = {"RAY_TRAIN_V2_ENABLED": "0"},
|
|
tags = [
|
|
"exclusive",
|
|
"team:ml",
|
|
],
|
|
deps = [
|
|
":conftest",
|
|
":train_lib",
|
|
],
|
|
)
|
|
|
|
py_test(
|
|
name = "test_storage",
|
|
size = "small",
|
|
srcs = ["tests/test_storage.py"],
|
|
env = {"RAY_TRAIN_V2_ENABLED": "0"},
|
|
tags = [
|
|
"exclusive",
|
|
"team:ml",
|
|
],
|
|
deps = [
|
|
":conftest",
|
|
":train_lib",
|
|
],
|
|
)
|
|
|
|
py_test(
|
|
name = "test_state",
|
|
size = "medium",
|
|
srcs = ["tests/test_state.py"],
|
|
env = {"RAY_TRAIN_V2_ENABLED": "0"},
|
|
tags = [
|
|
"exclusive",
|
|
"team:ml",
|
|
],
|
|
deps = [
|
|
":conftest",
|
|
":train_lib",
|
|
],
|
|
)
|
|
|
|
py_test(
|
|
name = "test_state_export",
|
|
size = "medium",
|
|
srcs = ["tests/test_state_export.py"],
|
|
env = {"RAY_TRAIN_V2_ENABLED": "0"},
|
|
tags = [
|
|
"exclusive",
|
|
"team:ml",
|
|
],
|
|
deps = [
|
|
":conftest",
|
|
":train_lib",
|
|
],
|
|
)
|
|
|
|
py_test(
|
|
name = "test_tensorflow_checkpoint",
|
|
size = "medium",
|
|
srcs = ["tests/test_tensorflow_checkpoint.py"],
|
|
env = {
|
|
"RAY_TRAIN_V2_ENABLED": "1",
|
|
"TF_USE_LEGACY_KERAS": "1",
|
|
},
|
|
tags = [
|
|
"exclusive",
|
|
"team:ml",
|
|
"train_v2",
|
|
],
|
|
deps = [":train_lib"],
|
|
)
|
|
|
|
py_test(
|
|
name = "test_tensorflow_trainer",
|
|
size = "medium",
|
|
srcs = ["tests/test_tensorflow_trainer.py"],
|
|
env = {
|
|
"RAY_TRAIN_V2_ENABLED": "0",
|
|
"TF_USE_LEGACY_KERAS": "1",
|
|
},
|
|
tags = [
|
|
"exclusive",
|
|
"team:ml",
|
|
],
|
|
deps = [":train_lib"],
|
|
)
|
|
|
|
py_test(
|
|
name = "test_torch_checkpoint",
|
|
size = "small",
|
|
srcs = ["tests/test_torch_checkpoint.py"],
|
|
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
|
tags = [
|
|
"exclusive",
|
|
"team:ml",
|
|
"train_v2",
|
|
],
|
|
deps = [
|
|
":conftest",
|
|
":train_lib",
|
|
],
|
|
)
|
|
|
|
py_test(
|
|
name = "test_torch_device_manager",
|
|
size = "medium",
|
|
srcs = ["tests/test_torch_device_manager.py"],
|
|
# TODO: Fix accelerator integrations and move over.
|
|
env = {"RAY_TRAIN_V2_ENABLED": "0"},
|
|
tags = [
|
|
"exclusive",
|
|
"gpu",
|
|
"team:ml",
|
|
],
|
|
deps = [
|
|
":conftest",
|
|
":train_lib",
|
|
],
|
|
)
|
|
|
|
py_test(
|
|
name = "test_torch_trainer",
|
|
size = "large",
|
|
srcs = ["tests/test_torch_trainer.py"],
|
|
env = {"RAY_TRAIN_V2_ENABLED": "0"},
|
|
tags = [
|
|
"exclusive",
|
|
"team:ml",
|
|
],
|
|
deps = [":train_lib"],
|
|
)
|
|
|
|
py_test(
|
|
name = "test_torch_utils",
|
|
size = "small",
|
|
srcs = ["tests/test_torch_utils.py"],
|
|
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
|
tags = [
|
|
"exclusive",
|
|
"team:ml",
|
|
"train_v2",
|
|
],
|
|
deps = [":train_lib"],
|
|
)
|
|
|
|
py_test(
|
|
name = "test_train_usage",
|
|
size = "medium",
|
|
srcs = ["tests/test_train_usage.py"],
|
|
env = {"RAY_TRAIN_V2_ENABLED": "0"},
|
|
tags = [
|
|
"exclusive",
|
|
"team:ml",
|
|
],
|
|
deps = [":train_lib"],
|
|
)
|
|
|
|
py_test(
|
|
name = "test_training_iterator",
|
|
size = "large",
|
|
srcs = ["tests/test_training_iterator.py"],
|
|
env = {
|
|
"RAY_TRAIN_V2_ENABLED": "0",
|
|
"TF_USE_LEGACY_KERAS": "1",
|
|
},
|
|
tags = [
|
|
"exclusive",
|
|
"team:ml",
|
|
],
|
|
deps = [":train_lib"],
|
|
)
|
|
|
|
py_test(
|
|
name = "test_tune",
|
|
size = "large",
|
|
srcs = ["tests/test_tune.py"],
|
|
env = {
|
|
"RAY_TRAIN_V2_ENABLED": "0",
|
|
"TF_USE_LEGACY_KERAS": "1",
|
|
},
|
|
tags = [
|
|
"exclusive",
|
|
"team:ml",
|
|
"tune",
|
|
],
|
|
deps = [
|
|
":conftest",
|
|
":train_lib",
|
|
],
|
|
)
|
|
|
|
py_test(
|
|
name = "test_e2e_wandb_integration",
|
|
size = "small",
|
|
srcs = ["tests/test_e2e_wandb_integration.py"],
|
|
env = {"RAY_TRAIN_V2_ENABLED": "0"},
|
|
tags = [
|
|
"exclusive",
|
|
"team:ml",
|
|
],
|
|
deps = [":train_lib"],
|
|
)
|
|
|
|
py_test(
|
|
name = "test_worker_group",
|
|
size = "medium",
|
|
srcs = ["tests/test_worker_group.py"],
|
|
env = {"RAY_TRAIN_V2_ENABLED": "0"},
|
|
tags = [
|
|
"exclusive",
|
|
"team:ml",
|
|
],
|
|
deps = [":train_lib"],
|
|
)
|
|
|
|
py_test(
|
|
name = "test_windows",
|
|
size = "small",
|
|
srcs = ["tests/test_windows.py"],
|
|
env = {"RAY_TRAIN_V2_ENABLED": "1"},
|
|
tags = [
|
|
"exclusive",
|
|
"team:ml",
|
|
],
|
|
deps = [":train_lib"],
|
|
)
|
|
|
|
py_test(
|
|
name = "test_xgboost_trainer",
|
|
size = "medium",
|
|
srcs = ["tests/test_xgboost_trainer.py"],
|
|
env = {"RAY_TRAIN_V2_ENABLED": "0"},
|
|
tags = [
|
|
"exclusive",
|
|
"team:ml",
|
|
],
|
|
deps = [":train_lib"],
|
|
)
|
|
|
|
py_test(
|
|
name = "test_trainer_restore",
|
|
size = "large",
|
|
srcs = ["tests/test_trainer_restore.py"],
|
|
env = {"RAY_TRAIN_V2_ENABLED": "0"},
|
|
tags = [
|
|
"exclusive",
|
|
"team:ml",
|
|
],
|
|
deps = [
|
|
":conftest",
|
|
":train_lib",
|
|
],
|
|
)
|
|
|
|
py_test(
|
|
name = "test_telemetry",
|
|
size = "medium",
|
|
srcs = ["tests/test_telemetry.py"],
|
|
env = {"RAY_TRAIN_V2_ENABLED": "0"},
|
|
tags = ["team:ml"],
|
|
deps = [":train_lib"],
|
|
)
|
|
|
|
py_test(
|
|
name = "test_train_head",
|
|
size = "small",
|
|
srcs = ["tests/test_train_head.py"],
|
|
env = {"RAY_TRAIN_V2_ENABLED": "0"},
|
|
tags = [
|
|
"exclusive",
|
|
"team:ml",
|
|
],
|
|
deps = [
|
|
":conftest",
|
|
":train_lib",
|
|
],
|
|
)
|
|
|
|
# This is a dummy test dependency that causes the above tests to be
|
|
# re-run if any of these files changes.
|
|
py_library(
|
|
name = "train_lib",
|
|
srcs = glob(
|
|
["**/*.py"],
|
|
exclude = ["tests/*.py"],
|
|
),
|
|
visibility = [
|
|
"//python/ray/air:__pkg__",
|
|
"//python/ray/air:__subpackages__",
|
|
"//python/ray/train:__pkg__",
|
|
"//python/ray/train:__subpackages__",
|
|
],
|
|
)
|