chore: import upstream snapshot with attribution

python/ray/train/v2/_internal/constants.py

@@ -0,0 +1,159 @@
+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