def test_dequeue(self): """Test that we can enqueue and dequeue on different threads""" queue = ScriptRequestQueue() # This should return immediately self.assertEqual((None, None), queue.dequeue()) lock = Lock() dequeued_evt = [None] def get_event(): with lock: return dequeued_evt[0] def set_event(value): with lock: dequeued_evt[0] = value def do_dequeue(): event = None while event is None: event, _ = queue.dequeue() set_event(event) thread = Thread(target=do_dequeue, name='test_dequeue') thread.start() self.assertIsNone(get_event()) queue.enqueue(ScriptRequest.STOP) time.sleep(0.1) self.assertEqual(ScriptRequest.STOP, get_event()) thread.join(timeout=0.25) self.assertFalse(thread.is_alive())
class ReportSession(object): """ Contains session data for a single "user" of an active report (that is, a connected browser tab). Each ReportSession has its own Report, root DeltaGenerator, ScriptRunner, and widget state. A ReportSession is attached to each thread involved in running its Report. """ _next_id = 0 def __init__(self, ioloop, script_path, script_argv): """Initialize the ReportSession. Parameters ---------- ioloop : tornado.ioloop.IOLoop The Tornado IOLoop that we're running within script_path : str Path of the Python file from which this report is generated. script_argv : list of str Command-line arguments to run the script with. """ # Each ReportSession gets a unique ID self.id = ReportSession._next_id ReportSession._next_id += 1 self._ioloop = ioloop self._report = Report(script_path, script_argv) self._state = ReportSessionState.REPORT_NOT_RUNNING self._main_dg = DeltaGenerator(self.enqueue, container=BlockPath.MAIN) self._sidebar_dg = DeltaGenerator(self.enqueue, container=BlockPath.SIDEBAR) self._widget_states = WidgetStates() self._local_sources_watcher = LocalSourcesWatcher( self._report, self._on_source_file_changed) self._sent_initialize_message = False self._storage = None self._maybe_reuse_previous_run = False self._run_on_save = config.get_option('server.runOnSave') # The ScriptRequestQueue is the means by which we communicate # with the active ScriptRunner. self._script_request_queue = ScriptRequestQueue() self._scriptrunner = None LOGGER.debug('ReportSession initialized (id=%s)', self.id) def flush_browser_queue(self): """Clears the report queue and returns the messages it contained. The Server calls this periodically to deliver new messages to the browser connected to this report. Returns ------- list[ForwardMsg] The messages that were removed from the queue and should be delivered to the browser. """ return self._report.flush_browser_queue() def shutdown(self): """Shuts down the ReportSession. It's an error to use a ReportSession after it's been shut down. """ if self._state != ReportSessionState.SHUTDOWN_REQUESTED: LOGGER.debug('Shutting down (id=%s)', self.id) # Shut down the ScriptRunner, if one is active. # self._state must not be set to SHUTDOWN_REQUESTED until # after this is called. if self._scriptrunner is not None: self._enqueue_script_request(ScriptRequest.SHUTDOWN) self._state = ReportSessionState.SHUTDOWN_REQUESTED self._local_sources_watcher.close() def enqueue(self, msg): """Enqueues a new ForwardMsg to our browser queue. This can be called on both the main thread and a ScriptRunner run thread. Parameters ---------- msg : ForwardMsg The message to enqueue Returns ------- bool True if the message was enqueued, or False if client.displayEnabled is not set """ if not config.get_option('client.displayEnabled'): return False # If we have an active ScriptRunner, signal that it can handle # an execution control request. (Copy the scriptrunner reference # to avoid it being unset from underneath us, as this function can # be called outside the main thread.) scriptrunner = self._scriptrunner if scriptrunner is not None: scriptrunner.maybe_handle_execution_control_request() self._report.enqueue(msg) return True def enqueue_exception(self, e): """Enqueues an Exception message. Parameters ---------- e : BaseException """ # This does a few things: # 1) Clears the current report in the browser. # 2) Marks the current report as "stopped" in the browser. # 3) HACK: Resets any script params that may have been broken (e.g. the # command-line when rerunning with wrong argv[0]) self._on_scriptrunner_event( ScriptRunnerEvent.SCRIPT_STOPPED_WITH_SUCCESS) self._on_scriptrunner_event(ScriptRunnerEvent.SCRIPT_STARTED) self._on_scriptrunner_event( ScriptRunnerEvent.SCRIPT_STOPPED_WITH_SUCCESS) msg = ForwardMsg() msg.delta.id = 0 exception_proto.marshall(msg.delta.new_element.exception, e) self.enqueue(msg) def request_rerun(self, argv=None, widget_state=None): """Signal that we're interested in running the script. If the script is not already running, it will be started immediately. Otherwise, a rerun will be requested. Parameters ---------- argv : dict | None The new command line arguments to run the script with, or None to use the argv from the previous run of the script. widget_state : dict | None The widget state dictionary to run the script with, or None to use the widget state from the previous run of the script. """ self._enqueue_script_request(ScriptRequest.RERUN, RerunData(argv, widget_state)) def _on_source_file_changed(self): """One of our source files changed. Schedule a rerun if appropriate.""" if self._run_on_save: self.request_rerun() else: self._enqueue_file_change_message() def _clear_queue(self): self._report.clear() def _on_scriptrunner_event(self, event, exception=None, widget_states=None): """Called when our ScriptRunner emits an event. This is *not* called on the main thread. Parameters ---------- event : ScriptRunnerEvent exception : BaseException | None An exception thrown during compilation. Set only for the SCRIPT_STOPPED_WITH_COMPILE_ERROR event. widget_states : streamlit.proto.Widget_pb2.WidgetStates | None The ScriptRunner's final WidgetStates. Set only for the SHUTDOWN event. """ LOGGER.debug('OnScriptRunnerEvent: %s', event) prev_state = self._state if event == ScriptRunnerEvent.SCRIPT_STARTED: if self._state != ReportSessionState.SHUTDOWN_REQUESTED: self._state = ReportSessionState.REPORT_IS_RUNNING if config.get_option('server.liveSave'): # Enqueue into the IOLoop so it runs without blocking AND runs # on the main thread. self._ioloop.spawn_callback(self._save_running_report) self._clear_queue() self._maybe_enqueue_initialize_message() self._enqueue_new_report_message() elif (event == ScriptRunnerEvent.SCRIPT_STOPPED_WITH_SUCCESS or event == ScriptRunnerEvent.SCRIPT_STOPPED_WITH_COMPILE_ERROR): if self._state != ReportSessionState.SHUTDOWN_REQUESTED: self._state = ReportSessionState.REPORT_NOT_RUNNING self._enqueue_report_finished_message() if config.get_option('server.liveSave'): # Enqueue into the IOLoop so it runs without blocking AND runs # on the main thread. self._ioloop.spawn_callback(self._save_final_report_and_quit) script_succeeded = \ event == ScriptRunnerEvent.SCRIPT_STOPPED_WITH_SUCCESS if script_succeeded: # When a script completes successfully, we update our # LocalSourcesWatcher to account for any source code changes # that change which modules should be watched. (This is run on # the main thread, because LocalSourcesWatcher is not # thread safe.) self._ioloop.spawn_callback( self._local_sources_watcher.update_watched_modules) else: # When a script fails to compile, we send along the exception. from streamlit.elements import exception_proto msg = ForwardMsg() exception_proto.marshall( msg.session_event.script_compilation_exception, exception) self.enqueue(msg) elif event == ScriptRunnerEvent.SHUTDOWN: # When ScriptRunner shuts down, update our local reference to it, # and check to see if we need to spawn a new one. (This is run on # the main thread.) def on_shutdown(): self._widget_states = widget_states self._scriptrunner = None # Because a new ScriptEvent could have been enqueued while the # scriptrunner was shutting down, we check to see if we should # create a new one. (Otherwise, a newly-enqueued ScriptEvent # won't be processed until another event is enqueued.) self._maybe_create_scriptrunner() self._ioloop.spawn_callback(on_shutdown) # Send a message if our run state changed report_was_running = prev_state == ReportSessionState.REPORT_IS_RUNNING report_is_running = self._state == ReportSessionState.REPORT_IS_RUNNING if report_is_running != report_was_running: self._enqueue_session_state_changed_message() def _enqueue_session_state_changed_message(self): msg = ForwardMsg() msg.session_state_changed.run_on_save = self._run_on_save msg.session_state_changed.report_is_running = \ self._state == ReportSessionState.REPORT_IS_RUNNING self.enqueue(msg) def _enqueue_file_change_message(self): LOGGER.debug('Enqueuing report_changed message (id=%s)', self.id) msg = ForwardMsg() msg.session_event.report_changed_on_disk = True self.enqueue(msg) def _maybe_enqueue_initialize_message(self): if self._sent_initialize_message: return self._sent_initialize_message = True msg = ForwardMsg() imsg = msg.initialize imsg.config.sharing_enabled = (config.get_option('global.sharingMode') != 'off') LOGGER.debug('New browser connection: sharing_enabled=%s', imsg.config.sharing_enabled) imsg.config.gather_usage_stats = ( config.get_option('browser.gatherUsageStats')) LOGGER.debug('New browser connection: gather_usage_stats=%s', imsg.config.gather_usage_stats) imsg.environment_info.streamlit_version = __version__ imsg.environment_info.python_version = ('.'.join( map(str, sys.version_info))) imsg.session_state.run_on_save = self._run_on_save imsg.session_state.report_is_running = ( self._state == ReportSessionState.REPORT_IS_RUNNING) imsg.user_info.installation_id = __installation_id__ imsg.user_info.email = Credentials.get_current().activation.email self.enqueue(msg) def _enqueue_new_report_message(self): self._report.generate_new_id() msg = ForwardMsg() msg.new_report.id = self._report.report_id msg.new_report.command_line.extend(self._report.argv) msg.new_report.name = self._report.name msg.new_report.script_path = self._report.script_path self.enqueue(msg) def _enqueue_report_finished_message(self): msg = ForwardMsg() msg.report_finished = True self.enqueue(msg) def handle_rerun_script_request(self, command_line=None, widget_state=None, is_preheat=False): """Tells the ScriptRunner to re-run its report. Parameters ---------- command_line : str | None The new command line arguments to run the script with, or None to use its previous command line value. widget_state : WidgetStates | None The WidgetStates protobuf to run the script with, or None to use its previous widget states. is_preheat: boolean True if this ReportSession should run the script immediately, and then ignore the next rerun request if it matches the already-ran argv and widget state. """ old_argv = self._report.argv if command_line is not None: self._report.parse_argv_from_command_line(command_line) if is_preheat: self._maybe_reuse_previous_run = True # For next time. elif self._maybe_reuse_previous_run: # If this is a "preheated" ReportSession, reuse the previous run if # the argv and widget state matches. But only do this one time # ever. self._maybe_reuse_previous_run = False has_widget_state = (widget_state is not None and len(widget_state.widgets) > 0) has_new_argv = old_argv != self._report.argv if not has_widget_state and not has_new_argv: LOGGER.debug( 'Skipping rerun since the preheated run is the same') return self.request_rerun(self._report.argv, widget_state) def handle_stop_script_request(self): """Tells the ScriptRunner to stop running its report.""" self._enqueue_script_request(ScriptRequest.STOP) def handle_clear_cache_request(self): """Clears this report's cache. Because this cache is global, it will be cleared for all users. """ # Setting verbose=True causes clear_cache to print to stdout. # Since this command was initiated from the browser, the user # doesn't need to see the results of the command in their # terminal. caching.clear_cache() def handle_set_run_on_save_request(self, new_value): """Changes our run_on_save flag to the given value. The browser will be notified of the change. Parameters ---------- new_value : bool New run_on_save value """ self._run_on_save = new_value self._enqueue_session_state_changed_message() def _enqueue_script_request(self, request, data=None): """Enqueue a ScriptEvent into our ScriptEventQueue. If a script thread is not already running, one will be created to handle the event. Parameters ---------- request : ScriptRequest The type of request. data : Any Data associated with the request, if any. """ if self._state == ReportSessionState.SHUTDOWN_REQUESTED: LOGGER.warning('Discarding %s request after shutdown' % request) return self._script_request_queue.enqueue(request, data) self._maybe_create_scriptrunner() def _maybe_create_scriptrunner(self): """Create a new ScriptRunner if we have unprocessed script requests. This is called every time a ScriptRequest is enqueued, and also after a ScriptRunner shuts down, in case new requests were enqueued during its termination. This function should only be called on the main thread. """ if (self._state == ReportSessionState.SHUTDOWN_REQUESTED or self._scriptrunner is not None or not self._script_request_queue.has_request): return # Create the ScriptRunner, attach event handlers, and start it self._scriptrunner = ScriptRunner( report=self._report, main_dg=self._main_dg, sidebar_dg=self._sidebar_dg, widget_states=self._widget_states, request_queue=self._script_request_queue) self._scriptrunner.on_event.connect(self._on_scriptrunner_event) self._scriptrunner.start() @tornado.gen.coroutine def handle_save_request(self, ws): """Save serialized version of report deltas to the cloud.""" @tornado.gen.coroutine def progress(percent): progress_msg = ForwardMsg() progress_msg.upload_report_progress = percent yield ws.write_message(progress_msg.SerializeToString(), binary=True) # Indicate that the save is starting. try: yield progress(0) url = yield self._save_final_report(progress) # Indicate that the save is done. progress_msg = ForwardMsg() progress_msg.report_uploaded = url yield ws.write_message(progress_msg.SerializeToString(), binary=True) except Exception as e: # Horrible hack to show something if something breaks. err_msg = '%s: %s' % (type(e).__name__, str(e) or 'No further details.') progress_msg = ForwardMsg() progress_msg.report_uploaded = err_msg yield ws.write_message(progress_msg.SerializeToString(), binary=True) raise e @tornado.gen.coroutine def _save_running_report(self): files = self._report.serialize_running_report_to_files() url = yield self._get_storage().save_report_files( self._report.report_id, files) if config.get_option('server.liveSave'): util.print_url('Saved running report', url) raise tornado.gen.Return(url) @tornado.gen.coroutine def _save_final_report(self, progress=None): files = self._report.serialize_final_report_to_files() url = yield self._get_storage().save_report_files( self._report.report_id, files, progress) if config.get_option('server.liveSave'): util.print_url('Saved final report', url) raise tornado.gen.Return(url) @tornado.gen.coroutine def _save_final_report_and_quit(self): yield self._save_final_report() self._ioloop.stop() def _get_storage(self): if self._storage is None: self._storage = Storage() return self._storage
class TestScriptRunner(ScriptRunner): """Subclasses ScriptRunner to provide some testing features.""" def __init__(self, script_name): """Initializes the ScriptRunner for the given script_name""" # DeltaGenerator deltas will be enqueued into self.report_queue. self.report_queue = ReportQueue() def enqueue_fn(msg): self.report_queue.enqueue(msg) self.maybe_handle_execution_control_request() return True self.main_dg = DeltaGenerator(enqueue=enqueue_fn, container=BlockPath.MAIN) self.sidebar_dg = DeltaGenerator( enqueue=enqueue_fn, container=BlockPath.SIDEBAR ) self.script_request_queue = ScriptRequestQueue() script_path = os.path.join(os.path.dirname(__file__), "test_data", script_name) super(TestScriptRunner, self).__init__( report=Report(script_path, "test command line"), main_dg=self.main_dg, sidebar_dg=self.sidebar_dg, widget_states=WidgetStates(), request_queue=self.script_request_queue, ) # Accumulates uncaught exceptions thrown by our run thread. self.script_thread_exceptions = [] # Accumulates all ScriptRunnerEvents emitted by us. self.events = [] def record_event(event, **kwargs): self.events.append(event) self.on_event.connect(record_event, weak=False) @property def widget_states(self): """ Returns ------- WidgetStates A WidgetStates protobuf object """ return self._widgets.get_state() def enqueue_rerun(self, argv=None, widget_state=None): self.script_request_queue.enqueue( ScriptRequest.RERUN, RerunData(widget_state=widget_state) ) def enqueue_stop(self): self.script_request_queue.enqueue(ScriptRequest.STOP) def enqueue_shutdown(self): self.script_request_queue.enqueue(ScriptRequest.SHUTDOWN) def _process_request_queue(self): try: super(TestScriptRunner, self)._process_request_queue() except BaseException as e: self.script_thread_exceptions.append(e) def join(self): """Joins the run thread, if it was started""" if self._script_thread is not None: self._script_thread.join() def deltas(self): """Returns the delta messages in our ReportQueue""" return [msg.delta for msg in self.report_queue._queue if msg.HasField("delta")] def get_widget_id(self, widget_type, label): """Returns the id of the widget with the specified type and label""" for delta in self.deltas(): new_element = getattr(delta, "new_element", None) widget = getattr(new_element, widget_type, None) widget_label = getattr(widget, "label", None) if widget_label == label: return widget.id return None
def test_rerun_data_coalescing(self): """Test that multiple RERUN requests get coalesced with expected values. (This is similar to widgets_test.test_coalesce_widget_states - it's testing the same thing, but through the ScriptEventQueue interface.) """ queue = ScriptRequestQueue() states = WidgetStates() _create_widget('trigger', states).trigger_value = True _create_widget('int', states).int_value = 123 queue.enqueue(ScriptRequest.RERUN, RerunData(argv=None, widget_state=states)) states = WidgetStates() _create_widget('trigger', states).trigger_value = False _create_widget('int', states).int_value = 456 queue.enqueue(ScriptRequest.RERUN, RerunData(argv=None, widget_state=states)) event, data = queue.dequeue() self.assertEqual(event, ScriptRequest.RERUN) widgets = Widgets() widgets.set_state(data.widget_state) # Coalesced triggers should be True if either the old or # new value was True self.assertEqual(True, widgets.get_widget_value('trigger')) # Other widgets should have their newest value self.assertEqual(456, widgets.get_widget_value('int')) # We should have no more events self.assertEqual((None, None), queue.dequeue(), 'Expected empty event queue') # Test that we can coalesce if previous widget state is None queue.enqueue(ScriptRequest.RERUN, RerunData(argv=None, widget_state=None)) queue.enqueue(ScriptRequest.RERUN, RerunData(argv=None, widget_state=None)) states = WidgetStates() _create_widget('int', states).int_value = 789 queue.enqueue(ScriptRequest.RERUN, RerunData(argv=None, widget_state=states)) event, data = queue.dequeue() widgets = Widgets() widgets.set_state(data.widget_state) self.assertEqual(event, ScriptRequest.RERUN) self.assertEqual(789, widgets.get_widget_value('int')) # We should have no more events self.assertEqual((None, None), queue.dequeue(), 'Expected empty event queue') # Test that we can coalesce if our *new* widget state is None states = WidgetStates() _create_widget('int', states).int_value = 101112 queue.enqueue(ScriptRequest.RERUN, RerunData(argv=None, widget_state=states)) queue.enqueue(ScriptRequest.RERUN, RerunData(argv=None, widget_state=None)) event, data = queue.dequeue() widgets = Widgets() widgets.set_state(data.widget_state) self.assertEqual(event, ScriptRequest.RERUN) self.assertEqual(101112, widgets.get_widget_value('int')) # We should have no more events self.assertEqual((None, None), queue.dequeue(), 'Expected empty event queue')
class ReportSession(object): """ Contains session data for a single "user" of an active report (that is, a connected browser tab). Each ReportSession has its own Report, root DeltaGenerator, ScriptRunner, and widget state. A ReportSession is attached to each thread involved in running its Report. """ _next_id = 0 def __init__(self, ioloop, script_path, command_line): """Initialize the ReportSession. Parameters ---------- ioloop : tornado.ioloop.IOLoop The Tornado IOLoop that we're running within. script_path : str Path of the Python file from which this report is generated. command_line : str Command line as input by the user. """ # Each ReportSession gets a unique ID self.id = ReportSession._next_id ReportSession._next_id += 1 self._ioloop = ioloop self._report = Report(script_path, command_line) self._state = ReportSessionState.REPORT_NOT_RUNNING self._uploaded_file_mgr = UploadedFileManager() self._widget_states = WidgetStates() self._local_sources_watcher = LocalSourcesWatcher( self._report, self._on_source_file_changed) self._sent_initialize_message = False self._storage = None self._maybe_reuse_previous_run = False self._run_on_save = config.get_option("server.runOnSave") # The ScriptRequestQueue is the means by which we communicate # with the active ScriptRunner. self._script_request_queue = ScriptRequestQueue() self._scriptrunner = None LOGGER.debug("ReportSession initialized (id=%s)", self.id) def flush_browser_queue(self): """Clears the report queue and returns the messages it contained. The Server calls this periodically to deliver new messages to the browser connected to this report. Returns ------- list[ForwardMsg] The messages that were removed from the queue and should be delivered to the browser. """ return self._report.flush_browser_queue() def shutdown(self): """Shuts down the ReportSession. It's an error to use a ReportSession after it's been shut down. """ if self._state != ReportSessionState.SHUTDOWN_REQUESTED: LOGGER.debug("Shutting down (id=%s)", self.id) self._uploaded_file_mgr.delete_all_files() # Shut down the ScriptRunner, if one is active. # self._state must not be set to SHUTDOWN_REQUESTED until # after this is called. if self._scriptrunner is not None: self._enqueue_script_request(ScriptRequest.SHUTDOWN) self._state = ReportSessionState.SHUTDOWN_REQUESTED self._local_sources_watcher.close() def enqueue(self, msg): """Enqueues a new ForwardMsg to our browser queue. This can be called on both the main thread and a ScriptRunner run thread. Parameters ---------- msg : ForwardMsg The message to enqueue """ if not config.get_option("client.displayEnabled"): return # Avoid having two maybe_handle_execution_control_request running on # top of each other when tracer is installed. This leads to a lock # contention. if not config.get_option("runner.installTracer"): # If we have an active ScriptRunner, signal that it can handle an # execution control request. (Copy the scriptrunner reference to # avoid it being unset from underneath us, as this function can be # called outside the main thread.) scriptrunner = self._scriptrunner if scriptrunner is not None: scriptrunner.maybe_handle_execution_control_request() self._report.enqueue(msg) def enqueue_exception(self, e): """Enqueues an Exception message. Parameters ---------- e : BaseException """ # This does a few things: # 1) Clears the current report in the browser. # 2) Marks the current report as "stopped" in the browser. # 3) HACK: Resets any script params that may have been broken (e.g. the # command-line when rerunning with wrong argv[0]) self._on_scriptrunner_event( ScriptRunnerEvent.SCRIPT_STOPPED_WITH_SUCCESS) self._on_scriptrunner_event(ScriptRunnerEvent.SCRIPT_STARTED) self._on_scriptrunner_event( ScriptRunnerEvent.SCRIPT_STOPPED_WITH_SUCCESS) msg = ForwardMsg() msg.metadata.delta_id = 0 exception_proto.marshall(msg.delta.new_element.exception, e) self.enqueue(msg) def request_rerun(self, widget_state=None): """Signal that we're interested in running the script. If the script is not already running, it will be started immediately. Otherwise, a rerun will be requested. Parameters ---------- widget_state : dict | None The widget state dictionary to run the script with, or None to use the widget state from the previous run of the script. """ self._enqueue_script_request(ScriptRequest.RERUN, RerunData(widget_state)) def _on_source_file_changed(self): """One of our source files changed. Schedule a rerun if appropriate.""" if self._run_on_save: self.request_rerun() else: self._enqueue_file_change_message() def _clear_queue(self): self._report.clear() def _on_scriptrunner_event(self, event, exception=None, widget_states=None): """Called when our ScriptRunner emits an event. This is *not* called on the main thread. Parameters ---------- event : ScriptRunnerEvent exception : BaseException | None An exception thrown during compilation. Set only for the SCRIPT_STOPPED_WITH_COMPILE_ERROR event. widget_states : streamlit.proto.Widget_pb2.WidgetStates | None The ScriptRunner's final WidgetStates. Set only for the SHUTDOWN event. """ LOGGER.debug("OnScriptRunnerEvent: %s", event) prev_state = self._state if event == ScriptRunnerEvent.SCRIPT_STARTED: if self._state != ReportSessionState.SHUTDOWN_REQUESTED: self._state = ReportSessionState.REPORT_IS_RUNNING if config.get_option("server.liveSave"): # Enqueue into the IOLoop so it runs without blocking AND runs # on the main thread. self._ioloop.spawn_callback(self._save_running_report) self._clear_queue() self._maybe_enqueue_initialize_message() self._enqueue_new_report_message() elif (event == ScriptRunnerEvent.SCRIPT_STOPPED_WITH_SUCCESS or event == ScriptRunnerEvent.SCRIPT_STOPPED_WITH_COMPILE_ERROR): if self._state != ReportSessionState.SHUTDOWN_REQUESTED: self._state = ReportSessionState.REPORT_NOT_RUNNING script_succeeded = event == ScriptRunnerEvent.SCRIPT_STOPPED_WITH_SUCCESS self._enqueue_report_finished_message( ForwardMsg.FINISHED_SUCCESSFULLY if script_succeeded else ForwardMsg.FINISHED_WITH_COMPILE_ERROR) if config.get_option("server.liveSave"): # Enqueue into the IOLoop so it runs without blocking AND runs # on the main thread. self._ioloop.spawn_callback(self._save_final_report_and_quit) if script_succeeded: # When a script completes successfully, we update our # LocalSourcesWatcher to account for any source code changes # that change which modules should be watched. (This is run on # the main thread, because LocalSourcesWatcher is not # thread safe.) self._ioloop.spawn_callback( self._local_sources_watcher.update_watched_modules) else: # When a script fails to compile, we send along the exception. from streamlit.elements import exception_proto msg = ForwardMsg() exception_proto.marshall( msg.session_event.script_compilation_exception, exception) self.enqueue(msg) elif event == ScriptRunnerEvent.SHUTDOWN: # When ScriptRunner shuts down, update our local reference to it, # and check to see if we need to spawn a new one. (This is run on # the main thread.) def on_shutdown(): self._widget_states = widget_states self._scriptrunner = None # Because a new ScriptEvent could have been enqueued while the # scriptrunner was shutting down, we check to see if we should # create a new one. (Otherwise, a newly-enqueued ScriptEvent # won't be processed until another event is enqueued.) self._maybe_create_scriptrunner() self._ioloop.spawn_callback(on_shutdown) # Send a message if our run state changed report_was_running = prev_state == ReportSessionState.REPORT_IS_RUNNING report_is_running = self._state == ReportSessionState.REPORT_IS_RUNNING if report_is_running != report_was_running: self._enqueue_session_state_changed_message() def _enqueue_session_state_changed_message(self): msg = ForwardMsg() msg.session_state_changed.run_on_save = self._run_on_save msg.session_state_changed.report_is_running = ( self._state == ReportSessionState.REPORT_IS_RUNNING) self.enqueue(msg) def _enqueue_file_change_message(self): LOGGER.debug("Enqueuing report_changed message (id=%s)", self.id) msg = ForwardMsg() msg.session_event.report_changed_on_disk = True self.enqueue(msg) def _maybe_enqueue_initialize_message(self): if self._sent_initialize_message: return self._sent_initialize_message = True msg = ForwardMsg() imsg = msg.initialize imsg.config.sharing_enabled = config.get_option( "global.sharingMode") != "off" imsg.config.gather_usage_stats = config.get_option( "browser.gatherUsageStats") imsg.config.max_cached_message_age = config.get_option( "global.maxCachedMessageAge") imsg.config.mapbox_token = config.get_option("mapbox.token") LOGGER.debug( "New browser connection: " "gather_usage_stats=%s, " "sharing_enabled=%s, " "max_cached_message_age=%s", imsg.config.gather_usage_stats, imsg.config.sharing_enabled, imsg.config.max_cached_message_age, ) imsg.environment_info.streamlit_version = __version__ imsg.environment_info.python_version = ".".join( map(str, sys.version_info)) imsg.session_state.run_on_save = self._run_on_save imsg.session_state.report_is_running = ( self._state == ReportSessionState.REPORT_IS_RUNNING) imsg.user_info.installation_id = __installation_id__ if Credentials.get_current().activation: imsg.user_info.email = Credentials.get_current().activation.email else: imsg.user_info.email = "" imsg.command_line = self._report.command_line self.enqueue(msg) def _enqueue_new_report_message(self): self._report.generate_new_id() msg = ForwardMsg() msg.new_report.id = self._report.report_id msg.new_report.name = self._report.name msg.new_report.script_path = self._report.script_path self.enqueue(msg) def _enqueue_report_finished_message(self, status): """Enqueues a report_finished ForwardMsg. Parameters ---------- status : ReportFinishedStatus """ msg = ForwardMsg() msg.report_finished = status self.enqueue(msg) def handle_rerun_script_request(self, command_line=None, widget_state=None, is_preheat=False): """Tells the ScriptRunner to re-run its report. Parameters ---------- command_line : str | None The new command line arguments to run the script with, or None to use its previous command line value. widget_state : WidgetStates | None The WidgetStates protobuf to run the script with, or None to use its previous widget states. is_preheat: boolean True if this ReportSession should run the script immediately, and then ignore the next rerun request if it matches the already-ran widget state. """ if is_preheat: self._maybe_reuse_previous_run = True # For next time. elif self._maybe_reuse_previous_run: # If this is a "preheated" ReportSession, reuse the previous run if # the widget state matches. But only do this one time ever. self._maybe_reuse_previous_run = False has_widget_state = (widget_state is not None and len(widget_state.widgets) > 0) if not has_widget_state: LOGGER.debug( "Skipping rerun since the preheated run is the same") return self.request_rerun(widget_state) def handle_upload_file(self, upload_file): self._uploaded_file_mgr.create_or_clear_file( widget_id=upload_file.widget_id, name=upload_file.name, size=upload_file.size, last_modified=upload_file.lastModified, chunks=upload_file.chunks, ) self.handle_rerun_script_request(widget_state=self._widget_states) def handle_upload_file_chunk(self, upload_file_chunk): progress = self._uploaded_file_mgr.process_chunk( widget_id=upload_file_chunk.widget_id, index=upload_file_chunk.index, data=upload_file_chunk.data, ) if progress == 1: self.handle_rerun_script_request(widget_state=self._widget_states) def handle_delete_uploaded_file(self, delete_uploaded_file): self._uploaded_file_mgr.delete_file( widget_id=delete_uploaded_file.widget_id) self.handle_rerun_script_request(widget_state=self._widget_states) def handle_stop_script_request(self): """Tells the ScriptRunner to stop running its report.""" self._enqueue_script_request(ScriptRequest.STOP) def handle_clear_cache_request(self): """Clears this report's cache. Because this cache is global, it will be cleared for all users. """ # Setting verbose=True causes clear_cache to print to stdout. # Since this command was initiated from the browser, the user # doesn't need to see the results of the command in their # terminal. caching.clear_cache() def handle_set_run_on_save_request(self, new_value): """Changes our run_on_save flag to the given value. The browser will be notified of the change. Parameters ---------- new_value : bool New run_on_save value """ self._run_on_save = new_value self._enqueue_session_state_changed_message() def _enqueue_script_request(self, request, data=None): """Enqueue a ScriptEvent into our ScriptEventQueue. If a script thread is not already running, one will be created to handle the event. Parameters ---------- request : ScriptRequest The type of request. data : Any Data associated with the request, if any. """ if self._state == ReportSessionState.SHUTDOWN_REQUESTED: LOGGER.warning("Discarding %s request after shutdown" % request) return self._script_request_queue.enqueue(request, data) self._maybe_create_scriptrunner() def _maybe_create_scriptrunner(self): """Create a new ScriptRunner if we have unprocessed script requests. This is called every time a ScriptRequest is enqueued, and also after a ScriptRunner shuts down, in case new requests were enqueued during its termination. This function should only be called on the main thread. """ if (self._state == ReportSessionState.SHUTDOWN_REQUESTED or self._scriptrunner is not None or not self._script_request_queue.has_request): return # Create the ScriptRunner, attach event handlers, and start it self._scriptrunner = ScriptRunner( report=self._report, enqueue_forward_msg=self.enqueue, widget_states=self._widget_states, request_queue=self._script_request_queue, uploaded_file_mgr=self._uploaded_file_mgr, ) self._scriptrunner.on_event.connect(self._on_scriptrunner_event) self._scriptrunner.start() @tornado.gen.coroutine def handle_save_request(self, ws): """Save serialized version of report deltas to the cloud. "Progress" ForwardMsgs will be sent to the client during the upload. These messages are sent "out of band" - that is, they don't get enqueued into the ReportQueue (because they're not part of the report). Instead, they're written directly to the report's WebSocket. Parameters ---------- ws : _BrowserWebSocketHandler The report's websocket handler. """ @tornado.gen.coroutine def progress(percent): progress_msg = ForwardMsg() progress_msg.upload_report_progress = percent yield ws.write_message(serialize_forward_msg(progress_msg), binary=True) # Indicate that the save is starting. try: yield progress(0) url = yield self._save_final_report(progress) # Indicate that the save is done. progress_msg = ForwardMsg() progress_msg.report_uploaded = url yield ws.write_message(serialize_forward_msg(progress_msg), binary=True) except Exception as e: # Horrible hack to show something if something breaks. err_msg = "%s: %s" % (type(e).__name__, str(e) or "No further details.") progress_msg = ForwardMsg() progress_msg.report_uploaded = err_msg yield ws.write_message(serialize_forward_msg(progress_msg), binary=True) LOGGER.warning("Failed to save report:", exc_info=e) @tornado.gen.coroutine def _save_running_report(self): files = self._report.serialize_running_report_to_files() url = yield self._get_storage().save_report_files( self._report.report_id, files) if config.get_option("server.liveSave"): url_util.print_url("Saved running app", url) raise tornado.gen.Return(url) @tornado.gen.coroutine def _save_final_report(self, progress_coroutine=None): files = self._report.serialize_final_report_to_files() url = yield self._get_storage().save_report_files( self._report.report_id, files, progress_coroutine) if config.get_option("server.liveSave"): url_util.print_url("Saved final app", url) raise tornado.gen.Return(url) @tornado.gen.coroutine def _save_final_report_and_quit(self): yield self._save_final_report() self._ioloop.stop() def _get_storage(self): if self._storage is None: sharing_mode = config.get_option("global.sharingMode") if sharing_mode == "s3": self._storage = S3Storage() elif sharing_mode == "file": self._storage = FileStorage() else: raise RuntimeError("Unsupported sharing mode '%s'" % sharing_mode) return self._storage