diff options
author | Angelos Evripiotis <jevripiotis@bloomberg.net> | 2019-10-17 18:14:11 +0100 |
---|---|---|
committer | bst-marge-bot <marge-bot@buildstream.build> | 2019-10-29 14:17:17 +0000 |
commit | 5f48e8aa6b6a1ddf6a9d5d93fac74207f0946b30 (patch) | |
tree | df16e7a8a1c8ff662cdb52282546e9eaec6c6405 /src | |
parent | a274fb067d694de967b460df58a8e15d78f2b032 (diff) | |
download | buildstream-5f48e8aa6b6a1ddf6a9d5d93fac74207f0946b30.tar.gz |
job pickling: also pickle global state in node.pyx
Diffstat (limited to 'src')
-rw-r--r-- | src/buildstream/_scheduler/jobs/job.py | 37 | ||||
-rw-r--r-- | src/buildstream/_scheduler/jobs/jobpickler.py | 27 | ||||
-rw-r--r-- | src/buildstream/node.pyx | 27 |
3 files changed, 72 insertions, 19 deletions
diff --git a/src/buildstream/_scheduler/jobs/job.py b/src/buildstream/_scheduler/jobs/job.py index 940b7d2cd..b36167a8f 100644 --- a/src/buildstream/_scheduler/jobs/job.py +++ b/src/buildstream/_scheduler/jobs/job.py @@ -33,9 +33,9 @@ import traceback from ..._exceptions import ImplError, BstError, set_last_task_error, SkipJob from ..._message import Message, MessageType, unconditional_messages from ...types import FastEnum -from ... import _signals, utils +from ... import _signals, utils, node -from .jobpickler import pickle_child_job +from .jobpickler import pickle_child_job_data # Return code values shutdown of job handling child processes @@ -87,6 +87,27 @@ class _MessageType(FastEnum): SUBCLASS_CUSTOM_MESSAGE = 5 +# _pickle_child_job() +# +# Perform the special case pickling required to pickle a child job for +# unpickling in a child process. +# +# Args: +# child_job (ChildJob): The job to pickle. +# projects (List[Project]): The list of loaded projects, so we can get the +# relevant factories. +# +def _pickle_child_job(child_job, projects): + # Note that we need to consider all the state of the program that's + # necessary for the job, this includes e.g. the global state of the node + # module. + node_module_state = node._get_state_for_pickling() + return pickle_child_job_data( + (child_job, node_module_state), + projects, + ) + + # _do_pickled_child_job() # # Unpickle the supplied 'pickled' job and call 'child_action' on it. @@ -95,14 +116,14 @@ class _MessageType(FastEnum): # such it will fixup any globals to be in the expected state. # # Args: -# pickled (BytesIO): The pickled job to execute. +# pickled (BytesIO): The pickled data, and job to execute. # *child_args (any) : Any parameters to be passed to `child_action`. # def _do_pickled_child_job(pickled, *child_args): - utils._is_main_process = _not_main_process - child_job = pickle.load(pickled) + child_job, node_module_state = pickle.load(pickled) + node._set_state_from_pickling(node_module_state) return child_job.child_action(*child_args) @@ -212,8 +233,10 @@ class Job(): ) if self._scheduler.context.platform.does_multiprocessing_start_require_pickling(): - pickled = pickle_child_job( - child_job, self._scheduler.context.get_projects()) + pickled = _pickle_child_job( + child_job, + self._scheduler.context.get_projects(), + ) self._process = Process( target=_do_pickled_child_job, args=[pickled, self._queue], diff --git a/src/buildstream/_scheduler/jobs/jobpickler.py b/src/buildstream/_scheduler/jobs/jobpickler.py index d514e87ae..9cc8c5e06 100644 --- a/src/buildstream/_scheduler/jobs/jobpickler.py +++ b/src/buildstream/_scheduler/jobs/jobpickler.py @@ -42,18 +42,21 @@ _PROTO_CLASS_TO_NAME = { } -# pickle_child_job() +# pickle_child_job_data() # # Perform the special case pickling required to pickle a child job for # unpickling in a child process. # -# Note that we don't need an `unpickle_child_job`, as regular `pickle.load()` -# will do everything required. +# Note that this just enables the pickling of things that contain ChildJob-s, +# the thing to be pickled doesn't have to be a ChildJob. +# +# Note that we don't need an `unpickle_child_job_data`, as regular +# `pickle.load()` will do everything required. # # Args: -# child_job (ChildJob): The job to be pickled. -# projects (List[Project]): The list of loaded projects, so we can get the -# relevant factories. +# child_job_data (ChildJob): The job to be pickled. +# projects (List[Project]): The list of loaded projects, so we can get the +# relevant factories. # # Returns: # An `io.BytesIO`, with the pickled contents of the ChildJob and everything it @@ -77,7 +80,7 @@ _PROTO_CLASS_TO_NAME = { # below. Some state in plugins is not necessary for child jobs, and comes # with a heavy cost; we also need to remove this before pickling. # -def pickle_child_job(child_job, projects): +def pickle_child_job_data(child_job_data, projects): factory_list = [ factory @@ -97,8 +100,8 @@ def pickle_child_job(child_job, projects): for cls, _ in factory.all_loaded_plugins() } - data = io.BytesIO() - pickler = pickle.Pickler(data) + pickled_data = io.BytesIO() + pickler = pickle.Pickler(pickled_data) pickler.dispatch_table = copyreg.dispatch_table.copy() def reduce_plugin(plugin): @@ -111,10 +114,10 @@ def pickle_child_job(child_job, projects): pickler.dispatch_table[Loader] = _reduce_object pickler.dispatch_table[Messenger] = _reduce_object - pickler.dump(child_job) - data.seek(0) + pickler.dump(child_job_data) + pickled_data.seek(0) - return data + return pickled_data def _reduce_object(instance): diff --git a/src/buildstream/node.pyx b/src/buildstream/node.pyx index bf6ce3a32..58fd0f33d 100644 --- a/src/buildstream/node.pyx +++ b/src/buildstream/node.pyx @@ -1596,6 +1596,33 @@ def _reset_global_state(): __counter = 0 +# _get_state_for_pickling() +# +# This gets the global variables necessary to preserve in a child process when +# e.g. running a ChildJob. Things that are pickled from the parent process to +# the child process will expect this module to be in the same state as in the +# parent. +# +# Returns: +# (object): The state to supply to a call of _get_state_for_pickling(). +# +def _get_state_for_pickling(): + return __FILE_LIST, __counter + + +# _set_state_from_pickling() +# +# This restores the global variables saved from _get_state_for_pickling(). See +# that function for more details. +# +# Args: +# state (object): The result of calling _get_state_for_pickling(). +# +def _set_state_from_pickling(state): + global __FILE_LIST, __counter + __FILE_LIST, __counter = state + + ############################################################# # Module local helper Methods # ############################################################# |