chore: import upstream snapshot with attribution

python/ray/train/tests/test_e2e_wandb_integration.py

@@ -0,0 +1,73 @@
+"""
+If a user uses Trainer API directly with wandb integration, they expect to see
+* train_loop_config to show up in wandb.config.
+
+This test uses mocked call into wandb API.
+"""
+
+import pytest
+
+import ray
+from ray.air.integrations.wandb import WANDB_ENV_VAR
+from ray.air.tests.mocked_wandb_integration import WandbTestExperimentLogger
+from ray.train import RunConfig, ScalingConfig
+from ray.train.examples.pytorch.torch_linear_example import (
+    train_func as linear_train_func,
+)
+from ray.train.torch import TorchTrainer
+
+
+@pytest.fixture
+def ray_start_4_cpus():
+    address_info = ray.init(num_cpus=4)
+    yield address_info
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+CONFIG = {"lr": 1e-2, "hidden_size": 1, "batch_size": 4, "epochs": 3}
+
+
+@pytest.mark.parametrize("with_train_loop_config", (True, False))
+def test_trainer_wandb_integration(
+    ray_start_4_cpus, with_train_loop_config, monkeypatch
+):
+    monkeypatch.setenv(WANDB_ENV_VAR, "9012")
+
+    def train_func(config=None):
+        config = config or CONFIG
+        result = linear_train_func(config)
+        assert len(result) == config["epochs"]
+        assert result[-1]["loss"] < result[0]["loss"]
+
+    scaling_config = ScalingConfig(num_workers=2)
+
+    logger = WandbTestExperimentLogger(project="test_project")
+    if with_train_loop_config:
+        trainer = TorchTrainer(
+            train_loop_per_worker=train_func,
+            train_loop_config=CONFIG,
+            scaling_config=scaling_config,
+            run_config=RunConfig(callbacks=[logger]),
+        )
+    else:
+        trainer = TorchTrainer(
+            train_loop_per_worker=train_func,
+            scaling_config=scaling_config,
+            run_config=RunConfig(callbacks=[logger]),
+        )
+    trainer.fit()
+    config = list(logger.trial_logging_actor_states.values())[0].config
+
+    if with_train_loop_config:
+        assert "train_loop_config" in config
+    else:
+        assert "train_loop_config" not in config
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))