Merge remote-tracking branch 'origin/master' into rc1.6.1

.github/workflows/test.yml

@@ -105,6 +105,7 @@ jobs:
     strategy:
       matrix:
         test:
+          - disk
           - unittest gen-rest-docs gen-cli-docs gen-readthedocs basic auth status batch anonymous competition unicode rest1 upload1 upload2 upload3 upload4 download
           - refs binary rm make worksheet_search worksheet_tags bundle_freeze_unfreeze worksheet_freeze_unfreeze detach perm search_time groups
           - worker_manager service
@@ -292,6 +293,7 @@ jobs:
     strategy:
       matrix:
         test:
+          - disk
           - basic status batch anonymous unicode rest1 upload1 download
           - refs binary rm make worksheet_search worksheet_tags bundle_freeze_unfreeze worksheet_freeze_unfreeze detach perm search_time groups
           - run
@@ -504,6 +506,7 @@ jobs:
     strategy:
       matrix:
         test:
+          - disk
           - unittest gen-rest-docs gen-cli-docs gen-readthedocs basic auth status batch anonymous competition unicode rest1 upload1 upload2 upload3 upload4 download
           - refs binary rm make worksheet_search worksheet_tags bundle_freeze_unfreeze worksheet_freeze_unfreeze detach perm search_time groups
           - worker_manager service

codalab/client/json_api_client.py

@@ -634,7 +634,13 @@ def fetch_contents_blob(self, target, range_=None, head=None, tail=None, truncat
 
     @wrap_exception('Unable to upload contents of bundle {1}')
     def upload_contents_blob(
-        self, bundle_id, fileobj=None, pass_self=False, params=None, progress_callback=None
+        self,
+        bundle_id,
+        fileobj=None,
+        pass_self=False,
+        bundle_uuid=None,
+        params=None,
+        progress_callback=None,
     ):
         """
         Uploads the contents of the given fileobj as the contents of specified
@@ -660,6 +666,7 @@ def upload_contents_blob(
                 query_params=params,
                 fileobj=fileobj,
                 pass_self=pass_self,
+                bundle_uuid=bundle_uuid,
                 progress_callback=progress_callback,
             )
 

codalab/lib/upload_manager.py

@@ -75,6 +75,7 @@ def write_fileobj(
         unpack_archive: bool,
         bundle_conn_str=None,
         index_conn_str=None,
+        bundle_uuid=None,
         progress_callback=None,
     ):
         """Writes fileobj indicated, unpacks if specified, and uploads it to the path at bundle_path.
@@ -195,6 +196,7 @@ def write_fileobj(
         unpack_archive: bool,
         bundle_conn_str=None,
         index_conn_str=None,
+        bundle_uuid=None,
         progress_callback=None,
     ):
         if unpack_archive:
@@ -227,6 +229,7 @@ def write_fileobj(
         unpack_archive: bool,
         bundle_conn_str=None,
         index_conn_str=None,
+        bundle_uuid=None,
         progress_callback=None,
     ):
         if unpack_archive:
@@ -256,7 +259,8 @@ def write_fileobj(
                     # Update disk and check if client has gone over disk usage.
                     if self._client and iteration % ITERATIONS_PER_DISK_CHECK == 0:
                         self._client.update(
-                            'user/increment_disk_used', {'disk_used_increment': len(to_send)}
+                            'user/increment_disk_used',
+                            {'disk_used_increment': len(to_send), 'bundle_uuid': bundle_uuid},
                         )
                         user_info = self._client.fetch('user')
                         if user_info['disk_used'] >= user_info['disk_quota']:
@@ -362,10 +366,12 @@ def has_contents(self, bundle):
         return os.path.exists(self._bundle_store.get_bundle_location(bundle.uuid))
 
     def cleanup_existing_contents(self, bundle):
-        self._bundle_store.cleanup(bundle.uuid, dry_run=False)
+        data_size = self._bundle_model.get_bundle_metadata(bundle.uuid, 'data_size')[bundle.uuid]
+        removed = self._bundle_store.cleanup(bundle.uuid, dry_run=False)
         bundle_update = {'data_hash': None, 'metadata': {'data_size': 0}}
         self._bundle_model.update_bundle(bundle, bundle_update)
-        self._bundle_model.update_user_disk_used(bundle.owner_id)
+        if removed:
+            self._bundle_model.increment_user_disk_used(bundle.owner_id, -data_size)
 
     def get_bundle_sas_token(self, path, **kwargs):
         """
@@ -484,6 +490,7 @@ def upload_to_bundle_store(
                         source_ext=source_ext,
                         should_unpack=unpack_before_upload,
                         json_api_client=self._client,
+                        bundle_uuid=bundle['id'],
                         progress_callback=progress.update,
                     )
                 self._client.update_bundle_state(bundle['id'], params={'success': True})
@@ -493,8 +500,6 @@ def upload_to_bundle_store(
                     params={'success': False, 'error_msg': f'Bypass server upload error. {err}',},
                 )
                 raise err
-            else:
-                self._client.update_bundle_state(bundle['id'], params={'success': True})
         else:
             # 5) Otherwise, upload the bundle directly through the server.
             progress = FileTransferProgress('Sent ', packed_source['filesize'], f=self.stderr)
@@ -512,6 +517,7 @@ def upload_to_bundle_store(
                     },
                     progress_callback=progress.update,
                     pass_self=True,
+                    bundle_uuid=bundle['id'],
                 )
 
     def upload_Azure_blob_storage(
@@ -524,6 +530,7 @@ def upload_Azure_blob_storage(
         source_ext,
         should_unpack,
         json_api_client,
+        bundle_uuid,
         progress_callback=None,
     ):
         """
@@ -552,6 +559,7 @@ def upload_Azure_blob_storage(
             should_unpack,
             bundle_conn_str,
             index_conn_str,
+            bundle_uuid,
             progress_callback,
         )
 
@@ -565,6 +573,7 @@ def upload_GCS_blob_storage(
         source_ext,
         should_unpack,
         json_api_client,
+        bundle_uuid,
         progress_callback=None,
     ):
         from codalab.lib import zip_util
@@ -582,6 +591,7 @@ def upload_GCS_blob_storage(
             fileobj=output_fileobj,
             query_params={},
             progress_callback=progress_callback,
+            bundle_uuid=bundle_uuid,
             json_api_client=json_api_client,
         )
         # upload the index file
@@ -602,5 +612,6 @@ def upload_GCS_blob_storage(
                 query_params={},
                 fileobj=open(tmp_index_file.name, "rb"),
                 progress_callback=None,
+                bundle_uuid=bundle_uuid,
                 json_api_client=self._client,
             )

codalab/model/bundle_model.py

@@ -1143,17 +1143,27 @@ def update_disk_metadata(self, bundle, bundle_location, enforce_disk_quota=False
         # TODO(Ashwin): make this non-fs specific
         data_hash = '0x%s' % (path_util.hash_directory(bundle_location, dirs_and_files))
         data_size = path_util.get_size(bundle_location, dirs_and_files)
+        try:
+            if 'data_size' in bundle.metadata.__dict__:
+                current_data_size = bundle.metadata.data_size
+            else:
+                current_data_size = int(
+                    self.get_bundle_metadata([bundle.uuid], 'data_size')[bundle.uuid]
+                )
+        except Exception:
+            current_data_size = 0
+        disk_increment = data_size - current_data_size
         if enforce_disk_quota:
             disk_left = self.get_user_disk_quota_left(bundle.owner_id)
-            if data_size > disk_left:
+            if disk_increment > disk_left:
                 raise UsageError(
                     "Can't save bundle, bundle size %s greater than user's disk quota left: %s"
                     % (data_size, disk_left)
                 )
 
         bundle_update = {'data_hash': data_hash, 'metadata': {'data_size': data_size}}
         self.update_bundle(bundle, bundle_update)
-        self.update_user_disk_used(bundle.owner_id)
+        self.increment_user_disk_used(bundle.owner_id, disk_increment)
 
     def bundle_checkin(self, bundle, worker_run, user_id, worker_id):
         """
@@ -2773,12 +2783,6 @@ def get_user_disk_quota_left(self, user_id, user_info=None):
             user_info = self.get_user_info(user_id)
         return user_info['disk_quota'] - user_info['disk_used']
 
-    def update_user_disk_used(self, user_id):
-        user_info = self.get_user_info(user_id)
-        # Compute from scratch for simplicity
-        user_info['disk_used'] = self._get_disk_used(user_id)
-        self.update_user_info(user_info)
-
     # ===========================================================================
     # OAuth-related methods follow!
     # ===========================================================================

codalab/rest/bundles.py

@@ -916,13 +916,15 @@ def _fetch_bundle_contents_blob(uuid, path=''):
     - `Content-Disposition: inline; filename=<bundle name or target filename>`
     - `Content-Type: <guess of mimetype based on file extension>`
     - `Content-Encoding: [gzip|identity]`
+    - `Access-Control-Allow-Origin: *`
     - `Target-Type: file`
     - `X-CodaLab-Target-Size: <size of the target>`
 
     HTTP Response headers (for directories):
     - `Content-Disposition: attachment; filename=<bundle or directory name>.tar.gz`
     - `Content-Type: application/gzip`
     - `Content-Encoding: identity`
+    - `Access-Control-Allow-Origin: *`
     - `Target-Type: directory`
     - `X-CodaLab-Target-Size: <size of the target>`
 
@@ -1038,6 +1040,7 @@ def _fetch_bundle_contents_blob(uuid, path=''):
         response.set_header('Content-Disposition', 'inline; filename="%s"' % filename)
     else:
         response.set_header('Content-Disposition', 'attachment; filename="%s"' % filename)
+    response.set_header('Access-Control-Allow-Origin', '*')
     response.set_header('Target-Type', target_info['type'])
     if target_info['type'] == 'file':
         size = target_info['size']

codalab/rest/users.py

@@ -260,12 +260,20 @@ def increment_user_disk_used():
     # TODO(agaut): Potentially convert the below to use a Schema (like those in schemas.py)
     # (Although, that does have downsides in this case.)
     disk_used_increment = request.json['data'][0]['attributes']['disk_used_increment']
+    bundle_uuid = request.json['data'][0]['attributes']['bundle_uuid']
 
     # only allow positive disk increments so that users can't abuse this endpoint.
     if disk_used_increment <= 0:
         abort(http.client.BAD_REQUEST, "Only positive disk increments are allowed.")
 
     local.model.increment_user_disk_used(request.user.user_id, disk_used_increment)
+    try:
+        data_size = int(local.model.get_bundle_metadata([bundle_uuid], 'data_size')[bundle_uuid])
+    except Exception:
+        data_size = 0
+    new_data_size = data_size + disk_used_increment
+    bundle = local.model.get_bundle(bundle_uuid)
+    local.model.update_bundle(bundle, {'metadata': {'data_size': new_data_size}})
     return (
         AuthenticatedUserSchema(many=True).dump([local.model.get_user(request.user.user_id)]).data
     )

codalab/worker/rest_client.py

@@ -114,7 +114,14 @@ def _make_request(
                     raise RestClientException('Invalid JSON: ' + response_data, False)
 
     def _upload_with_chunked_encoding(
-        self, method, url, query_params, fileobj, pass_self=False, progress_callback=None
+        self,
+        method,
+        url,
+        query_params,
+        fileobj,
+        pass_self=False,
+        bundle_uuid=None,
+        progress_callback=None,
     ):
         """
         Uploads the fileobj to url using method with query_params,
@@ -148,4 +155,5 @@ def _upload_with_chunked_encoding(
             url=url,
             progress_callback=progress_callback,
             json_api_client=json_api_client,
+            bundle_uuid=bundle_uuid,
         )

codalab/worker/upload_util.py

@@ -15,6 +15,7 @@ def upload_with_chunked_encoding(
     need_response=False,
     url="",
     progress_callback=None,
+    bundle_uuid=None,
     json_api_client=None,
 ):
     """
@@ -76,7 +77,8 @@ def upload_with_chunked_encoding(
             # Update disk and check if client has gone over disk usage.
             if json_api_client and iteration % ITERATIONS_PER_DISK_CHECK == 0:
                 json_api_client.update(
-                    'user/increment_disk_used', {'disk_used_increment': len(to_send)}
+                    'user/increment_disk_used',
+                    {'disk_used_increment': len(to_send), 'bundle_uuid': bundle_uuid},
                 )
                 user_info = json_api_client.fetch('user')
                 if user_info['disk_used'] >= user_info['disk_quota']:

docs/REST-API-Reference.md

@@ -678,13 +678,15 @@ HTTP Response headers (for single-file targets):
 - `Content-Disposition: inline; filename=<bundle name or target filename>`
 - `Content-Type: <guess of mimetype based on file extension>`
 - `Content-Encoding: [gzip|identity]`
+- `Access-Control-Allow-Origin: *`
 - `Target-Type: file`
 - `X-CodaLab-Target-Size: <size of the target>`
 
 HTTP Response headers (for directories):
 - `Content-Disposition: attachment; filename=<bundle or directory name>.tar.gz`
 - `Content-Type: application/gzip`
 - `Content-Encoding: identity`
+- `Access-Control-Allow-Origin: *`
 - `Target-Type: directory`
 - `X-CodaLab-Target-Size: <size of the target>`
 
@@ -729,13 +731,15 @@ HTTP Response headers (for single-file targets):
 - `Content-Disposition: inline; filename=<bundle name or target filename>`
 - `Content-Type: <guess of mimetype based on file extension>`
 - `Content-Encoding: [gzip|identity]`
+- `Access-Control-Allow-Origin: *`
 - `Target-Type: file`
 - `X-CodaLab-Target-Size: <size of the target>`
 
 HTTP Response headers (for directories):
 - `Content-Disposition: attachment; filename=<bundle or directory name>.tar.gz`
 - `Content-Type: application/gzip`
 - `Content-Encoding: identity`
+- `Access-Control-Allow-Origin: *`
 - `Target-Type: directory`
 - `X-CodaLab-Target-Size: <size of the target>`
 

docs/Worker-Managers.md

@@ -375,10 +375,9 @@ through the worker manager.
 
 To create a service account:
 
-1. Run `kubectl create -f service-account.yaml --namespace default`
+1. Run `cd codalab-worksheets && kubectl apply -f docs/gcp/service-account.yaml --namespace default`
 2. Then, run `kubectl get secrets --namespace default`
-3. Get the auth token by first finding the name of the first secret (in the form `codalab-token-<random string>`) and 
-   then use the name to get the token by running: `kubectl describe secret/codalab-token-<random string>`.
+3. Get the auth token by running: `kubectl describe secret/codalab-service-account-secret`.
 
 To get the cluster certificate:
 

docs/gcp/service-account.yaml

@@ -38,3 +38,13 @@ roleRef:
 subjects:
 - kind: ServiceAccount
   name: codalab
+
+---
+
+apiVersion: v1
+kind: Secret
+metadata:
+  name: codalab-service-account-secret
+  annotations:
+    kubernetes.io/service-account.name: codalab
+type: kubernetes.io/service-account-token