From 0db2ad213293f81a3057267e53509313128698c5 Mon Sep 17 00:00:00 2001 From: rkuo-danswer Date: Thu, 1 May 2025 10:47:26 -0700 Subject: [PATCH] memory optimize task generation for connector deletion (#4645) * memory optimize task generation for connector deletion * test * fix up integration test docker file * more no-cache --------- Co-authored-by: Richard Kuo (Onyx) --- .../actions/custom-build-and-push/action.yml | 7 +++++++ ...er-build-push-backend-container-on-tag.yml | 1 + ...-build-push-cloud-web-container-on-tag.yml | 4 ++-- ...ild-push-model-server-container-on-tag.yml | 2 ++ ...docker-build-push-web-container-on-tag.yml | 1 + .github/workflows/pr-integration-tests.yml | 9 ++++++--- .../workflows/pr-mit-integration-tests.yml | 9 ++++++--- backend/onyx/db/document.py | 15 ++++++++++++++ backend/onyx/redis/redis_connector_delete.py | 20 +++++++++---------- backend/tests/integration/Dockerfile | 14 +++++++------ 10 files changed, 58 insertions(+), 24 deletions(-) diff --git a/.github/actions/custom-build-and-push/action.yml b/.github/actions/custom-build-and-push/action.yml index fbee0554995c..2849d4a150c0 100644 --- a/.github/actions/custom-build-and-push/action.yml +++ b/.github/actions/custom-build-and-push/action.yml @@ -25,6 +25,10 @@ inputs: tags: description: 'Image tags' required: true + no-cache: + description: 'Read from cache' + required: false + default: 'false' cache-from: description: 'Cache sources' required: false @@ -55,6 +59,7 @@ runs: push: ${{ inputs.push }} load: ${{ inputs.load }} tags: ${{ inputs.tags }} + no-cache: ${{ inputs.no-cache }} cache-from: ${{ inputs.cache-from }} cache-to: ${{ inputs.cache-to }} @@ -77,6 +82,7 @@ runs: push: ${{ inputs.push }} load: ${{ inputs.load }} tags: ${{ inputs.tags }} + no-cache: ${{ inputs.no-cache }} cache-from: ${{ inputs.cache-from }} cache-to: ${{ inputs.cache-to }} @@ -99,6 +105,7 @@ runs: push: ${{ inputs.push }} load: ${{ inputs.load }} tags: ${{ inputs.tags }} + no-cache: ${{ inputs.no-cache }} cache-from: ${{ inputs.cache-from }} cache-to: ${{ inputs.cache-to }} diff --git a/.github/workflows/docker-build-push-backend-container-on-tag.yml b/.github/workflows/docker-build-push-backend-container-on-tag.yml index 162c02213083..9ee9d3152fca 100644 --- a/.github/workflows/docker-build-push-backend-container-on-tag.yml +++ b/.github/workflows/docker-build-push-backend-container-on-tag.yml @@ -45,6 +45,7 @@ jobs: ${{ env.LATEST_TAG == 'true' && format('{0}:latest', env.REGISTRY_IMAGE) || '' }} build-args: | ONYX_VERSION=${{ github.ref_name }} + no-cache: true # trivy has their own rate limiting issues causing this action to flake # we worked around it by hardcoding to different db repos in env diff --git a/.github/workflows/docker-build-push-cloud-web-container-on-tag.yml b/.github/workflows/docker-build-push-cloud-web-container-on-tag.yml index a6ff2f51c5af..65e72c4aa7c1 100644 --- a/.github/workflows/docker-build-push-cloud-web-container-on-tag.yml +++ b/.github/workflows/docker-build-push-cloud-web-container-on-tag.yml @@ -70,10 +70,10 @@ jobs: NEXT_PUBLIC_FORGOT_PASSWORD_ENABLED=true NEXT_PUBLIC_INCLUDE_ERROR_POPUP_SUPPORT_LINK=true NODE_OPTIONS=--max-old-space-size=8192 - # needed due to weird interactions with the builds for different platforms - no-cache: true labels: ${{ steps.meta.outputs.labels }} outputs: type=image,name=${{ env.REGISTRY_IMAGE }},push-by-digest=true,name-canonical=true,push=true + # no-cache needed due to weird interactions with the builds for different platforms + no-cache: true - name: Export digest run: | diff --git a/.github/workflows/docker-build-push-model-server-container-on-tag.yml b/.github/workflows/docker-build-push-model-server-container-on-tag.yml index 2c152f46658f..6f663192e388 100644 --- a/.github/workflows/docker-build-push-model-server-container-on-tag.yml +++ b/.github/workflows/docker-build-push-model-server-container-on-tag.yml @@ -86,6 +86,7 @@ jobs: DANSWER_VERSION=${{ github.ref_name }} outputs: type=registry provenance: false + no-cache: true build-arm64: needs: [check_model_server_changes] @@ -127,6 +128,7 @@ jobs: DANSWER_VERSION=${{ github.ref_name }} outputs: type=registry provenance: false + no-cache: true merge-and-scan: needs: [build-amd64, build-arm64, check_model_server_changes] diff --git a/.github/workflows/docker-build-push-web-container-on-tag.yml b/.github/workflows/docker-build-push-web-container-on-tag.yml index d6cbc4c7071f..605c54521c0a 100644 --- a/.github/workflows/docker-build-push-web-container-on-tag.yml +++ b/.github/workflows/docker-build-push-web-container-on-tag.yml @@ -66,6 +66,7 @@ jobs: no-cache: true labels: ${{ steps.meta.outputs.labels }} outputs: type=image,name=${{ env.REGISTRY_IMAGE }},push-by-digest=true,name-canonical=true,push=true + no-cache: true - name: Export digest run: | diff --git a/.github/workflows/pr-integration-tests.yml b/.github/workflows/pr-integration-tests.yml index 593a9adfd24b..d95c4dd81845 100644 --- a/.github/workflows/pr-integration-tests.yml +++ b/.github/workflows/pr-integration-tests.yml @@ -61,7 +61,8 @@ jobs: tags: onyxdotapp/onyx-backend:test push: false load: true - cache-from: type=s3,prefix=cache/${{ github.repository }}/integration-tests/backend/,region=${{ env.RUNS_ON_AWS_REGION }},bucket=${{ env.RUNS_ON_S3_BUCKET_CACHE }} + no-cache: true +# cache-from: type=s3,prefix=cache/${{ github.repository }}/integration-tests/backend/,region=${{ env.RUNS_ON_AWS_REGION }},bucket=${{ env.RUNS_ON_S3_BUCKET_CACHE }} cache-to: type=s3,prefix=cache/${{ github.repository }}/integration-tests/backend/,region=${{ env.RUNS_ON_AWS_REGION }},bucket=${{ env.RUNS_ON_S3_BUCKET_CACHE }},mode=max - name: Build Model Server Docker image @@ -73,7 +74,8 @@ jobs: tags: onyxdotapp/onyx-model-server:test push: false load: true - cache-from: type=s3,prefix=cache/${{ github.repository }}/integration-tests/model-server/,region=${{ env.RUNS_ON_AWS_REGION }},bucket=${{ env.RUNS_ON_S3_BUCKET_CACHE }} + no-cache: true +# cache-from: type=s3,prefix=cache/${{ github.repository }}/integration-tests/model-server/,region=${{ env.RUNS_ON_AWS_REGION }},bucket=${{ env.RUNS_ON_S3_BUCKET_CACHE }} cache-to: type=s3,prefix=cache/${{ github.repository }}/integration-tests/model-server/,region=${{ env.RUNS_ON_AWS_REGION }},bucket=${{ env.RUNS_ON_S3_BUCKET_CACHE }},mode=max - name: Build integration test Docker image @@ -85,7 +87,8 @@ jobs: tags: onyxdotapp/onyx-integration:test push: false load: true - cache-from: type=s3,prefix=cache/${{ github.repository }}/integration-tests/integration/,region=${{ env.RUNS_ON_AWS_REGION }},bucket=${{ env.RUNS_ON_S3_BUCKET_CACHE }} + no-cache: true +# cache-from: type=s3,prefix=cache/${{ github.repository }}/integration-tests/integration/,region=${{ env.RUNS_ON_AWS_REGION }},bucket=${{ env.RUNS_ON_S3_BUCKET_CACHE }} cache-to: type=s3,prefix=cache/${{ github.repository }}/integration-tests/integration/,region=${{ env.RUNS_ON_AWS_REGION }},bucket=${{ env.RUNS_ON_S3_BUCKET_CACHE }},mode=max # Start containers for multi-tenant tests diff --git a/.github/workflows/pr-mit-integration-tests.yml b/.github/workflows/pr-mit-integration-tests.yml index 2c65c73e7397..2f4ede4a697f 100644 --- a/.github/workflows/pr-mit-integration-tests.yml +++ b/.github/workflows/pr-mit-integration-tests.yml @@ -61,7 +61,8 @@ jobs: tags: onyxdotapp/onyx-backend:test push: false load: true - cache-from: type=s3,prefix=cache/${{ github.repository }}/integration-tests/backend/,region=${{ env.RUNS_ON_AWS_REGION }},bucket=${{ env.RUNS_ON_S3_BUCKET_CACHE }} + no-cache: true +# cache-from: type=s3,prefix=cache/${{ github.repository }}/integration-tests/backend/,region=${{ env.RUNS_ON_AWS_REGION }},bucket=${{ env.RUNS_ON_S3_BUCKET_CACHE }} cache-to: type=s3,prefix=cache/${{ github.repository }}/integration-tests/backend/,region=${{ env.RUNS_ON_AWS_REGION }},bucket=${{ env.RUNS_ON_S3_BUCKET_CACHE }},mode=max - name: Build Model Server Docker image @@ -73,7 +74,8 @@ jobs: tags: onyxdotapp/onyx-model-server:test push: false load: true - cache-from: type=s3,prefix=cache/${{ github.repository }}/integration-tests/model-server/,region=${{ env.RUNS_ON_AWS_REGION }},bucket=${{ env.RUNS_ON_S3_BUCKET_CACHE }} + no-cache: true +# cache-from: type=s3,prefix=cache/${{ github.repository }}/integration-tests/model-server/,region=${{ env.RUNS_ON_AWS_REGION }},bucket=${{ env.RUNS_ON_S3_BUCKET_CACHE }} cache-to: type=s3,prefix=cache/${{ github.repository }}/integration-tests/model-server/,region=${{ env.RUNS_ON_AWS_REGION }},bucket=${{ env.RUNS_ON_S3_BUCKET_CACHE }},mode=max - name: Build integration test Docker image @@ -85,7 +87,8 @@ jobs: tags: onyxdotapp/onyx-integration:test push: false load: true - cache-from: type=s3,prefix=cache/${{ github.repository }}/integration-tests/integration/,region=${{ env.RUNS_ON_AWS_REGION }},bucket=${{ env.RUNS_ON_S3_BUCKET_CACHE }} + no-cache: true +# cache-from: type=s3,prefix=cache/${{ github.repository }}/integration-tests/integration/,region=${{ env.RUNS_ON_AWS_REGION }},bucket=${{ env.RUNS_ON_S3_BUCKET_CACHE }} cache-to: type=s3,prefix=cache/${{ github.repository }}/integration-tests/integration/,region=${{ env.RUNS_ON_AWS_REGION }},bucket=${{ env.RUNS_ON_S3_BUCKET_CACHE }},mode=max # NOTE: Use pre-ping/null pool to reduce flakiness due to dropped connections diff --git a/backend/onyx/db/document.py b/backend/onyx/db/document.py index 0c0b225372c2..b40ea47970f0 100644 --- a/backend/onyx/db/document.py +++ b/backend/onyx/db/document.py @@ -141,6 +141,21 @@ def get_all_documents_needing_vespa_sync_for_cc_pair( return list(db_session.scalars(stmt).all()) +def construct_document_id_select_for_connector_credential_pair( + connector_id: int, credential_id: int | None = None +) -> Select: + initial_doc_ids_stmt = select(DocumentByConnectorCredentialPair.id).where( + and_( + DocumentByConnectorCredentialPair.connector_id == connector_id, + DocumentByConnectorCredentialPair.credential_id == credential_id, + ) + ) + stmt = ( + select(DbDocument.id).where(DbDocument.id.in_(initial_doc_ids_stmt)).distinct() + ) + return stmt + + def construct_document_select_for_connector_credential_pair( connector_id: int, credential_id: int | None = None ) -> Select: diff --git a/backend/onyx/redis/redis_connector_delete.py b/backend/onyx/redis/redis_connector_delete.py index 98a36fe78823..70be982192f7 100644 --- a/backend/onyx/redis/redis_connector_delete.py +++ b/backend/onyx/redis/redis_connector_delete.py @@ -16,8 +16,7 @@ from onyx.configs.constants import OnyxCeleryQueues from onyx.configs.constants import OnyxCeleryTask from onyx.configs.constants import OnyxRedisConstants from onyx.db.connector_credential_pair import get_connector_credential_pair_from_id -from onyx.db.document import construct_document_select_for_connector_credential_pair -from onyx.db.models import Document as DbDocument +from onyx.db.document import construct_document_id_select_for_connector_credential_pair class RedisConnectorDeletePayload(BaseModel): @@ -116,7 +115,6 @@ class RedisConnectorDelete: Otherwise, returns an int with the number of generated tasks.""" last_lock_time = time.monotonic() - async_results = [] cc_pair = get_connector_credential_pair_from_id( db_session=db_session, cc_pair_id=int(self.id), @@ -124,11 +122,13 @@ class RedisConnectorDelete: if not cc_pair: return None - stmt = construct_document_select_for_connector_credential_pair( + num_tasks_sent = 0 + + stmt = construct_document_id_select_for_connector_credential_pair( cc_pair.connector_id, cc_pair.credential_id ) - for doc_temp in db_session.scalars(stmt).yield_per(DB_YIELD_PER_DEFAULT): - doc: DbDocument = doc_temp + for doc_id in db_session.scalars(stmt).yield_per(DB_YIELD_PER_DEFAULT): + doc_id = cast(str, doc_id) current_time = time.monotonic() if current_time - last_lock_time >= ( CELERY_VESPA_SYNC_BEAT_LOCK_TIMEOUT / 4 @@ -143,10 +143,10 @@ class RedisConnectorDelete: self.redis.sadd(self.taskset_key, custom_task_id) # Priority on sync's triggered by new indexing should be medium - result = celery_app.send_task( + celery_app.send_task( OnyxCeleryTask.DOCUMENT_BY_CC_PAIR_CLEANUP_TASK, kwargs=dict( - document_id=doc.id, + document_id=doc_id, connector_id=cc_pair.connector_id, credential_id=cc_pair.credential_id, tenant_id=self.tenant_id, @@ -157,9 +157,9 @@ class RedisConnectorDelete: ignore_result=True, ) - async_results.append(result) + num_tasks_sent += 1 - return len(async_results) + return num_tasks_sent def reset(self) -> None: self.redis.srem(OnyxRedisConstants.ACTIVE_FENCES, self.fence_key) diff --git a/backend/tests/integration/Dockerfile b/backend/tests/integration/Dockerfile index 9164d1a8ab9e..af9fbcd96940 100644 --- a/backend/tests/integration/Dockerfile +++ b/backend/tests/integration/Dockerfile @@ -15,14 +15,16 @@ RUN apt-get update && \ curl \ zip \ ca-certificates \ - libgnutls30=3.7.9-2+deb12u3 \ - libblkid1=2.38.1-5+deb12u1 \ - libmount1=2.38.1-5+deb12u1 \ - libsmartcols1=2.38.1-5+deb12u1 \ - libuuid1=2.38.1-5+deb12u1 \ + libgnutls30 \ + libblkid1 \ + libmount1 \ + libsmartcols1 \ + libuuid1 \ libxmlsec1-dev \ pkg-config \ - gcc && \ + gcc \ + nano \ + vim && \ rm -rf /var/lib/apt/lists/* && \ apt-get clean