[AIR] Update WandbLoggerCallback example

.buildkite/pipeline.build.yml

@@ -573,6 +573,7 @@
     - pip uninstall -y ray && rm -rf /ray/python/ray/thirdparty_files
     - pushd /ray && git clean -f -f -x -d -e .whl -e python/ray/dashboard/client && popd
     - bazel clean --expunge
+    - export WANDB_MODE=offline
     - ./ci/ci.sh build
 
 - label: ":octopus: Tune multinode tests"
@@ -623,7 +624,7 @@
     - pip install -U typing-extensions
     - HOROVOD_WITH_GLOO=1 HOROVOD_WITHOUT_MPI=1 HOROVOD_WITHOUT_MXNET=1 HOROVOD_WITH_TENSORFLOW=1 HOROVOD_WITH_PYTORCH=1 pip install horovod
     - ./ci/env/env_info.sh
-    - bazel test --config=ci $(./scripts/bazel_export_options) --test_tag_filters=compat_py36 
+    - bazel test --config=ci $(./scripts/bazel_export_options) --test_tag_filters=compat_py36
       python/ray/tests/horovod/...
       python/ray/tests/lightgbm/...
       python/ray/tests/ml_py36_compat/...

python/ray/air/integrations/wandb.py

@@ -37,6 +37,7 @@
 WANDB_ENV_VAR = "WANDB_API_KEY"
 WANDB_PROJECT_ENV_VAR = "WANDB_PROJECT_NAME"
 WANDB_GROUP_ENV_VAR = "WANDB_GROUP_NAME"
+WANDB_MODE_ENV_VAR = "WANDB_MODE"
 # Hook that is invoked before wandb.init in the setup method of WandbLoggerCallback
 # to populate the API key if it isn't already set when initializing the callback.
 # It doesn't take in any arguments and returns the W&B API key.
@@ -295,7 +296,11 @@ def _set_api_key(api_key_file: Optional[str] = None, api_key: Optional[str] = No
       1) From `api_key` or `api_key_file` arguments
       2) From WANDB_API_KEY environment variables
       3) User already logged in to W&B (wandb.api.api_key set)
-      4) From external hook WANDB_SETUP_API_KEY_HOOK"""
+      4) From external hook WANDB_SETUP_API_KEY_HOOK
+    """
+    if os.environ.get(WANDB_MODE_ENV_VAR) in {"offline", "disabled"}:
+        return
+
     if api_key_file:
         if api_key:
             raise ValueError("Both WandB `api_key_file` and `api_key` set.")
@@ -448,6 +453,43 @@ class WandbLoggerCallback(LoggerCallback):
     ``LoggerCallback`` sends metrics to Wandb for automatic tracking and
     visualization.
 
+    Example:
+
+        .. testcode::
+
+            import random
+
+            from ray import tune
+            from ray.air import session, RunConfig
+            from ray.air.integrations.wandb import WandbLoggerCallback
+
+
+            def train_func(config):
+                offset = random.random() / 5
+                for epoch in range(2, config["epochs"]):
+                    acc = 1 - (2 + config["lr"]) ** -epoch - random.random() / epoch - offset
+                    loss = (2 + config["lr"]) ** -epoch + random.random() / epoch + offset
+                    session.report({"acc": acc, "loss": loss})
+
+
+            tuner = tune.Tuner(
+                train_func,
+                param_space={
+                    "lr": tune.grid_search([0.001, 0.01, 0.1, 1.0]),
+                    "epochs": 10,
+                },
+                run_config=RunConfig(
+                    callbacks=[WandbLoggerCallback(project="Optimization_Project")]
+                ),
+            )
+            results = tuner.fit()
+
+        .. testoutput::
+            :hide:
+            :options: +ELLIPSIS
+
+            ...
+
     Args:
         project: Name of the Wandb project. Mandatory.
         group: Name of the Wandb group. Defaults to the trainable
@@ -472,26 +514,7 @@ class WandbLoggerCallback(LoggerCallback):
 
     Please see here for all other valid configuration settings:
     https://docs.wandb.ai/library/init
-
-    Example:
-
-    .. code-block:: python
-
-        from ray.tune.logger import DEFAULT_LOGGERS
-        from ray.air.integrations.wandb import WandbLoggerCallback
-        tune.run(
-            train_fn,
-            config={
-                # define search space here
-                "parameter_1": tune.choice([1, 2, 3]),
-                "parameter_2": tune.choice([4, 5, 6]),
-            },
-            callbacks=[WandbLoggerCallback(
-                project="Optimization_Project",
-                api_key_file="/path/to/file",
-                log_config=True)])
-
-    """
+    """  # noqa: E501
 
     # Do not log these result keys
     _exclude_results = ["done", "should_checkpoint"]