add nim plugin

[samhita-alla]

Tracking issue

Fixes flyteorg/flyte#5478

Why are the changes needed?

This PR adds NIM to flytekit plugins to enable serving optimized model containers, which can include NVIDIA CUDA software, NVIDIA Triton Inference Server, and NVIDIA TensorRT-LLM software. Since each Flyte task runs in a Kubernetes pod, a NIM container can be easily deployed as a sidecar service, allowing users to invoke the model's endpoint as if it is hosted locally, thus minimizing network overhead. Additionally, Flyte can manage data pre-processing and post-processing, facilitating the creation of end-to-end batch inference pipelines with optimized models.

What changes were proposed in this pull request?

• A new ModelInferenceTemplate has been added, providing a base for creating other inference solutions like Ollama and Hugging Face.
• The NIM class, which subclasses ModelInferenceTemplate, has been introduced as well.

from flytekit import ImageSpec, Secret, task, Resources
from flytekit.extras.accelerators import A10G
from flytekitplugins.inference import NIM, NIMSecrets
from openai import OpenAI

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

nim_instance = NIM(
    image="nvcr.io/nim/meta/llama3-8b-instruct:1.0.0",
    secrets=NIMSecrets(
        ngc_image_secret="nvcrio-cred",
        ngc_secret_key=NGC_KEY,
        secrets_prefix="_FSEC_",
    ),
)


@task(
    container_image=image,
    pod_template=nim_instance.pod_template,
    accelerator=A10G,
    secret_requests=[
        Secret(
            key="ngc_api_key", mount_requirement=Secret.MountType.ENV_VAR
        )  # must be mounted as an env var
    ],
    requests=Resources(gpu="0"),
)
def model_serving() -> str:
    client = OpenAI(
        base_url=f"{nim_instance.base_url}/v1", api_key="nim"
    )  # api key required but ignored

    completion = client.chat.completions.create(
        model="meta/llama3-8b-instruct",
        messages=[
            {
                "role": "user",
                "content": "Write a limerick about the wonders of GPU computing.",
            }
        ],
        temperature=0.5,
        top_p=1,
        max_tokens=1024,
    )

    return completion.choices[0].message.content

How was this patch tested?

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

Attention: Patch coverage is 83.09859% with 12 lines in your changes missing coverage. Please review.

Project coverage is 76.21%. Comparing base (097e9e8) to head (7e62555).
Report is 4 commits behind head on master.

Files	Patch %	Lines
flytekit/core/inference.py	80.39%	5 Missing and 5 partials ⚠️
flytekit/core/utils.py	90.00%	1 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #2475      +/-   ##
==========================================
- Coverage   76.22%   76.21%   -0.01%     
==========================================
  Files         187      186       -1     
  Lines       18938    18903      -35     
  Branches     3706     3719      +13     
==========================================
- Hits        14435    14407      -28     
+ Misses       3870     3851      -19     
- Partials      633      645      +12     

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

[kumare3]

How about we move this pattern to core flytekit. I don't see any extra requirements?

[pingsutw]

It looks awesome! I think having ModelInferenceTemplate is really great, so users don't need to deal with the K8s pod spec

[eapolinario]

This is looking good. I have only minor comments.

[eapolinario]

one thing to consider is what sort of inference plugins we're going to add to flytekitplugins-inference. Right now there's no extra dependencies (since NIM is essentially a pod template wrapper), but that might not be the case in all plugins. That said, we'll cross that bridge when we get there.

[eapolinario]

the plugin doesn't really require openai, right?

[samhita-alla]

what sort of inference plugins we're going to add to flytekitplugins-inference

I've got ollama working with the sidecar pattern, so will be adding that next. in flytekitplugins-inference, i believe we need to stick to serving plugins that can be implemented with the sidecar pattern. if not, they should be standalone plugins, such as the sagemaker deployment agent.

this parent plugin should essentially capture inference providers that allow self-hosting and serving models within a flyte task. the ideal goal is to make it work with actors to enable near real-time inference.

Right now there's no extra dependencies (since NIM is essentially a pod template wrapper), but that might not be the case in all plugins.

if we derive from the sidecar pattern, we don't really need to worry about extra dependencies. a typical approach is for the plugin to use a docker image to serve models.

the plugin doesn't really require openai, right?

you're right. however, since most inference providers now support the openai api, it'd be beneficial to include it as part of the plugin.

[eapolinario]

Thank you. This looks great.