diff --git a/README.md b/README.md index b5eb6a33..0f48cf8d 100644 --- a/README.md +++ b/README.md @@ -250,6 +250,26 @@ ollama.embed(model='gemma3', input=['The sky is blue because of rayleigh scatter ollama.ps() ``` +### GPU Selection + +`generate()` and `chat()` accept a `num_gpu` argument (number of model layers to offload to the GPU), merged into `options` alongside any other options you pass: + +```python +ollama.generate(model='gemma3', prompt='Why is the sky blue?', num_gpu=1) +``` + +For hard isolation across multiple GPUs on a shared server (e.g. pinning separate notebooks/processes to different physical GPUs), run one `ollama serve` process per GPU, each with its own `CUDA_VISIBLE_DEVICES`, and point a separate `Client(host=...)` at each: + +```python +# Terminal 1: CUDA_VISIBLE_DEVICES=0 OLLAMA_HOST=127.0.0.1:11434 ollama serve +# Terminal 2: CUDA_VISIBLE_DEVICES=1 OLLAMA_HOST=127.0.0.1:11435 ollama serve + +from ollama import Client + +gpu0_client = Client(host='http://127.0.0.1:11434') +gpu1_client = Client(host='http://127.0.0.1:11435') +``` + ## Errors Errors are raised if requests return an error status or if an error is detected while streaming. diff --git a/examples/README.md b/examples/README.md index 1df713ea..3405cd2d 100644 --- a/examples/README.md +++ b/examples/README.md @@ -102,6 +102,10 @@ Configuration to use with an MCP client: - [ps.py](ps.py) +### GPU Selection - Offload a request to the GPU with num_gpu + +- [gpu-selection.py](gpu-selection.py) + ### Ollama Pull - Pull a model from Ollama Requirement: `pip install tqdm` diff --git a/examples/gpu-selection.py b/examples/gpu-selection.py new file mode 100644 index 00000000..5cfd6c0f --- /dev/null +++ b/examples/gpu-selection.py @@ -0,0 +1,16 @@ +from ollama import Client + +# num_gpu controls how many model layers are offloaded to the GPU for this request. +client = Client() +response = client.generate('gemma3', 'Why is the sky blue?', num_gpu=1) +print(response['response']) + +# For hard isolation across multiple GPUs (e.g. pinning separate processes to +# different physical GPUs on a shared server), run one `ollama serve` per GPU +# with its own CUDA_VISIBLE_DEVICES and point a separate Client(host=...) at each: +# +# CUDA_VISIBLE_DEVICES=0 OLLAMA_HOST=127.0.0.1:11434 ollama serve +# CUDA_VISIBLE_DEVICES=1 OLLAMA_HOST=127.0.0.1:11435 ollama serve +# +# gpu0_client = Client(host='http://127.0.0.1:11434') +# gpu1_client = Client(host='http://127.0.0.1:11435') diff --git a/ollama/_client.py b/ollama/_client.py index 8dfce824..d9c0eded 100644 --- a/ollama/_client.py +++ b/ollama/_client.py @@ -220,6 +220,7 @@ def generate( width: Optional[int] = None, height: Optional[int] = None, steps: Optional[int] = None, + num_gpu: Optional[int] = None, ) -> GenerateResponse: ... @overload @@ -244,6 +245,7 @@ def generate( width: Optional[int] = None, height: Optional[int] = None, steps: Optional[int] = None, + num_gpu: Optional[int] = None, ) -> Iterator[GenerateResponse]: ... def generate( @@ -267,10 +269,17 @@ def generate( width: Optional[int] = None, height: Optional[int] = None, steps: Optional[int] = None, + num_gpu: Optional[int] = None, ) -> Union[GenerateResponse, Iterator[GenerateResponse]]: """ Create a response using the requested model. + Args: + num_gpu: Number of layers to offload to the GPU. Merged into `options`. + For hard isolation across multiple GPUs, run a separate `ollama serve` + per GPU (each with its own `CUDA_VISIBLE_DEVICES`) and point a separate + `Client(host=...)` at each. + Raises `RequestError` if a model is not provided. Raises `ResponseError` if the request could not be fulfilled. @@ -296,7 +305,7 @@ def generate( raw=raw, format=format, images=list(_copy_images(images)) if images else None, - options=options, + options=_merge_options(options, num_gpu=num_gpu), keep_alive=keep_alive, width=width, height=height, @@ -319,6 +328,7 @@ def chat( format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, options: Optional[Union[Mapping[str, Any], Options]] = None, keep_alive: Optional[Union[float, str]] = None, + num_gpu: Optional[int] = None, ) -> ChatResponse: ... @overload @@ -335,6 +345,7 @@ def chat( format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, options: Optional[Union[Mapping[str, Any], Options]] = None, keep_alive: Optional[Union[float, str]] = None, + num_gpu: Optional[int] = None, ) -> Iterator[ChatResponse]: ... def chat( @@ -350,6 +361,7 @@ def chat( format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, options: Optional[Union[Mapping[str, Any], Options]] = None, keep_alive: Optional[Union[float, str]] = None, + num_gpu: Optional[int] = None, ) -> Union[ChatResponse, Iterator[ChatResponse]]: """ Create a chat response using the requested model. @@ -361,6 +373,10 @@ def chat( For more information, see: https://google.github.io/styleguide/pyguide.html#38-comments-and-docstrings stream: Whether to stream the response. format: The format of the response. + num_gpu: Number of layers to offload to the GPU. Merged into `options`. + For hard isolation across multiple GPUs, run a separate `ollama serve` + per GPU (each with its own `CUDA_VISIBLE_DEVICES`) and point a separate + `Client(host=...)` at each. Example: def add_two_numbers(a: int, b: int) -> int: @@ -397,7 +413,7 @@ def add_two_numbers(a: int, b: int) -> int: logprobs=logprobs, top_logprobs=top_logprobs, format=format, - options=options, + options=_merge_options(options, num_gpu=num_gpu), keep_alive=keep_alive, ).model_dump(exclude_none=True), stream=stream, @@ -853,6 +869,7 @@ async def generate( width: Optional[int] = None, height: Optional[int] = None, steps: Optional[int] = None, + num_gpu: Optional[int] = None, ) -> GenerateResponse: ... @overload @@ -877,6 +894,7 @@ async def generate( width: Optional[int] = None, height: Optional[int] = None, steps: Optional[int] = None, + num_gpu: Optional[int] = None, ) -> AsyncIterator[GenerateResponse]: ... async def generate( @@ -900,10 +918,17 @@ async def generate( width: Optional[int] = None, height: Optional[int] = None, steps: Optional[int] = None, + num_gpu: Optional[int] = None, ) -> Union[GenerateResponse, AsyncIterator[GenerateResponse]]: """ Create a response using the requested model. + Args: + num_gpu: Number of layers to offload to the GPU. Merged into `options`. + For hard isolation across multiple GPUs, run a separate `ollama serve` + per GPU (each with its own `CUDA_VISIBLE_DEVICES`) and point a separate + `Client(host=...)` at each. + Raises `RequestError` if a model is not provided. Raises `ResponseError` if the request could not be fulfilled. @@ -928,7 +953,7 @@ async def generate( raw=raw, format=format, images=list(_copy_images(images)) if images else None, - options=options, + options=_merge_options(options, num_gpu=num_gpu), keep_alive=keep_alive, width=width, height=height, @@ -951,6 +976,7 @@ async def chat( format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, options: Optional[Union[Mapping[str, Any], Options]] = None, keep_alive: Optional[Union[float, str]] = None, + num_gpu: Optional[int] = None, ) -> ChatResponse: ... @overload @@ -967,6 +993,7 @@ async def chat( format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, options: Optional[Union[Mapping[str, Any], Options]] = None, keep_alive: Optional[Union[float, str]] = None, + num_gpu: Optional[int] = None, ) -> AsyncIterator[ChatResponse]: ... async def chat( @@ -982,6 +1009,7 @@ async def chat( format: Optional[Union[Literal['', 'json'], JsonSchemaValue]] = None, options: Optional[Union[Mapping[str, Any], Options]] = None, keep_alive: Optional[Union[float, str]] = None, + num_gpu: Optional[int] = None, ) -> Union[ChatResponse, AsyncIterator[ChatResponse]]: """ Create a chat response using the requested model. @@ -993,6 +1021,10 @@ async def chat( For more information, see: https://google.github.io/styleguide/pyguide.html#38-comments-and-docstrings stream: Whether to stream the response. format: The format of the response. + num_gpu: Number of layers to offload to the GPU. Merged into `options`. + For hard isolation across multiple GPUs, run a separate `ollama serve` + per GPU (each with its own `CUDA_VISIBLE_DEVICES`) and point a separate + `Client(host=...)` at each. Example: def add_two_numbers(a: int, b: int) -> int: @@ -1030,7 +1062,7 @@ def add_two_numbers(a: int, b: int) -> int: logprobs=logprobs, top_logprobs=top_logprobs, format=format, - options=options, + options=_merge_options(options, num_gpu=num_gpu), keep_alive=keep_alive, ).model_dump(exclude_none=True), stream=stream, @@ -1313,6 +1345,12 @@ async def ps(self) -> ProcessResponse: ) +def _merge_options(options: Optional[Union[Mapping[str, Any], Options]], **overrides: Any) -> Optional[Dict[str, Any]]: + merged = options.model_dump(exclude_none=True) if isinstance(options, Options) else dict(options or {}) + merged.update({k: v for k, v in overrides.items() if v is not None}) + return merged or None + + def _copy_images(images: Optional[Sequence[Union[Image, Any]]]) -> Iterator[Image]: for image in images or []: yield image if isinstance(image, Image) else Image(value=image) diff --git a/tests/test_client.py b/tests/test_client.py index 7b7ab38e..d655b644 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -100,6 +100,58 @@ def test_client_chat_with_logprobs(httpserver: HTTPServer): assert response['logprobs'][0]['top_logprobs'][1]['token'] == 'Hi' +def test_client_chat_with_num_gpu(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/chat', + method='POST', + json={ + 'model': 'dummy', + 'messages': [{'role': 'user', 'content': 'Hi'}], + 'tools': [], + 'stream': False, + 'options': {'num_gpu': 2}, + }, + ).respond_with_json( + { + 'model': 'dummy', + 'message': { + 'role': 'assistant', + 'content': 'Hello', + }, + } + ) + + client = Client(httpserver.url_for('/')) + response = client.chat('dummy', messages=[{'role': 'user', 'content': 'Hi'}], num_gpu=2) + assert response['message']['content'] == 'Hello' + + +def test_client_chat_with_num_gpu_merges_options(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/chat', + method='POST', + json={ + 'model': 'dummy', + 'messages': [{'role': 'user', 'content': 'Hi'}], + 'tools': [], + 'stream': False, + 'options': {'temperature': 0.5, 'num_gpu': 2}, + }, + ).respond_with_json( + { + 'model': 'dummy', + 'message': { + 'role': 'assistant', + 'content': 'Hello', + }, + } + ) + + client = Client(httpserver.url_for('/')) + response = client.chat('dummy', messages=[{'role': 'user', 'content': 'Hi'}], options={'temperature': 0.5}, num_gpu=2) + assert response['message']['content'] == 'Hello' + + def test_client_chat_stream(httpserver: HTTPServer): def stream_handler(_: Request): def generate(): @@ -333,6 +385,29 @@ def test_client_generate(httpserver: HTTPServer): assert response['response'] == 'Because it is.' +def test_client_generate_with_num_gpu(httpserver: HTTPServer): + httpserver.expect_ordered_request( + '/api/generate', + method='POST', + json={ + 'model': 'dummy', + 'prompt': 'Why is the sky blue?', + 'stream': False, + 'options': {'num_gpu': 1}, + }, + ).respond_with_json( + { + 'model': 'dummy', + 'response': 'Because it is.', + } + ) + + client = Client(httpserver.url_for('/')) + response = client.generate('dummy', 'Why is the sky blue?', num_gpu=1) + assert response['model'] == 'dummy' + assert response['response'] == 'Because it is.' + + def test_client_generate_with_logprobs(httpserver: HTTPServer): httpserver.expect_ordered_request( '/api/generate',