[PT/XLA] Use on-demand capacity for v4 tests (#170)

Change-Id: I976a2b1c09d577392fb65f940ae03848dfca06a7

Author: will-cromar

dags/pytorch_xla/nightly.py

@@ -55,13 +55,14 @@ def torchvision():
   ).run()
   resnet_v2_8 = task.TpuQueuedResourceTask(
       test_config.JSonnetTpuVmTest.from_pytorch(
-          "pt-nightly-resnet50-pjrt-fake-v2-8-1vm"
+          "pt-nightly-resnet50-pjrt-fake-v2-8-1vm",
+          reserved=True,
       ),
       US_CENTRAL1_C,
   ).run()
   resnet_v3_8_tests = [
       task.TpuQueuedResourceTask(
-          test_config.JSonnetTpuVmTest.from_pytorch(test),
+          test_config.JSonnetTpuVmTest.from_pytorch(test, reserved=True),
           US_CENTRAL1_C,
       ).run()
       for test in (
@@ -92,6 +93,7 @@ def torchvision():
           "pt-nightly-resnet50-pjrt-fake-v5litepod-4-1vm",
           network=V5_NETWORKS,
           subnetwork=V5E_SUBNETWORKS,
+          reserved=True,
       ),
       US_EAST1_C,
   ).run()
@@ -137,7 +139,9 @@ def torchvision():
 @task_group(prefix_group_id=False)
 def huggingface():
   accelerate_v2_8 = task.TpuQueuedResourceTask(
-      test_config.JSonnetTpuVmTest.from_pytorch("pt-nightly-accelerate-smoke-v2-8-1vm"),
+      test_config.JSonnetTpuVmTest.from_pytorch(
+          "pt-nightly-accelerate-smoke-v2-8-1vm", reserved=True
+      ),
       US_CENTRAL1_C,
   ).run()
   accelerate_v4_8 = task.TpuQueuedResourceTask(

xlml/apis/test_config.py

@@ -313,7 +313,7 @@ def _from_json_helper(
 
   @staticmethod
   def from_jax(
-      test_name: str, reserved_tpu: bool = True, network='default', subnetwork='default'
+      test_name: str, reserved: bool = False, network='default', subnetwork='default'
   ):
     """Parses a compiled legacy JSonnet config test from `tests/jax`."""
     test = _load_compiled_jsonnet(test_name)
@@ -323,14 +323,14 @@ def from_jax(
         setup=test['setup'],
         exports='',
         test_command=['bash', '-c', test['runTest']],
-        reserved=reserved_tpu,
+        reserved=reserved,
         network=network,
         subnetwork=subnetwork,
     )
 
   @staticmethod
   def from_pytorch(
-      test_name: str, reserved_tpu: bool = True, network='default', subnetwork='default'
+      test_name: str, reserved: bool = False, network='default', subnetwork='default'
   ):
     """Parses a compiled legacy JSonnet test config from `tests/pytorch`."""
     test = _load_compiled_jsonnet(test_name)
@@ -341,7 +341,7 @@ def from_pytorch(
         + '\ncd ~\n' + test['tpuSettings']['tpuVmExtraSetup'],
         exports=test['tpuSettings']['tpuVmExports'],
         test_command=test['command'],
-        reserved=reserved_tpu,
+        reserved=reserved,
         network=network,
         subnetwork=subnetwork,
     )