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,113 @@ 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]
+
+    for farg in fn_args.items():
+        if farg[0] not in func_argspec.args and farg[0] not in func_argspec.kwonlyargs:
+            raise TypeError(
+                "Using non-explicitly declared arguments in graph-defining function is not allowed. "
+                "Please remove `{}` argument or declare it explicitly in the function signature.".format(
+                    farg[0]))
+
+    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.
+
+    .. note::
+
+        Using non-explicitly declared arguments in graph-defining function is not allowed.
+        They may result in unwanted, silent hijacking of some arguments of the same name by
+        Pipeline constructor. Code written this way may cease to work 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

dali/test/python/test_pipeline_decorator.py

@@ -0,0 +1,146 @@
+# Copyright (c) 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.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from nvidia.dali import Pipeline, pipeline_def
+from nose.tools import nottest, raises
+import nvidia.dali.fn as fn
+from test_utils import get_dali_extra_path, compare_pipelines
+import os
+
+data_root = get_dali_extra_path()
+images_dir = os.path.join(data_root, 'db', 'single', 'jpeg')
+
+N_ITER = 7
+
+max_batch_size = 16
+num_threads = 4
+device_id = 0
+
+
+def reference_pipeline(flip_vertical, flip_horizontal, ref_batch_size=max_batch_size):
+    pipeline = Pipeline(ref_batch_size, num_threads, device_id)
+    with pipeline:
+        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)
+        pipeline.set_outputs(flipped, img)
+    return pipeline
+
+
+@nottest
+@pipeline_def(batch_size=max_batch_size, num_threads=num_threads, device_id=device_id)
+def pipeline_static(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
+
+
+@nottest
+@pipeline_def
+def pipeline_runtime(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
+
+
+@nottest
+def test_pipeline_static(flip_vertical, flip_horizontal):
+    put_args = pipeline_static(flip_vertical, flip_horizontal)
+    ref = reference_pipeline(flip_vertical, flip_horizontal)
+    compare_pipelines(put_args, ref, batch_size=max_batch_size, N_iterations=N_ITER)
+
+
+@nottest
+def test_pipeline_runtime(flip_vertical, flip_horizontal):
+    put_combined = pipeline_runtime(flip_vertical, flip_horizontal, batch_size=max_batch_size,
+                                    num_threads=num_threads, device_id=device_id)
+    ref = reference_pipeline(flip_vertical, flip_horizontal)
+    compare_pipelines(put_combined, ref, batch_size=max_batch_size, N_iterations=N_ITER)
+
+
+@nottest
+def test_pipeline_override(flip_vertical, flip_horizontal, batch_size):
+    put_combined = pipeline_static(flip_vertical, flip_horizontal, batch_size=batch_size,
+                                   num_threads=num_threads, device_id=device_id)
+    ref = reference_pipeline(flip_vertical, flip_horizontal, ref_batch_size=batch_size)
+    compare_pipelines(put_combined, ref, batch_size=batch_size, N_iterations=N_ITER)
+
+
+def test_pipeline_decorator():
+    for vert in [0, 1]:
+        for hori in [0, 1]:
+            yield test_pipeline_static, vert, hori
+            yield test_pipeline_runtime, vert, hori
+            yield test_pipeline_override, vert, hori, 16
+    yield test_pipeline_runtime, fn.random.coin_flip(seed=123), fn.random.coin_flip(seed=234)
+    yield test_pipeline_static, fn.random.coin_flip(seed=123), fn.random.coin_flip(seed=234)
+
+
+def test_duplicated_argument():
+    @pipeline_def(batch_size=max_batch_size, num_threads=num_threads, device_id=device_id)
+    def ref_pipeline(val):
+        data, _ = fn.file_reader(file_root=images_dir)
+        return data + val
+
+    @pipeline_def(batch_size=max_batch_size, num_threads=num_threads, device_id=device_id)
+    def pipeline_duplicated_arg(max_streams):
+        data, _ = fn.file_reader(file_root=images_dir)
+        return data + max_streams
+
+    pipe = pipeline_duplicated_arg(max_streams=42)
+    assert pipe._max_streams == -1
+    ref = ref_pipeline(42)
+    compare_pipelines(pipe, ref, batch_size=max_batch_size, N_iterations=N_ITER)
+
+
+# test_kwargs_exception tests against user introducing arguments,
+# that are not explicitly declared in function signature
+
+@pipeline_def
+def pipeline_kwargs(arg1, arg2, *args, **kwargs):
+    pass
+
+
+@pipeline_def
+def pipeline_kwonlyargs(arg1, *, arg2, **kwargs):
+    pass
+
+
+def test_kwargs_exception_1():
+    pipeline_kwargs(1, arg2=2)
+
+
+def test_kwargs_exception_2():
+    pipeline_kwonlyargs(1, arg2=2)
+
+
+@raises(TypeError)
+def test_kwargs_exception_3():
+    pipeline_kwargs(arg1=1, arg2=2, arg3=3)
+
+
+def test_kwargs_exception_4():
+    pipeline_kwargs(1, 2, 3, 4, 5)
+
+
+@raises(TypeError)
+def test_kwargs_exception_5():
+    pipeline_kwonlyargs(1, arg2=2, arg3=3)
+
+
+@raises(TypeError)
+def test_kwargs_exception_6():
+    pipeline_kwargs(1, arg2=2, arg3=3)