Add skeleton of HookManager (#1460)

Author: cbalioglu

src/fairseq2/recipe/base.py

@@ -37,6 +37,7 @@
 from fairseq2.recipe.internal.dataset import _DatasetHolder
 from fairseq2.recipe.internal.evaluator import _EvaluatorFactory
 from fairseq2.recipe.internal.generator import _GeneratorFactory
+from fairseq2.recipe.internal.hook import _TrainHookManager
 from fairseq2.recipe.internal.model import _ModelHolder
 from fairseq2.recipe.internal.reference_model import _ReferenceModelBootstrapper
 from fairseq2.recipe.internal.tokenizer import _TokenizerHolder
@@ -187,7 +188,13 @@ def create_trainer(
 
         trainer_factory = self._resolver.resolve(_TrainerFactory)
 
-        return trainer_factory.create(unit, data_reader, validator)
+        trainer = trainer_factory.create(unit, data_reader, validator)
+
+        hook_manager = self._resolver.resolve(_TrainHookManager)
+
+        hook_manager.maybe_register_trainer_hooks(trainer)
+
+        return trainer
 
     def create_evaluator(
         self,

src/fairseq2/recipe/composition/root.py

@@ -48,6 +48,7 @@
     _GangsFactory,
     _warmup_gangs,
 )
+from fairseq2.recipe.internal.hook import _HookManager, _TrainHookManager
 from fairseq2.recipe.internal.log import _log_ranks, _LogHelper, _StandardLogHelper
 from fairseq2.recipe.internal.logging import _DistributedLogConfigurer
 from fairseq2.recipe.internal.output_dir import _OutputDirectoryCreator
@@ -58,7 +59,11 @@
 from fairseq2.recipe.internal.task import _TaskRunner
 from fairseq2.recipe.internal.torch import _TorchConfigurer
 from fairseq2.recipe.run import _RecipeConfigDumper
-from fairseq2.runtime.dependency import DependencyContainer, DependencyResolver
+from fairseq2.runtime.dependency import (
+    DependencyContainer,
+    DependencyResolver,
+    wire_object,
+)
 from fairseq2.task import Task
 from fairseq2.utils.stopwatch import Stopwatch
 
@@ -107,6 +112,8 @@ def create_gangs(resolver: DependencyResolver) -> Gangs:
         CheckpointManager, StandardCheckpointManager, singleton=True
     )
 
+    container.register_type(_TrainHookManager, singleton=True)
+
     _register_data_parallel_wrappers(container)
     _register_lr_schedulers(container)
     _register_optim(container)
@@ -176,16 +183,27 @@ def _register_recipe_common(container: DependencyContainer) -> None:
     _register_sampling(container)
     _register_seq_generators(container)
 
+    def create_task_runner(resolver: DependencyResolver) -> _TaskRunner:
+        task_runner = wire_object(resolver, _TaskRunner)
+
+        hook_manager = resolver.resolve(_HookManager)
+
+        hook_manager.maybe_register_task_hooks(task_runner)
+
+        return task_runner
+
+    container.register(_TaskRunner, create_task_runner)
+
     container.register_type(_AssetConfigOverrider, _StandardAssetConfigOverrider)
     container.register_type(_ClusterPreparer)
     container.register_type(ComponentManager, _StandardComponentManager, singleton=True)
     container.register_type(_DistributedLogConfigurer)
     container.register_type(_GangsFactory)
+    container.register_type(_HookManager, singleton=True)
     container.register_type(_LogHelper, _StandardLogHelper)
     container.register_type(_OutputDirectoryCreator)
     container.register_type(_RecipeConfigDumper)
     container.register_type(_SweepTagGenerator, _StandardSweepTagGenerator)
-    container.register_type(_TaskRunner)
     container.register_type(_TorchConfigurer)
 
     container.collection.register_type(AssetMetadataSource, _ExtraAssetMetadataSource)

src/fairseq2/recipe/internal/hook.py

@@ -0,0 +1,23 @@
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from __future__ import annotations
+
+from typing import final
+
+from fairseq2.recipe.internal.task import _TaskRunner
+from fairseq2.trainer import Trainer
+
+
+@final
+class _TrainHookManager:
+    def maybe_register_trainer_hooks(self, trainer: Trainer) -> None:
+        pass
+
+
+@final
+class _HookManager:
+    def maybe_register_task_hooks(self, task_runner: _TaskRunner) -> None:
+        pass

src/fairseq2/recipe/internal/task.py

@@ -6,9 +6,12 @@
 
 from __future__ import annotations
 
+from collections import OrderedDict
 from signal import SIGUSR1, signal
 from types import FrameType
-from typing import final
+from typing import Protocol, final
+
+from torch.utils.hooks import RemovableHandle
 
 from fairseq2.error import raise_operational_system_error
 from fairseq2.gang import GangError, Gangs, raise_operational_gang_error
@@ -17,11 +20,16 @@
 from fairseq2.utils.stopwatch import Stopwatch
 
 
+class _TaskStartHook(Protocol):
+    def __call__(self) -> None: ...
+
+
 @final
 class _TaskRunner:
     def __init__(self, gangs: Gangs, wall_watch: Stopwatch) -> None:
         self._gangs = gangs
         self._wall_watch = wall_watch
+        self._start_hooks: dict[int, _TaskStartHook] = OrderedDict()
 
     def run(self, task: Task) -> None:
         log.info("Running on {} process(es).", self._gangs.root.size)
@@ -34,6 +42,9 @@ def request_stop(signum: int, frame: FrameType | None) -> None:
 
         original_signal_handler = signal(SIGUSR1, request_stop)
 
+        for hook in self._start_hooks.values():
+            hook()
+
         try:
             task.run()
         except OSError as ex:
@@ -69,3 +80,10 @@ def request_stop(signum: int, frame: FrameType | None) -> None:
             task.close()
 
             signal(SIGUSR1, original_signal_handler)
+
+    def register_start_hook(self, hook: _TaskStartHook) -> RemovableHandle:
+        handle = RemovableHandle(self._start_hooks)
+
+        self._start_hooks[handle.id] = hook
+
+        return handle