vllm inference plugin

[dansola]

Why are the changes needed?

A vllm addition to the existing flytekitplugins-inference plugin which already has NIM and ollama.

What changes were proposed in this pull request?

A vllm plugin that lets you easily create a pod template to serve a vllm in an init container for a flyte task. User passes a hugging face secret name and the model in hugging face they want to serve.

import flytekit as fl
from openai import OpenAI

model_name = "google/gemma-2b-it"
hf_token_key = "vllm_hf_token"

vllm_args = {
    "model": model_name,
    "dtype": "half",
    "max-model-len": 2000,
}

hf_secrets = HFSecret(
    secrets_prefix="_FSEC_",
    hf_token_key=hf_token_key
)

vllm_instance = VLLM(
    hf_secret=hf_secrets,
    arg_dict=vllm_args
)

image = fl.ImageSpec(
    name="vllm_serve",
    registry="...",
    packages=["flytekitplugins-inference"],
)


@fl.task(
    pod_template=vllm_instance.pod_template,
    container_image=image,
    secret_requests=[
        fl.Secret(
            key=hf_token_key, mount_requirement=fl.Secret.MountType.ENV_VAR  # must be mounted as an env var
        )
    ],
)
def model_serving() -> str:
    client = OpenAI(
        base_url=f"{vllm_instance.base_url}/v1", api_key="vllm"  # api key required but ignored
    )

    completion = client.chat.completions.create(
        model=model_name,
        messages=[
            {
                "role": "user",
                "content": "Compose a haiku about the power of AI.",
            }
        ],
        temperature=0.5,
        top_p=1,
        max_tokens=1024,
    )
    return completion.choices[0].message.content

How was this patch tested?

Unit tests and running a remote workflow from the README.

Setup process

Screenshots

Check all the applicable boxes

• I updated the documentation accordingly.
• All new and existing tests passed.
• All commits are signed-off.

Related PRs

Docs link

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 76.45%. Comparing base (e24b9c1) to head (acb1639).
Report is 1 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #2967      +/-   ##
==========================================
+ Coverage   75.71%   76.45%   +0.73%     
==========================================
  Files         214      200      -14     
  Lines       21598    20922     -676     
  Branches     2693     2694       +1     
==========================================
- Hits        16352    15995     -357     
+ Misses       4489     4202     -287     
+ Partials      757      725      -32     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[Future-Outlier]

This is huge!

[samhita-alla]

Suggested change

	Initialize NIM class for managing a Kubernetes pod template.
	Initialize VLLM class for managing a Kubernetes pod template.

[samhita-alla]

lovely!