wehub-resource-sync
117 lines
3.8 KiB
Python
117 lines
3.8 KiB
Python
import threading
|
|
from typing import Any, Dict, Optional
|
|
|
|
from ray.train._internal import session
|
|
from ray.train.constants import (
|
|
V2_MIGRATION_GUIDE_MESSAGE,
|
|
_v2_migration_warnings_enabled,
|
|
)
|
|
from ray.train.context import TrainContext as TrainV1Context
|
|
from ray.train.utils import _copy_doc
|
|
from ray.tune.execution.placement_groups import PlacementGroupFactory
|
|
from ray.util.annotations import Deprecated, PublicAPI
|
|
|
|
# The context singleton on this process.
|
|
_tune_context: Optional["TuneContext"] = None
|
|
_tune_context_lock = threading.Lock()
|
|
|
|
|
|
_TRAIN_SPECIFIC_CONTEXT_DEPRECATION_MESSAGE = (
|
|
"`{}` is deprecated for Ray Tune because there is no concept of worker ranks "
|
|
"for Ray Tune, so these methods only make sense to use in the context of "
|
|
f"a Ray Train worker. {V2_MIGRATION_GUIDE_MESSAGE}"
|
|
)
|
|
|
|
|
|
@PublicAPI(stability="beta")
|
|
class TuneContext(TrainV1Context):
|
|
"""Context to access metadata within Ray Tune functions."""
|
|
|
|
# NOTE: These methods are deprecated on the TrainContext, but are still
|
|
# available on the TuneContext. Re-defining them here to avoid the
|
|
# deprecation warnings.
|
|
|
|
@_copy_doc(session.get_trial_name)
|
|
def get_trial_name(self) -> str:
|
|
return session.get_trial_name()
|
|
|
|
@_copy_doc(session.get_trial_id)
|
|
def get_trial_id(self) -> str:
|
|
return session.get_trial_id()
|
|
|
|
@_copy_doc(session.get_trial_resources)
|
|
def get_trial_resources(self) -> PlacementGroupFactory:
|
|
return session.get_trial_resources()
|
|
|
|
@_copy_doc(session.get_trial_dir)
|
|
def get_trial_dir(self) -> str:
|
|
return session.get_trial_dir()
|
|
|
|
# Deprecated APIs
|
|
|
|
@Deprecated
|
|
def get_metadata(self) -> Dict[str, Any]:
|
|
raise DeprecationWarning(
|
|
"`get_metadata` is deprecated for Ray Tune, as it has never been usable."
|
|
)
|
|
|
|
@Deprecated(
|
|
message=_TRAIN_SPECIFIC_CONTEXT_DEPRECATION_MESSAGE.format("get_world_size"),
|
|
warning=_v2_migration_warnings_enabled(),
|
|
)
|
|
@_copy_doc(TrainV1Context.get_world_size)
|
|
def get_world_size(self) -> int:
|
|
return session.get_world_size()
|
|
|
|
@Deprecated(
|
|
message=_TRAIN_SPECIFIC_CONTEXT_DEPRECATION_MESSAGE.format("get_world_rank"),
|
|
warning=_v2_migration_warnings_enabled(),
|
|
)
|
|
@_copy_doc(TrainV1Context.get_world_rank)
|
|
def get_world_rank(self) -> int:
|
|
return session.get_world_rank()
|
|
|
|
@Deprecated(
|
|
message=_TRAIN_SPECIFIC_CONTEXT_DEPRECATION_MESSAGE.format("get_local_rank"),
|
|
warning=_v2_migration_warnings_enabled(),
|
|
)
|
|
@_copy_doc(TrainV1Context.get_local_rank)
|
|
def get_local_rank(self) -> int:
|
|
return session.get_local_rank()
|
|
|
|
@Deprecated(
|
|
message=_TRAIN_SPECIFIC_CONTEXT_DEPRECATION_MESSAGE.format(
|
|
"get_local_world_size"
|
|
),
|
|
warning=_v2_migration_warnings_enabled(),
|
|
)
|
|
@_copy_doc(TrainV1Context.get_local_world_size)
|
|
def get_local_world_size(self) -> int:
|
|
return session.get_local_world_size()
|
|
|
|
@Deprecated(
|
|
message=_TRAIN_SPECIFIC_CONTEXT_DEPRECATION_MESSAGE.format("get_node_rank"),
|
|
warning=_v2_migration_warnings_enabled(),
|
|
)
|
|
@_copy_doc(TrainV1Context.get_node_rank)
|
|
def get_node_rank(self) -> int:
|
|
return session.get_node_rank()
|
|
|
|
|
|
@PublicAPI(stability="beta")
|
|
def get_context() -> TuneContext:
|
|
"""Get or create a singleton Ray Tune context.
|
|
|
|
The context is only available in a tune function passed to the `ray.tune.Tuner`.
|
|
|
|
See the :class:`~ray.tune.TuneContext` API reference to see available methods.
|
|
"""
|
|
global _tune_context
|
|
|
|
with _tune_context_lock:
|
|
if _tune_context is None:
|
|
# TODO(justinvyu): This default should be a dummy context
|
|
# that is only used for testing / running outside of Tune.
|
|
_tune_context = TuneContext()
|
|
return _tune_context
|