fix(launch): Fix race condition in agent thread clean up

[KyleGoyette]

Fixes

• Fixes WB-NNNNN
• Fixes #NNNN

Description

There was a race condition between the main agent loops update_finished function and the within thread, thread clean up. The agent could try to access a thread that was deleted in update_finished but this would lead to a KeyError and crash the whole agent.

The solution is to move all thread clean up out of the main agent loop, threads all clean up after themselves now.

🤖 Generated by Copilot at 2e337a1

Improved the launch agent logic and fixed a TypeError bug. The file wandb/sdk/launch/agent/agent.py was refactored to use a generator for thread ids and to simplify the job status checking and thread finishing.

Testing

How was this PR tested?

Checklist

• Include reference to internal ticket "Fixes WB-NNNNN" and/or GitHub issue "Fixes #NNNN" (if applicable)
• Ensure PR title compliance with the conventional commits standards

🤖 Generated by Copilot at 2e337a1

Sing, O Muse, of the skillful refactorer
Who tamed the agent logic with his code
And made a generator for thread ids
To spawn and join the workers as they strode.

[codecov]

Codecov Report

Merging #6352 (df6fdcd) into main (34c41a5) will decrease coverage by 0.04%.
The diff coverage is 100.00%.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #6352      +/-   ##
==========================================
- Coverage   77.24%   77.21%   -0.04%     
==========================================
  Files         387      387              
  Lines       44441    44437       -4     
==========================================
- Hits        34329    34310      -19     
- Misses      10059    10074      +15     
  Partials       53       53              

Flag	Coverage Δ
unittest	80.97% <100.00%> (-0.04%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

Files	Coverage Δ
wandb/sdk/launch/agent/agent.py	81.57% <100.00%> (+2.93%)	⬆️

... and 16 files with indirect coverage changes

[bcsherma]

Looks good, one nit

And I think we might want to acquire the lock before we read the job here https://github.com/wandb/wandb/blob/4a71959e0c9d440b8fdf5567c8db9ee6ca3101ad/wandb/sdk/launch/agent/agent.py#L285C51-L285C51 just to be safe