diff --git a/src/lighteval/tasks/registry.py b/src/lighteval/tasks/registry.py index e7c4e9eb6..84d2a4484 100644 --- a/src/lighteval/tasks/registry.py +++ b/src/lighteval/tasks/registry.py @@ -295,7 +295,7 @@ def create_custom_tasks_module(custom_tasks: str | Path | ModuleType) -> ModuleT """Creates a custom task module to load tasks defined by the user in their own file. Args: - custom_tasks (Optional[Union[str, ModuleType]]): Path to the custom tasks file or name of a module to import containing custom tasks or the module itself + custom_tasks: Path to the custom tasks file, name of a module containing custom tasks, or the module itself. Returns: ModuleType: The newly imported/created custom tasks modules @@ -303,6 +303,19 @@ def create_custom_tasks_module(custom_tasks: str | Path | ModuleType) -> ModuleT if isinstance(custom_tasks, ModuleType): return custom_tasks if isinstance(custom_tasks, (str, Path)) and os.path.exists(custom_tasks): + custom_tasks_path = Path(custom_tasks).resolve() + for module in tuple(sys.modules.values()): + if not isinstance(module, ModuleType): + continue + module_file = getattr(module, "__file__", None) + if module_file is None: + continue + try: + if Path(module_file).resolve() == custom_tasks_path: + return module + except (OSError, RuntimeError, TypeError): + continue + module_name = os.path.splitext(os.path.basename(custom_tasks))[0] spec = importlib.util.spec_from_file_location(module_name, custom_tasks) diff --git a/tests/unit/tasks/test_registry.py b/tests/unit/tasks/test_registry.py index cc98cb212..08c87a7e9 100644 --- a/tests/unit/tasks/test_registry.py +++ b/tests/unit/tasks/test_registry.py @@ -20,12 +20,50 @@ # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. +import importlib.util +import sys +from types import ModuleType + import pytest from lighteval.tasks.lighteval_task import LightevalTask, LightevalTaskConfig from lighteval.tasks.registry import Registry +def test_create_custom_tasks_module_reuses_previously_imported_file(tmp_path, monkeypatch): + custom_tasks_path = tmp_path / "custom_tasks.py" + custom_tasks_path.write_text( + "import custom_tasks_import_state\ncustom_tasks_import_state.import_count += 1\nTASKS_TABLE = []\n", + encoding="utf-8", + ) + + import_state = ModuleType("custom_tasks_import_state") + import_state.import_count = 0 + monkeypatch.setitem(sys.modules, import_state.__name__, import_state) + + spec = importlib.util.spec_from_file_location("preloaded_custom_tasks", custom_tasks_path) + assert spec is not None + assert spec.loader is not None + + preloaded_module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, spec.name, preloaded_module) + spec.loader.exec_module(preloaded_module) + + loaded_module = Registry.create_custom_tasks_module(custom_tasks_path) + + assert loaded_module is preloaded_module + assert import_state.import_count == 1 + + +def test_create_custom_tasks_module_loads_new_file(tmp_path): + custom_tasks_path = tmp_path / "custom_tasks.py" + custom_tasks_path.write_text("TASKS_TABLE = []\n", encoding="utf-8") + + loaded_module = Registry.create_custom_tasks_module(custom_tasks_path) + + assert loaded_module.TASKS_TABLE == [] + + def test_superset_expansion(): """ Tests that task info selector correctly handles supersets.