Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Make get_iree_devices read IREE_DEVICE env var if provided #891

Merged
merged 2 commits into from
Mar 7, 2025
Merged

Make get_iree_devices read IREE_DEVICE env var if provided #891

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
4624180
Make get_iree_devices read IREE_DEVICE env var if provided
sogartar Jan 30, 2025
55e3a69
Add docstring to get_iree_devices
sogartar Feb 21, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 35 additions & 1 deletion sharktank/sharktank/utils/iree.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
from typing import List, Tuple, Optional, Union
from pathlib import Path
import torch
import os
import numpy as np
import collections.abc
from collections import OrderedDict
Expand All @@ -22,7 +23,40 @@
from .tree import Tree


def get_iree_devices(driver: str, device_count: int) -> List[iree.runtime.HalDevice]:
def get_iree_devices(
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would, but unfortunately our code is fragmented and not all tests use this function to create the devices.
I would need to refactor them first to use this function.

*, driver: str | None = None, device_count: int = 1
) -> List[iree.runtime.HalDevice]:
Comment on lines +26 to +28
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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)

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done.

"""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 ...
```
"""
if "IREE_DEVICE" in os.environ:
device_uris = [d.strip() for d in os.environ["IREE_DEVICE"].split(",")]
driver_names = [n.split("://")[0] for n in device_uris]
if driver is not None:
if any(driver != driver_name for driver_name in driver_names):
ValueError(
f'Inconsistent IREE driver, expected "{driver}" for all devices f{device_uris}'
)
if device_count > len(device_uris):
raise ValueError(
"Environment variable IREE_DEVICE provides less devices than requested."
)
return [
iree.runtime.get_driver(driver_names[i]).create_device_by_uri(
device_uris[i]
)
for i in range(device_count)
]

hal_driver = iree.runtime.get_driver(driver)
available_devices = hal_driver.query_available_devices()
if driver in ["local-task", "local-sync"]:
Expand Down
Loading