Method get_task_batch_size

label_studio/projects/models.py:1264–1310

Calculate optimal batch size based on task data size and annotation result size

(self)

            return row[1]

    def get_task_batch_size(self):
        """Calculate optimal batch size based on task data size and annotation result size"""
        # For SQLite, use default MAX_TASK_BATCH_SIZE
        if settings.DJANGO_DB == settings.DJANGO_DB_SQLITE:
            return settings.MAX_TASK_BATCH_SIZE

        # Get maximum task data size using the optimized index
        max_task_size = 0
        with connection.cursor() as cursor:
            cursor.execute(
                """
                SELECT id,
                       octet_length(data::text) AS bytes
                FROM   task
                WHERE  project_id = %s
                ORDER  BY octet_length(data::text) DESC
                LIMIT  1
            """,
                [self.id],
            )

            row = cursor.fetchone()
            if row and row[1]:
                max_task_size = row[1]

        # Get maximum annotation result size using the new optimized index
        max_annotation_size = self.get_max_annotation_result_size()

        # Use the larger of the two sizes for batch calculation
        max_data_size = max(max_task_size, max_annotation_size)

        if max_data_size == 0:
            return settings.MAX_TASK_BATCH_SIZE

        batch_size = settings.TASK_DATA_PER_BATCH // max_data_size

        if batch_size > settings.MAX_TASK_BATCH_SIZE:
            batch_size = settings.MAX_TASK_BATCH_SIZE
        elif batch_size < 1:
            batch_size = 1

        logger.info(
            f'Project {self.id}: max task size {max_task_size} bytes, '
            f'max annotation size {max_annotation_size} bytes, '
            f'calculated batch size {batch_size}'
        )
        return batch_size

    def __str__(self):
        return f'{self.title} (id={self.id})' or _('Business number %d') % self.pk