fix: flush asynchronous JSONL logs on training end

[Alwin4Zhang]

What does this PR do?

• Add flush() and close() lifecycle methods to JsonlWriter.
• Wait for pending asynchronous writes before shutting down the writer thread.
• Propagate asynchronous write exceptions back to the caller.
• Close the Megatron logging writer when training ends.
• Add a regression test covering asynchronous writer shutdown.

Motivation

Megatron writes logging.jsonl through a daemon thread. The writer was not flushed or closed when training finished, so the process could exit before pending log entries were written, resulting in incomplete logging files.

Tests

• Added test_async_jsonl_close to verify that all queued records are persisted before close() returns.
• Verified that appending after the writer is closed raises RuntimeError.
• Passed Python compilation and git diff --check.

[gemini-code-assist]

Code Review

This pull request introduces proper lifecycle management for JsonlWriter by adding flush and close methods, implementing asynchronous closing with a worker thread sentinel, and raising errors when appending to a closed writer. It also updates the training callback to close the writer at the end of training and adds corresponding unit tests. The review feedback suggests several robustness improvements: adding a null check for self.jsonl_writer in the callback to prevent potential AttributeError crashes, preserving only the first background worker exception to avoid overwriting the root cause, and checking for worker exceptions during append to enable fail-fast behavior.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

If on_train_begin fails or is not called before on_train_end, self.jsonl_writer will remain None, causing an AttributeError when calling close(). Adding a null check prevents this potential crash during cleanup.

Suggested change

	self.jsonl_writer.close()
	if self.jsonl_writer is not None:
	self.jsonl_writer.close()

[gemini-code-assist]

Storing only the first exception that occurs in the background worker thread is a best practice. This ensures that the root cause of the failure is preserved and reported, rather than being overwritten by subsequent errors.

Suggested change

	except Exception as e:
	self._worker_exception = e
	except Exception as e:
	if self._worker_exception is None:
	self._worker_exception = e

[gemini-code-assist]

Checking self._worker_exception inside append allows the writer to fail-fast. If a background write fails (e.g., due to disk full), subsequent calls to append will raise the exception immediately rather than silently queueing items in memory until flush or close is called.

        if self._closed:
            raise RuntimeError('Cannot append to a closed JsonlWriter')
        if self._worker_exception is not None:
            raise self._worker_exception