Support inlining task arguments #72

omus · 2023-08-31T19:44:27Z

Should improve performance for Ray task calls which use small sized arguments.

codecov-commenter · 2023-08-31T19:51:11Z

Codecov Report

Merging #72 (c6b0298) into main (5b92046) will increase coverage by 2.01%.
The diff coverage is 92.30%.

@@            Coverage Diff             @@
##             main      #72      +/-   ##
==========================================
+ Coverage   69.26%   71.27%   +2.01%     
==========================================
  Files           6        6              
  Lines         257      282      +25     
==========================================
+ Hits          178      201      +23     
- Misses         79       81       +2

Flag	Coverage Δ
Ray.jl	`71.27% <92.30%> (+2.01%)`	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

Files Changed	Coverage Δ
Ray.jl/src/Ray.jl	`100.00% <ø> (ø)`
Ray.jl/src/runtime.jl	`51.53% <92.30%> (+7.33%)`	⬆️

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more

omus · 2023-08-31T19:51:55Z

I've been attempting to construct the TaskArgByValue and TaskArgByReference instances in Julia and then pass them directly into
submit_task. However, this is challenging as I need to pass in a std::vector<std::unique_ptr<TaskArg>>.

Compiling this example fails:

size_t demo(const std::vector<std::unique_ptr<std::string>> &vector) {
    return vector.size();
}

JLCXX_MODULE define_julia_module(jlcxx::Module& mod)
{
    mod.method("demo", &demo);
}

Failure:

bazel-out/darwin_arm64-opt/bin/external/libcxxwrap_julia/_virtual_includes/headers/jlcxx/stl.hpp:173:89: error: object of type 'std::__deque_base<std::unique_ptr<std::string>, std::allocator<std::unique_ptr<std::string>>>::value_type' (aka 'std::unique_ptr<std::string>') cannot be assigned because its copy assignment operator is implicitly deleted
    wrapped.method("cxxsetindex!", [](WrappedT& v, const T& val, cxxint_t i) { v[i - 1] = val; });

As reported in: JuliaInterop/CxxWrap.jl#370

I've worked around this issue for now by creating a StdVector of TaskArg pointers and then converting these pointers to std::unique_ptr inside of the C++ code.

One addition challenge was around TaskArg being an abstract C++ class. In particular attempting to use a Julia Vector{CxxRef{TaskArg}} and pushing CxxRef{TaskArgByValueAllocated} values fails as it tries to convert to an abstract type. Using a Vector{Any} does work but when converting to a StdVector we get a StdVector{Any} which is an issue for calling the C++ code. The solution here was to start with an empty StdVector allocated on the Julia side and then push elements to it. This works but required addtional work since CxxWrap was automatically dereferencing my values.

omus · 2023-09-01T17:08:14Z

I merged in #74 here as there was some overlap which makes the diff messy (TIL). Once we merge #74 I'll clean this PR up. Besides that this PR is ready for review.

Trying to get `push!` for free but converting a shared_ptr to a unique_ptr is trouble

omus · 2023-09-01T19:55:30Z

I've rebased these changes to no longer depend on #74

omus · 2023-09-01T20:01:30Z

Ray.jl/src/runtime.jl

+function transform_task_args(task_args)
+    task_arg_ptrs = StdVector{CxxPtr{ray_jll.TaskArg}}()
+    for task_arg in task_args
+        push!(task_arg_ptrs, CxxRef(task_arg))
+    end
+    return task_arg_ptrs
 end


One liner doesn't currently work:

julia> using Ray, CxxWrap julia> Ray.init() julia> task_args = Ray.serialize_args([1,2,3]) 3-element Vector{Any}: ray_core_worker_julia_jll.TaskArgByValueAllocated(Ptr{Nothing} @0x0000600000042100) ray_core_worker_julia_jll.TaskArgByValueAllocated(Ptr{Nothing} @0x0000600000043be0) ray_core_worker_julia_jll.TaskArgByValueAllocated(Ptr{Nothing} @0x0000600000042560) julia> CxxRef.(convert.(ray_jll.TaskArg, task_args)) 3-element Vector{CxxRef{ray_core_worker_julia_jll.TaskArg}}: CxxRef{ray_core_worker_julia_jll.TaskArg}(Ptr{ray_core_worker_julia_jll.TaskArg} @0x0000600000051300) CxxRef{ray_core_worker_julia_jll.TaskArg}(Ptr{ray_core_worker_julia_jll.TaskArg} @0x000060000004de60) CxxRef{ray_core_worker_julia_jll.TaskArg}(Ptr{ray_core_worker_julia_jll.TaskArg} @0x000060000004d0e0) julia> StdVector{CxxPtr{ray_jll.TaskArg}}(CxxRef.(convert.(ray_jll.TaskArg, task_args))) ERROR: MethodError: no method matching StdVector{CxxPtr{ray_core_worker_julia_jll.TaskArg}}(::Vector{CxxRef{ray_core_worker_julia_jll.TaskArg}}) Closest candidates are: StdVector{CxxPtr{ray_core_worker_julia_jll.TaskArg}}() @ ray_core_worker_julia_jll ~/.julia/packages/CxxWrap/aXNBY/src/CxxWrap.jl:624 Stacktrace: [1] top-level scope @ REPL[15]:1

kleinschmidt

I don't have any objections here; a few tiny suggestions and clarification questions but otherwise LGTM

Ray.jl/src/runtime.jl

kleinschmidt · 2023-09-05T15:35:24Z

deps/wrapper.cc

+        .method("max_direct_call_object_size", [](RayConfig &config) {
+            return config.max_direct_call_object_size();
+        })
+        .method("task_rpc_inlined_bytes_limit", [](RayConfig &config) {
+            return config.task_rpc_inlined_bytes_limit();
+        })
+        .method("record_ref_creation_sites", [](RayConfig &config) {
+            return config.record_ref_creation_sites();
+        });


why do these need to be lambdas? seem pretty straightforward...

I'm not sure now. I'll try this without the lambdas

Fails with this:

wrapper.cc:582:59: error: call to non-static member function without an object argument .method("max_direct_call_object_size", RayConfig::max_direct_call_object_size) ~~~~~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~~~

kleinschmidt · 2023-09-05T15:38:06Z

src/wrappers/any.jl

-function _submit_task(fd, oids::AbstractVector, serialized_runtime_env_info, resources)
-    # https://github.com/JuliaInterop/CxxWrap.jl/issues/367
-    args = isempty(oids) ? StdVector{ObjectID}() : StdVector(oids)
+function _submit_task(fd, args, serialized_runtime_env_info, resources::AbstractDict)


any reason to remove the AbstractVector restriction here?

The args passed in here are a StdVector so if we'd leave the type assertion in we'd have to update it. Mainly this assertion was dropped as this method should only be called if the resources argument is a AbstractDict

kleinschmidt · 2023-09-05T15:38:59Z

deps/wrapper.cc

+    for (auto &task_arg : task_args) {
+        args.emplace_back(task_arg);


since the element type is unique_ptr already I'm assuming this will auto-cast to unique_ptr?

ah wait I always forget, emplace_back constructs and pushes...

since the element type is unique_ptr already I'm assuming this will auto-cast to unique_ptr?

Just to avoid any possible confusion the task_arg element is just a raw pointer and is not already a unique_ptr

Ray.jl/test/runtime.jl

kleinschmidt · 2023-09-05T15:45:06Z

Ray.jl/src/runtime.jl

+        # Note: The Python `prepare_args_internal` function checks if the `arg` is an
+        # `ObjectRef` and in that case uses the object ID to directly make a
+        # `TaskArgByReference`. However, as the `args` here are flattened the `arg` will
+        # always be a `Pair` (or a list in Python). I suspect this Python code path just
+        # dead code so we'll exclude it from ours.


should we be doing this though? the values will be automatically dereferenced on the the task execution side so I'm not so sure; then again, if we don't do it then users may be surprised if they do have to de-reference in their work functions.

The Python code I'm referring to is attempting to avoid creating a new ObjectRef here if the user has passed in one as an argument. In all cases Python will dereference the arguments inside the called task. What I was specifically calling out here is that due to flatten_args this optimization will never occur as an arg is never a ObjectRef. At best it would be a list containing an ObjectRef as the second value.

Ray.jl/src/runtime.jl

Co-authored-by: Dave Kleinschmidt <dave.f.kleinschmidt@gmail.com>

omus marked this pull request as ready for review September 1, 2023 17:08

omus added 20 commits September 1, 2023 14:34

Support RayConfig

30d2abe

Rough first pass on prepare_task_args

1dd4e56

Support rpc::Address

57cb3fa

Support TaskArg

f36b06a

Support constructable RayObject

5185e71

Support GetCoreWorker

c8e7875

Attempt to pass back unique_ptr TaskArg

f572822

Lots of experimentation

1f40349

Functional with segfaults

feb71cf

Disable default finalizer

a0852b0

Use TaskArgByReference in prepare_task_args

a48db48

Cleanup

9e27f4c

Experiment with shared_ptr

35b7e4c

Trying to get `push!` for free but converting a shared_ptr to a unique_ptr is trouble

Extend push! with our custom _push_back

68b5372

Working submit_task

2344b1f

Refactoring

026b152

Generate call_site string

b5738cb

Add tests

2727fd3

Drop unused GetCoreWorker

48b1b65

fixup! Add tests

7623be0

omus force-pushed the cv/task-arg-by-value branch from d171b9d to 7623be0 Compare September 1, 2023 19:54

omus commented Sep 1, 2023

View reviewed changes

omus requested a review from kleinschmidt September 1, 2023 20:02

omus mentioned this pull request Sep 5, 2023

Utilize Vector{TaskArg} for serialize_args #79

Open

Add TODO about task_args being an Any vector

64c71cb

Manual cleanup of TaskArgs

327dd2b

omus mentioned this pull request Sep 5, 2023

Unable to construct StdVector{CxxPtr{TaskArg}} from Vector{CxxRef{TaskArg}} #80

Open

kleinschmidt approved these changes Sep 5, 2023

View reviewed changes

omus and others added 5 commits September 5, 2023 10:56

Merge branch 'main' into cv/task-arg-by-value

cf8da68

Add comment about using lambdas for RayConfig

8eec7a3

Comment about use of understore in serialize_args tests

c66bea3

Indentation

399c917

Co-authored-by: Dave Kleinschmidt <dave.f.kleinschmidt@gmail.com>

Push CxxPtr instead of CxxRef

987e22a

omus merged commit c630b0f into main Sep 5, 2023
4 checks passed

omus deleted the cv/task-arg-by-value branch September 5, 2023 17:13

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Support inlining task arguments #72

Support inlining task arguments #72

omus commented Aug 31, 2023

codecov-commenter commented Aug 31, 2023 •

edited

Loading

omus commented Aug 31, 2023

omus commented Sep 1, 2023

omus commented Sep 1, 2023

omus Sep 1, 2023 •

edited

Loading

omus Sep 5, 2023

kleinschmidt left a comment

kleinschmidt Sep 5, 2023

omus Sep 5, 2023

omus Sep 5, 2023

kleinschmidt Sep 5, 2023

omus Sep 5, 2023

kleinschmidt Sep 5, 2023

kleinschmidt Sep 5, 2023

omus Sep 5, 2023

kleinschmidt Sep 5, 2023

omus Sep 5, 2023

		for (auto &task_arg : task_args) {
		args.emplace_back(task_arg);

Support inlining task arguments #72

Support inlining task arguments #72

Conversation

omus commented Aug 31, 2023

codecov-commenter commented Aug 31, 2023 • edited Loading

Codecov Report

omus commented Aug 31, 2023

omus commented Sep 1, 2023

omus commented Sep 1, 2023

omus Sep 1, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

kleinschmidt left a 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

codecov-commenter commented Aug 31, 2023 •

edited

Loading

omus Sep 1, 2023 •

edited

Loading