Fix: dictionary changed size during iteration in GCSObjectMetadataClient (#468)

* Fix: dictionary changed size during iteration in GCSObjectMetadataClient

* chroe:ruff

* refactor: Convert lists to tuples to reduce memory usage

* fix: Use constant for GCS metadata size

* fix: Remove unnecessary result assignment in metadata size adjustment test

---------

Co-authored-by: Masayuki-Kamoda <[email protected]>

Author: ikyasam18

gokart/gcs_obj_metadata_client.py

@@ -6,7 +6,7 @@
 import re
 from collections.abc import Iterable
 from logging import getLogger
-from typing import Any
+from typing import Any, Final
 from urllib.parse import urlsplit
 
 from googleapiclient.model import makepatch
@@ -24,6 +24,9 @@ class GCSObjectMetadataClient:
     This class used for adding metadata as labels.
     """
 
+    # Maximum metadata size for GCS objects (8 KiB)
+    MAX_GCS_METADATA_SIZE: Final[int] = 8 * 1024
+
     @staticmethod
     def _is_log_related_path(path: str) -> bool:
         return re.match(r'^gs://.+?/log/(processing_time/|task_info/|task_log/|module_versions/|random_seed/|task_params/).+', path) is not None
@@ -154,15 +157,20 @@ def _get_label_size(label_name: str, label_value: str) -> int:
 
         labels = copy.deepcopy(_labels)
         max_gcs_metadata_size, current_total_metadata_size = (
-            8 * 1024,
+            GCSObjectMetadataClient.MAX_GCS_METADATA_SIZE,
             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()):
+        # NOTE: remove labels to stay within max metadata size.
+        to_remove = []
+        for label_name, label_value in reversed(tuple(labels.items())):
             size = _get_label_size(label_name, label_value)
-            del labels[label_name]
+            to_remove.append(label_name)
             current_total_metadata_size -= size
             if current_total_metadata_size <= max_gcs_metadata_size:
                 break
+
+        for key in to_remove:
+            del labels[key]
         return labels

test/test_gcs_obj_metadata_client.py

@@ -140,6 +140,13 @@ def test_get_patched_obj_metadata_with_nested_required_task_outputs(self):
             got['__required_task_outputs'], '{"nested_task": {"nest": {"__gokart_task_name": "task1", "__gokart_output_path": "path/to/output1"}}}'
         )
 
+    def test_adjust_gcs_metadata_limit_size_runtime_error(self):
+        large_labels = {}
+        for i in range(100):
+            large_labels[f'key_{i}'] = 'x' * 1000
+
+        GCSObjectMetadataClient._adjust_gcs_metadata_limit_size(large_labels)
+
 
 class TestGokartTask(unittest.TestCase):
     @patch.object(_DummyTaskOnKart, '_get_output_target')