Skip to content

Commit

Permalink
[AutoTVM] Fix database APIs (apache#3821)
Browse files Browse the repository at this point in the history
* [AutoTVM] Fix database APIs

* Refactor the byte conversion
  • Loading branch information
comaniac authored and wweic committed Sep 16, 2019
1 parent 68ab209 commit 118e151
Show file tree
Hide file tree
Showing 2 changed files with 22 additions and 11 deletions.
21 changes: 10 additions & 11 deletions python/tvm/autotvm/database.py
Original file line number Diff line number Diff line change
Expand Up @@ -117,12 +117,12 @@ def set(self, key, value):
self.db.set(key, value)

def get(self, key):
return self.db.get(key)
current = self.db.get(key)
return current.decode() if isinstance(current, bytes) else current

def load(self, inp, get_all=False):
current = self.get(measure_str_key(inp))
if current is not None:
current = str(current)
records = [decode(x) for x in current.split(RedisDatabase.MAGIC_SPLIT)]
results = [rec[1] for rec in records]
if get_all:
Expand All @@ -142,29 +142,31 @@ def save(self, inp, res, extend=False):

def filter(self, func):
"""
Dump all of the records for a particular target
Dump all of the records that match the given rule
Parameters
----------
func: callable
The signature of the function is bool (MeasureInput, Array of MeasureResult)
The signature of the function is (MeasureInput, [MeasureResult]) -> bool
Returns
-------
list of records (inp, result) matching the target
list of records in tuple (MeasureInput, MeasureResult) matching the rule
Examples
--------
get records for a target
>>> db.filter(lambda inp, resulst: "cuda" in inp.target.keys)
get records with errors
>>> db.filter(lambda inp, results: any(r.error_no != 0 for r in results))
"""
matched_records = list()
# may consider filtering in iterator in the future
for key in self.db:
for key in self.db.keys():
current = self.get(key)
try:
records = [decode(x) for x in current.spilt(RedisDatabase.MAGIC_SPLIT)]
except TypeError: # got a badly formatted/old format record
records = [decode(x) for x in current.split(RedisDatabase.MAGIC_SPLIT)]
except TypeError: # got a badly formatted/old format record
continue

inps, results = zip(*records)
Expand All @@ -190,8 +192,5 @@ def __init__(self):
def set(self, key, value):
self.db[key] = value

def get(self, key):
return self.db.get(key)

def flush(self):
self.db = {}
12 changes: 12 additions & 0 deletions tests/python/unittest/test_autotvm_database.py
Original file line number Diff line number Diff line change
Expand Up @@ -99,8 +99,20 @@ def test_db_latest_all():
assert encode(inp1, load4[1]) == encode(inp1, res2)
assert encode(inp1, load4[2]) == encode(inp1, res3)

def test_db_filter():
logging.info("test db filter ...")
records = get_sample_records(5)
_db = database.DummyDatabase()
_db.flush()
for inp, result in records:
_db.save(inp, result)

records = _db.filter(lambda inp, ress: any(r.costs[0] <= 2 for r in ress))
assert len(records) == 2

if __name__ == '__main__':
logging.basicConfig(level=logging.INFO)
test_save_load()
test_db_hash()
test_db_latest_all()
test_db_filter()

0 comments on commit 118e151

Please sign in to comment.