cleanup

Signed-off-by: Jeev B <jeevb@users.noreply.github.com>
flyteorg · Sep 20, 2023 · e98517b · e98517b
1 parent ab1e0d6
commit e98517b
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 13 deletions.
diff --git a/plugins/flytekit-bigquery/tests/test_agent.py b/plugins/flytekit-bigquery/tests/test_agent.py
@@ -6,10 +6,12 @@
 
 import grpc
 from flyteidl.admin.agent_pb2 import SUCCEEDED
+from flyteidl.core import tasks_pb2 as _core_task
 from flytekitplugins.bigquery.agent import Metadata
 
 import flytekit.models.interface as interface_models
 from flytekit.extend.backend.base_agent import AgentRegistry
+from flytekit.extras.accelerators import NvidiaTeslaT4
 from flytekit.interfaces.cli_identifiers import Identifier
 from flytekit.models import literals, task, types
 from flytekit.models.core.identifier import ResourceType
@@ -59,7 +61,7 @@ def __init__(self):
         "This is deprecated!",
         True,
         "A",
-        None,
+        _core_task.ResourceMetadata(gpu_accelerator=NvidiaTeslaT4.to_flyte_idl()),
     )
     task_config = {
         "Location": "us-central1",

diff --git a/tests/flytekit/unit/extend/test_agent.py b/tests/flytekit/unit/extend/test_agent.py
@@ -20,6 +20,7 @@
     GetTaskResponse,
     Resource,
 )
+from flyteidl.core import tasks_pb2 as _core_task
 
 import flytekit.models.interface as interface_models
 from flytekit import PythonFunctionTask
@@ -31,6 +32,7 @@
     convert_to_flyte_state,
     is_terminal_state,
 )
+from flytekit.extras.accelerators import NvidiaTeslaT4
 from flytekit.models import literals, task, types
 from flytekit.models.core.identifier import Identifier, ResourceType
 from flytekit.models.literals import LiteralMap
@@ -78,7 +80,7 @@ def delete(self, context: grpc.ServicerContext, resource_meta: bytes) -> DeleteT
     "This is deprecated!",
     True,
     "A",
-    None,
+    _core_task.ResourceMetadata(gpu_accelerator=NvidiaTeslaT4.to_flyte_idl()),
 )
 
 int_type = types.LiteralType(types.SimpleType.INTEGER)

diff --git a/tests/flytekit/unit/models/core/test_workflow.py b/tests/flytekit/unit/models/core/test_workflow.py
@@ -2,7 +2,7 @@
 
 from flyteidl.core import tasks_pb2 as _core_task
 
-from flytekit.extras.accelerators import NvidiaTeslaA100
+from flytekit.extras.accelerators import NvidiaTeslaT4
 from flytekit.models import interface as _interface
 from flytekit.models import literals as _literals
 from flytekit.models import types as _types
@@ -304,11 +304,7 @@ def test_task_node_overrides():
             requests=[Resources.ResourceEntry(Resources.ResourceName.CPU, "1")],
             limits=[Resources.ResourceEntry(Resources.ResourceName.CPU, "2")],
         ),
-        _core_task.ResourceMetadata(
-            gpu_accelerator=NvidiaTeslaA100.with_partition_size(
-                NvidiaTeslaA100.partition_sizes.PARTITION_1G_5GB
-            ).to_flyte_idl()
-        ),
+        _core_task.ResourceMetadata(gpu_accelerator=NvidiaTeslaT4.to_flyte_idl()),
     )
     assert overrides.resources.requests == [Resources.ResourceEntry(Resources.ResourceName.CPU, "1")]
     assert overrides.resources.limits == [Resources.ResourceEntry(Resources.ResourceName.CPU, "2")]
@@ -325,11 +321,7 @@ def test_task_node_with_overrides():
                 requests=[Resources.ResourceEntry(Resources.ResourceName.CPU, "1")],
                 limits=[Resources.ResourceEntry(Resources.ResourceName.CPU, "2")],
             ),
-            _core_task.ResourceMetadata(
-                gpu_accelerator=NvidiaTeslaA100.with_partition_size(
-                    NvidiaTeslaA100.partition_sizes.PARTITION_1G_5GB
-                ).to_flyte_idl()
-            ),
+            _core_task.ResourceMetadata(gpu_accelerator=NvidiaTeslaT4.to_flyte_idl()),
         ),
     )