Support functionalities to enhance task traceability with metadata for dependency search.

gokart/gcs_obj_metadata_client.py

@@ -9,6 +9,7 @@
 from googleapiclient.model import makepatch
 
 from gokart.gcs_config import GCSConfig
+from gokart.utils import FlattenableItems
 
 logger = getLogger(__name__)
 
@@ -32,7 +33,12 @@ def _path_to_bucket_and_key(path: str) -> tuple[str, str]:
         return netloc, path_without_initial_slash
 
     @staticmethod
-    def add_task_state_labels(path: str, task_params: dict[str, str] | None = None, custom_labels: dict[str, Any] | None = None) -> None:
+    def add_task_state_labels(
+        path: str,
+        task_params: dict[str, str] | None = None,
+        custom_labels: dict[str, Any] | None = None,
+        required_task_outputs: FlattenableItems[str] | None = None,
+    ) -> None:
         if GCSObjectMetadataClient._is_log_related_path(path):
             return
         # In gokart/object_storage.get_time_stamp, could find same call.
@@ -42,20 +48,17 @@ def add_task_state_labels(path: str, task_params: dict[str, str] | None = None,
         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,
             custom_labels,
         )
-
         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.
@@ -71,7 +74,6 @@ def add_task_state_labels(path: str, task_params: dict[str, str] | None = None,
                 )
                 .execute()
             )
-
             if update_response is None:
                 logger.error(f'failed to patch object {obj} in bucket {bucket} and object {obj}.')
 
@@ -90,7 +92,6 @@ def _get_patched_obj_metadata(
         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 and not custom_labels:
             return metadata
         # Maximum size of metadata for each object is 8 KiB.
@@ -132,10 +133,8 @@ def _get_label_size(label_name: str, label_value: str) -> int:
             8 * 1024,
             sum(_get_label_size(label_name, label_value) for label_name, label_value in labels.items()),
         )
-
         if current_total_metadata_size <= max_gcs_metadata_size:
             return labels
-
         for label_name, label_value in reversed(labels.items()):
             size = _get_label_size(label_name, label_value)
             del labels[label_name]

gokart/in_memory/target.py

@@ -5,6 +5,7 @@
 
 from gokart.in_memory.repository import InMemoryCacheRepository
 from gokart.target import TargetOnKart, TaskLockParams
+from gokart.utils import FlattenableItems
 
 _repository = InMemoryCacheRepository()
 
@@ -26,7 +27,13 @@ def _get_task_lock_params(self) -> TaskLockParams:
     def _load(self) -> Any:
         return _repository.get_value(self._data_key)
 
-    def _dump(self, obj: Any, task_params: dict[str, str] | None = None, custom_labels: dict[str, Any] | None = None) -> None:
+    def _dump(
+        self,
+        obj: Any,
+        task_params: dict[str, str] | None = None,
+        custom_labels: dict[str, Any] | None = None,
+        required_task_outputs: FlattenableItems[str] | None = None,
+    ) -> None:
         return _repository.set_value(self._data_key, obj)
 
     def _remove(self) -> None:

gokart/target.py

@@ -18,6 +18,7 @@
 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.utils import FlattenableItems
 from gokart.zip_client_util import make_zip_client
 
 logger = getLogger(__name__)
@@ -30,13 +31,23 @@ 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, task_params: dict[str, str] | None = None, custom_labels: dict[str, Any] | None = None) -> None:
+    def dump(
+        self,
+        obj,
+        lock_at_dump: bool = True,
+        task_params: dict[str, str] | None = None,
+        custom_labels: dict[str, Any] | None = None,
+        required_task_outputs: FlattenableItems | None = 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=obj, task_params=task_params, custom_labels=custom_labels
+                obj=obj,
+                task_params=task_params,
+                custom_labels=custom_labels,
+                required_task_outputs=required_task_outputs,
             )
         else:
-            self._dump(obj=obj, task_params=task_params, custom_labels=custom_labels)
+            self._dump(obj=obj, task_params=task_params, custom_labels=custom_labels, required_task_outputs=required_task_outputs)
 
     def remove(self) -> None:
         if self.exists():
@@ -61,7 +72,13 @@ def _load(self) -> Any:
         pass
 
     @abstractmethod
-    def _dump(self, obj, task_params: dict[str, str] | None = None, custom_labels: dict[str, Any] | None = None) -> None:
+    def _dump(
+        self,
+        obj,
+        task_params: Optional[dict[str, str]] = None,
+        custom_labels: dict[str, Any] | None = None,
+        required_task_outputs: FlattenableItems[str] | None = None,
+    ) -> None:
         pass
 
     @abstractmethod
@@ -98,11 +115,19 @@ def _load(self) -> Any:
         with self._target.open('r') as f:
             return self._processor.load(f)
 
-    def _dump(self, obj, task_params: dict[str, str] | None = None, custom_labels: dict[str, Any] | None = None) -> None:
+    def _dump(
+        self,
+        obj,
+        task_params: dict[str, str] | None = None,
+        custom_labels: dict[str, Any] | None = None,
+        required_task_outputs: FlattenableItems[str] | None = 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, custom_labels=custom_labels)
+            GCSObjectMetadataClient.add_task_state_labels(
+                path=self.path(), task_params=task_params, custom_labels=custom_labels, required_task_outputs=required_task_outputs
+            )
 
     def _remove(self) -> None:
         self._target.remove()
@@ -142,10 +167,18 @@ def _load(self) -> Any:
         self._remove_temporary_directory()
         return model
 
-    def _dump(self, obj, task_params: dict[str, str] | None = None, custom_labels: dict[str, Any] | None = None) -> None:
+    def _dump(
+        self,
+        obj,
+        task_params: dict[str, str] | None = None,
+        custom_labels: dict[str, Any] | None = None,
+        required_task_outputs: FlattenableItems[str] | None = None,
+    ) -> None:
         self._make_temporary_directory()
         self._save_function(obj, self._model_path())
-        make_target(self._load_function_path()).dump(self._load_function, task_params=task_params)
+        make_target(self._load_function_path()).dump(
+            self._load_function, task_params=task_params, custom_labels=custom_labels, required_task_outputs=required_task_outputs
+        )
         self._zip_client.make_archive()
         self._remove_temporary_directory()
 

gokart/task.py

@@ -11,6 +11,8 @@
 from logging import getLogger
 from typing import Any, Callable, Dict, Generator, Generic, Iterable, List, Optional, Set, TypeVar, Union, overload
 
+from gokart.utils import map_flattenable_items
+
 if sys.version_info < (3, 13):
     from typing_extensions import deprecated
 else:
@@ -367,6 +369,10 @@ def dump(self, obj: Any, target: Union[None, str, TargetOnKart] = None, custom_l
             lock_at_dump=self._lock_at_dump,
             task_params=super().to_str_params(only_significant=True, only_public=True),
             custom_labels=custom_labels,
+            required_task_outputs=map_flattenable_items(
+                self.requires(),
+                func=lambda task: map_flattenable_items(task.output(), func=lambda output: output.path()),
+            ),
         )
 
     @staticmethod

gokart/utils.py

@@ -3,7 +3,7 @@
 import os
 import sys
 from io import BytesIO
-from typing import Any, Iterable, Protocol, TypeVar, Union
+from typing import Any, Callable, Iterable, Protocol, TypeVar, Union
 
 import dill
 import luigi
@@ -71,6 +71,19 @@ def flatten(targets: FlattenableItems[T]) -> list[T]:
     return flat
 
 
+K = TypeVar('K')
+
+
+def map_flattenable_items(items: FlattenableItems[T], func: Callable[[T], K]) -> FlattenableItems[K]:
+    if isinstance(items, dict):
+        return {k: map_flattenable_items(v, func) for k, v in items.items()}
+    if isinstance(items, str):
+        return items  # type: ignore
+    if isinstance(items, Iterable):
+        return [map_flattenable_items(i, func) for i in items]
+    return func(items)
+
+
 def load_dill_with_pandas_backward_compatibility(file: Union[FileLike, BytesIO]) -> Any:
     """Load binary dumped by dill with pandas backward compatibility.
     pd.read_pickle can load binary dumped in backward pandas version, and also any objects dumped by pickle.

test/test_gcs_obj_metadata_client.py

@@ -120,7 +120,10 @@ def test_mock_target_on_kart(self, mock_get_output_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={}, custom_labels=None)
+
+        mock_target.dump.assert_called_once_with(
+            {'key': 'value'}, lock_at_dump=task._lock_at_dump, task_params={}, custom_labels=None, required_task_outputs=[]
+        )
 
 
 if __name__ == '__main__':

test/test_utils.py

@@ -1,6 +1,6 @@
 import unittest
 
-from gokart.utils import flatten
+from gokart.utils import flatten, map_flattenable_items
 
 
 class TestFlatten(unittest.TestCase):
@@ -18,3 +18,12 @@ def test_flatten_int(self):
 
     def test_flatten_none(self):
         self.assertEqual(flatten(None), [])
+
+
+class TestMapFlatten(unittest.TestCase):
+    def test_map_flattenable_items(self):
+        self.assertEqual(map_flattenable_items({'a': 1, 'b': 2}, func=lambda x: str(x)), {'a': '1', 'b': '2'})
+        self.assertEqual(
+            map_flattenable_items({'a': [1, 2, 3, '4'], 'b': {'c': True, 'd': {'e': 5}}}, func=lambda x: str(x)),
+            {'a': ['1', '2', '3', '4'], 'b': {'c': 'True', 'd': {'e': '5'}}},
+        )