From 7513c16d1cabfc108fe05afb0eae7297d26ca994 Mon Sep 17 00:00:00 2001 From: Ramakrishnan Sivakumar Date: Fri, 3 May 2024 13:35:34 -0700 Subject: [PATCH] Verify model deletion after testing (#6127) ### Description Deletes model downloaded to disk after testing to avoid storage issues. --------- Signed-off-by: Ramakrishnan Sivakumar --- workflow_scripts/test_model_zoo.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/workflow_scripts/test_model_zoo.py b/workflow_scripts/test_model_zoo.py index dc0702f93f0..8f218fde4f6 100644 --- a/workflow_scripts/test_model_zoo.py +++ b/workflow_scripts/test_model_zoo.py @@ -6,6 +6,7 @@ import argparse import gc import os +import shutil import sys import time @@ -95,8 +96,10 @@ def main(): ) # remove the model to save space in CIs - if os.path.exists(model_name): - os.remove(model_name) + full_model_path = os.path.join(hub._ONNX_HUB_DIR, model_path) + parent_dir = os.path.dirname(full_model_path) + if os.path.exists(parent_dir): + shutil.rmtree(parent_dir) except Exception as e: # noqa: BLE001 print(f"[FAIL]: {e}")