Repository: tenstorrent/tt-forge-fe
Pull request #1467: Add local model caching to fix transient CI failures
Author: login: vkovinicTT
Created: 2025-03-17T15:36:54Z | Merged: 2025-03-25T17:47:14Z

## Problem
Sometimes on CI some tests fail to download model locally resulting in failing a test and not passing CI checks.

## Solution
This workaround enables us to cache model with method:

- `fetch_model` which downloads model with retry mechanisms from given url to our model cache repository if it's not present already in the model repo.  Cache repo location depends on whether env variable `FORGE_MODELS_CACHE` is set or not (having default location in latter case).

The changes help stabilize our CI pipeline and prevent unnecessary test failures.


Diff:
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 1185cb2de8..d4c8ef31da 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -123,6 +123,7 @@ jobs:
       env:
         HF_TOKEN: ${{ secrets.HF_TOKEN }}
         HF_HOME: /mnt/dockercache/huggingface
+        FORGE_MODELS_CACHE: /mnt/dockercache/forge_models_cache
         HF_HUB_DISABLE_PROGRESS_BARS: 1
         FORGE_DISABLE_REPORTIFY_DUMP: 1
         OPERATORS: ${{ inputs.operators }}
diff --git a/.gitignore b/.gitignore
index 142bc6ca07..260a602435 100644
--- a/.gitignore
+++ b/.gitignore
@@ -69,3 +69,6 @@ pipeline_*.json
 
 coverage.info
 coverage.xml
+
+# Exclude cached models
+.forge_models_cache/
diff --git a/forge/test/models/pytorch/vision/yolo/test_yolo_v5.py b/forge/test/models/pytorch/vision/yolo/test_yolo_v5.py
index 807bf5aa9b..a6576a5e16 100644
--- a/forge/test/models/pytorch/vision/yolo/test_yolo_v5.py
+++ b/forge/test/models/pytorch/vision/yolo/test_yolo_v5.py
@@ -8,13 +8,15 @@
 from forge.verify.verify import verify
 
 from test.models.utils import Framework, build_module_name
-from test.utils import download_model
+from test.utils import fetch_model, yolov5_loader
+
+base_url = "https://github.com/ultralytics/yolov5/releases/download/v7.0"
 
 
 def generate_model_yoloV5I320_imgcls_torchhub_pytorch(variant, size):
     name = "yolov5" + size
 
-    model = download_model(torch.hub.load, variant, name, pretrained=True)
+    model = fetch_model(name, f"{base_url}/{name}.pt", yolov5_loader, variant=variant)
 
     input_shape = (1, 3, 320, 320)
     input_tensor = torch.rand(input_shape)
@@ -66,7 +68,7 @@ def test_yolov5_320x320(forge_property_recorder, size):
 
 def generate_model_yoloV5I640_imgcls_torchhub_pytorch(variant, size):
     name = "yolov5" + size
-    model = download_model(torch.hub.load, variant, name, pretrained=True)
+    model = fetch_model(name, f"{base_url}/{name}.pt", yolov5_loader, variant=variant)
 
     input_shape = (1, 3, 640, 640)
     input_tensor = torch.rand(input_shape)
@@ -126,7 +128,7 @@ def test_yolov5_640x640(forge_property_recorder, size):
 
 def generate_model_yoloV5I480_imgcls_torchhub_pytorch(variant, size):
     name = "yolov5" + size
-    model = download_model(torch.hub.load, variant, name, pretrained=True)
+    model = fetch_model(name, f"{base_url}/{name}.pt", yolov5_loader, variant=variant)
     input_shape = (1, 3, 480, 480)
     input_tensor = torch.rand(input_shape)
     return model, [input_tensor], {}
@@ -202,7 +204,12 @@ def test_yolov5_1280x1280(forge_property_recorder, variant):
     forge_property_recorder.record_group("generality")
     forge_property_recorder.record_model_name(module_name)
 
-    framework_model = download_model(torch.hub.load, "ultralytics/yolov5", variant, pretrained=True)
+    framework_model = fetch_model(
+        variant,
+        f"{base_url}/{variant}.pt",
+        yolov5_loader,
+        variant="ultralytics/yolov5",
+    )
 
     input_shape = (1, 3, 1280, 1280)
     input_tensor = torch.rand(input_shape)
diff --git a/forge/test/utils.py b/forge/test/utils.py
index 53a427af6f..feface0c85 100644
--- a/forge/test/utils.py
+++ b/forge/test/utils.py
@@ -1,6 +1,7 @@
 # SPDX-FileCopyrightText: © 2024 Tenstorrent AI ULC
 
 # SPDX-License-Identifier: Apache-2.0
+from typing import Any, Callable, Optional
 from loguru import logger
 import random
 import requests
@@ -8,10 +9,12 @@
 import os
 import shutil
 import urllib
+from filelock import FileLock
 
 import numpy as np
 import torch
 import tensorflow as tf
+from forge.module import FrameworkModule
 
 
 def download_model(download_func, *args, num_retries=3, timeout=180, **kwargs):
@@ -35,6 +38,95 @@ def download_model(download_func, *args, num_retries=3, timeout=180, **kwargs):
     assert False, "Failed to download the model after multiple retries."
 
 
+def get_cache_dir() -> str:
+    """Get models cache directory from env var or use local default."""
+    cache_dir = os.environ.get("FORGE_MODELS_CACHE")
+    if not cache_dir:
+        cache_dir = os.path.join(os.getcwd(), ".forge_models_cache")
+        os.makedirs(cache_dir, exist_ok=True)
+    return cache_dir
+
+
+def default_loader(path: str):
+    """Load model with PyTorch."""
+    try:
+        return torch.load(path, map_location="cpu")
+    except Exception as e:
+        print(f"PyTorch loading error: {e}")
+        return None
+
+
+def yolov5_loader(path: str, variant: str = "ultralytics/yolov5"):
+    try:
+        model = torch.hub.load(variant, "custom", path=path)
+        return model
+    except Exception as e:
+        print(f"YOLOv5 loading error: {e}")
+        return None
+
+
+def fetch_model(
+    model_name: str,
+    url: str,
+    loader: Optional[Callable] = default_loader,
+    max_retries: int = 3,
+    timeout: int = 30,
+    **kwargs: Any,
+) -> FrameworkModule:
+    """Fetch model from URL, cache it, and load it."""
+
+    model_file = model_name + ".pt"
+
+    model_path = os.path.join(get_cache_dir(), model_file)
+
+    # Download if needed
+    if not os.path.exists(model_path):
+        os.makedirs(os.path.dirname(model_path), exist_ok=True)
+        lock_path = model_path + ".lock"
+        lock = FileLock(lock_path)
+
+        with lock:
+            # Check again after acquiring the lock to handle concurrent processes
+            if not os.path.exists(model_path):
+                for attempt in range(1, max_retries + 1):
+                    try:
+                        print(f"Downloading {model_name}, attempt {attempt}/{max_retries}...")
+                        response = requests.get(url, timeout=timeout)
+                        response.raise_for_status()
+
+                        # Write to temporary file first
+                        temp_path = model_path + ".tmp"
+                        with open(temp_path, "wb") as f:
+                            f.write(response.content)
+
+                        # Atomic rename after successful download
+                        try:
+                            os.rename(temp_path, model_path)
+                        except OSError as e:
+                            print(f"Error during rename: {e}")
+                            # Clean up the temp file
+                            if os.path.exists(temp_path):
+                                os.remove(temp_path)
+                            raise  # Let this be caught by the outer exception handler
+
+                        break  # Successfully downloaded and renamed
+
+                    except (requests.exceptions.RequestException, OSError) as e:
+                        print(f"Attempt {attempt} failed: {e}")
+                        # Clean up temp file if it exists
+                        if os.path.exists(temp_path):
+                            os.remove(temp_path)
+
+                        if attempt < max_retries:
+                            time.sleep(2**attempt)  # Exponential backoff
+                        else:
+                            raise RuntimeError(f"Failed to download {model_name} after {max_retries} attempts.")
+
+    # Load model
+    model = loader(model_path, **kwargs) if loader else None
+    return model
+
+
 def reset_seeds():
     random.seed(0)
     np.random.seed(0)


Reviews:
user: login: mstojkovicTT
body: LGTM
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-03-18T04:21:59Z
user: login: pilkicTT
body: How is this caching anything? You should look into `/mnt/dockercache` which we use in the CI for caching (for now it looks like we only cache models from huggingface). We should do something similar for this.
state: CHANGES_REQUESTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-24T10:58:35Z
user: login: pilkicTT
body: Minor comments, looks good otherwise... Have you ran the nightlies with this change? 
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-25T08:51:07Z
user: login: vmilosevic
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-25T09:52:33Z
user: login: vmilosevic
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-25T09:57:00Z
user: login: vmilosevic
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-25T09:59:47Z
user: login: vmilosevic
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-25T10:09:29Z
user: login: vmilosevic
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-25T10:09:51Z
user: login: vkovinicTT
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-25T10:31:09Z
user: login: vkovinicTT
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-25T13:02:03Z
user: login: pilkicTT
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-25T13:10:10Z
user: login: vkovinicTT
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-25T13:53:16Z
user: login: pilkicTT
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-03-25T14:41:32Z
user: login: vmilosevic
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-25T14:56:37Z
user: login: vmilosevic
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-03-25T14:56:45Z

Review comments:
diff_hunk: @@ -124,6 +124,7 @@ jobs:
       env:
         HF_TOKEN: ${{ secrets.HF_TOKEN }}
         HF_HOME: /mnt/dockercache/huggingface
+        FORGE_MODELS_HUB: /mnt/dockercache/custom_models
path: .github/workflows/test.yml
user: login: pilkicTT
body: I would prefer something like:
```yaml
FORGE_MODELS_CACHE: /mnt/dockercache/forge_models_cache
```
created_at: 2025-03-25T08:44:50Z
author_association: CONTRIBUTOR
diff_hunk: @@ -35,6 +37,74 @@ def download_model(download_func, *args, num_retries=3, timeout=180, **kwargs):
     assert False, "Failed to download the model after multiple retries."
 
 
+def get_cache_dir() -> str:
+    """Get models cache directory from env var or use local default."""
+    cache_dir = os.environ.get("FORGE_MODELS_HUB")
+    if not cache_dir:
+        cache_dir = os.path.join(os.getcwd(), "cached_forge_models")
path: forge/test/utils.py
user: login: pilkicTT
body: make this directory `.forge_models_cache`
additionally you can add it to `.gitignore` as well...
created_at: 2025-03-25T08:46:29Z
author_association: CONTRIBUTOR
diff_hunk: @@ -8,13 +8,15 @@
 from forge.verify.verify import verify
 
 from test.models.utils import Framework, build_module_name
-from test.utils import download_model
+from test.utils import fetch_model, yolov5_loader
 
 
 def generate_model_yoloV5I320_imgcls_torchhub_pytorch(variant, size):
     name = "yolov5" + size
 
-    model = download_model(torch.hub.load, variant, name, pretrained=True)
+    model = fetch_model(
+        name, f"https://github.com/ultralytics/yolov5/releases/download/v7.0/{name}.pt", yolov5_loader, variant=variant
path: forge/test/models/pytorch/vision/yolo/test_yolo_v5.py
user: login: vmilosevic
body: Can we move this to a constant as it repeats a few times
created_at: 2025-03-25T09:52:33Z
author_association: COLLABORATOR
diff_hunk: @@ -124,6 +124,7 @@ jobs:
       env:
         HF_TOKEN: ${{ secrets.HF_TOKEN }}
         HF_HOME: /mnt/dockercache/huggingface
+        FORGE_MODELS_HUB: /mnt/dockercache/custom_models
path: .github/workflows/test.yml
user: login: vmilosevic
body: If this is specific to tt-forge-fe can it be under path /mnt/dockercache/tt-forge-fe/ ...
created_at: 2025-03-25T09:57:00Z
author_association: COLLABORATOR
diff_hunk: @@ -35,6 +37,74 @@ def download_model(download_func, *args, num_retries=3, timeout=180, **kwargs):
     assert False, "Failed to download the model after multiple retries."
 
 
+def get_cache_dir() -> str:
+    """Get models cache directory from env var or use local default."""
+    cache_dir = os.environ.get("FORGE_MODELS_HUB")
+    if not cache_dir:
+        cache_dir = os.path.join(os.getcwd(), "cached_forge_models")
+        os.makedirs(cache_dir, exist_ok=True)
+    return cache_dir
+
+
+def default_loader(path: str):
+    """Load model with PyTorch."""
+    try:
+        return torch.load(path, map_location="cpu")
+    except Exception as e:
+        print(f"PyTorch loading error: {e}")
+        return None
+
+
+def yolov5_loader(path: str, variant: str = "ultralytics/yolov5"):
+    try:
+        model = torch.hub.load(variant, "custom", path=path)
+        return model
+    except Exception as e:
+        print(f"YOLOv5 loading error: {e}")
+        return None
+
+
+def fetch_model(
+    model_name: str,
+    url: str,
+    loader: Optional[Callable] = default_loader,
+    max_retries: int = 3,
+    timeout: int = 30,
+    **kwargs: Any,
+) -> FrameworkModule:
+    """Fetch model from URL, cache it, and load it."""
+
+    model_file = model_name + ".pt"
+
+    model_path = os.path.join(get_cache_dir(), model_file)
+
+    # Download if needed
+    if not os.path.exists(model_path):
path: forge/test/utils.py
user: login: vmilosevic
body: Will the "fetch_model" recover in next run if download fails? It looks like that folder will be crated and it will not attempt the download again
created_at: 2025-03-25T09:59:46Z
author_association: COLLABORATOR
diff_hunk: @@ -35,6 +37,74 @@ def download_model(download_func, *args, num_retries=3, timeout=180, **kwargs):
     assert False, "Failed to download the model after multiple retries."
 
 
+def get_cache_dir() -> str:
+    """Get models cache directory from env var or use local default."""
+    cache_dir = os.environ.get("FORGE_MODELS_HUB")
+    if not cache_dir:
+        cache_dir = os.path.join(os.getcwd(), "cached_forge_models")
+        os.makedirs(cache_dir, exist_ok=True)
+    return cache_dir
+
+
+def default_loader(path: str):
+    """Load model with PyTorch."""
+    try:
+        return torch.load(path, map_location="cpu")
+    except Exception as e:
+        print(f"PyTorch loading error: {e}")
+        return None
+
+
+def yolov5_loader(path: str, variant: str = "ultralytics/yolov5"):
+    try:
+        model = torch.hub.load(variant, "custom", path=path)
+        return model
+    except Exception as e:
+        print(f"YOLOv5 loading error: {e}")
+        return None
+
+
+def fetch_model(
+    model_name: str,
+    url: str,
+    loader: Optional[Callable] = default_loader,
+    max_retries: int = 3,
+    timeout: int = 30,
+    **kwargs: Any,
+) -> FrameworkModule:
+    """Fetch model from URL, cache it, and load it."""
+
+    model_file = model_name + ".pt"
+
+    model_path = os.path.join(get_cache_dir(), model_file)
+
+    # Download if needed
+    if not os.path.exists(model_path):
+        os.makedirs(os.path.dirname(model_path), exist_ok=True)
+
+        for attempt in range(1, max_retries + 1):
path: forge/test/utils.py
user: login: vmilosevic
body: We are using shared (nfs) folder for file download we need to make sure that we handle concurrent "fetch_model" calls from multiple runners. Look into acquiring file lock before downloading
created_at: 2025-03-25T10:09:29Z
author_association: COLLABORATOR
diff_hunk: @@ -35,6 +37,74 @@ def download_model(download_func, *args, num_retries=3, timeout=180, **kwargs):
     assert False, "Failed to download the model after multiple retries."
 
 
+def get_cache_dir() -> str:
+    """Get models cache directory from env var or use local default."""
+    cache_dir = os.environ.get("FORGE_MODELS_HUB")
+    if not cache_dir:
+        cache_dir = os.path.join(os.getcwd(), "cached_forge_models")
+        os.makedirs(cache_dir, exist_ok=True)
+    return cache_dir
+
+
+def default_loader(path: str):
+    """Load model with PyTorch."""
+    try:
+        return torch.load(path, map_location="cpu")
+    except Exception as e:
+        print(f"PyTorch loading error: {e}")
+        return None
+
+
+def yolov5_loader(path: str, variant: str = "ultralytics/yolov5"):
+    try:
+        model = torch.hub.load(variant, "custom", path=path)
+        return model
+    except Exception as e:
+        print(f"YOLOv5 loading error: {e}")
+        return None
+
+
+def fetch_model(
+    model_name: str,
+    url: str,
+    loader: Optional[Callable] = default_loader,
+    max_retries: int = 3,
+    timeout: int = 30,
+    **kwargs: Any,
+) -> FrameworkModule:
+    """Fetch model from URL, cache it, and load it."""
+
+    model_file = model_name + ".pt"
+
+    model_path = os.path.join(get_cache_dir(), model_file)
+
+    # Download if needed
+    if not os.path.exists(model_path):
+        os.makedirs(os.path.dirname(model_path), exist_ok=True)
+
+        for attempt in range(1, max_retries + 1):
path: forge/test/utils.py
user: login: vmilosevic
body: https://chatgpt.com/share/67e28097-819c-8000-b35b-52f38087ca24
created_at: 2025-03-25T10:09:51Z
author_association: COLLABORATOR
diff_hunk: @@ -35,6 +37,74 @@ def download_model(download_func, *args, num_retries=3, timeout=180, **kwargs):
     assert False, "Failed to download the model after multiple retries."
 
 
+def get_cache_dir() -> str:
+    """Get models cache directory from env var or use local default."""
+    cache_dir = os.environ.get("FORGE_MODELS_HUB")
+    if not cache_dir:
+        cache_dir = os.path.join(os.getcwd(), "cached_forge_models")
+        os.makedirs(cache_dir, exist_ok=True)
+    return cache_dir
+
+
+def default_loader(path: str):
+    """Load model with PyTorch."""
+    try:
+        return torch.load(path, map_location="cpu")
+    except Exception as e:
+        print(f"PyTorch loading error: {e}")
+        return None
+
+
+def yolov5_loader(path: str, variant: str = "ultralytics/yolov5"):
+    try:
+        model = torch.hub.load(variant, "custom", path=path)
+        return model
+    except Exception as e:
+        print(f"YOLOv5 loading error: {e}")
+        return None
+
+
+def fetch_model(
+    model_name: str,
+    url: str,
+    loader: Optional[Callable] = default_loader,
+    max_retries: int = 3,
+    timeout: int = 30,
+    **kwargs: Any,
+) -> FrameworkModule:
+    """Fetch model from URL, cache it, and load it."""
+
+    model_file = model_name + ".pt"
+
+    model_path = os.path.join(get_cache_dir(), model_file)
+
+    # Download if needed
+    if not os.path.exists(model_path):
path: forge/test/utils.py
user: login: vkovinicTT
body: You are right, I changed that now.
created_at: 2025-03-25T10:31:09Z
author_association: MEMBER
diff_hunk: @@ -35,6 +37,74 @@ def download_model(download_func, *args, num_retries=3, timeout=180, **kwargs):
     assert False, "Failed to download the model after multiple retries."
 
 
+def get_cache_dir() -> str:
+    """Get models cache directory from env var or use local default."""
+    cache_dir = os.environ.get("FORGE_MODELS_HUB")
+    if not cache_dir:
+        cache_dir = os.path.join(os.getcwd(), "cached_forge_models")
+        os.makedirs(cache_dir, exist_ok=True)
+    return cache_dir
+
+
+def default_loader(path: str):
+    """Load model with PyTorch."""
+    try:
+        return torch.load(path, map_location="cpu")
+    except Exception as e:
+        print(f"PyTorch loading error: {e}")
+        return None
+
+
+def yolov5_loader(path: str, variant: str = "ultralytics/yolov5"):
+    try:
+        model = torch.hub.load(variant, "custom", path=path)
+        return model
+    except Exception as e:
+        print(f"YOLOv5 loading error: {e}")
+        return None
+
+
+def fetch_model(
+    model_name: str,
+    url: str,
+    loader: Optional[Callable] = default_loader,
+    max_retries: int = 3,
+    timeout: int = 30,
+    **kwargs: Any,
+) -> FrameworkModule:
+    """Fetch model from URL, cache it, and load it."""
+
+    model_file = model_name + ".pt"
+
+    model_path = os.path.join(get_cache_dir(), model_file)
+
+    # Download if needed
+    if not os.path.exists(model_path):
+        os.makedirs(os.path.dirname(model_path), exist_ok=True)
+
+        for attempt in range(1, max_retries + 1):
path: forge/test/utils.py
user: login: vkovinicTT
body: I've made the change already, thanks for the tips though :))
created_at: 2025-03-25T13:02:03Z
author_association: MEMBER
diff_hunk: @@ -35,6 +38,95 @@ def download_model(download_func, *args, num_retries=3, timeout=180, **kwargs):
     assert False, "Failed to download the model after multiple retries."
 
 
+def get_cache_dir() -> str:
+    """Get models cache directory from env var or use local default."""
+    cache_dir = os.environ.get("FORGE_MODELS_CACHE")
+    if not cache_dir:
+        cache_dir = os.path.join(os.getcwd(), ".forge_models_cache")
+        os.makedirs(cache_dir, exist_ok=True)
+    return cache_dir
+
+
+def default_loader(path: str):
+    """Load model with PyTorch."""
+    try:
+        return torch.load(path, map_location="cpu")
+    except Exception as e:
+        print(f"PyTorch loading error: {e}")
+        return None
+
+
+def yolov5_loader(path: str, variant: str = "ultralytics/yolov5"):
+    try:
+        model = torch.hub.load(variant, "custom", path=path)
+        return model
+    except Exception as e:
+        print(f"YOLOv5 loading error: {e}")
+        return None
+
+
+def fetch_model(
+    model_name: str,
+    url: str,
+    loader: Optional[Callable] = default_loader,
+    max_retries: int = 3,
+    timeout: int = 30,
+    **kwargs: Any,
+) -> FrameworkModule:
+    """Fetch model from URL, cache it, and load it."""
+
+    model_file = model_name + ".pt"
+
+    model_path = os.path.join(get_cache_dir(), model_file)
+
+    # Download if needed
+    if not os.path.exists(model_path):
+        os.makedirs(os.path.dirname(model_path), exist_ok=True)
+        lock_path = model_path + ".lock"
+        lock = FileLock(lock_path)
+
+        with lock:
+            # Check again after acquiring the lock to handle concurrent processes
+            if not os.path.exists(model_path):
+                for attempt in range(1, max_retries + 1):
+                    try:
+                        print(f"Downloading {model_name}, attempt {attempt}/{max_retries}...")
+                        response = requests.get(url, timeout=timeout)
+                        response.raise_for_status()
+
+                        # Write to temporary file first
+                        temp_path = model_path + ".tmp"
+                        with open(temp_path, "wb") as f:
+                            f.write(response.content)
+
+                        # Atomic rename after successful download
path: forge/test/utils.py
user: login: pilkicTT
body: Why the rename step? why not directly writing it to the desired file?
created_at: 2025-03-25T13:10:03Z
author_association: CONTRIBUTOR
diff_hunk: @@ -35,6 +38,95 @@ def download_model(download_func, *args, num_retries=3, timeout=180, **kwargs):
     assert False, "Failed to download the model after multiple retries."
 
 
+def get_cache_dir() -> str:
+    """Get models cache directory from env var or use local default."""
+    cache_dir = os.environ.get("FORGE_MODELS_CACHE")
+    if not cache_dir:
+        cache_dir = os.path.join(os.getcwd(), ".forge_models_cache")
+        os.makedirs(cache_dir, exist_ok=True)
+    return cache_dir
+
+
+def default_loader(path: str):
+    """Load model with PyTorch."""
+    try:
+        return torch.load(path, map_location="cpu")
+    except Exception as e:
+        print(f"PyTorch loading error: {e}")
+        return None
+
+
+def yolov5_loader(path: str, variant: str = "ultralytics/yolov5"):
+    try:
+        model = torch.hub.load(variant, "custom", path=path)
+        return model
+    except Exception as e:
+        print(f"YOLOv5 loading error: {e}")
+        return None
+
+
+def fetch_model(
+    model_name: str,
+    url: str,
+    loader: Optional[Callable] = default_loader,
+    max_retries: int = 3,
+    timeout: int = 30,
+    **kwargs: Any,
+) -> FrameworkModule:
+    """Fetch model from URL, cache it, and load it."""
+
+    model_file = model_name + ".pt"
+
+    model_path = os.path.join(get_cache_dir(), model_file)
+
+    # Download if needed
+    if not os.path.exists(model_path):
+        os.makedirs(os.path.dirname(model_path), exist_ok=True)
+        lock_path = model_path + ".lock"
+        lock = FileLock(lock_path)
+
+        with lock:
+            # Check again after acquiring the lock to handle concurrent processes
+            if not os.path.exists(model_path):
+                for attempt in range(1, max_retries + 1):
+                    try:
+                        print(f"Downloading {model_name}, attempt {attempt}/{max_retries}...")
+                        response = requests.get(url, timeout=timeout)
+                        response.raise_for_status()
+
+                        # Write to temporary file first
+                        temp_path = model_path + ".tmp"
+                        with open(temp_path, "wb") as f:
+                            f.write(response.content)
+
+                        # Atomic rename after successful download
path: forge/test/utils.py
user: login: vkovinicTT
body: This handles cases where a process might be terminated unexpectedly (e.g., by OOM killer, system shutdown, or network interruption) during download.
For example, if process A starts downloading a large model file and is killed halfway through, process B might see the incomplete file as valid and try to use it, causing errors. With temp file, process B would correctly detect that the download needs to be attempted again.
created_at: 2025-03-25T13:53:16Z
author_association: MEMBER
diff_hunk: @@ -123,6 +123,7 @@ jobs:
       env:
         HF_TOKEN: ${{ secrets.HF_TOKEN }}
         HF_HOME: /mnt/dockercache/huggingface
+        FORGE_MODELS_CACHE: /mnt/dockercache/forge_models_cache
path: .github/workflows/test.yml
user: login: vmilosevic
body: If this is cache specific to tt-forge-fe project please set path to 
FORGE_MODELS_CACHE: /mnt/dockercache/tt-forge-fe/models_cache
created_at: 2025-03-25T14:56:36Z
author_association: COLLABORATOR

PR comments:
user: login: codecov-commenter
created_at: 2025-03-17T15:52:41Z
body: ## [Codecov](https://app.codecov.io/gh/tenstorrent/tt-forge-fe/pull/1467?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=tenstorrent) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 43.49%. Comparing base [(`6f6614d`)](https://app.codecov.io/gh/tenstorrent/tt-forge-fe/commit/6f6614da4128d89ad668c26e1f3609bd629ce5c9?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=tenstorrent) to head [(`4445581`)](https://app.codecov.io/gh/tenstorrent/tt-forge-fe/commit/444558105db3a477524179d25d52ef673cddfb2b?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=tenstorrent).

<details><summary>Additional details and impacted files</summary>


```diff
@@           Coverage Diff           @@
##             main    #1467   +/-   ##
=======================================
  Coverage   43.49%   43.49%           
=======================================
  Files          49       49           
  Lines        7843     7843           
=======================================
  Hits         3411     3411           
  Misses       4432     4432           
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/tenstorrent/tt-forge-fe/pull/1467?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=tenstorrent).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=tenstorrent).

author_association: NONE
user: login: github-actions[bot]
created_at: 2025-03-17T16:19:44Z
body: <table><tr><th><th>Tests</th><th>Passed ✅</th><th>Skipped ⚠️</th><th>Failed</th></tr><tr><td>TT-Forge-FE Tests</td><td>640 ran</td><td>499 passed</td><td>141 skipped</td><td>0 failed</td></tr></table>

<table><tr><th>Test</th><th>Result</th></tr><tr><td colspan="2">No test annotations available</td></tr></table>

<!-- Summary comment for ["TT-Forge-FE Tests"] by mikepenz/action-junit-report -->
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-03-17T16:23:14Z
body: <table><tr><th><th>Tests</th><th>Passed ✅</th><th>Skipped ⚠️</th><th>Failed</th></tr><tr><td>TT-Forge-FE Tests</td><td>640 ran</td><td>499 passed</td><td>141 skipped</td><td>0 failed</td></tr></table>

<table><tr><th>Test</th><th>Result</th></tr><tr><td colspan="2">No test annotations available</td></tr></table>

<!-- Summary comment for ["TT-Forge-FE Tests"] by mikepenz/action-junit-report -->
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-03-17T16:33:14Z
body: <table><tr><th><th>Tests</th><th>Passed ✅</th><th>Skipped ⚠️</th><th>Failed</th></tr><tr><td>TT-Forge-FE Tests</td><td>699 ran</td><td>561 passed</td><td>138 skipped</td><td>0 failed</td></tr></table>

<table><tr><th>Test</th><th>Result</th></tr><tr><td colspan="2">No test annotations available</td></tr></table>

<!-- Summary comment for ["TT-Forge-FE Tests"] by mikepenz/action-junit-report -->
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-03-17T16:37:35Z
body: <table><tr><th><th>Tests</th><th>Passed ✅</th><th>Skipped ⚠️</th><th>Failed</th></tr><tr><td>TT-Forge-FE Tests</td><td>699 ran</td><td>561 passed</td><td>138 skipped</td><td>0 failed</td></tr></table>

<table><tr><th>Test</th><th>Result</th></tr><tr><td colspan="2">No test annotations available</td></tr></table>

<!-- Summary comment for ["TT-Forge-FE Tests"] by mikepenz/action-junit-report -->
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-03-24T09:09:45Z
body: <table><tr><th><th>Tests</th><th>Passed ✅</th><th>Skipped ⚠️</th><th>Failed</th></tr><tr><td>TT-Forge-FE Tests</td><td>640 ran</td><td>492 passed</td><td>148 skipped</td><td>0 failed</td></tr></table>

<table><tr><th>Test</th><th>Result</th></tr><tr><td colspan="2">No test annotations available</td></tr></table>

<!-- Summary comment for ["TT-Forge-FE Tests"] by mikepenz/action-junit-report -->
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-03-24T09:25:56Z
body: <table><tr><th><th>Tests</th><th>Passed ✅</th><th>Skipped ⚠️</th><th>Failed</th></tr><tr><td>TT-Forge-FE Tests</td><td>640 ran</td><td>492 passed</td><td>148 skipped</td><td>0 failed</td></tr></table>

<table><tr><th>Test</th><th>Result</th></tr><tr><td colspan="2">No test annotations available</td></tr></table>

<!-- Summary comment for ["TT-Forge-FE Tests"] by mikepenz/action-junit-report -->
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-03-24T09:27:02Z
body: <table><tr><th><th>Tests</th><th>Passed ✅</th><th>Skipped ⚠️</th><th>Failed</th></tr><tr><td>TT-Forge-FE Tests</td><td>699 ran</td><td>555 passed</td><td>144 skipped</td><td>0 failed</td></tr></table>

<table><tr><th>Test</th><th>Result</th></tr><tr><td colspan="2">No test annotations available</td></tr></table>

<!-- Summary comment for ["TT-Forge-FE Tests"] by mikepenz/action-junit-report -->
author_association: NONE
user: login: vkovinicTT
created_at: 2025-03-24T11:49:39Z
body: > How is this caching anything? You should look into `/mnt/dockercache` which we use in the CI for caching (for now it looks like we only cache models from huggingface). We should do something similar for this.

Sorry, you are right, we are not caching anything, instead we are just downloading model everytime we run the test.
author_association: MEMBER
user: login: github-actions[bot]
created_at: 2025-03-25T09:06:04Z
body: <table><tr><th><th>Tests</th><th>Passed ❌️</th><th>Skipped</th><th>Failed</th></tr><tr><td>TT-Forge-FE Tests</td><td>0 ran</td><td>0 passed</td><td>0 skipped</td><td>0 failed</td></tr></table>

<table><tr><th>Test</th><th>Result</th></tr><tr><td colspan="2">No test annotations available</td></tr></table>

<!-- Summary comment for ["TT-Forge-FE Tests"] by mikepenz/action-junit-report -->
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-03-25T12:49:45Z
body: <table><tr><th><th>Tests</th><th>Passed ✅</th><th>Skipped ⚠️</th><th>Failed</th></tr><tr><td>TT-Forge-FE Tests</td><td>700 ran</td><td>555 passed</td><td>145 skipped</td><td>0 failed</td></tr></table>

<table><tr><th>Test</th><th>Result</th></tr><tr><td colspan="2">No test annotations available</td></tr></table>

<!-- Summary comment for ["TT-Forge-FE Tests"] by mikepenz/action-junit-report -->
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-03-25T13:03:51Z
body: <table><tr><th><th>Tests</th><th>Passed ✅</th><th>Skipped ⚠️</th><th>Failed</th></tr><tr><td>TT-Forge-FE Tests</td><td>642 ran</td><td>495 passed</td><td>147 skipped</td><td>0 failed</td></tr></table>

<table><tr><th>Test</th><th>Result</th></tr><tr><td colspan="2">No test annotations available</td></tr></table>

<!-- Summary comment for ["TT-Forge-FE Tests"] by mikepenz/action-junit-report -->
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-03-25T14:02:32Z
body: <table><tr><th><th>Tests</th><th>Passed ✅</th><th>Skipped ⚠️</th><th>Failed</th></tr><tr><td>TT-Forge-FE Tests</td><td>700 ran</td><td>555 passed</td><td>145 skipped</td><td>0 failed</td></tr></table>

<table><tr><th>Test</th><th>Result</th></tr><tr><td colspan="2">No test annotations available</td></tr></table>

<!-- Summary comment for ["TT-Forge-FE Tests"] by mikepenz/action-junit-report -->
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-03-25T15:23:43Z
body: <table><tr><th><th>Tests</th><th>Passed ✅</th><th>Skipped ⚠️</th><th>Failed</th></tr><tr><td>TT-Forge-FE Tests</td><td>700 ran</td><td>555 passed</td><td>145 skipped</td><td>0 failed</td></tr></table>

<table><tr><th>Test</th><th>Result</th></tr><tr><td colspan="2">No test annotations available</td></tr></table>

<!-- Summary comment for ["TT-Forge-FE Tests"] by mikepenz/action-junit-report -->
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-03-25T15:32:42Z
body: <table><tr><th><th>Tests</th><th>Passed ✅</th><th>Skipped ⚠️</th><th>Failed</th></tr><tr><td>TT-Forge-FE Tests</td><td>642 ran</td><td>495 passed</td><td>147 skipped</td><td>0 failed</td></tr></table>

<table><tr><th>Test</th><th>Result</th></tr><tr><td colspan="2">No test annotations available</td></tr></table>

<!-- Summary comment for ["TT-Forge-FE Tests"] by mikepenz/action-junit-report -->
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-03-25T16:06:17Z
body: <table><tr><th><th>Tests</th><th>Passed ✅</th><th>Skipped ⚠️</th><th>Failed</th></tr><tr><td>TT-Forge-FE Tests</td><td>642 ran</td><td>495 passed</td><td>147 skipped</td><td>0 failed</td></tr></table>

<table><tr><th>Test</th><th>Result</th></tr><tr><td colspan="2">No test annotations available</td></tr></table>

<!-- Summary comment for ["TT-Forge-FE Tests"] by mikepenz/action-junit-report -->
author_association: NONE
user: login: github-actions[bot]
created_at: 2025-03-25T16:23:25Z
body: <table><tr><th><th>Tests</th><th>Passed ✅</th><th>Skipped ⚠️</th><th>Failed</th></tr><tr><td>TT-Forge-FE Tests</td><td>700 ran</td><td>555 passed</td><td>145 skipped</td><td>0 failed</td></tr></table>

<table><tr><th>Test</th><th>Result</th></tr><tr><td colspan="2">No test annotations available</td></tr></table>

<!-- Summary comment for ["TT-Forge-FE Tests"] by mikepenz/action-junit-report -->
author_association: NONE