feat: ROOT-11: Support reading JSONL from source cloud storages

label_studio/core/settings/base.py

@@ -598,6 +598,7 @@
 MEMBER_PERM = 'core.api_permissions.MemberHasOwnerPermission'
 RECALCULATE_ALL_STATS = None
 GET_STORAGE_LIST = 'io_storages.functions.get_storage_list'
+STORAGE_LOAD_TASKS_JSON = 'io_storages.utils.load_tasks_json_lso'
 STORAGE_ANNOTATION_SERIALIZER = 'io_storages.serializers.StorageAnnotationSerializer'
 TASK_SERIALIZER_BULK = 'tasks.serializers.BaseTaskSerializerBulk'
 PREPROCESS_FIELD_NAME = 'data_manager.functions.preprocess_field_name'

label_studio/io_storages/azure_blob/models.py

@@ -25,7 +25,12 @@
     ImportStorageLink,
     ProjectStorageMixin,
 )
-from io_storages.utils import parse_range, storage_can_resolve_bucket_url
+from io_storages.utils import (
+    StorageObject,
+    load_tasks_json,
+    parse_range,
+    storage_can_resolve_bucket_url,
+)
 from tasks.models import Annotation
 
 from label_studio.io_storages.azure_blob.utils import AZURE
@@ -209,28 +214,16 @@ def iterkeys(self):
                 continue
             yield file.name
 
-    def get_data(self, key) -> Union[dict, list[dict]]:
+    def get_data(self, key) -> list[StorageObject]:
         if self.use_blob_urls:
             data_key = settings.DATA_UNDEFINED_NAME
-            return {data_key: f'{self.url_scheme}://{self.container}/{key}'}
+            task = {data_key: f'{self.url_scheme}://{self.container}/{key}'}
+            return [StorageObject(key=key, task_data=task)]
 
         container = self.get_container()
         blob = container.download_blob(key)
-        blob_str = blob.content_as_text()
-        value = json.loads(blob_str)
-        if isinstance(value, dict):
-            return value
-        elif isinstance(value, list):
-            for idx, item in enumerate(value):
-                if not isinstance(item, dict):
-                    raise ValueError(
-                        f'Error on key {key} item {idx}: For {self.__class__.__name__} your JSON file must be a dictionary with one task, or a list of dictionaries with one task each'
-                    )
-            return value
-        else:
-            raise ValueError(
-                f'Error on key {key}: For {self.__class__.__name__} your JSON file must be a dictionary with one task, or a list of dictionaries with one task each'
-            )
+        blob_str = blob.content_as_bytes()
+        return load_tasks_json(blob_str, key)
 
     def scan_and_create_links(self):
         return self._scan_and_create_links(AzureBlobImportStorageLink)

label_studio/io_storages/base_models.py

@@ -8,6 +8,7 @@
 import os
 import traceback as tb
 from concurrent.futures import ThreadPoolExecutor
+from dataclasses import asdict
 from datetime import datetime
 from typing import Union
 from urllib.parse import urljoin
@@ -27,7 +28,7 @@
 from django.utils import timezone
 from django.utils.translation import gettext_lazy as _
 from django_rq import job
-from io_storages.utils import get_uri_via_regex, parse_bucket_uri
+from io_storages.utils import StorageObject, get_uri_via_regex, parse_bucket_uri
 from rq.job import Job
 from tasks.models import Annotation, Task
 from tasks.serializers import AnnotationSerializer, PredictionSerializer
@@ -230,7 +231,7 @@
     def iterkeys(self):
         return iter(())
 
-    def get_data(self, key) -> list[dict]:
+    def get_data(self, key) -> list[StorageObject]:
         raise NotImplementedError
 
     def generate_http_url(self, url):
@@ -341,17 +342,20 @@
         raise NotImplementedError
 
     @classmethod
-    def add_task(cls, data, project, maximum_annotations, max_inner_id, storage, key, row_index, link_class):
+    def add_task(cls, project, maximum_annotations, max_inner_id, storage, link_object: StorageObject, link_class):
+        link_kwargs = asdict(link_object)
+        data = link_kwargs.pop('task_data', None)
+
         # predictions
-        predictions = data.get('predictions', [])
+        predictions = data.get('predictions') or []
         if predictions:
             if 'data' not in data:
                 raise ValueError(
                     'If you use "predictions" field in the task, ' 'you must put "data" field in the task too'
                 )
 
         # annotations
-        annotations = data.get('annotations', [])
+        annotations = data.get('annotations') or []
         cancelled_annotations = 0
         if annotations:
             if 'data' not in data:
@@ -361,7 +365,10 @@
             cancelled_annotations = len([a for a in annotations if a.get('was_cancelled', False)])
 
         if 'data' in data and isinstance(data['data'], dict):
-            data = data['data']
+            if data['data'] is not None:
+                data = data['data']
+            else:
+                data.pop('data')
 
         with transaction.atomic():
             task = Task.objects.create(
@@ -375,8 +382,8 @@
                 inner_id=max_inner_id,
             )
 
-            link_class.create(task, key, storage, row_index=row_index)
-            logger.debug(f'Create {storage.__class__.__name__} link with {key=} and {row_index=} for {task=}')
+            link_class.create(task, storage=storage, **link_kwargs)
+            logger.debug(f'Create {storage.__class__.__name__} link with {link_kwargs} for {task=}')
 
             raise_exception = not flag_set(
                 'ff_fix_back_dev_3342_storage_scan_with_invalid_annotations', user=AnonymousUser()
@@ -431,7 +438,7 @@
 
             logger.debug(f'{self}: found new key {key}')
             try:
-                tasks_data = self.get_data(key)
+                link_objects = self.get_data(key)
             except (UnicodeDecodeError, json.decoder.JSONDecodeError) as exc:
                 logger.debug(exc, exc_info=True)
                 raise ValueError(
@@ -440,19 +447,19 @@
                     f'"Treat every bucket object as a source file"'
                 )
 
-            if isinstance(tasks_data, dict):
-                tasks_data = [tasks_data]
-                row_indices = [None]
-            else:
-                if not flag_set('fflag_feat_dia_2092_multitasks_per_storage_link'):
-                    tasks_data = tasks_data[:1]
-                row_indices = range(len(tasks_data))
+            if not flag_set('fflag_feat_dia_2092_multitasks_per_storage_link'):
+                link_objects = link_objects[:1]
 
-            for row_index, task_data in zip(row_indices, tasks_data):
+            for link_object in link_objects:
                 # TODO: batch this loop body with add_task -> add_tasks in a single bulk write.
                 # See DIA-2062 for prerequisites
                 task = self.add_task(
-                    task_data, self.project, maximum_annotations, max_inner_id, self, key, row_index, link_class
+                    self.project,
+                    maximum_annotations,
+                    max_inner_id,
+                    self,
+                    link_object,
+                    link_class=link_class,
                 )
                 max_inner_id += 1
 
@@ -515,6 +522,8 @@
                 self.info_set_queued()
                 import_sync_background(self.__class__, self.id)
             except Exception:
+                # needed to facilitate debugging storage-related testcases, since otherwise no exception is logged
+                logger.debug(f'Storage {self} failed', exc_info=True)
                 storage_background_failure(self)
 
     class Meta:

label_studio/io_storages/gcs/models.py

@@ -23,7 +23,12 @@
     ProjectStorageMixin,
 )
 from io_storages.gcs.utils import GCS
-from io_storages.utils import parse_range, storage_can_resolve_bucket_url
+from io_storages.utils import (
+    StorageObject,
+    load_tasks_json,
+    parse_range,
+    storage_can_resolve_bucket_url,
+)
 from tasks.models import Annotation
 
 logger = logging.getLogger(__name__)
@@ -180,28 +185,16 @@ def iterkeys(self):
             return_key=True,
         )
 
-    def get_data(self, key) -> Union[dict, list[dict]]:
+    def get_data(self, key) -> list[StorageObject]:
         if self.use_blob_urls:
-            return {settings.DATA_UNDEFINED_NAME: GCS.get_uri(self.bucket, key)}
-        data = GCS.read_file(
+            task = {settings.DATA_UNDEFINED_NAME: GCS.get_uri(self.bucket, key)}
+            return [StorageObject(key=key, task_data=task)]
+        blob_str = GCS.read_file(
             client=self.get_client(),
             bucket_name=self.bucket,
             key=key,
-            convert_to=GCS.ConvertBlobTo.JSON,
         )
-        if isinstance(data, dict):
-            return data
-        elif isinstance(data, list):
-            for idx, item in enumerate(data):
-                if not isinstance(item, dict):
-                    raise ValueError(
-                        f'Error on key {key} item {idx}: For {self.__class__.__name__} your JSON file must be a dictionary with one task, or a list of dictionaries with one task each'
-                    )
-            return data
-        else:
-            raise ValueError(
-                f'Error on key {key}: For {self.__class__.__name__} your JSON file must be a dictionary with one task, or a list of dictionaries with one task each'
-            )
+        return load_tasks_json(blob_str, key)
 
     def generate_http_url(self, url):
         return GCS.generate_http_url(

label_studio/io_storages/gcs/utils.py

@@ -250,34 +250,15 @@ def iter_images_filename(cls, client, bucket_name, max_files):
     def get_uri(cls, bucket_name, key):
         return f'gs://{bucket_name}/{key}'
 
-    @classmethod
-    def _try_read_json(cls, blob_str):
-        try:
-            data = json.loads(blob_str)
-        except ValueError:
-            logger.error(f"Can't parse JSON from {blob_str}")
-            return
-        return data
-
     @classmethod
     def read_file(
         cls, client: gcs.Client, bucket_name: str, key: str, convert_to: ConvertBlobTo = ConvertBlobTo.NOTHING
     ):
         bucket = client.get_bucket(bucket_name)
         blob = bucket.blob(key)
         blob_str = blob.download_as_bytes()
-        if convert_to == cls.ConvertBlobTo.NOTHING:
-            return blob_str
-        elif convert_to == cls.ConvertBlobTo.JSON:
-            return cls._try_read_json(blob_str)
-        elif convert_to == cls.ConvertBlobTo.JSON_DICT:
-            json_data = cls._try_read_json(blob_str)
-            if not isinstance(json_data, dict):
-                raise ValueError(
-                    f'Error on key {key}: For {cls.__name__} your JSON file must be a dictionary with one task.'
-                )
-            return json_data
-        elif convert_to == cls.ConvertBlobTo.BASE64:
+
+        if convert_to == cls.ConvertBlobTo.BASE64:
             return base64.b64encode(blob_str)
 
         return blob_str

label_studio/io_storages/localfiles/models.py

@@ -20,6 +20,7 @@
     ImportStorageLink,
     ProjectStorageMixin,
 )
+from io_storages.utils import StorageObject, load_tasks_json
 from rest_framework.exceptions import ValidationError
 from tasks.models import Annotation
 
@@ -78,39 +79,24 @@
                     continue
                 yield str(file)
 
-    def get_data(self, key) -> dict | list[dict]:
+    def get_data(self, key) -> list[StorageObject]:
         path = Path(key)
         if self.use_blob_urls:
             # include self-hosted links pointed to local resources via
             # {settings.HOSTNAME}/data/local-files?d=<path/to/local/dir>
             document_root = Path(settings.LOCAL_FILES_DOCUMENT_ROOT)
             relative_path = str(path.relative_to(document_root))
-            return {
+            task = {
                 settings.DATA_UNDEFINED_NAME: f'{settings.HOSTNAME}/data/local-files/?d={quote(str(relative_path))}'
             }
+            return [StorageObject(key=key, task_data=task)]
 
         try:
-            with open(path, encoding='utf8') as f:
-                value = json.load(f)
-        except (UnicodeDecodeError, json.decoder.JSONDecodeError):
-            raise ValueError(
-                f"Can't import JSON-formatted tasks from {key}. If you're trying to import binary objects, "
-                f'perhaps you\'ve forgot to enable "Treat every bucket object as a source file" option?'
-            )
-
-        if isinstance(value, dict):
-            return value
-        elif isinstance(value, list):
-            for idx, item in enumerate(value):
-                if not isinstance(item, dict):
-                    raise ValueError(
-                        f'Error on key {key} item {idx}: For {self.__class__.__name__} your JSON file must be a dictionary with one task, or a list of dictionaries with one task each'
-                    )
-            return value
-        else:
-            raise ValueError(
-                f'Error on key {key}: For {self.__class__.__name__} your JSON file must be a dictionary with one task, or a list of dictionaries with one task each'
-            )
+            with open(path, 'rb') as f:
+                blob_str = f.read()
+                return load_tasks_json(blob_str, key)
+        except OSError as e:
+            raise ValueError(f'Failed to read file {path}: {str(e)}')
 
     def scan_and_create_links(self):
         return self._scan_and_create_links(LocalFilesImportStorageLink)

label_studio/io_storages/redis/models.py

@@ -3,7 +3,6 @@
 
 import json
 import logging
-from typing import Union
 
 import redis
 from django.db import models
@@ -17,6 +16,7 @@
     ImportStorageLink,
     ProjectStorageMixin,
 )
+from io_storages.utils import StorageObject, load_tasks_json
 from tasks.models import Annotation
 
 logger = logging.getLogger(__name__)
@@ -90,32 +90,12 @@ def iterkeys(self):
         for key in client.keys(path + '*'):
             yield key
 
-    def get_data(self, key) -> Union[dict, list[dict]]:
+    def get_data(self, key) -> list[StorageObject]:
         client = self.get_client()
         value_str = client.get(key)
         if not value_str:
             return []
-        try:
-            value = json.loads(value_str)
-            # NOTE: this validation did not previously exist, we were accepting any JSON values
-            if isinstance(value, dict):
-                return value
-            elif isinstance(value, list):
-                for idx, item in enumerate(value):
-                    if not isinstance(item, dict):
-                        raise ValueError(
-                            f'Error on key {key} item {idx}: For {self.__class__.__name__} your JSON file must be a dictionary with one task, or a list of dictionaries with one task each'
-                        )
-                return value
-            else:
-                raise ValueError(
-                    f'Error on key {key}: For {self.__class__.__name__} your JSON file must be a dictionary with one task, or a list of dictionaries with one task each'
-                )
-        except json.decoder.JSONDecodeError:
-            raise ValueError(
-                f"Can't import JSON-formatted tasks from {key}. If you're trying to import binary objects, "
-                f'perhaps you\'ve forgot to enable "Treat every bucket object as a source file" option?'
-            )
+        return load_tasks_json(value_str, key)
 
     def scan_and_create_links(self):
         return self._scan_and_create_links(RedisImportStorageLink)