Allow a custom_load function

Author: basnijholt

adaptive_scheduler/_server_support/job_manager.py

@@ -10,6 +10,10 @@
 from .common import MaxRestartsReachedError, log
 
 if TYPE_CHECKING:
+    from typing import Callable
+
+    import adaptive
+
     from adaptive_scheduler.scheduler import BaseScheduler
     from adaptive_scheduler.utils import (
         _DATAFRAME_FORMATS,
@@ -31,6 +35,7 @@ def command_line_options(
     save_dataframe: bool = True,
     dataframe_format: _DATAFRAME_FORMATS = "parquet",
     loky_start_method: LOKY_START_METHODS = "loky",
+    custom_load: Callable[[adaptive.BaseLearner, str], None] | None = None,
 ) -> dict[str, Any]:
     """Return the command line options for the job_script."""
     if runner_kwargs is None:
@@ -49,6 +54,9 @@ def command_line_options(
         "--save-interval": save_interval,
         "--serialized-runner-kwargs": base64_runner_kwargs,
     }
+    if custom_load:
+        base64_custom_load = _serialize_to_b64(custom_load)
+        opts["--custom-load"] = base64_custom_load
     if scheduler.executor_type == "loky":
         opts["--loky-start-method"] = loky_start_method
     if save_dataframe:

adaptive_scheduler/_server_support/launcher.py

@@ -74,7 +74,6 @@ def _parse_args() -> argparse.Namespace:
     parser.add_argument("--job-id", action="store", type=str)
     parser.add_argument("--name", action="store", dest="name", type=str, required=True)
     parser.add_argument("--url", action="store", type=str, required=True)
-
     parser.add_argument("--save-dataframe", action="store_true", default=False)
     parser.add_argument(
         "--dataframe-format",
@@ -103,6 +102,7 @@ def _parse_args() -> argparse.Namespace:
         default=120,
     )
     parser.add_argument("--serialized-runner-kwargs", action="store", type=str)
+    parser.add_argument("--custom-load", action="store", type=str, default=None)
     return parser.parse_args()
 
 
@@ -121,8 +121,15 @@ def main() -> None:
     if args.executor_type == "process-pool":
         learner.function = WrappedFunction(learner.function)
 
+    if args.custom_load is not None:
+        custom_load = _deserialize_from_b64(args.custom_load)
+
     with suppress(Exception):
-        learner.load(fname)
+        if args.custom_load is not None:
+            custom_load(learner, fname)
+        else:
+            learner.load(fname)
+
     npoints_start = learner.npoints
 
     executor = _get_executor(