chore: import upstream snapshot with attribution

python/ray/tune/execution/class_cache.py

@@ -0,0 +1,68 @@
+import os
+
+import ray
+from ray.air.constants import COPY_DIRECTORY_CHECKPOINTS_INSTEAD_OF_MOVING_ENV
+from ray.train.constants import (
+    ENABLE_V2_MIGRATION_WARNINGS_ENV_VAR,
+    RAY_CHDIR_TO_TRIAL_DIR,
+)
+from ray.train.v2._internal.constants import (
+    ENV_VARS_TO_PROPAGATE as TRAIN_ENV_VARS_TO_PROPAGATE,
+)
+
+DEFAULT_ENV_VARS = {
+    # https://github.com/ray-project/ray/issues/28197
+    "PL_DISABLE_FORK": "1"
+}
+ENV_VARS_TO_PROPAGATE = (
+    {
+        COPY_DIRECTORY_CHECKPOINTS_INSTEAD_OF_MOVING_ENV,
+        RAY_CHDIR_TO_TRIAL_DIR,
+        ENABLE_V2_MIGRATION_WARNINGS_ENV_VAR,
+        "AWS_ACCESS_KEY_ID",
+        "AWS_SECRET_ACCESS_KEY",
+        "AWS_SECURITY_TOKEN",
+        "AWS_SESSION_TOKEN",
+    }
+    # Propagate the Ray Train environment variables from the driver process
+    # to the trainable process so that Tune + Train v2 can be used together.
+    | TRAIN_ENV_VARS_TO_PROPAGATE
+)
+
+
+class _ActorClassCache:
+    """Caches actor classes.
+
+    ray.remote is a registration call. It sends the serialized object to the
+    key value store (redis), and will be fetched at an arbitrary worker
+    later. Registration does not use any Ray scheduling resources.
+
+    Later, class.remote() actually creates the remote actor. The
+    actor will be instantiated on some arbitrary machine,
+    according to the underlying Ray scheduler.
+
+    Without this cache, you would register the same serialized object
+    over and over again. Naturally, since redis doesn’t spill to disk,
+    this can easily nuke the redis instance (and basically blow up Ray).
+    This cache instead allows us to register once and only once.
+
+    Note that we assume there can be multiple trainables in the
+    system at once.
+    """
+
+    def __init__(self):
+        self._cache = {}
+
+    def get(self, trainable_cls):
+        """Gets the wrapped trainable_cls, otherwise calls ray.remote."""
+        env_vars = DEFAULT_ENV_VARS.copy()
+
+        for env_var_to_propagate in ENV_VARS_TO_PROPAGATE:
+            if env_var_to_propagate in os.environ:
+                env_vars[env_var_to_propagate] = os.environ[env_var_to_propagate]
+
+        runtime_env = {"env_vars": env_vars}
+        if trainable_cls not in self._cache:
+            remote_cls = ray.remote(runtime_env=runtime_env)(trainable_cls)
+            self._cache[trainable_cls] = remote_cls
+        return self._cache[trainable_cls]