Add asl module lib

Makefile

@@ -32,7 +32,7 @@ clean:
 install:
 	@pip install --user -U pip
 	@pip install --user "Cython<3"
-	@pip install --user -r requirements.txt
+	@pip install --user -e .
 	@pip install --user --no-deps -r requirements-without-deps.txt
 	@./scripts/setup_on_jupyterlab.sh
 	@pre-commit install
@@ -61,3 +61,7 @@ object_detection_kernel:
 .PHONY: pytorch_kfp_kernel
 pytorch_kfp_kernel:
 	./kernels/pytorch_kfp.sh
+
+.PHONY: tests
+tests:
+	pytest tests/unit

asl/lib.py

@@ -0,0 +1,52 @@
+""" Utils and GCP APIs wrappers.
+"""
+
+import os
+
+from google.cloud import storage
+
+
+def extract_relative_path(root_path, path):
+    """If root_path=/tmp/a and path=/tmp/a/b then return a/b."""
+    root_path = os.path.abspath(root_path)
+    subpath = path.replace(root_path, "").strip("/")
+    basedir = os.path.basename(root_path)
+    relative_path = os.path.join(basedir, subpath)
+    return relative_path
+
+
+def upload_directory_to_gcs(bucket_name, folder_path, bucket_prefix=None):
+    """Upload directory to GCP bucket at bucket_prefix."""
+    storage_client = storage.Client()
+    bucket = storage_client.get_bucket(bucket_name)
+    for path, _, files in os.walk(folder_path):
+        for name in files:
+            path_local = os.path.join(path, name)
+            blob_path = extract_relative_path(folder_path, path_local)
+            if bucket_prefix:
+                blob_path = os.path.join(bucket_prefix, blob_path)
+            blob = bucket.blob(blob_path)
+            blob.upload_from_filename(path_local)
+
+
+def download_directory_from_gcs(bucket_name, local_folder, bucket_prefix):
+    """Download gs://bucket_name/a/b to local_folder/b (bucket_prefix=a/b)."""
+    storage_client = storage.Client()
+    local_folder = os.path.abspath(local_folder)
+    if os.path.isdir(local_folder) is False:
+        os.makedirs(local_folder)
+
+    bucket = storage_client.bucket(bucket_name=bucket_name)
+    blobs = bucket.list_blobs(prefix=bucket_prefix)
+
+    for blob in blobs:
+        blob_name = blob.name
+        dst_file = os.path.join(
+            local_folder,
+            os.path.basename(bucket_prefix),
+            blob_name.replace(bucket_prefix, "").strip("/"),
+        )
+        dst_dir = os.path.dirname(dst_file)
+        if not os.path.isdir(dst_dir):
+            os.makedirs(dst_dir)
+        blob.download_to_filename(dst_file)

requirements.txt

@@ -15,3 +15,4 @@ keras_nlp==0.5.1
 
 # For development work
 pre-commit
+pytest

setup.py

@@ -0,0 +1,15 @@
+""" Setup for pip install -e .
+"""
+
+from setuptools import find_packages, setup
+
+REQS = None
+with open("requirements.txt", encoding="utf-8") as fp:
+    REQS = [req.strip() for req in fp.readlines()]
+
+setup(
+    name="asl",
+    version="0.1",
+    packages=find_packages(),
+    install_requires=REQS,
+)

tests/unit/lib_test.py

@@ -0,0 +1,10 @@
+""" Unit tests for lib.py
+"""
+
+from asl import lib
+
+
+def test_extract_relative_path():
+    test_root_path = "/tmp/a"
+    test_path = "/tmp/a/b"
+    assert "a/b" == lib.extract_relative_path(test_root_path, test_path)