Add pre-call callback.

Signed-off-by: Michal Zientkiewicz <[email protected]>

Author: mzient

dali/python/nvidia/dali/experimental/dali2/_op_builder.py

@@ -200,6 +200,7 @@ def build_call_function(schema, op_class):
     header = f"__call__({', '.join(['self'] + inputs + call_args)})"
 
     def call(self, *raw_args, batch_size=None, **raw_kwargs):
+        self._pre_call(*raw_args, **raw_kwargs)
         is_batch = batch_size is not None
         if batch_size is None:
             for i, x in enumerate(list(raw_args) + list(raw_kwargs.values())):

dali/python/nvidia/dali/experimental/dali2/ops.py

@@ -105,6 +105,9 @@ def infer_output_devices(self, *inputs, **args):
         self._init_spec(inputs, args)
         return self._output_devices
 
+    def _pre_call(self, *inputs, **args):
+        pass
+
     def _is_backend_initialized(self):
         return self._op_backend is not None
 
@@ -302,16 +305,14 @@ def __init__(
             self._actual_batch_size, name, device, num_inputs, call_arg_names, **kwargs
         )
 
-    def __call__(self, ctx=None, *inputs, **args):
+    def _pre_call(self, *inputs, **args):
         if self._api_type is None:
             self._api_type = "run"
         elif self._api_type != "run":
             raise RuntimeError(
                 "Cannot mix `samples`, `batches` and `run`/`__call__` on the same reader until the end of the epoch."
             )
 
-        return super()(ctx, *inputs, **args)
-
     def run(self, ctx=None, *inputs, **args):
         if self._api_type is None:
             self._api_type = "run"
@@ -342,7 +343,7 @@ def samples(self, ctx: Optional[_eval_context.EvalContext] = None):
             meta = self._op_backend.GetReaderMeta()
             idx = 0
             while idx < meta["epoch_size_padded"]:
-                outputs = super()(ctx)
+                outputs = super().run(ctx)
                 batch_size = len(outputs[0])
                 idx += batch_size
                 for x in zip(*outputs):
@@ -378,7 +379,7 @@ def batches(self, batch_size=None, ctx: Optional[_eval_context.EvalContext] = No
             meta = self._op_backend.GetReaderMeta()
             idx = 0
             while idx < meta["epoch_size_padded"]:
-                outputs = super()(ctx)
+                outputs = super().run(ctx)
                 batch_size = len(outputs[0])
                 idx += batch_size
                 yield outputs