Pipeline decorator

dali/python/__init__.py.in

@@ -1,4 +1,4 @@
-# Copyright (c) 2017-2018, 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.
@@ -27,3 +27,4 @@ from . import tfrecord
 from . import types
 from . import plugin_manager
 from . import sysconfig
+from .pipeline import Pipeline, pipeline_def

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,101 @@ def iter_setup(self):
         For example, one can use this function to feed the input
         data from NumPy arrays."""
         pass
+
+
+def _discriminate_args(func, **func_kwargs):
+    """Split args on those applicable to Pipeline constructor and the decorated function."""
+    func_argspec = inspect.getfullargspec(func)
+    ctor_argspec = inspect.getfullargspec(Pipeline.__init__)
+
+    ctor_args = {}
+    fn_args = {}
+
+    for farg in func_kwargs.items():
+        is_ctor_arg = farg[0] in ctor_argspec.args or farg[0] in ctor_argspec.kwonlyargs
+        is_fn_arg = farg[0] in func_argspec.args or farg[0] in func_argspec.kwonlyargs
+        if is_fn_arg:
+            fn_args[farg[0]] = farg[1]
+            if is_ctor_arg:
+                print(
+                    "Warning: the argument `{}` shadows a Pipeline constructor argument of the same name.".format(
+                        farg[0]))
+        elif is_ctor_arg:
+            ctor_args[farg[0]] = farg[1]
+        else:
+            fn_args[farg[0]] = farg[1]
+
+    return ctor_args, fn_args
+
+
+def pipeline_def(fn=None, **pipeline_kwargs):
+    """
+    Decorator that converts a graph definition function into a DALI pipeline factory.
+
+    A graph definition function is a function that returns intended pipeline outputs.
+    You can decorate this function with ``@pipeline_def``::
+
+        @pipeline_def
+        def my_pipe(flip_vertical, flip_horizontal):
+            ''' Creates a DALI pipeline, which returns flipped and original images '''
+            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
+
+    The decorated function returns a DALI Pipeline object::
+
+        pipe = my_pipe(True, False)
+        # pipe.build()  # the pipeline is not configured properly yet
+
+    A pipeline requires additional parameters such as batch size, number of worker threads,
+    GPU device id and so on (see :meth:`Pipeline.__init__` for a complete list of pipeline parameters).
+    These parameters can be supplied as additional keyword arguments, passed to the decorated function::
+
+        pipe = my_pipe(True, False, batch_size=32, num_threads=1, device_id=0)
+        pipe.build()  # the pipeline is properly configured, we can build it now
+
+    The outputs from the original function became the outputs of the Pipeline::
+
+        flipped, img = pipe.run()
+
+    When some of the pipeline parameters are fixed, they can be specified by name in the decorator::
+
+        @pipeline_def(batch_size=42, num_threads=3)
+        def my_pipe(flip_vertical, flip_horizontal):
+            ...
+
+    Any Pipeline constructor parameter passed later when calling the decorated function will
+    override the decorator-defined params::
+
+        @pipeline_def(batch_size=32, num_threads=3)
+        def my_pipe():
+            data = fn.external_source(source=my_generator)
+            return data
+
+        pipe = my_pipe(batch_size=128)  # batch_size=128 overrides batch_size=32
+
+    .. warning::
+
+        The arguments of the function being decorated can shadow pipeline constructor arguments -
+        in which case there's no way to alter their values. Be especially mindful about using
+        ``**kwargs``, since code written this way may break with future versions of DALI, when
+        new parameters are added to the ``Pipeline`` constructor.
+    """
+    def actual_decorator(func):
+        @functools.wraps(func)
+        def create_pipeline(*args, **kwargs):
+            ctor_args, fn_kwargs = _discriminate_args(func, **kwargs)
+            pipe = Pipeline(**{**pipeline_kwargs, **ctor_args})  # Merge and overwrite dict
+            with pipe:
+                pipe_outputs = func(*args, **fn_kwargs)
+                if isinstance(pipe_outputs, tuple):
+                    po = pipe_outputs
+                elif pipe_outputs is None:
+                    po = ()
+                else:
+                    po = (pipe_outputs,)
+                pipe.set_outputs(*po)
+            return pipe
+        return create_pipeline
+    return actual_decorator(fn) if fn else actual_decorator