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)