Skip to content

Commit

Permalink
[BugFix] Make program thread-local to support multi-threading (#338)
Browse files Browse the repository at this point in the history
* make program thread local

* doc string
  • Loading branch information
lingfanyu authored Jan 13, 2019
1 parent 19e1b14 commit ed1948b
Showing 1 changed file with 18 additions and 5 deletions.
23 changes: 18 additions & 5 deletions python/dgl/runtime/ir/program.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
from __future__ import absolute_import

from contextlib import contextmanager
import threading

from .registry import IR_REGISTRY

Expand Down Expand Up @@ -44,18 +45,30 @@ def pprint(self):
for exe in self.execs:
self.pprint_exe(exe)

class CurrentProgram(threading.local):
"""Thread local storage to keep the reference of current thread's program"""
def __init__(self):
super(CurrentProgram, self).__init__()
self.prog = None

def get_prog(self):
"""Get program"""
return self.prog

def set_prog(self, program):
"""Set program"""
self.prog = program

# current program
CURRENT_PROG = None
CURRENT_PROG = CurrentProgram()

def get_current_prog():
"""Get the current program."""
global CURRENT_PROG
return CURRENT_PROG
return CURRENT_PROG.get_prog()

def set_current_prog(program):
"""Set the current program."""
global CURRENT_PROG
CURRENT_PROG = program
CURRENT_PROG.set_prog(program)

@contextmanager
def prog():
Expand Down

0 comments on commit ed1948b

Please sign in to comment.