From 72e97b93c57d9c6f704d07b6654ef550235dfc7b Mon Sep 17 00:00:00 2001
From: Han Yin <han.yin@arm.com>
Date: Thu, 17 Apr 2025 11:42:57 -0700
Subject: [PATCH] feature: check for available space before copying local model

---
 .../revamp/data/repository/ModelRepository.kt | 19 ++++++++++++++++++-
 .../viewmodel/ModelsManagementViewModel.kt    |  5 +++++
 2 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/examples/llama.android/app/src/main/java/com/example/llama/revamp/data/repository/ModelRepository.kt b/examples/llama.android/app/src/main/java/com/example/llama/revamp/data/repository/ModelRepository.kt
index bc76ad8f9a..24ee752658 100644
--- a/examples/llama.android/app/src/main/java/com/example/llama/revamp/data/repository/ModelRepository.kt
+++ b/examples/llama.android/app/src/main/java/com/example/llama/revamp/data/repository/ModelRepository.kt
@@ -13,6 +13,7 @@ import com.example.llama.revamp.util.copyWithChannels
 import com.example.llama.revamp.util.extractModelTypeFromFilename
 import com.example.llama.revamp.util.extractParametersFromFilename
 import com.example.llama.revamp.util.extractQuantizationFromFilename
+import com.example.llama.revamp.util.formatSize
 import com.example.llama.revamp.util.getFileNameFromUri
 import com.example.llama.revamp.util.getFileSizeFromUri
 import dagger.hilt.android.qualifiers.ApplicationContext
@@ -79,6 +80,8 @@ interface ModelRepository {
     suspend fun deleteModels(modelIds: List<String>)
 }
 
+class InsufficientStorageException(message: String) : IOException(message)
+
 @Singleton
 class ModelRepositoryImpl @Inject constructor(
     @ApplicationContext private val context: Context,
@@ -131,8 +134,14 @@ class ModelRepositoryImpl @Inject constructor(
             throw IllegalStateException("Another import is already in progress!")
         }
 
-        val fileName = name ?: getFileNameFromUri(context, uri) ?: throw FileNotFoundException("Filename N/A")
         val fileSize = size ?: getFileSizeFromUri(context, uri) ?: throw FileNotFoundException("File size N/A")
+        if (!hasEnoughSpaceForImport(fileSize)) {
+            throw InsufficientStorageException(
+                "Not enough storage space. Required: ${formatSize(fileSize)}, Available: ${formatSize(availableSpaceBytes)}"
+            )
+        }
+
+        val fileName = name ?: getFileNameFromUri(context, uri) ?: throw FileNotFoundException("Filename N/A")
         val modelFile = File(modelsDir, fileName)
 
         importJob = coroutineContext[Job]
@@ -221,6 +230,13 @@ class ModelRepositoryImpl @Inject constructor(
         }
     }
 
+    // Add this method to ModelRepositoryImpl.kt
+    private fun hasEnoughSpaceForImport(fileSize: Long): Boolean {
+        val availableSpace = availableSpaceBytes
+        val requiredSpace = (fileSize * MODEL_IMPORT_SPACE_BUFFER_SCALE ).toLong()
+        return availableSpace >= requiredSpace
+    }
+
     override suspend fun cancelImport(): Boolean? = withContext(Dispatchers.IO) {
         val job = importJob
         val file = currentModelFile
@@ -312,6 +328,7 @@ class ModelRepositoryImpl @Inject constructor(
         private const val STORAGE_METRICS_UPDATE_INTERVAL = 5_000L
         private const val BYTES_IN_GB = 1024f * 1024f * 1024f
 
+        private const val MODEL_IMPORT_SPACE_BUFFER_SCALE = 1.2f
         private const val LARGE_MODEL_THRESHOLD_SIZE = 1024 * 1024 * 1024
         private const val NIO_BUFFER_SIZE = 32 * 1024 * 1024
         private const val NIO_YIELD_SIZE = 128 * 1024 * 1024
diff --git a/examples/llama.android/app/src/main/java/com/example/llama/revamp/viewmodel/ModelsManagementViewModel.kt b/examples/llama.android/app/src/main/java/com/example/llama/revamp/viewmodel/ModelsManagementViewModel.kt
index 94f7d32bae..9a6851c43b 100644
--- a/examples/llama.android/app/src/main/java/com/example/llama/revamp/viewmodel/ModelsManagementViewModel.kt
+++ b/examples/llama.android/app/src/main/java/com/example/llama/revamp/viewmodel/ModelsManagementViewModel.kt
@@ -5,6 +5,7 @@ import android.net.Uri
 import androidx.lifecycle.ViewModel
 import androidx.lifecycle.viewModelScope
 import com.example.llama.revamp.data.model.ModelInfo
+import com.example.llama.revamp.data.repository.InsufficientStorageException
 import com.example.llama.revamp.data.repository.ModelRepository
 import com.example.llama.revamp.data.repository.StorageMetrics
 import com.example.llama.revamp.util.getFileNameFromUri
@@ -107,6 +108,10 @@ class ModelsManagementViewModel @Inject constructor(
                 _managementState.value = Importation.Importing(progress, fileName, fileSize)
             }
             _managementState.value = Importation.Success(model)
+        } catch (e: InsufficientStorageException) {
+            _managementState.value = Importation.Error(
+                message = e.message ?: "Insufficient storage space to import $uri",
+            )
         } catch (e: Exception) {
             _managementState.value = Importation.Error(
                 message = e.message ?: "Unknown error importing $uri",