[BUG] Fix bugs in dist env (#1251)

* add necessary synchronization for data_iter destroying and creating bug when in dist

* refine examples code

* print CPU mem info by default, add psutil to requirements.txt and refine install_mesh.sh

* fix train_LBFGS_epoch_func too

* Update install_mesh.sh

Co-authored-by: Copilot <[email protected]>

* refine code as copilot

* move uss&pss_msg to benchmark_flag branch

* Apply suggestion from @Copilot

Co-authored-by: Copilot <[email protected]>

* refine

* update API doc

---------

Co-authored-by: Copilot <[email protected]>

Author: HydrogenSulfate

docs/zh/api/utils/misc.md

@@ -9,6 +9,7 @@
         - Prettydefaultdict
         - RankZeroOnly
         - RankZeroFirst
+        - Synchronized
         - Timer
         - all_gather
         - concat_dict_list

examples/amgnet/amgnet_airfoil.py

@@ -86,6 +86,7 @@ def train(cfg: DictConfig):
         loss=ppsci.loss.FunctionalLoss(train_mse_func),
         name="Sup",
     )
+    cfg.TRAIN.iters_per_epoch = len(sup_constraint.data_loader)
     # wrap constraints together
     constraint = {sup_constraint.name: sup_constraint}
 
@@ -121,16 +122,9 @@ def train(cfg: DictConfig):
     solver = ppsci.solver.Solver(
         model,
         constraint,
-        cfg.output_dir,
-        optimizer,
-        None,
-        cfg.TRAIN.epochs,
-        cfg.TRAIN.iters_per_epoch,
-        save_freq=cfg.TRAIN.save_freq,
-        eval_during_train=cfg.TRAIN.eval_during_train,
-        eval_freq=cfg.TRAIN.eval_freq,
+        optimizer=optimizer,
         validator=validator,
-        eval_with_no_grad=cfg.EVAL.eval_with_no_grad,
+        cfg=cfg,
     )
     # train model
     solver.train()

examples/biharmonic2d/biharmonic2d.py

@@ -413,10 +413,9 @@ def inference(cfg: DictConfig):
     x_flatten = x_grad.reshape(-1, 1)
     y_flatten = y_grad.reshape(-1, 1)
 
-    with ppsci.misc.Timer("infer"):
-        output_dict = predictor.predict(
-            {"x": x_flatten, "y": y_flatten}, cfg.INFER.batch_size
-        )
+    output_dict = predictor.predict(
+        {"x": x_flatten, "y": y_flatten}, cfg.INFER.batch_size
+    )
 
     # mapping data to cfg.INFER.output_keys
     output_dict = {

install_mesh.sh

@@ -8,28 +8,22 @@ echo "==== Python Environment Verified ===="
 
 
 echo "==== Step 2: Checking CMake Environment ===="
-# Verify if cmake exists in current environment. If not, download a temporary version.
-if command -v cmake >/dev/null 2>&1; then
-    echo "CMake found: $(cmake --version | head -n 1)"
-else
-    echo "CMake not found. Installing temporary CMake 3.23.0 ..."
-    wget -c https://paddle-org.bj.bcebos.com/paddlescience/cmake-3.23.0-linux-x86_64.tar.gz
-    tar -zxf cmake-3.23.0-linux-x86_64.tar.gz --checkpoint=.100 --totals
-    rm -f cmake-3.23.0-linux-x86_64.tar.gz
-    export PATH=$PWD/cmake-3.23.0-linux-x86_64/bin:$PATH
-    echo "Temporary CMake installed: $(cmake --version | head -n 1)"
-fi
+# Always install a temporary version of CMake 3.23.0, regardless of existing installations.
+echo "Installing temporary CMake 3.23.0 ..."
+wget -c https://paddle-org.bj.bcebos.com/paddlescience/cmake-3.23.0-linux-x86_64.tar.gz
+tar -zxf cmake-3.23.0-linux-x86_64.tar.gz --checkpoint=.100 --totals
+rm -f cmake-3.23.0-linux-x86_64.tar.gz
+export PATH=$PWD/cmake-3.23.0-linux-x86_64/bin:$PATH
+echo "Temporary CMake installed: $(cmake --version | head -n 1)"
 echo "CMake environment ready."
 
-
 echo "==== Step 3: Downloading PyMesh Package ===="
 # Download PyMesh package if not already present.
 wget -c https://paddle-org.bj.bcebos.com/paddlescience/PyMesh.tar.gz
 echo "Download completed. Extracting package..."
 tar -zxf PyMesh.tar.gz --checkpoint=.1000 --totals
 echo "PyMesh package extracted."
 
-
 echo "==== Step 4: Entering PyMesh Directory ===="
 cd PyMesh
 export PYMESH_PATH=$(pwd)

ppsci/solver/printer.py

@@ -17,6 +17,7 @@
 from typing import Dict
 from typing import Optional
 
+import psutil
 from paddle import device
 
 from ppsci.utils import logger
@@ -80,13 +81,21 @@ def log_train_info(
         + ", ".join(filter(None, [metric_msg, time_msg, ips_msg, eta_msg]))
     )
     if solver.benchmark_flag:
+        # GPU memory
         max_mem_reserved_msg = (
             f"max_mem_reserved: {device.max_memory_reserved() // (1 << 20)} MB"
         )
         max_mem_allocated_msg = (
             f"max_mem_allocated: {device.max_memory_allocated() // (1 << 20)} MB"
         )
         log_str += f", {max_mem_reserved_msg}, {max_mem_allocated_msg}"
+
+        # CPU memory
+        _process = psutil.Process()
+        _mem_full = _process.memory_full_info()
+        uss_msg = f"USS: {_mem_full.uss // (1 << 20)} MB"  # Unique Set Size (USS) memory in MB
+        pss_msg = f"PSS: {_mem_full.pss // (1 << 20)} MB"  # Proportional Set Size (PSS) memory in MB
+        log_str += f", {uss_msg}, {pss_msg}"
     logger.info(log_str)
 
     # reset time information after printing

ppsci/solver/train.py

@@ -14,6 +14,7 @@
 
 from __future__ import annotations
 
+import gc
 import sys
 import time
 from typing import TYPE_CHECKING
@@ -88,7 +89,11 @@ def train_epoch_func(solver: "solver.Solver", epoch_id: int, log_freq: int):
             try:
                 input_dict, label_dict, weight_dict = next(_constraint.data_iter)
             except StopIteration:
-                _constraint.data_iter = iter(_constraint.data_loader)
+                with misc.Synchronized(solver.world_size > 1):
+                    if hasattr(_constraint, "data_iter"):
+                        del _constraint.data_iter
+                        gc.collect()
+                    _constraint.data_iter = iter(_constraint.data_loader)
                 input_dict, label_dict, weight_dict = next(_constraint.data_iter)
 
             if solver.nvtx_flag:  # only for nsight analysis
@@ -245,7 +250,11 @@ def train_LBFGS_epoch_func(solver: "solver.Solver", epoch_id: int, log_freq: int
             try:
                 input_dict, label_dict, weight_dict = next(_constraint.data_iter)
             except StopIteration:
-                _constraint.data_iter = iter(_constraint.data_loader)
+                with misc.Synchronized(solver.world_size > 1):
+                    if hasattr(_constraint, "data_iter"):
+                        del _constraint.data_iter
+                        gc.collect()
+                    _constraint.data_iter = iter(_constraint.data_loader)
                 input_dict, label_dict, weight_dict = next(_constraint.data_iter)
             reader_cost += time.perf_counter() - reader_tic
 

ppsci/utils/misc.py

@@ -43,6 +43,7 @@
     "Prettydefaultdict",
     "RankZeroOnly",
     "RankZeroFirst",
+    "Synchronized",
     "Timer",
     "all_gather",
     "concat_dict_list",
@@ -226,6 +227,43 @@ def __exit__(self, type, value, traceback):
             dist.barrier()  # Allow others to proceed
 
 
+class Synchronized(ContextDecorator):
+    """
+    A context manager/decorator that ensures code blocks are synchronized across all processes.
+
+    It calls barrier before and after the code block execution to ensure synchronization
+    among all processes.
+
+    Args:
+        enabled (bool): Whether to enable synchronization. Defaults to True.
+
+    Examples:
+        >>> import paddle
+        >>> from paddle import distributed as dist
+        >>> dist.init_parallel_env()
+        >>> with Synchronized(dist.get_world_size() > 1):
+        ...     x = paddle.randn(2, 2) * 2
+    """
+
+    def __init__(self, enabled: bool = True):
+        if enabled and not dist.is_initialized():
+            logger.warning(
+                "Distributed environment is not initialized, but `Synchronized` is enabled. "
+                "This may cause unexpected behavior. Please ensure distributed environment is properly initialized "
+                "when using `Synchronized` with `enabled=True`.",
+            )
+        self.enabled = enabled
+
+    def __enter__(self):
+        if self.enabled and dist.is_initialized():
+            dist.barrier()
+        return self
+
+    def __exit__(self, type, value, traceback):
+        if self.enabled and dist.is_initialized():
+            dist.barrier()
+
+
 class Timer(ContextDecorator):
     """Count time cost for code block within context.
 

requirements.txt

@@ -6,6 +6,7 @@ imageio
 matplotlib
 meshio==5.3.4
 numpy>=1.20.0,<2.0.0
+psutil
 pydantic>=2.5.0
 pyevtk
 pyyaml