Update decorators, add timeit

Author: alaydshah

python/fedml/computing/scheduler/master/base_master_job_runner.py

@@ -16,7 +16,7 @@
 from ....core.mlops.mlops_utils import MLOpsUtils
 from ..scheduler_core.log_manager import LogsManager
 from ..scheduler_core.metrics_manager import MetricsManager
-from fedml.utils.debugging import debug
+from fedml.utils.decorators import debug
 from ..scheduler_core.status_center import JobStatus
 from ..scheduler_core.compute_cache_manager import ComputeCacheManager
 from multiprocessing import Process, Queue

python/fedml/utils/debugging.py renamed to python/fedml/utils/decorators.py

@@ -1,9 +1,10 @@
 from typing import Callable
 import os
+import time
 import functools
 
-def debug(_func: Callable=None, *, output_file="output.txt"):
 
+def debug(_func: Callable = None, *, output_file="output.txt"):
     def decorator(func: Callable):
 
         @functools.wraps(func)
@@ -36,3 +37,16 @@ def wrapper(*args, **kwargs):
         return decorator
 
     return decorator(_func)
+
+def timeit(func: Callable):
+    """Print the runtime of the decorated function"""
+    functools.wraps(func)
+
+    def wrapper(*args, **kwargs):
+        start = time.perf_counter()
+        func(*args, **kwargs)
+        end = time.perf_counter()
+        run_time = end - start
+        print(f"Finished {func.__name__!r} in {run_time:.4f} seconds")
+
+    return wrapper