Repository: huggingface/accelerate
Pull request: Fix: Defer Tracker Initialization to Prevent Premature Distributed Setup
# What does this PR do?

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

<!-- Remove if not applicable -->

Fixes #3550

This issue was caused by the `@on_main_process` decorator being applied directly to the `__init__` method in `TensorBoardTracker` as well as other Tracker classes. The `@on_main_process` decorator would immediately attempt to access PartialState. This could lead to `torch.distributed` being initialized, before `Accelerator` had fully configured the distributed environment (especially when `InitProcessGroupKwargs` was used).

## Changes
**`__init__` for Argument Storage**: Tracker __init__ methods now only store arguments and no longer perform setup or use @on_main_process.
**`start()` Method for Setup**: A new start() method, decorated with @on_main_process, now handles all initialization logic, including resource allocation and initial logging.
**Deferred Execution**: Accelerator calls this start() method during init_trackers(), ensuring the distributed environment is ready.
**Testing**: `class TrackerDeferredInitializationTest(unittest.TestCase):` verifies this deferred initialization.


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/accelerate/blob/main/CONTRIBUTING.md#submitting-a-pull-request-pr),
      Pull Request section?
- [x] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/accelerate/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/accelerate/tree/main/docs#writing-documentation---specification).
- [x] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@SunMarc 

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.

- Big modeling: @SunMarc
- Fully-Sharded Data Parallism: @SunMarc @zach-huggingface
- DeepSpeed: @SunMarc @zach-huggingface
- Command Line Interface: @SunMarc @zach-huggingface
- Documentation: @SunMarc @zach-huggingface
- Core parts of the library: @BenjaminBossan @SunMarc @zach-huggingface 
- Maintained examples: @SunMarc or  @zach-huggingface

 -->

Problem statement:
Timeout does not work when used with TensorboardTracker
### System Info

```Shell
- `Accelerate` version: 1.6.0
- Platform: Linux-5.15.0-70-generic-x86_64-with-glibc2.31
- `accelerate` bash location: /data2/houxiuquan/envs/cp311pt211/bin/accelerate
- Python version: 3.11.9
- Numpy version: 1.26.4
- PyTorch version (GPU?): 2.1.1 (True)
- PyTorch XPU available: False
- PyTorch NPU available: False
- PyTorch MLU available: False
- PyTorch SDAA available: False
- PyTorch MUSA available: False
- System RAM: 503.30 GB
- GPU type: NVIDIA GeForce RTX 3090
- `Accelerate` default config:
        Not found
```

### Information

- [ ] The official example scripts
- [x] My own modified scripts

### Tasks

- [ ] One of the scripts in the examples/ folder of Accelerate or an officially supported `no_trainer` script in the `examples` folder of the `transformers` repo (such as `run_no_trainer_glue.py`)
- [x] My own task or dataset (give details below)

### Reproduction

Here are three minimal scripts to reproduce the bug. 

1. In the following script, I set timeout to `4` seconds and stop the main process for `8` seconds. As expected, it will raise `timeout` error by watchdog when `wait_for_everyone`.

```python
import time
from datetime import timedelta

from accelerate import Accelerator, InitProcessGroupKwargs
from accelerate.utils import ProjectConfiguration
from torch import tensor


accelerator = Accelerator(
    kwargs_handlers=[InitProcessGroupKwargs(timeout=timedelta(seconds=4))],
)
accelerator.init_trackers(project_name="my_project")

accelerator.wait_for_everyone()
if accelerator.is_main_process:
    t = tensor(0).to(accelerator.device)
    time.sleep(8)
else:
    t = tensor(0).to(accelerator.device)
accelerator.wait_for_everyone()

print("All called!")
```

2. However, when adding `TensorboardTracker` into `accelerator`, `timeout` seems not effective and the script ends successfully.

```python
import time
from datetime import timedelta

from accelerate import Accelerator, InitProcessGroupKwargs
from accelerate.tracking import TensorBoardTracker
from accelerate.utils import ProjectConfiguration
from torch import tensor

accelerator = Accelerator(
    log_with=TensorBoardTracker(run_name="tf_log", logging_dir="tmp"),
    kwargs_handlers=[InitProcessGroupKwargs(timeout=timedelta(seconds=4))],
)
accelerator.init_trackers(project_name="my_project")


accelerator.wait_for_everyone()
if accelerator.is_main_process:
    t = tensor(0).to(accelerator.device)
    time.sleep(8)
else:
    t = tensor(0).to(accelerator.device)
accelerator.wait_for_everyone()

print("All called!")
```

3. If replacing `TensorboardTracker` with `"tensorboard"`, `timeout` becomes effective and the script raises `timeout` error as expected again.

```python

import time
from datetime import timedelta

from accelerate import Accelerator, InitProcessGroupKwargs
from accelerate.utils import ProjectConfiguration
from torch import tensor


accelerator = Accelerator(
    log_with="tensorboard",
    project_config=ProjectConfiguration(
        project_dir="tmp",
        total_limit=5,
        automatic_checkpoint_naming=True,
    ),
    kwargs_handlers=[InitProcessGroupKwargs(timeout=timedelta(seconds=4))],
)
accelerator.init_trackers(project_name="my_project")

accelerator.wait_for_everyone()
if accelerator.is_main_process:
    t = tensor(0).to(accelerator.device)
    time.sleep(8)
else:
    t = tensor(0).to(accelerator.device)
accelerator.wait_for_everyone()

print("All called!")
```

### Expected behavior

All the three scripts should raise `timeout` error by watchdog.

Hints from issue discussion:
Thanks for the report ! I'll leave it as a good first issue so that the community can try to solve this if I don't do it in time 
I would like to work on this

Commit message: Improve test for #3550

Patch to review:
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index 5dd76f40904..ef25736c081 100755
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -2932,6 +2932,10 @@ def init_trackers(self, project_name: str, config: dict | None = None, init_kwar
                     )
                 else:
                     self.trackers.append(tracker_init(project_name, **init_kwargs.get(str(tracker), {})))
+
+        for tracker in self.trackers:
+            tracker.start()
+
         if config is not None:
             for tracker in self.trackers:
                 tracker.store_init_configuration(config)
diff --git a/src/accelerate/test_utils/testing.py b/src/accelerate/test_utils/testing.py
index 87ce5035c3e..1ea86d16ab7 100644
--- a/src/accelerate/test_utils/testing.py
+++ b/src/accelerate/test_utils/testing.py
@@ -40,6 +40,7 @@
     is_bnb_available,
     is_clearml_available,
     is_comet_ml_available,
+    is_aim_available,
     is_cuda_available,
     is_datasets_available,
     is_deepspeed_available,
@@ -460,6 +461,13 @@ def require_comet_ml(test_case):
     return unittest.skipUnless(is_comet_ml_available(), "test requires comet_ml")(test_case)
 
 
+def require_aim(test_case):
+    """
+    Decorator marking a test that requires aim installed. These tests are skipped when aim isn't installed
+    """
+    return unittest.skipUnless(is_aim_available(), "test requires aim")(test_case)
+
+
 def require_clearml(test_case):
     """
     Decorator marking a test that requires clearml installed. These tests are skipped when clearml isn't installed
diff --git a/src/accelerate/tracking.py b/src/accelerate/tracking.py
index 765f7adbf76..8d69c9982eb 100644
--- a/src/accelerate/tracking.py
+++ b/src/accelerate/tracking.py
@@ -131,6 +131,13 @@ def __init__(self, _blank=False):
                     f"{err}"
                 )
 
+    def start(self):
+        """
+        Lazy initialization of the tracker inside Accelerator to avoid initalizing PartialState before
+        InitProcessGroupKwargs.
+        """
+        pass
+
     def store_init_configuration(self, values: dict):
         """
         Logs `values` as hyperparameters for the run. Implementations should use the experiment configuration
@@ -180,16 +187,20 @@ class TensorBoardTracker(GeneralTracker):
     name = "tensorboard"
     requires_logging_directory = True
 
-    @on_main_process
     def __init__(self, run_name: str, logging_dir: Union[str, os.PathLike], **kwargs):
+        super().__init__()
+        self.run_name = run_name
+        self.logging_dir_param = logging_dir
+        self.init_kwargs = kwargs
+
+    @on_main_process
+    def start(self):
         try:
             from torch.utils import tensorboard
         except ModuleNotFoundError:
             import tensorboardX as tensorboard
-        super().__init__()
-        self.run_name = run_name
-        self.logging_dir = os.path.join(logging_dir, run_name)
-        self.writer = tensorboard.SummaryWriter(self.logging_dir, **kwargs)
+        self.logging_dir = os.path.join(self.logging_dir_param, self.run_name)
+        self.writer = tensorboard.SummaryWriter(self.logging_dir, **self.init_kwargs)
         logger.debug(f"Initialized TensorBoard project {self.run_name} logging to {self.logging_dir}")
         logger.debug(
             "Make sure to log any initial configurations with `self.store_init_configuration` before training!"
@@ -290,14 +301,16 @@ class WandBTracker(GeneralTracker):
     requires_logging_directory = False
     main_process_only = False
 
-    @on_main_process
     def __init__(self, run_name: str, **kwargs):
         super().__init__()
         self.run_name = run_name
+        self.init_kwargs = kwargs
 
+    @on_main_process
+    def start(self):
         import wandb
 
-        self.run = wandb.init(project=self.run_name, **kwargs)
+        self.run = wandb.init(project=self.run_name, **self.init_kwargs)
         logger.debug(f"Initialized WandB project {self.run_name}")
         logger.debug(
             "Make sure to log any initial configurations with `self.store_init_configuration` before training!"
@@ -419,19 +432,21 @@ class CometMLTracker(GeneralTracker):
     name = "comet_ml"
     requires_logging_directory = False
 
-    @on_main_process
     def __init__(self, run_name: str, **kwargs):
         super().__init__()
         self.run_name = run_name
+        self.init_kwargs = kwargs
 
+    @on_main_process
+    def start(self):
         import comet_ml
 
         comet_version = version.parse(comet_ml.__version__)
         if compare_versions(comet_version, ">=", "3.41.0"):
-            self.writer = comet_ml.start(project_name=run_name, **kwargs)
+            self.writer = comet_ml.start(project_name=self.run_name, **self.init_kwargs)
         else:
             logger.info("Update `comet_ml` (>=3.41.0) for experiment reuse and offline support.")
-            self.writer = comet_ml.Experiment(project_name=run_name, **kwargs)
+            self.writer = comet_ml.Experiment(project_name=self.run_name, **self.init_kwargs)
 
         logger.debug(f"Initialized CometML project {self.run_name}")
         logger.debug(
@@ -504,13 +519,17 @@ class AimTracker(GeneralTracker):
     name = "aim"
     requires_logging_directory = True
 
-    @on_main_process
     def __init__(self, run_name: str, logging_dir: Optional[Union[str, os.PathLike]] = ".", **kwargs):
+        super().__init__()
         self.run_name = run_name
+        self.aim_repo_path = logging_dir
+        self.init_kwargs = kwargs
 
+    @on_main_process
+    def start(self):
         from aim import Run
 
-        self.writer = Run(repo=logging_dir, **kwargs)
+        self.writer = Run(repo=self.aim_repo_path, **self.init_kwargs)
         self.writer.name = self.run_name
         logger.debug(f"Initialized Aim project {self.run_name}")
         logger.debug(
@@ -620,7 +639,6 @@ class MLflowTracker(GeneralTracker):
     name = "mlflow"
     requires_logging_directory = False
 
-    @on_main_process
     def __init__(
         self,
         experiment_name: str = None,
@@ -639,30 +657,40 @@ def __init__(
 
         nested_run = os.environ.get("MLFLOW_NESTED_RUN", nested_run)
 
+        self.experiment_name = experiment_name
+        self.logging_dir = logging_dir
+        self.run_id = run_id
+        self.tags = tags
+        self.nested_run = nested_run
+        self.run_name = run_name
+        self.description = description
+
+    @on_main_process
+    def start(self):
         import mlflow
 
-        exps = mlflow.search_experiments(filter_string=f"name = '{experiment_name}'")
+        exps = mlflow.search_experiments(filter_string=f"name = '{self.experiment_name}'")
         if len(exps) > 0:
             if len(exps) > 1:
                 logger.warning("Multiple experiments with the same name found. Using first one.")
             experiment_id = exps[0].experiment_id
         else:
             experiment_id = mlflow.create_experiment(
-                name=experiment_name,
-                artifact_location=logging_dir,
-                tags=tags,
+                name=self.experiment_name,
+                artifact_location=self.logging_dir,
+                tags=self.tags,
             )
 
         self.active_run = mlflow.start_run(
-            run_id=run_id,
+            run_id=self.run_id,
             experiment_id=experiment_id,
-            run_name=run_name,
-            nested=nested_run,
-            tags=tags,
-            description=description,
+            run_name=self.run_name,
+            nested=self.nested_run,
+            tags=self.tags,
+            description=self.description,
         )
 
-        logger.debug(f"Initialized mlflow experiment {experiment_name}")
+        logger.debug(f"Initialized mlflow experiment {self.experiment_name}")
         logger.debug(
             "Make sure to log any initial configurations with `self.store_init_configuration` before training!"
         )
@@ -802,20 +830,26 @@ class ClearMLTracker(GeneralTracker):
     name = "clearml"
     requires_logging_directory = False
 
-    @on_main_process
     def __init__(self, run_name: str = None, **kwargs):
+        super().__init__()
+        self.user_provided_run_name = run_name
+        self._initialized_externally = False
+        self.init_kwargs = kwargs
+
+    @on_main_process
+    def start(self):
         from clearml import Task
 
         current_task = Task.current_task()
-        self._initialized_externally = False
         if current_task:
             self._initialized_externally = True
             self.task = current_task
             return
 
-        kwargs.setdefault("project_name", os.environ.get("CLEARML_PROJECT", run_name))
-        kwargs.setdefault("task_name", os.environ.get("CLEARML_TASK", run_name))
-        self.task = Task.init(**kwargs)
+        task_init_args = {**self.init_kwargs}
+        task_init_args.setdefault("project_name", os.environ.get("CLEARML_PROJECT", self.user_provided_run_name))
+        task_init_args.setdefault("task_name", os.environ.get("CLEARML_TASK", self.user_provided_run_name))
+        self.task = Task.init(**task_init_args)
 
     @property
     def tracker(self):
@@ -962,12 +996,17 @@ class DVCLiveTracker(GeneralTracker):
     name = "dvclive"
     requires_logging_directory = False
 
-    @on_main_process
     def __init__(self, run_name: Optional[str] = None, live: Optional[Any] = None, **kwargs):
+        super().__init__()
+        self.live = live
+        self.init_kwargs = kwargs
+
+    @on_main_process
+    def start(self):
         from dvclive import Live
 
         super().__init__()
-        self.live = live if live is not None else Live(**kwargs)
+        self.live = self.live if self.live is not None else Live(**self.init_kwargs)
 
     @property
     def tracker(self):
diff --git a/src/accelerate/utils/imports.py b/src/accelerate/utils/imports.py
index 42511201122..d3bdfa92053 100644
--- a/src/accelerate/utils/imports.py
+++ b/src/accelerate/utils/imports.py
@@ -282,6 +282,10 @@ def is_comet_ml_available():
     return _is_package_available("comet_ml")
 
 
+def is_aim_available():
+    return _is_package_available("aim")
+
+
 def is_boto3_available():
     return _is_package_available("boto3")
 
diff --git a/tests/test_tracking.py b/tests/test_tracking.py
index dc627d80d83..37e4c9bf398 100644
--- a/tests/test_tracking.py
+++ b/tests/test_tracking.py
@@ -31,6 +31,7 @@
 
 # We use TF to parse the logs
 from accelerate import Accelerator
+from accelerate.state import PartialState
 from accelerate.test_utils.testing import (
     MockingTestCase,
     TempDirTestCase,
@@ -42,9 +43,19 @@
     require_pandas,
     require_tensorboard,
     require_wandb,
+    require_aim,
     skip,
 )
-from accelerate.tracking import CometMLTracker, GeneralTracker, MLflowTracker
+from accelerate.tracking import (
+    GeneralTracker,
+    TensorBoardTracker,
+    WandBTracker,
+    CometMLTracker,
+    AimTracker,
+    MLflowTracker,
+    ClearMLTracker,
+    DVCLiveTracker,
+)
 from accelerate.utils import (
     ProjectConfiguration,
     is_comet_ml_available,
@@ -525,10 +536,17 @@ class MyCustomTracker(GeneralTracker):
     name = "my_custom_tracker"
     requires_logging_directory = False
 
-    def __init__(self, dir: str):
-        self.f = open(f"{dir}/log.csv", "w+")
-        self.writer = csv.DictWriter(self.f, fieldnames=self._col_names)
-        self.writer.writeheader()
+    def __init__(self, dir: str, **kwargs):
+        super().__init__(**kwargs)
+        self.log_dir = dir
+        self.f = None
+        self.writer = None
+
+    def start(self):
+        if self.f is None:
+            self.f = open(os.path.join(self.log_dir, "log.csv"), "w+")
+            self.writer = csv.DictWriter(self.f, fieldnames=self._col_names)
+            self.writer.writeheader()
 
     @property
     def tracker(self):
@@ -591,6 +609,7 @@ def test_log(self):
                 assert data == truth
 
 
+
 @require_dvclive
 @mock.patch("dvclive.live.get_dvc_repo", return_value=None)
 class DVCLiveTrackingTest(unittest.TestCase):
@@ -634,3 +653,77 @@ def test_log(self, mock_repo):
                 val_path = os.path.join(scalars, f"{val}.tsv")
                 steps = [int(row["step"]) for row in logs[val_path]]
                 assert steps == [0, 1, 3]
+
+class TrackerDeferredInitializationTest(unittest.TestCase):
+    """
+    Tests tracker's deferred initialization via `start()` method, preventing
+    premature `PartialState` access (and `torch.distributed` init) before
+    `Accelerator` has configured the distributed environment, especially with
+    `InitProcessGroupKwargs`.
+    """
+    @require_tensorboard
+    def test_tensorboard_deferred_init(self):
+        """Test that TensorBoard tracker initialization doesn't initialize distributed"""
+        with tempfile.TemporaryDirectory() as temp_dir:
+            PartialState._reset_state()
+            tracker = TensorBoardTracker(run_name="test_tb", logging_dir=temp_dir)
+            self.assertEqual(PartialState._shared_state, {})
+            accelerator = Accelerator(log_with=tracker)
+            self.assertNotEqual(PartialState._shared_state, {})
+
+    @require_wandb
+    def test_wandb_deferred_init(self):
+        """Test that WandB tracker initialization doesn't initialize distributed"""
+        PartialState._reset_state()
+        tracker = WandBTracker(run_name="test_wandb")
+        self.assertEqual(PartialState._shared_state, {})
+        accelerator = Accelerator(log_with=tracker)
+        self.assertNotEqual(PartialState._shared_state, {})
+
+    @require_comet_ml
+    def test_comet_ml_deferred_init(self):
+        """Test that CometML tracker initialization doesn't initialize distributed"""
+        PartialState._reset_state()
+        tracker = CometMLTracker(run_name="test_comet")
+        self.assertEqual(PartialState._shared_state, {})
+        accelerator = Accelerator(log_with=tracker)
+        self.assertNotEqual(PartialState._shared_state, {})
+
+    @require_aim
+    def test_aim_deferred_init(self):
+        """Test that Aim tracker initialization doesn't initialize distributed"""
+        with tempfile.TemporaryDirectory() as temp_dir:
+            PartialState._reset_state()
+            tracker = AimTracker(run_name="test_aim", repo=temp_dir)
+            self.assertEqual(PartialState._shared_state, {})
+            accelerator = Accelerator(log_with=tracker)
+            self.assertNotEqual(PartialState._shared_state, {})
+
+    @require_mlflow
+    def test_mlflow_deferred_init(self):
+        """Test that MLflow tracker initialization doesn't initialize distributed"""
+        with tempfile.TemporaryDirectory() as temp_dir:
+            PartialState._reset_state()
+            tracker = MLflowTracker(experiment_name="test_mlflow", logging_dir=temp_dir)
+            self.assertEqual(PartialState._shared_state, {})
+            accelerator = Accelerator(log_with=tracker)
+            self.assertNotEqual(PartialState._shared_state, {})
+
+    @require_clearml
+    def test_clearml_deferred_init(self):
+        """Test that ClearML tracker initialization doesn't initialize distributed"""
+        PartialState._reset_state()
+        tracker = ClearMLTracker(run_name="test_clearml")
+        self.assertEqual(PartialState._shared_state, {})
+        accelerator = Accelerator(log_with=tracker)
+        self.assertNotEqual(PartialState._shared_state, {})
+
+    @require_dvclive
+    def test_dvclive_deferred_init(self):
+        """Test that DVCLive tracker initialization doesn't initialize distributed"""
+        with tempfile.TemporaryDirectory() as temp_dir:
+            PartialState._reset_state()
+            tracker = DVCLiveTracker(dir=temp_dir)
+            self.assertEqual(PartialState._shared_state, {})
+            accelerator = Accelerator(log_with=tracker)
+            self.assertNotEqual(PartialState._shared_state, {})
\ No newline at end of file


Code review comments:
diff_hunk: @@ -962,12 +996,17 @@ class DVCLiveTracker(GeneralTracker):
     name = "dvclive"
     requires_logging_directory = False
 
-    @on_main_process
     def __init__(self, run_name: Optional[str] = None, live: Optional[Any] = None, **kwargs):
+        super().__init__()
+        self.live = live
+        self.init_kwargs = kwargs
+
+    @on_main_process
+    def start(self):
         from dvclive import Live
 
         super().__init__()
path: src/accelerate/tracking.py
text: @user1:
```suggestion
```

@author:
Thanks!