finish

Signed-off-by: OneSizeFitsQuorum <txypotato@gmail.com>

Author: OneSizeFitsQuorum

python/ray/util/client/worker.py

@@ -24,6 +24,7 @@
 from ray._private.ray_constants import DEFAULT_CLIENT_RECONNECT_GRACE_PERIOD
 from ray._private.runtime_env.py_modules import upload_py_modules_if_needed
 from ray._private.runtime_env.working_dir import upload_working_dir_if_needed
+from ray._private.ray_constants import env_integer, env_float
 
 # Use cloudpickle's version of pickle for UnpicklingError
 from ray.cloudpickle.compat import pickle
@@ -52,13 +53,12 @@
 
 logger = logging.getLogger(__name__)
 
-INITIAL_TIMEOUT_SEC = 5
-MAX_TIMEOUT_SEC = 30
-
+INITIAL_TIMEOUT_SEC = env_integer("RAY_CLIENT_INITIAL_CONNECTION_TIMEOUT_S", 5)
+MAX_TIMEOUT_SEC =  env_integer("RAY_CLIENT_MAX_CONNECTION_TIMEOUT_S", 30)
 # The max amount of time an operation can run blocking in the server. This
 # allows for Ctrl-C of the client to work without explicitly cancelling server
 # operations.
-MAX_BLOCKING_OPERATION_TIME_S: float = 2.0
+MAX_BLOCKING_OPERATION_TIME_S: float = env_float("RAY_CLIENT_MAX_BLOCKING_OPERATION_TIME_S", 2.0)
 
 # If the total size (bytes) of all outbound messages to schedule tasks since
 # the connection began exceeds this value, a warning should be raised
@@ -416,19 +416,14 @@ def get(self, vals, *, timeout: Optional[float] = None) -> Any:
         else:
             deadline = time.monotonic() + timeout
 
-        max_blocking_operation_time = MAX_BLOCKING_OPERATION_TIME_S
-        if "RAY_CLIENT_MAX_BLOCKING_OPERATION_TIME_S" in os.environ:
-            max_blocking_operation_time = float(
-                os.environ["RAY_CLIENT_MAX_BLOCKING_OPERATION_TIME_S"]
-            )
         while True:
             if deadline:
                 op_timeout = min(
-                    max_blocking_operation_time,
+                    MAX_BLOCKING_OPERATION_TIME_S,
                     max(deadline - time.monotonic(), 0.001),
                 )
             else:
-                op_timeout = max_blocking_operation_time
+                op_timeout = MAX_BLOCKING_OPERATION_TIME_S
             try:
                 res = self._get(to_get, op_timeout)
                 break