Skip to content

Commit d469205

Browse files
committed
Update test_model_framework_tensorflow_model.py to use local file as lock.
1 parent cd2c8f3 commit d469205

File tree

1 file changed

+18
-2
lines changed

1 file changed

+18
-2
lines changed

tests/unitary/with_extras/model/test_model_framework_tensorflow_model.py

Lines changed: 18 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,8 +9,10 @@
99
import base64
1010
import os
1111
import shutil
12+
import time
1213
import uuid
1314
from io import BytesIO
15+
from pathlib import Path
1416

1517
import numpy as np
1618
import onnxruntime as rt
@@ -35,9 +37,23 @@ def setup_module():
3537

3638

3739
class MyTFModel:
38-
3940
mnist = tf.keras.datasets.mnist
40-
(x_train, y_train), (x_test, y_test) = mnist.load_data()
41+
lock_file = os.path.expanduser("~/.keras/datasets/mnist.npz.downloading")
42+
try:
43+
time_start = time.time()
44+
while os.path.exists(lock_file):
45+
time.sleep(10)
46+
elapsed = time.time() - time_start
47+
if elapsed > 60 * 15:
48+
raise TimeoutError(
49+
"Timeout when waiting for MNIST data to finished downloading in another process. "
50+
"Please check internet connection."
51+
)
52+
Path(lock_file).touch()
53+
(x_train, y_train), (x_test, y_test) = mnist.load_data()
54+
finally:
55+
os.unlink(lock_file)
56+
4157
x_train, x_test = x_train / 255.0, x_test / 255.0
4258

4359
x_train, y_train = x_train[:1000], y_train[:1000]

0 commit comments

Comments
 (0)