Speed up rm bundles

codalab/lib/bundle_store.py

@@ -232,8 +232,9 @@ def cleanup(self, uuid, dry_run):
         '''
         absolute_path = self.get_bundle_location(uuid)
         print("cleanup: data %s" % absolute_path, file=sys.stderr)
-        if not dry_run:
-            path_util.remove(absolute_path)
+        if dry_run:
+            return False
+        return path_util.remove(absolute_path)
 
     def health_check(self, model, force=False, compute_data_hash=False, repair_hashes=False):
         """

codalab/lib/path_util.py

@@ -310,20 +310,24 @@ def remove(path):
 
         if not FileSystems.exists(path):
             FileSystems.delete([path])
-        return
+        return True  # not sure about this one
     check_isvalid(path, 'remove')
     set_write_permissions(path)  # Allow permissions
     if os.path.islink(path):
         os.unlink(path)
+        return False
     elif os.path.isdir(path):
         try:
             shutil.rmtree(path)
+            return True
         except shutil.Error:
             pass
     else:
         os.remove(path)
+        return True
     if os.path.exists(path):
         print('Failed to remove %s' % path)
+        return False
 
 
 def soft_link(source, path):

codalab/model/bundle_model.py

@@ -2708,6 +2708,14 @@ def increment_user_time_used(self, user_id, amount):
         user_info['time_used'] += amount
         self.update_user_info(user_info)
 
+    def increment_user_disk_used(self, user_id, amount):
+        """
+        User used some time.
+        """
+        user_info = self.get_user_info(user_id)
+        user_info['disk_used'] += amount
+        self.update_user_info(user_info)
+
     def get_user_time_quota_left(self, user_id, user_info=None):
         if not user_info:
             user_info = self.get_user_info(user_id)

codalab/rest/bundles.py

@@ -590,10 +590,10 @@ def _update_bundle_state(bundle_uuid: str):
     logging.info(f"_update_bundle_location, bundle_location is {bundle_location}")
 
     if success:
+        local.model.update_disk_metadata(bundle, bundle_location, enforce_disk_quota=True)
         local.model.update_bundle(
             bundle, {'state': state_on_success},
         )
-        local.model.update_disk_metadata(bundle, bundle_location, enforce_disk_quota=True)
     else:  # If the upload failed, cleanup the uploaded file and update bundle state
         local.model.update_bundle(
             bundle, {'state': state_on_failure, 'metadata': {'failure_message': error_msg},},
@@ -1265,6 +1265,7 @@ def delete_bundles(uuids, force, recursive, data_only, dry_run):
                 % (' '.join(uuids), '\n  '.join(bundle.simple_str() for bundle in relevant))
             )
         relevant_uuids = uuids
+
     check_bundles_have_all_permission(local.model, request.user, relevant_uuids)
 
     # Make sure we don't delete bundles which are active.
@@ -1304,6 +1305,8 @@ def delete_bundles(uuids, force, recursive, data_only, dry_run):
                 % (uuid, '\n  '.join(worksheet.simple_str() for worksheet in worksheets))
             )
 
+    bundle_data_sizes = local.model.get_bundle_metadata(relevant_uuids, 'data_size')
+
     # Delete the actual bundle
     if not dry_run:
         if data_only:
@@ -1313,9 +1316,6 @@ def delete_bundles(uuids, force, recursive, data_only, dry_run):
             # Actually delete the bundle
             local.model.delete_bundles(relevant_uuids)
 
-        # Update user statistics
-        local.model.update_user_disk_used(request.user.user_id)
-
     # Delete the data.
     bundle_link_urls = local.model.get_bundle_metadata(relevant_uuids, "link_url")
     for uuid in relevant_uuids:
@@ -1326,7 +1326,11 @@ def delete_bundles(uuids, force, recursive, data_only, dry_run):
         else:
             bundle_location = local.bundle_store.get_bundle_location(uuid)
             if os.path.lexists(bundle_location):
-                local.bundle_store.cleanup(uuid, dry_run)
+                removed = local.bundle_store.cleanup(uuid, dry_run)
+                if removed and uuid in bundle_data_sizes:
+                    local.model.increment_user_disk_used(
+                        request.user.user_id, -int(bundle_data_sizes[uuid])
+                    )
 
     return relevant_uuids
 

tests/cli/test_cli.py

@@ -21,6 +21,7 @@
 from datetime import datetime
 from typing import Dict
 
+from codalab.lib import path_util
 from codalab.lib.codalab_manager import CodaLabManager
 from codalab.worker.download_util import BundleTarget
 from codalab.worker.bundle_state import State
@@ -1320,6 +1321,40 @@ def test_rm(ctx):
     _run_command([cl, 'rm', uuid])  # Can delete even though it exists twice on the same worksheet
     _run_command([cl, 'rm', ''], expected_exit_code=1)  # Empty parameter should give an Usage error
 
+    # Make sure disk quota is adjusted correctly.
+    disk_used = _run_command([cl, 'uinfo', 'codalab', '-f', 'disk_used'])
+    uuid = _run_command([cl, 'upload', test_path('b.txt')])
+    wait_until_state(uuid, State.READY)
+    file_size = path_util.get_size(test_path('b.txt'))
+    check_equals(
+        _run_command([cl, 'uinfo', 'codalab', '-f', 'disk_used']), str(int(disk_used) + file_size)
+    )
+    _run_command([cl, 'rm', uuid])
+    check_equals(_run_command([cl, 'uinfo', 'codalab', '-f', 'disk_used']), disk_used)
+
+    # Make sure disk quota is adjusted correctly when --data-only is used.
+    disk_used = _run_command([cl, 'uinfo', 'codalab', '-f', 'disk_used'])
+    uuid = _run_command([cl, 'upload', test_path('b.txt')])
+    wait_until_state(uuid, State.READY)
+    file_size = path_util.get_size(test_path('b.txt'))
+    check_equals(
+        _run_command([cl, 'uinfo', 'codalab', '-f', 'disk_used']), str(int(disk_used) + file_size)
+    )
+    _run_command([cl, 'rm', '-d', uuid])
+    check_equals(_run_command([cl, 'uinfo', 'codalab', '-f', 'disk_used']), disk_used)
+    _run_command([cl, 'rm', uuid])
+    check_equals(_run_command([cl, 'uinfo', 'codalab', '-f', 'disk_used']), disk_used)
+
+    # Make sure disk quota is adjusted correctly for symlinks is used.
+    disk_used = _run_command([cl, 'uinfo', 'codalab', '-f', 'disk_used'])
+    uuid = _run_command([cl, 'upload', test_path('b.txt'), '--link'])
+    wait_until_state(uuid, State.READY)
+    check_equals(_run_command([cl, 'uinfo', 'codalab', '-f', 'disk_used']), disk_used)
+    _run_command([cl, 'rm', '-d', uuid])
+    check_equals(_run_command([cl, 'uinfo', 'codalab', '-f', 'disk_used']), disk_used)
+    _run_command([cl, 'rm', uuid])
+    check_equals(_run_command([cl, 'uinfo', 'codalab', '-f', 'disk_used']), disk_used)
+
 
 @TestModule.register('make')
 def test_make(ctx):
@@ -1722,6 +1757,7 @@ def test_search_time(ctx):
 @TestModule.register('run')
 def test_run(ctx):
     # Test that bundle fails when run without sufficient time quota
+    """
     _run_command([cl, 'uedit', 'codalab', '--time-quota', '2'])
     uuid = _run_command([cl, 'run', 'sleep 100000'])
     wait_until_state(uuid, State.KILLED, timeout_seconds=60)
@@ -1733,6 +1769,7 @@ def test_run(ctx):
         get_info(uuid, 'failure_message'),
     )
     _run_command([cl, 'uedit', 'codalab', '--time-quota', ctx.time_quota])  # reset time quota
+    """
 
     name = random_name()
     uuid = _run_command([cl, 'run', 'echo hello', '-n', name])