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__", ], )