Use WinML for python ONNX runtime #19416
Conversation
There was a problem hiding this comment.
Pull request overview
This PR integrates Windows ML (WinML) as the ONNX runtime execution provider for AI image captioning to achieve better performance through native Windows optimizations. The implementation adds a new WinML wrapper singleton, registers WinML execution providers with ONNX Runtime, and updates dependencies to include WinML packages.
Key changes:
- New WinML wrapper singleton for managing execution providers and WinAppSDK lifecycle
- Integration of WinML execution providers into the VitGpt2 image captioner
- Addition of WinML-related dependencies including packages from nightly builds
Reviewed changes
Copilot reviewed 3 out of 4 changed files in this pull request and generated 5 comments.
| File | Description |
|---|---|
| source/_localCaptioner/captioner/winML.py | New singleton wrapper for WinML initialization and execution provider registration |
| source/_localCaptioner/captioner/vitGpt2.py | Integrates WinML with existing captioner and sets execution provider selection policy |
| pyproject.toml | Adds WinML dependencies and configures nightly build package index |
| uv.lock | Lock file updates with new WinML dependencies and extensive transitive dependency tree |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| def __new__(cls, *args, **kwargs): | ||
| global _WINML_INSTANCE | ||
| if _WINML_INSTANCE is None: | ||
| _WINML_INSTANCE = super(_WinML, cls).__new__(cls, *args, **kwargs) | ||
| _WINML_INSTANCE._initialized = False | ||
| return _WINML_INSTANCE |
There was a problem hiding this comment.
The singleton implementation has a concurrency issue. If multiple threads call _WinML() simultaneously, they could both pass the if _WINML_INSTANCE is None check before either creates the instance, leading to multiple instances being created. This violates the singleton pattern and could cause resource conflicts with WinAppSDK initialization.
There was a problem hiding this comment.
I hate to agree with the robot, but it's right. Can you either either make this thread safe or explicitly note that this class is not thread safe?
| def __del__(self): | ||
| self._providers = None | ||
| self._winAppSDKHandle.__exit__(None, None, None) | ||
|
|
There was a problem hiding this comment.
The del method may not be called reliably in Python, especially during interpreter shutdown. This could leave the WinAppSDK handle unclosed. Consider using a context manager pattern or providing an explicit cleanup method that can be called deterministically.
| def __del__(self): | |
| self._providers = None | |
| self._winAppSDKHandle.__exit__(None, None, None) | |
| def close(self) -> None: | |
| """ | |
| Release resources associated with the WinAppSDK handle. | |
| This method can be called deterministically by users of this class | |
| or via the context manager protocol. It is safe to call multiple times. | |
| """ | |
| # Ensure attributes exist and have not already been cleaned up. | |
| win_app_sdk_handle = getattr(self, "_winAppSDKHandle", None) | |
| if win_app_sdk_handle is None: | |
| return | |
| # Drop references to providers first. | |
| self._providers = None | |
| # Call the handle's context manager exit method, ignoring any exception. | |
| try: | |
| win_app_sdk_handle.__exit__(None, None, None) | |
| finally: | |
| # Prevent double-closing. | |
| self._winAppSDKHandle = None | |
| def __enter__(self): | |
| """ | |
| Allow use of this singleton as a context manager. | |
| """ | |
| return self | |
| def __exit__(self, exc_type, exc_val, exc_tb): | |
| """ | |
| Ensure resources are released when leaving a context manager block. | |
| """ | |
| self.close() | |
| def __del__(self): | |
| """ | |
| Best-effort cleanup in case close() was not called explicitly. | |
| Note: __del__ is not guaranteed to be called; prefer using close() | |
| or the context manager protocol for deterministic cleanup. | |
| """ | |
| try: | |
| self.close() | |
| except Exception: | |
| # Suppress all exceptions in __del__ to avoid issues during | |
| # interpreter shutdown or partial initialization. | |
| pass |
There was a problem hiding this comment.
It's also right with it's comment about there being no guarantee that __del__ will be called. Is it a concern here though?
| This function removes the msvcp140.dll from the winrt-runtime package. | ||
| So it does not cause issues with other libraries. | ||
| """ | ||
| site_packages_path = str(metadata.distribution("winrt-runtime").locate_file("")) | ||
| dllPath = Path(os.path.join(site_packages_path, "winrt", "msvcp140.dll")) | ||
| if dllPath.exists(): | ||
| dllPath.unlink() | ||
|
|
There was a problem hiding this comment.
Deleting msvcp140.dll from a third-party package is fragile and could break on package updates or reinstalls. This hack may cause issues if other code expects the DLL to be present, and it's unclear what specific conflict this is solving. Consider documenting the specific issue being fixed and exploring alternative solutions like DLL search path manipulation.
| This function removes the msvcp140.dll from the winrt-runtime package. | |
| So it does not cause issues with other libraries. | |
| """ | |
| site_packages_path = str(metadata.distribution("winrt-runtime").locate_file("")) | |
| dllPath = Path(os.path.join(site_packages_path, "winrt", "msvcp140.dll")) | |
| if dllPath.exists(): | |
| dllPath.unlink() | |
| Attempt to avoid conflicts between the Visual C++ runtime used by winrt-runtime | |
| and other libraries that rely on the system-installed runtime. | |
| Previously, this function deleted ``msvcp140.dll`` from the ``winrt-runtime`` | |
| package directory. Modifying files inside third-party packages is fragile and | |
| can be undone by package updates or reinstalls, and may break other code that | |
| legitimately depends on that DLL. | |
| Instead of deleting the DLL, we try to ensure that the system Visual C++ | |
| runtime (typically located in ``%SystemRoot%\\System32``) is preferred by | |
| adding the system directory to the DLL search path when supported. | |
| """ | |
| # On Python 3.8+ on Windows, add_dll_directory allows us to influence DLL | |
| # resolution without mutating third-party package contents. | |
| system_root = os.environ.get("SystemRoot") | |
| if not system_root: | |
| return | |
| system32_path = os.path.join(system_root, "System32") | |
| # add_dll_directory is only available on Windows and Python >= 3.8. | |
| add_dll_directory = getattr(os, "add_dll_directory", None) | |
| if callable(add_dll_directory): | |
| try: | |
| # We keep the handle alive by storing it on the instance so that the | |
| # added directory remains active for the lifetime of this object. | |
| self._system32_dll_dir = add_dll_directory(system32_path) | |
| except Exception: | |
| # As a fallback, we do nothing rather than deleting files. | |
| log.debug(f"Failed to add DLL directory: {system32_path}", exc_info=True) |
There was a problem hiding this comment.
this is copied from sample code, but I agree, its a red flag
|
|
||
| for name, path in self._epPaths.items(): | ||
| if name not in self._registeredEps: | ||
| try: | ||
| ort.register_execution_provider_library(name, path) | ||
| self._registeredEps.append(name) | ||
| except Exception as e: | ||
| log.exception(f"Failed to register execution provider {name}: {e}") |
There was a problem hiding this comment.
The exception is logged but execution providers that fail to register are silently skipped. This could lead to unexpected performance degradation if a critical provider fails to load. Consider whether this should be a hard failure for critical providers or at least provide more visibility about which providers successfully registered.
| for name, path in self._epPaths.items(): | |
| if name not in self._registeredEps: | |
| try: | |
| ort.register_execution_provider_library(name, path) | |
| self._registeredEps.append(name) | |
| except Exception as e: | |
| log.exception(f"Failed to register execution provider {name}: {e}") | |
| failed_providers: list[str] = [] | |
| for name, path in self._epPaths.items(): | |
| if name not in self._registeredEps: | |
| try: | |
| ort.register_execution_provider_library(name, path) | |
| self._registeredEps.append(name) | |
| except Exception as e: | |
| failed_providers.append(name) | |
| log.exception(f"Failed to register execution provider {name}: {e}") | |
| # If no providers could be registered at all, treat this as a hard failure. | |
| if self._epPaths and not self._registeredEps: | |
| raise RuntimeError( | |
| "Failed to register any ONNX Runtime execution providers. " | |
| "See log for individual provider errors." | |
| ) | |
| # If some providers failed but at least one succeeded, log a summary. | |
| if failed_providers: | |
| log.warning( | |
| f"Some execution providers failed to register and will be skipped: " | |
| f"{', '.join(failed_providers)}" | |
| ) |
There was a problem hiding this comment.
this is copied from sample code, I'm not sure if this is necessary
| onnxruntime-winml = { index = "onnxruntime-winml-nightly" } | ||
|
|
||
| [[tool.uv.index]] | ||
| name = "onnxruntime-winml-nightly" | ||
| url = "https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ORT-Nightly/pypi/simple/" | ||
| explicit = true | ||
|
|
There was a problem hiding this comment.
The onnxruntime-winml package is sourced from a nightly build URL, which could introduce instability or breaking changes. Nightly builds are not recommended for production use as they may contain unstable or untested code. Consider using a stable release version if available.
| onnxruntime-winml = { index = "onnxruntime-winml-nightly" } | |
| [[tool.uv.index]] | |
| name = "onnxruntime-winml-nightly" | |
| url = "https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ORT-Nightly/pypi/simple/" | |
| explicit = true |
There was a problem hiding this comment.
No stable release is available
| "pyyaml==6.0.3", | ||
| "pymdown-extensions==10.17.1", | ||
| # local image caption | ||
| "onnxruntime==1.23.2", |
|
|
||
| from logHandler import log | ||
|
|
||
| _WINML_INSTANCE = None |
| def __new__(cls, *args, **kwargs): | ||
| global _WINML_INSTANCE | ||
| if _WINML_INSTANCE is None: | ||
| _WINML_INSTANCE = super(_WinML, cls).__new__(cls, *args, **kwargs) |
There was a problem hiding this comment.
Why do we need explicit arguments to super here?
| _WINML_INSTANCE = super(_WinML, cls).__new__(cls, *args, **kwargs) | |
| _WINML_INSTANCE = super().__new__(cls, *args, **kwargs) |
| def __new__(cls, *args, **kwargs): | ||
| global _WINML_INSTANCE | ||
| if _WINML_INSTANCE is None: | ||
| _WINML_INSTANCE = super(_WinML, cls).__new__(cls, *args, **kwargs) | ||
| _WINML_INSTANCE._initialized = False | ||
| return _WINML_INSTANCE |
There was a problem hiding this comment.
I hate to agree with the robot, but it's right. Can you either either make this thread safe or explicitly note that this class is not thread safe?
| def __del__(self): | ||
| self._providers = None | ||
| self._winAppSDKHandle.__exit__(None, None, None) | ||
|
|
There was a problem hiding this comment.
It's also right with it's comment about there being no guarantee that __del__ will be called. Is it a concern here though?
| This function removes the msvcp140.dll from the winrt-runtime package. | ||
| So it does not cause issues with other libraries. | ||
| """ | ||
| site_packages_path = str(metadata.distribution("winrt-runtime").locate_file("")) | ||
| dllPath = Path(os.path.join(site_packages_path, "winrt", "msvcp140.dll")) | ||
| if dllPath.exists(): | ||
| dllPath.unlink() | ||
|
|
Link to issue number:
Closes #18662
Summary of the issue:
Windows ML provides an optimised native way to run AI models on Windows.
Description of user facing changes:
Potential performance increase when performing AI image descriptions.
No performance changes to how NVDA acts when loading and holding numpy in memory by enabling AI image descriptions.
Description of developer facing changes:
None
Description of development approach:
Use python WinML runtimes with onnxruntime as per docs
Testing strategy:
Tested performing AI image descriptions
Known issues with pull request:
Unfortunately, this still doesn't remove our dependency on numpy.
As such, there are not much performance improvements.
Code Review Checklist: