From 3f395466534050342f95d4fb837f27ee3cc1d688 Mon Sep 17 00:00:00 2001 From: Satvik Sharma <61220944+Satvik1701@users.noreply.github.com> Date: Thu, 25 Apr 2024 23:35:24 -0700 Subject: [PATCH] Adding lock for Viser Element callbacks (#3067) * add callback lock, prepare for testing Co-authored-by: Justin Kerr --- nerfstudio/scripts/viewer/run_viewer.py | 5 +++++ nerfstudio/viewer/viewer.py | 16 ++++++++++++++-- 2 files changed, 19 insertions(+), 2 deletions(-) diff --git a/nerfstudio/scripts/viewer/run_viewer.py b/nerfstudio/scripts/viewer/run_viewer.py index 79e7167f3..b4d0420ee 100644 --- a/nerfstudio/scripts/viewer/run_viewer.py +++ b/nerfstudio/scripts/viewer/run_viewer.py @@ -16,11 +16,13 @@ """ Starts viewer in eval mode. """ + from __future__ import annotations import time from dataclasses import dataclass, field, fields from pathlib import Path +from threading import Lock from typing import Literal import tyro @@ -89,12 +91,14 @@ def _start_viewer(config: TrainerConfig, pipeline: Pipeline, step: int): viewer_log_path = base_dir / config.viewer.relative_log_filename banner_messages = None viewer_state = None + viewer_callback_lock = Lock() if config.vis == "viewer_legacy": viewer_state = ViewerLegacyState( config.viewer, log_filename=viewer_log_path, datapath=pipeline.datamanager.get_datapath(), pipeline=pipeline, + train_lock=viewer_callback_lock, ) banner_messages = [f"Legacy viewer at: {viewer_state.viewer_url}"] if config.vis == "viewer": @@ -104,6 +108,7 @@ def _start_viewer(config: TrainerConfig, pipeline: Pipeline, step: int): datapath=pipeline.datamanager.get_datapath(), pipeline=pipeline, share=config.viewer.make_share_url, + train_lock=viewer_callback_lock, ) banner_messages = viewer_state.viewer_info diff --git a/nerfstudio/viewer/viewer.py b/nerfstudio/viewer/viewer.py index db84e9550..a5093f0da 100644 --- a/nerfstudio/viewer/viewer.py +++ b/nerfstudio/viewer/viewer.py @@ -12,9 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. -""" Manage the state of the viewer """ +"""Manage the state of the viewer""" + from __future__ import annotations +import contextlib import threading import time from pathlib import Path @@ -209,12 +211,22 @@ class Viewer: # Keep track of the pointers to generated GUI folders, because each generated folder holds a unique ID. viewer_gui_folders = dict() + def prev_cb_wrapper(prev_cb): + # We wrap the callbacks in the train_lock so that the callbacks are thread-safe with the + # concurrently executing render thread. This may block rendering, however this can be necessary + # if the callback uses get_outputs internally. + def cb_lock(element): + with self.train_lock if self.train_lock is not None else contextlib.nullcontext(): + prev_cb(element) + + return cb_lock + def nested_folder_install(folder_labels: List[str], prev_labels: List[str], element: ViewerElement): if len(folder_labels) == 0: element.install(self.viser_server) # also rewire the hook to rerender prev_cb = element.cb_hook - element.cb_hook = lambda element: [prev_cb(element), self._trigger_rerender()] + element.cb_hook = lambda element: [prev_cb_wrapper(prev_cb)(element), self._trigger_rerender()] else: # recursively create folders # If the folder name is "Custom Elements/a/b", then: