forked from streamlit/streamlit
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathscript_runner.py
More file actions
469 lines (365 loc) · 16 KB
/
Copy pathscript_runner.py
File metadata and controls
469 lines (365 loc) · 16 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
# Copyright 2018-2021 Streamlit Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import sys
import threading
import gc
from contextlib import contextmanager
from enum import Enum
from blinker import Signal
from streamlit import config
from streamlit import magic
from streamlit import source_util
from streamlit import util
from streamlit.error_util import handle_uncaught_app_exception
from streamlit.media_file_manager import media_file_manager
from streamlit.report_thread import ReportThread, ReportContext
from streamlit.report_thread import get_report_ctx
from streamlit.script_request_queue import ScriptRequest
from streamlit.logger import get_logger
from streamlit.proto.ClientState_pb2 import ClientState
from streamlit.widgets import WidgetStateManager
LOGGER = get_logger(__name__)
class ScriptRunnerEvent(Enum):
# The script started running.
SCRIPT_STARTED = "SCRIPT_STARTED"
# The script run stopped because of a compile error.
SCRIPT_STOPPED_WITH_COMPILE_ERROR = "SCRIPT_STOPPED_WITH_COMPILE_ERROR"
# The script run stopped because it ran to completion, or was
# interrupted by the user.
SCRIPT_STOPPED_WITH_SUCCESS = "SCRIPT_STOPPED_WITH_SUCCESS"
# The ScriptRunner is done processing the ScriptEventQueue and
# is shut down.
SHUTDOWN = "SHUTDOWN"
class ScriptRunner(object):
def __init__(
self,
session_id,
report,
enqueue_forward_msg,
client_state,
request_queue,
uploaded_file_mgr=None,
):
"""Initialize the ScriptRunner.
(The ScriptRunner won't start executing until start() is called.)
Parameters
----------
session_id : str
The ReportSession's id.
report : Report
The ReportSession's report.
client_state : streamlit.proto.ClientState_pb2.ClientState
The current state from the client (widgets and query params).
request_queue : ScriptRequestQueue
The queue that the ReportSession is publishing ScriptRequests to.
ScriptRunner will continue running until the queue is empty,
and then shut down.
uploaded_file_mgr : UploadedFileManager
The File manager to store the data uploaded by the file_uploader widget.
"""
self._session_id = session_id
self._report = report
self._enqueue_forward_msg = enqueue_forward_msg
self._request_queue = request_queue
self._uploaded_file_mgr = uploaded_file_mgr
self._client_state = client_state
self._widgets = WidgetStateManager()
self._widgets.set_state(client_state.widget_states)
self.on_event = Signal(
doc="""Emitted when a ScriptRunnerEvent occurs.
This signal is *not* emitted on the same thread that the
ScriptRunner was created on.
Parameters
----------
event : ScriptRunnerEvent
exception : BaseException | None
Our compile error. Set only for the
SCRIPT_STOPPED_WITH_COMPILE_ERROR event.
widget_states : streamlit.proto.WidgetStates_pb2.WidgetStates | None
The ScriptRunner's final WidgetStates. Set only for the
SHUTDOWN event.
"""
)
# Set to true when we process a SHUTDOWN request
self._shutdown_requested = False
# Set to true while we're executing. Used by
# maybe_handle_execution_control_request.
self._execing = False
# This is initialized in start()
self._script_thread = None
def __repr__(self) -> str:
return util.repr_(self)
def start(self):
"""Start a new thread to process the ScriptEventQueue.
This must be called only once.
"""
if self._script_thread is not None:
raise Exception("ScriptRunner was already started")
self._script_thread = ReportThread(
session_id=self._session_id,
enqueue=self._enqueue_forward_msg,
query_string=self._client_state.query_string,
widgets=self._widgets,
uploaded_file_mgr=self._uploaded_file_mgr,
target=self._process_request_queue,
name="ScriptRunner.scriptThread",
)
self._script_thread.start()
def _process_request_queue(self):
"""Process the ScriptRequestQueue and then exits.
This is run in a separate thread.
"""
LOGGER.debug("Beginning script thread")
while not self._shutdown_requested and self._request_queue.has_request:
request, data = self._request_queue.dequeue()
if request == ScriptRequest.STOP:
LOGGER.debug("Ignoring STOP request while not running")
elif request == ScriptRequest.SHUTDOWN:
LOGGER.debug("Shutting down")
self._shutdown_requested = True
elif request == ScriptRequest.RERUN:
self._run_script(data)
else:
raise RuntimeError("Unrecognized ScriptRequest: %s" % request)
# Send a SHUTDOWN event before exiting. This includes the widget values
# as they existed after our last successful script run, which the
# ReportSession will pass on to the next ScriptRunner that gets
# created.
client_state = ClientState()
client_state.query_string = self._client_state.query_string
self._widgets.marshall(client_state)
self.on_event.send(ScriptRunnerEvent.SHUTDOWN, client_state=client_state)
def _is_in_script_thread(self):
"""True if the calling function is running in the script thread"""
return self._script_thread == threading.current_thread()
def maybe_handle_execution_control_request(self):
if not self._is_in_script_thread():
# We can only handle execution_control_request if we're on the
# script execution thread. However, it's possible for deltas to
# be enqueued (and, therefore, for this function to be called)
# in separate threads, so we check for that here.
return
if not self._execing:
# If the _execing flag is not set, we're not actually inside
# an exec() call. This happens when our script exec() completes,
# we change our state to STOPPED, and a statechange-listener
# enqueues a new ForwardEvent
return
# Pop the next request from our queue.
request, data = self._request_queue.dequeue()
if request is None:
return
LOGGER.debug("Received ScriptRequest: %s", request)
if request == ScriptRequest.STOP:
raise StopException()
elif request == ScriptRequest.SHUTDOWN:
self._shutdown_requested = True
raise StopException()
elif request == ScriptRequest.RERUN:
raise RerunException(data)
else:
raise RuntimeError("Unrecognized ScriptRequest: %s" % request)
def _install_tracer(self):
"""Install function that runs before each line of the script."""
def trace_calls(frame, event, arg):
self.maybe_handle_execution_control_request()
return trace_calls
# Python interpreters are not required to implement sys.settrace.
if hasattr(sys, "settrace"):
sys.settrace(trace_calls)
@contextmanager
def _set_execing_flag(self):
"""A context for setting the ScriptRunner._execing flag.
Used by maybe_handle_execution_control_request to ensure that
we only handle requests while we're inside an exec() call
"""
if self._execing:
raise RuntimeError("Nested set_execing_flag call")
self._execing = True
try:
yield
finally:
self._execing = False
def _run_script(self, rerun_data):
"""Run our script.
Parameters
----------
rerun_data: RerunData
The RerunData to use.
"""
assert self._is_in_script_thread()
LOGGER.debug("Running script %s", rerun_data)
# Reset DeltaGenerators, widgets, media files.
media_file_manager.clear_session_files()
ctx = get_report_ctx()
if ctx is None:
# This should never be possible on the script_runner thread.
raise RuntimeError(
"ScriptRunner thread has a null ReportContext. Something has gone very wrong!"
)
ctx.reset(query_string=rerun_data.query_string)
self.on_event.send(ScriptRunnerEvent.SCRIPT_STARTED)
# Compile the script. Any errors thrown here will be surfaced
# to the user via a modal dialog in the frontend, and won't result
# in their previous report disappearing.
try:
with source_util.open_python_file(self._report.script_path) as f:
filebody = f.read()
if config.get_option("runner.magicEnabled"):
filebody = magic.add_magic(filebody, self._report.script_path)
code = compile(
filebody,
# Pass in the file path so it can show up in exceptions.
self._report.script_path,
# We're compiling entire blocks of Python, so we need "exec"
# mode (as opposed to "eval" or "single").
mode="exec",
# Don't inherit any flags or "future" statements.
flags=0,
dont_inherit=1,
# Use the default optimization options.
optimize=-1,
)
except BaseException as e:
# We got a compile error. Send an error event and bail immediately.
LOGGER.debug("Fatal script error: %s" % e)
self.on_event.send(
ScriptRunnerEvent.SCRIPT_STOPPED_WITH_COMPILE_ERROR, exception=e
)
return
# If we get here, we've successfully compiled our script. The next step
# is to run it. Errors thrown during execution will be shown to the
# user as ExceptionElements.
# Update the Widget object with the new widget_states.
# (The ReportContext has a reference to this object, so we just update it in-place)
if rerun_data.widget_states is not None:
self._widgets.set_state(rerun_data.widget_states)
if config.get_option("runner.installTracer"):
self._install_tracer()
# This will be set to a RerunData instance if our execution
# is interrupted by a RerunException.
rerun_with_data = None
try:
# Create fake module. This gives us a name global namespace to
# execute the code in.
module = _new_module("__main__")
# Install the fake module as the __main__ module. This allows
# the pickle module to work inside the user's code, since it now
# can know the module where the pickled objects stem from.
# IMPORTANT: This means we can't use "if __name__ == '__main__'" in
# our code, as it will point to the wrong module!!!
sys.modules["__main__"] = module
# Add special variables to the module's globals dict.
# Note: The following is a requirement for the CodeHasher to
# work correctly. The CodeHasher is scoped to
# files contained in the directory of __main__.__file__, which we
# assume is the main script directory.
module.__dict__["__file__"] = self._report.script_path
with modified_sys_path(self._report), self._set_execing_flag():
exec(code, module.__dict__)
except RerunException as e:
rerun_with_data = e.rerun_data
except StopException:
pass
except BaseException as e:
handle_uncaught_app_exception(e)
finally:
self._on_script_finished(ctx)
# Use _log_if_error() to make sure we never ever ever stop running the
# script without meaning to.
_log_if_error(_clean_problem_modules)
if rerun_with_data is not None:
self._run_script(rerun_with_data)
def _on_script_finished(self, ctx: ReportContext) -> None:
"""Called when our script finishes executing, even if it finished
early with an exception. We perform post-run cleanup here.
"""
self._widgets.reset_triggers()
self._widgets.cull_nonexistent(ctx.widget_ids_this_run.items())
# Signal that the script has finished. (We use SCRIPT_STOPPED_WITH_SUCCESS
# even if we were stopped with an exception.)
self.on_event.send(ScriptRunnerEvent.SCRIPT_STOPPED_WITH_SUCCESS)
# Delete expired files now that the script has run and files in use
# are marked as active.
media_file_manager.del_expired_files()
# Force garbage collection to run, to help avoid memory use building up
# This is usually not an issue, but sometimes GC takes time to kick in and
# causes apps to go over resource limits, and forcing it to run between
# script runs is low cost, since we aren't doing much work anyway.
if config.get_option("runner.postScriptGC"):
gc.collect(2)
class ScriptControlException(BaseException):
"""Base exception for ScriptRunner."""
pass
class StopException(ScriptControlException):
"""Silently stop the execution of the user's script."""
pass
class RerunException(ScriptControlException):
"""Silently stop and rerun the user's script."""
def __init__(self, rerun_data):
"""Construct a RerunException
Parameters
----------
rerun_data : RerunData
The RerunData that should be used to rerun the report
"""
self.rerun_data = rerun_data
def __repr__(self) -> str:
return util.repr_(self)
def _clean_problem_modules():
"""Some modules are stateful, so we have to clear their state."""
if "keras" in sys.modules:
try:
keras = sys.modules["keras"]
keras.backend.clear_session() # type: ignore[attr-defined]
except:
pass
if "matplotlib.pyplot" in sys.modules:
try:
plt = sys.modules["matplotlib.pyplot"]
plt.close("all") # type: ignore[attr-defined]
except:
pass
def _new_module(name):
"""Create a new module with the given name."""
import types
return types.ModuleType(name)
# Code modified from IPython (BSD license)
# Source: https://github.com/ipython/ipython/blob/master/IPython/utils/syspathcontext.py#L42
class modified_sys_path(object):
"""A context for prepending a directory to sys.path for a second."""
def __init__(self, report):
self._report = report
self._added_path = False
def __repr__(self) -> str:
return util.repr_(self)
def __enter__(self):
if self._report.script_path not in sys.path:
sys.path.insert(0, self._report.script_path)
self._added_path = True
def __exit__(self, type, value, traceback):
if self._added_path:
try:
sys.path.remove(self._report.script_path)
except ValueError:
pass
# Returning False causes any exceptions to be re-raised.
return False
# The reason this is not a decorator is because we want to make it clear at the
# calling location that this function is being used.
def _log_if_error(fn):
try:
fn()
except Exception as e:
LOGGER.warning(e)