feature: check for available space before copying local model

This commit is contained in:
Han Yin 2025-04-17 11:42:57 -07:00
parent 65d4a57a8b
commit 72e97b93c5
2 changed files with 23 additions and 1 deletions

View File

@ -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

View File

@ -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",