Fix comms logger KeyError when log_name is omitted

[jinyouzhi]

Summary

Fix a KeyError: 'log_name' raised by the DeepSpeed communication logger
when a wrapped collective is called without an explicit log_name.

"comms_logger": {
    "enabled": true,
    "prof_all": true,
    "debug": true
}

This is exposed by multi-rank AutoTP input consistency checks, which call
broadcast_object_list without passing profiling metadata. Single-rank TP
does not exercise this communication path.

Changes

• Add prof/log_name/debug for broadcast_object_list and all_to_all
• Use the func.__name__ as the default log_name to cover missing status
• Add a regression test

Validation

• python -m pytest -q tests/unit/comm/test_comms_logger.py

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.

[ebarkhordar]

prof_ops never matches an op that takes its log_name from the signature default. The gate at comm.py:111-113 tests 'log_name' in kwargs, so it fires only when a caller passes the name explicitly, while config-json.md documents "prof_ops": ["all_reduce", "all_gather"] against ordinary calls.

At 1f95164 in a clean container, CPU torch and a stub cdb:

prof_ops = ['all_reduce']  prof_all = False
A. dist.all_reduce(t)               -> comms_dict keys: []
B. dist.all_reduce(t, log_name=...) -> comms_dict keys: ['all_reduce']

Your setdefault is one step from covering this. Resolving the name once per op keeps the per-call fast path at two conditions:

def timed_op(func):
    default_log_name = get_default_args(func).get('log_name', func.__name__)

    def log_wrapper(*args, **kwargs):
        if comms_logger.enabled:
            selected = kwargs.get('log_name', default_log_name)
            if kwargs.get('prof') or comms_logger.prof_all or selected in comms_logger.prof_ops:

then func_args['log_name'] = selected in place of the setdefault, and the same condition in the finally gate. With that, A logs and tests/unit/comm/test_comms_logger.py is still 4 passed. It is a separate bug from the KeyError you are fixing, so it may belong in its own PR.

[jinyouzhi]

prof_ops never matches an op that takes its log_name from the signature default. The gate at comm.py:111-113 tests 'log_name' in kwargs, so it fires only when a caller passes the name explicitly, while config-json.md documents "prof_ops": ["all_reduce", "all_gather"] against ordinary calls.

At 1f95164 in a clean container, CPU torch and a stub cdb:

prof_ops = ['all_reduce']  prof_all = False
A. dist.all_reduce(t)               -> comms_dict keys: []
B. dist.all_reduce(t, log_name=...) -> comms_dict keys: ['all_reduce']

Your setdefault is one step from covering this. Resolving the name once per op keeps the per-call fast path at two conditions:

def timed_op(func):
    default_log_name = get_default_args(func).get('log_name', func.__name__)

    def log_wrapper(*args, **kwargs):
        if comms_logger.enabled:
            selected = kwargs.get('log_name', default_log_name)
            if kwargs.get('prof') or comms_logger.prof_all or selected in comms_logger.prof_ops:

then func_args['log_name'] = selected in place of the setdefault, and the same condition in the finally gate. With that, A logs and tests/unit/comm/test_comms_logger.py is still 4 passed. It is a separate bug from the KeyError you are fixing, so it may belong in its own PR.

Wow, that's a very insightful observation! I agree that prof_ops currently fails to match an operation when log_name comes from the function signature default rather than being explicitly passed in kwargs. Would you prefer that I fix this as part of this PR, or you will handle it separately in a follow-up PR?

[ebarkhordar]

Your call as the author, but I would take it in this PR. The repair replaces the setdefault line you just added, so as two PRs whichever one lands second has to rebase through the other for no benefit. That is the opposite of what I said above about a separate PR, and the overlap on that line is why I changed my mind.

Either way I am not going to open a competing PR for it.

One thing to keep if you do take it: the finally gate at comm.py:133 tests the same 'log_name' in kwargs expression, so it needs the same selected condition. Without it the start gate selects the call, the stop gate does not, and timers(log_name).stop() and comms_logger.append never run, so the op still does not appear.

I re-read comm.py at cbbf47b after your merge from master and both gates are unchanged there, so the differential above still stands at your current head.