chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,261 @@
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Any, Dict
|
||||
|
||||
from ray.train.v2._internal.execution.context import (
|
||||
get_train_context as get_internal_train_context,
|
||||
)
|
||||
from ray.util.annotations import Deprecated, DeveloperAPI, PublicAPI
|
||||
|
||||
|
||||
@PublicAPI(stability="stable")
|
||||
class TrainContext(ABC):
|
||||
"""Abstract interface for training context."""
|
||||
|
||||
@Deprecated
|
||||
def get_metadata(self) -> Dict[str, Any]:
|
||||
"""[Deprecated] User metadata dict passed to the Trainer constructor."""
|
||||
from ray.train.context import _GET_METADATA_DEPRECATION_MESSAGE
|
||||
|
||||
raise DeprecationWarning(_GET_METADATA_DEPRECATION_MESSAGE)
|
||||
|
||||
@Deprecated
|
||||
def get_trial_name(self) -> str:
|
||||
"""[Deprecated] Trial name for the corresponding trial."""
|
||||
from ray.train.context import _TUNE_SPECIFIC_CONTEXT_DEPRECATION_MESSAGE
|
||||
|
||||
raise DeprecationWarning(
|
||||
_TUNE_SPECIFIC_CONTEXT_DEPRECATION_MESSAGE.format("get_trial_name")
|
||||
)
|
||||
|
||||
@Deprecated
|
||||
def get_trial_id(self) -> str:
|
||||
"""[Deprecated] Trial id for the corresponding trial."""
|
||||
from ray.train.context import _TUNE_SPECIFIC_CONTEXT_DEPRECATION_MESSAGE
|
||||
|
||||
raise DeprecationWarning(
|
||||
_TUNE_SPECIFIC_CONTEXT_DEPRECATION_MESSAGE.format("get_trial_id")
|
||||
)
|
||||
|
||||
@Deprecated
|
||||
def get_trial_resources(self):
|
||||
"""[Deprecated] Trial resources for the corresponding trial."""
|
||||
from ray.train.context import _TUNE_SPECIFIC_CONTEXT_DEPRECATION_MESSAGE
|
||||
|
||||
raise DeprecationWarning(
|
||||
_TUNE_SPECIFIC_CONTEXT_DEPRECATION_MESSAGE.format("get_trial_resources")
|
||||
)
|
||||
|
||||
@Deprecated
|
||||
def get_trial_dir(self) -> str:
|
||||
"""[Deprecated] Log directory corresponding to the trial directory for a Tune session.
|
||||
This is deprecated for Ray Train and should no longer be called in Ray Train workers.
|
||||
|
||||
If this directory is needed, please pass it into the `train_loop_config` directly.
|
||||
"""
|
||||
from ray.train.context import _TUNE_SPECIFIC_CONTEXT_DEPRECATION_MESSAGE
|
||||
|
||||
raise DeprecationWarning(
|
||||
_TUNE_SPECIFIC_CONTEXT_DEPRECATION_MESSAGE.format("get_trial_dir")
|
||||
)
|
||||
|
||||
@abstractmethod
|
||||
def get_experiment_name(self) -> str:
|
||||
"""Experiment name for the corresponding trial."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_world_size(self) -> int:
|
||||
"""Get the current world size (i.e. total number of workers) for this run.
|
||||
|
||||
.. testcode::
|
||||
|
||||
import ray.train
|
||||
from ray.train.torch import TorchTrainer
|
||||
|
||||
NUM_WORKERS = 2
|
||||
|
||||
def train_fn_per_worker(config):
|
||||
assert ray.train.get_context().get_world_size() == NUM_WORKERS
|
||||
|
||||
trainer = TorchTrainer(
|
||||
train_fn_per_worker,
|
||||
scaling_config=ray.train.ScalingConfig(num_workers=NUM_WORKERS),
|
||||
)
|
||||
trainer.fit()
|
||||
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_world_rank(self) -> int:
|
||||
"""Get the world rank of this worker.
|
||||
|
||||
.. testcode::
|
||||
|
||||
import ray.train
|
||||
from ray.train.torch import TorchTrainer
|
||||
|
||||
def train_fn_per_worker(config):
|
||||
if ray.train.get_context().get_world_rank() == 0:
|
||||
print("Worker 0")
|
||||
|
||||
trainer = TorchTrainer(
|
||||
train_fn_per_worker,
|
||||
scaling_config=ray.train.ScalingConfig(num_workers=2),
|
||||
)
|
||||
trainer.fit()
|
||||
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_local_rank(self) -> int:
|
||||
"""Get the local rank of this worker (rank of the worker on its node).
|
||||
|
||||
.. testcode::
|
||||
|
||||
import ray.train
|
||||
from ray.train.torch import TorchTrainer
|
||||
|
||||
def train_fn_per_worker(config):
|
||||
if ray.train.get_context().get_local_rank() == 0:
|
||||
print("Local rank 0 worker")
|
||||
|
||||
trainer = TorchTrainer(
|
||||
train_fn_per_worker,
|
||||
scaling_config=ray.train.ScalingConfig(num_workers=2),
|
||||
)
|
||||
trainer.fit()
|
||||
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_local_world_size(self) -> int:
|
||||
"""Get the local world size of this node (i.e. number of workers on this node).
|
||||
|
||||
Example:
|
||||
|
||||
.. testcode::
|
||||
|
||||
import ray.train
|
||||
from ray.train.torch import TorchTrainer
|
||||
|
||||
def train_fn_per_worker():
|
||||
print(ray.train.get_context().get_local_world_size())
|
||||
|
||||
trainer = TorchTrainer(
|
||||
train_fn_per_worker,
|
||||
scaling_config=ray.train.ScalingConfig(num_workers=2),
|
||||
)
|
||||
trainer.fit()
|
||||
|
||||
Returns:
|
||||
The number of workers running on this node.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_node_rank(self) -> int:
|
||||
"""Get the rank of this node.
|
||||
|
||||
Example:
|
||||
|
||||
.. testcode::
|
||||
|
||||
import ray.train
|
||||
from ray.train.torch import TorchTrainer
|
||||
|
||||
def train_fn_per_worker():
|
||||
print(ray.train.get_context().get_node_rank())
|
||||
|
||||
trainer = TorchTrainer(
|
||||
train_fn_per_worker,
|
||||
scaling_config=ray.train.ScalingConfig(num_workers=1),
|
||||
)
|
||||
trainer.fit()
|
||||
|
||||
Returns:
|
||||
The rank of this node among the nodes participating in training.
|
||||
"""
|
||||
pass
|
||||
|
||||
@DeveloperAPI
|
||||
@abstractmethod
|
||||
def get_storage(self):
|
||||
"""Returns the :class:`~ray.train._internal.storage.StorageContext` storage
|
||||
context which gives advanced access to the filesystem and paths
|
||||
configured through `RunConfig`.
|
||||
|
||||
NOTE: This is a DeveloperAPI, and the `StorageContext` interface may change
|
||||
without notice between minor versions.
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
class DistributedTrainContext(TrainContext):
|
||||
"""Implementation of TrainContext for distributed mode."""
|
||||
|
||||
def get_experiment_name(self) -> str:
|
||||
return get_internal_train_context().get_experiment_name()
|
||||
|
||||
def get_world_size(self) -> int:
|
||||
return get_internal_train_context().get_world_size()
|
||||
|
||||
def get_world_rank(self) -> int:
|
||||
return get_internal_train_context().get_world_rank()
|
||||
|
||||
def get_local_rank(self) -> int:
|
||||
return get_internal_train_context().get_local_rank()
|
||||
|
||||
def get_local_world_size(self) -> int:
|
||||
return get_internal_train_context().get_local_world_size()
|
||||
|
||||
def get_node_rank(self) -> int:
|
||||
return get_internal_train_context().get_node_rank()
|
||||
|
||||
def get_storage(self):
|
||||
return get_internal_train_context().get_storage()
|
||||
|
||||
|
||||
@DeveloperAPI
|
||||
class LocalTrainContext(TrainContext):
|
||||
"""Implementation of TrainContext for local mode."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
experiment_name: str,
|
||||
world_size: int = 1,
|
||||
world_rank: int = 0,
|
||||
local_rank: int = 0,
|
||||
local_world_size: int = 1,
|
||||
node_rank: int = 0,
|
||||
):
|
||||
self.experiment_name = experiment_name
|
||||
self.world_size = world_size
|
||||
self.world_rank = world_rank
|
||||
self.local_rank = local_rank
|
||||
self.local_world_size = local_world_size
|
||||
self.node_rank = node_rank
|
||||
|
||||
def get_experiment_name(self) -> str:
|
||||
return self.experiment_name
|
||||
|
||||
def get_world_size(self) -> int:
|
||||
return self.world_size
|
||||
|
||||
def get_world_rank(self) -> int:
|
||||
return self.world_rank
|
||||
|
||||
def get_local_rank(self) -> int:
|
||||
return self.local_rank
|
||||
|
||||
def get_local_world_size(self) -> int:
|
||||
return self.local_world_size
|
||||
|
||||
def get_node_rank(self) -> int:
|
||||
return self.node_rank
|
||||
|
||||
def get_storage(self):
|
||||
raise NotImplementedError("Local storage context not yet implemented. ")
|
||||
Reference in New Issue
Block a user