[REVIEW] Faster Treelite serialization

[hcho3]

Summary

Speed up serialization of Treelite model objects and reduce overhead in multi-GPU RF prediction.

Features

• Flat object layout: now Tree is an array of POD objects (Node). As a consequence, all queries for a give node should now be made via methods of Tree class.
• Binary serialization without depending on Protobuf
• Eliminate the use of tempfile.

Benchmark setup

• Using this script by @Salonijain27. See the link for instructions.
• AWS EC2 instance g4dn.12xlarge, T4 GPU (16 GB GDDR6) X 4
• Benchmark setting: n_gpus=2, n_gb=2, n_features=20, depth=25, n_estimators=10. This leads to a forest consisting of 10 depth-25 trees, and we run through 25 million data rows.

Benchmark Results

Aggregate

Before	After
269.0 sec	109.4 sec (2.46x speedup)

Breakdown by components

Notice that some of the overhead is still yet to be explained. Also note that the actual prediction time is a small portion of the total run time. Due to the nature of distributed algorithm, timing measures are approximate.

Current cuML

	Master	Worker 0	Worker 1
RF->Treelite		32.6	32.8
Treelite->Protobuf		10.8	11.3
(Copy models workers -> master)
Protobuf->Treelite		26.6	23.8
Concatenate Treelite handles	23.3
(Copy model master -> workers)
Protobuf->Treelite		52.2	52.3
Treelite->FIL		8.6	8.5
FIL predict		2.9	1.1

This PR

	Master	Worker 0	Worker 1
RF->Treelite		22.1	22.5
Treelite->Bytes		< 0.001	< 0.001
(Copy models workers -> master)
Bytes->Treelite		< 0.001	< 0.001
Concatenate Treelite handles	3.1
(Copy model master -> workers)
Bytes->Treelite		< 0.001	< 0.001
Treelite->FIL		4.8	4.7
FIL predict		4.9	1.1

[JohnZed]

Very promising! I suspect that the RF -> treelite can be accelerated a lot in a future rev too... I don't know why it should have to be much longer than treelite-> fil in general if we move to convert the representation efficiently.

[JohnZed]

I did a fairly quick pass for the first rev... will review again in more detail as things get wrapped up.
Overall, I think it's great and looks pretty clear. The changes to the core cython code are smaller than I expected - it fits in well with the existing pattern, so I think it's pretty understandable.
I don't fully get the serialization to frames - looks like it simply passing through the binary format used within tl? Seems reasonable to me, but I haven't used this style of conversion before with Py_buffers etc.

[JohnZed]

In another pr, I proposed renaming this something like _alloc_and_convert_model to make it clear that the caller needs to free the result.

[JohnZed]

Maybe clarify the type and update the rest of docstring: "Returns the treelite binary format representation of this model."

[hcho3]

I don't fully get the serialization to frames - looks like it simply passing through the binary format used within tl?

Treelite objects now exposes the Python buffer protocol interface, so that we can transparently convert Treelite objects to memory views with zero overhead. In init_from_frames(), we fetch the buffer interface from the Treelite object and cast it into memory view. It is O(1) because it amounts to pointer casting.

@jakirkham gave me valuable advice for implementing the Python buffer protocol.

[Salonijain27]

Looks good, its great to see the time required for predict drop so much! I have a few suggestions and questions

[hcho3]

This is quite a strange error:

=================================== FAILURES ===================================

_________________________ test_real_algos_runner[FIL] __________________________



algo_name = 'FIL'



    @pytest.mark.parametrize('algo_name', ['UMAP-Supervised',

                                           'DBSCAN',

                                           'LogisticRegression',

                                           'ElasticNet',

                                           'FIL'])

    def test_real_algos_runner(algo_name):

        pair = algorithms.algorithm_by_name(algo_name)

    

        if (algo_name == 'UMAP' and not has_umap()) or \

           (algo_name == 'FIL' and not has_xgboost()):

            pytest.xfail()

    

        runner = AccuracyComparisonRunner(

            [20], [5], dataset_name='classification', test_fraction=0.20

        )

        results = runner.run(pair)[0]

        print(results)

>       assert results["cuml_acc"] is not None

E       KeyError: 'cuml_acc'



cuml/test/test_benchmark.py:190: KeyError

[hcho3]

I managed to fix the failing benchmark test. Marking this as ready for review.

[hcho3]

List of changes made to Treelite:

• Refactor CMake build; create treelite_runtime Python pkg (Refactor CMake build; create treelite_runtime Python pkg dmlc/treelite#167)
• Upgrade to PEP 8 style (Re-enable formatting check dmlc/treelite#175)
• Refactor Python tests with Pytest (Refactor Python tests with Pytest dmlc/treelite#176)
• Re-factor scikit-learn converter to use mixins (Refactor scikit-learn converter dmlc/treelite#179)
• Fast binary serialization of Treelite models (Fast binary serialization of Treelite models dmlc/treelite#178)

Once all tests pass, I will go ahead and release 0.92 version of Treelite.

[jakirkham]

Exciting to see all of the progress here, @hcho3! 😄

Added a couple of comments about how we might simplify things a bit here. Please let me know if you have any questions 🙂

[hcho3]

I submitted Treelite 0.92 to conda-forge: conda-forge/staged-recipes#11926. Fingers crossed.

[JohnZed]

Looks great! I have only small questions/suggestions.
Also, are there additional unit tests that would be helpful here? Serializing+deserializing model variants (e.g. classification/regression/multiclass) and ensuring we get the properties right? Not sure...

[JohnZed]

Isn't multiclass currently disabled until #2248 goes in?

[hcho3]

I think so. I included the line because it was already part of the current codebase.

[hcho3]

@JohnZed The Treelite repo contains several round-trip tests for the new serializer: https://github.com/dmlc/treelite/blob/master/tests/cpp/test_serializer.cc

[hcho3]

@jakirkham I addressed all your comments, except the one about casting buffer frames.

[jakirkham]

Grouping together suggested format_str encoding changes for clearer discussion.

[jakirkham]

Suggested change

	header = {'format_str': [x.format.encode('utf-8') for x in frames],
	header = {'format_str': [x.format for x in frames],

[hcho3]

I went ahead and applied your suggestion.

Is it preferable to pass str to pickle, rather than bytes? I'd like to understand your reasoning behind this suggestion.

[jakirkham]

Was thinking about this in the context of moving from pickle to Dask serialization down the road (assuming that is still the plan). Typically the header consists of things like int, str, dict, list, and tuple. Generally things that are MsgPack serializable. Typically bytes and memoryviews are reserved for frames instead.

Was unsure at first whether bytes would work in the header. However after playing with things a bit bytes may work. MsgPack is at least able to handle them with the flags that Dask is using.

[hcho3]

Good point. Thanks for your explanation. I agree that built-in types like str would be well supported by Dask serializer.

[jakirkham]

Thanks for all of the work here @hcho3! 😄 Looks good

Sounds like we are going to handle switching to the treelite Conda package in another PR. Is that right?

[hcho3]

Sounds like we are going to handle switching to the treelite Conda package in another PR. Is that right?

Yes, I’ll work on it after this PR is merged.

[hcho3]

The failing test in the CI should be fixed by #2432

[JohnZed]

Looks great!