Repository: codalab/codalab-worksheets

Issue:
Support GCP for bundle storage

Discussion:
epicfaace: @percyliang FYI this PR is about the unit of work required to support a new storage type that Apache Beam supports. Relatively straightforward and minimal -- a similar PR could add support for S3.

Patch:
diff --git a/.gitignore b/.gitignore
index e5eb72a8..9ed5eb27 100644
--- a/.gitignore
+++ b/.gitignore
@@ -45,4 +45,5 @@ codalab-worker-scratch
 
 .coverage
 
-sent.jsonl
\ No newline at end of file
+sent.jsonl
+google-credentials.json
\ No newline at end of file
diff --git a/alembic/versions/2021121222_rm_storage_type_enum_a325a0756797.py b/alembic/versions/2021121222_rm_storage_type_enum_a325a0756797.py
new file mode 100644
index 00000000..b8a056e2
--- /dev/null
+++ b/alembic/versions/2021121222_rm_storage_type_enum_a325a0756797.py
@@ -0,0 +1,28 @@
+"""rm storage type enum
+
+Revision ID: a325a0756797
+Revises: 685cb4bc0b79
+Create Date: 2021-12-12 22:04:38.060506
+
+"""
+
+# revision identifiers, used by Alembic.
+revision = 'a325a0756797'
+down_revision = '685cb4bc0b79'
+
+from alembic import op
+import sqlalchemy as sa
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.alter_column('bundle_store', 'storage_type', type_=sa.String(255), existing_type=sa.Enum("disk", "azure_blob"), nullable=True)
+    op.alter_column('bundle_store', 'storage_format', type_=sa.String(255), existing_type=sa.Enum("uncompressed", "compressed_v1"), nullable=True)
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.alter_column('bundle_store', 'storage_type', type_=sa.Enum("disk", "azure_blob"), existing_type=sa.String(255), nullable=True)
+    op.alter_column('bundle_store', 'storage_format', type_=sa.Enum("uncompressed", "compressed_v1"), existing_type=sa.String(255), nullable=True)
+    # ### end Alembic commands ###
diff --git a/codalab/common.py b/codalab/common.py
index f3d28681..c8b63311 100644
--- a/codalab/common.py
+++ b/codalab/common.py
@@ -172,6 +172,7 @@ class StorageType(Enum):
 
     DISK_STORAGE = "disk"
     AZURE_BLOB_STORAGE = "azure_blob"
+    GCS_STORAGE = "gcs"
 
 
 class StorageURLScheme(Enum):
@@ -182,6 +183,7 @@ class StorageURLScheme(Enum):
 
     DISK_STORAGE = ""
     AZURE_BLOB_STORAGE = "azfs://"
+    GCS_STORAGE = "gs://"
 
 
 class StorageFormat(Enum):
@@ -236,7 +238,7 @@ class LinkedBundlePath:
         """Generates a SAS URL that can be used to read the given blob for one hour."""
         if self.storage_type != StorageType.AZURE_BLOB_STORAGE.value:
             raise ValueError(
-                f"SAS URLs can only be retrieved for bundles on Blob Storage. Storage type is: {self.storage_type}."
+                f"SAS URLs can only be retrieved for bundles on Azure Blob Storage. Storage type is: {self.storage_type}."
             )
         blob_name = path.replace(
             f"azfs://{AZURE_BLOB_ACCOUNT_NAME}/{AZURE_BLOB_CONTAINER_NAME}/", ""
@@ -266,12 +268,23 @@ def parse_linked_bundle_url(url):
 
         Returns a LinkedBundlePath instance to encode this information.
     """
-    if url.startswith(StorageURLScheme.AZURE_BLOB_STORAGE.value):
+    if url.startswith(StorageURLScheme.AZURE_BLOB_STORAGE.value) or url.startswith(
+        StorageURLScheme.GCS_STORAGE.value
+    ):
         uses_beam = True
-        storage_type = StorageType.AZURE_BLOB_STORAGE.value
-        url = url[len(StorageURLScheme.AZURE_BLOB_STORAGE.value) :]
-        storage_account, container, bundle_uuid, contents_file, *remainder = url.split("/", 4)
-        bundle_path = f"{StorageURLScheme.AZURE_BLOB_STORAGE.value}{storage_account}/{container}/{bundle_uuid}/{contents_file}"
+        if url.startswith(StorageURLScheme.AZURE_BLOB_STORAGE.value):
+            storage_type = StorageType.AZURE_BLOB_STORAGE.value
+            url = url[len(StorageURLScheme.AZURE_BLOB_STORAGE.value) :]
+            storage_account, container, bundle_uuid, contents_file, *remainder = url.split("/", 4)
+            bundle_path = f"{StorageURLScheme.AZURE_BLOB_STORAGE.value}{storage_account}/{container}/{bundle_uuid}/{contents_file}"
+        if url.startswith(StorageURLScheme.GCS_STORAGE.value):
+            storage_type = StorageType.GCS_STORAGE.value
+            url = url[len(StorageURLScheme.GCS_STORAGE.value) :]
+            bucket_name, bundle_uuid, contents_file, *remainder = url.split("/", 3)
+            bundle_path = (
+                f"{StorageURLScheme.GCS_STORAGE.value}{bucket_name}/{bundle_uuid}/{contents_file}"
+            )
+
         is_archive = contents_file.endswith(".gz") or contents_file.endswith(".tar.gz")
         is_archive_dir = contents_file.endswith(".tar.gz")
         index_path = None
diff --git a/codalab/lib/bundle_store.py b/codalab/lib/bundle_store.py
index a4423cca..a173831b 100644
--- a/codalab/lib/bundle_store.py
+++ b/codalab/lib/bundle_store.py
@@ -403,7 +403,11 @@ class MultiDiskBundleStore(_MultiDiskBundleStoreBase):
 
     In Blob Storage, each bundle is stored in the format:
     azfs://{container name}/bundles/{bundle uuid}/contents.tar.gz if a directory,
-    azfs://{container name}/bundles/{bundle uuid}/contents.gz
+    azfs://{container name}/bundles/{bundle uuid}/contents.gz if a file.
+
+    In GCS, each bundle is stored in the format:
+    gs://{bucket name}/{bundle uuid}/contents.tar.gz if a directory,
+    gs://{bucket name}/{bundle uuid}/contents.gz if a file.
 
     If the bundle is a directory, the entire contents of the bundle is stored in the .tar.gz file;
     otherwise, if the bundle is a single file, the file is stored in the .gz file as an archive
@@ -448,7 +452,8 @@ class MultiDiskBundleStore(_MultiDiskBundleStoreBase):
                 # Next precedence: prefer blob storage over disk storage.
                 PRIORITY = 2
                 if (
-                    location["storage_type"] == StorageType.AZURE_BLOB_STORAGE.value
+                    location["storage_type"]
+                    in (StorageType.AZURE_BLOB_STORAGE.value, StorageType.GCS_STORAGE.value)
                     and PRIORITY < selected_location_priority
                 ):
                     selected_location = location
@@ -470,6 +475,14 @@ class MultiDiskBundleStore(_MultiDiskBundleStoreBase):
                 url = selected_location["url"]  # Format: "azfs://[container name]/bundles"
                 assert url.startswith("azfs://")
                 return f"{url}/{uuid}/{file_name}"
+            elif selected_location["storage_type"] == StorageType.GCS_STORAGE.value:
+                assert (
+                    selected_location["storage_format"] == StorageFormat.COMPRESSED_V1.value
+                )  # Only supported format on GCS
+                file_name = "contents.tar.gz" if is_dir else "contents.gz"
+                url = selected_location["url"]  # Format: "gs://[bucket name]"
+                assert url.startswith("gs://")
+                return f"{url}/{uuid}/{file_name}"
             else:
                 assert (
                     selected_location["storage_format"] == StorageFormat.UNCOMPRESSED.value
diff --git a/codalab/lib/upload_manager.py b/codalab/lib/upload_manager.py
index d19b1ec1..5992b0eb 100644
--- a/codalab/lib/upload_manager.py
+++ b/codalab/lib/upload_manager.py
@@ -237,7 +237,10 @@ class UploadManager(object):
         UploaderCls: Any = DiskStorageUploader
         if destination_bundle_store:
             # Set the uploader class based on which bundle store is specified.
-            if destination_bundle_store["storage_type"] == StorageType.AZURE_BLOB_STORAGE.value:
+            if destination_bundle_store["storage_type"] in (
+                StorageType.AZURE_BLOB_STORAGE.value,
+                StorageType.GCS_STORAGE.value,
+            ):
                 UploaderCls = BlobStorageUploader
         elif use_azure_blob_beta:
             # Legacy "-a" flag without specifying a bundle store.
diff --git a/codalab/model/tables.py b/codalab/model/tables.py
index 39b16608..5f76327b 100644
--- a/codalab/model/tables.py
+++ b/codalab/model/tables.py
@@ -19,7 +19,7 @@ from sqlalchemy.types import (
     Unicode,
 )
 from sqlalchemy.sql.schema import ForeignKeyConstraint
-from codalab.common import StorageType, StorageFormat
+from codalab.common import StorageType
 
 db_metadata = MetaData()
 
@@ -118,15 +118,9 @@ bundle_store = Table(
     # Name used to refer to this bundle store from the CLI.
     Column('name', String(255), nullable=False),
     # Storage type. This is usually redundant with information already in the url column, but set for efficiency reasons. When updating this column, sync it with codalab.common.StorageType.
-    Column(
-        'storage_type',
-        Enum(StorageType.DISK_STORAGE.value, StorageType.AZURE_BLOB_STORAGE.value),
-        nullable=True,
-    ),
+    Column('storage_type', String(255), nullable=True),
     # The format through which the bundle is stored.
-    Column(
-        'storage_format', Enum(StorageFormat.UNCOMPRESSED.value, StorageFormat.COMPRESSED_V1.value)
-    ),
+    Column('storage_format', String(255), nullable=True),
     # URL that uniquely identifies the bundle store.
     Column('url', String(255), nullable=True),
     # Authentication / password to access the bundle store.
diff --git a/codalab/rest/bundles.py b/codalab/rest/bundles.py
index 9b0c56f3..67364a8c 100644
--- a/codalab/rest/bundles.py
+++ b/codalab/rest/bundles.py
@@ -505,7 +505,7 @@ def _add_bundle_store():
     storage_type = new_bundle_store.get('storage_type')
     storage_format = new_bundle_store.get('storage_format')
     if storage_format is None:
-        if storage_type == StorageType.AZURE_BLOB_STORAGE.value:
+        if storage_type in (StorageType.AZURE_BLOB_STORAGE.value, StorageType.GCS_STORAGE.value):
             storage_format = StorageFormat.COMPRESSED_V1.value
         elif storage_type == StorageType.DISK_STORAGE.value:
             storage_format = StorageFormat.UNCOMPRESSED.value
diff --git a/codalab_service.py b/codalab_service.py
index dac5a207..18a33ad4 100755
--- a/codalab_service.py
+++ b/codalab_service.py
@@ -397,6 +397,12 @@ CODALAB_ARGUMENTS = [
         default='/tmp/codalab/link-mounts',
     ),
     CodalabArg(name='azure_blob_connection_string', help='Azure Blob storage connection string'),
+    CodalabArg(
+        name='google_application_credentials',
+        help='Path to Google Application Credentials file.',
+        # TODO: find a better default that is a no-op (we need a default value in order for docker-compose to mount this properly).
+        default='./codalab_service.py',
+    ),
     CodalabArg(
         name='always_use_azure_blob_beta',
         help='If set, Azure Blob Storage is always used to store uploads, regardless of the use_azure_blob_beta parameter.',
diff --git a/docker_config/compose_files/docker-compose.yml b/docker_config/compose_files/docker-compose.yml
index da40f648..120e1846 100644
--- a/docker_config/compose_files/docker-compose.yml
+++ b/docker_config/compose_files/docker-compose.yml
@@ -70,6 +70,7 @@ x-codalab-env: &codalab-env
   - COMPOSE_HTTP_TIMEOUT=${COMPOSE_HTTP_TIMEOUT}
   - DOCKER_CLIENT_TIMEOUT=${DOCKER_CLIENT_TIMEOUT}
   - CODALAB_AZURE_BLOB_CONNECTION_STRING=${CODALAB_AZURE_BLOB_CONNECTION_STRING}
+  - GOOGLE_APPLICATION_CREDENTIALS=/google-application-credentials.json
   - CODALAB_ALWAYS_USE_AZURE_BLOB_BETA=${CODALAB_ALWAYS_USE_AZURE_BLOB_BETA}
   - CODALAB_SENTRY_INGEST_URL=${CODALAB_SENTRY_INGEST_URL}
   - CODALAB_SENTRY_ENVIRONMENT=${CODALAB_SENTRY_ENVIRONMENT}
@@ -100,6 +101,7 @@ x-codalab-server: &codalab-server
   - "${CODALAB_HOME}:${CODALAB_HOME}"
   - "${CODALAB_BUNDLE_MOUNT}:${CODALAB_BUNDLE_MOUNT}"
   - "${CODALAB_MONITOR_DIR}:${CODALAB_MONITOR_DIR}"
+  - "${CODALAB_GOOGLE_APPLICATION_CREDENTIALS}:/google-application-credentials.json"
 
 ############################################################
 
diff --git a/docs/Multiple-Bundle-Stores.md b/docs/Multiple-Bundle-Stores.md
index 5115cafd..4760f387 100644
--- a/docs/Multiple-Bundle-Stores.md
+++ b/docs/Multiple-Bundle-Stores.md
@@ -36,9 +36,18 @@ cl store add --name blob --storage-type azure_blob --url azfs://devstoreaccount1
 cl upload --store blob mkdocs.yml
 ```
 
-## Steps to set up GCP
+## Steps to set up and test with GCP
+
+First, make sure the `CODALAB_GOOGLE_APPLICATION_CREDENTIALS` env var is set on the server that runs the REST server. It should be set to a path to a credentials JSON file.
+
+```
+export CODALAB_GOOGLE_APPLICATION_CREDENTIALS=/path/to/credentials.json
+codalab-service start -bd
+```
+
+Then make a bucket on your GCP account (in this case, it's called `ashwin123123123`). You can then store bundles on GCP:
 
 ```
-cl store add --name gcp --storage-type gcp --url gcs://bucket1
-cl upload --store gcp mkdocs.yml
+cl store add --name gcp2 --storage-type gcs --url gs://ashwin123123123
+cl upload --store gcp2 mkdocs.yml
 ```
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
index 4e3e1e85..2ed346f5 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,4 +1,4 @@
-apache-beam[azure]==2.32.0
+apache-beam[azure,gcp]==2.32.0
 bottle==0.12.19
 dataclasses==0.7;python_version<'3.7'
 docker==4.3.0


Test patch:
diff --git a/tests/unit/lib/path_util_test.py b/tests/unit/lib/path_util_test.py
index 5d3eafd1..757a2f52 100644
--- a/tests/unit/lib/path_util_test.py
+++ b/tests/unit/lib/path_util_test.py
@@ -118,7 +118,7 @@ class PathUtilTest(unittest.TestCase):
 
 class ParseBundleUrl(unittest.TestCase):
     def test_single_file(self):
-        """Parse a URL referring to a single file on Azure."""
+        """Parse a URL referring to a single file on Blob."""
         linked_bundle_path = parse_linked_bundle_url(
             "azfs://storageclwsdev0/bundles/uuid/contents.gz"
         )
@@ -135,6 +135,16 @@ class ParseBundleUrl(unittest.TestCase):
         self.assertEqual(linked_bundle_path.archive_subpath, None)
         self.assertEqual(linked_bundle_path.bundle_uuid, "uuid")
 
+        linked_bundle_path = parse_linked_bundle_url("gs://codalabbucket1/uuid/contents.gz")
+        self.assertEqual(linked_bundle_path.storage_type, StorageType.GCS_STORAGE.value)
+        self.assertEqual(linked_bundle_path.uses_beam, True)
+        self.assertEqual(linked_bundle_path.bundle_path, "gs://codalabbucket1/uuid/contents.gz")
+        self.assertEqual(linked_bundle_path.is_archive, True)
+        self.assertEqual(linked_bundle_path.is_archive_dir, False)
+        self.assertEqual(linked_bundle_path.index_path, "gs://codalabbucket1/uuid/index.sqlite")
+        self.assertEqual(linked_bundle_path.archive_subpath, None)
+        self.assertEqual(linked_bundle_path.bundle_uuid, "uuid")
+
     def test_directory(self):
         """Parse a URL referring to an archived directory."""
         linked_bundle_path = parse_linked_bundle_url(