Add torch profiler support

[akaitsuki-ii]

No description provided.

[gemini-code-assist]

Code Review

This pull request introduces a TorchProfiler utility and integrates it into the engine and worker modules to enable distributed profiling. The engine's worker communication logic was refactored to support multi-rank responses. Feedback suggests improving the portability of the trace compression by checking for the gzip utility, ensuring returned file paths accurately reflect whether compression succeeded, and removing a redundant variable initialization in the worker loop.

[gemini-code-assist]

The background compression relies on the system gzip utility, which may not be available in all environments (e.g., Windows or minimal containers). If gzip is missing, subprocess.Popen will raise a FileNotFoundError. While this is caught, the return values of start and stop will still incorrectly indicate a .gz extension. Consider checking for gzip availability using shutil.which and adjusting the returned file path accordingly, or use the built-in gzip module in a separate thread for better portability.

Suggested change

	try:
	subprocess.Popen(["gzip", "-f", json_file])
	logger.info(f"[Rank {rank}] Triggered background compression for {json_file}")
	# Update variable to point to the eventual file
	json_file = f"{json_file}.gz"
	except Exception as compress_err:
	logger.warning(f"[Rank {rank}] Background gzip failed to start: {compress_err}")
	import shutil
	if shutil.which("gzip"):
	try:
	subprocess.Popen(["gzip", "-f", json_file])
	logger.info(f"[Rank {rank}] Triggered background compression for {json_file}")
	json_file = f"{json_file}.gz"
	except Exception as compress_err:
	logger.warning(f"[Rank {rank}] Background gzip failed to start: {compress_err}")
	else:
	logger.warning(f"[Rank {rank}] gzip utility not found, skipping compression")

[gemini-code-assist]

The method returns a .gz path regardless of whether compression actually started or succeeded. This can lead to issues if the caller expects the file to exist at the returned path. It would be more robust to return the actual path based on the success of the compression trigger.

[gemini-code-assist]

The variable should_reply is initialized here but is immediately overwritten a few lines later (line 118) before being used. This initialization is redundant and can be removed to improve clarity.

Suggested change

	should_reply = rank == 0
	try: