wehub-resource-sync
69 lines
2.3 KiB
Python
69 lines
2.3 KiB
Python
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]
|