def spread(f: RemoteFunction, *args, **kwargs):
with spread_to_all_nodes(f) as (nodes, pg):
return [
f.options(placement_group=pg,
placement_group_bundle_index=index).remote(
*args, **kwargs) for index in range(nodes)
]
def update_workflow_options(f: RemoteFunction,
**workflow_options) -> RemoteFunction:
new_metadata = f._default_options.get("_metadata", {}).copy()
# copy again because the origina copy is shallow copy
new_workflow_options = new_metadata.get(WORKFLOW_OPTIONS, {}).copy()
new_workflow_options.update(**workflow_options)
new_metadata[WORKFLOW_OPTIONS] = new_workflow_options
return f.options(_metadata=new_metadata)