fix

neuro-inc · Dec 23, 2024 · 33817a9 · 33817a9
1 parent 70a3d26
commit 33817a9
Showing 1 changed file with 13 additions and 9 deletions.
diff --git a/tests/integration/test_jobs_service.py b/tests/integration/test_jobs_service.py
@@ -249,9 +249,9 @@ async def test_save_no_tag(
         resources = Resources(
             memory=32 * 1024**2,
             cpu=0.1,
-            gpu=None,
             shm=False,
-            gpu_model=None,
+            nvidia_gpu=None,
+            nvidia_gpu_model=None,
             tpu_type=None,
             tpu_software_version=None,
         )
@@ -287,9 +287,9 @@ async def test_save_pending_job(
         resources = Resources(
             memory=2**60,
             cpu=0.1,
-            gpu=None,
             shm=False,
-            gpu_model=None,
+            nvidia_gpu=None,
+            nvidia_gpu_model=None,
             tpu_type=None,
             tpu_software_version=None,
         )
@@ -313,9 +313,9 @@ async def test_save_push_failure(
         resources = Resources(
             memory=32 * 1024**2,
             cpu=0.1,
-            gpu=None,
             shm=False,
-            gpu_model=None,
+            nvidia_gpu=None,
+            nvidia_gpu_model=None,
             tpu_type=None,
             tpu_software_version=None,
         )
@@ -354,9 +354,9 @@ async def test_save_commit_fails_with_exception(
         resources = Resources(
             memory=32 * 1024**2,
             cpu=0.1,
-            gpu=None,
             shm=False,
-            gpu_model=None,
+            nvidia_gpu=None,
+            nvidia_gpu_model=None,
             tpu_type=None,
             tpu_software_version=None,
         )
@@ -383,7 +383,11 @@ async def test_mark_logs_dropped(
         jobs_service: JobsService,
     ) -> None:
         resources = Resources(
-            memory=16 * 1024**2, cpu=0.1, gpu=None, shm=False, gpu_model=None
+            memory=16 * 1024**2,
+            cpu=0.1,
+            shm=False,
+            nvidia_gpu=None,
+            nvidia_gpu_model=None,
         )
         job = await job_factory(
             "alpine:latest",