[batch] Dont box data_disk_space_remaining (#13968)

Very small change, something I noticed while working on something else entirely. Given how this is currently used I don't think it needs to be boxed anymore.
hail-is · Nov 2, 2023 · 2c1188c · 2c1188c
1 parent b4a3d57
commit 2c1188c
Showing 1 changed file with 5 additions and 6 deletions.
diff --git a/batch/batch/worker/worker.py b/batch/batch/worker/worker.py
@@ -85,7 +85,6 @@
 from ..publicly_available_images import publicly_available_images
 from ..resource_usage import ResourceUsageMonitor
 from ..semaphore import FIFOWeightedSemaphore
-from ..utils import Box
 from ..worker.worker_api import CloudDisk, CloudWorkerAPI, ContainerRegistryCredentials
 from .credentials import CloudUserCredentials
 from .jvm_entryway_protocol import EndOfStream, read_bool, read_int, read_str, write_int, write_str
@@ -1828,7 +1827,7 @@ def container_name(self, task_name: str):
     async def setup_io(self):
         assert instance_config
         if not instance_config.job_private:
-            if self.worker.data_disk_space_remaining.value < self.external_storage_in_gib:
+            if self.worker.data_disk_space_remaining < self.external_storage_in_gib:
                 log.info(
                     f'worker data disk storage is full: {self.external_storage_in_gib}Gi requested and {self.worker.data_disk_space_remaining}Gi remaining'
                 )
@@ -1849,7 +1848,7 @@ async def setup_io(self):
                 log.info(f'created disk {self.disk.name} for job {self.id}')
                 return
 
-            self.worker.data_disk_space_remaining.value -= self.external_storage_in_gib
+            self.worker.data_disk_space_remaining -= self.external_storage_in_gib
             log.info(
                 f'acquired {self.external_storage_in_gib}Gi from worker data disk storage with {self.worker.data_disk_space_remaining}Gi remaining'
             )
@@ -2015,7 +2014,7 @@ async def cleanup(self):
             except Exception:
                 log.exception(f'while detaching and deleting disk {self.disk.name} for job {self.id}')
         else:
-            self.worker.data_disk_space_remaining.value += self.external_storage_in_gib
+            self.worker.data_disk_space_remaining += self.external_storage_in_gib
 
         if self.cloudfuse:
             for config in self.cloudfuse:
@@ -2959,7 +2958,7 @@ def __init__(self, client_session: httpx.ClientSession):
         self.cores_mcpu = CORES * 1000
         self.last_updated = time_msecs()
         self.cpu_sem = FIFOWeightedSemaphore(self.cores_mcpu)
-        self.data_disk_space_remaining = Box(UNRESERVED_WORKER_DATA_DISK_SIZE_GB)
+        self.data_disk_space_remaining = UNRESERVED_WORKER_DATA_DISK_SIZE_GB
         self.pool = concurrent.futures.ThreadPoolExecutor()
         self.jobs: Dict[Tuple[int, int], Job] = {}
         self.stop_event = asyncio.Event()
@@ -3247,7 +3246,7 @@ async def run(self):
                         break
                     log.info(
                         f'n_jobs {len(self.jobs)} free_cores {self.cpu_sem.value / 1000} idle {idle_duration} '
-                        f'free worker data disk storage {self.data_disk_space_remaining.value}Gi'
+                        f'free worker data disk storage {self.data_disk_space_remaining}Gi'
                     )
         finally:
             self.active = False