[punet] Add integration tests. #84
Conversation
* Imports/validates the FP16 model running eagerly. * Imports/validates the int8 model running eagerly. * Exports the models.
stellaraccident
force-pushed
the
punet_integration_test
branch
from
4ec1e65
to
471473b
Compare
| REPO_ID = "amd-shark/sharktank-goldens" | ||
| REVISION = "230dad4d85fbcb8759a331dcf1d45f0562875abe" |
There was a problem hiding this comment.
Ooh using huggingface (instead of Azure, GCS, etc.)? Nice!
May want to link https://huggingface.co/amd-shark/sharktank-goldens somewhere for easy referencing
| def get_best_torch_device() -> str: | ||
| import torch | ||
|
|
||
| if torch.cuda.is_available() and torch.cuda.device_count() > 0: | ||
| return "cuda:0" | ||
| return "cpu" |
There was a problem hiding this comment.
May want a way to set this device index via environment variable or flag.
On some of our CI runners we try to only use a specific GPU (like 6), since we and possibly other groups are running multiple jobs on the same "machine" / node.
There was a problem hiding this comment.
Until this is running on CI, it would be useful to see the logs produced (e.g. in a gist).
Things I'd look for:
- Time taken for each step / test case
- Format of logs on success
- Format of logs on failure
- How easy the flow is to understand and reproduce outside of pytest
- What artifacts are downloaded
- What artifacts are produced
| return "cpu" | ||
|
|
||
|
|
||
| def assert_golden_safetensors(actual_path, ref_path): |
There was a problem hiding this comment.
This function is nice, giving a good reason to stay in Python instead of using bash or just the native tools like iree-run-module.
(I'm still slowly coming to terms with moving infrastructure from C/C++ to python :P)
Progress on #76