Walkthrough

The changes in the JobAnnotation class within the cvat/apps/dataset_manager/task.py file introduce a new variable, is_custom_queryset, to manage job retrieval based on the type of queryset provided. The control flow for job fetching is modified to handle different conditions, enhancing error handling with a try-except block for StopIteration. Additionally, the import_annotations method is updated to initialize job data with a new JobData instance, ensuring proper job annotation creation. The overall structure remains intact while improving robustness and clarity.

Changes

Poem

In the garden where jobs do play,
A custom path now lights the way.
With queries clear and errors caught,
Annotations bloom, as they ought.
So hop along, let changes ring,
For every job, a joyful spring! 🐇✨

Thank you for using CodeRabbit. We offer it for free to the OSS community and would appreciate your support in helping us grow. If you find it useful, would you consider giving us a shout-out on your favorite social media?

Codecov Report

Attention: Patch coverage is 93.75000% with 2 lines in your changes missing coverage. Please review.

Project coverage is 74.05%. Comparing base (7d0205b) to head (021de6e).
Report is 1 commits behind head on develop.

@@             Coverage Diff             @@
##           develop    #8676      +/-   ##
===========================================
- Coverage    74.05%   74.05%   -0.01%     
===========================================
  Files          409      409              
  Lines        43783    43792       +9     
  Branches      3984     3984              
===========================================
+ Hits         32425    32431       +6     
- Misses       11358    11361       +3     

@Eldies, Could you please provide the difference in memory usage and number of db queries (before/after the patch)?

Measuring consumed memory with this memray.patch
With the task and annotations attached to the issue (100 jobs, 4098 images)

1. run locally (in docker)
2. import task from backup
3. apply patch
4. python manage.py shell

# preparations
In [1]: import cvat
In [2]: from django.db import reset_queries, connections
In [3]: cvat.apps.dataset_manager.task.delete_task_data(<task id>)
In [4]: reset_queries()
# run import (and measure memory consumption)
In [5]: cvat.apps.dataset_manager.task.import_task_annotations(<path to annotations archive>, <task id>, 'Segmentation mask 1.1', True)
# count queries
In [6]: sum([len(conn.queries) for conn in connections.all()])

On develop:
Peak memory usage: 2.5 GiB
Number of queries: 12287

With this PR:
Peak memory usage: 2.0 GiB
Number of queries: 11287

Quality Gate passed

Issues
0 New issues
0 Accepted issues

Measures
0 Security Hotspots
0.0% Coverage on New Code
0.0% Duplication on New Code

See analysis details on SonarQube Cloud

I returned select_for_update because apparently it is needed on export. But I changed the argument name to more descriptive and removed unnecessary differences in the conditional logic for locking

do you have some breakdown on what is being queried? It's approximately 100 requests per job or ~3 per image.

In: Counter([
    m.groups()
    for query in connection.queries
    for m in [re.search('FROM "([^"]+)"', query['sql'])]
    if m and len(m.groups()) >= 1
])
Out: 
Counter({('engine_label',): 2002,
         ('engine_job',): 1526,
         ('engine_data',): 1001,
         ('engine_image',): 1001,
         ('engine_skeleton',): 1001,
         ('engine_attributespec',): 1001,
         ('auth_user',): 1001,
         ('organizations_organization',): 1000,
         ('webhooks_webhook',): 216,
         ('engine_labeledimage',): 200,
         ('engine_labeledshape',): 200,
         ('engine_labeledtrack',): 200,
         ('engine_task',): 109,
         ('engine_trackedshapeattributeval',): 100,
         ('engine_trackedshape',): 100,
         ('engine_labeledtrackattributeval',): 100,
         ('engine_labeledshapeattributeval',): 100,
         ('engine_labeledimageattributeval',): 100,
         ('engine_validationlayout',): 1,
         ('engine_video',): 1,
         ('engine_segment',): 1})

@Eldies, please try to enable the silk profiler, you'll get a picture like this:

Note that you might need to use @silk_profile or call the import or export function directly from some API endpoint to collect the metrics, instead of calling it in an rq job.

     ('auth_user',): 1001,
     ('organizations_organization',): 1000,

These 2 might be optimized by #8275 or by a similar approach with changing subj.obj.id to subj.obj_id.

('engine_label',): 2002,
('engine_job',): 1526,
('engine_data',): 1001,
('engine_image',): 1001,
('engine_skeleton',): 1001,
('engine_task',): 109,

These are the ones I'd look more closely, it feels like some prefetching could be added.

For every case when JobAnnotation is called from TaskAnnotation, passed a db_task so that all JobAnnotations could use the same, shared db_task.

Added prefetching to TaskAnnotation. But now TaskAnnotation.add_prefetch_info kinda duplicates JobAnnotation.add_prefetch_info. Something like this could remove duplication:

-        return queryset.select_related(
-            'segment',
-            'segment__task',
-        ).prefetch_related(
-            'segment__task__project',
-            'segment__task__owner',
-             ....
-       )
+        return queryset.select_related(
+            'segment',
+        ).prefetch_related(
+            Prefetch('segment__task', queryset=TaskAnnotation.add_prefetch_info(models.Task.objects))
+        )

but it will add one request, so I am not sure is it a good idea or not.

It is 5396 db requests now,

In: Counter([
    m.groups()
    for query in connection.queries
    for m in [re.search('FROM "([^"]+)"', query['sql'])]
    if m and len(m.groups()) >= 1
])
Out: 
Counter({('engine_job',): 1419,
         ('webhooks_webhook',): 216,
         ('engine_labeledimage',): 200,
         ('engine_labeledshape',): 200,
         ('engine_labeledtrack',): 200,
         ('engine_task',): 109,
         ('engine_trackedshapeattributeval',): 100,
         ('engine_trackedshape',): 100,
         ('engine_labeledtrackattributeval',): 100,
         ('engine_labeledshapeattributeval',): 100,
         ('engine_labeledimageattributeval',): 100,
         ('engine_label',): 2,
         ('auth_user',): 1,
         ('engine_data',): 1,
         ('engine_image',): 1,
         ('engine_skeleton',): 1,
         ('engine_attributespec',): 1,
         ('engine_validationlayout',): 1,
         ('engine_segment',): 1,
         ('organizations_organization',): 1})

It is 5396 db requests now

Looks great! However, I still can see some suspicious numbers in these lines:

('engine_task',): 109,
('engine_job',): 1419,

From the regex, I can guess tasks and jobs can be mixed into some other requests, but still, the numbers are big. Could you check it please?

And here:
('webhooks_webhook',): 216,

Probably, it should be like 100 (job updated) + 1 (task updated) + 1 (project updated), if I'm not missing something.

But now TaskAnnotation.add_prefetch_info kinda duplicates JobAnnotation.add_prefetch_info.

Actually, maybe we could call some functions from JobAnnotation in TaskAnnotation prefetch? I think we know there that we will need jobs at some point later.

but it will add one request, so I am not sure is it a good idea or not.

The proposed code snippet doesn't look right. We don't need all the task prefetch, if we're working with just 1 job. But we know that we're going to work with jobs, if we're working with TaskAnnotations.

Do you have updated memory metrics for the import use case after optimizations?

And here: ('webhooks_webhook',): 216, Probably, it should be like 100 (job updated)

On every job update there is also a task update - in JobAnnotation._set_updated_date both the job and its task are touched.
There are 100+ job updates - therefore, 100+ task updates.
Not exactly 100, because there are a lot of annotations and TaskAnnotation._patch_data is called several times, and some jobs are updated twice.
So, 200+ webhook reads.

('engine_task',): 109,

On every job update, the it's task is touched and when the task is touched, it is read from db - pre_save_resource_event reads old instance.

('engine_job',): 1419,

On each of 100+ job updates:
1 read in pre_save_resource_event
1 read in __save_job_handler
1 read in __save_job__update_quality_metrics

200 reads because in TaskData.meta_for_task jobs were not ordered when prefetced for segments, and each db_segment.job_set.first() caused a db query. Fixed it.

All the other reads are in JobAnnotation initializations.
Redesigned it - now TaskAnnotations loads all the jobs it needs along with all needed data once, and then passes jobs to JobAnnotation

Now there are 4401 db queries,

Counter({('engine_job',): 318,
         ('webhooks_webhook',): 216,
         ('engine_labeledimage',): 200,
         ('engine_labeledshape',): 200,
         ('engine_labeledtrack',): 200,
         ('engine_task',): 109,
         ('engine_trackedshapeattributeval',): 100,
         ('engine_trackedshape',): 100,
         ('engine_labeledtrackattributeval',): 100,
         ('engine_labeledshapeattributeval',): 100,
         ('engine_labeledimageattributeval',): 100,
         ('organizations_organization',): 100,
         ('engine_label',): 4,
         ('engine_data',): 2,
         ('engine_skeleton',): 2,
         ('engine_attributespec',): 2,
         ('auth_user',): 2,
         ('engine_image',): 1,
         ('engine_validationlayout',): 1,
         ('engine_video',): 1,
         ('engine_segment',): 1})

For some reason when I try to use silk, the ui does not show me my tasks. I'll try to investigate it

Ok, it feels like a good point to stop optimizing tasks for now. Do you have updated memory measurements for the import use case?

Memory consumption is the same as earlier, ~0.5 Gb less then on develop branch

For some reason when I try to use silk, the ui does not show me my tasks. I'll try to investigate it

It's only configured for working with development setup (VS Code debug tasks), so maybe this is the reason. Make sure you're connecting to the right DB and server. You'll need to start docker compose -f docker-compose.yml -f docker-compose.dev.yml up -d --build cvat_opa cvat_db cvat_redis_inmem cvat_redis_ondisk cvat_server cvat_vector, then launch the server debug task in VS code, launch the UI manually with (cd cvat-ui && yarn run start), then go to localhost:3000 in the browser.