Skip to content

Commit e2a0008

Browse files
Merge pull request #249 from bioimage-io/cache
Use new no-cache env var
2 parents 1c366cf + 8345acd commit e2a0008

File tree

2 files changed

+35
-16
lines changed

2 files changed

+35
-16
lines changed

bioimageio/core/resource_io/io_.py

Lines changed: 25 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import os
22
import pathlib
33
from copy import deepcopy
4+
from tempfile import TemporaryDirectory
45
from typing import Dict, Optional, Sequence, Union
56
from zipfile import ZIP_DEFLATED, ZipFile
67

@@ -10,7 +11,12 @@
1011
from bioimageio.core.resource_io.nodes import ResourceDescription
1112
from bioimageio.spec import load_raw_resource_description
1213
from bioimageio.spec.shared import raw_nodes
13-
from bioimageio.spec.shared.common import BIOIMAGEIO_CACHE_PATH, get_class_name_from_type
14+
from bioimageio.spec.shared.common import (
15+
BIOIMAGEIO_CACHE_PATH,
16+
BIOIMAGEIO_USE_CACHE,
17+
get_class_name_from_type,
18+
no_cache_tmp_list,
19+
)
1420
from bioimageio.spec.shared.raw_nodes import ResourceDescription as RawResourceDescription
1521
from . import nodes
1622
from .utils import resolve_raw_resource_description, resolve_source
@@ -134,21 +140,26 @@ def _get_package_base_name(raw_rd: RawResourceDescription, weights_priority_orde
134140

135141

136142
def _get_tmp_package_path(raw_rd: RawResourceDescription, weights_priority_order: Optional[Sequence[str]]):
137-
package_file_name = _get_package_base_name(raw_rd, weights_priority_order)
138-
139-
cache_folder = BIOIMAGEIO_CACHE_PATH / "packages"
140-
cache_folder.mkdir(exist_ok=True, parents=True)
141-
package_path = (cache_folder / package_file_name).with_suffix(".zip")
142-
max_cached_packages_with_same_name = 100
143-
for p in range(max_cached_packages_with_same_name):
144-
if package_path.exists():
145-
package_path = (cache_folder / f"{package_file_name}p{p}").with_suffix(".zip")
143+
if BIOIMAGEIO_USE_CACHE:
144+
package_file_name = _get_package_base_name(raw_rd, weights_priority_order)
145+
cache_folder = BIOIMAGEIO_CACHE_PATH / "packages"
146+
cache_folder.mkdir(exist_ok=True, parents=True)
147+
148+
package_path = (cache_folder / package_file_name).with_suffix(".zip")
149+
max_cached_packages_with_same_name = 100
150+
for p in range(max_cached_packages_with_same_name):
151+
if package_path.exists():
152+
package_path = (cache_folder / f"{package_file_name}p{p}").with_suffix(".zip")
153+
else:
154+
break
146155
else:
147-
break
156+
raise FileExistsError(
157+
f"Already caching {max_cached_packages_with_same_name} versions of {cache_folder / package_file_name}!"
158+
)
148159
else:
149-
raise FileExistsError(
150-
f"Already caching {max_cached_packages_with_same_name} versions of {cache_folder / package_file_name}!"
151-
)
160+
tmp_dir = TemporaryDirectory()
161+
no_cache_tmp_list.append(tmp_dir)
162+
package_path = pathlib.Path(tmp_dir.name) / "file"
152163

153164
return package_path
154165

tests/test_cli.py

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import os
12
import subprocess
23
from typing import Sequence
34

@@ -7,9 +8,9 @@
78
from bioimageio.core import load_resource_description
89

910

10-
def run_subprocess(commands: Sequence[str]) -> subprocess.CompletedProcess:
11+
def run_subprocess(commands: Sequence[str], **kwargs) -> subprocess.CompletedProcess:
1112
# return subprocess.run(commands, capture_output=True)
12-
return subprocess.run(commands, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, encoding="utf-8")
13+
return subprocess.run(commands, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, encoding="utf-8", **kwargs)
1314

1415

1516
def test_validate_model(unet2d_nuclei_broad_model):
@@ -24,6 +25,13 @@ def test_cli_package(unet2d_nuclei_broad_model, tmp_path):
2425
assert out_path.exists()
2526

2627

28+
def test_cli_package_wo_cache(unet2d_nuclei_broad_model):
29+
env = os.environ.copy()
30+
env["BIOIMAGEIO_USE_CACHE"] = "false"
31+
ret = run_subprocess(["bioimageio", "package", unet2d_nuclei_broad_model], env=env)
32+
assert ret.returncode == 0, ret.stdout
33+
34+
2735
def test_cli_test_model(unet2d_nuclei_broad_model):
2836
ret = run_subprocess(["bioimageio", "test-model", unet2d_nuclei_broad_model])
2937
assert ret.returncode == 0, ret.stdout

0 commit comments

Comments
 (0)