Skip to content

Commit

Permalink
fix thread local (#86)
Browse files Browse the repository at this point in the history
  • Loading branch information
zhreshold authored Jul 21, 2020
1 parent fa47442 commit 4fa8999
Show file tree
Hide file tree
Showing 2 changed files with 29 additions and 0 deletions.
11 changes: 11 additions & 0 deletions python/decord/bridge/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,9 +17,12 @@

_CURRENT_BRIDGE = threading.local()
_CURRENT_BRIDGE.type = 'native'
_GLOBAL_BRIDGE_TYPE = 'native' # child threads will derive from the global type but not overwrite

def reset_bridge():
_CURRENT_BRIDGE.type = 'native'
if threading.current_thread().name == 'MainThread':
_GLOBAL_BRIDGE_TYPE = 'native'

def set_bridge(new_bridge):
assert isinstance(new_bridge, str), (
Expand All @@ -28,11 +31,17 @@ def set_bridge(new_bridge):
"valid bridges: {}".format(_BRIDGE_TYPES.keys()))
global _CURRENT_BRIDGE
_CURRENT_BRIDGE.type = new_bridge
if threading.current_thread().name == 'MainThread':
_GLOBAL_BRIDGE_TYPE = new_bridge

def bridge_out(native_arr):
if not hasattr(_CURRENT_BRIDGE, 'type'):
_CURRENT_BRIDGE.type = _GLOBAL_BRIDGE_TYPE
return _BRIDGE_TYPES[_CURRENT_BRIDGE.type][0](native_arr)

def bridge_in(arr):
if not hasattr(_CURRENT_BRIDGE, 'type'):
_CURRENT_BRIDGE.type = _GLOBAL_BRIDGE_TYPE
return _BRIDGE_TYPES[_CURRENT_BRIDGE.type][1](arr)

class _BridgeScope(object):
Expand All @@ -42,6 +51,8 @@ def __init__(self, bridge_type='native'):

def __enter__(self):
global _CURRENT_BRIDGE
if not hasattr(_CURRENT_BRIDGE, 'type'):
_CURRENT_BRIDGE.type = _GLOBAL_BRIDGE_TYPE
try:
self._prev = _CURRENT_BRIDGE.type
except AttributeError:
Expand Down
18 changes: 18 additions & 0 deletions tests/python/unittests/test_bridges.py
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,24 @@ def test_tvm_bridge():
except ImportError:
print('Skip test tvm bridge as tvm is not found')

def test_threaded_bridge():
# issue #85
from decord import cpu, gpu
from multiprocessing.dummy import Pool as ThreadPool

video_paths = [
os.path.expanduser('~/Dev/decord/examples/flipping_a_pancake.mkv'), #list of paths to video
]

def process_path(path):
vr = VideoReader(path, ctx=cpu(0))

for i in range(len(vr)):
frame = vr[i]

pool = ThreadPool(1)
pool.map(process_path, video_paths)

if __name__ == '__main__':
import nose
nose.runmodule()

0 comments on commit 4fa8999

Please sign in to comment.