load("@rules_python//python:defs.bzl", "py_library", "py_test") load("//bazel:python.bzl", "doctest") doctest( name = "py_doctest[train_v2]", size = "large", env = { "RAY_TRAIN_V2_ENABLED": "1", "TF_USE_LEGACY_KERAS": "1", }, files = glob( ["**/*.py"], exclude = [ "tests/**", "horovod/**", "jax/**", "examples/pytorch/torchft_linear_example.py", ], ), tags = ["team:ml"], ) py_library( name = "conftest", srcs = ["tests/conftest.py"], ) py_test( name = "test_accelerator_utils", size = "small", srcs = ["tests/test_accelerator_utils.py"], env = {"RAY_TRAIN_V2_ENABLED": "1"}, tags = [ "exclusive", "team:ml", "train_v2", ], deps = [ ":conftest", "//:ray_lib", ], ) py_test( name = "test_async_checkpointing_validation", size = "large", srcs = ["tests/test_async_checkpointing_validation.py"], env = {"RAY_TRAIN_V2_ENABLED": "1"}, tags = [ "exclusive", "team:ml", "train_v2", ], deps = [ ":conftest", "//:ray_lib", ], ) py_test( name = "test_checkpoint_manager", size = "small", srcs = ["tests/test_checkpoint_manager.py"], env = {"RAY_TRAIN_V2_ENABLED": "1"}, tags = [ "exclusive", "team:ml", "train_v2", ], deps = [ ":conftest", "//:ray_lib", ], ) py_test( name = "test_config", size = "small", srcs = ["tests/test_config.py"], env = {"RAY_TRAIN_V2_ENABLED": "1"}, tags = [ "exclusive", "team:ml", "train_v2", ], deps = [ ":conftest", "//:ray_lib", ], ) py_test( name = "test_circular_imports", size = "small", srcs = ["tests/test_circular_imports.py"], env = {"RAY_TRAIN_V2_ENABLED": "1"}, tags = [ "exclusive", "team:ml", "train_v2", ], deps = [ ":conftest", "//:ray_lib", ], ) py_test( name = "test_circular_import_linter", size = "small", srcs = ["tests/test_circular_import_linter.py"], env = {"RAY_TRAIN_V2_ENABLED": "1"}, tags = [ "exclusive", "team:ml", "train_v2", ], deps = [ ":conftest", "//:ray_lib", ], ) py_test( name = "test_validation_manager", size = "small", srcs = ["tests/test_validation_manager.py"], env = {"RAY_TRAIN_V2_ENABLED": "1"}, tags = [ "exclusive", "team:ml", "train_v2", ], deps = [ ":conftest", "//:ray_lib", ], ) py_test( name = "test_collective", size = "medium", srcs = ["tests/test_collective.py"], env = {"RAY_TRAIN_V2_ENABLED": "1"}, tags = [ "exclusive", "team:ml", "train_v2", ], deps = [ ":conftest", "//:ray_lib", ], ) py_test( name = "test_callback_manager", size = "small", srcs = ["tests/test_callback_manager.py"], env = {"RAY_TRAIN_V2_ENABLED": "1"}, tags = [ "exclusive", "team:ml", "train_v2", ], deps = [ ":conftest", "//:ray_lib", ], ) py_test( name = "test_controller", size = "medium", srcs = ["tests/test_controller.py"], env = {"RAY_TRAIN_V2_ENABLED": "1"}, tags = [ "exclusive", "team:ml", "train_v2", ], deps = [ ":conftest", "//:ray_lib", ], ) py_test( name = "test_elastic_scaling_policy", size = "small", srcs = ["tests/test_elastic_scaling_policy.py"], env = {"RAY_TRAIN_V2_ENABLED": "1"}, tags = [ "exclusive", "team:ml", "train_v2", ], deps = [ ":conftest", "//:ray_lib", ], ) py_test( name = "test_elastic_e2e", size = "medium", srcs = ["tests/test_elastic_e2e.py"], env = {"RAY_TRAIN_V2_ENABLED": "1"}, tags = [ "exclusive", "team:ml", "train_v2", ], deps = [ ":conftest", "//:ray_lib", ], ) py_test( name = "test_controller_callback_behaviour", size = "medium", srcs = ["tests/test_controller_callback_behaviour.py"], env = {"RAY_TRAIN_V2_ENABLED": "1"}, tags = [ "exclusive", "team:ml", "train_v2", ], deps = [ ":conftest", "//:ray_lib", ], ) py_test( name = "test_data_integration", size = "large", srcs = ["tests/test_data_integration.py"], env = {"RAY_TRAIN_V2_ENABLED": "1"}, tags = [ "data_integration", "exclusive", "team:ml", "train_v2", ], deps = [ ":conftest", "//:ray_lib", ], ) py_test( name = "test_data_parallel_trainer", size = "large", srcs = ["tests/test_data_parallel_trainer.py"], env = {"RAY_TRAIN_V2_ENABLED": "1"}, tags = [ "exclusive", "team:ml", "train_v2", ], deps = [ ":conftest", "//:ray_lib", ], ) py_test( name = "test_report_fault_tolerance", size = "medium", srcs = ["tests/test_report_fault_tolerance.py"], env = {"RAY_TRAIN_V2_ENABLED": "1"}, tags = [ "exclusive", "team:ml", "train_v2", ], deps = [ ":conftest", "//:ray_lib", ], ) py_test( name = "test_data_resource_cleanup", size = "medium", srcs = ["tests/test_data_resource_cleanup.py"], env = {"RAY_TRAIN_V2_ENABLED": "1"}, tags = [ "data_integration", "exclusive", "team:ml", "train_v2", ], deps = [ ":conftest", "//:ray_lib", ], ) py_test( name = "test_dataset_manager", size = "medium", srcs = ["tests/test_dataset_manager.py"], env = {"RAY_TRAIN_V2_ENABLED": "1"}, tags = [ "data_integration", "exclusive", "team:ml", "train_v2", ], deps = [ ":conftest", "//:ray_lib", ], ) py_test( name = "test_env_callbacks", size = "small", srcs = ["tests/test_env_callbacks.py"], env = {"RAY_TRAIN_V2_ENABLED": "1"}, tags = [ "exclusive", "team:ml", "train_v2", ], deps = [ ":conftest", "//:ray_lib", ], ) py_test( name = "test_failure_policy", size = "small", srcs = ["tests/test_failure_policy.py"], env = {"RAY_TRAIN_V2_ENABLED": "1"}, tags = [ "exclusive", "team:ml", "train_v2", ], deps = [ ":conftest", "//:ray_lib", ], ) py_test( name = "test_jax_elastic_e2e", size = "medium", srcs = ["tests/test_jax_elastic_e2e.py"], env = {"RAY_TRAIN_V2_ENABLED": "1"}, tags = [ "exclusive", "team:ml", "train_v2", ], deps = [ ":conftest", "//:ray_lib", ], ) py_test( name = "test_jax_trainer", size = "medium", srcs = ["tests/test_jax_trainer.py"], env = {"RAY_TRAIN_V2_ENABLED": "1"}, tags = [ "exclusive", "team:ml", "train_v2", ], deps = [ ":conftest", "//:ray_lib", ], ) py_test( name = "test_jax_gpu", size = "medium", srcs = ["tests/test_jax_gpu.py"], env = {"RAY_TRAIN_V2_ENABLED": "1"}, tags = [ # Temporarily disabled in CI: Ray GPU CI is on CUDA 12.1, but JAX wheels # do not support CUDA 12.1 (and older JAX versions were removed from PyPI). # Tagging as 'manual' excludes it from Ray's default Bazel CI runs. "manual", "exclusive", "team:ml", "train_v2_gpu", ], deps = [ ":conftest", "//:ray_lib", ], ) py_test( name = "test_lightgbm_trainer", size = "small", srcs = ["tests/test_lightgbm_trainer.py"], env = {"RAY_TRAIN_V2_ENABLED": "1"}, tags = [ "exclusive", "team:ml", "train_v2", ], deps = [ ":conftest", "//:ray_lib", ], ) py_test( name = "test_lightning_integration", size = "medium", srcs = ["tests/test_lightning_integration.py"], env = {"RAY_TRAIN_V2_ENABLED": "1"}, tags = [ "exclusive", "team:ml", "train_v2", ], deps = [ ":conftest", "//:ray_lib", ], ) py_test( name = "test_logging", size = "medium", srcs = ["tests/test_logging.py"], env = {"RAY_TRAIN_V2_ENABLED": "1"}, tags = [ "exclusive", "team:ml", "train_v2", ], deps = [ ":conftest", "//:ray_lib", ], ) py_test( name = "test_metrics", size = "small", srcs = ["tests/test_metrics.py"], env = {"RAY_TRAIN_V2_ENABLED": "1"}, tags = [ "exclusive", "team:ml", "train_v2", ], deps = [ ":conftest", "//:ray_lib", ], ) py_test( name = "test_persistence", size = "medium", srcs = ["tests/test_persistence.py"], env = {"RAY_TRAIN_V2_ENABLED": "1"}, tags = [ "exclusive", "team:ml", "train_v2", ], deps = [ ":conftest", "//:ray_lib", ], ) py_test( name = "test_report_handler", size = "small", srcs = ["tests/test_report_handler.py"], env = {"RAY_TRAIN_V2_ENABLED": "1"}, tags = [ "exclusive", "team:ml", "train_v2", ], deps = [ ":conftest", "//:ray_lib", ], ) py_test( name = "test_placement_group_cleaner", size = "small", srcs = ["tests/test_placement_group_cleaner.py"], env = {"RAY_TRAIN_V2_ENABLED": "1"}, tags = [ "exclusive", "team:ml", "train_v2", ], deps = [ ":conftest", "//:ray_lib", ], ) py_test( name = "test_placement_group_handle", size = "medium", srcs = ["tests/test_placement_group_handle.py"], env = {"RAY_TRAIN_V2_ENABLED": "1"}, tags = [ "exclusive", "team:ml", "train_v2", ], deps = [ ":conftest", "//:ray_lib", ], ) py_test( name = "test_preemption_watcher", size = "small", srcs = ["tests/test_preemption_watcher.py"], env = {"RAY_TRAIN_V2_ENABLED": "1"}, tags = [ "exclusive", "team:ml", "train_v2", ], deps = [ ":conftest", "//:ray_lib", ], ) py_test( name = "test_result", size = "medium", srcs = ["tests/test_result.py"], env = {"RAY_TRAIN_V2_ENABLED": "1"}, tags = [ "exclusive", "team:ml", "train_v2", ], deps = [ ":conftest", "//:ray_lib", ], ) py_test( name = "test_scheduling", size = "medium", srcs = ["tests/test_scheduling.py"], env = {"RAY_TRAIN_V2_ENABLED": "1"}, tags = [ "exclusive", "team:ml", "train_v2", ], deps = [ ":conftest", "//:ray_lib", ], ) py_test( name = "test_serialization", size = "small", srcs = ["tests/test_serialization.py"], env = {"RAY_TRAIN_V2_ENABLED": "1"}, tags = [ "exclusive", "team:ml", "train_v2", ], deps = [ ":conftest", "//:ray_lib", ], ) py_test( name = "test_state", size = "medium", srcs = ["tests/test_state.py"], env = {"RAY_TRAIN_V2_ENABLED": "1"}, tags = [ "exclusive", "team:ml", "train_v2", ], deps = [ ":conftest", "//:ray_lib", ], ) py_test( name = "test_state_export", size = "medium", srcs = ["tests/test_state_export.py"], env = {"RAY_TRAIN_V2_ENABLED": "1"}, tags = [ "exclusive", "team:ml", "train_v2", ], deps = [ ":conftest", "//:ray_lib", ], ) py_test( name = "test_storage", size = "small", srcs = ["tests/test_storage.py"], env = {"RAY_TRAIN_V2_ENABLED": "1"}, tags = [ "exclusive", "team:ml", "train_v2", ], deps = [ ":conftest", "//:ray_lib", ], ) py_test( name = "test_sync_actor", size = "small", srcs = ["tests/test_sync_actor.py"], env = {"RAY_TRAIN_V2_ENABLED": "1"}, tags = [ "exclusive", "team:ml", "train_v2", ], deps = [ ":conftest", "//:ray_lib", ], ) py_test( name = "test_telemetry", size = "medium", srcs = ["tests/test_telemetry.py"], env = {"RAY_TRAIN_V2_ENABLED": "1"}, tags = [ "exclusive", "team:ml", "train_v2", ], deps = [ ":conftest", "//:ray_lib", ], ) py_test( name = "test_tensorflow_trainer", size = "medium", srcs = ["tests/test_tensorflow_trainer.py"], env = { "RAY_TRAIN_V2_ENABLED": "1", "TF_USE_LEGACY_KERAS": "1", }, tags = [ "exclusive", "team:ml", "train_v2", ], deps = [ ":conftest", "//:ray_lib", ], ) py_test( name = "test_thread_runner", size = "small", srcs = ["tests/test_thread_runner.py"], env = {"RAY_TRAIN_V2_ENABLED": "1"}, tags = [ "exclusive", "team:ml", "train_v2", ], deps = [ ":conftest", "//:ray_lib", ], ) py_test( name = "test_torch_gpu", size = "large", srcs = ["tests/test_torch_gpu.py"], env = {"RAY_TRAIN_V2_ENABLED": "1"}, tags = [ "exclusive", "team:ml", "train_v2_gpu", ], deps = [ ":conftest", "//:ray_lib", ], ) py_test( name = "test_torch_trainer", size = "medium", srcs = ["tests/test_torch_trainer.py"], env = {"RAY_TRAIN_V2_ENABLED": "1"}, tags = [ "exclusive", "team:ml", "torchft", "train_v2", ], deps = [ ":conftest", "//:ray_lib", ], ) py_test( name = "test_torch_transformers_train", size = "medium", srcs = ["tests/test_torch_transformers_train.py"], env = {"RAY_TRAIN_V2_ENABLED": "1"}, tags = [ "exclusive", "team:ml", "train_v2_gpu", ], deps = [ ":conftest", "//:ray_lib", ], ) py_test( name = "test_util", size = "medium", srcs = ["tests/test_util.py"], env = {"RAY_TRAIN_V2_ENABLED": "1"}, tags = [ "exclusive", "team:ml", "train_v2", ], deps = [ ":conftest", "//:ray_lib", ], ) py_test( name = "test_v2_api", size = "small", srcs = ["tests/test_v2_api.py"], env = {"RAY_TRAIN_V2_ENABLED": "1"}, tags = [ "exclusive", "team:ml", "train_v2", ], deps = [ ":conftest", "//:ray_lib", ], ) py_test( name = "test_worker", size = "small", srcs = ["tests/test_worker.py"], env = {"RAY_TRAIN_V2_ENABLED": "1"}, tags = [ "exclusive", "team:ml", "train_v2", ], deps = [ ":conftest", "//:ray_lib", ], ) py_test( name = "test_worker_group", size = "medium", srcs = ["tests/test_worker_group.py"], env = {"RAY_TRAIN_V2_ENABLED": "1"}, tags = [ "exclusive", "team:ml", "train_v2", ], deps = [ ":conftest", "//:ray_lib", ], ) py_test( name = "test_worker_group_poll_status", size = "small", srcs = ["tests/test_worker_group_poll_status.py"], env = {"RAY_TRAIN_V2_ENABLED": "1"}, tags = [ "exclusive", "team:ml", "train_v2", ], deps = [ ":conftest", "//:ray_lib", ], ) py_test( name = "test_xgboost_trainer", size = "small", srcs = ["tests/test_xgboost_trainer.py"], env = {"RAY_TRAIN_V2_ENABLED": "1"}, tags = [ "exclusive", "team:ml", "train_v2", ], deps = [ ":conftest", "//:ray_lib", ], ) py_test( name = "test_local_mode", size = "medium", srcs = ["tests/test_local_mode.py"], env = { "RAY_TRAIN_V2_ENABLED": "1", "TF_USE_LEGACY_KERAS": "1", }, tags = [ "exclusive", "team:ml", "train_v2", ], deps = [ ":conftest", "//:ray_lib", ], ) py_test( name = "test_data_config", size = "medium", srcs = ["tests/test_data_config.py"], env = {"RAY_TRAIN_V2_ENABLED": "1"}, tags = [ "data_integration", "exclusive", "team:ml", "train_v2", ], deps = [ ":conftest", "//:ray_lib", ], )