[shortfin] Unhelpful error when specifying too many devices

[Alex-Vasile]

When running a sharded model through shortfin, the error messages associated with incorrectly specifying devices are unhelpful.

For example, if running a model sharded over 2 devices, but shortfin is invoked with --device=hip --device_ids 0 1 2 3 4, then shortfin will provide the following error:

ERROR:shortfin_apps.llm.components.service:Fatal error in prefetch invocation
Traceback (most recent call last):
  File "/shark-ai/shortfin/python/shortfin_apps/llm/components/service.py", line 414, in run
    (logits,) = await fn(*args, fiber=self.fiber)
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^
ValueError: shortfin_iree-src/runtime/src/iree/base/allocator.c:21: INVALID_ARGUMENT; allocator has no control routine
[2025-03-14 08:44:26.866] [error] [service.py:439] Fatal error in prefetch invocation
Traceback (most recent call last):
  File "/shark-ai/shortfin/python/shortfin_apps/llm/components/service.py", line 414, in run
    (logits,) = await fn(*args, fiber=self.fiber)
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^
ValueError: shortfin_iree-src/runtime/src/iree/base/allocator.c:21: INVALID_ARGUMENT; allocator has no control routine
/shark-ai/shortfin/src/shortfin/support/iree_helpers.h:315: UNKNOWN; Unhandled exception: Traceback (most recent call last):
  File "/shark-ai/shortfin/python/shortfin_apps/llm/components/generate.py", line 83, in run
RuntimeError: Async exception on <Worker 'default-inference'>): argmax(): incompatible function arguments. The following argument types are supported:
    1. argmax(input: _shortfin_default.lib.array.device_array, axis: int = -1, out: _shortfin_default.lib.array.device_array | None = None, *, keepdims: bool = False, device_visible: bool = False) -> _shortfin_default.lib.array.device_array

Invoked with types: NoneType

It's not obvious from this error that the issue is with specifying too many devices. If too few devices are specified the error is very actionable. As is the error if we specify more devices than there are on the host.