Expose num_gpu as a first-class argument on generate/chat

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.

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`

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')

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)

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',