def f(block, accepted_resources):
true_resources = {
resource: value[0][1]
for resource, value in ray.get_resource_ids().items()
}
if block:
ray.get(g.remote())
return dicts_equal(true_resources, accepted_resources)
def cpu_affinity(name=None):
resources = ray.get_resource_ids()
if 'CPU' in resources:
cpus = [v[0] for v in resources['CPU']]
else:
raise ValueError(f'No cpu is available')
psutil.Process().cpu_affinity(cpus)
if name:
print(f'CPUs corresponding to {name}: {cpus}')
def pin_to_core(worker):
"""
Worker is a handle to a remote Ray actor with a `run_fn` method.
"""
resource_ids = ray.get(
worker.run_fn.remote((lambda: ray.get_resource_ids())))
cpu_idx = resource_ids["CPU"][0][0]
ray.get(
worker.run_fn.remote(
(lambda: psutil.Process().cpu_affinity([cpu_idx]))))
def _set_device_from_fluid_res():
for node in ray.nodes():
num_gpus = node["Resources"].get("GPU", 0)
if num_gpus == 0:
return
for k, _ in ray.get_resource_ids().items():
if not k.startswith("fluid_GPU_"):
continue
gpu_id = int(k.split("_")[2])
torch.cuda.set_device(gpu_id)
break
def get_pin_list():
"""Get list of CPU IDs to pin to, using Ray's CPU allocation."""
resources = ray.get_resource_ids()
cpu_ids = []
for cpu_id, cpu_frac in resources['CPU']:
# sanity check: we should have 100% of each CPU
assert abs(cpu_frac - 1.0) < 1e-5, \
"for some reason I have fraction %f of CPU %d (??)" \
% (cpu_id, cpu_frac)
cpu_ids.append(cpu_id)
assert len(cpu_ids) > 0, \
"Ray returned no CPU IDs (was num_cpus=0 accidentally specified " \
"for this task?)"
return cpu_ids
def ray_pin_to_core():
resources = ray.get_resource_ids()
core_ids = [core for core, _frac in resources.get("CPU", [])]
curr_process = psutil.Process()
if len(core_ids) == 0:
logger.error(
"Failed to pin current process to cores because no CPU assigned")
return
if sys.platform != "linux":
logger.error(
"Current platform doesn't support pinnning to cores, only Linux "
"is suppported")
return
try:
curr_process.cpu_affinity(core_ids)
except Exception:
logger.error("Cannot pin current process to cores")
return
logger.debug(f"Pinning current process to cores {core_ids}")
def f():
return ray.get_resource_ids()
def f(self):
return ray.get_resource_ids()
def get_actor_method_resources(self):
return ray.get_resource_ids()
def __init__(self):
self.resources = ray.get_resource_ids()
def _ray_get_actor_cpus():
# Get through resource IDs
resource_ids = ray.get_resource_ids()
if "CPU" in resource_ids:
return sum(cpu[1] for cpu in resource_ids["CPU"])
return None
def get_num_cpus():
return len(ray.get_resource_ids()['CPU'])