Skip to content

fix: batch process dataset function raises out of memory with large feeds#1320

Merged
cka-y merged 10 commits into
mainfrom
feat/1297
Aug 19, 2025
Merged

fix: batch process dataset function raises out of memory with large feeds#1320
cka-y merged 10 commits into
mainfrom
feat/1297

Conversation

@cka-y

@cka-y cka-y commented Aug 18, 2025

Copy link
Copy Markdown
Contributor

Summary

Closes #1297

This PR addresses a bug related to GCP’s handling of function memory and local storage, where large file downloads consume function memory. To mitigate this, we’ve introduced an explicit in-memory storage limit and increased allocated resources to ensure more reliable and predictable behavior.

Important

This change requires manual setup across envs, as the configuration is not currently supported by Terraform or the Google API.

Key Changes

  • Batch Datasets Function Enhancements

    • Added support for limiting the feeds to process by accepting an optional feed_stable_ids list in the request body, enabling targeted processing of specific feeds.
  • Dataset Processing Refactor

    • Split dataset processing into two clearer methods:

      • process_from_producer_url: downloads directly from producer URLs.
      • process_from_bucket_latest: rebuilds datasets from the latest version in GCS.
    • This improves code clarity, maintainability, and separation of responsibilities.

  • Resource Allocation Updates

    • Increased memory allocation from 2Gi → 8Gi and CPU from 1 → 2 for the batch dataset processing function.
    • Introduced an in-memory storage limit of 7Gi to ensure large file handling does not cause uncontrolled retries.
    • Verified locally that memory usage stays well below this threshold (typically <200 MiB), so the new limit provides sufficient headroom while still enabling catchable errors when exceeded.
  • File System Adjustments

    • Updated temporary file paths to use /in-memory instead of /tmp to better align with GCP function storage semantics.

Please make sure these boxes are checked before submitting your pull request - thanks!

  • Run the unit tests with ./scripts/api-tests.sh to make sure you didn't break anything
  • Add or update any needed documentation to the repo
  • Format the title like "feat: [new feature short description]". Title must follow the Conventional Commit Specification(https://www.conventionalcommits.org/en/v1.0.0/).
  • Linked all relevant issues
  • Include screenshot(s) showing how this pull request works and fixes the issue(s)

@cka-y cka-y requested a review from Copilot August 18, 2025 19:25

This comment was marked as off-topic.

@cka-y cka-y requested a review from Copilot August 18, 2025 19:41
@cka-y cka-y marked this pull request as ready for review August 18, 2025 19:41
@cka-y cka-y changed the title Feat/1297 fix: batch process dataset function raises out of memory with large feeds Aug 18, 2025

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull Request Overview

This PR implements enhancements to batch dataset processing to address memory issues and improve reliability. It introduces targeted feed processing, refactors dataset handling into cleaner separation of concerns, and increases resource allocation to handle large file downloads more effectively.

Key changes:

  • Added support for limiting feeds to process via optional feed_stable_ids parameter in batch datasets function
  • Split dataset processing into two distinct methods: process_from_producer_url for direct downloads and process_from_bucket_latest for rebuilding from GCS
  • Increased function memory allocation from 2Gi to 8Gi and CPU from 1 to 2 cores with 7Gi in-memory storage limit

Reviewed Changes

Copilot reviewed 12 out of 12 changed files in this pull request and generated 4 comments.

Show a summary per file
File Description
infra/functions-python/main.tf Updated environment variables for Pub/Sub topic names
functions-python/batch_process_dataset/function_config.json Increased memory and CPU allocation for better performance
functions-python/batch_process_dataset/src/main.py Added new processing method and refactored file upload logic
functions-python/batch_datasets/src/main.py Added optional feed filtering capability
functions-python/tasks_executor/src/tasks/dataset_files/rebuild_missing_dataset_files.py Refactored to use Pub/Sub messages instead of direct processing
functions-python/helpers/utils.py Added GCS download utility function

Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.

Comment thread functions-python/big_query_ingestion/tests/test_common.py Outdated
Comment thread functions-python/batch_process_dataset/src/main.py
Comment thread functions-python/batch_process_dataset/src/main.py
Comment thread functions-python/batch_datasets/src/main.py

@davidgamez davidgamez left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM!

Comment on lines +93 to 97
if feed_stable_ids:
# If feed_stable_ids are provided, filter the query by stable IDs
query = query.filter(Gtfsfeed.stable_id.in_(feed_stable_ids))
# Limit the query to 10 feeds (or FEEDS_LIMIT param) for testing purposes and lower environments
if os.getenv("ENVIRONMENT", "").lower() in ("dev", "test", "qa"):

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nitpick: I would have the limit no applied when the feed_stable_ids are passed.

@cka-y cka-y merged commit ef2d1f4 into main Aug 19, 2025
3 checks passed
@cka-y cka-y deleted the feat/1297 branch August 19, 2025 17:49
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Bash process dataset function raises out of memory with large feeds

3 participants