Support functionalities to provide metadata as labels to object stored in GCS as cache. (#445)

* WIP: implement GCSObjectMetadataClient class to attach custom-metadata to gcs-object.

* WIP: change dump method interface.

* WIP: add user_provided_gcs_labels parameter to TaskOnKart.

* Test: add test of GCSObjectMetadataClient.

* feat: dealed with nits PR comments.

* feat: Remove user_provided_labels feature. This feature will be supported another PR.

* for-PR: apply almost all comments.

* fix: change Dict to dict.

* feat: add gokart specific parameter serialize test.

* fix: fix testcases with literals and more meaningful assertion.

* feat: add mock testcase.

* fix: fix CI errors.

* feat: deal with pr comments, and modify testcases.

* feat: deal with kitagry comments.

* feat: deal with PR comments, typle annotation, sealed private parameters.

---------

Co-authored-by: 荒木 太一 <taichi.araki@m3-2024mac75.local>

Author: tyzerrr

gokart/gcs_obj_metadata_client.py

@@ -0,0 +1,100 @@
+from __future__ import annotations
+
+import copy
+from logging import getLogger
+from typing import Any, Optional, Union
+from urllib.parse import urlsplit
+
+from googleapiclient.model import makepatch
+
+from gokart.gcs_config import GCSConfig
+
+logger = getLogger(__name__)
+
+
+class GCSObjectMetadataClient:
+    """
+    This class is Utility-Class, so should not be initialized.
+    This class used for adding metadata as labels.
+    """
+
+    # This is the copied method of luigi.gcs._path_to_bucket_and_key(path).
+    @staticmethod
+    def path_to_bucket_and_key(path: str) -> tuple[str, str]:
+        (scheme, netloc, path, _, _) = urlsplit(path)
+        assert scheme == 'gs'
+        path_without_initial_slash = path[1:]
+        return netloc, path_without_initial_slash
+
+    @staticmethod
+    def add_task_state_labels(
+        path: str,
+        task_params: Optional[dict[Any, str]] = None,
+    ) -> None:
+        # In gokart/object_storage.get_time_stamp, could find same call.
+        # _path_to_bucket_and_key is a private method, so, this might not be acceptable.
+        bucket, obj = GCSObjectMetadataClient.path_to_bucket_and_key(path)
+
+        _response = GCSConfig().get_gcs_client().client.objects().get(bucket=bucket, object=obj).execute()
+        if _response is None:
+            logger.error(f'failed to get object from GCS bucket {bucket} and object {obj}.')
+            return
+
+        response: dict[str, Any] = dict(_response)
+        original_metadata: dict[Any, Any] = {}
+        if 'metadata' in response.keys():
+            _metadata = response.get('metadata')
+            if _metadata is not None:
+                original_metadata = dict(_metadata)
+
+        patched_metadata = GCSObjectMetadataClient._get_patched_obj_metadata(
+            copy.deepcopy(original_metadata),
+            task_params,
+        )
+
+        if original_metadata != patched_metadata:
+            # If we use update api, existing object metadata are removed, so should use patch api.
+            # See the official document descriptions.
+            # [Link] https://cloud.google.com/storage/docs/viewing-editing-metadata?hl=ja#rest-set-object-metadata
+            update_response = (
+                GCSConfig()
+                .get_gcs_client()
+                .client.objects()
+                .patch(
+                    bucket=bucket,
+                    object=obj,
+                    body=makepatch({'metadata': original_metadata}, {'metadata': patched_metadata}),
+                )
+                .execute()
+            )
+
+            if update_response is None:
+                logger.error(f'failed to patch object {obj} in bucket {bucket} and object {obj}.')
+
+    @staticmethod
+    def _get_patched_obj_metadata(
+        metadata: Any,
+        task_params: Optional[dict[Any, str]] = None,
+    ) -> Union[dict, Any]:
+        # If metadata from response when getting bucket and object information is not dictionary,
+        # something wrong might be happened, so return original metadata, no patched.
+        if not isinstance(metadata, dict):
+            logger.warning(f'metadata is not a dict: {metadata}, something wrong was happened when getting response when get bucket and object information.')
+            return metadata
+
+        if not task_params:
+            return metadata
+        # Maximum size of metadata for each object is 8 KiB.
+        # [Link]: https://cloud.google.com/storage/quotas#objects
+        max_gcs_metadata_size, total_metadata_size, labels = 8 * 1024, 0, []
+        for label_name, label_value in task_params.items():
+            if len(label_value) == 0:
+                logger.warning(f'value of label_name={label_name} is empty. So skip to add as a metadata.')
+                continue
+            size = len(str(label_name).encode('utf-8')) + len(str(label_value).encode('utf-8'))
+            if total_metadata_size + size > max_gcs_metadata_size:
+                logger.warning(f'current metadata total size is {total_metadata_size} byte, and no more labels would be added.')
+                break
+            total_metadata_size += size
+            labels.append((label_name, label_value))
+        return dict(metadata) | dict(labels)

gokart/in_memory/target.py

@@ -1,5 +1,5 @@
 from datetime import datetime
-from typing import Any
+from typing import Any, Optional
 
 from gokart.in_memory.repository import InMemoryCacheRepository
 from gokart.target import TargetOnKart, TaskLockParams
@@ -24,7 +24,7 @@ def _get_task_lock_params(self) -> TaskLockParams:
     def _load(self) -> Any:
         return _repository.get_value(self._data_key)
 
-    def _dump(self, obj: Any) -> None:
+    def _dump(self, obj: Any, task_params: Optional[dict[str, str]] = None) -> None:
         return _repository.set_value(self._data_key, obj)
 
     def _remove(self) -> None:

gokart/target.py

@@ -14,6 +14,7 @@
 from gokart.conflict_prevention_lock.task_lock import TaskLockParams, make_task_lock_params
 from gokart.conflict_prevention_lock.task_lock_wrappers import wrap_dump_with_lock, wrap_load_with_lock, wrap_remove_with_lock
 from gokart.file_processor import FileProcessor, make_file_processor
+from gokart.gcs_obj_metadata_client import GCSObjectMetadataClient
 from gokart.object_storage import ObjectStorage
 from gokart.zip_client_util import make_zip_client
 
@@ -27,11 +28,11 @@ def exists(self) -> bool:
     def load(self) -> Any:
         return wrap_load_with_lock(func=self._load, task_lock_params=self._get_task_lock_params())()
 
-    def dump(self, obj, lock_at_dump: bool = True) -> None:
+    def dump(self, obj, lock_at_dump: bool = True, task_params: Optional[dict[str, str]] = None) -> None:
         if lock_at_dump:
-            wrap_dump_with_lock(func=self._dump, task_lock_params=self._get_task_lock_params(), exist_check=self.exists)(obj)
+            wrap_dump_with_lock(func=self._dump, task_lock_params=self._get_task_lock_params(), exist_check=self.exists)(obj=obj, task_params=task_params)
         else:
-            self._dump(obj)
+            self._dump(obj=obj, task_params=task_params)
 
     def remove(self) -> None:
         if self.exists():
@@ -56,7 +57,7 @@ def _load(self) -> Any:
         pass
 
     @abstractmethod
-    def _dump(self, obj) -> None:
+    def _dump(self, obj, task_params: Optional[dict[str, str]] = None) -> None:
         pass
 
     @abstractmethod
@@ -93,9 +94,11 @@ def _load(self) -> Any:
         with self._target.open('r') as f:
             return self._processor.load(f)
 
-    def _dump(self, obj) -> None:
+    def _dump(self, obj, task_params: Optional[dict[str, str]] = None) -> None:
         with self._target.open('w') as f:
             self._processor.dump(obj, f)
+        if self.path().startswith('gs://'):
+            GCSObjectMetadataClient.add_task_state_labels(path=self.path(), task_params=task_params)
 
     def _remove(self) -> None:
         self._target.remove()
@@ -135,10 +138,10 @@ def _load(self) -> Any:
         self._remove_temporary_directory()
         return model
 
-    def _dump(self, obj) -> None:
+    def _dump(self, obj, task_params: Optional[dict[str, str]] = None) -> None:
         self._make_temporary_directory()
         self._save_function(obj, self._model_path())
-        make_target(self._load_function_path()).dump(self._load_function)
+        make_target(self._load_function_path()).dump(self._load_function, task_params=task_params)
         self._zip_client.make_archive()
         self._remove_temporary_directory()
 

gokart/task.py

@@ -359,7 +359,8 @@ def dump(self, obj: Any, target: Union[None, str, TargetOnKart] = None) -> None:
         PandasTypeConfigMap().check(obj, task_namespace=self.task_namespace)
         if self.fail_on_empty_dump and isinstance(obj, pd.DataFrame):
             assert not obj.empty
-        self._get_output_target(target).dump(obj, lock_at_dump=self._lock_at_dump)
+
+        self._get_output_target(target).dump(obj, lock_at_dump=self._lock_at_dump, task_params=super().to_str_params(only_significant=True, only_public=True))
 
     @staticmethod
     def get_code(target_class) -> Set[str]:

test/test_gcs_obj_metadata_client.py

@@ -0,0 +1,62 @@
+import datetime
+import unittest
+from typing import Any
+from unittest.mock import MagicMock, patch
+
+import gokart
+from gokart.gcs_obj_metadata_client import GCSObjectMetadataClient
+from gokart.target import TargetOnKart
+
+
+class _DummyTaskOnKart(gokart.TaskOnKart):
+    task_namespace = __name__
+
+    def run(self):
+        self.dump('Dummy TaskOnKart')
+
+
+class TestGCSObjectMetadataClient(unittest.TestCase):
+    def test_get_patched_obj_metadata(self):
+        task_params: dict[Any, str] = {
+            'param1': 'a' * 1000,
+            'param2': str(1000),
+            'param3': str({'key1': 'value1', 'key2': True, 'key3': 2}),
+            'param4': str([1, 2, 3, 4, 5]),
+            'param5': str(datetime.datetime(year=2025, month=1, day=2, hour=3, minute=4, second=5)),
+            'param6': '',
+        }
+        got = GCSObjectMetadataClient._get_patched_obj_metadata({}, task_params=task_params)
+        self.assertIsInstance(got, dict)
+        self.assertIn('param1', got)
+        self.assertIn('param2', got)
+        self.assertIn('param3', got)
+        self.assertIn('param4', got)
+        self.assertIn('param5', got)
+        self.assertNotIn('param6', got)
+
+    def test_get_patched_obj_metadata_with_exceeded_size_metadata(self):
+        task_params = {
+            'param1': 'a' * 5000,
+            'param2': 'b' * 5000,
+        }
+        want = {
+            'param1': 'a' * 5000,
+        }
+        got = GCSObjectMetadataClient._get_patched_obj_metadata({}, task_params=task_params)
+        self.assertEqual(got, want)
+
+
+class TestGokartTask(unittest.TestCase):
+    @patch.object(_DummyTaskOnKart, '_get_output_target')
+    def test_mock_target_on_kart(self, mock_get_output_target):
+        mock_target = MagicMock(spec=TargetOnKart)
+        mock_get_output_target.return_value = mock_target
+
+        task = _DummyTaskOnKart()
+        task.dump({'key': 'value'}, mock_target)
+
+        mock_target.dump.assert_called_once_with({'key': 'value'}, lock_at_dump=task._lock_at_dump, task_params={})
+
+
+if __name__ == '__main__':
+    unittest.main()