diff options
-rw-r--r-- | src/buildstream/_scheduler/jobs/jobpickler.py | 27 |
1 files changed, 21 insertions, 6 deletions
diff --git a/src/buildstream/_scheduler/jobs/jobpickler.py b/src/buildstream/_scheduler/jobs/jobpickler.py index 0edf88c10..6a4a8213b 100644 --- a/src/buildstream/_scheduler/jobs/jobpickler.py +++ b/src/buildstream/_scheduler/jobs/jobpickler.py @@ -23,12 +23,25 @@ import io import pickle from ..._protos.buildstream.v2.artifact_pb2 import Artifact as ArtifactProto +from ..._protos.build.bazel.remote.execution.v2.remote_execution_pb2 import Digest as DigestProto # BuildStream toplevel imports from ..._loader import Loader from ..._messenger import Messenger +# Note that `str(type(proto_class))` results in `GeneratedProtocolMessageType` +# instead of the concrete type, so we come up with our own names here. +_NAME_TO_PROTO_CLASS = { + "artifact": ArtifactProto, + "digest": DigestProto, +} + +_PROTO_CLASS_TO_NAME = { + cls: name for name, cls in _NAME_TO_PROTO_CLASS.items() +} + + # pickle_child_job() # # Perform the special case pickling required to pickle a child job for @@ -87,7 +100,8 @@ def pickle_child_job(child_job, projects): pickler.dispatch_table[cls] = _reduce_plugin for cls in source_classes: pickler.dispatch_table[cls] = _reduce_plugin - pickler.dispatch_table[ArtifactProto] = _reduce_artifact_proto + pickler.dispatch_table[ArtifactProto] = _reduce_proto + pickler.dispatch_table[DigestProto] = _reduce_proto pickler.dispatch_table[Loader] = _reduce_object pickler.dispatch_table[Messenger] = _reduce_object @@ -103,14 +117,15 @@ def _reduce_object(instance): return (cls.__new__, (cls,), state) -def _reduce_artifact_proto(instance): - assert isinstance(instance, ArtifactProto) +def _reduce_proto(instance): + name = _PROTO_CLASS_TO_NAME[type(instance)] data = instance.SerializeToString() - return (_new_artifact_proto_from_reduction_args, (data,)) + return (_new_proto_from_reduction_args, (name, data)) -def _new_artifact_proto_from_reduction_args(data): - instance = ArtifactProto() +def _new_proto_from_reduction_args(name, data): + cls = _NAME_TO_PROTO_CLASS[name] + instance = cls() instance.ParseFromString(data) return instance |