logging working

Author: lessw2020

distributed/__init__.py

@@ -6,6 +6,7 @@
 
 from distributed.parallelize_llama import parallelize_llama
 from distributed.parallel_config import ParallelDims
-from distributed.utils import init_distributed, logger
+from distributed.utils import init_distributed
 from distributed.checkpoint import load_checkpoints_to_model
 from distributed.world_maker import launch_distributed 
+from distributed.logging_utils import logger

distributed/config_manager.py

@@ -9,7 +9,7 @@
 from collections import defaultdict
 from typing import Tuple, Union
 import os
-from distributed.utils import logger
+from distributed.logging_utils import logger
 from pathlib import Path
 
 import torch
@@ -64,7 +64,7 @@ def parse_args(self, config_file):
         full_path = os.path.join(os.getcwd(), local_path)
         file_path = Path(full_path)
 
-        print(f"Loading config file {file_path}")
+        logger.info(f"Loading config file {config_file}")
 
         if not file_path.is_file():
             raise FileNotFoundError(f"Config file {full_path} does not exist")
@@ -87,8 +87,7 @@ def parse_args(self, config_file):
         # for section, section_args in cmd_args_dict.items():
         #     for k, v in section_args.items():
         #         args_dict[section][k] = v
-        print(f"args_dict: {args_dict}")
-
+        
         for k, v in args_dict.items():
             class_type = type(k.title(), (), v)
             setattr(self, k, class_type())

distributed/parallelize_llama.py

@@ -16,7 +16,7 @@
 from torch.distributed._tensor import Replicate, Shard
 from distributed.parallel_config import ParallelDims
 from torch.distributed.device_mesh import DeviceMesh
-from distributed.utils import logger
+from distributed.logging_utils import logger
 
 
 def apply_tp(

distributed/utils.py

@@ -9,7 +9,7 @@
 
 import torch
 from dataclasses import dataclass, field
-
+from distributed.logging_utils import logger
 
 def _warn_overwrite_env(env, val):
     if env in os.environ:

distributed/world_maker.py

@@ -50,20 +50,20 @@ def launch_distributed(
             - The second element is an optional ParallelDims object, 
             which represents the parallel dimensions configuration.
     """
-    init_logger()
+    #init_logger()  TODO - do we want formatted logging? 
     world_size = int(os.environ["WORLD_SIZE"])
     config = InferenceConfig()
     config.parse_args(toml_config)
-    
-    print(f"logging here...")
-    logger.info(f"***************** from logger")
 
-    assert False, "check"
+    
+    logger.info(f"toml parsing completed.  Launching with {world_size} GPUs")
 
+    
     parallel_dims = ParallelDims(
         tp=8,
         pp=1,
         world_size=world_size,
     )
     init_distributed()
     world_mesh = parallel_dims.build_mesh(device_type="cuda")
+    assert False, "--- function end"