Attempt at making viewer robust to random errors (#429)

* catch more specific errors

* reduce complication; simple handle of oom error

* lints

nerfactory/engine/trainer.py

@@ -20,6 +20,7 @@ from __future__ import annotations
 import functools
 import logging
 import os
+import time
 from pathlib import Path
 from typing import Any, Dict, List
 
@@ -120,7 +121,6 @@ class Trainer:
         with TimeWriter(writer, EventName.TOTAL_TRAIN_TIME):
             num_iterations = self.config.trainer.max_num_iterations
             for step in range(self.start_step, self.start_step + num_iterations):
-
                 # if the visualizer used, the rendering of the visualizer will be included in the iteration train time
                 with TimeWriter(writer, EventName.ITER_TRAIN_TIME, step=step) as train_t:
 
@@ -138,7 +138,14 @@ class Trainer:
 
                     with TimeWriter(writer, EventName.ITER_VIS_TIME, step=step) as vis_t:
                         num_rays_per_batch = self.config.pipeline.datamanager.train_num_rays_per_batch
-                        self.visualizer_state.update_scene(step, self.pipeline.model, num_rays_per_batch)
+                        try:
+                            self.visualizer_state.update_scene(step, self.pipeline.model, num_rays_per_batch)
+                        except RuntimeError:
+                            time.sleep(0.03)  # sleep to allow buffer to reset
+                            assert self.visualizer_state.vis is not None
+                            self.visualizer_state.vis["renderingState/log_errors"].write(
+                                "Error: GPU out of memory. Reduce resolution to prevent viewer from crashing."
+                            )
 
                 if step != 0 and step % self.config.logging.steps_per_log == 0:
                     writer.put_dict(name="Loss/train-loss_metrics_dict", scalar_dict=loss_metric_dict, step=step)

nerfactory/viewer/app/src/modules/LogPanel/LogPanel.jsx

@@ -0,0 +1,91 @@
+import { useContext, useEffect } from 'react';
+import { useDispatch, useSelector } from 'react-redux';
+
+import { WebSocketContext } from '../WebSocket/WebSocket';
+
+const msgpack = require('msgpack-lite');
+
+export function LogPanel() {
+  const websocket = useContext(WebSocketContext).socket;
+  const dispatch = useDispatch();
+  const gpu_oom_error_msg = "GPU out of memory";
+  const resolved_msg = "resolved";
+  let local_error = resolved_msg;
+  // connection status indicators
+
+  const set_max_train_util = () => {
+    if (websocket.readyState === WebSocket.OPEN) {
+      dispatch({
+        type: 'write',
+        path: 'renderingState/targetTrainUtil',
+        data: 0.9,
+      });
+      const cmd = 'write';
+      const path = 'renderingState/targetTrainUtil';
+      const data = {
+        type: cmd,
+        path,
+        data: 0.9,
+      };
+      const message = msgpack.encode(data);
+      websocket.send(message);
+    }
+  };
+
+  const set_small_resolution = () => {
+    if (websocket.readyState === WebSocket.OPEN) {
+      dispatch({
+        type: 'write',
+        path: 'renderingState/maxResolution',
+        data: 512,
+      });
+      const cmd = 'write';
+      const path = 'renderingState/maxResolution';
+      const data = {
+        type: cmd,
+        path,
+        data: 512,
+      };
+      const message = msgpack.encode(data);
+      websocket.send(message);
+    }
+  };
+
+  const set_log_message = () => {
+    if (websocket.readyState === WebSocket.OPEN) {
+      dispatch({
+        type: 'write',
+        path: 'renderingState/log_errors',
+        data: resolved_msg,
+      });
+      const cmd = 'write';
+      const path = 'renderingState/log_errors';
+      const data = {
+        type: cmd,
+        path,
+        data: resolved_msg,
+      };
+      const message = msgpack.encode(data);
+      websocket.send(message);
+    }
+  };
+
+  const check_error = useSelector(
+    (state) => {
+        local_error = state.renderingState.log_errors;
+        if(local_error.includes(gpu_oom_error_msg)){
+            console.log(local_error);
+            set_log_message();
+            set_small_resolution();
+            set_max_train_util();
+        }
+    }
+  );
+
+  useEffect(() => {}, [
+    check_error,
+    local_error,
+  ]);
+
+  return null;
+}

nerfactory/viewer/app/src/modules/SidePanel/SidePanel.jsx

@@ -20,6 +20,7 @@ import SceneNode from '../../SceneNode';
 import LevaTheme from '../ConfigPanel/leva_theme.json';
 import CameraPanel from './CameraPanel';
 import { RenderControls } from '../ConfigPanel/ConfigPanel';
+import { LogPanel } from '../LogPanel/LogPanel';
 
 interface TabPanelProps {
   children: React.ReactNode;
@@ -182,7 +183,7 @@ export function BasicTabs(props: BasicTabsProps) {
         </TabPanel>
 
         <TabPanel value={value} index={3}>
-          TODO: add a logging panel
+          <LogPanel>TODO: Something maybe?</LogPanel>
         </TabPanel>
       </Box>
     </div>

nerfactory/viewer/app/src/reducer.js

@@ -33,6 +33,7 @@ const initialState = {
     eval_fps: 'Paused',
     train_eta: 'Paused',
     vis_train_ratio: 'Paused',
+    log_errors: '',
   },
   // the scene state
   sceneState: {

nerfactory/viewer/server/viewer_utils.py

@@ -20,6 +20,7 @@ import logging
 import os
 import sys
 import threading
+import time
 from pathlib import Path
 from typing import Any, Dict, Optional
 
@@ -106,8 +107,9 @@ class RenderThread(threading.Thread):
         outputs = None
         try:
             with SetTrace(self.state.check_interrupt):
-                outputs = self.graph.get_outputs_for_camera_ray_bundle(self.camera_ray_bundle)
-        except IOChangeException as e:
+                with torch.no_grad():
+                    outputs = self.graph.get_outputs_for_camera_ray_bundle(self.camera_ray_bundle)
+        except Exception as e:  # pylint: disable=broad-except
             self.exc = e
 
         if outputs:
@@ -119,7 +121,6 @@ class RenderThread(threading.Thread):
 
     def join(self, timeout=None):
         threading.Thread.join(self)
-        torch.cuda.empty_cache()
         if self.exc:
             raise self.exc
 
@@ -544,7 +545,14 @@ class VisualizerState:
         self.prev_colormap_type = colormap_type
 
         # Calculate camera pose and intrinsics
-        image_height = self._calculate_image_height(camera_object, is_training)
+        try:
+            image_height = self._calculate_image_height(camera_object, is_training)
+        except ZeroDivisionError as e:
+            self.vis["renderingState/log_errors"].write("Error: Screen too small; no rays intersecting scene.")
+            time.sleep(0.03)  # sleep to allow buffer to reset
+            print(f"Error: {e}")
+            return
+
         if image_height is None:
             return
 
@@ -585,8 +593,17 @@ class VisualizerState:
             try:
                 render_thread.join()
                 check_thread.join()
-            except Exception:  # pylint: disable=broad-except
+            except IOChangeException:
                 pass
+            except RuntimeError as e:
+                del camera_ray_bundle
+                torch.cuda.empty_cache()
+                time.sleep(0.5)  # sleep to allow buffer to reset
+                self.vis["renderingState/log_errors"].write(
+                    "Error: GPU out of memory. Reduce resolution to prevent viewer from crashing."
+                )
+                print(f"Error: {e}")
+
         graph.train()
         outputs = render_thread.vis_outputs
         if outputs is not None: