Skip to content

Add an environment variable to disable ort session cache if necessary | fix(atenlib) #685

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

Merged
merged 6 commits into from
Apr 29, 2023
Merged

Add an environment variable to disable ort session cache if necessary | fix(atenlib) #685

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
d70ff44
Add an environment variable to disable ort session cache for windows …
fatcat-z Apr 27, 2023
1fadd93
Fix a lint issue.
fatcat-z Apr 27, 2023
776c041
Merge branch 'main' into remove_cache
fatcat-z Apr 28, 2023
16244e7
Add a centralized place for possible feature switches.
fatcat-z Apr 28, 2023
4083b07
Address comments.
fatcat-z Apr 28, 2023
a3a4952
Merge branch 'main' into remove_cache
fatcat-z Apr 29, 2023
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
1 change: 1 addition & 0 deletions .github/workflows/main.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,7 @@ jobs:
- name: Run tests
run: nox -t ${{ matrix.nox-tag }} --forcecolor -- -v --cov=onnxscript --cov-report=xml --cov-append --cov-branch -n=auto
env:
CACHE_ORT_SESSIONS: "${{ matrix.os == 'windows-latest' && '0' || '1' }}"
CATCH_ORT_SEGFAULT: "${{ matrix.os == 'ubuntu-latest' && '1' || '0' }}"
- name: Upload coverage to Codecov
uses: codecov/codecov-action@v3
Expand Down
10 changes: 10 additions & 0 deletions onnxscript/_internal/feature_switch.py
View file
Open in desktop
Copy link
Collaborator

Choose a reason for hiding this comment

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

consider flags.py as in "feature flags"

Copy link
Contributor Author

Choose a reason for hiding this comment

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

In my past projects, such things which determines if a feature is enable or not were named like "xxxx_switch", so I followed them.
In addition, I think 'switch' implies this can control something while 'flag' is more like a static label.

Original file line number Diff line number Diff line change
@@ -0,0 +1,10 @@
# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# --------------------------------------------------------------------------
"""Switches to determine if the corresponding feature of onnxscript is enabled or not."""

import os

# By default: Enable
CACHE_ORT_SESSIONS: bool = os.getenv("CACHE_ORT_SESSIONS", "1") != "0"
Copy link
Collaborator

Choose a reason for hiding this comment

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

Suggested change
CACHE_ORT_SESSIONS: bool = os.getenv("CACHE_ORT_SESSIONS", "1") != "0"
cache_ort_sessions: bool = os.getenv("CACHE_ORT_SESSIONS", "1") != "0"

sorry - I changed my mind because it looks like it is a flag that can be changed in runtime (not constant), so snake case instead?

Copy link
Collaborator

Choose a reason for hiding this comment

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

Is this in _internal by design? I can see arguments both ways. It should ideally be exposed to the user, but if we think this API will change, I can see a point in leaving it here.

Copy link
Contributor Author

@fatcat-z fatcat-z Apr 29, 2023 •
edited
Loading

Choose a reason for hiding this comment

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

Is this in _internal by design? I can see arguments both ways. It should ideally be exposed to the user, but if we think this API will change, I can see a point in leaving it here.

My opinion is: Such switches should not be exposed to users calling onnx-script ops or torch_libs to construct a model. Because these switches control some internal features of onnx-script only. These internal features are not necessary to be known by those users.

Copy link
Contributor Author

@fatcat-z fatcat-z Apr 29, 2023 •
edited
Loading

Choose a reason for hiding this comment

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

sorry - I changed my mind because it looks like it is a flag that can be changed in runtime (not constant), so snake case instead?

I think a variable in upper case is more like a constant which should not be changed in runtime. Isn't it? And I think this is aligned with the assumption that these switches are designed to be controlled by the environment variable only, so we don't need to change them to snake case.

18 changes: 11 additions & 7 deletions onnxscript/evaluator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@
from typing_extensions import TypeAlias

from onnxscript import autocast, irbuilder, onnx_opset, tensor, utils, values
from onnxscript._internal import param_manipulation
from onnxscript._internal import feature_switch, param_manipulation

Check notice

Code scanning / CodeQL

Cyclic import

Import of module [onnxscript._internal.param_manipulation](1) begins an import cycle.

if typing.TYPE_CHECKING:
import onnxruntime as ort
Expand Down Expand Up @@ -354,12 +354,16 @@ def _cache_(model, providers):
import onnxruntime as ort # pylint: disable=import-outside-toplevel

serialized = model.SerializeToString()
key = serialized, tuple(providers)
if key in _cache_models:
return _cache_models[key]
session = ort.InferenceSession(serialized, providers=providers)
_cache_models[key] = session
return session
if feature_switch.CACHE_ORT_SESSIONS:
key = serialized, tuple(providers)
if key in _cache_models:
return _cache_models[key]
session = ort.InferenceSession(serialized, providers=providers)
_cache_models[key] = session

return session

return ort.InferenceSession(serialized, providers=providers)


def _os_to_ort_value(v):
Expand Down