From 0bd9c05c356af899ca8646a8f1f2b007fccb214f Mon Sep 17 00:00:00 2001 From: Yifan Mai Date: Tue, 7 Feb 2023 09:40:28 -0800 Subject: [PATCH 1/3] Add Access-Control-Allow-Origin response header (#4367) --- codalab/rest/bundles.py | 3 +++ docs/REST-API-Reference.md | 4 ++++ 2 files changed, 7 insertions(+) diff --git a/codalab/rest/bundles.py b/codalab/rest/bundles.py index 841e568b6..9ef10ce4d 100644 --- a/codalab/rest/bundles.py +++ b/codalab/rest/bundles.py @@ -916,6 +916,7 @@ def _fetch_bundle_contents_blob(uuid, path=''): - `Content-Disposition: inline; filename=` - `Content-Type: ` - `Content-Encoding: [gzip|identity]` + - `Access-Control-Allow-Origin: *` - `Target-Type: file` - `X-CodaLab-Target-Size: ` @@ -923,6 +924,7 @@ def _fetch_bundle_contents_blob(uuid, path=''): - `Content-Disposition: attachment; filename=.tar.gz` - `Content-Type: application/gzip` - `Content-Encoding: identity` + - `Access-Control-Allow-Origin: *` - `Target-Type: directory` - `X-CodaLab-Target-Size: ` @@ -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'] diff --git a/docs/REST-API-Reference.md b/docs/REST-API-Reference.md index 2a8da1d9e..f58e21315 100644 --- a/docs/REST-API-Reference.md +++ b/docs/REST-API-Reference.md @@ -678,6 +678,7 @@ HTTP Response headers (for single-file targets): - `Content-Disposition: inline; filename=` - `Content-Type: ` - `Content-Encoding: [gzip|identity]` +- `Access-Control-Allow-Origin: *` - `Target-Type: file` - `X-CodaLab-Target-Size: ` @@ -685,6 +686,7 @@ HTTP Response headers (for directories): - `Content-Disposition: attachment; filename=.tar.gz` - `Content-Type: application/gzip` - `Content-Encoding: identity` +- `Access-Control-Allow-Origin: *` - `Target-Type: directory` - `X-CodaLab-Target-Size: ` @@ -729,6 +731,7 @@ HTTP Response headers (for single-file targets): - `Content-Disposition: inline; filename=` - `Content-Type: ` - `Content-Encoding: [gzip|identity]` +- `Access-Control-Allow-Origin: *` - `Target-Type: file` - `X-CodaLab-Target-Size: ` @@ -736,6 +739,7 @@ HTTP Response headers (for directories): - `Content-Disposition: attachment; filename=.tar.gz` - `Content-Type: application/gzip` - `Content-Encoding: identity` +- `Access-Control-Allow-Origin: *` - `Target-Type: directory` - `X-CodaLab-Target-Size: ` From 7fc08ccda2d0c6a0208be8c327efbedf3ae7a53d Mon Sep 17 00:00:00 2001 From: Ashwin Ramaswami Date: Tue, 7 Feb 2023 14:05:15 -0500 Subject: [PATCH 2/3] Update Worker-Managers.md (#4371) * Update Worker-Managers.md Fix path to file * Update Worker-Managers.md * Update service-account.yaml --------- Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- docs/Worker-Managers.md | 5 ++--- docs/gcp/service-account.yaml | 10 ++++++++++ 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/docs/Worker-Managers.md b/docs/Worker-Managers.md index f6a3294c4..8772bfe51 100644 --- a/docs/Worker-Managers.md +++ b/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-`) and - then use the name to get the token by running: `kubectl describe secret/codalab-token-`. +3. Get the auth token by running: `kubectl describe secret/codalab-service-account-secret`. To get the cluster certificate: diff --git a/docs/gcp/service-account.yaml b/docs/gcp/service-account.yaml index d360fecfa..f98b9deeb 100644 --- a/docs/gcp/service-account.yaml +++ b/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 From e40009b72973667d3c01051c03e5be44618b7ef0 Mon Sep 17 00:00:00 2001 From: AndrewJGaut <35617203+AndrewJGaut@users.noreply.github.com> Date: Wed, 8 Feb 2023 00:37:36 -0800 Subject: [PATCH 3/3] Speed up disk quota computation (#4354) * Created a bunch more disk tests and got them to pass * update the updaet_disk_metadata function to update disk incrementally, which is mcuh faster * Removed all calls of update_user_disk_used * Remove update_user_disk_used function because it was slow * add in new module to Github actions workflow * minor change * Modify tests. They now pass. However, something weird is going on with the data_size stuff. I'm not surprised that it's hard to predict for a running bundle, but I am surprised that something weird happens when uploading two things at once. The bundle size isn't the sum of the two archives. Very weird * add in cl search to try and find the issue * add in cl search to try and find the issue * put disk in its own test just to see * Attempt to fix the Azure issue with the disk storage stuff * Clean up repo * Add in a new test case to verify that cleanup_existing_contents works as intended * formatting changes * Added changes to make the parallel disk updates work with the disk incremental updates * minor formatting changes --- .github/workflows/test.yml | 3 + codalab/client/json_api_client.py | 9 ++- codalab/lib/upload_manager.py | 21 +++++-- codalab/model/bundle_model.py | 20 ++++--- codalab/rest/users.py | 8 +++ codalab/worker/rest_client.py | 10 +++- codalab/worker/upload_util.py | 4 +- tests/cli/test_cli.py | 98 ++++++++++++++++++++++++++++++- 8 files changed, 155 insertions(+), 18 deletions(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 00cc9ec5d..83320d137 100644 --- a/.github/workflows/test.yml +++ b/.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 diff --git a/codalab/client/json_api_client.py b/codalab/client/json_api_client.py index 44c1edf33..e58e11e2a 100644 --- a/codalab/client/json_api_client.py +++ b/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, ) diff --git a/codalab/lib/upload_manager.py b/codalab/lib/upload_manager.py index d0d6edb80..28d9490d9 100644 --- a/codalab/lib/upload_manager.py +++ b/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, ) diff --git a/codalab/model/bundle_model.py b/codalab/model/bundle_model.py index a22d743d1..0444e96a2 100644 --- a/codalab/model/bundle_model.py +++ b/codalab/model/bundle_model.py @@ -1143,9 +1143,19 @@ 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) @@ -1153,7 +1163,7 @@ def update_disk_metadata(self, bundle, bundle_location, enforce_disk_quota=False 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! # =========================================================================== diff --git a/codalab/rest/users.py b/codalab/rest/users.py index 15a040b89..05337260e 100644 --- a/codalab/rest/users.py +++ b/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 ) diff --git a/codalab/worker/rest_client.py b/codalab/worker/rest_client.py index 60a1d65f9..4dd65d8f7 100644 --- a/codalab/worker/rest_client.py +++ b/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, ) diff --git a/codalab/worker/upload_util.py b/codalab/worker/upload_util.py index 83e168f8d..aa9548753 100644 --- a/codalab/worker/upload_util.py +++ b/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']: diff --git a/tests/cli/test_cli.py b/tests/cli/test_cli.py index 5a7ae9324..26c8b4b0f 100644 --- a/tests/cli/test_cli.py +++ b/tests/cli/test_cli.py @@ -22,6 +22,7 @@ from typing import Dict from codalab.lib import path_util +from codalab.lib.zip_util import pack_files_for_upload from codalab.lib.codalab_manager import CodaLabManager from codalab.worker.download_util import BundleTarget from codalab.worker.bundle_state import State @@ -1354,7 +1355,40 @@ def test_rm(ctx): _run_command([cl, 'rm', uuid]) check_equals(disk_used, _run_command([cl, 'uinfo', 'codalab', '-f', 'disk_used'])) - # Make sure disk quota is adjusted correctly when --data-only is used. + +@TestModule.register('disk') +def test_disk(ctx): + # Basic test. + 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( + str(int(disk_used) + file_size), _run_command([cl, 'uinfo', 'codalab', '-f', 'disk_used']) + ) + _run_command([cl, 'rm', uuid]) + check_equals(disk_used, _run_command([cl, 'uinfo', 'codalab', '-f', 'disk_used'])) + + # Test with directory upload + disk_used = _run_command([cl, 'uinfo', 'codalab', '-f', 'disk_used']) + uuid = _run_command([cl, 'upload', test_path('dir1')]) + wait_until_state(uuid, State.READY) + # Directories are stored in their gzipped format when uploaded to Azure/GCS + # but are stored in their full file size format on disk. + if os.environ.get("CODALAB_ALWAYS_USE_AZURE_BLOB_BETA") == '1': + tarred_dir = pack_files_for_upload([test_path('dir1')], True, False)['fileobj'] + file_size = len(tarred_dir.read()) + else: + file_size = path_util.get_size(test_path('dir1')) + data_size = _run_command([cl, 'info', uuid, '-f', 'data_size']) + check_equals( + str(int(disk_used) + int(data_size)), + _run_command([cl, 'uinfo', 'codalab', '-f', 'disk_used']), + ) + _run_command([cl, 'rm', uuid]) + check_equals(disk_used, _run_command([cl, 'uinfo', 'codalab', '-f', 'disk_used'])) + + # Test rm --data-only. 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) @@ -1367,7 +1401,7 @@ def test_rm(ctx): _run_command([cl, 'rm', uuid]) check_equals(disk_used, _run_command([cl, 'uinfo', 'codalab', '-f', 'disk_used'])) - # Make sure disk quota is adjusted correctly for symlinks is used. + # Test with symlinks 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) @@ -1377,6 +1411,66 @@ def test_rm(ctx): _run_command([cl, 'rm', uuid]) check_equals(disk_used, _run_command([cl, 'uinfo', 'codalab', '-f', 'disk_used'])) + # Test with running bundle + disk_used = _run_command([cl, 'uinfo', 'codalab', '-f', 'disk_used']) + uuid = _run_command([cl, 'run', 'head -c 50 /dev/zero > test.txt']) + wait_until_state(uuid, State.READY) + data_size = _run_command([cl, 'info', uuid, '-f', 'data_size']) + check_equals( + str(int(data_size) + int(disk_used)), + _run_command([cl, 'uinfo', 'codalab', '-f', 'disk_used']), + ) + _run_command([cl, 'rm', uuid]) + check_equals(disk_used, _run_command([cl, 'uinfo', 'codalab', '-f', 'disk_used'])) + + # Test with archive upload + disk_used = _run_command([cl, 'uinfo', 'codalab', '-f', 'disk_used']) + archive_paths = [temp_path(''), temp_path('')] + archive_exts = [p + '.tar.gz' for p in archive_paths] + contents_paths = [test_path('dir1'), test_path('a.txt')] + for (archive, content) in zip(archive_exts, contents_paths): + _run_command( + ['tar', 'cfz', archive, '-C', os.path.dirname(content), os.path.basename(content)] + ) + uuid = _run_command([cl, 'upload'] + archive_exts) + wait_until_state(uuid, State.READY) + data_size = _run_command([cl, 'info', uuid, '-f', 'data_size']) + check_equals( + str(int(disk_used) + int(data_size)), + _run_command([cl, 'uinfo', 'codalab', '-f', 'disk_used']), + ) + _run_command([cl, 'rm', uuid]) + check_equals(disk_used, _run_command([cl, 'uinfo', 'codalab', '-f', 'disk_used'])) + + # Test with make. + disk_used = _run_command([cl, 'uinfo', 'codalab', '-f', 'disk_used']) + uuid1 = _run_command([cl, 'upload', test_path('a.txt')]) + uuid2 = _run_command([cl, 'upload', test_path('b.txt')]) + uploaded_files_size = path_util.get_size(test_path('a.txt')) + path_util.get_size( + test_path('b.txt') + ) + uuid3 = _run_command([cl, 'make', 'dep1:' + uuid1, 'dep2:' + uuid2]) + wait_until_state(uuid3, State.READY) + data_size = _run_command([cl, 'info', uuid3, '-f', 'data_size']) + check_equals( + str(int(disk_used) + int(data_size) + uploaded_files_size), + _run_command([cl, 'uinfo', 'codalab', '-f', 'disk_used']), + ) + _run_command([cl, 'rm', uuid3]) + _run_command([cl, 'rm', uuid2]) + _run_command([cl, 'rm', uuid1]) + check_equals(disk_used, _run_command([cl, 'uinfo', 'codalab', '-f', 'disk_used'])) + + # Test cleanup_existing_contents. + disk_used = _run_command([cl, 'uinfo', 'codalab', '-f', 'disk_used']) + _run_command([cl, 'uedit', 'codalab', '--disk-quota', f'{int(disk_used) + 10}']) + uuid = _run_command( + [cl, 'run', 'head -c 1000 /dev/zero > test.txt',], request_disk=None, request_memory=None, + ) + wait_until_state(uuid, State.FAILED) + _run_command([cl, 'uedit', 'codalab', '--disk-quota', ctx.disk_quota]) # reset disk quota + check_equals(disk_used, _run_command([cl, 'uinfo', 'codalab', '-f', 'disk_used'])) + @TestModule.register('make') def test_make(ctx):