[AutoScheduler] Make SearchTask and ComputeDAG serializable #6842
Conversation
| return {"compute": SaveJSON(self.compute), "sche": SaveJSON(self.sche)} | ||
|
|
||
| def __setstate__(self, state): | ||
| self.compute = LoadJSON(state["compute"]) # pylint: disable=assignment-from-no-return |
There was a problem hiding this comment.
@comaniac Do we need to explicitly call SaveJSON, LoadJSON here?
I find that you don't call these functions in the SearchTask::__getstate__
There was a problem hiding this comment.
We don't call them in SearchTask because all its members have setstate and getstate implemented correctly. This function provides the correct implementation for ConouteDAG and it will be called when processing ComputeDAG in SearchTask.
For Load/SaveJSON here, alternatively we can use pickle.loads/dumps and let them call Load/SaveJSON via Objects. However, this will introduce the dependency of pickle in this function, and will somehow trigger a bug in the unit test. You can test it by replacing Load/SaveJSON with pickle calls. In the unit test I add, you will find that after loading the DAG back, the placeholder A appears twice.
This PR adds required
__setstate__and__getstate__to makeSearchTaskandComputeDAGserializable. The idea is to serialize the list of tensors or the schedule used to construct a ComputeDAG, so that we can create an identical ComputeDAG bypickle.loads.cc @merrymercy @tqchen @junrushao1994