W&B: Enable login timeout (#4843)

Dockerfile

@@ -10,7 +10,7 @@ RUN apt update && apt install -y zip htop screen libgl1-mesa-glx
 COPY requirements.txt .
 RUN python -m pip install --upgrade pip
 RUN pip uninstall -y nvidia-tensorboard nvidia-tensorboard-plugin-dlprof
-RUN pip install --no-cache -r requirements.txt coremltools onnx gsutil notebook
+RUN pip install --no-cache -r requirements.txt coremltools onnx gsutil notebook wandb>=0.12.2
 RUN pip install --no-cache -U torch torchvision numpy
 # RUN pip install --no-cache torch==1.9.0+cu111 torchvision==0.10.0+cu111 -f https://download.pytorch.org/whl/torch_stable.html
 

utils/loggers/wandb/wandb_utils.py

@@ -5,6 +5,7 @@ import os
 import sys
 from contextlib import contextmanager
 from pathlib import Path
+import pkg_resources as pkg
 
 import yaml
 from tqdm import tqdm
@@ -20,6 +21,8 @@ try:
     import wandb
 
     assert hasattr(wandb, '__version__')  # verify package import not local dir
+    if pkg.parse_version(wandb.__version__) >= pkg.parse_version('0.12.2'):
+        wandb.login(timeout=30)
 except (ImportError, AssertionError):
     wandb = None