Memory estimate fixes

[magdyksaleh]

• Update mem wiggle room to 0.9 (from 0.8)
• Ensure free_memory is never negative
• Lazy load graph for larger compile batch sizes
• Add new argument compile_batch_size (default 32) which is the batch size we will initially use for cuda compilation of models
• make the precommit use ruff instead of flake8

[magdyksaleh]

lorax/server/lorax_server/utils/graph.py

Lines 478 to 498 in 12e530a

	def can_use_graph(
	self,
	batch: "FlashCausalLMBatch",
	adapter_data: AdapterBatchData,
	) -> bool:
	ranks = adapter_data.ranks()
	nranks = len(ranks)
	max_rank = max(ranks) if len(ranks) > 0 else 0
	batch_size = batch.input_ids.shape[0]
	max_s = batch.max_current_length
	# TODO(travis): allow using CUDA graphs with multi-rank batches
	return (
	torch.cuda.is_available()
	and batch_size <= MAX_BATCH_SIZE
	and max_s <= self.max_total_tokens
	and max_rank <= MAX_RANK
	and nranks <= 1
	and max_rank in _allowed_ranks
	)

This ensures that we stay within the range

[tgaddair]

LGTM!

[tgaddair]

Actually, I think we need changes to warmup.

[tgaddair]

LGTM!