[Feature Request] support for custom logic in submitit plugin's checkpoint function

[tesfaldet]

🚀 Feature Request

Currently, when preemption/timeout occurs, the submitit plugin resubmits the job with the same initial arguments. There is no clear way of adding custom logic within the checkpoint function of the plugin.

hydra/plugins/hydra_submitit_launcher/hydra_plugins/hydra_submitit_launcher/submitit_launcher.py

Lines 79 to 84 in 046837b

	def checkpoint(self, *args: Any, **kwargs: Any) -> Any:
	"""Resubmit the current callable at its current state with the same initial arguments."""
	# lazy import to ensure plugin discovery remains fast
	import submitit
	return submitit.helpers.DelayedSubmission(self, *args, **kwargs)

Motivation

My train loop saves (pytorch) checkpoints every N iterations. When preemption/timeout occurs and the submitit plugin resubmits the job, after the job starts it will continue from the most recently saved checkpoint, which could be several hundred (or thousands) of iterations old. The job will have to re-do training iterations that were already done before, which is a waste of time and resources.

It would be nice to have it such that once the submitit plugin's checkpoint function is called (in reaction to preemption/timeout), it could re-use the checkpoint save code that my train loop uses (or whatever logic I'd want, really), that way upon job restart it will start from a checkpoint at most 1 iteration old.

Pitch

Describe the solution you'd like
From what I see so far, the checkpoint function is exposed to the current callable's context (via self):

hydra/plugins/hydra_submitit_launcher/hydra_plugins/hydra_submitit_launcher/submitit_launcher.py

Lines 45 to 77 in 046837b

	def __call__(
	self,
	sweep_overrides: List[str],
	job_dir_key: str,
	job_num: int,
	job_id: str,
	singleton_state: Dict[type, Singleton],
	) -> JobReturn:
	# lazy import to ensure plugin discovery remains fast
	import submitit
	assert self.hydra_context is not None
	assert self.config is not None
	assert self.task_function is not None
	Singleton.set_state(singleton_state)
	setup_globals()
	sweep_config = self.hydra_context.config_loader.load_sweep_config(
	self.config, sweep_overrides
	)
	with open_dict(sweep_config.hydra.job) as job:
	# Populate new job variables
	job.id = submitit.JobEnvironment().job_id # type: ignore
	sweep_config.hydra.job.num = job_num
	return run_job(
	hydra_context=self.hydra_context,
	task_function=self.task_function,
	config=sweep_config,
	job_dir_key=job_dir_key,
	job_subdir_key="hydra.sweep.subdir",
	)

The callable returns a JobReturn object, which is where the task_function(...) (i.e., the train loop which contains all the relevant train loop data, such as the model, its state dict, the optimizer's state dict, etc.) is executed. Specifically, it's executed within run_job(...):

hydra/hydra/core/utils.py

Line 160 in 046837b

ret.return_value = task_function(task_cfg)

An issue I see here is that the task_function's context is inaccessible, so it's not possible to be able to pass through objects such as the model's current state, etc. Meaning that the checkpoint function also won't have access to that data. You'd need to figure out how to pass the task_function's context up to __call__(...) so that when checkpoint(...) is called it will have access to job-specific data and the user can inject custom logic (maybe through a yaml config with _target_ set to a save() function) that will be able to act upon this job-specific data, such as checkpointing/dumping it.

Here are some steps forward I think would make sense (correct me if I'm wrong):

1. It looks like you could provide a TaskFunction class with a __call__ function instead of a task_function from what I see here:

   hydra/hydra/main.py

   Lines 15 to 18 in b16baa7

   	def main(
   	config_path: Optional[str] = _UNSPECIFIED_,
   	config_name: Optional[str] = None,
   	) -> Callable[[TaskFunction], Any]:

   
   That way, you can treat each task_function as an instance and save your relevant train data using self.<whatever> = <whatever> in __call__. From there, you can access each task_function's context after it's been called. More importantly, run_job(...) will have access to its context.
2. Have run_job(...) optionally accept a Launcher context so that it could modify it. This would allow it to pass through task_function's context back up to BaseSubmititLauncher, which would make it available to its checkpoint function. The task_function's context can be saved in self.task_function_context.
3. Have the checkpoint function check for an on_job_preempt_or_timeout callback using _get_callbacks_for_run_job(...) and execute the callback while passing in self.task_function_context.

(Also, as a bonus feature request that could be tackled simultaneously with the above steps, it'd be nice to be able to pass a task_function's context to each of the callbacks in run_job(...). I could submit a separate feature request for this. The quick motivation is that you can make each task's wandb run object that was initialized within the task using self.run = wandb.init(...) accessible to callbacks so that you could make use of their Alerts interface to send Slack alerts. Currently there are ways to do this already but it involves reinitializing a wandb run within each callback separately, with the same run ID you used before, just to send an alert, which introduces unnecessary overhead due to the init process)

Describe alternatives you've considered
I've considered saving checkpoints every iteration but that's just time and space consuming...

Are you willing to open a pull request? (See CONTRIBUTING)
I could be willing to fork and try out the steps above then submit a PR, but I'd like to get a feel for what y'all think about this request before I try implementing it myself.

Additional context

None.

[Jasha10]

Hi @tesfaldet,

I'm not an expert in using submitit, but let me nevertheless share an idea:
If your main concern is being able to inject custom logic into the submitit launcher's checkpoint method, might this be achieved by defining a custom subclass of the submitit launcher plugin?

Here is what I am thinking:
Next to your python script (say my_app.py), you could create a file hydra_plugins/custom_submitit_logic.py which defines subclasses of the BaseSubmititLauncher class from Hydra's submitit plugin:

$ tree
.
├── conf
│   └── config.yaml
├── hydra_plugins
│   └── custom_submitit_logic.py
└── my_app.py

Here is the main python script:

$ cat my_app.py
import hydra
from omegaconf import DictConfig


@hydra.main(config_path="conf", config_name="config")
def app(cfg: DictConfig) -> None:
    print(cfg)


if __name__ == "__main__":
    app()

Here is the yaml configuration file:

$ cat conf/config.yaml
defaults:
  - override hydra/launcher: submitit_local
  # - override hydra/launcher: submitit_slurm
  - _self_

hydra:
  launcher:
    _target_: hydra_plugins.custom_submitit_logic.MyLocalLauncher
    # _target_: hydra_plugins.custom_submitit_logic.MySlurmLauncher

And here is the custom plugin subclassing BaseSubmititLauncher:

$ cat hydra_plugins/custom_submitit_logic.py
from typing import Any

from hydra_plugins.hydra_submitit_launcher.submitit_launcher import BaseSubmititLauncher


class MyBaseSubmititLauncher(BaseSubmititLauncher):
    def __init__(self, *args, **kwargs) -> None:
        print("INITIALIZING CUSTOM SUBMITIT LAUNCHER")
        ...
        super().__init__(*args, **kwargs)

    def checkpoint(self, *args: Any, **kwargs: Any) -> Any:
        """This method is a modified version of the BaseSubmititLauncher.checkpoint method"""

        #########################
        ### CUSTOM LOGIC HERE ###
        #########################

        super().checkpoint(*modified_args, **modified_kwargs)


class MyLocalLauncher(MyBaseSubmititLauncher):
    _EXECUTOR = "local"


class MySlurmLauncher(MyBaseSubmititLauncher):
    _EXECUTOR = "slurm"

When running my_app.py in --multirun mode, the call to print in MyBaseSubmititLauncher.__init__ verifies that the custom subclass is being used:

$ python my_app.py --multirun
INITIALIZING CUSTOM SUBMITIT LAUNCHER
[2022-02-18 17:15:33,990][HYDRA] Submitit 'local' sweep output dir : multirun/2022-02-18/17-15-33
[2022-02-18 17:15:33,991][HYDRA]        #0 :

Would subclassing the submitit plugin in this way be sufficiently flexible for your use-case?

[Jasha10]

I have tried passing a callable class instance (instead of passing a function) to hydra.main:

from typing import Any

import hydra
from omegaconf import DictConfig


class MyTaskFunction:
    def __init__(self, context: Any) -> None:
        self._context = context

    def __call__(self, cfg: DictConfig) -> None:
        print(cfg)

if __name__ == "__main__":
    my_task_function = MyTaskFunction(context=123)
    app = hydra.main(config_path="conf", config_name="config")(my_task_function)
    app()

It is not working currently, but should be easy to support.
For now, would you be able to use e.g. a global variable to store the context that is relevant for checkpointing?

[tesfaldet]

Thanks for the quick response @Jasha10.

Your suggestion is a great one and I think it could work.

Specifically, my suggestion of using a TaskFunction class with a __call__ function that's decorated with the main decorator combined with your suggestion could work. I see that you already tried the idea and that it wasn't working for you, but I think I may have an MVP below:

from omegaconf import DictConfig, OmegaConf
from typing import Any, Optional, Callable
from functools import wraps
from hydra.types import TaskFunction
from hydra.core.utils import JobReturn


TaskFunction = Callable[[Any], Any]

_UNSPECIFIED_: Any = object()


def main(config_path: Optional[str] = _UNSPECIFIED_,
         config_name: Optional[str] = None) -> Callable[[TaskFunction], Any]:
  def main_decorator(task_function: TaskFunction) -> Callable[[], None]:
    @wraps(task_function)
    def decorator_main(cfg_passthrough: Optional[DictConfig] = None, checkpoint: Optional[bool] = False) -> Any:
      if cfg_passthrough is not None:
        return task_function(cfg_passthrough)
      else:
        conf = OmegaConf.create({'config_name': config_name, 'config_path': config_path})
        launcher = BaseSubmititLauncher()
        launcher.setup(task_function=task_function, config=conf)
        ret = launcher.launch()
        if checkpoint:
          launcher.checkpoint()
        return ret
    return decorator_main
  return main_decorator


def run_job(task_function: TaskFunction,
            config: DictConfig):
  return task_function(config)


class BaseSubmititLauncher():
  def __init__(self) -> None:
    self.config: Optional[DictConfig] = None
    self.task_function: Optional[TaskFunction] = None

  def setup(
    self,
    *,
    task_function: TaskFunction,
    config: DictConfig
  ) -> None:
    self.config = config
    self.task_function = task_function

  def __call__(self) -> JobReturn:
    assert self.config is not None
    assert self.task_function is not None

    return run_job(
      task_function=self.task_function,
      config=self.config
    )

  def checkpoint(self, *args: Any, **kwargs: Any) -> Any:
    """Resubmit the current callable at its current state with the same initial arguments."""
    print(f'in checkpoint with model {self.task_function.model}')

  def launch(self) -> JobReturn:
    return self()


class Train(TaskFunction):
  def __init__(self):
    self.model = None
    self.config = None

  # main train loop code
  def __call__(self, cfg: DictConfig) -> float:
    self.model = 'PyTorchModel'
    self.config = cfg
    return 3.3

  def __str__(self) -> str:
    return f'Train(self.model = {self.model}, self.config = {self.config})'


if __name__ == '__main__':
  train = Train()
  print(f'in __main__ after instantiating Train object: {train}')

  wrapped = main(config_path='./conf', config_name='config')(train)
  print(f'in __main__ after wrapping Train object with main decorator: {train}')

  wrapped()
  print(f'in __main__ after calling train() with no config passthrough: {train}')

  wrapped(checkpoint=True)
  print(f'in __main__ after calling train() with no config passthrough and calling checkpoint: {train}')

  wrapped(OmegaConf.create({'foo': 'bar'}))
  print(f'in __main__ after calling train() with config passthrough: {train}')

outputs:

in __main__ after instantiating Train object: Train(self.model = None, self.config = None)
in __main__ after wrapping Train object with main decorator: Train(self.model = None, self.config = None)
in __main__ after calling train() with no config passthrough: Train(self.model = PyTorchModel, self.config = {'config_name': 'config', 'config_path': './conf'})
in checkpoint with model PyTorchModel
in __main__ after calling train() with no config passthrough and calling checkpoint: Train(self.model = PyTorchModel, self.config = {'config_name': 'config', 'config_path': './conf'})
in __main__ after calling train() with config passthrough: Train(self.model = PyTorchModel, self.config = {'foo': 'bar'})

What do you think? It follows a similar execution path as Hydra to show that this might be possible without any change to Hydra's code :)

Basically, since the launcher already has access to task_function, then by making task_function a Callable instance (i.e., a TaskFunction object) I could access its attributes. These attributes would be the relevant pytorch data. Combine that with a custom launcher that's a subclass of BaseSubmititLauncher that has its own checkpoint function that accesses task_function and its attributes and I'm pretty sure we're golden.

[tesfaldet]

I just took a closer look at your example and it's basically the same as mine 😅 I'm surprised it didn't work for you! I haven't tested mine with Hydra's actual main function since I made my own version of it, but it's fairly similar and it worked...

I'm not at my computer for the weekend so I can't test it out myself with Hydra's main at the moment but I'm curious what error you got.

[tesfaldet]

For now, would you be able to use e.g. a global variable to store the context that is relevant for checkpointing?

Correct me if I'm wrong, but doesn't the submitit plug-in launch a subprocess for each list of sweep params, with each subprocess executing your app's task function with those params? Wouldn't that mean that a global variable outside of the app's task function won't be accessible to each subprocess unless there's some IPC implemented?

[Jasha10]

I'm not at my computer for the weekend so I can't test it out myself with Hydra's main at the moment but I'm curious what error you got.

When I run the example from my previous comment, I get the following error:

$ python my_app.py
Traceback (most recent call last):
  File "my_app.py", line 17, in <module>
    app()
  File "/home/jasha10/hydra.git/hydra/main.py", line 52, in decorated_main
    config_name=config_name,
  File "/home/jasha10/hydra.git/hydra/_internal/utils.py", line 320, in _run_hydra
    ) = detect_calling_file_or_module_from_task_function(task_function)
  File "/home/jasha10/hydra.git/hydra/_internal/utils.py", line 51, in detect_calling_file_or_module_from_task_function
    calling_file = task_function.__code__.co_filename
AttributeError: 'MyTaskFunction' object has no attribute '__code__'

I am almost done with a PR to patch this so that a callable can be passed to hydra.main.
Edit: See PR #2044.

[Jasha10]

Wouldn't that mean that a global variable outside of the app's task function won't be accessible to each subprocess unless there's some IPC implemented?

Aah yes, good point. I haven't tried it myself, but I suspect you're right.

[tesfaldet]

@Jasha10 I believe the error you're experiencing occurs because you're passing the callable object and not its function, since only functions have a __code__ attribute. So perhaps the PR is not necessary?

import hydra
from omegaconf import DictConfig


class MyTaskFunction:
	def __init__(self) -> None:
		self._context = None

	def __call__(self, cfg: DictConfig) -> None:
		self._context = 123
		print(cfg, self._context)

if __name__ == "__main__":
	my_task_function = MyTaskFunction()
	print(my_task_function._context)
	app = hydra.main(config_path=None, config_name=None)(my_task_function.__call__)
	app()
	print(my_task_function._context)

outputs:

None
{} 123
123

[tesfaldet]

Nvm, the PR is super necessary haha. Although the above works, the problem with passing in my_task_function.__call__ is that you're not passing in the object, which is what's needed for the custom MyBaseSubmititLauncher's checkpoint function to access the context of my_task_function. So forget what I said!

[tesfaldet]

the problem with passing in my_task_function.call is that you're not passing in the object

I just realized that you could access the object through my_task_function.__call__.__self__. Welp, I'm gonna test this out right now with a custom MyBaseSubmititLauncher and report back :)

[tesfaldet]

I finished trying it out and it worked :) It was a hassle to get working properly for a couple of reasons:

1. Providing the path to the file that has MyBaseSubmititLauncher (e.g., src.launchers.MyBaseSubmititLauncher) as part of hydra.launcher._target_ only works if the path starts with hydra_plugins, so hydra_plugins.launchers.MyBaseSubmititLauncher. However, during instantiation, hydra complains that hydra_plugins is not a module. But if you put an __init__.py file there it'll clash with the import line from hydra_plugins.hydra_submitit_launcher.submitit_launcher import BaseSubmititLauncher that's inside launchers.py. So you have to instead make a folder within hydra_plugins and put the __init__.py there as well as the custom plugin. So hydra_plugins.launchers.my_submitit_launcher.MyBaseSubmititLauncher and MyBaseSubmititLauncher is located within my_submitit_launcher.py, where launchers.py got turned into a folder /launchers.
2. After submitit's pickle save-then-load-context process before starting each job (essentially sending pickled contexts to each node, which will unpickle the context and run the task), if your cwd (hydra.sweep.dir) is not in the same location as where the main app is located, then when hydra tries to load the hydra_plugins.launchers module you made in the previous step, it won't be able to find it. A workaround is:

# otherwise pickle.load from submitit launcher won't find it since cloudpickle registers by reference by default
import hydra_plugins.launchers
cloudpickle.register_pickle_by_value(hydra_plugins.launchers)

in the top your main app file. Matter of fact, you need to do this register_pickle_by_value for any module that you've created in your project folder hierarchy if you want to use a hydra working directory that's not the same as the project directory.

There are a couple of weird gotchas but I'm too tired to list them out right now. I've managed to get it working with my own custom submitit launcher with its own checkpoint function and had it send me a wandb alert during pre-emption, re-using the run = wandb.init() context that was only available within the task function. The only issue right now is that upon job restart, specifically during the pickle load process of submitit when it's reloading the associated job's pickle, it fails with a ModuleNotFoundError for a module that exists in my project hierarchy. Which is very odd considering I already used cloudpickle.register_pickle_by_value on this module.