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:
@@ -106,6 +114,7 @@ def __init__(
         save_interval: int | float = 300,
         runner_kwargs: dict[str, Any] | None = None,
         goal: GoalTypes = None,
+        custom_load: Callable[[adaptive.BaseLearner, str], None] | None = None,
     ) -> None:
         super().__init__()
         self.job_names = job_names
@@ -127,6 +136,7 @@ def __init__(
         self.save_interval = save_interval
         self.runner_kwargs = runner_kwargs
         self.goal = goal
+        self.custom_load = custom_load
 
     @property
     def max_job_starts(self) -> int:
@@ -152,6 +162,7 @@ def _setup(self) -> None:
             dataframe_format=self.dataframe_format,
             goal=self.goal,
             loky_start_method=self.loky_start_method,
+            custom_load=self.custom_load,
         )
         self.scheduler.write_job_script(name_prefix, options)
 

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(

adaptive_scheduler/_server_support/run_manager.py

@@ -64,6 +64,10 @@ class RunManager(BaseManager):
     runner_kwargs : dict, default: None
         Extra keyword argument to pass to the `adaptive.Runner`. Note that this dict
         will be serialized and pasted in the ``job_script``.
+    custom_load : callable, default: None
+        A function that is called to load the learner. It is called as
+        ``custom_load(learner, fname)``. If None, the learner is loaded with
+        ``learner.load(fname)``.
     url : str, default: None
         The url of the database manager, with the format
         ``tcp://ip_of_this_machine:allowed_port.``. If None, a correct url will be chosen.
@@ -149,7 +153,7 @@ class RunManager(BaseManager):
 
     """
 
-    def __init__(
+    def __init__(  # noqa: PLR0915
         self,
         scheduler: BaseScheduler,
         learners: list[adaptive.BaseLearner],
@@ -158,6 +162,7 @@ def __init__(
         goal: GoalTypes = None,
         check_goal_on_start: bool = True,
         runner_kwargs: dict | None = None,
+        custom_load: Callable[[adaptive.BaseLearner, str], None] | None = None,
         url: str | None = None,
         save_interval: int | float = 300,
         log_interval: int | float = 300,
@@ -185,6 +190,7 @@ def __init__(
         self.goal = smart_goal(goal, learners)
         self.check_goal_on_start = check_goal_on_start
         self.runner_kwargs = runner_kwargs
+        self.custom_load = custom_load
         self.save_interval = save_interval
         self.log_interval = log_interval
         self.job_name = job_name
@@ -264,6 +270,7 @@ def __init__(
             save_interval=self.save_interval,
             runner_kwargs=self.runner_kwargs,
             goal=self.goal,
+            custom_load=self.custom_load,
             **self.job_manager_kwargs,
         )
         self.kill_manager: KillManager | None