Resolve run teardown on individual slice retry exhaustion (#432)

Author: HudsonGraeme

Cargo.lock

@@ -50,7 +50,7 @@ futures-util = "0.3"
 shellexpand = "3.1.2"
 openssl = { version = "0.10", features = ["vendored"] }
 flate2 = "1"
-dsperse = { git = "https://github.com/inference-labs-inc/dsperse", rev = "176ed08ee9552ac23f4a6470d1e7b53b67305b08" }
+dsperse = { git = "https://github.com/inference-labs-inc/dsperse", rev = "3d1db73f" }
 ndarray = { version = "0.17", features = ["serde"] }
 zip = { version = "2", default-features = false, features = ["deflate"] }
 bittensor-drand = { git = "https://github.com/inference-labs-inc/bittensor-drand.git", rev = "e55ed98", default-features = false }

Cargo.toml

@@ -189,10 +189,6 @@ impl IncrementalRunManager {
         self.evicted.contains(run_uid)
     }
 
-    pub fn get_run_source(&self, run_uid: &str) -> Option<RunSource> {
-        self.runs.get(run_uid).map(|r| r.run_source)
-    }
-
     pub fn all_circuit_work(&self, run_uid: &str) -> anyhow::Result<Vec<SliceWork>> {
         let run = self
             .runs
@@ -323,6 +319,34 @@ impl IncrementalRunManager {
         false
     }
 
+    pub fn mark_slice_failed(&mut self, run_uid: &str, slice_id: &str) -> usize {
+        self.tile_counters
+            .remove(&(run_uid.to_string(), slice_id.to_string()));
+        if let Some(run) = self.runs.get_mut(run_uid) {
+            run.last_activity = Instant::now();
+            if let Some(ref mut combined) = run.combined {
+                combined.mark_slice_failed(slice_id);
+                return combined.failed_count();
+            }
+        }
+        0
+    }
+
+    pub fn is_slice_failed(&self, run_uid: &str, slice_id: &str) -> bool {
+        self.runs
+            .get(run_uid)
+            .and_then(|r| r.combined.as_ref())
+            .is_some_and(|c| c.is_slice_failed(slice_id))
+    }
+
+    pub fn failed_slice_count(&self, run_uid: &str) -> usize {
+        self.runs
+            .get(run_uid)
+            .and_then(|r| r.combined.as_ref())
+            .map(|c| c.failed_count())
+            .unwrap_or(0)
+    }
+
     pub fn is_run_complete(&self, run_uid: &str) -> bool {
         self.runs
             .get(run_uid)

crates/sn2-validator/src/incremental_runner.rs

@@ -377,7 +377,8 @@ impl ValidatorLoop {
         self.dslice_input_scales
             .retain(|(uid, _), _| uid != run_uid);
 
-        info!(run_uid = %run_uid, "combined run complete");
+        let failed_count = self.run_manager.failed_slice_count(run_uid);
+        info!(run_uid = %run_uid, failed_count, "combined run complete");
 
         let final_output = self.run_manager.final_output_json(run_uid);
         let mut active_run = self.run_manager.remove_run(run_uid);
@@ -389,7 +390,7 @@ impl ValidatorLoop {
 
         let relay_output = final_output.clone();
         self.spawn_artifact_upload(run_uid, &mut active_run, final_output);
-        self.notify_run_completed(run_uid, &active_run, relay_output)
+        self.notify_run_completed(run_uid, &active_run, relay_output, failed_count)
             .await;
     }
 

crates/sn2-validator/src/validator_loop/dslice.rs

@@ -125,26 +125,41 @@ impl ValidatorLoop {
         run_uid: &str,
         active_run: &Option<crate::incremental_runner::ActiveRun>,
         final_output: Option<serde_json::Value>,
+        failed_count: usize,
     ) {
         let notify_circuit_id = active_run
             .as_ref()
             .map(|r| r.circuit_id.as_str())
             .unwrap_or_default()
             .to_string();
-        let mut result = serde_json::json!({"run_uid": run_uid, "status": "complete"});
+        let status = if failed_count > 0 {
+            "partial"
+        } else {
+            "complete"
+        };
+        let mut result = serde_json::json!({"run_uid": run_uid, "status": status});
         if let Some(output) = final_output {
             result["output"] = output;
         }
+        if failed_count > 0 {
+            result["failed_slices"] = serde_json::json!(failed_count);
+        }
         self.relay_set_request_result(run_uid, result).await;
-        self.relay_send_notification(
-            "subnet-2.batch_completed",
-            serde_json::json!({
-                "run_uid": run_uid,
-                "circuit_id": notify_circuit_id,
-                "status": "completed",
-            }),
-        )
-        .await;
+        let notification_status = if failed_count > 0 {
+            "partial"
+        } else {
+            "completed"
+        };
+        let mut notification = serde_json::json!({
+            "run_uid": run_uid,
+            "circuit_id": notify_circuit_id,
+            "status": notification_status,
+        });
+        if failed_count > 0 {
+            notification["failed_slices"] = serde_json::json!(failed_count);
+        }
+        self.relay_send_notification("subnet-2.batch_completed", notification)
+            .await;
     }
 
     pub(super) fn report_dsperse_completion(&self, run: &crate::incremental_runner::ActiveRun) {

crates/sn2-validator/src/validator_loop/relay.rs

@@ -1,5 +1,5 @@
 use sn2_types::*;
-use tracing::{info, warn};
+use tracing::{debug, info, warn};
 
 use super::{event_slice_num, RetryPayload, ValidatorLoop};
 use crate::incremental_runner::SliceArtifact;
@@ -126,27 +126,32 @@ impl ValidatorLoop {
         reason: &str,
     ) {
         if let Some(run_uid) = run_uid {
+            if !self.run_manager.has_run(run_uid) {
+                return;
+            }
+
             if let Some(snum) = slice_num {
-                let ruid = run_uid.clone();
-                let event_snum = event_slice_num(snum, is_tile, tile_idx);
-                let err = reason.to_string();
-                self.emit_event(move |ev| async move {
-                    ev.emit_slice_failed(&ruid, &event_snum, &err).await;
-                });
+                if !self.run_manager.is_slice_failed(run_uid, snum) {
+                    let ruid = run_uid.clone();
+                    let event_snum = event_slice_num(snum, is_tile, tile_idx);
+                    let err = reason.to_string();
+                    self.emit_event(move |ev| async move {
+                        ev.emit_slice_failed(&ruid, &event_snum, &err).await;
+                    });
+
+                    let failed_count = self.run_manager.mark_slice_failed(run_uid, snum);
+                    warn!(
+                        run_uid = %run_uid,
+                        slice = %snum,
+                        failed_count,
+                        "slice max retries exceeded, continuing run"
+                    );
+                }
             }
-            warn!(run_uid = %run_uid, "dslice max retries exceeded, removing run");
-            if self.run_manager.get_run_source(run_uid) == Some(RunSource::Api) {
-                self.relay_set_request_result(
-                    run_uid,
-                    serde_json::json!({
-                        "run_uid": run_uid,
-                        "status": "failed",
-                        "error": "max retries exceeded",
-                    }),
-                )
-                .await;
+
+            if self.run_manager.is_run_complete(run_uid) {
+                self.finalize_combined_run(run_uid).await;
             }
-            self.teardown_run(run_uid).await;
         }
     }
 
@@ -180,6 +185,15 @@ impl ValidatorLoop {
             return;
         }
 
+        if self.run_manager.is_slice_failed(&run_uid, &slice_num) {
+            debug!(
+                run_uid = %run_uid,
+                slice = %slice_num,
+                "ignoring late success for failed slice"
+            );
+            return;
+        }
+
         self.run_manager.push_artifact(
             &run_uid,
             SliceArtifact {