Clean up dependency registrations (#1444)

Author: cbalioglu

src/fairseq2/composition/assets.py

@@ -47,46 +47,46 @@
 def register_file_assets(
     container: DependencyContainer, path: Path, *, not_exist_ok: bool = False
 ) -> None:
-    def get_source(resolver: DependencyResolver) -> AssetMetadataSource:
+    def create_source(resolver: DependencyResolver) -> AssetMetadataSource:
         return wire_object(
             resolver, FileAssetMetadataSource, path=path, not_exist_ok=not_exist_ok
         )
 
-    container.collection.register(AssetMetadataSource, get_source)
+    container.collection.register(AssetMetadataSource, create_source)
 
 
 def register_package_assets(container: DependencyContainer, package: str) -> None:
-    def get_source(resolver: DependencyResolver) -> AssetMetadataSource:
+    def create_source(resolver: DependencyResolver) -> AssetMetadataSource:
         return wire_object(resolver, PackageAssetMetadataSource, package=package)
 
-    container.collection.register(AssetMetadataSource, get_source)
+    container.collection.register(AssetMetadataSource, create_source)
 
 
 def register_in_memory_assets(
     container: DependencyContainer, source: str, entries: Sequence[dict[str, object]]
 ) -> None:
-    def get_source(resolver: DependencyResolver) -> AssetMetadataSource:
+    def create_source(resolver: DependencyResolver) -> AssetMetadataSource:
         return wire_object(
             resolver, InMemoryAssetMetadataSource, name=source, entries=entries
         )
 
-    container.collection.register(AssetMetadataSource, get_source)
+    container.collection.register(AssetMetadataSource, create_source)
 
 
 def register_checkpoint_models(
     container: DependencyContainer, checkpoint_dir: Path
 ) -> None:
-    def get_source(resolver: DependencyResolver) -> AssetMetadataSource:
+    def create_source(resolver: DependencyResolver) -> AssetMetadataSource:
         return wire_object(resolver, ModelMetadataSource, checkpoint_dir=checkpoint_dir)
 
-    container.collection.register(AssetMetadataSource, get_source)
+    container.collection.register(AssetMetadataSource, create_source)
 
 
 def _register_asset(container: DependencyContainer) -> None:
     container.register_type(AssetDirectoryAccessor, StandardAssetDirectoryAccessor)
 
     # Store
-    def get_asset_store(resolver: DependencyResolver) -> AssetStore:
+    def load_asset_store(resolver: DependencyResolver) -> AssetStore:
         sources = resolver.collection.resolve(AssetMetadataSource)
 
         def load_providers() -> Iterator[AssetMetadataProvider]:
@@ -97,14 +97,16 @@ def load_providers() -> Iterator[AssetMetadataProvider]:
 
         env = env_detector.detect()
 
+        metadata_providers = load_providers()
+
         return wire_object(
             resolver,
             StandardAssetStore,
-            metadata_providers=load_providers(),
+            metadata_providers=metadata_providers,
             default_env=env,
         )
 
-    container.register(AssetStore, get_asset_store, singleton=True)
+    container.register(AssetStore, load_asset_store, singleton=True)
 
     container.register_type(AssetEnvironmentDetector)
 
@@ -128,7 +130,7 @@ def load_providers() -> Iterator[AssetMetadataProvider]:
     container.collection.register_type(AssetDownloadManager, LocalAssetDownloadManager)
     container.collection.register_type(AssetDownloadManager, HuggingFaceHub)
 
-    def get_standard_asset_download_manager(
+    def create_standard_asset_download_manager(
         resolver: DependencyResolver,
     ) -> AssetDownloadManager:
         dirs = resolver.resolve(AssetDirectoryAccessor)
@@ -150,5 +152,5 @@ def get_standard_asset_download_manager(
         )
 
     container.collection.register(
-        AssetDownloadManager, get_standard_asset_download_manager
+        AssetDownloadManager, create_standard_asset_download_manager
     )

src/fairseq2/composition/datasets.py

@@ -49,7 +49,7 @@ def register_dataset_family(
     elif opener is None:
         raise ValueError("`opener` or `advanced_opener` must be specified.")
 
-    def get_family(resolver: DependencyResolver) -> DatasetFamily:
+    def create_family(resolver: DependencyResolver) -> DatasetFamily:
         nonlocal opener
 
         if advanced_opener is not None:
@@ -70,7 +70,7 @@ def open_dataset(config: DatasetConfigT) -> DatasetT:
             opener=opener,
         )
 
-    container.register(DatasetFamily, get_family, key=name)
+    container.register(DatasetFamily, create_family, key=name)
 
 
 def _register_dataset_families(container: DependencyContainer) -> None:

src/fairseq2/composition/lib.py

@@ -128,7 +128,7 @@ def _register_library(
     else:
         container.register_type(ProgressReporter, RichProgressReporter, singleton=True)
 
-        def get_download_progress_reporter(
+        def create_download_progress_reporter(
             resolver: DependencyResolver,
         ) -> ProgressReporter:
             columns = _create_rich_download_progress_columns()
@@ -137,7 +137,7 @@ def get_download_progress_reporter(
 
         container.register(
             ProgressReporter,
-            get_download_progress_reporter,
+            create_download_progress_reporter,
             key="download_reporter",
             singleton=True,
         )
@@ -151,30 +151,29 @@ def get_world_info(resolver: DependencyResolver) -> WorldInfo:
     container.register(WorldInfo, get_world_info, singleton=True)
 
     # Device
-    def get_default_device(resolver: DependencyResolver) -> Device:
+    def detect_default_device(resolver: DependencyResolver) -> Device:
         device_detector = resolver.resolve(DefaultDeviceDetector)
 
         return device_detector.detect()
 
-    container.register(Device, get_default_device, singleton=True)
+    container.register(Device, detect_default_device, singleton=True)
 
     # ThreadPool
-    def get_default_thread_pool(resolver: DependencyResolver) -> ThreadPool:
+    def create_thread_pool(resolver: DependencyResolver) -> ThreadPool:
         world_info = resolver.resolve(WorldInfo)
 
         return StandardThreadPool.create_default(world_info.local_size)
 
-    container.register(ThreadPool, get_default_thread_pool, singleton=True)
+    container.register(ThreadPool, create_thread_pool, singleton=True)
 
     # RngBag
-    def get_default_rng_bag(resolver: DependencyResolver) -> RngBag:
+    def create_rng_bag(resolver: DependencyResolver) -> RngBag:
         device = resolver.resolve(Device)
 
         return RngBag.from_device_defaults(CPU, device)
 
-    container.register(RngBag, get_default_rng_bag, singleton=True)
+    container.register(RngBag, create_rng_bag, singleton=True)
 
-    # fmt: off
     container.register_type(AssetConfigLoader, StandardAssetConfigLoader)
     container.register_type(ClusterResolver, StandardClusterResolver)
     container.register_type(ConfigMerger, StandardConfigMerger)
@@ -184,13 +183,17 @@ def get_default_rng_bag(resolver: DependencyResolver) -> RngBag:
     container.register_type(FileSystem, LocalFileSystem, singleton=True)
     container.register_type(GlobalModelLoader, singleton=True)
     container.register_type(GlobalTokenizerLoader, singleton=True)
-    container.register_type(ModelCheckpointLoader, DelegatingModelCheckpointLoader, singleton=True)
+    container.register_type(
+        ModelCheckpointLoader, DelegatingModelCheckpointLoader, singleton=True
+    )
     container.register_type(ModelMetadataDumper, StandardModelMetadataDumper)
     container.register_type(ModelMetadataLoader, StandardModelMetadataLoader)
     container.register_type(ModelSharder, StandardModelSharder, singleton=True)
     container.register_type(ObjectValidator, StandardObjectValidator, singleton=True)
     container.register_type(SafetensorsLoader, HuggingFaceSafetensorsLoader)
-    container.register_type(SentencePieceModelLoader, StandardSentencePieceModelLoader, singleton=True)
+    container.register_type(
+        SentencePieceModelLoader, StandardSentencePieceModelLoader, singleton=True
+    )
     container.register_type(TensorDumper, TorchTensorDumper, singleton=True)
     container.register_type(TensorLoader, TorchTensorLoader, singleton=True)
     container.register_type(ValueConverter, StandardValueConverter, singleton=True)
@@ -203,13 +206,18 @@ def get_default_rng_bag(resolver: DependencyResolver) -> RngBag:
     container.collection.register_type(ModuleSharder, LinearSharder)
     container.collection.register_type(ModuleSharder, MoESharder)
 
-    container.collection.register_type(ModelCheckpointLoader, BasicModelCheckpointLoader)
-    container.collection.register_type(ModelCheckpointLoader, NativeModelCheckpointLoader)
-    container.collection.register_type(ModelCheckpointLoader, SafetensorsCheckpointLoader)
+    container.collection.register_type(
+        ModelCheckpointLoader, BasicModelCheckpointLoader
+    )
+    container.collection.register_type(
+        ModelCheckpointLoader, NativeModelCheckpointLoader
+    )
+    container.collection.register_type(
+        ModelCheckpointLoader, SafetensorsCheckpointLoader
+    )
     container.collection.register_type(ModelCheckpointLoader, LLaMACheckpointLoader)
 
     container.collection.register_type(ConfigDirective, ReplaceEnvDirective)
-    # fmt: on
 
     _register_asset(container)
 

src/fairseq2/composition/models.py

@@ -176,7 +176,7 @@ def register_model_family(
     elif factory is None:
         raise ValueError("`factory` or `advanced_factory` must be specified.")
 
-    def get_family(resolver: DependencyResolver) -> ModelFamily:
+    def create_family(resolver: DependencyResolver) -> ModelFamily:
         nonlocal factory
 
         if advanced_factory is not None:
@@ -207,7 +207,7 @@ def create_model(config: ModelConfigT) -> ModelT:
             hg_exporter=hg_exporter,
         )
 
-    container.register(ModelFamily, get_family, key=name)
+    container.register(ModelFamily, create_family, key=name)
 
 
 def _register_model_families(container: DependencyContainer) -> None:

src/fairseq2/composition/tokenizers.py

@@ -84,7 +84,7 @@ def register_tokenizer_family(
     elif loader is None:
         raise ValueError("`loader` or `advanced_loader` must be specified.")
 
-    def get_family(resolver: DependencyResolver) -> TokenizerFamily:
+    def create_family(resolver: DependencyResolver) -> TokenizerFamily:
         nonlocal loader
 
         if advanced_loader is not None:
@@ -105,7 +105,7 @@ def load_tokenizer(path: Path, config: TokenizerConfigT) -> Tokenizer:
             loader=loader,
         )
 
-    container.register(TokenizerFamily, get_family, key=name)
+    container.register(TokenizerFamily, create_family, key=name)
 
 
 def _register_tokenizer_families(container: DependencyContainer) -> None:

src/fairseq2/recipe/cli.py

@@ -202,7 +202,6 @@ def get_config(resolver: DependencyResolver) -> _RecipeConfigHolder:
     container.register(_RecipeConfigHolder, get_config)
 
     container.register_type(_RecipeConfigLoader)
-    container.register_type(_RecipeConfigPrinter)
     container.register_type(_RecipeConfigStructurer, _StandardRecipeConfigStructurer)
 
     # Recipe Output Directory
@@ -213,6 +212,8 @@ def get_output_dir(resolver: DependencyResolver) -> Path:
 
     container.register(Path, get_output_dir)
 
+    container.register_type(_RecipeConfigPrinter)
+
     # CLI Errors
     _register_cli_errors(container)
 

src/fairseq2/recipe/composition/data_parallel.py

@@ -11,10 +11,8 @@
 from fairseq2.recipe.base import Recipe, RecipeContext
 from fairseq2.recipe.internal.data_parallel import (
     _DataParallelModelWrapper,
-    _DDPFactory,
     _DDPModelWrapper,
     _DelegatingDPModelWrapper,
-    _FSDPFactory,
     _FSDPModelWrapper,
 )
 from fairseq2.runtime.dependency import (
@@ -29,20 +27,21 @@ def _register_data_parallel_wrappers(container: DependencyContainer) -> None:
     container.register_type(_DataParallelModelWrapper, _DelegatingDPModelWrapper)
 
     # DDP
-    def get_ddp_wrapper(resolver: DependencyResolver) -> _DataParallelModelWrapper:
+    def create_ddp_wrapper(resolver: DependencyResolver) -> _DataParallelModelWrapper:
         train_recipe = resolver.resolve(Recipe)
 
         context = RecipeContext(resolver)
 
         static_graph = train_recipe.has_static_autograd_graph(context)
 
-        return wire_object(resolver, _DDPModelWrapper, static_graph=static_graph)
+        return wire_object(
+            resolver, _DDPModelWrapper, ddp_factory=to_ddp, static_graph=static_graph
+        )
 
-    container.register(_DataParallelModelWrapper, get_ddp_wrapper, key="ddp")
-
-    container.register_instance(_DDPFactory, to_ddp)  # type: ignore[arg-type]
+    container.register(_DataParallelModelWrapper, create_ddp_wrapper, key="ddp")
 
     # FSDP
-    container.register_type(_DataParallelModelWrapper, _FSDPModelWrapper, key="fsdp")
+    def create_fsdp_wrapper(resolver: DependencyResolver) -> _FSDPModelWrapper:
+        return wire_object(resolver, _FSDPModelWrapper, fsdp_factory=to_fsdp)
 
-    container.register_instance(_FSDPFactory, to_fsdp)  # type: ignore[arg-type]
+    container.register(_DataParallelModelWrapper, create_fsdp_wrapper, key="fsdp")

src/fairseq2/recipe/composition/dataset.py

@@ -8,23 +8,21 @@
 
 from fairseq2.recipe.composition.config import register_config_section
 from fairseq2.recipe.config import DatasetSection
-from fairseq2.recipe.internal.dataset import _DatasetHolder, _RecipeDatasetOpener
+from fairseq2.recipe.internal.dataset import _DatasetHolder, _DatasetOpener
 from fairseq2.runtime.dependency import DependencyContainer, DependencyResolver
 
 
 def register_dataset(container: DependencyContainer, section_name: str) -> None:
     register_config_section(container, section_name, DatasetSection, keyed=True)
 
-    def get_dataset_holder(resolver: DependencyResolver) -> _DatasetHolder:
+    def open_dataset(resolver: DependencyResolver) -> _DatasetHolder:
         section = resolver.resolve(DatasetSection, key=section_name)
 
-        dataset_opener = resolver.resolve(_RecipeDatasetOpener)
+        dataset_opener = resolver.resolve(_DatasetOpener)
 
         return dataset_opener.open(section_name, section)
 
-    container.register(
-        _DatasetHolder, get_dataset_holder, key=section_name, singleton=True
-    )
+    container.register(_DatasetHolder, open_dataset, key=section_name, singleton=True)
 
     def get_dataset(resolver: DependencyResolver) -> object:
         dataset_holder = resolver.resolve(_DatasetHolder, key=section_name)
@@ -34,11 +32,11 @@ def get_dataset(resolver: DependencyResolver) -> object:
     container.register(object, get_dataset, key=section_name, singleton=True)
 
 
-def _register_datasets(container: DependencyContainer) -> None:
-    container.register_type(_RecipeDatasetOpener)
-
+def _register_default_dataset(container: DependencyContainer) -> None:
     register_dataset(container, section_name="dataset")
 
+    container.register_type(_DatasetOpener)
+
     # Default Dataset
     def get_dataset_holder(resolver: DependencyResolver) -> _DatasetHolder:
         return resolver.resolve(_DatasetHolder, key="dataset")

src/fairseq2/recipe/composition/evaluator.py

@@ -18,10 +18,10 @@
 
 
 def _register_evaluator_factory(container: DependencyContainer) -> None:
-    def get_evaluator_factory(resolver: DependencyResolver) -> _EvaluatorFactory:
+    def create_evaluator_factory(resolver: DependencyResolver) -> _EvaluatorFactory:
         def create_evaluator(**kwargs: Any) -> Evaluator:
             return wire_object(resolver, Evaluator, **kwargs)
 
-        return wire_object(resolver, _EvaluatorFactory, activator=create_evaluator)
+        return wire_object(resolver, _EvaluatorFactory, base_factory=create_evaluator)
 
-    container.register(_EvaluatorFactory, get_evaluator_factory)
+    container.register(_EvaluatorFactory, create_evaluator_factory)

src/fairseq2/recipe/composition/generator.py

@@ -18,10 +18,10 @@
 
 
 def _register_generator_factory(container: DependencyContainer) -> None:
-    def get_generator_factory(resolver: DependencyResolver) -> _GeneratorFactory:
+    def create_generator_factory(resolver: DependencyResolver) -> _GeneratorFactory:
         def create_generator(**kwargs: Any) -> Generator:
             return wire_object(resolver, Generator, **kwargs)
 
-        return wire_object(resolver, _GeneratorFactory, activator=create_generator)
+        return wire_object(resolver, _GeneratorFactory, base_factory=create_generator)
 
-    container.register(_GeneratorFactory, get_generator_factory)
+    container.register(_GeneratorFactory, create_generator_factory)