wehub-resource-sync
66 lines
2.5 KiB
Python
66 lines
2.5 KiB
Python
from typing import Dict, Optional, Union
|
|
|
|
import ray
|
|
|
|
|
|
def _get_node_id_from_node_ip(node_ip: str) -> Optional[str]:
|
|
"""Returns the node ID for the first alive node with the input IP."""
|
|
for node in ray.nodes():
|
|
if node["Alive"] and node["NodeManagerAddress"] == node_ip:
|
|
return node["NodeID"]
|
|
|
|
return None
|
|
|
|
|
|
def _force_on_node(
|
|
node_id: str,
|
|
remote_func_or_actor_class: Optional[
|
|
Union[ray.remote_function.RemoteFunction, ray.actor.ActorClass]
|
|
] = None,
|
|
) -> Union[Union[ray.remote_function.RemoteFunction, ray.actor.ActorClass], Dict]:
|
|
"""Schedule a remote function or actor class on a given node.
|
|
|
|
Args:
|
|
node_id: The node to schedule on.
|
|
remote_func_or_actor_class: A Ray remote function or actor class
|
|
to schedule on the input node. If None, this function will directly
|
|
return the options dict to pass to another remote function or actor class
|
|
as remote options.
|
|
Returns:
|
|
The provided remote function or actor class, but with options modified to force
|
|
placement on the input node. If remote_func_or_actor_class is None,
|
|
the options dict to pass to another remote function or
|
|
actor class as remote options kwargs.
|
|
"""
|
|
|
|
options = {"label_selector": {ray._raylet.RAY_NODE_ID_KEY: node_id}}
|
|
|
|
if remote_func_or_actor_class is None:
|
|
return options
|
|
|
|
return remote_func_or_actor_class.options(**options)
|
|
|
|
|
|
def _force_on_current_node(
|
|
remote_func_or_actor_class: Optional[
|
|
Union[ray.remote_function.RemoteFunction, ray.actor.ActorClass]
|
|
] = None
|
|
) -> Union[Union[ray.remote_function.RemoteFunction, ray.actor.ActorClass], Dict]:
|
|
"""Schedule a remote function or actor class on the current node.
|
|
|
|
If using Ray Client, the current node is the client server node.
|
|
|
|
Args:
|
|
remote_func_or_actor_class: A Ray remote function or actor class
|
|
to schedule on the current node. If None, this function will directly
|
|
return the options dict to pass to another remote function or actor class
|
|
as remote options.
|
|
Returns:
|
|
The provided remote function or actor class, but with options modified to force
|
|
placement on the current node. If remote_func_or_actor_class is None,
|
|
the options dict to pass to another remote function or
|
|
actor class as remote options kwargs.
|
|
"""
|
|
current_node_id = ray.get_runtime_context().get_node_id()
|
|
return _force_on_node(current_node_id, remote_func_or_actor_class)
|