[Code Health] Consuming a producer erases a newer same-run TensorMap mapping for the same key

[ChaoWao]

Category

Robustness (potential edge-case failure)

Component

TensorMap

Description

TensorMap::erase_task_outputs(run_id, keys) erases every key unconditionally:

void TensorMap::erase_task_outputs(RunId run_id, const std::vector<TensorKey> &keys) {
    std::lock_guard<std::mutex> lk(mu_);
    for (const auto &key : keys)
        map_.erase(RunTensorKey{run_id, key});
}

It is called from Orchestrator::on_consumed() with the consumed task's output_keys. If a later task in the same run has since overwritten the mapping for one of those keys, the erase drops the newer producer's entry, and a subsequent consumer infers no dependency at all.

The reachable shape needs two writers with no ordering edge between them, which OUTPUT / OUTPUT_EXISTING produce because that path inserts without looking up:

case TensorArgType::OUTPUT:
case TensorArgType::OUTPUT_EXISTING: {
    tensormap_->insert(run_id, key, slot);   // no lookup -> no producer edge
    output_keys.push_back(key);
    break;
}

Scenario, all in one run:

1. Task A: OUTPUT on key K → insert(run, K, A); A.output_keys = [K]
2. Task B: OUTPUT on key K → insert(run, K, B) overwrites; no edge A→B, so A and B are unordered
3. A reaches CONSUMED (it need not wait for B) → erase_task_outputs(run, [K]) removes the mapping that now points at B
4. Task C: INPUT on key K → lookup(run, K) returns INVALID_SLOT → no dependency on B

If B is still executing, C can be dispatched early and read K before B has written it.

The INOUT case is safe today: INOUT looks up before inserting, so B depends on A, and A cannot reach CONSUMED until B has released its reference — by then the erase is harmless. The hazard is specific to two OUTPUT-only writers of the same key.

Not observed in practice — found by reading the code while reviewing the worker async-preparation series, and verified against main. Pre-existing: the unconditional erase predates the run-scoping added in #1459, which only added the RunId to the key.

Location

• src/common/hierarchical/tensormap.cpp — TensorMap::erase_task_outputs
• src/common/hierarchical/orchestrator.cpp — Orchestrator::on_consumed, the tensormap_->erase_task_outputs(run_id, s.output_keys) call
• src/common/hierarchical/orchestrator.cpp — Orchestrator::infer_deps, the OUTPUT / OUTPUT_EXISTING case

Proposed Fix

Make the cleanup ownership-aware: pass the consuming task's slot to erase_task_outputs and erase a key only when the current mapping for that run still points at that slot, so a newer same-run producer survives. Add a same-run write-after-write regression test alongside the existing cross-run isolation test in tests/ut/cpp/hierarchical/test_tensormap.cpp.

Worth deciding separately whether two unordered OUTPUT writers to one key should be a diagnosed error rather than a silently unordered WAW.

Priority

Medium (minor risk, should fix in next few releases)

---

Related: #1388 (reader tracking for WAR edges), #1375 (WAIT / RETAIN dependency semantics).

Found while reviewing #1462; out of scope for that PR (the code landed in #1459).