Pipeline decorator

dali/python/nvidia/dali/pipeline.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2017-2020, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2017-2021, NVIDIA CORPORATION. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -12,20 +12,25 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-#pylint: disable=no-member
+# pylint: disable=no-member
 from collections import deque
 from nvidia.dali import backend as b
 from nvidia.dali import tensors as Tensors
 from nvidia.dali import types
 from nvidia.dali.backend import CheckDLPackCapsule
 from threading import local as tls
 from . import data_node as _data_node
+import functools
+import inspect
 import warnings
 import ctypes
+
 pipeline_tls = tls()
 
 from .data_node import DataNode
-DataNode.__module__ = __name__      # move to pipeline
+
+DataNode.__module__ = __name__  # move to pipeline
+
 
 def _show_deprecation_warning(deprecated, in_favor_of):
     # show only this warning
@@ -34,6 +39,7 @@ def _show_deprecation_warning(deprecated, in_favor_of):
         warnings.warn("{} is deprecated, please use {} instead".format(deprecated, in_favor_of),
                       Warning, stacklevel=2)
 
+
 def _get_default_stream_for_array(array):
     if types._is_torch_tensor(array):
         import torch
@@ -990,3 +996,63 @@ def iter_setup(self):
         For example, one can use this function to feed the input
         data from NumPy arrays."""
         pass
+
+
+def _discriminate_args(**kwargs):
+    """Split args on those applicable to Pipeline constructor and the rest."""
+    fca = inspect.getfullargspec(Pipeline.__init__)  # Full Ctor Args
+    ctor_args = dict(filter(lambda x: x[0] in fca.args or x[0] in fca.kwonlyargs, kwargs.items()))
+    fn_args = dict(filter(lambda x: x[0] not in ctor_args.keys(), kwargs.items()))
+    return ctor_args, fn_args
+
+
+def pipeline(fn=None, **pipeline_kwargs):
+    """
+    Decorator for wrapping functions that define DALI pipelines.
+
+    **Usage**
+
+    First, implement a function, that defines a DALI pipeline.
+    Such function shall return DALI's DataNodes. These return
+    values will denote output from the pipeline. You can
+    decorate this function with ``@pipeline``::
+
+        @nvidia.dali.pipeline.pipeline
+        def pipe(flip_vertical, flip_horizontal):
+            data, _ = fn.file_reader(file_root=images_dir)
+            img = fn.image_decoder(data, device="mixed")
+            flipped = fn.flip(img, horizontal=flip_horizontal, vertical=flip_vertical)
+            return flipped, img
+
+    When creating a Pipeline object using the decorated function,
+    you can pass any number of ``Pipeline.__init__`` arguments as
+    keyword-args to the call::
+
+        my_pipe = pipe(0, batch_size=32, num_threads=1, device_id=0, flip_horizontal=1)
+        my_pipe.build()
+
+    Additionally, decorator can accept Pipeline constructor
+    parameters. Please note, that in this case you must use
+    keyword args only. Any Pipeline constructor parameter
+    passed later at pipeline instantiation will overwrite
+    the decorator-defined params::
+
+        @nvidia.dali.pipeline.pipeline(batch_size=32, num_threads=3)
+        def pipe():
+            data = fn.external_source(source=my_generator)
+            return data
+
+        my_pipe = pipe(batch_size=128)  # batch_size=128 overwrites batch_size=32
+    """
+    def actual_decorator(func):
+        @functools.wraps(func)
+        def create_pipeline(*args, **kwargs):
+            ctor_args, fn_kwargs = _discriminate_args(**kwargs)
+            pipe = Pipeline(**{**pipeline_kwargs, **ctor_args})  # Merge and overwrite dict
+            with pipe:
+                pipe_outputs = func(*args, **fn_kwargs)
+                _ = pipe_outputs if isinstance(pipe_outputs, tuple) else pipe_outputs,
+                pipe.set_outputs(*_)
+            return pipe
+        return create_pipeline
+    return actual_decorator(fn) if fn else actual_decorator