bugfix: wait for model to load before navigating to benchmark screen; use NavigationActions instead of raw navController

2025-04-13 22:14:09 -07:00 · 2025-04-13 22:14:09 -07:00 · 511df35704
parent ea11ee3c94
commit 511df35704
9 changed files with 38 additions and 63 deletions
--- a/examples/llama.android/app/src/main/java/com/example/llama/revamp/MainActivity.kt
+++ b/examples/llama.android/app/src/main/java/com/example/llama/revamp/MainActivity.kt
@ -97,18 +97,16 @@ fun AppContent(

    // Model unloading confirmation
    var showUnloadDialog by remember { mutableStateOf(false) }
-
-    // Helper function to handle back press with model unloading check
    val handleBackWithModelCheck = {
        if (isModelLoading) {
            // If model is still loading, ignore the request
            true // Mark as handled
        } else if (isModelLoaded) {
            showUnloadDialog = true
-            pendingNavigation = { navController.popBackStack() }
+            pendingNavigation = { navigationActions.navigateUp() }
            true // Mark as handled
        } else {
-            navController.popBackStack()
+            navigationActions.navigateUp()
            true // Mark as handled
        }
    }
@ -138,14 +136,9 @@ fun AppContent(
    val drawerGesturesEnabled by remember(currentRoute, drawerState.currentValue) {
        derivedStateOf {
            // Always allow gesture dismissal when drawer is open
-            if (drawerState.currentValue == DrawerValue.Open) {
-                true
-            } else {
-                // Only enable drawer opening by gesture on these screens
-                currentRoute == AppDestinations.MODEL_SELECTION_ROUTE ||
-                    currentRoute == AppDestinations.SETTINGS_GENERAL_ROUTE ||
-                    currentRoute == AppDestinations.MODELS_MANAGEMENT_ROUTE
-            }
+            if (drawerState.currentValue == DrawerValue.Open) true
+            // Only enable drawer opening by gesture on these screens
+            else currentRoute == AppDestinations.MODEL_SELECTION_ROUTE
        }
    }

@ -191,8 +184,6 @@ fun AppContent(
                        navigationActions.navigateToModelsManagement()
                    },
                    onMenuClicked = openDrawer,
-                    drawerState = drawerState,
-                    navigationActions = navigationActions
                )
            }

@ -201,8 +192,21 @@ fun AppContent(
                ModelLoadingScreen(
                    engineState = engineState,
                    onBenchmarkSelected = {
-                        mainVewModel.prepareForBenchmark()
-                        navigationActions.navigateToBenchmark()
+                        // Store a reference to the loading job
+                        val loadingJob = coroutineScope.launch {
+                            mainVewModel.prepareForBenchmark()
+                            // Check if the job wasn't cancelled before navigating
+                            if (isActive) {
+                                navigationActions.navigateToBenchmark()
+                            }
+                        }
+
+
+                        // Update the pendingNavigation handler to cancel any ongoing loading
+                        pendingNavigation = {
+                            loadingJob.cancel()
+                            navigationActions.navigateUp()
+                        }
                    },
                    onConversationSelected = { systemPrompt ->
                        // Store a reference to the loading job
@ -216,15 +220,13 @@ fun AppContent(
                        // Update the pendingNavigation handler to cancel any ongoing loading
                        pendingNavigation = {
                            loadingJob.cancel()
-                            navController.popBackStack()
+                            navigationActions.navigateUp()
                        }
                    },
                    onBackPressed = {
                        // Need to unload model before going back
                        handleBackWithModelCheck()
                    },
-                    drawerState = drawerState,
-                    navigationActions = navigationActions
                )
            }

@ -252,8 +254,6 @@ fun AppContent(
                    onSharePressed = {
                        // Stub for sharing functionality
                    },
-                    drawerState = drawerState,
-                    navigationActions = navigationActions,
                    viewModel = mainVewModel
                )
            }
@ -261,9 +261,7 @@ fun AppContent(
            // Settings General Screen
            composable(AppDestinations.SETTINGS_GENERAL_ROUTE) {
                SettingsGeneralScreen(
-                    onBackPressed = { navController.popBackStack() },
-                    drawerState = drawerState,
-                    navigationActions = navigationActions,
+                    onBackPressed = { navigationActions.navigateUp() },
                    onMenuClicked = openDrawer
                )
            }
@ -271,7 +269,7 @@ fun AppContent(
            // Models Management Screen
            composable(AppDestinations.MODELS_MANAGEMENT_ROUTE) {
                ModelsManagementScreen(
-                    onBackPressed = { navController.popBackStack() },
+                    onBackPressed = { navigationActions.navigateUp() },
                )
            }
        }
--- a/examples/llama.android/app/src/main/java/com/example/llama/revamp/ui/components/AppScaffolds.kt
+++ b/examples/llama.android/app/src/main/java/com/example/llama/revamp/ui/components/AppScaffolds.kt
@ -11,7 +11,6 @@ import androidx.compose.runtime.remember
 import androidx.hilt.navigation.compose.hiltViewModel
 import com.example.llama.revamp.viewmodel.PerformanceViewModel

-// DefaultAppScaffold.kt
@Composable
 fun DefaultAppScaffold(
    title: String,
@ -35,7 +34,6 @@ fun DefaultAppScaffold(
    )
 }

-// PerformanceAppScaffold.kt
@Composable
 fun PerformanceAppScaffold(
    performanceViewModel: PerformanceViewModel = hiltViewModel(),
@ -68,7 +66,6 @@ fun PerformanceAppScaffold(
    )
 }

-// StorageAppScaffold.kt
@Composable
 fun StorageAppScaffold(
    title: String,
--- a/examples/llama.android/app/src/main/java/com/example/llama/revamp/ui/components/TopAppBars.kt
+++ b/examples/llama.android/app/src/main/java/com/example/llama/revamp/ui/components/TopAppBars.kt
@ -26,7 +26,6 @@ import com.example.llama.revamp.monitoring.MemoryMetrics
 import com.example.llama.revamp.monitoring.TemperatureMetrics
 import com.example.llama.revamp.monitoring.TemperatureWarningLevel

-// TopAppBars.kt
@OptIn(ExperimentalMaterial3Api::class)
@Composable
 fun DefaultTopBar(
--- a/examples/llama.android/app/src/main/java/com/example/llama/revamp/ui/screens/BenchmarkScreen.kt
+++ b/examples/llama.android/app/src/main/java/com/example/llama/revamp/ui/screens/BenchmarkScreen.kt
@ -13,10 +13,10 @@ import androidx.compose.foundation.shape.RoundedCornerShape
 import androidx.compose.foundation.verticalScroll
 import androidx.compose.material3.Card
 import androidx.compose.material3.CircularProgressIndicator
-import androidx.compose.material3.DrawerState
 import androidx.compose.material3.MaterialTheme
 import androidx.compose.material3.Text
 import androidx.compose.runtime.Composable
+import androidx.compose.runtime.LaunchedEffect
 import androidx.compose.runtime.collectAsState
 import androidx.compose.runtime.getValue
 import androidx.compose.ui.Alignment
@ -24,7 +24,6 @@ import androidx.compose.ui.Modifier
 import androidx.compose.ui.unit.dp
 import androidx.hilt.navigation.compose.hiltViewModel
 import com.example.llama.revamp.engine.InferenceEngine
-import com.example.llama.revamp.navigation.NavigationActions
 import com.example.llama.revamp.ui.components.PerformanceAppScaffold
 import com.example.llama.revamp.ui.theme.MonospacedTextStyle
 import com.example.llama.revamp.viewmodel.MainViewModel
@ -34,14 +33,16 @@ fun BenchmarkScreen(
    onBackPressed: () -> Unit,
    onRerunPressed: () -> Unit,
    onSharePressed: () -> Unit,
-    drawerState: DrawerState,
-    navigationActions: NavigationActions,
    viewModel: MainViewModel = hiltViewModel()
 ) {
    val engineState by viewModel.engineState.collectAsState()
    val benchmarkResults by viewModel.benchmarkResults.collectAsState()
    val selectedModel by viewModel.selectedModel.collectAsState()

+    LaunchedEffect(selectedModel) {
+        viewModel.runBenchmark()
+    }
+
    PerformanceAppScaffold(
        title = "Chat",
        onNavigateBack = onBackPressed,
--- a/examples/llama.android/app/src/main/java/com/example/llama/revamp/ui/screens/ConversationScreen.kt
+++ b/examples/llama.android/app/src/main/java/com/example/llama/revamp/ui/screens/ConversationScreen.kt
@ -55,11 +55,11 @@ import androidx.compose.ui.Modifier
 import androidx.compose.ui.draw.clip
 import androidx.compose.ui.graphics.Color
 import androidx.compose.ui.graphics.StrokeCap
-import androidx.compose.ui.platform.LocalLifecycleOwner
 import androidx.compose.ui.unit.dp
 import androidx.hilt.navigation.compose.hiltViewModel
 import androidx.lifecycle.Lifecycle
 import androidx.lifecycle.LifecycleEventObserver
+import androidx.lifecycle.compose.LocalLifecycleOwner
 import com.example.llama.revamp.engine.InferenceEngine
 import com.example.llama.revamp.ui.components.PerformanceAppScaffold
 import com.example.llama.revamp.viewmodel.MainViewModel
@ -159,11 +159,12 @@ fun ConversationScreen(
 }

@Composable
-fun AnimatedSystemPrompt(modelName: String?, systemPrompt: String?) {
+fun AnimatedSystemPrompt(
+    modelName: String?,     // TODO-han.yin: add model name into this card, on top of system prompt!
+    systemPrompt: String?
+) {
    var expanded by remember { mutableStateOf(false) }

-    // TODO-han.yin: add model name into this card, on top of system prompt!
-
    if (!systemPrompt.isNullOrBlank()) {
        Card(
            modifier = Modifier
--- a/examples/llama.android/app/src/main/java/com/example/llama/revamp/ui/screens/ModelLoadingScreen.kt
+++ b/examples/llama.android/app/src/main/java/com/example/llama/revamp/ui/screens/ModelLoadingScreen.kt
@ -23,7 +23,6 @@ import androidx.compose.material.icons.filled.Check
 import androidx.compose.material3.Button
 import androidx.compose.material3.Card
 import androidx.compose.material3.CircularProgressIndicator
-import androidx.compose.material3.DrawerState
 import androidx.compose.material3.ExperimentalMaterial3Api
 import androidx.compose.material3.HorizontalDivider
 import androidx.compose.material3.Icon
@ -50,7 +49,6 @@ import androidx.compose.ui.unit.dp
 import androidx.hilt.navigation.compose.hiltViewModel
 import com.example.llama.revamp.data.model.SystemPrompt
 import com.example.llama.revamp.engine.InferenceEngine
-import com.example.llama.revamp.navigation.NavigationActions
 import com.example.llama.revamp.ui.components.PerformanceAppScaffold
 import com.example.llama.revamp.viewmodel.SystemPromptViewModel

@ -66,8 +64,6 @@ fun ModelLoadingScreen(
    onBenchmarkSelected: () -> Unit,
    onConversationSelected: (String?) -> Unit,
    onBackPressed: () -> Unit,
-    drawerState: DrawerState,
-    navigationActions: NavigationActions
 ) {
    val presetPrompts by viewModel.presetPrompts.collectAsState()
    val recentPrompts by viewModel.recentPrompts.collectAsState()
--- a/examples/llama.android/app/src/main/java/com/example/llama/revamp/ui/screens/ModelSelectionScreen.kt
+++ b/examples/llama.android/app/src/main/java/com/example/llama/revamp/ui/screens/ModelSelectionScreen.kt
@ -14,7 +14,6 @@ import androidx.compose.material.icons.Icons
 import androidx.compose.material.icons.filled.PlayArrow
 import androidx.compose.material3.Card
 import androidx.compose.material3.CardDefaults
-import androidx.compose.material3.DrawerState
 import androidx.compose.material3.ExperimentalMaterial3Api
 import androidx.compose.material3.Icon
 import androidx.compose.material3.IconButton
@ -27,7 +26,6 @@ import androidx.compose.ui.Alignment
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.unit.dp
 import com.example.llama.revamp.data.model.ModelInfo
-import com.example.llama.revamp.navigation.NavigationActions
 import com.example.llama.revamp.ui.components.PerformanceAppScaffold
 import java.text.SimpleDateFormat
 import java.util.Date
@ -39,8 +37,6 @@ fun ModelSelectionScreen(
    onModelSelected: (ModelInfo) -> Unit,
    onManageModelsClicked: () -> Unit,
    onMenuClicked: () -> Unit,
-    drawerState: DrawerState,
-    navigationActions: NavigationActions
 ) {
    // For demo purposes, we'll use sample models
    val models = remember { ModelInfo.getSampleModels() }
--- a/examples/llama.android/app/src/main/java/com/example/llama/revamp/ui/screens/SettingsGeneralScreen.kt
+++ b/examples/llama.android/app/src/main/java/com/example/llama/revamp/ui/screens/SettingsGeneralScreen.kt
@ -10,7 +10,6 @@ import androidx.compose.foundation.layout.padding
 import androidx.compose.foundation.rememberScrollState
 import androidx.compose.foundation.verticalScroll
 import androidx.compose.material3.Card
-import androidx.compose.material3.DrawerState
 import androidx.compose.material3.HorizontalDivider
 import androidx.compose.material3.MaterialTheme
 import androidx.compose.material3.Switch
@ -22,7 +21,6 @@ import androidx.compose.ui.Alignment
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.unit.dp
 import androidx.hilt.navigation.compose.hiltViewModel
-import com.example.llama.revamp.navigation.NavigationActions
 import com.example.llama.revamp.ui.components.DefaultAppScaffold
 import com.example.llama.revamp.viewmodel.PerformanceViewModel

@ -33,8 +31,6 @@ import com.example.llama.revamp.viewmodel.PerformanceViewModel
 fun SettingsGeneralScreen(
    performanceViewModel: PerformanceViewModel = hiltViewModel(),
    onBackPressed: () -> Unit,
-    drawerState: DrawerState,
-    navigationActions: NavigationActions,
    onMenuClicked: () -> Unit
 ) {
    // Collect state from ViewModel
--- a/examples/llama.android/app/src/main/java/com/example/llama/revamp/viewmodel/MainViewModel.kt
+++ b/examples/llama.android/app/src/main/java/com/example/llama/revamp/viewmodel/MainViewModel.kt
@ -57,30 +57,21 @@ class MainViewModel @Inject constructor (
    /**
     * Prepares the engine for benchmark mode.
     */
-    fun prepareForBenchmark() {
-        viewModelScope.launch {
-            _selectedModel.value?.let { model ->
-                inferenceEngine.loadModel(model.path)
-                runBenchmark()
-            }
+    suspend fun prepareForBenchmark() {
+        _selectedModel.value?.let { model ->
+            inferenceEngine.loadModel(model.path)
        }
    }

    /**
     * Runs the benchmark with current parameters.
     */
-    private suspend fun runBenchmark() {
-        inferenceEngine.bench(512, 128, 1, 3)
-    }
+    suspend fun runBenchmark() = inferenceEngine.bench(512, 128, 1, 3)

    /**
     * Reruns the benchmark.
     */
-    fun rerunBenchmark() {
-        viewModelScope.launch {
-            runBenchmark()
-        }
-    }
+    fun rerunBenchmark() = viewModelScope.launch { runBenchmark() }

    /**
     * Prepares the engine for conversation mode.