drop tqdm

per discussion.

Signed-off-by: Jinzhe Zeng <[email protected]>

backend/dynamic_metadata.py

@@ -90,6 +90,5 @@ def dynamic_metadata(
             ],
             "torch": [
                 "torch>=2a",
-                "tqdm",
             ],
         }

deepmd/pt/train/training.py

@@ -14,12 +14,6 @@
 
 import numpy as np
 import torch
-from tqdm import (
-    tqdm,
-)
-from tqdm.contrib.logging import (
-    logging_redirect_tqdm,
-)
 
 from deepmd.common import (
     symlink_prefix_files,
@@ -47,7 +41,6 @@
 )
 from deepmd.pt.utils.env import (
     DEVICE,
-    DISABLE_TQDM,
     JIT,
     LOCAL_RANK,
     NUM_WORKERS,
@@ -662,29 +655,24 @@ def log_loss_valid(_task_key="Default"):
                     f.write(str(self.latest_model))
 
         self.t0 = time.time()
-        with logging_redirect_tqdm():
-            for step_id in tqdm(
-                range(self.num_steps),
-                disable=(bool(dist.get_rank()) if dist.is_initialized() else False)
-                or DISABLE_TQDM,
-            ):  # set to None to disable on non-TTY; disable on not rank 0
-                if step_id < self.start_step:
-                    continue
-                if self.multi_task:
-                    chosen_index_list = dp_random.choice(
-                        np.arange(self.num_model),
-                        p=np.array(self.model_prob),
-                        size=self.world_size,
-                        replace=True,
-                    )
-                    assert chosen_index_list.size == self.world_size
-                    model_index = chosen_index_list[self.rank]
-                    model_key = self.model_keys[model_index]
-                else:
-                    model_key = "Default"
-                step(step_id, model_key)
-                if JIT:
-                    break
+        for step_id in range(self.num_steps):
+            if step_id < self.start_step:
+                continue
+            if self.multi_task:
+                chosen_index_list = dp_random.choice(
+                    np.arange(self.num_model),
+                    p=np.array(self.model_prob),
+                    size=self.world_size,
+                    replace=True,
+                )
+                assert chosen_index_list.size == self.world_size
+                model_index = chosen_index_list[self.rank]
+                model_key = self.model_keys[model_index]
+            else:
+                model_key = "Default"
+            step(step_id, model_key)
+            if JIT:
+                break
 
         if (
             self.rank == 0 or dist.get_rank() == 0

deepmd/pt/utils/dataset.py

@@ -13,9 +13,6 @@
 from torch.utils.data import (
     Dataset,
 )
-from tqdm import (
-    trange,
-)
 
 from deepmd.pt.utils import (
     dp_random,
@@ -506,7 +503,7 @@ def preprocess(self, batch):
         assert batch["atype"].max() < len(self._type_map)
         nlist, nlist_loc, nlist_type, shift, mapping = [], [], [], [], []
 
-        for sid in trange(n_frames, disable=env.DISABLE_TQDM):
+        for sid in range(n_frames):
             region = Region3D(box[sid])
             nloc = atype[sid].shape[0]
             _coord = normalize_coord(coord[sid], region, nloc)

deepmd/pt/utils/env.py

@@ -8,7 +8,6 @@
 GLOBAL_NP_FLOAT_PRECISION = getattr(np, PRECISION)
 GLOBAL_PT_FLOAT_PRECISION = getattr(torch, PRECISION)
 GLOBAL_ENER_FLOAT_PRECISION = getattr(np, PRECISION)
-DISABLE_TQDM = os.environ.get("DISABLE_TQDM", False)
 SAMPLER_RECORD = os.environ.get("SAMPLER_RECORD", False)
 try:
     # only linux

deepmd/pt/utils/stat.py

@@ -3,9 +3,6 @@
 
 import numpy as np
 import torch
-from tqdm import (
-    trange,
-)
 
 from deepmd.pt.utils import (
     env,
@@ -40,7 +37,7 @@ def make_stat_input(datasets, dataloaders, nbatches):
     if datasets[0].mixed_type:
         keys.append("real_natoms_vec")
     logging.info(f"Packing data for statistics from {len(datasets)} systems")
-    for i in trange(len(datasets), disable=env.DISABLE_TQDM):
+    for i in range(len(datasets)):
         sys_stat = {key: [] for key in keys}
         iterator = iter(dataloaders[i])
         for _ in range(nbatches):