Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[train v2+tune] Add an environment variable to disable running the TrainController as an actor #49522

Merged
merged 4 commits into from
Jan 4, 2025
Merged

[train v2+tune] Add an environment variable to disable running the TrainController as an actor #49522

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
6935ddc
add an env var for running the controller as an actor
justinvyu Jan 1, 2025
35f12f7
fix lint
justinvyu Jan 1, 2025
c50f16d
Merge branch 'master' of https://github.com/ray-project/ray into tune…
justinvyu Jan 4, 2025
af89b51
read env var on fit
justinvyu Jan 4, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions python/ray/train/v2/_internal/constants.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,10 @@
ENABLE_PRINT_PATCH_ENV_VAR = "RAY_TRAIN_ENABLE_PRINT_PATCH"
DEFAULT_ENABLE_PRINT_PATCH = "1"

# Whether or not to run the controller as an actor.
RUN_CONTROLLER_AS_ACTOR_ENV_VAR = "RAY_TRAIN_RUN_CONTROLLER_AS_ACTOR"
DEFAULT_RUN_CONTROLLER_AS_ACTOR = "1"

# V2 feature flag.
V2_ENABLED_ENV_VAR = "RAY_TRAIN_V2_ENABLED"

Expand Down
37 changes: 25 additions & 12 deletions python/ray/train/v2/api/data_parallel_trainer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,9 @@
)
from ray.train.v2._internal.constants import (
_UNSUPPORTED,
DEFAULT_RUN_CONTROLLER_AS_ACTOR,
METRICS_ENABLED_ENV_VAR,
RUN_CONTROLLER_AS_ACTOR_ENV_VAR,
get_env_vars_to_propagate,
)
from ray.train.v2._internal.execution.context import TrainRunContext
Expand Down Expand Up @@ -195,26 +197,15 @@ def fit(self) -> Result:

# TODO: Add support for user-defined callbacks

# By default, attach the controller to the node running the driver script.
controller_actor_cls = ray.remote(
num_cpus=0,
scheduling_strategy=NodeAffinitySchedulingStrategy(
node_id=ray.get_runtime_context().get_node_id(), soft=False
),
runtime_env={"env_vars": get_env_vars_to_propagate()},
)(TrainController)

controller = controller_actor_cls.remote(
result = self._initialize_and_run_controller(
train_fn=train_fn,
scaling_policy=create_scaling_policy(self.scaling_config),
failure_policy=DefaultFailurePolicy(self.run_config.failure_config),
train_run_context=self.train_run_context,
callbacks=callbacks,
resume_from_checkpoint=self.resume_from_checkpoint,
)
ray.get(controller.run.remote())

result = ray.get(controller.get_result.remote())
if result.error:
# NOTE: If the training run errored out, raise an error back to the
# user's driver script.
Expand All @@ -225,6 +216,28 @@ def fit(self) -> Result:

return result

def _initialize_and_run_controller(self, **controller_init_kwargs) -> Result:
run_controller_as_actor = env_bool(
RUN_CONTROLLER_AS_ACTOR_ENV_VAR, DEFAULT_RUN_CONTROLLER_AS_ACTOR
)
if run_controller_as_actor:
# Attach the controller to the node running the driver script.
controller_actor_cls = ray.remote(
num_cpus=0,
scheduling_strategy=NodeAffinitySchedulingStrategy(
node_id=ray.get_runtime_context().get_node_id(), soft=False
),
runtime_env={"env_vars": get_env_vars_to_propagate()},
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Unrelated to this PR, but would this merge with environment variables set by a user-defined runtime env?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't know actually, I'll test it out

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This will update a user-defined runtime env, which is as expected.

import ray
import os

ray.init(runtime_env={"env_vars": {"RAY_TRAIN_V2_ENABLED": "0"}})


import ray.train
from ray.train.v2.api.data_parallel_trainer import DataParallelTrainer

os.environ["RAY_TRAIN_V2_ENABLED"] = "1"

trainer = DataParallelTrainer(
    lambda: print(os.getenv("RAY_TRAIN_V2_ENABLED")),
    scaling_config=ray.train.ScalingConfig(num_workers=2)
)
trainer.fit()
(RayTrainWorker pid=69742) 1
(RayTrainWorker pid=69742) 
(RayTrainWorker pid=69743) 1
(RayTrainWorker pid=69743)

)(TrainController)

controller = controller_actor_cls.remote(**controller_init_kwargs)
ray.get(controller.run.remote())
return ray.get(controller.get_result.remote())
else:
controller = TrainController(**controller_init_kwargs)
controller.run()
return controller.get_result()

@classmethod
def restore(cls, *args, **kwargs):
raise DeprecationWarning(TRAINER_RESTORE_DEPRECATION_WARNING)
Expand Down
3 changes: 3 additions & 0 deletions python/ray/train/v2/tests/test_data_parallel_trainer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,10 +10,13 @@
from ray.train.backend import Backend
from ray.train.constants import RAY_CHDIR_TO_TRIAL_DIR, _get_ray_train_session_dir
from ray.train.tests.util import create_dict_checkpoint
from ray.train.v2._internal.constants import is_v2_enabled
from ray.train.v2._internal.exceptions import TrainingFailedError
from ray.train.v2.api.data_parallel_trainer import DataParallelTrainer
from ray.train.v2.api.result import Result

assert is_v2_enabled()


@pytest.fixture(scope="module", autouse=True)
def ray_start_4_cpus():
Expand Down
Loading