Support lambdas and local functions as callbacks in parallel ExternalSource #3269

stiepan · 2021-08-20T15:54:14Z

Description

Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds functionality)
Breaking change (fix or feature that would cause existing functionality to not work as expected)
Refactoring (Redesign of existing code that doesn't affect functionality)
Other (e.g. Documentation, Tests, Configuration)

What happened in this PR

It adds customized serialization of parallel externalsource callbacks so that lambdas and local functions are supported, usage of
functions defined inside Jupyter notebooks is easier and enables user to specify external package to be used to serialize callback if more specialized serialization is needed,

Additional information

Affected modules and functionalities:
pythonic parts of parallel external source implementation
Key points relevant for the review:

Checklist

Tests

Documentation

DALI team only

Requirements

Implements new requirements
Affects existing requirements
N/A

REQ IDs: N/A

JIRA TASK: DALI-2214

…kling module for ES callback Signed-off-by: Kamil Tokarski <[email protected]>

Signed-off-by: Kamil Tokarski <[email protected]>

…r callbacks instead Signed-off-by: Kamil Tokarski <[email protected]>

Signed-off-by: Kamil Tokarski <[email protected]>

stiepan · 2021-08-20T16:28:51Z

!build

dali-automaton · 2021-08-20T16:31:32Z

CI MESSAGE: [2808327]: BUILD STARTED

dali-automaton · 2021-08-20T17:52:15Z

CI MESSAGE: [2808327]: BUILD PASSED

mzient · 2021-08-22T07:13:31Z

dali/python/nvidia/dali/_multiproc/pool.py

+            if callback_pickler is None:
+                callbacks_ = callbacks
+            else:
+                callbacks_ = callback_pickler.dumps(callbacks)


So now callbacks contains either pickled or raw callbacks?

Yes, so in the first case it creates kind of additional layer of serialization where multiprocessing pickler just sees already serialized callbacks. At first I simply set Dali customized pickler to be used directly by multiprocessing but it is unfortunately a global change for the whole process and:

does not work if multiprocessing already started some actual work before DALI was able to modify the context

there's a question if that's not going to interfere with other packages using multiprocessing, like pytorch (it seemed to work okay, but I think that if we can avoid worrying about that let's avoid it.)

for user provided packages for serialization I still would go for this type of double serialization

mzient · 2021-08-22T07:15:04Z

dali/python/nvidia/dali/_multiproc/pool.py

@@ -202,10 +204,14 @@ def __init__(
            task_r, task_w = mp.Pipe(duplex=False)
            res_r, res_w = mp.Pipe(duplex=False)
            sock_reader, sock_writer = socket.socketpair()
+            if callback_pickler is None:
+                callbacks_ = callbacks


How about this?

Suggested change

callbacks_ = callbacks

callbacks_arg = callbacks

I'm not a fan of trailing underscore and leading underscore has a meaning in Python which I guess we should not abuse.

mzient · 2021-08-22T07:15:23Z

dali/python/nvidia/dali/pickling.py

+
+
+class CustomPickler:
+


Nitpick: weird empty line

mzient · 2021-08-22T07:16:28Z

dali/python/nvidia/dali/pickling.py

+
+def function_unpickle(name, qualname, code, closure):
+    code = marshal.loads(code)
+    globs = {'__builtins__': __builtins__}


Why "globs"?

Suggested change

globs = {'__builtins__': __builtins__}

builtins = {'__builtins__': __builtins__}

It is a dictionary containing "module global scope" of the function, the one available later on as fun.__globals__ immutable attribute (you cannot replace the dict but can update its contents). If one called globals() inside the instantiated function this is the object that would be returned. So it is rather a coincidence that this dictionary contains only builtins at this point - all other global references (if any) are added at the set state stage to handle possible cyclic references.
I didn't call it globals on the other hand to avoid shadowing globals() inside function_unpickle

mzient · 2021-08-22T07:19:30Z

dali/python/nvidia/dali/pickling.py

+
+
+class DaliCallbackPickler(pickle.Pickler):
+


Nitpick: weird empty line

The one line space between first method and the class ... :?
I can remove it if you think it somehow clashes with the code formatting in DALI though I tried to use it everywhere both in this and other PRs. D:
Most of the classes I can find have docstrings but those are followed with a single empty line too.

mzient · 2021-08-22T07:29:05Z

dali/python/nvidia/dali/pickling.py

+            if isinstance(obj, type(dummy_lambda)) and obj.__name__ == dummy_lambda.__name__ or \
+                    getattr(obj, '_dali_pickle_by_value', False):
+                return function_by_value_reducer(obj)
+            if '<locals>' in obj.__qualname__:


Does it work with regular local functions, such as:

def f(): def g(): return 42;

?
g.__qualname__ == 'f.<locals>.g', but it's otherwise quite trivial.
How about switching the logic to:

try: pickle.dumps(obj) except (whatever it raises): return function_by_value_reducer_(obj)

?

Yep, g will be serialized by value and plain pickling fails as pickle recursively steps down the attributes and raises AttributeError when encountering <locals>.

I'd rather limit this pickle.dumps test to the situations when I am nearly sure that it is going to fail (or maybe should get rid of the test whatsoever?) for efficiency reasons. If the test succeeded it serialized something only for us to discard that result and proceed with a single step of serialization.

It would be nice if we could catch the exception by overriding dump method of the pickler, catch AttributeError and only then fall back to obj.__dali_pickle_by_value = True. Unfortunately it doesn't seem to be (easily) doable as the pickle.Pickler class is in fact a wrapper around C implementation of pickle and overridden dump method would be called only once for a top level object; recursive steps for dependencies/contents will be serialized by C code unaware of override.

mzient · 2021-08-22T07:32:52Z

dali/test/python/test_external_source_parallel_custom_serialization.py

+def create_simple_pipeline(callback, py_callback_pickler, batch_size, parallel=True, py_num_workers=None):
+
+    extra = {}
+    if py_callback_pickler != UseDefault:


Why not:

Suggested change

if py_callback_pickler != UseDefault:

if py_callback_pickler is nt None:

like we do elsewhere?
I know that it would work with different defaults, but it seems a bit overengineered for the sake of just tests.

mzient · 2021-08-22T07:36:28Z

dali/python/nvidia/dali/pipeline.py

@@ -135,13 +136,34 @@ class Pipeline(object):
    you will need to call :meth:`start_py_workers` before calling :meth:`build` of any
    of the pipelines. You can find more details and caveats of both methods in Python's
    ``multiprocessing`` module documentation.
+`py_callback_pickler` : module or tuple, default = nvidia.dali.pickling
+    If `py_start_method` is set to *spawn* callback passed to parallel ExternalSource must be picklable.
+    If run in Python3.8 or newer, DALI uses customized pickle (`nvidia.dali.pickling`) when


Do we want to advertise this module to the users? I think it could be internal.

I decided to advertise it because apart from serving as default callback pickler it also contains @pickle_by_value decorator whose whole purpose is to be available for manual usage by users.

1.What about moving everything apart from the decorator, loads and dumps methods to some new module inside _multproc?
2. Or renaming all the things that would be moved in 1. to start with underscore?

dali/python/nvidia/dali/pipeline.py

mzient · 2021-08-22T07:39:14Z

dali/python/nvidia/dali/pipeline.py

+    external source callbacks.
+


Suggested change

external source callbacks.

external source callbacks. You can even pass a module directly as ``py_callback_pickler``::

import dill

src = fn.external_source(lambda sample_info: 42, batch=False, parallel=True, py_callback_pickler=dill)

Never mind, yo do. :D

mzient · 2021-08-22T10:49:42Z

dali/test/python/test_utils.py

+def _test_omitted():
+    print("Omitting tests suite for Python3.8+ serialization")


Suggested change

def _test_omitted():

print("Omitting tests suite for Python3.8+ serialization")

def _test_skipped(reason):

print("Test skipped." if reason is None else f"Test skipped: {reason}")

mzient · 2021-08-22T10:53:13Z

dali/test/python/test_utils.py

+        if version_info.major > major or \
+                (version_info.major == major and (minor is None or version_info.minor >= minor)):
+            return test_case
+        return _test_omitted


Suggested change

return _test_omitted

return lambda: test_skipped(f"Insufficient Python version {version_info.major}.{version_info.minor} - required {major}.{minor}")

mzient · 2021-08-22T10:56:29Z

dali/python/nvidia/dali/pickling.py

+        raise ValueError("Unsupported py_callback_pickler value provided.")
+
+    @classmethod
+    def of_reducer(cls, reducer, dumps_kwargs=None, loads_kwargs=None):


This is a weird name. Is it required by some API? What does this of stand for?

Nah, it's not required by any API, I meant create_instance_of_this_class_from_reducer and for some reason used of rather than from.

jantonguirao · 2021-08-23T09:54:15Z

dali/python/nvidia/dali/pipeline.py

+    that implements two methods: `dumps` and `loads` to make DALI use them to serialize
+    external source callbacks.
+
+    Valid value for `py_callback_pickler` is either a module/object implementing


Suggested change

Valid value for `py_callback_pickler` is either a module/object implementing

A valid value for `py_callback_pickler` is either a module/object implementing

jantonguirao · 2021-08-23T09:54:31Z

dali/python/nvidia/dali/pipeline.py

+    external source callbacks.
+
+    Valid value for `py_callback_pickler` is either a module/object implementing
+    dumps and loads methods or a tuple where first item is the module/object and the next


Suggested change

dumps and loads methods or a tuple where first item is the module/object and the next

``dumps`` and ``loads`` methods or a tuple where first item is the module/object and the next

jantonguirao · 2021-08-23T09:54:55Z

dali/python/nvidia/dali/pipeline.py

+    two optional parameters are extra kwargs to be passed when calling dumps and loads respectively.
+    Provided methods and kwargs must themselves be picklable.
+
+    If you run Python3.8 or newer and use default `nvidia.dali.pickling` you can hint DALI to serialize


Suggested change

If you run Python3.8 or newer and use default `nvidia.dali.pickling` you can hint DALI to serialize

If you run Python3.8 or newer and use the default `nvidia.dali.pickling` you can hint DALI to serialize

jantonguirao · 2021-08-23T09:55:49Z

dali/test/python/import_module_test_helper.py

@@ -0,0 +1,24 @@
+# Copyright (c) 2020-2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.


Suggested change

# Copyright (c) 2020-2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.

# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.

jantonguirao · 2021-08-23T09:56:04Z

dali/test/python/test_external_source_parallel_custom_serialization.py

@@ -0,0 +1,432 @@
+# Copyright (c) 2020-2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.


Suggested change

# Copyright (c) 2020-2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.

# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.

jantonguirao · 2021-08-23T09:57:13Z

dali/test/python/test_external_source_parallel_custom_serialization.py

+
+def dumps(obj, **kwargs):
+    if kwargs.get('special_dumps_param') != 42:
+        raise NoDumpsParam("Expected special_dumps_pram among kwargs, got {}".format(kwargs))


Suggested change

raise NoDumpsParam("Expected special_dumps_pram among kwargs, got {}".format(kwargs))

raise NoDumpsParam("Expected special_dumps_param among kwargs, got {}".format(kwargs))

jantonguirao · 2021-08-23T10:00:35Z

dali/test/python/test_external_source_parallel_custom_serialization.py

+
+    @pipeline_def(batch_size=batch_size, num_threads=2, device_id=0, py_num_workers=py_num_workers,
+        py_start_method="spawn", **extra)
+    def crate_pipline():


Suggested change

def crate_pipline():

def create_pipeline():

Signed-off-by: Kamil Tokarski <[email protected]>

stiepan · 2021-08-25T16:52:35Z

!build

dali-automaton · 2021-08-25T16:56:14Z

CI MESSAGE: [2840850]: BUILD STARTED

dali-automaton · 2021-08-25T18:14:51Z

CI MESSAGE: [2840850]: BUILD PASSED

Signed-off-by: Kamil Tokarski <[email protected]>

mzient · 2021-08-26T08:37:13Z

dali/test/python/test_external_source_parallel_custom_serialization.py

+jpeg_file = os.path.join(get_dali_extra_path(), 'db', 'single', 'jpeg', '510', 'ship-1083562_640.jpg')
+sequence_lenght = 4
+
+def create_callback_with_syntactically_nested_code_referencing_global_var():


Wow! That's a strong contender for the longest identifier in our code base!

Hoped to make the name self-explanatory

mzient · 2021-08-26T08:56:46Z

dali/test/python/test_utils.py

+    print("Test skipped." if reason is None else f"Test skipped: {reason}")
+
+
+def restrict_python_version(major, minor=None, reason=None):


Suggested change

def restrict_python_version(major, minor=None, reason=None):

def restrict_python_version(major, minor=None):

Not used?

mzient · 2021-08-26T09:01:10Z

dali/python/nvidia/dali/pipeline.py

 """
    def __init__(self, batch_size = -1, num_threads = -1, device_id = -1, seed = -1,
                 exec_pipelined=True, prefetch_queue_depth=2,
                 exec_async=True, bytes_per_sample=0,
                 set_affinity=False, max_streams=-1, default_cuda_stream_priority = 0,
                 *,
-                 enable_memory_stats=False, py_num_workers=1, py_start_method="fork"):
+                 enable_memory_stats=False, py_num_workers=1, py_start_method="fork",
+                 py_callback_pickler=dali_pickle):


Suggested change

py_callback_pickler=dali_pickle):

py_callback_pickler=None):

dali/python/nvidia/dali/pipeline.py

Signed-off-by: Kamil Tokarski <[email protected]>

stiepan · 2021-08-26T15:29:47Z

!build

dali-automaton · 2021-08-26T15:36:19Z

CI MESSAGE: [2848855]: BUILD STARTED

dali-automaton · 2021-08-26T16:47:36Z

CI MESSAGE: [2848855]: BUILD PASSED

Signed-off-by: Kamil Tokarski <[email protected]>

stiepan · 2021-08-26T17:13:23Z

!build

dali-automaton · 2021-08-26T17:16:07Z

CI MESSAGE: [2849616]: BUILD STARTED

dali-automaton · 2021-08-26T18:07:37Z

CI MESSAGE: [2849616]: BUILD FAILED

stiepan · 2021-08-26T18:10:46Z

!build

dali-automaton · 2021-08-26T18:16:08Z

CI MESSAGE: [2850059]: BUILD STARTED

dali-automaton · 2021-08-26T19:39:34Z

CI MESSAGE: [2850059]: BUILD PASSED

This comment has been minimized.

Sign in to view

stiepan added 8 commits August 20, 2021 18:04

Support pickling local functions and lambas. Support user defined pic…

2ed1dc5

…kling module for ES callback Signed-off-by: Kamil Tokarski <[email protected]>

Add tests of functions pickling in parallel external source

7dd84d4

Signed-off-by: Kamil Tokarski <[email protected]>

Include the test in QA

8b66881

Signed-off-by: Kamil Tokarski <[email protected]>

Register reducer documentation and fixes

139dc45

Signed-off-by: Kamil Tokarski <[email protected]>

Adjust reducer registration

a4980a0

Signed-off-by: Kamil Tokarski <[email protected]>

Do not override multiprocessing reducer, use custom pickler solely fo…

9d443fe

…r callbacks instead Signed-off-by: Kamil Tokarski <[email protected]>

Fix imports and typos

72101cd

Signed-off-by: Kamil Tokarski <[email protected]>

Extend notebook example, fix test dependencies

383a9e1

Signed-off-by: Kamil Tokarski <[email protected]>

stiepan force-pushed the ext_source_pickle_lambdas branch from e72016a to 383a9e1 Compare August 20, 2021 16:05

Move notebook to another PR

01f06f5

Signed-off-by: Kamil Tokarski <[email protected]>

mzient reviewed Aug 22, 2021

View reviewed changes

dali/python/nvidia/dali/pickling.py

class CustomPickler:

Copy link

Contributor

mzient Aug 22, 2021

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nitpick: weird empty line

mzient reviewed Aug 22, 2021

View reviewed changes

dali/python/nvidia/dali/pipeline.py Show resolved Hide resolved

mzient reviewed Aug 22, 2021

View reviewed changes

mzient self-assigned this Aug 22, 2021

mzient reviewed Aug 22, 2021

View reviewed changes

jantonguirao self-assigned this Aug 23, 2021

jantonguirao approved these changes Aug 23, 2021

View reviewed changes

Apply review suggestions. Handle globals from nested code.

87e70f0

Signed-off-by: Kamil Tokarski <[email protected]>

Use callback pickler only in spawn mode

f7b9fdb

Signed-off-by: Kamil Tokarski <[email protected]>

mzient reviewed Aug 26, 2021

View reviewed changes

dali/python/nvidia/dali/pipeline.py Show resolved Hide resolved

Change py_callback_pickler default, make reducers private

ec3c2eb

Signed-off-by: Kamil Tokarski <[email protected]>

mzient approved these changes Aug 26, 2021

View reviewed changes

Remove duplicated test case

2040752

Signed-off-by: Kamil Tokarski <[email protected]>

stiepan merged commit b37d522 into NVIDIA:main Aug 26, 2021

JanuszL mentioned this pull request Sep 1, 2021

DALI 2021 roadmap #2978

Closed

	globs = {'__builtins__': __builtins__}
	builtins = {'__builtins__': __builtins__}

	if py_callback_pickler != UseDefault:
	if py_callback_pickler is nt None:

-    external source callbacks.
+    external source callbacks. You can even pass a module directly as ``py_callback_pickler``::
+        import dill
+        src = fn.external_source(lambda sample_info: 42, batch=False, parallel=True, py_callback_pickler=dill)

		def _test_omitted():
		print("Omitting tests suite for Python3.8+ serialization")

	return _test_omitted
	return lambda: test_skipped(f"Insufficient Python version {version_info.major}.{version_info.minor} - required {major}.{minor}")

	Valid value for `py_callback_pickler` is either a module/object implementing
	A valid value for `py_callback_pickler` is either a module/object implementing

	dumps and loads methods or a tuple where first item is the module/object and the next
	``dumps`` and ``loads`` methods or a tuple where first item is the module/object and the next

	If you run Python3.8 or newer and use default `nvidia.dali.pickling` you can hint DALI to serialize
	If you run Python3.8 or newer and use the default `nvidia.dali.pickling` you can hint DALI to serialize

		@@ -0,0 +1,24 @@
		# Copyright (c) 2020-2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.

	# Copyright (c) 2020-2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
	# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.

		@@ -0,0 +1,432 @@
		# Copyright (c) 2020-2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.

	raise NoDumpsParam("Expected special_dumps_pram among kwargs, got {}".format(kwargs))
	raise NoDumpsParam("Expected special_dumps_param among kwargs, got {}".format(kwargs))

		print("Test skipped." if reason is None else f"Test skipped: {reason}")


		def restrict_python_version(major, minor=None, reason=None):

Support lambdas and local functions as callbacks in parallel ExternalSource #3269

Support lambdas and local functions as callbacks in parallel ExternalSource #3269

Conversation

stiepan commented Aug 20, 2021 • edited Loading

Description

What happened in this PR

Additional information

Checklist

Tests

Documentation

DALI team only

Requirements

This comment has been minimized.

stiepan commented Aug 20, 2021

dali-automaton commented Aug 20, 2021

dali-automaton commented Aug 20, 2021

Choose a reason for hiding this comment

stiepan Aug 25, 2021 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

stiepan Aug 25, 2021 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

mzient Aug 22, 2021 • edited Loading

Choose a reason for hiding this comment

stiepan Aug 25, 2021 • edited Loading

Choose a reason for hiding this comment

mzient Aug 22, 2021 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

stiepan Aug 25, 2021 • edited Loading

Choose a reason for hiding this comment

mzient Aug 22, 2021 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

mzient Aug 22, 2021 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

stiepan commented Aug 25, 2021

dali-automaton commented Aug 25, 2021

dali-automaton commented Aug 25, 2021

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

stiepan commented Aug 26, 2021

dali-automaton commented Aug 26, 2021

dali-automaton commented Aug 26, 2021

stiepan commented Aug 26, 2021

dali-automaton commented Aug 26, 2021

dali-automaton commented Aug 26, 2021

stiepan commented Aug 26, 2021

dali-automaton commented Aug 26, 2021

dali-automaton commented Aug 26, 2021

stiepan commented Aug 20, 2021 •

edited

Loading

stiepan Aug 25, 2021 •

edited

Loading

stiepan Aug 25, 2021 •

edited

Loading

mzient Aug 22, 2021 •

edited

Loading

stiepan Aug 25, 2021 •

edited

Loading

mzient Aug 22, 2021 •

edited

Loading

stiepan Aug 25, 2021 •

edited

Loading

mzient Aug 22, 2021 •

edited

Loading

mzient Aug 22, 2021 •

edited

Loading