wehub-resource-sync
160 lines
6.5 KiB
Python
160 lines
6.5 KiB
Python
import os
|
|
from typing import Dict, Optional
|
|
|
|
from ray._common.constants import RAY_WARN_BLOCKING_GET_INSIDE_ASYNC_ENV_VAR
|
|
from ray._private.ray_constants import env_bool, env_set_by_user
|
|
|
|
# Unsupported configs can use this value to detect if the user has set it.
|
|
_UNSUPPORTED = "UNSUPPORTED"
|
|
_DEPRECATED = "DEPRECATED"
|
|
|
|
# The name of the file that is used to validate the storage.
|
|
VALIDATE_STORAGE_MARKER_FILENAME = ".validate_storage_marker"
|
|
# The name of the file that is used to store the checkpoint manager snapshot.
|
|
CHECKPOINT_MANAGER_SNAPSHOT_FILENAME = "checkpoint_manager_snapshot.json"
|
|
|
|
AWS_RETRYABLE_TOKENS = (
|
|
"AWS Error SLOW_DOWN",
|
|
"AWS Error INTERNAL_FAILURE",
|
|
"AWS Error SERVICE_UNAVAILABLE",
|
|
"AWS Error NETWORK_CONNECTION",
|
|
"AWS Error UNKNOWN",
|
|
)
|
|
|
|
# -----------------------------------------------------------------------
|
|
# Environment variables used in the controller, workers, and state actor.
|
|
#
|
|
# Be sure to update ENV_VARS_TO_PROPAGATE when adding new
|
|
# environment variables in this section.
|
|
# -----------------------------------------------------------------------
|
|
|
|
# Polling interval for the Train controller.
|
|
# This determines how many seconds the controller will wait between
|
|
# polling the worker group for its status.
|
|
HEALTH_CHECK_INTERVAL_S_ENV_VAR = "RAY_TRAIN_HEALTH_CHECK_INTERVAL_S"
|
|
DEFAULT_HEALTH_CHECK_INTERVAL_S: float = 2.0
|
|
|
|
# The time in seconds a worker health check must be hanging for
|
|
# before the controller marks the worker as dead and handles the failure.
|
|
WORKER_HEALTH_CHECK_TIMEOUT_S_ENV_VAR = "RAY_TRAIN_WORKER_HEALTH_CHECK_TIMEOUT_S"
|
|
DEFAULT_WORKER_HEALTH_CHECK_TIMEOUT_S: float = 10 * 60
|
|
|
|
# Timeout in seconds for the worker group to start.
|
|
WORKER_GROUP_START_TIMEOUT_S_ENV_VAR = "RAY_TRAIN_WORKER_GROUP_START_TIMEOUT_S"
|
|
DEFAULT_WORKER_GROUP_START_TIMEOUT_S: float = 60.0
|
|
|
|
# Time in seconds for collective operations before raising a timeout error.
|
|
COLLECTIVE_TIMEOUT_S_ENV_VAR = "RAY_TRAIN_COLLECTIVE_TIMEOUT_S"
|
|
# NOTE: Default to no timeout to avoid introducing more timeouts for users to configure.
|
|
# For example, users can already configure timeouts in torch distributed.
|
|
DEFAULT_COLLECTIVE_TIMEOUT_S: Optional[float] = None
|
|
# Interval in seconds to log a warning when waiting for a collective operation to complete.
|
|
COLLECTIVE_WARN_INTERVAL_S_ENV_VAR = "RAY_TRAIN_COLLECTIVE_WARN_INTERVAL_S"
|
|
DEFAULT_COLLECTIVE_WARN_INTERVAL_S: float = 60
|
|
|
|
# Interval in seconds to log a warning when waiting for a checkpoint upload fn operation to complete.
|
|
CHECKPOINT_UPLOAD_WARN_INTERVAL_S_ENV_VAR = (
|
|
"RAY_TRAIN_CHECKPOINT_UPLOAD_WARN_INTERVAL_S"
|
|
)
|
|
DEFAULT_CHECKPOINT_UPLOAD_WARN_INTERVAL_S: float = 60
|
|
|
|
# Feature flag for the preemption watcher. Default-on; provides a quick
|
|
# rollback path if the watcher actor misbehaves in a cluster.
|
|
ENABLE_PREEMPTION_WATCHER_ENV_VAR = "RAY_TRAIN_ENABLE_PREEMPTION_WATCHER"
|
|
DEFAULT_ENABLE_PREEMPTION_WATCHER: bool = True
|
|
|
|
# How often the preemption watcher polls Ray Core's drain state.
|
|
PREEMPTION_POLL_INTERVAL_S_ENV_VAR = "RAY_TRAIN_PREEMPTION_POLL_INTERVAL_S"
|
|
DEFAULT_PREEMPTION_POLL_INTERVAL_S: float = 5.0
|
|
|
|
# Environment variable to enable the print function patching.
|
|
ENABLE_PRINT_PATCH_ENV_VAR = "RAY_TRAIN_ENABLE_PRINT_PATCH"
|
|
DEFAULT_ENABLE_PRINT_PATCH = True
|
|
|
|
# V2 feature flag.
|
|
V2_ENABLED_ENV_VAR = "RAY_TRAIN_V2_ENABLED"
|
|
|
|
# Environment variables to enable/disable controller/worker structured logging.
|
|
ENABLE_CONTROLLER_STRUCTURED_LOGGING_ENV_VAR = (
|
|
"RAY_TRAIN_ENABLE_CONTROLLER_STRUCTURED_LOGGING"
|
|
)
|
|
ENABLE_WORKER_STRUCTURED_LOGGING_ENV_VAR = "RAY_TRAIN_ENABLE_WORKER_STRUCTURED_LOGGING"
|
|
DEFAULT_ENABLE_CONTROLLER_LOGGING = True
|
|
DEFAULT_ENABLE_WORKER_LOGGING = True
|
|
|
|
# Environment variables to configure reconciliation interval for Train state actor.
|
|
# This determines how many seconds the state actor will wait between
|
|
# polling the controller for its status.
|
|
ENABLE_STATE_ACTOR_RECONCILIATION_ENV_VAR = (
|
|
"RAY_TRAIN_ENABLE_STATE_ACTOR_RECONCILIATION"
|
|
)
|
|
DEFAULT_ENABLE_STATE_ACTOR_RECONCILIATION = True
|
|
STATE_ACTOR_RECONCILIATION_INTERVAL_S_ENV_VAR = (
|
|
"RAY_TRAIN_STATE_ACTOR_RECONCILIATION_INTERVAL_S"
|
|
)
|
|
DEFAULT_STATE_ACTOR_RECONCILIATION_INTERVAL_S: float = 30.0
|
|
# TODO: `ray.util.state.api.get_actor` typically takes 10-50ms but can take longer
|
|
# when there is high load on the cluster.
|
|
GET_ACTOR_TIMEOUT_S: int = 10
|
|
# GET_ACTOR_TIMEOUT_S * CONTROLLERS_TO_POLL_PER_ITERATION should be
|
|
# way less than STATE_ACTOR_RECONCILIATION_INTERVAL_S to give the state actor
|
|
# time to update live train run state.
|
|
CONTROLLERS_TO_POLL_PER_ITERATION: int = 1
|
|
|
|
# Environment variable for Train execution callbacks
|
|
RAY_TRAIN_CALLBACKS_ENV_VAR = "RAY_TRAIN_CALLBACKS"
|
|
|
|
# Ray Train does not warn by default when using blocking ray.get inside async actor.
|
|
DEFAULT_RAY_WARN_BLOCKING_GET_INSIDE_ASYNC_VALUE = "0"
|
|
|
|
# torchft lighthouse address
|
|
TORCHFT_LIGHTHOUSE_ADDR_ENV_VAR = "TORCHFT_LIGHTHOUSE"
|
|
|
|
# Environment variables to propagate from the driver to the controller,
|
|
# and then from the controller to the workers.
|
|
ENV_VARS_TO_PROPAGATE = {
|
|
V2_ENABLED_ENV_VAR,
|
|
HEALTH_CHECK_INTERVAL_S_ENV_VAR,
|
|
WORKER_HEALTH_CHECK_TIMEOUT_S_ENV_VAR,
|
|
WORKER_GROUP_START_TIMEOUT_S_ENV_VAR,
|
|
COLLECTIVE_TIMEOUT_S_ENV_VAR,
|
|
COLLECTIVE_WARN_INTERVAL_S_ENV_VAR,
|
|
CHECKPOINT_UPLOAD_WARN_INTERVAL_S_ENV_VAR,
|
|
ENABLE_PRINT_PATCH_ENV_VAR,
|
|
ENABLE_CONTROLLER_STRUCTURED_LOGGING_ENV_VAR,
|
|
ENABLE_WORKER_STRUCTURED_LOGGING_ENV_VAR,
|
|
ENABLE_STATE_ACTOR_RECONCILIATION_ENV_VAR,
|
|
STATE_ACTOR_RECONCILIATION_INTERVAL_S_ENV_VAR,
|
|
RAY_WARN_BLOCKING_GET_INSIDE_ASYNC_ENV_VAR,
|
|
TORCHFT_LIGHTHOUSE_ADDR_ENV_VAR,
|
|
ENABLE_PREEMPTION_WATCHER_ENV_VAR,
|
|
PREEMPTION_POLL_INTERVAL_S_ENV_VAR,
|
|
}
|
|
|
|
|
|
# ------------------------------------------------------------
|
|
# Environment variables used in the driver script only.
|
|
# ------------------------------------------------------------
|
|
|
|
# The environment variable to enable the Ray Train Metrics.
|
|
METRICS_ENABLED_ENV_VAR = "RAY_TRAIN_METRICS_ENABLED"
|
|
|
|
|
|
def is_v2_enabled() -> bool:
|
|
return env_bool(V2_ENABLED_ENV_VAR, True)
|
|
|
|
|
|
def get_env_vars_to_propagate() -> Dict[str, str]:
|
|
"""Returns a dictionary of environment variables that should be propagated
|
|
from the driver to the controller, and then from the controller
|
|
to each training worker.
|
|
|
|
This way, users only need to set environment variables in one place
|
|
when launching the script instead of needing to manually set a runtime environment.
|
|
"""
|
|
env_vars = {}
|
|
for env_var in ENV_VARS_TO_PROPAGATE:
|
|
if env_set_by_user(env_var):
|
|
env_vars[env_var] = os.environ[env_var]
|
|
return env_vars
|