test(e2e): run gpu workloads from manifest

Signed-off-by: Evan Lezar <elezar@nvidia.com>

Author: elezar

e2e/rust/Cargo.lock

@@ -109,7 +109,9 @@ sha1 = "0.10"
 sha2 = "0.10"
 hex = "0.4"
 rand = "0.9"
+serde = { version = "1", features = ["derive"] }
 serde_json = "1"
+serde_yaml = "0.9"
 
 [lints.rust]
 unsafe_code = "warn"

e2e/rust/Cargo.toml

@@ -11,11 +11,12 @@ set -euo pipefail
 ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
 E2E_TEST="${OPENSHELL_E2E_DOCKER_TEST:-smoke}"
 E2E_FEATURES="${OPENSHELL_E2E_DOCKER_FEATURES:-e2e,e2e-docker}"
+DEFAULT_WORKLOAD_MANIFEST="${ROOT}/e2e/gpu/images/.build/workloads.yaml"
 
 cargo build -p openshell-cli --features openshell-core/dev-settings
 
-if [ "${E2E_TEST}" = "gpu" ] && [ -z "${OPENSHELL_E2E_GPU_CUDA_WORKLOAD_IMAGE:-}" ]; then
-  echo "note: running GPU e2e without OPENSHELL_E2E_GPU_CUDA_WORKLOAD_IMAGE; CUDA workload validation will log an explicit skip"
+if [ "${E2E_TEST}" = "gpu" ] && [ -z "${OPENSHELL_E2E_WORKLOAD_MANIFEST:-}" ] && [ ! -f "${DEFAULT_WORKLOAD_MANIFEST}" ]; then
+  echo "note: running GPU e2e without a workload manifest; workload validation will log an explicit skip. Build one with 'mise run e2e:workloads:build' or set OPENSHELL_E2E_WORKLOAD_MANIFEST."
 fi
 
 exec "${ROOT}/e2e/with-docker-gateway.sh" \

e2e/rust/e2e-docker.sh

@@ -3,46 +3,179 @@
 
 //! GPU workload validation e2e tests.
 
+use std::fs;
+use std::path::{Path, PathBuf};
+
 use openshell_e2e::harness::output::strip_ansi;
 use openshell_e2e::harness::sandbox::SandboxGuard;
+use serde::Deserialize;
 
-const CUDA_WORKLOAD_IMAGE_ENV: &str = "OPENSHELL_E2E_GPU_CUDA_WORKLOAD_IMAGE";
+const WORKLOAD_MANIFEST_ENV: &str = "OPENSHELL_E2E_WORKLOAD_MANIFEST";
 const GPU_WORKLOAD_SUCCESS_MARKER: &str = "OPENSHELL_GPU_WORKLOAD_SUCCESS";
-const GPU_WORKLOAD_BINARY: &str = "/usr/local/bin/openshell-gpu-workload";
+const GPU_WORKLOAD_FAILURE_MARKER: &str = "OPENSHELL_GPU_WORKLOAD_FAILURE";
+
+#[derive(Debug, Deserialize)]
+struct WorkloadManifest {
+    workloads: Vec<WorkloadDefinition>,
+}
+
+#[derive(Clone, Debug, Deserialize)]
+struct WorkloadDefinition {
+    name: String,
+    image: String,
+    command: Vec<String>,
+    expect: WorkloadExpectation,
+    #[serde(default)]
+    requirements: WorkloadRequirements,
+}
 
-fn cuda_workload_image() -> Option<String> {
-    std::env::var(CUDA_WORKLOAD_IMAGE_ENV)
+#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq)]
+#[serde(rename_all = "lowercase")]
+enum WorkloadExpectation {
+    Pass,
+    Fail,
+}
+
+#[derive(Clone, Debug, Default, Deserialize)]
+struct WorkloadRequirements {
+    #[serde(default)]
+    gpu: bool,
+}
+
+fn default_workload_manifest_path() -> PathBuf {
+    Path::new(env!("CARGO_MANIFEST_DIR")).join("../gpu/images/.build/workloads.yaml")
+}
+
+fn workload_manifest_path() -> PathBuf {
+    std::env::var(WORKLOAD_MANIFEST_ENV)
         .ok()
         .map(|value| value.trim().to_string())
         .filter(|value| !value.is_empty())
+        .map(PathBuf::from)
+        .unwrap_or_else(default_workload_manifest_path)
 }
 
-#[tokio::test]
-async fn cuda_gpu_workload_validation_runs_explicit_workload_binary() {
-    let Some(image) = cuda_workload_image() else {
-        eprintln!("skipping CUDA GPU workload validation: {CUDA_WORKLOAD_IMAGE_ENV} is not set");
-        return;
+fn load_workload_manifest() -> Option<WorkloadManifest> {
+    let path = workload_manifest_path();
+    let explicit_override = std::env::var(WORKLOAD_MANIFEST_ENV)
+        .ok()
+        .map(|value| !value.trim().is_empty())
+        .unwrap_or(false);
+
+    let contents = match fs::read_to_string(&path) {
+        Ok(contents) => contents,
+        Err(err) if !explicit_override && err.kind() == std::io::ErrorKind::NotFound => {
+            eprintln!(
+                "skipping GPU workload validation: no workload manifest at {}. \
+                 Run `mise run e2e:workloads:build` to create the local manifest \
+                 or set {WORKLOAD_MANIFEST_ENV} to an external manifest.",
+                path.display()
+            );
+            return None;
+        }
+        Err(err) => panic!("failed to read workload manifest {}: {err}", path.display()),
     };
 
-    let mut guard = SandboxGuard::create(&[
-        "--gpu",
-        "--from",
-        image.as_str(),
-        "--",
-        GPU_WORKLOAD_BINARY,
-    ])
-    .await
-    .unwrap_or_else(|err| {
+    let manifest: WorkloadManifest = serde_yaml::from_str(&contents).unwrap_or_else(|err| {
         panic!(
-            "CUDA GPU workload sandbox create failed for image {image} with binary {GPU_WORKLOAD_BINARY}:\n{err}"
+            "failed to parse workload manifest {}: {err}",
+            path.display()
+        )
+    });
+    assert!(
+        !manifest.workloads.is_empty(),
+        "workload manifest {} contains no workloads",
+        path.display()
+    );
+    Some(manifest)
+}
+
+async fn assert_expected_pass(workload: &WorkloadDefinition) {
+    let mut args = vec![
+        "--gpu".to_string(),
+        "--from".to_string(),
+        workload.image.clone(),
+        "--".to_string(),
+    ];
+    args.extend(workload.command.clone());
+    let arg_refs = args.iter().map(String::as_str).collect::<Vec<_>>();
+
+    let mut guard = SandboxGuard::create(&arg_refs).await.unwrap_or_else(|err| {
+        panic!(
+            "GPU workload '{}' expected success but sandbox create failed:\n{err}",
+            workload.name
         )
     });
 
     let clean_output = strip_ansi(&guard.create_output);
     assert!(
         clean_output.contains(GPU_WORKLOAD_SUCCESS_MARKER),
-        "expected success marker {GPU_WORKLOAD_SUCCESS_MARKER} for image {image} in sandbox output:\n{clean_output}"
+        "expected success marker {GPU_WORKLOAD_SUCCESS_MARKER} for workload '{}' image {} in sandbox output:\n{clean_output}",
+        workload.name,
+        workload.image,
     );
 
     guard.cleanup().await;
 }
+
+async fn assert_expected_fail(workload: &WorkloadDefinition) {
+    let mut args = vec![
+        "--gpu".to_string(),
+        "--from".to_string(),
+        workload.image.clone(),
+        "--".to_string(),
+    ];
+    args.extend(workload.command.clone());
+    let arg_refs = args.iter().map(String::as_str).collect::<Vec<_>>();
+
+    match SandboxGuard::create(&arg_refs).await {
+        Ok(mut guard) => {
+            let clean_output = strip_ansi(&guard.create_output);
+            guard.cleanup().await;
+            panic!(
+                "GPU workload '{}' unexpectedly succeeded. Output:\n{clean_output}",
+                workload.name
+            );
+        }
+        Err(err) => {
+            let clean_output = strip_ansi(&err);
+            assert!(
+                clean_output.contains(GPU_WORKLOAD_FAILURE_MARKER),
+                "expected failure marker {GPU_WORKLOAD_FAILURE_MARKER} for workload '{}' image {} in failure output:\n{clean_output}",
+                workload.name,
+                workload.image,
+            );
+        }
+    }
+}
+
+#[tokio::test]
+async fn gpu_workload_manifest_runs_expected_workloads() {
+    let Some(manifest) = load_workload_manifest() else {
+        return;
+    };
+
+    let gpu_workloads = manifest
+        .workloads
+        .into_iter()
+        .filter(|workload| workload.requirements.gpu)
+        .collect::<Vec<_>>();
+
+    assert!(
+        !gpu_workloads.is_empty(),
+        "workload manifest contains no GPU-tagged workloads"
+    );
+
+    for workload in gpu_workloads {
+        assert!(
+            !workload.command.is_empty(),
+            "workload '{}' must declare a non-empty command",
+            workload.name
+        );
+
+        match workload.expect {
+            WorkloadExpectation::Pass => assert_expected_pass(&workload).await,
+            WorkloadExpectation::Fail => assert_expected_fail(&workload).await,
+        }
+    }
+}