wehub-resource-sync
492 lines
16 KiB
Python
492 lines
16 KiB
Python
import asyncio
|
|
import binascii
|
|
import errno
|
|
import importlib
|
|
import inspect
|
|
import logging
|
|
import os
|
|
import random
|
|
import string
|
|
import sys
|
|
import tempfile
|
|
import time
|
|
from abc import ABC, abstractmethod
|
|
from inspect import signature
|
|
from types import ModuleType
|
|
from typing import Any, Coroutine, Dict, Optional, Tuple
|
|
|
|
import ray
|
|
from ray._raylet import GcsClient, NodeID
|
|
from ray.core.generated.gcs_pb2 import GcsNodeInfo
|
|
from ray.core.generated.gcs_service_pb2 import GetAllNodeInfoRequest
|
|
|
|
import psutil
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def env_integer(key, default):
|
|
if key in os.environ:
|
|
value = os.environ[key]
|
|
try:
|
|
return int(value)
|
|
except ValueError:
|
|
logger.debug(
|
|
f"Found {key} in environment, but value must "
|
|
f"be an integer. Got: {value}. Returning "
|
|
f"provided default {default}."
|
|
)
|
|
return default
|
|
return default
|
|
|
|
|
|
def env_float(key, default):
|
|
if key in os.environ:
|
|
value = os.environ[key]
|
|
try:
|
|
return float(value)
|
|
except ValueError:
|
|
logger.debug(
|
|
f"Found {key} in environment, but value must "
|
|
f"be a float. Got: {value}. Returning "
|
|
f"provided default {default}."
|
|
)
|
|
return default
|
|
return default
|
|
|
|
|
|
def env_bool(key, default):
|
|
if key in os.environ:
|
|
val = os.environ[key].lower()
|
|
return val == "true" or val == "1"
|
|
return default
|
|
|
|
|
|
def import_module_and_attr(
|
|
full_path: str, *, reload_module: bool = False
|
|
) -> Tuple[ModuleType, Any]:
|
|
"""Given a full import path to a module attr, return the imported module and attr.
|
|
|
|
If `reload_module` is set, the module will be reloaded using `importlib.reload`.
|
|
|
|
Args:
|
|
full_path: The full import path to the module and attr.
|
|
reload_module: Whether to reload the module.
|
|
|
|
Returns:
|
|
A tuple of the imported module and attr.
|
|
"""
|
|
if ":" in full_path:
|
|
if full_path.count(":") > 1:
|
|
raise ValueError(
|
|
f'Got invalid import path "{full_path}". An '
|
|
"import path may have at most one colon."
|
|
)
|
|
module_name, attr_name = full_path.split(":")
|
|
else:
|
|
last_period_idx = full_path.rfind(".")
|
|
module_name = full_path[:last_period_idx]
|
|
attr_name = full_path[last_period_idx + 1 :]
|
|
module = importlib.import_module(module_name)
|
|
if reload_module:
|
|
importlib.reload(module)
|
|
return module, getattr(module, attr_name)
|
|
|
|
|
|
def import_attr(full_path: str, *, reload_module: bool = False) -> Any:
|
|
"""Given a full import path to a module attr, return the imported attr.
|
|
|
|
If `reload_module` is set, the module will be reloaded using `importlib.reload`.
|
|
|
|
For example, the following are equivalent:
|
|
MyClass = import_attr("module.submodule:MyClass")
|
|
MyClass = import_attr("module.submodule.MyClass")
|
|
from module.submodule import MyClass
|
|
|
|
Args:
|
|
full_path: The full import path to the module and attr.
|
|
reload_module: Whether to reload the module.
|
|
|
|
Returns:
|
|
Imported attr
|
|
"""
|
|
return import_module_and_attr(full_path, reload_module=reload_module)[1]
|
|
|
|
|
|
def get_or_create_event_loop() -> asyncio.AbstractEventLoop:
|
|
"""Get a running async event loop if one exists, otherwise create one.
|
|
|
|
This function serves as a proxy for the deprecating get_event_loop().
|
|
It tries to get the running loop first, and if no running loop
|
|
could be retrieved:
|
|
- For python version <3.10: it falls back to the get_event_loop
|
|
call.
|
|
- For python version >= 3.10: it uses the same python implementation
|
|
of _get_event_loop() at asyncio/events.py.
|
|
|
|
Ideally, one should use high level APIs like asyncio.run() with python
|
|
version >= 3.7, if not possible, one should create and manage the event
|
|
loops explicitly.
|
|
"""
|
|
vers_info = sys.version_info
|
|
if vers_info.major >= 3 and vers_info.minor >= 10:
|
|
# This follows the implementation of the deprecating `get_event_loop`
|
|
# in python3.10's asyncio. See python3.10/asyncio/events.py
|
|
# _get_event_loop()
|
|
try:
|
|
loop = asyncio.get_running_loop()
|
|
assert loop is not None
|
|
return loop
|
|
except RuntimeError as e:
|
|
# No running loop, relying on the error message as for now to
|
|
# differentiate runtime errors.
|
|
assert "no running event loop" in str(e)
|
|
try:
|
|
loop = asyncio.get_event_loop_policy().get_event_loop()
|
|
return loop
|
|
except RuntimeError:
|
|
# Python 3.14+: get_event_loop() no longer creates a loop automatically
|
|
# See: https://docs.python.org/3.14/library/asyncio-eventloop.html
|
|
loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(loop)
|
|
return loop
|
|
|
|
return asyncio.get_event_loop()
|
|
|
|
|
|
_BACKGROUND_TASKS = set()
|
|
|
|
|
|
def run_background_task(coroutine: Coroutine) -> asyncio.Task:
|
|
"""Schedule a task reliably to the event loop.
|
|
|
|
This API is used when you don't want to cache the reference of `asyncio.Task`.
|
|
For example,
|
|
|
|
```
|
|
get_event_loop().create_task(coroutine(*args))
|
|
```
|
|
|
|
The above code doesn't guarantee to schedule the coroutine to the event loops
|
|
|
|
When using create_task in a "fire and forget" way, we should keep the references
|
|
alive for the reliable execution. This API is used to fire and forget
|
|
asynchronous execution.
|
|
|
|
https://docs.python.org/3/library/asyncio-task.html#creating-tasks
|
|
"""
|
|
task = get_or_create_event_loop().create_task(coroutine)
|
|
# Add task to the set. This creates a strong reference.
|
|
_BACKGROUND_TASKS.add(task)
|
|
|
|
# To prevent keeping references to finished tasks forever,
|
|
# make each task remove its own reference from the set after
|
|
# completion:
|
|
task.add_done_callback(_BACKGROUND_TASKS.discard)
|
|
return task
|
|
|
|
|
|
# Used in gpu detection
|
|
RESOURCE_CONSTRAINT_PREFIX = "accelerator_type:"
|
|
PLACEMENT_GROUP_BUNDLE_RESOURCE_NAME = "bundle"
|
|
|
|
|
|
def resources_from_ray_options(options_dict: Dict[str, Any]) -> Dict[str, Any]:
|
|
"""Determine a task's resource requirements.
|
|
|
|
Args:
|
|
options_dict: The dictionary that contains resources requirements.
|
|
|
|
Returns:
|
|
A dictionary of the resource requirements for the task.
|
|
"""
|
|
resources = (options_dict.get("resources") or {}).copy()
|
|
|
|
if "CPU" in resources or "GPU" in resources:
|
|
raise ValueError(
|
|
"The resources dictionary must not contain the key 'CPU' or 'GPU'"
|
|
)
|
|
elif "memory" in resources or "object_store_memory" in resources:
|
|
raise ValueError(
|
|
"The resources dictionary must not "
|
|
"contain the key 'memory' or 'object_store_memory'"
|
|
)
|
|
elif PLACEMENT_GROUP_BUNDLE_RESOURCE_NAME in resources:
|
|
raise ValueError(
|
|
"The resource should not include `bundle` which "
|
|
f"is reserved for Ray. resources: {resources}"
|
|
)
|
|
|
|
num_cpus = options_dict.get("num_cpus")
|
|
num_gpus = options_dict.get("num_gpus")
|
|
memory = options_dict.get("memory")
|
|
object_store_memory = options_dict.get("object_store_memory")
|
|
accelerator_type = options_dict.get("accelerator_type")
|
|
|
|
if num_cpus is not None:
|
|
resources["CPU"] = num_cpus
|
|
if num_gpus is not None:
|
|
resources["GPU"] = num_gpus
|
|
if memory is not None:
|
|
resources["memory"] = int(memory)
|
|
if object_store_memory is not None:
|
|
resources["object_store_memory"] = object_store_memory
|
|
if accelerator_type is not None:
|
|
resources[f"{RESOURCE_CONSTRAINT_PREFIX}{accelerator_type}"] = 0.001
|
|
|
|
return resources
|
|
|
|
|
|
# Match the standard alphabet used for UUIDs.
|
|
RANDOM_STRING_ALPHABET = string.ascii_lowercase + string.digits
|
|
|
|
|
|
def get_random_alphanumeric_string(length: int):
|
|
"""Generates random string of length consisting exclusively of
|
|
- Lower-case ASCII chars
|
|
- Digits
|
|
"""
|
|
return "".join(random.choices(RANDOM_STRING_ALPHABET, k=length))
|
|
|
|
|
|
_PRINTED_WARNING = set()
|
|
|
|
|
|
def get_call_location(back: int = 1):
|
|
"""
|
|
Get the location (filename and line number) of a function caller, `back`
|
|
frames up the stack.
|
|
|
|
Args:
|
|
back: The number of frames to go up the stack, not including this
|
|
function.
|
|
|
|
Returns:
|
|
A string with the filename and line number of the caller.
|
|
For example, "myfile.py:123".
|
|
"""
|
|
stack = inspect.stack()
|
|
try:
|
|
frame = stack[back + 1]
|
|
return f"{frame.filename}:{frame.lineno}"
|
|
except IndexError:
|
|
return "UNKNOWN"
|
|
|
|
|
|
def resolve_user_ray_temp_dir(gcs_client: GcsClient, node_id: str):
|
|
"""
|
|
Get the ray temp directory.
|
|
|
|
If a temp dir was specified for this node, this function will
|
|
retrieve the information from GCS. Otherwise, it will fallback to the
|
|
default ray temp directory.
|
|
|
|
Args:
|
|
gcs_client: The GCS client.
|
|
node_id: The ID of the node to fetch the temp dir for.
|
|
E.g.: "1a9904d8aa3de65367830e2aef6313a5b2e9d4b0e3725e0dceeacb1b"
|
|
(hex string representation of the node ID)
|
|
|
|
Returns:
|
|
The path to the ray temp directory.
|
|
"""
|
|
# check if temp dir is available from runtime context
|
|
if ray.is_initialized() and ray.get_runtime_context().get_node_id() == node_id:
|
|
return ray.get_runtime_context().get_temp_dir()
|
|
|
|
# Fetch temp dir as specified by --temp-dir at creation time.
|
|
try:
|
|
# Create node selector for node_id filter
|
|
node_selector = GetAllNodeInfoRequest.NodeSelector()
|
|
node_selector.node_id = NodeID.from_hex(node_id).binary()
|
|
|
|
node_infos = gcs_client.get_all_node_info(
|
|
node_selectors=[node_selector],
|
|
state_filter=GcsNodeInfo.GcsNodeState.ALIVE,
|
|
).values()
|
|
except Exception as e:
|
|
raise Exception(
|
|
f"Failed to get node info from GCS when fetching tempdir for node {node_id}: {e}"
|
|
)
|
|
if not node_infos:
|
|
raise Exception(
|
|
f"No node info associated with ALIVE state found for node {node_id} in GCS"
|
|
)
|
|
|
|
node_info = next(iter(node_infos))
|
|
if node_info is not None:
|
|
temp_dir = getattr(node_info, "temp_dir", None)
|
|
if temp_dir is not None:
|
|
return temp_dir
|
|
else:
|
|
raise Exception(
|
|
"Node temp_dir was not found in NodeInfo. did the node's raylet start successfully?"
|
|
)
|
|
|
|
|
|
def get_default_system_temp_dir():
|
|
if "RAY_TMPDIR" in os.environ:
|
|
return os.environ["RAY_TMPDIR"]
|
|
elif sys.platform.startswith("linux") and "TMPDIR" in os.environ:
|
|
return os.environ["TMPDIR"]
|
|
elif sys.platform.startswith("darwin") or sys.platform.startswith("linux"):
|
|
# Ideally we wouldn't need this fallback, but keep it for now for
|
|
# for compatibility
|
|
tempdir = os.path.join(os.sep, "tmp")
|
|
else:
|
|
tempdir = tempfile.gettempdir()
|
|
|
|
return tempdir
|
|
|
|
|
|
def get_default_ray_temp_dir():
|
|
return os.path.join(get_default_system_temp_dir(), "ray")
|
|
|
|
|
|
def get_ray_address_file(temp_dir: Optional[str]):
|
|
if temp_dir is None:
|
|
temp_dir = get_default_ray_temp_dir()
|
|
return os.path.join(temp_dir, "ray_current_cluster")
|
|
|
|
|
|
def reset_ray_address(temp_dir: Optional[str] = None):
|
|
address_file = get_ray_address_file(temp_dir)
|
|
if os.path.exists(address_file):
|
|
try:
|
|
os.remove(address_file)
|
|
except OSError:
|
|
pass
|
|
|
|
|
|
def load_class(path):
|
|
"""Load a class at runtime given a full path.
|
|
|
|
Example of the path: mypkg.mysubpkg.myclass
|
|
"""
|
|
class_data = path.split(".")
|
|
if len(class_data) < 2:
|
|
raise ValueError("You need to pass a valid path like mymodule.provider_class")
|
|
module_path = ".".join(class_data[:-1])
|
|
class_str = class_data[-1]
|
|
module = importlib.import_module(module_path)
|
|
return getattr(module, class_str)
|
|
|
|
|
|
def get_system_memory(
|
|
# For cgroups v1:
|
|
memory_limit_filename: str = "/sys/fs/cgroup/memory/memory.limit_in_bytes",
|
|
# For cgroups v2:
|
|
memory_limit_filename_v2: str = "/sys/fs/cgroup/memory.max",
|
|
):
|
|
"""Return the total amount of system memory in bytes.
|
|
|
|
Args:
|
|
memory_limit_filename: The path to the file that contains the memory
|
|
limit for the Docker container. Defaults to
|
|
/sys/fs/cgroup/memory/memory.limit_in_bytes.
|
|
memory_limit_filename_v2: The path to the file that contains the memory
|
|
limit for the Docker container in cgroups v2. Defaults to
|
|
/sys/fs/cgroup/memory.max.
|
|
|
|
Returns:
|
|
The total amount of system memory in bytes.
|
|
"""
|
|
# Try to accurately figure out the memory limit if we are in a docker
|
|
# container. Note that this file is not specific to Docker and its value is
|
|
# often much larger than the actual amount of memory.
|
|
docker_limit = None
|
|
if os.path.exists(memory_limit_filename):
|
|
with open(memory_limit_filename, "r") as f:
|
|
docker_limit = int(f.read().strip())
|
|
elif os.path.exists(memory_limit_filename_v2):
|
|
with open(memory_limit_filename_v2, "r") as f:
|
|
# Don't forget to strip() the newline:
|
|
max_file = f.read().strip()
|
|
if max_file.isnumeric():
|
|
docker_limit = int(max_file)
|
|
else:
|
|
# max_file is "max", i.e. is unset.
|
|
docker_limit = None
|
|
|
|
# Use psutil if it is available.
|
|
psutil_memory_in_bytes = psutil.virtual_memory().total
|
|
|
|
if docker_limit is not None:
|
|
# We take the min because the cgroup limit is very large if we aren't
|
|
# in Docker.
|
|
return min(docker_limit, psutil_memory_in_bytes)
|
|
|
|
return psutil_memory_in_bytes
|
|
|
|
|
|
def binary_to_hex(identifier):
|
|
hex_identifier = binascii.hexlify(identifier)
|
|
hex_identifier = hex_identifier.decode()
|
|
return hex_identifier
|
|
|
|
|
|
def hex_to_binary(hex_identifier):
|
|
return binascii.unhexlify(hex_identifier)
|
|
|
|
|
|
def try_make_directory_shared(directory_path):
|
|
try:
|
|
os.chmod(directory_path, 0o0777)
|
|
except OSError as e:
|
|
# Silently suppress the PermissionError that is thrown by the chmod.
|
|
# This is done because the user attempting to change the permissions
|
|
# on a directory may not own it. The chmod is attempted whether the
|
|
# directory is new or not to avoid race conditions.
|
|
# ray-project/ray/#3591
|
|
if e.errno in [errno.EACCES, errno.EPERM]:
|
|
pass
|
|
else:
|
|
raise
|
|
|
|
|
|
def try_to_create_directory(directory_path):
|
|
# Attempt to create a directory that is globally readable/writable.
|
|
directory_path = os.path.expanduser(directory_path)
|
|
os.makedirs(directory_path, exist_ok=True)
|
|
# Change the log directory permissions so others can use it. This is
|
|
# important when multiple people are using the same machine.
|
|
try_make_directory_shared(directory_path)
|
|
|
|
|
|
def get_function_args(callable):
|
|
all_parameters = frozenset(signature(callable).parameters)
|
|
return list(all_parameters)
|
|
|
|
|
|
def decode(byte_str: str, allow_none: bool = False, encode_type: str = "utf-8"):
|
|
"""Make this unicode in Python 3, otherwise leave it as bytes.
|
|
|
|
Args:
|
|
byte_str: The byte string to decode.
|
|
allow_none: If true, then we will allow byte_str to be None in which
|
|
case we will return an empty string. TODO(rkn): Remove this flag.
|
|
This is only here to simplify upgrading to flatbuffers 1.10.0.
|
|
encode_type: The encoding type to use for decoding. Defaults to "utf-8".
|
|
|
|
Returns:
|
|
A byte string in Python 2 and a unicode string in Python 3.
|
|
"""
|
|
if byte_str is None and allow_none:
|
|
return ""
|
|
|
|
if not isinstance(byte_str, bytes):
|
|
raise ValueError(f"The argument {byte_str} must be a bytes object.")
|
|
return byte_str.decode(encode_type)
|
|
|
|
|
|
class TimerBase(ABC):
|
|
@abstractmethod
|
|
def time(self) -> float:
|
|
"""Return the current time."""
|
|
raise NotImplementedError
|
|
|
|
|
|
class Timer(TimerBase):
|
|
def time(self) -> float:
|
|
return time.time()
|