vllm_backend

vllm_backend

vLLM-based generation backend for tabular data synthesis.

Classes:

Name	Description
VllmBackend	Generation backend using vLLM for high-throughput inference.

VllmBackend(config, model_metadata, workdir, **kwargs)

Bases: GeneratorBackend

Generation backend using vLLM for high-throughput inference.

Loads the base model with a LoRA adapter via vLLM and generates synthetic records in batches. Supports optional structured generation (regex or JSON schema) to constrain outputs.

LoRARequest("lora", 1, str(adapter_path)) is passed to llm.generate when an adapter is available. The vLLM engine uses config.training.lora_r as max_lora_rank.

Parameters:

Name	Type	Description	Default
config	SafeSynthesizerParameters	Pipeline configuration.	required
model_metadata	ModelMetadata	Model metadata (prompt template, adapter path, sequence length, etc.).	required
workdir	Workdir	Working directory containing the adapter and schema.	required
**kwargs		Additional options. use_detailed_logs (bool) enables verbose error messages (disabled by default to avoid leaking sensitive data).	{}

Methods:

Name	Description
teardown	Release GPU memory and distributed resources. Idempotent -- safe to call multiple times.
initialize	Initialize and load the model into memory.
prepare_params	Parse parameters and configure the generation method.
generate	Generate synthetic tabular data in batches until the target count is reached.

Source code in src/nemo_safe_synthesizer/generation/vllm_backend.py

def __init__(
    self,
    config: SafeSynthesizerParameters,
    model_metadata: ModelMetadata,
    workdir: Workdir,
    **kwargs,
):
    self.model_metadata = model_metadata
    self.config = config
    self.remote = False
    self.workdir = workdir
    self.schema = load_json(self.workdir.schema_file)
    self.columns = list(self.schema["properties"].keys())
    self.prompt = utils.create_schema_prompt(
        self.columns,
        instruction=self.model_metadata.instruction,
        prompt_template=self.model_metadata.prompt_config.template,
    )
    self.llm: vLLM | None = None
    self._prompt_token_count: int | None = None
    # Populated in ``initialize()`` after engine build; pre-declared
    # here so ``generate()`` can always read it (tests that mock the
    # backend without calling ``initialize()`` see the empty dict).
    self._engine_runtime_config: dict[str, Any] = {}

    # Do not generate detailed error messages in production to avoid leaking sensitive data.
    self.use_detailed_logs = kwargs.pop("use_detailed_logs", False)
    self.gen_method: partial | None = None
    self._gen_method: partial | None = None
    # Initial processor without a tokenizer; replaced in ``initialize()`` with a
    # tokenizer-aware processor once the vLLM engine (and its tokenizer) exists.
    # This lets callers introspect the processor type before ``initialize()``,
    # at the cost of token counts being zero until the tokenizer is attached.
    self.processor: Processor = create_processor(self.schema, self.model_metadata, self.config)
    adapter_path = self.workdir.adapter_path if self.workdir.adapter_path else self.model_metadata.adapter_path
    self.lora_req = LoRARequest("lora", 1, str(adapter_path)) if adapter_path else None
    self._torn_down = False

teardown()

Release GPU memory and distributed resources. Idempotent -- safe to call multiple times.

Source code in src/nemo_safe_synthesizer/generation/vllm_backend.py

def teardown(self) -> None:
    """Release GPU memory and distributed resources. Idempotent -- safe to call multiple times."""
    if self._torn_down:
        return
    self._torn_down = True

    try:
        cleanup_dist_env_and_memory()
    except Exception:
        logger.debug("cleanup_dist_env_and_memory failed during teardown", exc_info=True)

    self.llm = None
    self._gen_method = None
    self.gen_method = None

    try:
        cleanup_memory()
    except Exception:
        logger.debug("cleanup_memory failed during teardown", exc_info=True)

initialize(**kwargs)

Initialize and load the model into memory.

Creates the vLLM engine and then builds the record processor with the engine's tokenizer so that exact token counts are available during generation.

Source code in src/nemo_safe_synthesizer/generation/vllm_backend.py

def initialize(self, **kwargs) -> None:
    """Initialize and load the model into memory.

    Creates the vLLM engine and then builds the record processor
    with the engine's tokenizer so that exact token counts are
    available during generation.
    """
    self._torn_down = False

    # vLLM 0.12+ accepts attention_config as a constructor arg (replaces the
    # VLLM_ATTENTION_BACKEND env var used in 0.11.x).
    attn_backend = self.config.generation.attention_backend
    attention_config = {"backend": attn_backend} if attn_backend not in (None, "auto") else None

    max_vram = get_max_vram()
    # note this only works for single GPU setups
    max_vram = max_vram.get(0, 0.8)

    # vllm requires this "config" to set the backend ahead of time.
    structured_outputs_config = StructuredOutputsConfig(
        backend=self.config.generation.structured_generation.backend,
    )
    model_ref = ModelRef.parse(self.config.training.pretrained_model)

    with heartbeat("Model loading", logger_name=__name__, model=self.config.training.pretrained_model):
        self.llm = vLLM(
            model=model_ref.target(),
            gpu_memory_utilization=max_vram,
            enable_lora=True,
            max_lora_rank=self.config.training.lora_r,
            structured_outputs_config=structured_outputs_config,
            attention_config=attention_config,
            trust_remote_code=model_ref.trust_remote_code,
        )

    # Cache the engine's *effective* runtime config once at init. Read by
    # ``generate()`` to populate the generation-complete observability
    # event; used by the benchmark harness (and any other intent-comparing
    # caller) to detect "flag didn't engage" mismatches against what was
    # asked for.
    self._engine_runtime_config = probe_engine_runtime_config(self.llm)

    # vLLM's get_tokenizer() returns a wider union than HF's PreTrainedTokenizerBase;
    # in practice it's always a HF tokenizer subclass, so cast for the processor.
    tokenizer = cast(PreTrainedTokenizerBase, self.llm.get_tokenizer())
    self.processor = create_processor(
        self.schema,
        self.model_metadata,
        self.config,
        tokenizer=tokenizer,
    )

prepare_params(**kwargs)

Parse parameters and configure the generation method.

Parses a dictionary of parameters into SamplingParams, applying necessary transformations from the Safe Synthesizer API to vLLM's API. num_beams is mapped to beam_width only when greater than 1; otherwise it is omitted.

Parameters:

Name	Type	Description	Default
**kwargs		Sampling parameters to configure.	{}

Source code in src/nemo_safe_synthesizer/generation/vllm_backend.py

def prepare_params(self, **kwargs) -> None:
    """Parse parameters and configure the generation method.

    Parses a dictionary of parameters into ``SamplingParams``, applying
    necessary transformations from the Safe Synthesizer API to vLLM's API.
    ``num_beams`` is mapped to ``beam_width`` only when greater than 1;
    otherwise it is omitted.

    Args:
        **kwargs: Sampling parameters to configure.
    """
    structured_output_params = self._build_structured_output_params()
    kwargs |= {"structured_outputs": structured_output_params}

    resolved_temperature = self._resolve_temperature(kwargs)
    api_mapping = self._get_api_param_mapping(resolved_temperature)
    sampling_params = self._transform_kwargs_to_sampling_params(kwargs, api_mapping)

    real_params = SamplingParams(**sampling_params)
    logger.debug(f"SamplingParams: {real_params!r}")

    # Create a partially parametrized version of the underlying vllm.LLM.generate
    # method that is immediately callable downstream.
    if self.llm is None:
        raise InternalError(
            "VllmBackend._configure_sampling_params() called before initialize() -- self.llm is None."
        )
    self._gen_method = partial(
        self.llm.generate,
        sampling_params=real_params,
        lora_request=self.lora_req,
        # Show vLLM's tqdm progress bar only when debug logging is enabled.
        use_tqdm=logger.isEnabledFor(logging.DEBUG),
    )

generate(data_actions_fn=None)

Generate synthetic tabular data in batches until the target count is reached.

Iterates over generation batches, applying the processor to each LLM output, until the configured num_records target is met or a stopping condition fires.

Non-tabular processors need BOS/EOS delimiters in the raw text, so generation keeps special tokens for those processors and strips them only for TabularDataProcessor. Native EOS stopping remains enabled through ignore_eos=False.

Wrapped in an :class:NvmlPeakSampler context so a generation-complete observability event is emitted at end of the call carrying peak device VRAM, host loadavg pre/post, vLLM's kv_cache_usage_perc / prefix_cache_hit_rate / spec_accept_rate (read at end-of-generation), and the engine's effective runtime config probed at engine-init time. Sampler is degraded-mode when NVML is unavailable; the event still emits with peak_vram_gb=None.

Parameters:

Name	Type	Description	Default
data_actions_fn	DataActionsFn | None	Optional post-processing / validation function applied to each batch of generated records.	None

Returns:

Type	Description
GenerateJobResults	Results containing the generated DataFrame and statistics.

Source code in src/nemo_safe_synthesizer/generation/vllm_backend.py

def generate(
    self,
    data_actions_fn: utils.DataActionsFn | None = None,
) -> GenerateJobResults:
    """Generate synthetic tabular data in batches until the target count is reached.

    Iterates over generation batches, applying the processor to each
    LLM output, until the configured ``num_records`` target is met or
    a stopping condition fires.

    Non-tabular processors need BOS/EOS delimiters in the raw text, so
    generation keeps special tokens for those processors and strips them
    only for ``TabularDataProcessor``. Native EOS stopping remains enabled
    through ``ignore_eos=False``.

    Wrapped in an :class:`NvmlPeakSampler` context so a
    generation-complete observability event is emitted at end of
    the call carrying peak device VRAM, host loadavg pre/post, vLLM's
    kv_cache_usage_perc / prefix_cache_hit_rate / spec_accept_rate
    (read at end-of-generation), and the engine's effective runtime
    config probed at engine-init time. Sampler is degraded-mode when
    NVML is unavailable; the event still emits with ``peak_vram_gb=None``.

    Args:
        data_actions_fn: Optional post-processing / validation function
            applied to each batch of generated records.

    Returns:
        Results containing the generated DataFrame and statistics.
    """
    loadavg_pre = read_loadavg()
    sampler = NvmlPeakSampler()
    try:
        with sampler:
            self._run_generation(data_actions_fn)
    finally:
        # Emit regardless of success: the sampler thread has been joined
        # by ``with`` exit, so ``sampler.peak_gb`` is the final peak, and
        # the event carries whatever measurements were captured up to a
        # failure point.
        self._emit_generation_observability(sampler, loadavg_pre)

    return self.gen_results