[Tune] Add option to not override the working directory (ray-project#29258)

See ray-project#29128 for more context on the problem.

This PR does the following:
1. **Fix `TUNE_ORIG_WORKING_DIR`** to pull the correct current working directory within the worker process when the Tune trial logdir is being created. This PR deprecates this environment variable as it's confusing to get Tune metadata from sources other than `session`.
2. **Introduce a `chdir_to_trial_dir` flag** in the TuneConfig that defaults to `True`, which configures whether or not Tune should change the working directory of each worker to its corresponding trial directory.
    - If this flag is set to False, the user may still want to access the Tune trial directory. This can be done with a **newly added `session.get_trial_dir()` API.**
3. Make the `TUNE_ORIG_WORKING_DIR` deprecation, `chdir_to_trial_dir` flag, and `session.get_trial_dir()` more visible in the documentation with an **example in the Tune FAQ.**

Signed-off-by: Justin Yu <justinvyu@berkeley.edu>
Signed-off-by: Weichen Xu <weichen.xu@databricks.com>

Author: justinvyu

doc/source/tune/api_docs/trainable.rst

@@ -7,7 +7,7 @@
 Training (tune.Trainable, session.report)
 ==========================================
 
-Training can be done with either a **Class API** (``tune.Trainable``) or **function API** (``session.report``).
+Training can be done with either a **Function API** (``session.report``) or **Class API** (``tune.Trainable``).
 
 For the sake of example, let's maximize this objective function:
 
@@ -327,6 +327,9 @@ session (Function API)
 .. autofunction:: ray.air.session.get_trial_resources
     :noindex:
 
+.. autofunction:: ray.air.session.get_trial_dir
+    :noindex:
+
 tune.Trainable (Class API)
 --------------------------
 

doc/source/tune/doc_code/faq.py

@@ -411,3 +411,28 @@ def wait(self):
     ),
 )
 # __grid_search_2_end__
+
+if not MOCK:
+    import os
+    from pathlib import Path
+
+    # __no_chdir_start__
+    def train_func(config):
+        # Read from relative paths
+        print(open("./read.txt").read())
+
+        # The working directory shouldn't have changed from the original
+        # NOTE: The `TUNE_ORIG_WORKING_DIR` environment variable is deprecated.
+        assert os.getcwd() == os.environ["TUNE_ORIG_WORKING_DIR"]
+
+        # Write to the Tune trial directory, not the shared working dir
+        tune_trial_dir = Path(session.get_trial_dir())
+        with open(tune_trial_dir / "write.txt", "w") as f:
+            f.write("trial saved artifact")
+
+    tuner = tune.Tuner(
+        train_func,
+        tune_config=tune.TuneConfig(..., chdir_to_trial_dir=False),
+    )
+    tuner.fit()
+    # __no_chdir_end__

doc/source/tune/faq.rst

@@ -733,3 +733,35 @@ If `grid_search` is provided as an argument, the grid will be repeated ``num_sam
 Note that search spaces may not be interoperable across different search algorithms.
 For example, for many search algorithms, you will not be able to use a ``grid_search`` or ``sample_from`` parameters.
 Read about this in the :ref:`Search Space API <tune-search-space>` page.
+
+.. _tune-working-dir:
+
+How do I access relative filepaths in my Tune training function?
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Let's say you launch a Tune experiment from ``~/code/my_script.py``. By default, Tune
+changes the working directory of each worker from ``~/code`` to its corresponding trial
+directory (e.g. ``~/ray_results/exp_name/trial_0000x``). This default
+guarantees separate working directories for each worker process, avoiding conflicts when
+saving trial-specific outputs.
+
+You can configure this by setting `chdir_to_trial_dir=False` in `tune.TuneConfig`.
+This explicitly tells Tune to not change the working directory
+to the trial directory, giving access to paths relative to the original working directory.
+One caveat is that the working directory is now shared between workers, so the
+:meth:`session.get_trial_dir() <ray.air.session.get_trial_dir>`
+API should be used to get the path for saving trial-specific outputs.
+
+.. literalinclude:: doc_code/faq.py
+    :dedent:
+    :emphasize-lines: 3, 10, 11, 12, 16
+    :language: python
+    :start-after: __no_chdir_start__
+    :end-before: __no_chdir_end__
+
+.. warning::
+
+    The `TUNE_ORIG_WORKING_DIR` environment variable was the original workaround for
+    accessing paths relative to the original working directory. This environment
+    variable is deprecated, and the `chdir_to_trial_dir` flag described above should be
+    used instead.

python/ray/air/_internal/session.py

@@ -74,6 +74,11 @@ def trial_resources(self) -> "PlacementGroupFactory":
         """Trial resources for the corresponding trial."""
         raise NotImplementedError
 
+    @property
+    def trial_dir(self) -> str:
+        """Trial-level log directory for the corresponding trial."""
+        raise NotImplementedError
+
 
 def _get_session() -> Optional[Session]:
     from ray.train._internal.session import _session_v2 as train_session

python/ray/air/session.py

@@ -68,44 +68,43 @@ def get_checkpoint() -> Optional[Checkpoint]:
         Checkpoint object if the session is currently being resumed.
             Otherwise, return None.
 
-    Example:
-        .. code-block: python
-
-            ######## Using it in the *per worker* train loop (TrainSession) ######
-            from ray.air import session
-            from ray.air.checkpoint import Checkpoint
-            from ray.air.config import ScalingConfig
-            def train_func():
-                ckpt = session.get_checkpoint()
-                if ckpt:
-                    with ckpt.as_directory() as loaded_checkpoint_dir:
-                        import tensorflow as tf
-
-                        model = tf.keras.models.load_model(loaded_checkpoint_dir)
-                else:
-                    model = build_model()
+    .. code-block:: python
 
-                model.save("my_model", overwrite=True)
-                session.report(
-                    metrics={"iter": 1},
-                    checkpoint=Checkpoint.from_directory("my_model")
-                )
+        ######## Using it in the *per worker* train loop (TrainSession) ######
+        from ray.air import session
+        from ray.air.checkpoint import Checkpoint
+        from ray.air.config import ScalingConfig
+        def train_func():
+            ckpt = session.get_checkpoint()
+            if ckpt:
+                with ckpt.as_directory() as loaded_checkpoint_dir:
+                    import tensorflow as tf
+
+                    model = tf.keras.models.load_model(loaded_checkpoint_dir)
+            else:
+                model = build_model()
 
-            scaling_config = ScalingConfig(num_workers=2)
-            trainer = TensorflowTrainer(
-                train_loop_per_worker=train_func, scaling_config=scaling_config
+            model.save("my_model", overwrite=True)
+            session.report(
+                metrics={"iter": 1},
+                checkpoint=Checkpoint.from_directory("my_model")
             )
-            result = trainer.fit()
 
-            # trainer2 will pick up from the checkpoint saved by trainer1.
-            trainer2 = TensorflowTrainer(
-                train_loop_per_worker=train_func,
-                scaling_config=scaling_config,
-                # this is ultimately what is accessed through
-                # ``Session.get_checkpoint()``
-                resume_from_checkpoint=result.checkpoint,
-            )
-            result2 = trainer2.fit()
+        scaling_config = ScalingConfig(num_workers=2)
+        trainer = TensorflowTrainer(
+            train_loop_per_worker=train_func, scaling_config=scaling_config
+        )
+        result = trainer.fit()
+
+        # trainer2 will pick up from the checkpoint saved by trainer1.
+        trainer2 = TensorflowTrainer(
+            train_loop_per_worker=train_func,
+            scaling_config=scaling_config,
+            # this is ultimately what is accessed through
+            # ``Session.get_checkpoint()``
+            resume_from_checkpoint=result.checkpoint,
+        )
+        result2 = trainer2.fit()
     """
 
     return _get_session().loaded_checkpoint
@@ -126,6 +125,27 @@ def get_trial_resources() -> "PlacementGroupFactory":
     return _get_session().trial_resources
 
 
+def get_trial_dir() -> str:
+    """Log directory corresponding to the trial directory for a Tune session.
+    If calling from a Train session, this will give the trial directory of its parent
+    Tune session.
+
+    .. code-block:: python
+
+        from ray import tune
+        from ray.air import session
+
+        def train_func():
+            # Example:
+            # >>> session.get_trial_dir()
+            # ~/ray_results/<exp-name>/<trial-dir>
+
+        tuner = tune.Tuner(train_func)
+        tuner.fit()
+    """
+    return _get_session().trial_dir
+
+
 def get_world_size() -> int:
     """Get the current world size (i.e. total number of workers) for this run.
 

python/ray/train/data_parallel_trainer.py

@@ -1,6 +1,5 @@
 import inspect
 import logging
-import os
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Tuple, Type, Union
 from tabulate import tabulate
@@ -338,7 +337,7 @@ def training_loop(self) -> None:
             name=session.get_trial_name(),
             id=session.get_trial_id(),
             resources=session.get_trial_resources(),
-            logdir=os.getcwd(),
+            logdir=session.get_trial_dir(),
         )
 
         backend_executor = self._backend_executor_cls(

python/ray/train/session.py

@@ -45,6 +45,10 @@ def trial_id(self) -> str:
     def trial_resources(self) -> "PlacementGroupFactory":
         return self._session.trial_info.resources
 
+    @property
+    def trial_dir(self) -> str:
+        return self._session.trial_info.logdir
+
     @property
     def world_size(self) -> int:
         return self._session.world_size

python/ray/tune/execution/ray_trial_executor.py

@@ -72,7 +72,6 @@ def __init__(self):
     def get(self, trainable_cls):
         """Gets the wrapped trainable_cls, otherwise calls ray.remote."""
         env_vars = DEFAULT_ENV_VARS.copy()
-        env_vars["TUNE_ORIG_WORKING_DIR"] = os.getcwd()
 
         runtime_env = {"env_vars": env_vars}
         if trainable_cls not in self._cache:
@@ -145,11 +144,17 @@ def is_empty(self):
         return len(self._future_to_insert_time) == 0
 
 
-def _noop_logger_creator(config, logdir):
-    # Set the working dir in the remote process, for user file writes
+def _noop_logger_creator(config, logdir, should_chdir: bool = True):
+    # Upon remote process setup, record the actor's original working dir before
+    # changing the working dir to the Tune logdir
+    os.environ["TUNE_ORIG_WORKING_DIR"] = os.getcwd()
+
     os.makedirs(logdir, exist_ok=True)
-    if not ray._private.worker._mode() == ray._private.worker.LOCAL_MODE:
-        os.chdir(logdir)
+    if should_chdir:
+        # Set the working dir to the trial directory in the remote process,
+        # for user file writes
+        if not ray._private.worker._mode() == ray._private.worker.LOCAL_MODE:
+            os.chdir(logdir)
     return NoopLogger(config, logdir)
 
 
@@ -205,6 +210,7 @@ def __init__(
         reuse_actors: bool = False,
         result_buffer_length: Optional[int] = None,
         refresh_period: Optional[float] = None,
+        chdir_to_trial_dir: bool = False,
     ):
         self._cached_trial_state = {}
         self._trials_to_cache = set()
@@ -245,6 +251,8 @@ def __init__(
         )
         self._trainable_kwargs = {}
 
+        self._chdir_to_trial_dir = chdir_to_trial_dir
+
     def setup(
         self, max_pending_trials: int, trainable_kwargs: Optional[Dict] = None
     ) -> None:
@@ -335,7 +343,11 @@ def _setup_remote_runner(self, trial):
         trial.init_logdir()
         # We checkpoint metadata here to try mitigating logdir duplication
         self._trials_to_cache.add(trial)
-        logger_creator = partial(_noop_logger_creator, logdir=trial.logdir)
+        logger_creator = partial(
+            _noop_logger_creator,
+            logdir=trial.logdir,
+            should_chdir=self._chdir_to_trial_dir,
+        )
 
         if len(self._cached_actor_pg) > 0:
             assert self._reuse_actors

python/ray/tune/impl/tuner_internal.py

@@ -22,6 +22,7 @@
 if TYPE_CHECKING:
     from ray.train.trainer import BaseTrainer
 
+
 _TRAINABLE_PKL = "trainable.pkl"
 _TUNER_PKL = "tuner.pkl"
 _TRAINABLE_KEY = "_trainable"
@@ -382,6 +383,7 @@ def _get_tune_run_arguments(self, trainable) -> Dict[str, Any]:
             reuse_actors=self._tune_config.reuse_actors,
             max_concurrent_trials=self._tune_config.max_concurrent_trials,
             time_budget_s=self._tune_config.time_budget_s,
+            chdir_to_trial_dir=self._tune_config.chdir_to_trial_dir,
         )
 
     def _fit_internal(self, trainable, param_space) -> ExperimentAnalysis: