fix: silence instrumentor noise during traincheck-collect

- collect_trace: use force=True + WARNING default level so basicConfig
  format actually applies and INFO chatter is suppressed at runtime
- source_file: demote all annotate_stage insertion logger.info → debug
- dumper: demote attribute-dump failure logger.warning → debug (torch
  internals routinely fail, not actionable)
- call_graph_parser: convert all print() → logger.debug(); add logger

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: Essoz

traincheck/collect_trace.py

@@ -431,13 +431,13 @@ def main():
         if args.warm_up_steps is None:
             args.warm_up_steps = config.INSTRUMENTATION_POLICY["warm_up"]
 
-    # set up logging
+    # set up logging (force=True overrides any handler already added by imported libs)
     _log_fmt = "[TrainCheck] %(levelname)s: %(message)s"
     if args.debug_mode:
-        logging.basicConfig(level=logging.DEBUG, format=_log_fmt)
+        logging.basicConfig(level=logging.DEBUG, format=_log_fmt, force=True)
         os.environ["TRAINCHECK_DEBUG"] = "1"
     else:
-        logging.basicConfig(level=logging.INFO, format=_log_fmt)
+        logging.basicConfig(level=logging.WARNING, format=_log_fmt, force=True)
 
     logger = logging.getLogger(__name__)
 

traincheck/instrumentor/dumper.py

@@ -392,7 +392,7 @@ def convert_var_to_dict(var, include_tensor_data=True, dump_config=None) -> dict
                 and isinstance(var, torch.Tensor)
                 and not include_tensor_data
             ):
-                logger.warning(
+                logger.debug(
                     f"Failed to get attribute {attr_name} of object type {type(var)}, skipping it for all following dumps for this attribute."
                 )
                 if var_type not in skip_attrs_due_to_errs:

traincheck/instrumentor/source_file.py

@@ -736,7 +736,7 @@ def has_stage(src: str, name: str) -> bool:
 
     for stage_name, present in orig_has.items():
         if present:
-            logger.info(
+            logger.debug(
                 _ctx(
                     f"Stage '{stage_name}' already present in source; skip adding this stage."
                 )
@@ -819,13 +819,13 @@ def at_attr(name: str) -> bool:
                 indent = m.group(0)
                 new_lines.append(f'{indent}annotate_stage("{stage}")\n')
                 inserted_count[stage] += 1
-                logger.info(
+                logger.debug(
                     _ctx(
                         f"Inserted stage '{stage}' before line {lineno}: {line.strip()}"
                     )
                 )
             else:
-                logger.info(
+                logger.debug(
                     _ctx(
                         f"Skip inserting '{stage}' at line {lineno} (previous non-empty line already has it)."
                     )
@@ -858,7 +858,7 @@ def _find_annotate_import_idx(lines):
         lines_list.insert(insert_idx, "from traincheck import annotate_stage\n")
         annot_import_idx = insert_idx
         inserted_count["import"] += 1
-        logger.info(
+        logger.debug(
             _ctx(
                 f"Inserted import 'from traincheck import annotate_stage' at line {annot_import_idx + 1}."
             )
@@ -917,13 +917,13 @@ def is_single_line_triple_quoted_string(line: str, quote: str) -> bool:
             if not (("annotate_stage" in prev) and ("init" in prev)):
                 nl.insert(insert_at, f'{body_indent}annotate_stage("init")\n')
                 inserted_count["init"] += 1
-                logger.info(
+                logger.debug(
                     _ctx(
                         f"Inserted stage 'init' at start of main() body (line {insert_at + 1})."
                     )
                 )
             else:
-                logger.info(
+                logger.debug(
                     _ctx(
                         "Skip inserting 'init' inside main(): previous non-empty line already has it."
                     )
@@ -966,13 +966,13 @@ def is_single_line_triple_quoted_string(line: str, quote: str) -> bool:
             if not (("annotate_stage" in next_line) and ("init" in next_line)):
                 lines2.insert(insert_at, 'annotate_stage("init")\n')
                 inserted_count["init"] += 1
-                logger.info(
+                logger.debug(
                     _ctx(
                         f"Inserted stage 'init' right after annotate_stage import at line {insert_at + 1}."
                     )
                 )
             else:
-                logger.info(
+                logger.debug(
                     _ctx(
                         "Skip inserting 'init': next non-empty line after annotate_stage import is already init."
                     )

traincheck/static_analyzer/graph_generator/call_graph_parser.py

@@ -1,11 +1,14 @@
 ## Step 1: filer out the lines from func_level.log with the format <Node function:module_name.function_name> - function_depth, e.g. <Node function:torch.optim.lr_scheduler._format_param._copy> - 2
 
 import importlib
+import logging
 import os
 import re
 
 from traincheck.instrumentor.proxy_wrapper.proxy_observer import add_observer_to_func
 
+logger = logging.getLogger(__name__)
+
 
 def unparse_module(module_name, level=0):
     if level > 3:
@@ -19,10 +22,9 @@ def unparse_module(module_name, level=0):
         last_name = module_name.split(".")[-1]
         try:
             func_obj = getattr(module, last_name)
-            # print(f"object {last_name} found in module {'.'.join(module_name.split('.')[:-1])}")
             return func_obj
         except AttributeError:
-            print(
+            logger.debug(
                 f"object {last_name} not found in module {'.'.join(module_name.split('.')[:-1])}"
             )
             # Ziming: from out observation this typically just mean a function call contains a local class or function call, so we can just pass
@@ -39,7 +41,9 @@ def add_observer(module_name, function_name, observe_then_unproxy=False):
         # module could be a class here, load the class and get the function
         module = unparse_module(module_name)
         if module is None:
-            print(f"error finding {function_name}: module {module_name} not found")
+            logger.debug(
+                f"error finding {function_name}: module {module_name} not found"
+            )
 
     try:
         # Retrieve the function or property
@@ -48,23 +52,22 @@ def add_observer(module_name, function_name, observe_then_unproxy=False):
 
         # Check if it's a property before proceeding
         if isinstance(function, property):
-            print(
+            logger.debug(
                 f"Skipping property function: {function_name} in module: {module_name}"
             )
             return
 
         # Apply observer to non-property functions
-        print(f"Observe function: {function_name} found in module: {module}")
+        logger.debug(f"Observe function: {function_name} found in module: {module}")
         setattr(
             module,
             function_name,
             add_observer_to_func(function, observe_then_unproxy),
         )
 
     except AttributeError:
-        print(f"function {function_name} not found in module {module_name}")
+        logger.debug(f"function {function_name} not found in module {module_name}")
         return
-    # print(f'function: {function} found in module: {module}')
 
 
 # read the func_level.log file
@@ -104,11 +107,9 @@ def add_observer_given_call_graph(
                 # save those with function_depth <= depth
                 if observe_up_to_depth:
                     if int(function_depth) <= depth:
-                        # print(f'module_name: {module_name}, function_name: {function_name}, function_depth: {function_depth}')
                         add_observer(module_name, function_name, observe_then_unproxy)
                 else:
                     if int(function_depth) == depth:
-                        # print(f'module_name: {module_name}, function_name: {function_name}, function_depth: {function_depth}')
                         add_observer(module_name, function_name, observe_then_unproxy)