diff --git a/tests/aignostics/platform/e2e_test.py b/tests/aignostics/platform/e2e_test.py index e266684a6..35efdab32 100644 --- a/tests/aignostics/platform/e2e_test.py +++ b/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