diff --git a/src/forge/actors/__init__.py b/src/forge/actors/__init__.py index 84937b36b..b0c7cb695 100644 --- a/src/forge/actors/__init__.py +++ b/src/forge/actors/__init__.py @@ -4,6 +4,16 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -from .collector import Collector +__all__ = ["Policy", "PolicyRouter"] -__all__ = ["Collector"] + +def __getattr__(name): + if name == "Policy": + from .policy import Policy + + return Policy + if name == "PolicyRouter": + from .policy import PolicyRouter + + return PolicyRouter + raise AttributeError("fmodule {__name__} has no attribute {name}") diff --git a/src/forge/controller/__init__.py b/src/forge/controller/__init__.py index 66906fd6a..5e3285960 100644 --- a/src/forge/controller/__init__.py +++ b/src/forge/controller/__init__.py @@ -3,11 +3,6 @@ # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -from .actor import ForgeActor -from .interface import ServiceInterface, Session, SessionContext -from .proc_mesh import get_proc_mesh, spawn_actors -from .service import Service, ServiceConfig -from .spawn import spawn_service __all__ = [ "Service", @@ -20,3 +15,44 @@ "get_proc_mesh", "ForgeActor", ] + + +def __getattr__(name): + if name == "Service": + from .service import Service + + return Service + elif name == "ServiceConfig": + from .service import ServiceConfig + + return ServiceConfig + elif name == "ServiceInterface": + from .interface import ServiceInterface + + return ServiceInterface + elif name == "Session": + from .interface import Session + + return Session + elif name == "SessionContext": + from .interface import SessionContext + + return SessionContext + elif name == "spawn_service": + from .spawn import spawn_service + + return spawn_service + elif name == "spawn_actors": + from .proc_mesh import spawn_actors + + return spawn_actors + elif name == "get_proc_mesh": + from .proc_mesh import get_proc_mesh + + return get_proc_mesh + elif name == "ForgeActor": + from .actor import ForgeActor + + return ForgeActor + else: + raise AttributeError(f"module {__name__} has no attribute {name}") diff --git a/src/forge/data/__init__.py b/src/forge/data/__init__.py index 4347199b9..71ea6b8c9 100644 --- a/src/forge/data/__init__.py +++ b/src/forge/data/__init__.py @@ -3,7 +3,18 @@ # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -from .collate import collate_packed -from .utils import CROSS_ENTROPY_IGNORE_IDX __all__ = ["collate_packed", "CROSS_ENTROPY_IGNORE_IDX"] + + +def __getattr__(name): + if name == "collate_packed": + from .collate import collate_packed + + return collate_packed + elif name == "CROSS_ENTROPY_IGNORE_IDX": + from .utils import CROSS_ENTROPY_IGNORE_IDX + + return CROSS_ENTROPY_IGNORE_IDX + else: + raise AttributeError(f"module {__name__} has no attribute {name}")