chore: formatting

Author: sehoffmann

dmlcloud/core/distributed.py

@@ -105,12 +105,12 @@ def _initialize_via_tcp(
     msg += f'\n  master port: {port}'
     print(msg, flush=True)
 
-    # TODO: Add check that ip == rank0 host 
+    # TODO: Add check that ip == rank0 host
 
     store = dist.TCPStore(
-        host_name=ip, 
+        host_name=ip,
         port=port,
-        world_size=_WorkerInfo.WORLD_SIZE, 
+        world_size=_WorkerInfo.WORLD_SIZE,
         is_master=_WorkerInfo.RANK == 0,
     )
 
@@ -144,7 +144,6 @@ def has_environment():
     return 'MASTER_PORT' in os.environ
 
 
-
 def rank():
     """
     Returns the rank of the current process.
@@ -200,7 +199,6 @@ def local_node():
     return _WorkerInfo.NODE_ID
 
 
-
 def _init_process_group_env(**kwargs):
     """
     Intialize using "env://" method.
@@ -210,7 +208,7 @@ def _init_process_group_env(**kwargs):
     """
     if not has_environment():
         raise RuntimeError('Environment variables for env:// initialization not found')
-    
+
     _initialize_via_tcp(
         ip=os.environ['MASTER_ADDR'],
         port=int(os.environ['MASTER_PORT']),
@@ -356,7 +354,6 @@ def deinitialize_torch_distributed(fail_silently=False):
     dist.destroy_process_group()
 
 
-
 def is_root(group: dist.ProcessGroup = None):
     """
     Check if the current rank is the root rank (rank 0).
@@ -518,7 +515,6 @@ def root_first(group: dist.ProcessGroup = None):
             pass
 
 
-
 def all_gather_object(obj, group=None):
     """
     Gather objects from all ranks in the group.

dmlcloud/core/pipeline.py

@@ -66,7 +66,9 @@ def current_stage() -> Stage | None:
         return current_pipe().current_stage
 
 
-def log_metric(name: str, value: Any, reduction: str = 'mean', prefixed: bool = True, ignore_missing_stage: bool = True):
+def log_metric(
+    name: str, value: Any, reduction: str = 'mean', prefixed: bool = True, ignore_missing_stage: bool = True
+):
     """
     Shorthand for current_stage().log
 

dmlcloud/core/stage.py

@@ -167,13 +167,21 @@ def add_callback(self, callback: 'Callback', priority: int = 1):
         """
         self.callbacks.append(callback, priority)
 
-    def log(self, name: str, value: Any, reduction: str = 'mean', prefixed: bool = True, log_step: bool = True, synchronize: bool = True):
+    def log(
+        self,
+        name: str,
+        value: Any,
+        reduction: str = 'mean',
+        prefixed: bool = True,
+        log_step: bool = True,
+        synchronize: bool = True,
+    ):
         """
         Logs a metric for the current step and epoch.
 
         If `synchronize` is True, the metric will be (all-)reduced across distributed processes before being logged.
         Care must be taken to ensure that every process participates in this reduction to avoid hangs and failures.
-        
+
         Args:
             name (str): The name of the metric.
             value (Any): The value of the metric.