Repository: wandb/wandb
Pull request: [WB-5919][CLI-2349] fix: watch compatible with disabled mode
<!--
  Name your PR: (Use one of the below formats)
  - [CLI-NNNN] Brief description of changes if jira ticket
  - [WB-NNNN] Brief description of changes if jira ticket
  - Brief description of changes

  Also:
  - Mark your PR as a Draft if it isnt ready for merge yet
-->

<!-- Include one or more of the following issue URLs if applicable -->
https://wandb.atlassian.net/browse/WB-5919
https://github.com/wandb/client/issues/2349

Description
-----------

What does the PR do?
Fixes #2349 
Make `watch` method compatible with `disabled` mode.
Also it is now possible to disable logging of graph topology when using `watch`.

Testing
-------

How was this PR tested?
Manually

Problem statement:
[CLI] Error when trying to disable wandb
**Description**
I am getting
`torch.nn.modules.module.ModuleAttributeError: '<my model name>' object has no attribute '_wandb_hook_names' `
when I try to disable wand either with the env variable WANDB_MODE=disabled or with wandb.init(mode='disabled')


**How to reproduce**
Minimal example
```
import wandb
import torch as th

wandb.init(mode='disabled')
# wandb.init()


class Model(th.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = th.nn.Conv2d(3, 9, 3)

    def forward(self, x):
        x = self.conv(x)
        return x


model = Model()

wandb.watch(model, log='all', log_freq=10)
```

**Environment**
- Ubuntu 20.04
- Python Version: 3.7
- Pytorch version 1.7.1
- Wandb 0.10.33

Hints from issue discussion:
@borisdayma could you please take a look

Commit message: feat: make watch more robust

Patch to review:
diff --git a/wandb/sdk/wandb_run.py b/wandb/sdk/wandb_run.py
index 37af8650e55..9916c2a3a9f 100644
--- a/wandb/sdk/wandb_run.py
+++ b/wandb/sdk/wandb_run.py
@@ -2075,8 +2075,8 @@ def define_metric(
         return m
 
     # TODO(jhr): annotate this
-    def watch(self, models, criterion=None, log="gradients", log_freq=100, idx=None) -> None:  # type: ignore
-        wandb.watch(models, criterion, log, log_freq, idx)
+    def watch(self, models, criterion=None, log="gradients", log_freq=100, idx=None, log_graph=True) -> None:  # type: ignore
+        wandb.watch(models, criterion, log, log_freq, idx, log_graph)
 
     # TODO(jhr): annotate this
     def use_artifact(self, artifact_or_name, type=None, aliases=None):  # type: ignore
diff --git a/wandb/sdk/wandb_watch.py b/wandb/sdk/wandb_watch.py
index bade4e2e8be..8209b33fd64 100644
--- a/wandb/sdk/wandb_watch.py
+++ b/wandb/sdk/wandb_watch.py
@@ -14,7 +14,9 @@
 _global_watch_idx = 0
 
 
-def watch(models, criterion=None, log="gradients", log_freq=1000, idx=None):
+def watch(
+    models, criterion=None, log="gradients", log_freq=1000, idx=None, log_graph=True
+):
     """Hooks into the torch model to collect gradients and the topology.
 
     Should be extended to accept arbitrary ML models.
@@ -25,6 +27,7 @@ def watch(models, criterion=None, log="gradients", log_freq=1000, idx=None):
         log: (str) One of "gradients", "parameters", "all", or None
         log_freq: (int) log gradients and parameters every N batches
         idx: (int) an index to be used when calling wandb.watch on multiple models
+        log_graph: (boolean) log graph topology
 
     Returns:
         `wandb.Graph` The graph object that will populate after the first backward pass
@@ -73,6 +76,7 @@ def watch(models, criterion=None, log="gradients", log_freq=1000, idx=None):
 
     graphs = []
     prefix = ""
+
     if idx is None:
         idx = _global_watch_idx
     for local_idx, model in enumerate(models):
@@ -91,11 +95,12 @@ def watch(models, criterion=None, log="gradients", log_freq=1000, idx=None):
             jupyter_run=wandb.run if in_jupyter else None,
         )
 
-        graph = wandb.wandb_torch.TorchGraph.hook_torch(
-            model, criterion, graph_idx=global_idx
-        )
-        graphs.append(graph)
-        # NOTE: the graph is set in run.summary by hook_torch on the backward pass
+        if log_graph:
+            graph = wandb.run.history.torch.hook_torch(
+                model, criterion, graph_idx=global_idx
+            )
+            graphs.append(graph)
+            # NOTE: the graph is set in run.summary by hook_torch on the backward pass
     return graphs
 
 
diff --git a/wandb/sdk_py27/wandb_run.py b/wandb/sdk_py27/wandb_run.py
index 6801ca0f9d5..82a253bd953 100644
--- a/wandb/sdk_py27/wandb_run.py
+++ b/wandb/sdk_py27/wandb_run.py
@@ -2075,8 +2075,8 @@ def define_metric(
         return m
 
     # TODO(jhr): annotate this
-    def watch(self, models, criterion=None, log="gradients", log_freq=100, idx=None):  # type: ignore
-        wandb.watch(models, criterion, log, log_freq, idx)
+    def watch(self, models, criterion=None, log="gradients", log_freq=100, idx=None, log_graph=True):  # type: ignore
+        wandb.watch(models, criterion, log, log_freq, idx, log_graph)
 
     # TODO(jhr): annotate this
     def use_artifact(self, artifact_or_name, type=None, aliases=None):  # type: ignore
diff --git a/wandb/sdk_py27/wandb_watch.py b/wandb/sdk_py27/wandb_watch.py
index 8956f4e6ef0..eb0bab978b6 100644
--- a/wandb/sdk_py27/wandb_watch.py
+++ b/wandb/sdk_py27/wandb_watch.py
@@ -14,7 +14,9 @@
 _global_watch_idx = 0
 
 
-def watch(models, criterion=None, log="gradients", log_freq=1000, idx=None):
+def watch(
+    models, criterion=None, log="gradients", log_freq=1000, idx=None, log_graph=True
+):
     """Hooks into the torch model to collect gradients and the topology.
 
     Should be extended to accept arbitrary ML models.
@@ -25,6 +27,7 @@ def watch(models, criterion=None, log="gradients", log_freq=1000, idx=None):
         log: (str) One of "gradients", "parameters", "all", or None
         log_freq: (int) log gradients and parameters every N batches
         idx: (int) an index to be used when calling wandb.watch on multiple models
+        log_graph: (boolean) log graph topology
 
     Returns:
         `wandb.Graph` The graph object that will populate after the first backward pass
@@ -73,6 +76,7 @@ def watch(models, criterion=None, log="gradients", log_freq=1000, idx=None):
 
     graphs = []
     prefix = ""
+
     if idx is None:
         idx = _global_watch_idx
     for local_idx, model in enumerate(models):
@@ -91,11 +95,12 @@ def watch(models, criterion=None, log="gradients", log_freq=1000, idx=None):
             jupyter_run=wandb.run if in_jupyter else None,
         )
 
-        graph = wandb.wandb_torch.TorchGraph.hook_torch(
-            model, criterion, graph_idx=global_idx
-        )
-        graphs.append(graph)
-        # NOTE: the graph is set in run.summary by hook_torch on the backward pass
+        if log_graph:
+            graph = wandb.run.history.torch.hook_torch(
+                model, criterion, graph_idx=global_idx
+            )
+            graphs.append(graph)
+            # NOTE: the graph is set in run.summary by hook_torch on the backward pass
     return graphs
 
 
diff --git a/wandb/wandb_torch.py b/wandb/wandb_torch.py
index 737d422c77b..4df30fd8037 100644
--- a/wandb/wandb_torch.py
+++ b/wandb/wandb_torch.py
@@ -79,6 +79,7 @@ def __init__(self, history):
         self._num_bins = 64
         self._is_cuda_histc_supported = None
         self._jupyter_run = None
+        self.hook_torch = TorchGraph.hook_torch
 
     def add_log_hooks_to_pytorch_module(
         self,
@@ -309,6 +310,7 @@ def __init__(self):
 
     @classmethod
     def hook_torch(cls, model, criterion=None, graph_idx=0):
+        print("wandb: logging graph, to disable use `wandb.watch(log_graph=False)`")
         graph = TorchGraph()
         graph.hook_torch_modules(model, criterion, graph_idx=graph_idx)
         return graph
@@ -410,6 +412,9 @@ def hook_torch_modules(
                 wandb.run.history.torch._hook_handles[
                     "topology/" + str(id(graph_hook))
                 ] = graph_hook
+                if not hasattr(parent, "_wandb_hook_names"):
+                    # should never happen but let's be extra safe
+                    parent._wandb_hook_names = []
                 parent._wandb_hook_names.append("topology/" + str(id(graph_hook)))
 
     @classmethod


Code review comments:
diff_hunk: @@ -309,6 +310,7 @@ def __init__(self):
 
     @classmethod
     def hook_torch(cls, model, criterion=None, graph_idx=0):
+        print("wandb: logging graph, to disable use `wandb.watch(log_graph=False)`")
path: wandb/wandb_torch.py
text: @user1:
Should like y just use `wandb.termlog("logging graph...` here.