[air] Fix broken release tests from head node sync deprecation (ray-p…

…roject#37613)

Signed-off-by: Justin Yu <justinvyu@anyscale.com>
Signed-off-by: e428265 <arvind.chandramouli@lmco.com>

python/ray/train/BUILD

@@ -246,15 +246,6 @@ py_test(
 # Please keep these sorted alphabetically.
 # --------------------------------------------------------------------
 
-py_test(
-    name = "pytorch_pbt_failure",
-    size = "small",
-    srcs = ["tests/pytorch_pbt_failure.py"],
-    tags = ["team:ml", "exlusive", "no_main"],
-    deps = [":train_lib"],
-    args = ["--smoke-test"]
-)
-
 py_test(
     name = "test_accelerate_trainer_gpu",
     size = "large",

release/air_tests/air_benchmarks/workloads/pytorch_training_e2e.py

@@ -16,7 +16,7 @@
 from ray import train
 from ray.air import session
 from ray.train.torch import TorchTrainer
-from ray.air.config import ScalingConfig
+from ray.air.config import RunConfig, ScalingConfig
 
 
 def add_fake_labels(batch: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
@@ -108,6 +108,7 @@ def main(data_size_gb: int, num_epochs=2, num_workers=1, smoke_test: bool = Fals
         scaling_config=ScalingConfig(
             num_workers=num_workers, use_gpu=int(not smoke_test)
         ),
+        run_config=RunConfig(storage_path="/mnt/cluster_storage"),
     )
     trainer.fit()
 

release/long_running_distributed_tests/workloads/pytorch_pbt_failure.py

@@ -72,6 +72,7 @@
         failure_config=FailureConfig(max_failures=-1),
         checkpoint_config=CheckpointConfig(num_to_keep=10),
         callbacks=[FailureInjectorCallback(time_between_checks=90), ProgressCallback()],
+        storage_path="/mnt/cluster_storage",
     ),
 )