fix: make wandb imports optional (vLLM 0.22.0 broke wandb proto)

Author: xrsrke

nemo_rl/experience/rollout_manager.py

@@ -19,7 +19,10 @@
 
 import torch
 from transformers import PreTrainedTokenizerBase
-from wandb import Table
+try:
+    from wandb import Table
+except ImportError:
+    Table = None
 
 from nemo_rl.data.interfaces import DatumSpec
 from nemo_rl.distributed.batched_data_dict import BatchedDataDict

nemo_rl/experience/rollouts.py

@@ -30,7 +30,10 @@
 import torch
 from pydantic import BaseModel
 from transformers import PreTrainedTokenizerBase
-from wandb import Histogram, Table
+try:
+    from wandb import Histogram, Table
+except ImportError:
+    Histogram = Table = None
 
 from nemo_rl.data.interfaces import (
     DatumSpec,

nemo_rl/utils/logger.py

@@ -30,7 +30,10 @@
 import requests
 import swanlab
 import torch
-import wandb
+try:
+    import wandb
+except ImportError:
+    wandb = None
 from matplotlib import pyplot as plt
 from prometheus_client.parser import text_string_to_metric_families
 from prometheus_client.samples import Sample