wehub-resource-sync
545 lines
17 KiB
Python
545 lines
17 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
import asyncio
|
|
import datetime as dt
|
|
import fnmatch
|
|
import glob
|
|
import importlib
|
|
import json
|
|
import json_repair
|
|
import numpy as np
|
|
import os
|
|
import random
|
|
import re
|
|
import shutil
|
|
import socket
|
|
import subprocess
|
|
import sys
|
|
import threading
|
|
import time
|
|
import torch
|
|
import torch.distributed as dist
|
|
from contextlib import contextmanager
|
|
from functools import wraps
|
|
from transformers import HfArgumentParser, enable_full_determinism, set_seed
|
|
from transformers.utils import strtobool
|
|
from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Tuple, Type, TypeVar, Union
|
|
|
|
from .env import is_dist, is_master
|
|
from .logger import get_logger
|
|
from .np_utils import stat_array
|
|
|
|
logger = get_logger()
|
|
|
|
|
|
def check_json_format(obj: Any, token_safe: bool = True) -> Any:
|
|
if obj is None or isinstance(obj, (int, float, str, complex)): # bool is a subclass of int
|
|
return obj
|
|
if isinstance(obj, bytes):
|
|
return '<<<bytes>>>'
|
|
if isinstance(obj, (torch.dtype, torch.device)):
|
|
obj = str(obj)
|
|
return obj[len('torch.'):] if obj.startswith('torch.') else obj
|
|
|
|
if isinstance(obj, Sequence):
|
|
res = []
|
|
for x in obj:
|
|
res.append(check_json_format(x, token_safe))
|
|
elif isinstance(obj, Mapping):
|
|
res = {}
|
|
for k, v in obj.items():
|
|
if token_safe and isinstance(k, str) and '_token' in k and isinstance(v, str):
|
|
res[k] = None
|
|
else:
|
|
res[k] = check_json_format(v, token_safe)
|
|
else:
|
|
if token_safe:
|
|
unsafe_items = {}
|
|
for k, v in obj.__dict__.items():
|
|
if '_token' in k:
|
|
unsafe_items[k] = v
|
|
setattr(obj, k, None)
|
|
res = repr(obj)
|
|
# recover
|
|
for k, v in unsafe_items.items():
|
|
setattr(obj, k, v)
|
|
else:
|
|
res = repr(obj) # e.g. function, object
|
|
return res
|
|
|
|
|
|
def _get_version(work_dir: str) -> int:
|
|
if os.path.isdir(work_dir):
|
|
fnames = os.listdir(work_dir)
|
|
else:
|
|
fnames = []
|
|
v_list = [-1]
|
|
for fname in fnames:
|
|
m = re.match(r'v(\d+)', fname)
|
|
if m is None:
|
|
continue
|
|
v = m.group(1)
|
|
v_list.append(int(v))
|
|
return max(v_list) + 1
|
|
|
|
|
|
def format_time(seconds):
|
|
days = int(seconds // (24 * 3600))
|
|
hours = int((seconds % (24 * 3600)) // 3600)
|
|
minutes = int((seconds % 3600) // 60)
|
|
seconds = round(seconds % 60)
|
|
|
|
if days > 0:
|
|
time_str = f'{days}d {hours}h {minutes}m {seconds}s'
|
|
elif hours > 0:
|
|
time_str = f'{hours}h {minutes}m {seconds}s'
|
|
elif minutes > 0:
|
|
time_str = f'{minutes}m {seconds}s'
|
|
else:
|
|
time_str = f'{seconds}s'
|
|
|
|
return time_str
|
|
|
|
|
|
def deep_getattr(obj, attr: str, default=None):
|
|
attrs = attr.split('.')
|
|
for a in attrs:
|
|
if obj is None:
|
|
break
|
|
if isinstance(obj, dict):
|
|
obj = obj.get(a, default)
|
|
else:
|
|
obj = getattr(obj, a, default)
|
|
return obj
|
|
|
|
|
|
def seed_everything(seed: Optional[int] = None, full_determinism: bool = False) -> int:
|
|
|
|
if seed is None:
|
|
seed_max = np.iinfo(np.int32).max
|
|
seed = random.randint(0, seed_max)
|
|
|
|
if full_determinism:
|
|
enable_full_determinism(seed)
|
|
else:
|
|
set_seed(seed)
|
|
return seed
|
|
|
|
|
|
def add_version_to_work_dir(work_dir: str) -> str:
|
|
"""add version"""
|
|
version = _get_version(work_dir)
|
|
time = dt.datetime.now().strftime('%Y%m%d-%H%M%S')
|
|
sub_folder = f'v{version}-{time}'
|
|
if dist.is_initialized() and is_dist():
|
|
obj_list = [sub_folder]
|
|
dist.broadcast_object_list(obj_list)
|
|
sub_folder = obj_list[0]
|
|
|
|
work_dir = os.path.join(work_dir, sub_folder)
|
|
return work_dir
|
|
|
|
|
|
_T = TypeVar('_T')
|
|
|
|
|
|
def _patch_args(class_type):
|
|
try:
|
|
for k, v in class_type.__annotations__.items():
|
|
if v == Union[str, dict, type(None)]:
|
|
class_type.__annotations__[k] = Union[dict, str, type(None)]
|
|
except Exception:
|
|
logger.warning('patch args failed')
|
|
|
|
|
|
@contextmanager
|
|
def _patch_get_type_hints():
|
|
# Fix parsing string arguments into dicts
|
|
from transformers import hf_argparser
|
|
origin_get_type_hints = hf_argparser.get_type_hints
|
|
|
|
def get_type_hints(*args, **kwargs):
|
|
kwargs = origin_get_type_hints(*args, **kwargs)
|
|
for k, v in kwargs.items():
|
|
if v == Union[str, dict, type(None)]:
|
|
kwargs[k] = Union[dict, str, type(None)]
|
|
return kwargs
|
|
|
|
hf_argparser.get_type_hints = get_type_hints
|
|
try:
|
|
yield
|
|
finally:
|
|
hf_argparser.get_type_hints = origin_get_type_hints
|
|
|
|
|
|
def parse_args(class_type: Type[_T], argv: Optional[List[str]] = None) -> Tuple[_T, List[str]]:
|
|
with _patch_get_type_hints():
|
|
parser = HfArgumentParser([class_type])
|
|
_ray_args = os.environ.get('RAY_SWIFT_ARGS')
|
|
if _ray_args:
|
|
argv = json.loads(_ray_args)
|
|
elif argv is None:
|
|
argv = sys.argv[1:]
|
|
args, remaining_args = parser.parse_args_into_dataclasses(argv, return_remaining_strings=True)
|
|
return args, remaining_args
|
|
|
|
|
|
def parse_args_from_dict(class_type: Type[_T], args: Dict[str, Any]) -> _T:
|
|
with _patch_get_type_hints():
|
|
parser = HfArgumentParser([class_type])
|
|
return parser.parse_dict(args, allow_extra_keys=True)[0]
|
|
|
|
|
|
def lower_bound(lo: int, hi: int, cond: Callable[[int], bool]) -> int:
|
|
# The lower bound satisfying the condition "cond".
|
|
while lo < hi:
|
|
mid = (lo + hi) >> 1
|
|
if cond(mid):
|
|
hi = mid
|
|
else:
|
|
lo = mid + 1
|
|
return lo
|
|
|
|
|
|
def upper_bound(lo: int, hi: int, cond: Callable[[int], bool]) -> int:
|
|
# The upper bound satisfying the condition "cond".
|
|
while lo < hi:
|
|
mid = (lo + hi + 1) >> 1 # lo + (hi-lo+1)>>1
|
|
if cond(mid):
|
|
lo = mid
|
|
else:
|
|
hi = mid - 1
|
|
return lo
|
|
|
|
|
|
def test_time(func: Callable[[], _T],
|
|
number: int = 1,
|
|
warmup: int = 0,
|
|
timer: Optional[Callable[[], float]] = None) -> _T:
|
|
# timer: e.g. time_synchronize
|
|
timer = timer if timer is not None else time.perf_counter
|
|
|
|
ts = []
|
|
res = None
|
|
# warmup
|
|
for _ in range(warmup):
|
|
res = func()
|
|
|
|
for _ in range(number):
|
|
t1 = timer()
|
|
res = func()
|
|
t2 = timer()
|
|
ts.append(t2 - t1)
|
|
|
|
ts = np.array(ts)
|
|
_, stat_str = stat_array(ts)
|
|
# print
|
|
logger.info(f'time[number={number}]: {stat_str}')
|
|
return res
|
|
|
|
|
|
def read_multi_line(addi_prompt: str = '') -> str:
|
|
res = []
|
|
prompt = f'<<<{addi_prompt} '
|
|
while True:
|
|
text = input(prompt) + '\n'
|
|
prompt = ''
|
|
res.append(text)
|
|
if text.endswith('#\n'):
|
|
res[-1] = text[:-2]
|
|
break
|
|
return ''.join(res)
|
|
|
|
|
|
def subprocess_run(command: List[str], env: Optional[Dict[str, str]] = None, stdout=None, stderr=None):
|
|
# stdout stderr: e.g. subprocess.PIPE.
|
|
import shlex
|
|
command_str = ' '.join(shlex.quote(a) for a in command)
|
|
logger.info_if(f'Run the command: `{command_str}`', is_master())
|
|
resp = subprocess.run(command, env=env, stdout=stdout, stderr=stderr)
|
|
resp.check_returncode()
|
|
return resp
|
|
|
|
|
|
def get_env_args(args_name: str, type_func: Callable[[str], _T], default_value: Optional[_T]) -> Optional[_T]:
|
|
args_name_upper = args_name.upper()
|
|
value = os.getenv(args_name_upper)
|
|
if value is None:
|
|
value = default_value
|
|
log_info = (f'Setting {args_name}: {default_value}. '
|
|
f'You can adjust this hyperparameter through the environment variable: `{args_name_upper}`.')
|
|
else:
|
|
if type_func is bool:
|
|
value = strtobool(value)
|
|
value = type_func(value)
|
|
log_info = f'Using environment variable `{args_name_upper}`, Setting {args_name}: {value}.'
|
|
logger.info_once(log_info)
|
|
return value
|
|
|
|
|
|
def find_node_ip() -> Optional[str]:
|
|
import psutil
|
|
main_ip, virtual_ip = None, None
|
|
for name, addrs in sorted(psutil.net_if_addrs().items()):
|
|
for addr in addrs:
|
|
if addr.family.name == 'AF_INET' and not addr.address.startswith('127.'):
|
|
# Heuristic to prefer non-virtual interfaces
|
|
if any(s in name for s in ['lo', 'docker', 'veth', 'vmnet']):
|
|
if virtual_ip is None:
|
|
virtual_ip = addr.address
|
|
else:
|
|
if main_ip is None:
|
|
main_ip = addr.address
|
|
return main_ip or virtual_ip
|
|
|
|
|
|
def find_free_port(start_port: Optional[int] = None, retry: int = 100) -> int:
|
|
if start_port is None:
|
|
start_port = 0
|
|
for port in range(start_port, start_port + retry):
|
|
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
|
|
try:
|
|
sock.bind(('', port))
|
|
port = sock.getsockname()[1]
|
|
break
|
|
except OSError:
|
|
pass
|
|
return port
|
|
|
|
|
|
def copy_files_by_pattern(source_dir, dest_dir, patterns, exclude_patterns=None):
|
|
if not os.path.exists(dest_dir):
|
|
os.makedirs(dest_dir)
|
|
|
|
if isinstance(patterns, str):
|
|
patterns = [patterns]
|
|
|
|
if exclude_patterns is None:
|
|
exclude_patterns = []
|
|
elif isinstance(exclude_patterns, str):
|
|
exclude_patterns = [exclude_patterns]
|
|
|
|
def should_exclude_file(file_path, file_name):
|
|
for exclude_pattern in exclude_patterns:
|
|
if fnmatch.fnmatch(file_name, exclude_pattern):
|
|
return True
|
|
rel_file_path = os.path.relpath(file_path, source_dir)
|
|
if fnmatch.fnmatch(rel_file_path, exclude_pattern):
|
|
return True
|
|
return False
|
|
|
|
for pattern in patterns:
|
|
pattern_parts = pattern.split(os.path.sep)
|
|
if len(pattern_parts) > 1:
|
|
subdir_pattern = os.path.sep.join(pattern_parts[:-1])
|
|
file_pattern = pattern_parts[-1]
|
|
|
|
for root, dirs, files in os.walk(source_dir):
|
|
rel_path = os.path.relpath(root, source_dir)
|
|
if rel_path == '.' or (rel_path != '.' and not fnmatch.fnmatch(rel_path, subdir_pattern)):
|
|
continue
|
|
|
|
for file in files:
|
|
if fnmatch.fnmatch(file, file_pattern):
|
|
file_path = os.path.join(root, file)
|
|
|
|
if should_exclude_file(file_path, file):
|
|
continue
|
|
|
|
target_dir = os.path.join(dest_dir, rel_path)
|
|
if not os.path.exists(target_dir):
|
|
os.makedirs(target_dir)
|
|
dest_file = os.path.join(target_dir, file)
|
|
|
|
if not os.path.exists(dest_file):
|
|
shutil.copy2(file_path, dest_file)
|
|
else:
|
|
search_path = os.path.join(source_dir, pattern)
|
|
matched_files = glob.glob(search_path)
|
|
|
|
for file_path in matched_files:
|
|
if os.path.isfile(file_path):
|
|
file_name = os.path.basename(file_path)
|
|
|
|
if should_exclude_file(file_path, file_name):
|
|
continue
|
|
|
|
destination = os.path.join(dest_dir, file_name)
|
|
if not os.path.exists(destination):
|
|
shutil.copy2(file_path, destination)
|
|
|
|
|
|
def split_list(ori_list: List[_T], num_shards: int, contiguous=True) -> List[List[_T]]:
|
|
shard = []
|
|
if contiguous:
|
|
idx_list = np.linspace(0, len(ori_list), num_shards + 1, dtype=np.int64)
|
|
for i in range(len(idx_list) - 1):
|
|
shard.append(ori_list[idx_list[i]:idx_list[i + 1]])
|
|
else:
|
|
ori_list = np.array(ori_list)
|
|
for i in range(num_shards):
|
|
shard.append(ori_list[np.arange(i, len(ori_list), num_shards)].tolist())
|
|
return shard
|
|
|
|
|
|
def patch_getattr(obj_cls, item_name: str):
|
|
if hasattr(obj_cls, '_patch'): # avoid double patch
|
|
return
|
|
|
|
def __new_getattr__(self, key: str):
|
|
try:
|
|
return super(self.__class__, self).__getattr__(key)
|
|
except AttributeError:
|
|
if item_name in dir(self):
|
|
item = getattr(self, item_name)
|
|
return getattr(item, key)
|
|
raise
|
|
|
|
obj_cls.__getattr__ = __new_getattr__
|
|
obj_cls._patch = True
|
|
|
|
|
|
def import_external_file(file_path: str):
|
|
file_path = os.path.abspath(os.path.expanduser(file_path))
|
|
py_dir, py_file = os.path.split(file_path)
|
|
assert os.path.isdir(py_dir), f'py_dir: {py_dir}'
|
|
sys.path.insert(0, py_dir)
|
|
return importlib.import_module(py_file.split('.', 1)[0])
|
|
|
|
|
|
def json_parse_to_dict(value: Union[str, Dict, None], strict: bool = True) -> Union[str, Dict]:
|
|
"""Convert a JSON string or JSON file into a dict"""
|
|
# If the value could potentially be a string, it is generally advisable to set strict to False.
|
|
if value is None:
|
|
value = {}
|
|
elif isinstance(value, str):
|
|
if os.path.exists(value): # local path
|
|
with open(value, 'r', encoding='utf-8') as f:
|
|
value = json.load(f)
|
|
else: # json str
|
|
try:
|
|
value = json.loads(value)
|
|
except json.JSONDecodeError:
|
|
if strict:
|
|
try:
|
|
# fix malformed json string, e.g., incorrect quotation marks
|
|
old_value = value
|
|
value = json_repair.repair_json(value)
|
|
logger.warning(f'Unable to parse json string, try to repair it, '
|
|
f"the string before and after repair are '{old_value}' | '{value}'")
|
|
value = json.loads(value)
|
|
except Exception:
|
|
logger.error(f"Unable to parse json string: '{value}', and try to repair failed")
|
|
raise
|
|
return value
|
|
|
|
|
|
def retry_decorator(retry=3):
|
|
|
|
def _retry(func):
|
|
|
|
@wraps(func)
|
|
def new_func(*args, **kwargs):
|
|
i = 1
|
|
while True:
|
|
try:
|
|
return func(*args, **kwargs)
|
|
except Exception:
|
|
if i == retry:
|
|
raise
|
|
i += 1
|
|
|
|
return new_func
|
|
|
|
return _retry
|
|
|
|
|
|
def start_event_loop_in_daemon(name: str = None) -> Tuple[threading.Thread, asyncio.AbstractEventLoop, threading.Event]:
|
|
"""Create a new daemon thread that runs an asyncio event loop.
|
|
|
|
Args:
|
|
name: Name of the thread. If None, the default thread naming will be used.
|
|
|
|
Returns:
|
|
tuple: (thread, loop, loop_ready_event)
|
|
- thread: The thread running the event loop.
|
|
- loop: The event loop being run in the thread.
|
|
- loop_ready_event: An event that is set when the loop is ready.
|
|
"""
|
|
loop = asyncio.new_event_loop()
|
|
loop_ready_event = threading.Event()
|
|
|
|
def run_loop():
|
|
asyncio.set_event_loop(loop)
|
|
loop_ready_event.set()
|
|
loop.run_forever()
|
|
|
|
thread = threading.Thread(target=run_loop, name=name, daemon=True)
|
|
thread.start()
|
|
return thread, loop, loop_ready_event
|
|
|
|
|
|
def shutdown_event_loop_in_daemon(thread: threading.Thread = None, loop: asyncio.AbstractEventLoop = None) -> None:
|
|
"""Shutdown an asyncio event loop running in a separate thread.
|
|
|
|
This function stops the event loop and waits for the associated thread to finish execution.
|
|
|
|
Args:
|
|
thread: The thread running the event loop.
|
|
loop: The asyncio event loop to shut down.
|
|
"""
|
|
if loop is None or thread is None:
|
|
return
|
|
loop.call_soon_threadsafe(loop.stop)
|
|
thread.join(timeout=5)
|
|
|
|
|
|
def remove_response(messages) -> Optional[str]:
|
|
"""
|
|
Removes and returns the content of the last message if its role is 'assistant'.
|
|
|
|
Args:
|
|
messages (List[Dict]):
|
|
A list of message dictionaries, each typically containing a 'role' and 'content' key.
|
|
|
|
Returns:
|
|
Optional[str]:
|
|
The content of the removed 'assistant' message if present;
|
|
otherwise, returns None. The original messages list is modified in place.
|
|
"""
|
|
last_role = messages[-1]['role'] if messages else None
|
|
if last_role == 'assistant':
|
|
return messages.pop()['content']
|
|
|
|
|
|
def to_abspath(path: Union[str, List[str], None], check_path_exist: bool = False) -> Union[str, List[str], None]:
|
|
"""Check the path for validity and convert it to an absolute path.
|
|
|
|
Args:
|
|
path: The path to be checked/converted
|
|
check_path_exist: Whether to check if the path exists
|
|
|
|
Returns:
|
|
Absolute path
|
|
"""
|
|
if path is None:
|
|
return
|
|
elif isinstance(path, str):
|
|
# Remove user path prefix and convert to absolute path.
|
|
path = os.path.abspath(os.path.expanduser(path))
|
|
if check_path_exist and not os.path.exists(path):
|
|
raise FileNotFoundError(f"path: '{path}'")
|
|
return path
|
|
assert isinstance(path, list), f'path: {path}'
|
|
res = []
|
|
for v in path:
|
|
res.append(to_abspath(v, check_path_exist))
|
|
return res
|
|
|
|
|
|
def swanlab_get_run():
|
|
try:
|
|
import swanlab
|
|
return swanlab.get_run()
|
|
except (RuntimeError, ImportError):
|
|
return None
|