add NVIDIA NIM inference adapter

[mattf]

What does this PR do?

this PR adds a basic inference adapter to NVIDIA NIMs

what it does -

• chat completion api
  • tool calls
  • streaming
  • structured output
  • logprobs
• support hosted NIM on integrate.api.nvidia.com
• support downloaded NIM containers

what it does not do -

• completion api
• embedding api
• vision models
• builtin tools
• have certainty that sampling strategies are correct

Feature/Issue validation/testing/test plan

pytest -s -v --providers inference=nvidia llama_stack/providers/tests/inference/ --env NVIDIA_API_KEY=...

all tests should pass. there are pydantic v1 warnings.

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes?
• Did you write any new necessary tests?

Thanks for contributing 🎉!

[ashwinb]

this should be a dynamic import within get_adapter_impl() -- we want configs to be manipulated without needing implementation dependencies.

[mattf]

ptal

[ashwinb]

Thank you for this PR. So good!

Re: testing, I'd like to have a reproducible e2e test (ala what we have in providers/tests/inference/test_text_inference.py and providers/tests/inference/test_vision_inference.py) -- just having an nvidia specific fixture there which could then be invoked as

pytest -s -v --providers inference=nvidia test_text_inference.py --env ...

would be great.

[ashwinb]

this could be is_nvidia_hosted perhaps?

[mattf]

it's really an internal thing. i've removed it from the NVIDIAConfig api entirely.

[ashwinb]

this is a nitpick, you can ignore it you feel strongly. we don't usually do underscores in files in the repo - at least not yet. we don't even strongly enforce what symbols get exported out a module (that part is a bit sad admittedly.) could you make the files not have starting underscores?

[mattf]

my inclination is to be cautious about the exported symbols, but it's important to be cohesive w/ the project. i'll change these. ptal.

[ashwinb]

thanks for the explicit imports. we will be code-modding all our other code to do this sane thing soon :)

[mattf]

i spent so much time trying to figure out which classes were coming from which packages 😆

[ashwinb]

is there a "base" llama model this model would correspond most closely with? we like to know it because we try to format tools, etc. in a way which the model will work best with. this isn't strictly necessary if the provider / API works very robustly with tool calling, etc. but so far given our experience with various "openai" wrapper APIs, it has been spotty.

[mattf]

nvidia/llama-3.1-nemotron-51b-instruct (typo in my comment) is https://build.nvidia.com/nvidia/llama-3_1-nemotron-51b-instruct/modelcard

there's now a 70b variant at https://build.nvidia.com/nvidia/llama-3_1-nemotron-70b-instruct/modelcard

[ashwinb]

any chance this could be done? it's OK if not, but we have gone back and filled up many of the missing completion() methods also now

[mattf]

let me come back and add it in another PR, same for embedding

[aidando73]

@ashwinb - Are we planning to support that endpoint long-term?

It seems the industry is moving away from it. E.g., OpenAI marked their completions api as legacy https://platform.openai.com/docs/api-reference/completions and Anthropic too: https://docs.anthropic.com/en/api/complete

I'm wondering if it's a good idea to move away from the endpoint as well while llama-stack is still early days.

Otherwise might be painful to maintain it for a long time + deprecate + remove

[aidando73]

Reposting response here for visibility:

Thanks for all your contributions @aidan Do -- we want to keep supporting completions at least for now because we believe having raw access to a model is as important. Unlike other providers, Llamas are open source and people play with them and iterate in a variety of ways. The kinds of manipulations we do with a chat_completion endpoint internally may not be what users intend sometimes. Sometimes, they just want a carefully formatted prompt to hit the model directly.

On that theme, I think it would be great if Groq could build completions endpoints on their end too. But until that time, NotImplementedError() would have to do.

[ashwinb]

❤️

[mattf]

Thank you for this PR. So good!

Re: testing, I'd like to have a reproducible e2e test (ala what we have in providers/tests/inference/test_text_inference.py and providers/tests/inference/test_vision_inference.py) -- just having an nvidia specific fixture there which could then be invoked as

pytest -s -v --providers inference=nvidia test_text_inference.py --env ...

would be great.

i've updated the PR description to note that this does not cover structured output, vision models, embedding or completion apis.

if it's ok, i'll follow up with PRs to add those features.

➜ pytest -s -v --providers inference=nvidia llama_stack/providers/tests/inference/test_{text,vision}_inference.py --env NVIDIA_API_KEY=... --inference-model Llama3.1-8B-Instruct
/home/matt/.conda/envs/stack/lib/python3.10/site-packages/pytest_asyncio/plugin.py:208: PytestDeprecationWarning: The configuration option "asyncio_default_fixture_loop_scope" is unset.
The event loop scope for asynchronous fixtures will default to the fixture caching scope. Future versions of pytest-asyncio will default the loop scope for asynchronous fixtures to function scope. Set the default fixture loop scope explicitly in order to avoid unexpected behavior in the future. Valid fixture loop scopes are: "function", "class", "module", "package", "session"

  warnings.warn(PytestDeprecationWarning(_DEFAULT_FIXTURE_LOOP_SCOPE_UNSET))
===================================================================================== test session starts ======================================================================================
platform linux -- Python 3.10.15, pytest-8.3.3, pluggy-1.5.0 -- /home/matt/.conda/envs/stack/bin/python
cachedir: .pytest_cache
rootdir: /home/matt/Documents/Repositories/meta-llama/llama-stack
configfile: pyproject.toml
plugins: anyio-4.6.2.post1, asyncio-0.24.0, httpx-0.34.0
asyncio: mode=strict, default_loop_scope=None
collected 11 items                                                                                                                                                                             

llama_stack/providers/tests/inference/test_text_inference.py::TestInference::test_model_list[-nvidia] Resolved 4 providers
 inner-inference => nvidia
 models => __routing_table__
 inference => __autorouted__
 inspect => __builtin__

Initializing NVIDIAInferenceAdapter(https://integrate.api.nvidia.com)...
Models: Llama3.1-8B-Instruct served by nvidia

PASSED
llama_stack/providers/tests/inference/test_text_inference.py::TestInference::test_completion[-nvidia] SKIPPED (Other inference providers don't support completion() yet)
llama_stack/providers/tests/inference/test_text_inference.py::TestInference::test_completions_structured_output[-nvidia] SKIPPED (This test is not quite robust)
llama_stack/providers/tests/inference/test_text_inference.py::TestInference::test_chat_completion_non_streaming[-nvidia] PASSED
llama_stack/providers/tests/inference/test_text_inference.py::TestInference::test_structured_output[-nvidia] SKIPPED (Other inference providers don't support structured output yet)
llama_stack/providers/tests/inference/test_text_inference.py::TestInference::test_chat_completion_streaming[-nvidia] PASSED
llama_stack/providers/tests/inference/test_text_inference.py::TestInference::test_chat_completion_with_tool_calling[-nvidia] PASSED
llama_stack/providers/tests/inference/test_text_inference.py::TestInference::test_chat_completion_with_tool_calling_streaming[-nvidia] PASSED
llama_stack/providers/tests/inference/test_vision_inference.py::TestVisionModelInference::test_vision_chat_completion_non_streaming[-nvidia-image0-expected_strings0] SKIPPED (Other...)
llama_stack/providers/tests/inference/test_vision_inference.py::TestVisionModelInference::test_vision_chat_completion_non_streaming[-nvidia-image1-expected_strings1] SKIPPED (Other...)
llama_stack/providers/tests/inference/test_vision_inference.py::TestVisionModelInference::test_vision_chat_completion_streaming[-nvidia] SKIPPED (Other inference providers don't su...)

========================================================================== 5 passed, 6 skipped, 10 warnings in 29.23s ==========================================================================

[mattf]

@ashwinb i find test_structured_output to be flakey. it's both a functionality and accuracy test -

        answer = AnswerFormat.model_validate_json(response.completion_message.content)
        assert answer.first_name == "Michael"
        assert answer.last_name == "Jordan"
        assert answer.year_of_birth == 1963
        assert answer.num_seasons_in_nba == 15

it's an accuracy test because it checks the value of first/last name, birth year, and num seasons.

i find that -

• llama-3.1-8b-instruct and llama-3.2-3b-instruct pass the functionality portion
• llama-3.2-3b-instruct consistently fails the accuracy portion (thinking MJ was in the NBA for 14 seasons)
• llama-3.1-8b-instruct occasionally fails the accuracy portion

suggestions (not mutually exclusive) -

1. turn the test into functionality only, skip the value checks
2. split the test into a functionality version and an xfail accuracy version
3. add context to the prompt so the llm can answer without accessing embedded memory

[ashwinb]

@mattf I agree with your comments on test_structured_output completely. I think the third option makes the most sense. I will update the test pronto.

[ashwinb]

Looks good to me. Merging!

[aidando73]

I'm implementing a Groq adapter and wondering this as well

[aidando73]

This is such a cool trick