pytorch · vmoens · Jan 9, 2024 · Dec 21, 2023 · Dec 21, 2023 · Jan 9, 2024
diff --git a/.github/unittest/helpers/coverage_run_parallel.py b/.github/unittest/helpers/coverage_run_parallel.py
@@ -11,7 +11,7 @@
 nevertheless. It writes temporary coverage config files on the fly and
 invokes coverage with proper arguments
 """
-
+import logging
 import os
 import shlex
 import subprocess
@@ -45,7 +45,9 @@ def write_config(config_path: Path, argv: List[str]) -> None:
 
 def main(argv: List[str]) -> int:
     if len(argv) < 1:
-        print("Usage: 'python coverage_run_parallel.py <command> [command arguments]'")
+        logging.info(
+            "Usage: 'python coverage_run_parallel.py <command> [command arguments]'"
+        )
         sys.exit(1)
     # The temporary config is written into a temp dir that will be deleted
     # including all contents on context exit.

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -28,6 +28,7 @@ repos:
           - flake8-bugbear==22.10.27
           - flake8-comprehensions==3.10.1
           - torchfix==0.0.2
+          - flake8-print==5.0.0
 
   - repo: https://github.com/PyCQA/pydocstyle
     rev: 6.1.1

diff --git a/benchmarks/benchmark_batched_envs.py b/benchmarks/benchmark_batched_envs.py
@@ -68,8 +68,8 @@ def run_env(env):
         devices.append("cuda")
     for device in devices:
         for num_workers in [1, 4, 16]:
-            print(f"With num_workers={num_workers}, {device}")
-            print("Multithreaded...")
+            logging.info(f"With num_workers={num_workers}, {device}")
+            logging.info("Multithreaded...")
             env_multithreaded = create_multithreaded(num_workers, device)
             res_multithreaded = Timer(
                 stmt="run_env(env)",
@@ -78,7 +78,7 @@ def run_env(env):
             )
             time_multithreaded = res_multithreaded.blocked_autorange().mean
 
-            print("Serial...")
+            logging.info("Serial...")
             env_serial = create_serial(num_workers, device)
             res_serial = Timer(
                 stmt="run_env(env)",
@@ -87,7 +87,7 @@ def run_env(env):
             )
             time_serial = res_serial.blocked_autorange().mean
 
-            print("Parallel...")
+            logging.info("Parallel...")
             env_parallel = create_parallel(num_workers, device)
             res_parallel = Timer(
                 stmt="run_env(env)",
@@ -96,7 +96,6 @@ def run_env(env):
             )
             time_parallel = res_parallel.blocked_autorange().mean
 
-            print(time_serial, time_parallel, time_multithreaded)
             res[f"num_workers_{num_workers}_{device}"] = {
                 "Serial, s": time_serial,
                 "Parallel, s": time_parallel,

diff --git a/benchmarks/conftest.py b/benchmarks/conftest.py
@@ -2,7 +2,7 @@
 #
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
-
+import logging
 import os
 import time
 import warnings
@@ -32,7 +32,7 @@ def pytest_sessionfinish(maxprint=50):
         out_str += f"\t{key}{spaces}{item: 4.4f}s\n"
         if i == maxprint - 1:
             break
-    print(out_str)
+    logging.info(out_str)
 
 
 @pytest.fixture(autouse=True)

diff --git a/benchmarks/ecosystem/gym_env_throughput.py b/benchmarks/ecosystem/gym_env_throughput.py
@@ -63,7 +63,6 @@ def make(envname=envname, gym_backend=gym_backend):
                 global_step = 0
                 times = []
                 start = time.time()
-                print("Timer started.")
                 for _ in tqdm.tqdm(range(total_frames // num_workers)):
                     env.step(env.action_space.sample())
                     global_step += num_workers

diff --git a/benchmarks/ecosystem/vmas_rllib_vs_torchrl_sampling_performance.py b/benchmarks/ecosystem/vmas_rllib_vs_torchrl_sampling_performance.py
@@ -4,6 +4,7 @@
 # LICENSE file in the root directory of this source tree.
 
 
+import logging
 import os
 import pickle
 
@@ -164,11 +165,11 @@ def run_comparison_torchrl_rllib(
         evaluation = {}
     for framework in ["TorchRL", "RLlib"]:
         if framework not in evaluation.keys():
-            print(f"\nFramework {framework}")
+            logging.info(f"\nFramework {framework}")
             vmas_times = []
             for n_envs in list_n_envs:
                 n_envs = int(n_envs)
-                print(f"Running {n_envs} environments")
+                logging.info(f"Running {n_envs} environments")
                 if framework == "TorchRL":
                     vmas_times.append(
                         (n_envs * n_steps)
@@ -189,7 +190,7 @@ def run_comparison_torchrl_rllib(
                             device=device,
                         )
                     )
-                print(f"fps {vmas_times[-1]}s")
+                logging.info(f"fps {vmas_times[-1]}s")
             evaluation[framework] = vmas_times
 
     store_pickled_evaluation(name=figure_name_pkl, evaluation=evaluation)

diff --git a/benchmarks/storage/benchmark_sample_latency_over_rpc.py b/benchmarks/storage/benchmark_sample_latency_over_rpc.py
@@ -14,6 +14,7 @@
 This code is based on examples/distributed/distributed_replay_buffer.py.
 """
 import argparse
+import logging
 import os
 import pickle
 import sys
@@ -105,10 +106,10 @@ def _create_replay_buffer(self) -> rpc.RRef:
                 buffer_rref = rpc.remote(
                     replay_buffer_info, ReplayBufferNode, args=(1000000,)
                 )
-                print(f"Connected to replay buffer {replay_buffer_info}")
+                logging.info(f"Connected to replay buffer {replay_buffer_info}")
                 return buffer_rref
             except Exception:
-                print("Failed to connect to replay buffer")
+                logging.info("Failed to connect to replay buffer")
                 time.sleep(RETRY_DELAY_SECS)
 
 
@@ -143,7 +144,7 @@ def __init__(self, capacity: int):
     rank = args.rank
     storage_type = args.storage
 
-    print(f"Rank: {rank}; Storage: {storage_type}")
+    logging.info(f"Rank: {rank}; Storage: {storage_type}")
 
     os.environ["MASTER_ADDR"] = "localhost"
     os.environ["MASTER_PORT"] = "29500"
@@ -166,7 +167,7 @@ def __init__(self, capacity: int):
             if i == 0:
                 continue
             results.append(result)
-            print(i, results[-1])
+            logging.info(i, results[-1])
 
         with open(
             f'./benchmark_{datetime.now().strftime("%d-%m-%Y%H:%M:%S")};batch_size={BATCH_SIZE};tensor_size={TENSOR_SIZE};repeat={REPEATS};storage={storage_type}.pkl',
@@ -175,7 +176,7 @@ def __init__(self, capacity: int):
             pickle.dump(results, f)
 
         tensor_results = torch.tensor(results)
-        print(f"Mean: {torch.mean(tensor_results)}")
+        logging.info(f"Mean: {torch.mean(tensor_results)}")
         breakpoint()
     elif rank == 1:
         # rank 1 is the replay buffer

diff --git a/examples/a2c/a2c_atari.py b/examples/a2c/a2c_atari.py
@@ -2,6 +2,7 @@
 #
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
+import logging
 
 import hydra
 
@@ -212,7 +213,7 @@ def main(cfg: "DictConfig"):  # noqa: F821
 
     end_time = time.time()
     execution_time = end_time - start_time
-    print(f"Training took {execution_time:.2f} seconds to finish")
+    logging.info(f"Training took {execution_time:.2f} seconds to finish")
 
 
 if __name__ == "__main__":

diff --git a/examples/a2c/a2c_mujoco.py b/examples/a2c/a2c_mujoco.py
@@ -2,6 +2,7 @@
 #
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
+import logging
 
 import hydra
 
@@ -197,7 +198,7 @@ def main(cfg: "DictConfig"):  # noqa: F821
 
     end_time = time.time()
     execution_time = end_time - start_time
-    print(f"Training took {execution_time:.2f} seconds to finish")
+    logging.info(f"Training took {execution_time:.2f} seconds to finish")
 
 
 if __name__ == "__main__":

diff --git a/examples/cql/cql_offline.py b/examples/cql/cql_offline.py
@@ -9,7 +9,7 @@
 The helper functions are coded in the utils.py associated with this script.
 
 """
-
+import logging
 import time
 
 import hydra
@@ -145,7 +145,7 @@ def main(cfg: "DictConfig"):  # noqa: F821
         log_metrics(logger, to_log, i)
 
     pbar.close()
-    print(f"Training time: {time.time() - start_time}")
+    logging.info(f"Training time: {time.time() - start_time}")
 
 
 if __name__ == "__main__":

diff --git a/examples/cql/cql_online.py b/examples/cql/cql_online.py
@@ -11,7 +11,7 @@
 The helper functions are coded in the utils.py associated with this script.
 
 """
-
+import logging
 import time
 
 import hydra
@@ -206,7 +206,7 @@ def main(cfg: "DictConfig"):  # noqa: F821
     collector.shutdown()
     end_time = time.time()
     execution_time = end_time - start_time
-    print(f"Training took {execution_time:.2f} seconds to finish")
+    logging.info(f"Training took {execution_time:.2f} seconds to finish")
 
     collector.shutdown()
 

diff --git a/examples/cql/discrete_cql_online.py b/examples/cql/discrete_cql_online.py
@@ -10,7 +10,7 @@
 
 The helper functions are coded in the utils.py associated with this script.
 """
-
+import logging
 import time
 
 import hydra
@@ -192,7 +192,7 @@ def main(cfg: "DictConfig"):  # noqa: F821
     collector.shutdown()
     end_time = time.time()
     execution_time = end_time - start_time
-    print(f"Training took {execution_time:.2f} seconds to finish")
+    logging.info(f"Training took {execution_time:.2f} seconds to finish")
 
 
 if __name__ == "__main__":

diff --git a/examples/ddpg/ddpg.py b/examples/ddpg/ddpg.py
@@ -10,7 +10,7 @@
 
 The helper functions are coded in the utils.py associated with this script.
 """
-
+import logging
 import time
 
 import hydra
@@ -192,7 +192,7 @@ def main(cfg: "DictConfig"):  # noqa: F821
     collector.shutdown()
     end_time = time.time()
     execution_time = end_time - start_time
-    print(f"Training took {execution_time:.2f} seconds to finish")
+    logging.info(f"Training took {execution_time:.2f} seconds to finish")
 
 
 if __name__ == "__main__":

diff --git a/examples/decision_transformer/dt.py b/examples/decision_transformer/dt.py
@@ -6,6 +6,7 @@
 This is a self-contained example of an offline Decision Transformer training script.
 The helper functions are coded in the utils.py associated with this script.
 """
+import logging
 import time
 
 import hydra
@@ -78,7 +79,7 @@ def main(cfg: "DictConfig"):  # noqa: F821
     pretrain_log_interval = cfg.logger.pretrain_log_interval
     reward_scaling = cfg.env.reward_scaling
 
-    print(" ***Pretraining*** ")
+    logging.info(" ***Pretraining*** ")
     # Pretraining
     start_time = time.time()
     for i in range(pretrain_gradient_steps):
@@ -115,7 +116,7 @@ def main(cfg: "DictConfig"):  # noqa: F821
             log_metrics(logger, to_log, i)
 
     pbar.close()
-    print(f"Training time: {time.time() - start_time}")
+    logging.info(f"Training time: {time.time() - start_time}")
 
 
 if __name__ == "__main__":

diff --git a/examples/decision_transformer/online_dt.py b/examples/decision_transformer/online_dt.py
@@ -6,7 +6,7 @@
 This is a self-contained example of an Online Decision Transformer training script.
 The helper functions are coded in the utils.py associated with this script.
 """
-
+import logging
 import time
 
 import hydra
@@ -81,7 +81,7 @@ def main(cfg: "DictConfig"):  # noqa: F821
     pretrain_log_interval = cfg.logger.pretrain_log_interval
     reward_scaling = cfg.env.reward_scaling
 
-    print(" ***Pretraining*** ")
+    logging.info(" ***Pretraining*** ")
     # Pretraining
     start_time = time.time()
     for i in range(pretrain_gradient_steps):
@@ -132,7 +132,7 @@ def main(cfg: "DictConfig"):  # noqa: F821
             log_metrics(logger, to_log, i)
 
     pbar.close()
-    print(f"Training time: {time.time() - start_time}")
+    logging.info(f"Training time: {time.time() - start_time}")
 
 
 if __name__ == "__main__":

diff --git a/examples/discrete_sac/discrete_sac.py b/examples/discrete_sac/discrete_sac.py
@@ -10,7 +10,7 @@
 
 The helper functions are coded in the utils.py associated with this script.
 """
-
+import logging
 import time
 
 import hydra
@@ -208,7 +208,7 @@ def main(cfg: "DictConfig"):  # noqa: F821
     collector.shutdown()
     end_time = time.time()
     execution_time = end_time - start_time
-    print(f"Training took {execution_time:.2f} seconds to finish")
+    logging.info(f"Training took {execution_time:.2f} seconds to finish")
 
 
 if __name__ == "__main__":

diff --git a/examples/distributed/collectors/multi_nodes/delayed_dist.py b/examples/distributed/collectors/multi_nodes/delayed_dist.py
@@ -23,6 +23,7 @@
   and DEFAULT_SLURM_CONF_MAIN dictionaries below).
 
 """
+import logging
 import time
 from argparse import ArgumentParser
 
@@ -149,7 +150,7 @@ def make_env():
         if i == 10:
             t0 = time.time()
     t1 = time.time()
-    print(f"time elapsed: {t1-t0}s, rate: {counter/(t1-t0)} fps")
+    logging.info(f"time elapsed: {t1-t0}s, rate: {counter/(t1-t0)} fps")
     collector.shutdown()
     exit()
 

diff --git a/examples/distributed/collectors/multi_nodes/delayed_rpc.py b/examples/distributed/collectors/multi_nodes/delayed_rpc.py
@@ -23,6 +23,7 @@
   and DEFAULT_SLURM_CONF_MAIN dictionaries below).
 
 """
+import logging
 import time
 from argparse import ArgumentParser
 
@@ -147,7 +148,7 @@ def make_env():
         if i == 10:
             t0 = time.time()
     t1 = time.time()
-    print(f"time elapsed: {t1-t0}s, rate: {counter/(t1-t0)} fps")
+    logging.info(f"time elapsed: {t1-t0}s, rate: {counter/(t1-t0)} fps")
     collector.shutdown()
     exit()
 

diff --git a/examples/distributed/collectors/multi_nodes/generic.py b/examples/distributed/collectors/multi_nodes/generic.py
@@ -2,6 +2,7 @@
 #
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
+import logging
 import time
 from argparse import ArgumentParser
 
@@ -127,5 +128,5 @@ def gym_make():
             t0 = time.time()
     collector.shutdown()
     t1 = time.time()
-    print(f"time elapsed: {t1-t0}s, rate: {counter/(t1-t0)} fps")
+    logging.info(f"time elapsed: {t1-t0}s, rate: {counter/(t1-t0)} fps")
     exit()