search

ray-project / ray

Ray is a unified framework for scaling AI and Python applications. Ray consists of a core distributed runtime and a set of AI Libraries for accelerating ML workloads.
https://ray.io
Apache License 2.0
33.05k stars 5.59k forks source link

[aDAG] More intuitive API for (NCCL) type hints #47258

Open ruisearch42 opened 3 weeks ago

ruisearch42 commented 3 weeks ago

Description

The end user may have an impression that type hint is applied to a DAG node, as opposed to the edge between DAG nodes/tasks.

This might be partially due to that the way we name the variables in DAG construction, e.g., 

     dag = sender.send.bind(shape, dtype, inp)
     dag = dag.with_type_hint(TorchTensorType(shape, dtype, transport="nccl"))
     dag = receiver.recv.bind(dag)

The semantics becomes clearer if we have:

output = sender.send.bind(shape, dtype, inp)
output_decorated = output.with_type_hint(TorchTensorType(shape, dtype, transport="nccl"))
dag = receiver.recv.bind(output_decorated)

However, the API can potentially be improved to make it more intuitive. Say having something like:

sender.send.bind(inp, channel=‘nccl’)

cc: @woshiyyya

Use case

No response

kevin85421 commented 3 weeks ago

I added the accelerated-dag tag