Make get_iree_devices read IREE_DEVICE env var if provided

[sogartar]

This allows to inject what exact IREE device(s) are to be used without propagating all the way to program arguments.
Example:

IREE_DEVICE=hip://5,hip://3 pytest sharktank/tests

[archana-ramalingam]

If this change is specifically targeting CIs, our runners pick whichever gpu is available atm. We had a PR making this change for llama benchmark/ppl CIs. @ScottTodd can confirm.

[sogartar]

It is not just for CIs, but when running locally as well. There should be some way of overriding the device list.
Are CI machines actually allocating devices to different jobs? Where is this code hosted?

[ScottTodd]

It is not just for CIs, but when running locally as well. There should be some way of overriding the device list. Are CI machines actually allocating devices to different jobs? Where is this code hosted?

I believe on servers with multiple GPUs, we have one runner per GPU, using ROCR_VISIBLE_DEVICES=1 or equivalent on runner 1, ROCR_VISIBLE_DEVICES=2 on runner 2, etc. We also have a cluster of runners with labels like linux-mi300-gpu-1 getting set up as part of #793. I think we'd create a linux-mi300-gpu-2 label for workflows that request two GPUs?

That's my understanding anyways. @yamiyysu and @saienduri would know more, and we should write that all down somewhere and get the code used to configure the runners checked in too.

[ScottTodd]

Can you add a docstring to this function and mention the IREE_DEVICE environment variable with a few examples? You can lift some text from the PR description

This allows to inject what exact IREE device(s) are to be used without propagating all the way to program arguments. Example:

IREE_DEVICE=hip://5,hip://3 pytest sharktank/tests

def get_iree_devices(
    *, driver: str | None = None, device_count: int = 1
) -> List[iree.runtime.HalDevice]:
    """Gets a list of IREE HAL devices for the given driver.
    
    The first available device_count devices will be created,
    unless the IREE_DEVICE environment variable is set to an
    explicit list of device URIs.

    For example, to select HIP devices 5 and 3:
    ```
    export IREE_DEVICE=hip://5,hip://3
    python ... device=hip # (something here)
    ```
    """

(could wordsmith that a bit, just getting the general idea across)

[ScottTodd]

May also want to put pytest details in https://github.com/nod-ai/shark-ai/blob/main/docs/developer_guide.md#running-tests