[tx] Add GPU memory benchmark script

[raulchen]

Summary

• Add GPU memory benchmark script for measuring peak GRAM usage across different batch sizes and sequence lengths
• Fix API server to exit when the background engine subprocess crashes (prevents hanging benchmarks)

Changes

New: benchmarks/benchmark_memory.py

A comprehensive benchmark tool for profiling GPU memory consumption:

• Test modes: sampling (inference), training (forward-backward), or both
• Sweep parameters: configurable batch sizes and sequence lengths
• Early termination: skips remaining batch sizes if OOM occurs
• GPU monitoring: polls nvidia-smi for peak memory usage
• Server lifecycle: automatic start/stop per test with log capture
• Output: CSV results, per-test server logs, optional XLA HLO dumps
• Server-only mode: launch server for manual testing/debugging

Fix: tx/tinker/api.py

The API server spawns a background engine subprocess. Previously, if the engine crashed, the API server stayed alive but coul
dn't process requests, causing benchmarks to hang indefinitely.

Fixed by adding a monitor task that:

1. Waits on the engine subprocess in a background asyncio task
2. Exits the API server immediately if the engine crashes unexpectedly
3. Distinguishes crash vs graceful shutdown via a flag

Test Plan

Will be used for benchmarking various optimizations for this issue #891

[gemini-code-assist]

Code Review

This pull request introduces a comprehensive GPU memory benchmark script and an important fix to prevent the API server from hanging when its background engine crashes. The new benchmark script is well-designed and feature-rich. The fix in the API server correctly addresses the hanging issue. My review focuses on making the server shutdown process more graceful, improving the robustness and error reporting of the new benchmark script, and suggesting a minor refactoring for better maintainability.

[gemini-code-assist]

If nvidia-smi is not found, the script silently fails to monitor GPU memory and reports 0 MiB, which can be misleading for a memory benchmark. It would be more robust to check for nvidia-smi's availability at the start of monitoring. If it's not present or fails to run, a warning should be printed to stderr and the monitoring thread should not be started. This makes the script's behavior more transparent to the user.

Suggested change

	def start(self) -> None:
	"""Start background GPU monitoring thread."""
	self._stop_event.clear()
	self._peak_memory = 0
	self._thread = threading.Thread(target=self._monitor_loop, daemon=True)
	self._thread.start()
	def start(self) -> None:
	"""Start background GPU monitoring thread."""
	# Check for nvidia-smi once at the beginning
	try:
	subprocess.run(["nvidia-smi"], capture_output=True, check=True, timeout=5.0)
	except (FileNotFoundError, subprocess.CalledProcessError, subprocess.TimeoutExpired):
	print("WARNING: `nvidia-smi` not found or failed to run. GPU memory usage will be reported as 0.", file=sys.stderr)
	# Don't start the thread if nvidia-smi is not available.
	# The stop() method will then return the initial peak_memory of 0.
	return
	self._stop_event.clear()
	self._peak_memory = 0
	self._thread = threading.Thread(target=self._monitor_loop, daemon=True)
	self._thread.start()

[gemini-code-assist]

When an exception occurs during a test (e.g., an out-of-memory error), the peak GPU memory is not recorded. The gpu_monitor has been running and likely captured the peak memory usage right before the failure, which is valuable information for debugging. You should record this peak memory value in the except block.

Suggested change

	result.status = "ERROR"
	result.peak_gpu_mem_mib = gpu_monitor.stop()

[gemini-code-assist]

Using os._exit(1) causes a hard exit, bypassing cleanup handlers (like finally blocks) and potentially leading to resource leaks or data corruption. A more graceful shutdown can be achieved by sending a SIGTERM signal to the process. uvicorn will catch this signal and shut down cleanly. Note that this will require importing the signal module.

Suggested change

	os._exit(1)
	os.kill(os.getpid(), signal.SIGTERM)

[gemini-code-assist]

Using pkill -f is a powerful command that can be risky on a shared machine. It could unintentionally terminate processes belonging to other users if their command line also contains tx.tinker.api. For a more targeted and safer approach, consider managing the server's process ID (PID) using a PID file. The server could write its PID to a file on startup, and this script could then read that file to kill the specific process. While the current implementation is likely fine for a controlled benchmark environment, this change would make it more robust for shared systems.

[gemini-code-assist]

The manual creation of config_dict for JSON serialization is verbose and duplicates the structure of the BenchmarkConfig dataclass. This can lead to inconsistencies if BenchmarkConfig is updated but this dictionary is not. A more maintainable approach is to use dataclasses.asdict to convert the config object to a dictionary, then filter and format it for saving. This ensures the saved configuration always reflects the BenchmarkConfig definition.

Suggested change

	# Save config to output directory
	config_dict = {
	"base_model": config.base_model,
	"tp_size": config.tp_size,
	"max_lora_adapters": config.max_lora_adapters,
	"gradient_checkpointing": config.gradient_checkpointing,
	"extra_backend_config": config.extra_backend_config,
	"test_mode": config.test_mode,
	"batch_sizes": config.batch_sizes,
	"seq_lens": config.seq_lens,
	"server_only": config.server_only,
	"output_root": str(config.output_root),
	"xla_preallocate": config.xla_preallocate,
	"gpu_allocator": config.gpu_allocator,
	"jax_log_compiles": config.jax_log_compiles,
	"dump_xla": config.dump_xla,
	"timestamp": datetime.now().isoformat(),
	}
	with open(config.output_dir / "config.json", "w") as f:
	# Save config to output directory
	import dataclasses
	config_dict = dataclasses.asdict(config)
	# Filter and format for JSON serialization
	keys_to_save = [
	"base_model", "tp_size", "max_lora_adapters", "gradient_checkpointing",
	"extra_backend_config", "test_mode", "batch_sizes", "seq_lens",
	"server_only", "output_root", "xla_preallocate", "gpu_allocator",
	"jax_log_compiles", "dump_xla"
	]
	json_config = {k: config_dict[k] for k in keys_to_save if k in config_dict}
	json_config["output_root"] = str(json_config["output_root"])
	json_config["timestamp"] = datetime.now().isoformat()
	with open(config.output_dir / "config.json", "w") as f:
	json.dump(json_config, f, indent=2)

[raulchen]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces a comprehensive GPU memory benchmark script and a crucial fix to the API server to prevent it from hanging when the background engine crashes. The benchmark script is well-structured and feature-rich, covering various test modes and parameters. The fix in the API server correctly uses an asyncio task to monitor the engine subprocess and gracefully shut down the server on a crash.

My review includes a few suggestions to improve maintainability and adhere to modern Python best practices. Specifically, I've recommended moving a local import to the top level in the benchmark script, simplifying configuration object creation for better maintainability, and using the more recent asyncio.to_thread function in the API server for cleaner asynchronous code.

[gemini-code-assist]

For better code organization and readability, it's recommended to place all imports at the top of the file. Please move import httpx to the top-level imports section of this file (e.g., around line 71). This makes dependencies clearer and avoids potential issues with delayed import errors.

[gemini-code-assist]

The instantiation of BenchmarkConfig is verbose as it manually lists all arguments from args. This can be simplified and made more maintainable by creating a dictionary from vars(args) and unpacking it. This way, if you add new arguments to parse_args and BenchmarkConfig, they will be handled automatically without needing to update this section.

    # Build configuration with derived paths
    config_data = {k: v for k, v in vars(args).items() if hasattr(BenchmarkConfig, k)}
    config = BenchmarkConfig(
        **config_data,
        output_dir=output_dir,
        db_path=output_dir / "tinker.db",
        csv_path=output_dir / "results.csv",
    )

[gemini-code-assist]

Using loop.run_in_executor is correct for running blocking code in an async application. However, since Python 3.9, asyncio.to_thread provides a higher-level and more convenient API for this purpose. It's generally recommended to use asyncio.to_thread for better readability and conciseness.

Suggested change

	loop = asyncio.get_event_loop()
	exit_code = await loop.run_in_executor(None, background_engine.wait)
	exit_code = await asyncio.to_thread(background_engine.wait)