File size: 2,847 Bytes
8c9ba62 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 | # -*- coding: utf-8 -*-
"""Constants."""
from enum import Enum, EnumMeta
# names
EXPLORER_NAME = "explorer"
TRAINER_NAME = "trainer"
SELECTOR_METRIC = "selector_metric"
ROLLOUT_WEIGHT_SYNC_GROUP_NAME = "rollout_weight_sync"
DEBUG_NAMESPACE = "TRINITY_DEBUG_NAMESPACE"
# trinity env var names
CHECKPOINT_ROOT_DIR_ENV_VAR = "TRINITY_CHECKPOINT_ROOT_DIR"
PREVIOUS_STAGE_CHECKPOINT_DIR_ENV_VAR = "TRINITY_PREV_STAGE_CKPT_DIR"
MODEL_PATH_ENV_VAR = "TRINITY_MODEL_PATH"
TASKSET_PATH_ENV_VAR = "/nas/wjq/Trinity-RFT/tasksets"
BUFFER_PATH_ENV_VAR = "TRINITY_BUFFER_PATH"
PLUGIN_DIRS_ENV_VAR = "TRINITY_PLUGIN_DIRS"
LOG_DIR_ENV_VAR = "TRINITY_LOG_DIR" # log dir
LOG_LEVEL_ENV_VAR = "TRINITY_LOG_LEVEL" # global log level
LOG_NODE_IP_ENV_VAR = "TRINITY_LOG_NODE_IP" # whether to organize logs by node IP
# constants
MAX_MODEL_LEN = 4096
# enumerate types
class CaseInsensitiveEnumMeta(EnumMeta):
def __getitem__(cls, name):
return super().__getitem__(name.upper())
def __getattr__(cls, name):
if not name.startswith("_"):
return cls[name.upper()]
return super().__getattr__(name)
def __call__(cls, value, *args, **kwargs):
return super().__call__(value.lower(), *args, **kwargs)
class CaseInsensitiveEnum(Enum, metaclass=CaseInsensitiveEnumMeta):
pass
class PromptType(CaseInsensitiveEnum):
"""Prompt Type."""
MESSAGES = "messages" # a list of message dict
PLAINTEXT = "plaintext" # user prompt text and assistant response text
class StorageType(CaseInsensitiveEnum):
"""Storage Type."""
SQL = "sql"
QUEUE = "queue"
FILE = "file"
class SyncMethodEnumMeta(CaseInsensitiveEnumMeta):
def __call__(cls, value, *args, **kwargs):
if value == "online":
value = "nccl"
elif value == "offline":
value = "checkpoint"
try:
return super().__call__(value, *args, **kwargs)
except Exception:
raise ValueError(f"Invalid SyncMethod: {value}")
class SyncMethod(CaseInsensitiveEnum, metaclass=SyncMethodEnumMeta):
"""Sync Method."""
NCCL = "nccl"
CHECKPOINT = "checkpoint"
MEMORY = "memory"
class RunningStatus(Enum):
"""Running status of explorer and trainer."""
RUNNING = "running"
REQUIRE_SYNC = "require_sync"
WAITING_SYNC = "waiting_sync"
STOPPED = "stopped"
class OpType(Enum):
"""Operator type for reward shaping."""
ADD = "add"
SUB = "sub"
MUL = "mul"
DIV = "div"
class SyncStyle(CaseInsensitiveEnum):
FIXED = "fixed"
DYNAMIC_BY_TRAINER = "dynamic_by_trainer"
DYNAMIC_BY_EXPLORER = "dynamic_by_explorer"
class SaveStrategy(CaseInsensitiveEnum):
SINGLE_THREAD = "single_thread"
SINGLE_PROCESS = "single_process"
SINGLE_NODE = "single_node"
UNRESTRICTED = "unrestricted"
|