Add possibility to add a context function to run before each worker.

Author: fraimondo

README.md

@@ -95,8 +95,10 @@ These are the current HTCondor Backend parameters that can be set in the `parall
 - `log_dir_prefix`: Prefix for each of the HTCondor log files. If not specified, it will create a `logs` directory in the `initial_dir`.
 - `poll_interval`: Minimum time (in seconds) between polls to the HTCondor scheduler. Defaults to 5 seconds. A lower value will increase the load on the HTCondor collector as well as the filesystem, but will increase reactivity of the backend. A higher value will decrease the load, but the backend will take longer to react to changes in the queue. Important: there will be a poller for each nested parallel call, so the load on the system will be multiplied by the number of nested parallel calls.
 - `shared_data_dir`: Directory where the tasks and results will be saved. This directory must be shared across all the nodes in the HTCondor pool. If not specified, it will use a `joblib_htcondor_shared_data` directory inside the current working directory.
+- `delete_task_file_on_load`: Whether to delete the task file once the job enters RUNNING state. This can be useful when the tasks are data-intensive, and the shared disk space is limited. However, this will have the effect of not being able to load the task file in case of failure, thus making it impossible to retry the task (default False).
 - `worker_log_level`: Log level for the worker. Defaults to `INFO`.
 - `throttle`: Throttle the number of jobs submitted at once. If list, the first element is the throttle for the current level and the rest are for the nested levels (default None).
 - `batch_size`: Currently under development
 - `max_recursion_level`: Maximum recursion level for nested parallel calls. Defaults to 0 (no nested parallel calls allowed).
 - `export_metadata`: Export metadata to be used with the UI. Defaults to False.
+- `context_func`: A function to be called in the worker before running the actual function. This can be used to set up the context for the worker, for example by setting some global variables or importing some modules. The function will be serialized, so it can't rely on global variables. If None, no function will be called (default None).

joblib_htcondor/backend.py

@@ -358,6 +358,11 @@ class _HTCondorBackend(ParallelBackendBase):
         Export metadata to a file, to be used with the HTCondor Joblib Monitor.
         This increases the load on the filesystem considerably if the number
         of jobs is high and the duration is short (default False).
+    context_func : Callable or None, optional
+        A function to call in the worker before running the actual function.
+        This can be used to set global configuration variables in the worker,
+        for example. The function will be serialized, so it can't rely on
+        global variables. If None, no function will be called (default None).
 
     Raises
     ------
@@ -388,6 +393,7 @@ def __init__(  # noqa: C901
         batch_size: int = 1,
         max_recursion_level: int = 0,
         export_metadata: bool = False,
+        context_func: Optional[Callable] = None,
     ) -> None:
         super().__init__()
 
@@ -426,6 +432,7 @@ def __init__(  # noqa: C901
         self._batch_size = batch_size
         self._max_recursion_level = max_recursion_level
         self._export_metadata = export_metadata
+        self._context_func = context_func
 
         self._recursion_level = 0
         self._parent_uuid = None
@@ -458,6 +465,7 @@ def __init__(  # noqa: C901
         logger.debug(f"Batch Size: {self._batch_size}")
         logger.debug(f"Max recursion level: {self._max_recursion_level}")
         logger.debug(f"Export metadata: {self._export_metadata}")
+        logger.debug(f"Context function: {self._context_func}")
 
         logger.debug(f"Recursion level: {self._recursion_level}")
         logger.debug(f"Parent UUID: {self._parent_uuid}")
@@ -591,6 +599,7 @@ def get_nested_backend(self) -> tuple["ParallelBackendBase", int]:
             batch_size=self._batch_size,
             max_recursion_level=self._max_recursion_level,
             export_metadata=self._export_metadata,
+            context_func=self._context_func,
             recursion_level=self._recursion_level + 1,
             parent_uuid=self._this_batch_name,
         ), self._n_jobs
@@ -617,6 +626,7 @@ def __reduce__(self) -> tuple[Callable, tuple]:
                 self._batch_size,
                 self._max_recursion_level,
                 self._export_metadata,
+                self._context_func,
                 self._recursion_level,
                 self._parent_uuid,
             ),
@@ -793,6 +803,8 @@ def _submit(
 
         # Create the DelayedSubmission object
         ds = DelayedSubmission(func)
+        if self._context_func is not None:
+            ds.set_context_func(self._context_func)
         delete_file_param = (
             "--delete-file-on-load" if self._delete_task_file_on_load else ""
         )
@@ -1201,6 +1213,7 @@ def build(
         batch_size: int = 1,
         max_recursion_level: int = -1,
         export_metadata: bool = False,
+        context_func: Optional[Callable] = None,
         recursion_level: int = 0,
         parent_uuid: Optional[str] = None,
     ) -> _HTCondorBackend:
@@ -1261,6 +1274,13 @@ def build(
             Monitor. This increases the load on the filesystem considerably if
             the number of jobs is high and the duration is short
             (default False).
+        context_func : callable or None, optional
+            A function to be called in the worker before running the actual
+            function. This can be used to set up the context for the worker,
+            for example by setting some global variables or importing some
+            modules. The function will be serialized, so it can't rely on
+            global variables. If None, no function will be called (default
+            None).
         recursion_level : int, optional
             Recursion level of the backend. With each nested
             call, the recursion level increases by 1 (default 0).
@@ -1297,6 +1317,7 @@ def build(
             batch_size=batch_size,
             max_recursion_level=max_recursion_level,
             export_metadata=export_metadata,
+            context_func=context_func,
         )
         out._recursion_level = recursion_level
         out._parent_uuid = parent_uuid  # type: ignore

joblib_htcondor/delayed_submission.py

@@ -74,9 +74,12 @@ def __init__(
         self._done = False
         self._error = False
         self._done_timestamp = None
+        self.context_func = None
 
     def run(self) -> None:
         """Run the function with the arguments and store the result."""
+        if self.context_func is not None:
+            self.context_func()
         try:
             self._result = self.func(*self.args, **self.kwargs)  # type: ignore
         except BaseException as e:  # noqa: BLE001
@@ -88,6 +91,17 @@ def run(self) -> None:
         self._done_timestamp = datetime.now()
         self._done = True
 
+    def set_context_func(self, context_func: Callable) -> None:
+        """Set a context function to be called prior to the main function.
+
+        Parameters
+        ----------
+        context_func : callable
+            The context function to call before running the main function.
+
+        """
+        self.context_func = context_func
+
     def done(self) -> bool:
         """Return whether the function has been run.