chore: Allow kwargs only for batch-load function (#3037)

lablup · Nov 7, 2024 · 7b6072a · 7b6072a
1 parent 6e542a1
commit 7b6072a
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/src/ai/backend/manager/models/base.py b/src/ai/backend/manager/models/base.py
@@ -753,8 +753,9 @@ def get_loader_by_func(
         self,
         context: ContextT,
         batch_load_func: Callable[[ContextT, Sequence[LoaderKeyT]], Awaitable[LoaderResultT]],
-        *args,
-        **kwargs,
+        # Using kwargs-only to prevent argument position confusion
+        # when DataLoader calls `batch_load_func(keys)` which is `partial(batch_load_func, **extra_args)(keys)`.
+        **kwargs: Any,
     ) -> DataLoader:
         key = self._get_func_key(batch_load_func)
         loader = self.cache.get(key)
@@ -763,7 +764,6 @@ def get_loader_by_func(
                 functools.partial(
                     batch_load_func,
                     context,
-                    *args,
                     **kwargs,
                 ),
                 max_batch_size=128,