test: use GPUConfig in E2E test (#393)

Author: olivermeyer

tests/aignostics/platform/e2e_test.py

@@ -25,6 +25,7 @@
 
 from aignostics import platform
 from aignostics.platform import Run, RunSdkMetadata
+from aignostics.platform._sdk_metadata import GPUConfig
 from tests.constants_test import (
     HETA_APPLICATION_ID,
     HETA_APPLICATION_VERSION,
@@ -280,31 +281,33 @@ def _submit_and_validate(  # noqa: PLR0913, PLR0917
 
     logger.trace(f"Submitting application run for {application_id} version {application_version}")
     client = platform.Client()
+    gpu_config = GPUConfig(
+        gpu_type=PIPELINE_GPU_TYPE,
+        provisioning_mode=PIPELINE_GPU_PROVISIONING_MODE,
+        max_gpus_per_slide=PIPELINE_MAX_GPUS_PER_SLIDE,
+        flex_start_max_run_duration_minutes=PIPELINE_GPU_FLEX_START_MAX_RUN_DURATION_MINUTES,
+    )
+    custom_metadata = {
+        "sdk": {
+            "tags": tags or set(),
+            "scheduling": {
+                "due_date": (datetime.now(tz=UTC) + timedelta(seconds=due_date_seconds)).isoformat(),
+                "deadline": deadline.isoformat(),
+            },
+            "pipeline": {
+                "gpu": gpu_config.model_dump(),
+                "cpu": {
+                    "provisioning_mode": PIPELINE_CPU_PROVISIONING_MODE,
+                },
+                "node_acquisition_timeout_minutes": PIPELINE_NODE_ACQUISITION_TIMEOUT_MINUTES,
+            },
+        }
+    }
     run = client.runs.submit(
         application_id=application_id,
         application_version=application_version,
         items=payload,
-        custom_metadata={
-            "sdk": {
-                "tags": tags or set(),
-                "scheduling": {
-                    "due_date": (datetime.now(tz=UTC) + timedelta(seconds=due_date_seconds)).isoformat(),
-                    "deadline": deadline.isoformat(),
-                },
-                "pipeline": {
-                    "gpu": {
-                        "gpu_type": PIPELINE_GPU_TYPE,
-                        "provisioning_mode": PIPELINE_GPU_PROVISIONING_MODE,
-                        "flex_start_max_run_duration_minutes": PIPELINE_GPU_FLEX_START_MAX_RUN_DURATION_MINUTES,
-                        "max_gpus_per_slide": PIPELINE_MAX_GPUS_PER_SLIDE,
-                    },
-                    "cpu": {
-                        "provisioning_mode": PIPELINE_CPU_PROVISIONING_MODE,
-                    },
-                    "node_acquisition_timeout_minutes": PIPELINE_NODE_ACQUISITION_TIMEOUT_MINUTES,
-                },
-            }
-        },
+        custom_metadata=custom_metadata,
     )
 
     # Let's validate we can fiond the run by id