diff --git a/.gitignore b/.gitignore index 01d1986a1..d8e2547b0 100644 --- a/.gitignore +++ b/.gitignore @@ -87,5 +87,9 @@ functions-python/**/*.csv # Project files *.code-workspace +# Local build output, including the default destination of +# scripts/parquet-generate-local.sh +.dist/ + # Ignore OpenApi local backup files *.yaml.bak \ No newline at end of file diff --git a/README.md b/README.md index cf74ea15e..00799efdf 100644 --- a/README.md +++ b/README.md @@ -181,6 +181,20 @@ Note: the tests rely on having an empty local test DB instance. If you have data ./scripts/docker-localdb-rebuild-data.sh --use-test-db ``` +### Browsing a feed locally (Parquet) + +To convert a GTFS feed to Parquet on this machine - no bucket, database or GCP +credentials - and optionally serve it to a browser: + +```bash +scripts/parquet-generate-local.sh mdb-1210 --serve +``` + +The script accepts a feed or dataset stable id, a `.zip`, or an unpacked folder. See +[functions-python/parquet_builder](functions-python/parquet_builder/README.md) for the +full local testing guide, including running it against the Operations API and the +operations web app. + ## Running with Docker diff --git a/api/src/shared/common/gcp_memory_utils.py b/api/src/shared/common/gcp_memory_utils.py index 10019ed17..25e93f8bf 100644 --- a/api/src/shared/common/gcp_memory_utils.py +++ b/api/src/shared/common/gcp_memory_utils.py @@ -3,10 +3,31 @@ import resource import shutil import sys +from dataclasses import dataclass +from typing import Optional MB_MULTIPLIER = 1024**2 +@dataclass(frozen=True) +class MemoryBudget: + """What this container was given, as opposed to what it went on to use. + + These three numbers are computed on the way to setting RLIMIT_AS and were previously + only logged. A process that records how much memory a job needed cannot say whether + that was comfortable without also recording what it had, and reconstructing the + budget afterwards from the deployment means trusting that the deployment has not + moved since. + + Any field may be None where the figure could not be read, which is normal off + Cloud Run. + """ + + cgroup_limit_bytes: Optional[int] = None + volume_bytes: Optional[int] = None + rlimit_as_bytes: Optional[int] = None + + def is_filesystem_tmpfs(mount_point): """ Check if the given mount_point is a tmpfs filesystem in /proc/mounts. @@ -139,12 +160,20 @@ def limit_gcp_memory(mount_point): Environment Variables: MEMORY_MARGIN_MB: Safety margin in megabytes (default: 200) + + Returns: + A MemoryBudget describing what was found and what was set. Fields are None where + a figure could not be read or no limit was applied; callers that only want the + side effect can ignore it. """ + cgroup_limit_bytes = get_memory_limit_cgroup_bytes() + volume_bytes = get_tmpfs_size_bytes(mount_point) + # Calculate available memory: cgroup limit - tmpfs size available_memory_bytes = get_available_process_memory_bytes(mount_point) if not available_memory_bytes or available_memory_bytes <= 0: logging.info("Could not find the total memory of the process. Memory limit not set.") - return + return MemoryBudget(cgroup_limit_bytes=cgroup_limit_bytes, volume_bytes=volume_bytes) # Parse and validate the memory margin memory_margin_mb = 200 @@ -177,9 +206,14 @@ def limit_gcp_memory(mount_point): "Computed RLIMIT_AS <= 0 (%.2f MiB). Skipping setrlimit.", mem_limit / MB_MULTIPLIER, ) - return + return MemoryBudget(cgroup_limit_bytes=cgroup_limit_bytes, volume_bytes=volume_bytes) # Set RLIMIT_AS (address space limit) to prevent OOM kills # When this limit is exceeded, Python will raise MemoryError instead of being killed resource.setrlimit(resource.RLIMIT_AS, (mem_limit, mem_limit)) logging.info("RLIMIT_AS set to %.2f MiB", mem_limit / MB_MULTIPLIER) + return MemoryBudget( + cgroup_limit_bytes=cgroup_limit_bytes, + volume_bytes=volume_bytes, + rlimit_as_bytes=mem_limit, + ) diff --git a/api/src/shared/common/gcp_utils.py b/api/src/shared/common/gcp_utils.py index faeb98166..d0efdef2d 100644 --- a/api/src/shared/common/gcp_utils.py +++ b/api/src/shared/common/gcp_utils.py @@ -258,8 +258,13 @@ def create_http_task_with_name( task_time, http_method: any, # tasks_v2.HttpMethod timeout_s: int = 1800, # 30 minutes + raise_on_error: bool = False, ): - """Creates a GCP Cloud Task.""" + """Creates a GCP Cloud Task. + + `raise_on_error` is for callers that record state on the strength of the enqueue. + Off by default: the original callers are fire-and-forget. + """ from google.cloud import tasks_v2 from google.protobuf import duration_pb2 @@ -290,3 +295,5 @@ def create_http_task_with_name( logging.info("Task already exists for %s, skipping.", task_name) else: logging.error("Error creating task: %s", e) + if raise_on_error: + raise diff --git a/docs/OperationsAPI.yaml b/docs/OperationsAPI.yaml index ffd2cdf10..b1fb6f501 100644 --- a/docs/OperationsAPI.yaml +++ b/docs/OperationsAPI.yaml @@ -24,6 +24,9 @@ tags: - name: "early-access" description: >- Early access programs (e.g. a Summit or GBFS validator beta) grant feature flags in bulk, including to members who have never signed in. This is the only enrollment path — there is no self-service join endpoint anywhere in the public API. + - name: "parquet" + description: >- + On-demand Parquet renderings of GTFS datasets. A dataset is converted once into one Parquet file per GTFS table, published publicly with a manifest describing the set, and served over HTTP range requests so a client can query it in place without downloading or unpacking the feed. Conversion is asynchronous: ask for the state, start one if it is absent, and poll until it is ready. paths: /v1/operations/feeds: get: @@ -940,6 +943,153 @@ paths: description: Unauthorized. "404": description: Program not found. + /v1/operations/gtfs_feeds/{id}/execution_attempts: + parameters: + - $ref: "#/components/parameters/feed_id_path_param" + get: + summary: What background tasks have tried to do with this feed, and how it went + description: >- + One row per attempt a background task has made against any of this feed's datasets, newest first. Each row records which worker ran it, why that worker was chosen, what it cost in memory and time, and - when it failed - what ran out and whether the feed was moved to a larger worker as a result. + + + This is the audit trail behind automatic worker sizing. A build that exhausts its worker's memory or disk is retried one size up and the feed is given a size override, so the next run starts there; these rows are what show that happening, and why a feed is on a bigger worker than its size alone would suggest. + + + The filters combine: `failure_kind=resource_memory` with a date range answers "did this feed start running out of memory after we changed something", which is the question that usually prompts opening this page. + operationId: getGtfsFeedExecutionAttempts + tags: + - "operations" + security: + - Authentication: [] + parameters: + - $ref: "#/components/parameters/execution_attempt_task_name" + - $ref: "#/components/parameters/execution_attempt_status" + - $ref: "#/components/parameters/execution_attempt_failure_kind" + - $ref: "#/components/parameters/execution_attempt_from" + - $ref: "#/components/parameters/execution_attempt_to" + - $ref: "#/components/parameters/limit_query_param_availability_endpoint" + - $ref: "#/components/parameters/offset" + responses: + 200: + description: Attempt history for the feed, newest first. + content: + application/json: + schema: + $ref: "#/components/schemas/ExecutionAttemptsResponse" + 404: + description: GTFS feed not found. + 500: + description: Internal server error. + /v1/operations/gtfs_feeds/{id}/parquet: + parameters: + - $ref: "#/components/parameters/feed_id_path_param" + get: + summary: Where the Parquet for a GTFS feed's latest dataset stands + description: >- + Whether this feed's latest dataset has been rendered as Parquet - one file per GTFS table, every column typed as text - and if it is being prepared, how far along it is. A client can then query the feed in place over HTTP range requests instead of downloading and unpacking it. Read-only: this never starts a conversion, so it is safe to poll - the viewer polls it roughly twice a second while one runs. Use POST to begin one. + operationId: getGtfsFeedParquet + tags: + - "parquet" + security: + - Authentication: [] + responses: + "200": + description: >- + The dataset's current state. Returned for every state including `absent`, which means no conversion has been requested yet rather than that anything is missing. + content: + application/json: + schema: + $ref: "#/components/schemas/ParquetDatasetState" + "404": + description: >- + GTFS feed not found, or it has no dataset yet. Note this does **not** cover a dataset whose Parquet has never been generated - that is `absent` with a 200. + "500": + description: Internal server error. + post: + summary: Start preparing the Parquet for a GTFS feed's latest dataset + description: >- + Begin rendering this feed's latest dataset as Parquet, if it has not been already. Idempotent, and safe to call against a conversion that is already running: that one is reported rather than duplicated, and a dataset is never converted more than once at a time. Send `force` to replace an existing set. Work happens in the background - poll the GET on this path until it reports `ready`. + operationId: generateGtfsFeedParquet + tags: + - "parquet" + security: + - ApiKeyAuth: [] + requestBody: + required: false + content: + application/json: + schema: + $ref: "#/components/schemas/ParquetGenerateRequest" + responses: + "202": + description: >- + Accepted. The body carries the state as it stands, in the same shape the GET returns; a caller that only needs to start the work can ignore it and poll instead. + content: + application/json: + schema: + $ref: "#/components/schemas/ParquetDatasetState" + "404": + description: GTFS feed not found, or it has no dataset yet. + "422": + description: >- + Invalid request body - `retention_days` outside 1..60, or not an integer. + "500": + description: Internal server error. + /v1/operations/gtfs_datasets/{id}/parquet: + parameters: + - $ref: "#/components/parameters/dataset_id_path_param" + get: + summary: Where the Parquet for a GTFS dataset stands + description: >- + Whether this dataset has been rendered as Parquet - one file per GTFS table, every column typed as text - and if it is being prepared, how far along it is. A client can then query it in place over HTTP range requests instead of downloading and unpacking it. Read-only: this never starts a conversion, so it is safe to poll - the viewer polls it roughly twice a second while one runs. Use POST to begin one. + operationId: getGtfsDatasetParquet + tags: + - "parquet" + security: + - Authentication: [] + responses: + "200": + description: >- + The dataset's current state. Returned for every state including `absent`, which means no conversion has been requested yet rather than that anything is missing. + content: + application/json: + schema: + $ref: "#/components/schemas/ParquetDatasetState" + "404": + description: >- + GTFS dataset not found. Note this does **not** cover a dataset whose Parquet has never been generated - that is `absent` with a 200. + "500": + description: Internal server error. + post: + summary: Start preparing the Parquet for a GTFS dataset + description: >- + Begin rendering this dataset as Parquet, if it has not been already. Idempotent, and safe to call against a conversion that is already running: that one is reported rather than duplicated, and a dataset is never converted more than once at a time. Send `force` to replace an existing set. Work happens in the background - poll the GET on this path until it reports `ready`. + operationId: generateGtfsDatasetParquet + tags: + - "parquet" + security: + - ApiKeyAuth: [] + requestBody: + required: false + content: + application/json: + schema: + $ref: "#/components/schemas/ParquetGenerateRequest" + responses: + "202": + description: >- + Accepted. The body carries the state as it stands, in the same shape the GET returns; a caller that only needs to start the work can ignore it and poll instead. + content: + application/json: + schema: + $ref: "#/components/schemas/ParquetDatasetState" + "404": + description: GTFS dataset not found. + "422": + description: >- + Invalid request body - `retention_days` outside 1..60, or not an integer. + "500": + description: Internal server error. components: schemas: Redirect: @@ -3506,6 +3656,211 @@ components: value: true - feature_flag_id: "max_results" value: 50 + ExecutionAttempt: + x-operation: true + type: object + description: >- + One attempt by a background task against one dataset. Append-only: an attempt is never rewritten, so the sequence for a dataset is its full history rather than its current state. + properties: + task_name: + type: string + example: parquet_generation + dataset_stable_id: + type: string + nullable: true + example: mdb-1210-202402121801 + attempt: + description: 1-based, counted since this dataset last succeeded. + type: integer + example: 1 + variant: + description: >- + The worker that ran it. Read from the running service rather than from what was requested, so it is what actually happened. + type: string + nullable: true + example: s + variant_basis: + description: >- + Where that size came from. `measured` from the dataset's own size, `auto` from an override the builder set after an earlier failure, `operator` from one a person set by hand. + type: string + enum: [measured, auto, operator] + nullable: true + example: measured + override_at_attempt: + description: The size override in force for this feed when the attempt was routed, if any. + type: string + nullable: true + example: m + escalated_to: + description: >- + The worker this attempt moved the feed to. Null when nothing was escalated, which includes the cases where escalation was considered and declined - the failure was not a resource one, the largest worker was already in use, or the retry cap was spent. + type: string + nullable: true + example: m + status: + type: string + enum: [completed, failed] + example: failed + failure_kind: + description: What ran out. Null on success. + type: string + enum: [resource_memory, resource_disk, other] + nullable: true + example: resource_memory + error_type: + description: >- + The exception class. Recorded separately from the message because the most important one, `MemoryError`, carries an empty message. + type: string + nullable: true + example: MemoryError + error_message: + type: string + nullable: true + example: "Out of Memory Error: failed to allocate block of 262144 bytes" + started_at: + type: string + format: date-time + finished_at: + type: string + format: date-time + duration_ms: + type: integer + format: int64 + nullable: true + example: 4950 + peak_rss_bytes: + description: Peak resident memory. + type: integer + format: int64 + nullable: true + peak_vms_bytes: + description: >- + Peak address space, which is what the worker's limit actually caps. It runs two to three times resident memory, so a worker sized from `peak_rss_bytes` alone is under-provisioned by roughly half. + type: integer + format: int64 + nullable: true + ExecutionAttemptsResponse: + x-operation: true + type: object + properties: + feed_stable_id: + type: string + example: mdb-1210 + current_size_override: + description: >- + The worker size currently configured for this feed, if any. An override decides which worker runs the build, in place of the size the dataset would otherwise measure into. It is set either by an operator or automatically, after a build runs out of memory or disk; `variant_basis` on each attempt says which. + type: string + nullable: true + example: m + current_size_locked: + description: >- + Whether that override is held. An override is normally a starting point the builder may raise after a build runs out of resources, or lower after several that finished with room to spare. A locked one it will not move in either direction, so a feed locked below what it needs fails and stays failed, with no escalation, until a person changes the value - which is what this field is for. False when there is no override. + type: boolean + default: false + example: false + total: + type: integer + example: 2 + attempts: + type: array + items: + $ref: "#/components/schemas/ExecutionAttempt" + ParquetDatasetState: + x-operation: true + type: object + description: >- + Where a dataset's Parquet rendering stands. Which fields are set depends on `status`: `base_url` for `ready`, the progress fields for `preparing`, `message` for `failed`, and nothing beyond the identifiers for `absent`. What the dataset contains - its tables, their row and column counts, and what each weighed before and after conversion - is read from `{base_url}/manifest.json`, which is the one place that answer lives. The progress fields are named as the GTFS Garage viewer names its own load reports, so a client can pass them straight through. + required: + - status + - feed_stable_id + properties: + status: + description: >- + `absent` when no conversion has ever been requested - an ordinary starting state, not an error. `preparing` while one is queued or running, `ready` once the files are published, `failed` when the last attempt did not finish. + type: string + enum: [absent, preparing, ready, failed] + example: preparing + feed_stable_id: + description: The stable ID of the feed the dataset belongs to. + type: string + example: mdb-1210 + dataset_stable_id: + description: >- + The stable ID of the dataset. When the request addressed a feed, this is the feed's latest dataset, resolved at request time. + type: string + nullable: true + example: mdb-1210-202402121801 + base_url: + description: >- + Public URL prefix the Parquet files are served from, without a trailing slash. `{base_url}/manifest.json` is the dataset's description of itself - which tables exist, their row and column counts, and what each weighed before and after conversion - and each table is `{base_url}/{name}.parquet`. Null unless `status` is `ready`. The files are public and answer HTTP range requests; no query string is ever appended, so this can be used verbatim as a client's base URL. This response deliberately does not repeat the table list: a reader given one skips the manifest, and with it the sizes it carries. + type: string + format: url + nullable: true + example: https://files.mobilitydatabase.org/mdb-1210/mdb-1210-202402121801/parquet + phase: + description: >- + The step currently running, in the order they happen: `download` fetches the dataset archive, `extract` unpacks it, `convert` rewrites each table as Parquet, `upload` publishes the results, `summarise` records the manifest. Null unless `status` is `preparing`. + type: string + enum: + - start + - download + - upload + - extract + - convert + - summarise + - done + nullable: true + example: convert + done: + description: >- + Progress within the current phase: **bytes transferred** during `download`, and a count of files or tables otherwise. Null unless `status` is `preparing`. + type: integer + format: int64 + nullable: true + example: 12 + total: + description: >- + The target `done` is counting towards, in the same unit, or 0 when it is not knowable in advance - a download whose response carries no Content-Length, for instance. Null unless `status` is `preparing`. + type: integer + format: int64 + nullable: true + example: 32 + detail: + description: >- + The file or table currently being worked on; empty when there is none. Null unless `status` is `preparing`. + type: string + nullable: true + example: stop_times + message: + description: >- + Why the last attempt failed, in terms an operator can act on. Shown verbatim by the viewer, so it should name the actual cause. Null unless `status` is `failed`. + type: string + nullable: true + example: Conversion ran out of memory + generated_at: + description: When the current set of files finished generating. Null unless `status` is `ready`. + type: string + format: date-time + nullable: true + example: "2026-09-16T14:22:31Z" + ParquetGenerateRequest: + x-operation: true + type: object + description: Options for a Parquet generation request. + properties: + force: + description: >- + Regenerate even if the files already exist, replacing them. Has no effect on a conversion that is already running: that one is reported as-is rather than restarted. + type: boolean + default: false + example: false + retention_days: + description: >- + How long the generated files should be kept, counted from when they are built and not extended by use. Omit it and the builder's own default of 30 days applies; that default deliberately lives in one place rather than being restated here. On that date the dataset reports `absent` again, so the next request rebuilds it. The files themselves are removed by the storage bucket's own lifecycle rule shortly afterwards, on its schedule rather than at a guaranteed moment. + type: integer + minimum: 1 + maximum: 60 + example: 7 parameters: early_access_format_param: x-operation: true @@ -3621,6 +3976,67 @@ components: - desc default: desc example: asc + execution_attempt_status: + x-operation: true + name: status + in: query + description: Return only attempts that ended this way. + required: false + schema: + type: string + enum: [completed, failed] + example: failed + execution_attempt_failure_kind: + x-operation: true + name: failure_kind + in: query + description: >- + Return only failures of this kind. `resource_memory` and `resource_disk` are the two that trigger an automatic move to a larger worker; `other` never does. Implies `status=failed`, since a successful attempt has no kind. + required: false + schema: + type: string + enum: [resource_memory, resource_disk, other] + example: resource_memory + execution_attempt_from: + x-operation: true + name: from + in: query + description: >- + Return attempts that started at or after this timestamp, in ISO 8601 date-time format. Filtered on when the attempt began rather than when it ended, so a long build is found by when it was triggered. + required: false + schema: + type: string + format: date-time + example: "2026-04-01T00:00:00Z" + execution_attempt_to: + x-operation: true + name: to + in: query + description: Return attempts that started at or before this timestamp, in ISO 8601 date-time format. + required: false + schema: + type: string + format: date-time + example: "2026-04-30T23:59:59Z" + execution_attempt_task_name: + x-operation: true + name: task_name + in: query + description: >- + Limit the history to one task, for example `parquet_generation`. Omit it to see every task that has touched the feed. + required: false + schema: + type: string + example: parquet_generation + dataset_id_path_param: + x-operation: true + name: id + in: path + description: The dataset ID of the requested GTFS dataset. + required: True + schema: + type: string + example: mdb-1210-202402121801 securitySchemes: ApiKeyAuth: type: apiKey diff --git a/docs/parquet-feed-browsing.md b/docs/parquet-feed-browsing.md new file mode 100644 index 000000000..a694e9ab2 --- /dev/null +++ b/docs/parquet-feed-browsing.md @@ -0,0 +1,483 @@ +# Browsing a GTFS Feed - Operations API Flow + +This document describes the calls a client makes to browse a GTFS feed: how it asks the +Operations API whether a dataset has been rendered as Parquet, how it starts a rendering +and follows its progress, and what it fetches from the storage bucket once one exists. + +It is written for developers working in this repository. The only client today is the +operations web app, which lives in a separate repo and is referred to here as "the +viewer". + +For how the conversion itself works - DuckDB, the claim on `task_execution_log`, memory +behaviour, retention, local generation - see +[functions-python/parquet_builder/README.md](../functions-python/parquet_builder/README.md). +This document is about **what a client calls and in what order**, not how the files are +produced. + +Throughout, "**dataset**" means a `gtfsdataset` row addressed by its `stable_id` (for +example `mdb-1210-202402121801`), and "**feed**" means a `gtfsfeed` row addressed by its +own `stable_id` (`mdb-1210`). Every path parameter in this flow is a stable id, never a +UUID. + +## Table of Contents + +1. [Authentication](#authentication) +2. [Endpoint inventory](#endpoint-inventory) +3. [The state machine](#the-state-machine) +4. [Flow 1 - a feed that has never been converted](#flow-1---a-feed-that-has-never-been-converted) +5. [Flow 2 - a feed already converted](#flow-2---a-feed-already-converted) +6. [Flow 3 - a build that fails](#flow-3---a-build-that-fails) +7. [Polling](#polling) +8. [Progress fields](#progress-fields) +9. [Reading the output](#reading-the-output) +10. [Why the API does not return the table list](#why-the-api-does-not-return-the-table-list) +11. [Transport requirements for the files](#transport-requirements-for-the-files) +12. [Expiry](#expiry) +13. [Known deviations from the spec](#known-deviations-from-the-spec) +14. [Running it locally](#running-it-locally) + +--- + +## Authentication + +This is the part the spec will mislead you about, so it comes first. + +**Every call, GET and POST alike, needs a Google OAuth2 bearer token:** + +``` +Authorization: Bearer +Content-Type: application/json # POST with a body only +``` + +The token's `audience` must equal the `GOOGLE_CLIENT_ID` the function was deployed with +(`infra/functions-python/main.tf:879`). Validation happens in `RequestContextMiddleware` +(`functions-python/operations_api/src/middleware/request_context_middleware.py`), which +calls Google's `tokeninfo` endpoint and caches the result for the token's lifetime +(`functions-python/operations_api/src/middleware/request_context_oauth2.py:138-207`). +The Cloud Function is invokable by `allUsers`, so this middleware is the only gate. + +What the spec says instead: + +- the GETs declare `security: - Authentication: []`, and `Authentication` is **never + defined** under `securitySchemes` (`docs/OperationsAPI.yaml:3842-3846`). It is a + dangling reference, so the generated code fell back to the global scheme. +- the POSTs declare `ApiKeyAuth`, an `x-api-key` header. **Nothing reads it.** + `get_token_ApiKeyAuth` (`functions-python/operations_api/src/feeds_gen/security_api.py:22-38`) + is a generated stub whose body is `...`; it returns `None` and rejects nothing. + +So do not send `x-api-key` expecting it to authenticate anything, and do not expect a +GET to be unauthenticated. + +Preflight is handled before auth: `CORSMiddleware` is registered after +`RequestContextMiddleware` and therefore wraps it, so an `OPTIONS` request is answered +without a token (`functions-python/operations_api/src/main.py:44-58`). + +--- + +## Endpoint inventory + +### Resolving ids + +There is no dataset list or search endpoint in the Operations API. Two calls can find a +feed: + +| Method + path | operationId | Returns | +|---|---|---| +| `GET /v1/operations/feeds` | `getFeeds` | Paged feed list. `search_query` matches stable id, name and provider; `operation_status`, `data_type`, `offset`, `limit` | +| `GET /v1/operations/gtfs_feeds/{id}` | `getGtfsFeed` | One feed, including `latest_dataset` - the only Operations endpoint that hands back a dataset id | + +In practice a browse client needs neither. The feed-scoped Parquet endpoints resolve +`feed.latest_dataset` server-side and echo the result back as `dataset_stable_id`, so a +client holding only a feed id can do the whole flow. + +### The Parquet endpoints + +| Method + path | operationId | Purpose | +|---|---|---| +| `GET /v1/operations/gtfs_feeds/{id}/parquet` | `getGtfsFeedParquet` | State of the feed's latest dataset | +| `POST /v1/operations/gtfs_feeds/{id}/parquet` | `generateGtfsFeedParquet` | Start a build for it | +| `GET /v1/operations/gtfs_datasets/{id}/parquet` | `getGtfsDatasetParquet` | State of one specific dataset | +| `POST /v1/operations/gtfs_datasets/{id}/parquet` | `generateGtfsDatasetParquet` | Start a build for it | + +The feed and dataset variants are otherwise identical: same request body, same response +shape, same semantics. Pick the feed one unless you are deliberately addressing an older +dataset. + +**GET never starts work.** It is side-effect free and safe to poll. **POST** is what +begins a conversion, and it is idempotent: called against a build already running it +reports that build rather than queuing a second one. + +POST takes an optional body: + +```json +{ "force": false, "retention_days": 7 } +``` + +- `force` (default `false`) rebuilds a dataset that is already `ready`, replacing the + files. It has **no effect on a build in flight** - that one is reported as-is rather + than restarted. +- `retention_days` (1 to 60) overrides how long the output is kept. Omit it and the + builder applies its own default of 30 days. The schema deliberately declares no + default so the value lives in one place. + +A `retention_days` outside the range is a `422` from FastAPI's own body validation. + +--- + +## The state machine + +`GET` returns **200 for every state**, including `absent`. A `404` means something else +entirely: the feed or dataset does not exist. A client has to distinguish them, because +`absent` is the ordinary starting point of the whole flow. + +```mermaid +stateDiagram-v2 + [*] --> absent: no build ever requested + absent --> preparing: POST + preparing --> ready: build succeeded + preparing --> failed: build raised + failed --> preparing: POST (retry) + ready --> preparing: POST force=true + ready --> absent: retention expired +``` + +Which fields carry a value depends entirely on `status`. Every key is present in the +response; the ones that do not apply are `null`. + +| Field | `absent` | `preparing` | `ready` | `failed` | +|---|---|---|---|---| +| `status` | yes | yes | yes | yes | +| `feed_stable_id` | yes | yes | yes | yes | +| `dataset_stable_id` | yes | yes | yes | yes | +| `base_url` | - | - | **yes** | - | +| `generated_at` | - | - | **yes** | - | +| `phase` | - | **yes** | - | - | +| `done` | - | **yes** | - | - | +| `total` | - | **yes** | - | - | +| `detail` | - | **yes** | - | - | +| `message` | - | - | - | **yes** | + +The state is derived from a single `task_execution_log` row keyed +`(task_name="parquet_generation", entity_id=, run_id=)` +in `_state_of` +(`functions-python/operations_api/src/feeds_operations/impl/parquet_api_impl.py`): + +| Row | Reported as | +|---|---| +| no row | `absent` | +| `completed`, `expires_at` in the past | `absent` | +| `completed`, not expired | `ready` | +| `failed` | `failed` | +| `triggered` or `in_progress` | `preparing` | + +Two consequences worth knowing. A converter version bump changes `run_id`, so every +dataset built under the old version reads as `absent` and rebuilds on first request. And +a feed with no dataset at all is a **404**, not `absent`: `"GTFS feed has no dataset yet"`. + +The `failed` message is the builder's exception string, surfaced verbatim so an operator +can act on it. When the row carries no message the API substitutes +`"The conversion failed."`. + +--- + +## Flow 1 - a feed that has never been converted + +The common case on first browse. + +```mermaid +sequenceDiagram + autonumber + participant V as Viewer (browser) + participant API as Operations API + participant Q as Cloud Tasks + participant B as parquet_builder + participant GCS as Storage bucket + + V->>API: GET /v1/operations/gtfs_feeds/mdb-1210/parquet + API-->>V: 200 {status: "absent", dataset_stable_id: "mdb-1210-2024..."} + + V->>API: POST /v1/operations/gtfs_feeds/mdb-1210/parquet + API->>API: mark the dataset triggered (before dispatch) + API->>Q: enqueue build task + API-->>V: 200 {status: "preparing", phase: "start"} + + Q->>B: deliver task + B->>B: claim the dataset (try_acquire) + B->>GCS: read the dataset's files + B->>B: convert each table to Parquet + B->>GCS: upload tables, then manifest.json last + + loop every ~500ms until terminal + V->>API: GET .../parquet + API-->>V: 200 {status: "preparing", phase, done, total, detail} + end + + B->>API: (writes completion to task_execution_log) + V->>API: GET .../parquet + API-->>V: 200 {status: "ready", base_url, generated_at} + + V->>GCS: GET {base_url}/manifest.json + GCS-->>V: table list, row and column counts, sizes + V->>GCS: HEAD + ranged GET {base_url}/stops.parquet + GCS-->>V: 206 Partial Content +``` + +Points that matter to a client: + +- The POST response already carries the state, in the same shape the GET returns. A + client that only wants to start the work can ignore the body and poll. +- The dataset is marked `triggered` **before** the task is dispatched, so the very next + poll reads `preparing` rather than `absent` again. +- If the enqueue fails, the POST is a `500` and the dataset is left `failed` rather than + stuck `preparing`, so it stays retriable. + +--- + +## Flow 2 - a feed already converted + +```mermaid +sequenceDiagram + autonumber + participant V as Viewer (browser) + participant API as Operations API + participant GCS as Storage bucket + + V->>API: GET /v1/operations/gtfs_feeds/mdb-1210/parquet + API-->>V: 200 {status: "ready", base_url, generated_at} + V->>GCS: GET {base_url}/manifest.json + V->>GCS: HEAD + ranged GET {base_url}/.parquet +``` + +No POST, no polling. A client should always GET first and only POST on `absent` or +`failed`. + +A POST here returns the ready state untouched unless `force: true` is sent, in which case +the set is rebuilt and replaced. `force` cannot interrupt a build that is already +running. + +--- + +## Flow 3 - a build that fails + +A failed build is reported through the same GET, as `failed` with a `message`. Note that +the builder function itself returns HTTP 200 to Cloud Tasks even when the conversion +fails: nothing in the conversion fails in a way a retry would fix, so the reason is +recorded in the database rather than signalled by a status code. + +```mermaid +sequenceDiagram + autonumber + participant V as Viewer (browser) + participant API as Operations API + participant B as parquet_builder + + V->>API: POST .../parquet + API-->>V: 200 {status: "preparing"} + B->>B: build raises + B->>B: record failed + release the claim + V->>API: GET .../parquet + API-->>V: 200 {status: "failed", message: "Dataset archive not found at gs://..."} + V->>API: POST .../parquet (retry, no force needed) + API-->>V: 200 {status: "preparing"} +``` + +`failed` is re-triggerable without `force`. The two causes a client will actually see: + +- **`Dataset archive not found at gs://...`** - the dataset has no extracted files and no + `.zip` in the bucket. The API cannot detect this up front; it checks only that the feed + and dataset rows exist, so the error arrives through the build rather than from the POST. +- **`No GTFS tables could be converted for `** - the archive held nothing + usable. + +--- + +## Polling + +- Poll the GET. There is no webhook, no `Retry-After` and no rate limiting on this path. +- Roughly **twice a second** is the documented expectation and what the viewer does. +- The server throttles its own progress writes to at most one per second, except that a + phase change is always written immediately. Polling faster than that just returns the + same reading. +- A progress write that fails is logged and swallowed. A `preparing` reading can + therefore go stale without the build having died, so do not treat an unchanging reading + as failure. A build that really dies leaves its claim to expire and the row stays + `in_progress`, which still reads as `preparing`. + +--- + +## Progress fields + +`phase` tells the client which step is running. The spec lists the enum in a different +order from the order they occur; the order below is what actually happens. + +| `phase` | When | `done` / `total` | `detail` | +|---|---|---|---| +| `start` | Claim taken, nothing done yet | `0 / 0` | `""` | +| `convert` | Rewriting each table as Parquet | table index / table count | table name, e.g. `stop_times` | +| `upload` | Publishing `manifest.json` | `1 / 1` | `manifest.json` | +| `summarise` | Recording the result | `0 / 0` | `""` | + +`total: 0` means "not knowable in advance", not "nothing to do". + +Three phases a client will never observe, despite being in the enum: + +- `download` - the archive is read over the network with ranged requests as members are + needed, so there is no stage during which it is being fetched and nothing else. +- `extract` - members are unpacked lazily inside `convert`. +- `done` - written only together with completion, at which point the API reports `ready` + and suppresses `phase` altogether. + +--- + +## Reading the output + +`base_url` is a public URL prefix with **no trailing slash and no query string**: + +``` +https://files.mobilitydatabase.org/mdb-1210/mdb-1210-202402121801/parquet +``` + +It can be used verbatim as a reader's base URL. That is deliberate: the reader rebuilds +each file's URL from the origin and path only, discarding any query string, so a signed +URL could not survive the round trip - which is why the objects are public instead. + +Fetch `{base_url}/manifest.json` first. It is the dataset's description of itself and the +only place the table list, counts and sizes live. + +```json +{ + "version": 2, + "generated_at": "2026-09-17T22:41:03+00:00", + "converter_version": "2", + "source": { "kind": "zip", "bytes": 4821334 }, + "totals": { "uncompressed_bytes": 18422910, "stored_bytes": 903411 }, + "tables": [ + { "name": "stops", "file": "stops.parquet", "rows": 4821, "columns": 12, + "bytes": 481223, "compressed_bytes": 92210, "parquet_bytes": 41880 } + ] +} +``` + +| Field | Meaning | +|---|---| +| `version` | Manifest format version, integer, currently `2` | +| `generated_at` | ISO-8601 UTC, seconds precision | +| `converter_version` | The builder's converter version, a string. Matches the `run_id` of the tracking row | +| `source.kind` | `zip` or `folder` | +| `source.bytes` | Size of the originating archive, nullable | +| `totals.uncompressed_bytes` | Sum of table `bytes` where known | +| `totals.stored_bytes` | Sum of table `parquet_bytes` | +| `tables[].name` | Table name: the GTFS file stem (`stops.txt` to `stops`), plus `locations` for `locations.geojson` | +| `tables[].file` | `{name}.parquet` | +| `tables[].rows` | Row count, read from the Parquet footer | +| `tables[].columns` | Column count | +| `tables[].bytes` | **Source** file size, nullable | +| `tables[].compressed_bytes` | Size inside the zip. `null` for a folder source | +| `tables[].parquet_bytes` | Size of the written Parquet file | + +**Version 2 is not a superset of version 1.** In v1 `bytes` meant the Parquet size; in v2 +it means the source size and the converted size is `parquet_bytes`. Check `version` +before reading either, and do not add fields under an old number. + +Each table is then `{base_url}/{file}`, that is `{base_url}/{name}.parquet`. + +Two properties of the Parquet the reader depends on: **every column is text**, and an +**empty CSV field becomes NULL** rather than an empty string. A reader that filters with +`ILIKE` or `= ''` relies on both. + +**The publish order is a contract.** Tables are uploaded first, `manifest.json` last, and +stale objects are pruned only afterwards. Treat the manifest's presence as the signal that +the set is complete, and do not probe the bucket for table names: on a first build, before +the API reports `ready`, probing can catch a partially published set. A client driven by +the API or by the manifest never sees one. + +--- + +## Why the API does not return the table list + +`ready` carries only `base_url`, never the tables. This is deliberate, not an omission: a +reader handed a table list skips `manifest.json` entirely, and the manifest is where the +row counts and sizes live - and with them the load report the viewer draws from them. The +builder does record the tables on the tracking row, for diagnostics, but they are not +served. + +--- + +## Transport requirements for the files + +The Parquet files are read directly by the browser, not proxied through the API, so the +bucket itself has to cooperate. + +- **The objects are public**, set per object at upload. The datasets bucket uses object + ACLs rather than uniform bucket-level access, so a refused ACL fails the build rather + than being ignored - a `ready` state never advertises a `base_url` that answers 403. +- **The reader needs `HEAD` as well as ranged `GET`.** It probes each table with `HEAD` + before reading the footer. +- **The bucket CORS config must allow both methods** (`infra/batch/main.tf`): + + ```hcl + cors { + origin = ["*"] + method = ["GET", "HEAD"] + response_header = ["*"] + } + ``` + + GCS matches CORS methods literally and does not treat `HEAD` as a kind of `GET`. With + `GET` alone the `HEAD` preflight comes back with no `Access-Control-Allow-*` headers at + all and the browser blocks the request. The failure surfaces as an opaque + `NetworkError: Failed to execute 'send' on 'XMLHttpRequest'` with no mention of CORS, + and the viewer reports it as "no tables found" - so if a freshly deployed environment + cannot load any table while the objects are plainly present, check this first. + + The CORS block lives in the `infra/batch` module, so it is applied by the + `datasets-batch-deployer` workflow, not by a functions deploy. + +--- + +## Expiry + +Each build stamps every object it publishes with a GCS `customTime` of +`build time + retention_days`, and the bucket carries one lifecycle rule that deletes an +object once its own `customTime` has passed. The same instant is recorded on the tracking +row, which is what the API reads to report `absent` again on that date. + +For a client this means a `base_url` is not permanent. A dataset that was `ready` last +month can read `absent` today, and the next POST rebuilds it. Lifecycle deletion is +asynchronous and GCS gives no promptness guarantee, so there is a window where the API +reports `absent` while the files are still present - harmless, since a rebuild overwrites +them. + +--- + +## Known deviations from the spec + +Live behaviour that contradicts `docs/OperationsAPI.yaml`. Trust this table over the spec +until they are reconciled. + +| Deviation | Why | +|---|---| +| POST responds **200**, not the `202` the spec declares | The generated router sets no `status_code`, so FastAPI's default applies | +| GETs require a bearer token despite declaring an undefined `Authentication` scheme | `Authentication` is never defined under `securitySchemes`, so the generated code applied the global `ApiKeyAuth` dependency instead | +| `x-api-key` authenticates nothing | `get_token_ApiKeyAuth` is a generated stub with an empty body | +| An unauthenticated call returns **500**, not `401` | `RequestContextMiddleware` is raw ASGI middleware outside Starlette's exception handling, so its `HTTPException` reaches the Cloud Function wrapper and is returned as a 500 | +| `phase: "extract"` and `phase: "done"` never appear | `extract` is emitted only by a function with no remaining call sites; `done` is written with completion, when `phase` is suppressed | + +--- + +## Running it locally + +`functions-python/parquet_builder/README.md` has the full local walkthrough, including the +case where the Operations API itself is what you want to exercise. The short version is +that Cloud Tasks does not dispatch locally - the `PARQUET_BUILDER_QUEUE_M`/`_L` vars are +unset, so the enqueue is a logged no-op and the dataset sits at `preparing` - which means +the builder has to be invoked by hand in the queue's place. + +To skip the API entirely and just get Parquet files a browser can read: + +```bash +scripts/parquet-generate-local.sh mdb-1210 --serve # http://localhost:8090 +``` + +That serves the output with range support and the CORS headers a reader needs, neither of +which `python -m http.server` provides. diff --git a/functions-python/pmtiles_builder/src/ephemeral_workdir.py b/functions-python/helpers/ephemeral_workdir.py similarity index 84% rename from functions-python/pmtiles_builder/src/ephemeral_workdir.py rename to functions-python/helpers/ephemeral_workdir.py index e7f974a14..110cec79b 100644 --- a/functions-python/pmtiles_builder/src/ephemeral_workdir.py +++ b/functions-python/helpers/ephemeral_workdir.py @@ -15,23 +15,30 @@ class EphemeralOrDebugWorkdir: removes sibling directories older than a TTL (default 3600s / override via WORKDIR_MAX_AGE_SECONDS), and deletes the created directory on exit. - Only directories whose names start with the fixed CLEANUP_PREFIX are considered for cleanup + Only directories whose names start with this instance's `owner_prefix` are considered + for cleanup to avoid deleting unrelated folders that might exist under the same root. The final on-disk directory name always starts with the hardcoded prefix 'pmtiles_'. The caller-supplied prefix (if any) is appended verbatim after that. """ - CLEANUP_PREFIX = "pmtiles_" - def __init__( self, + owner_prefix: str, dir: str | None = None, prefix: str | None = None, logger: logging.Logger | None = None, ): + """`owner_prefix` names the function these directories belong to. + + Required rather than defaulted, and the only thing stopping one function from + deleting another's work: the root is a shared convention (`/tmp/in-memory`), so + a sweep that matched everything under it would be a cross-function data race. + """ import tempfile + self.owner_prefix = owner_prefix self._debug_dir = os.getenv("DEBUG_WORKDIR") or None self._root = dir or os.getenv("WORKDIR_ROOT", "/tmp/in-memory") self._logger = logger or get_logger("Workdir") @@ -50,14 +57,14 @@ def __init__( self._cleanup_old() # Simple prefix: fixed manager prefix + raw user prefix (if any) - combined_prefix = self.CLEANUP_PREFIX + (prefix or "") + combined_prefix = self.owner_prefix + (prefix or "") self._temp = tempfile.TemporaryDirectory(dir=self._root, prefix=combined_prefix) self.name = self._temp.name def _cleanup_old(self): """ - Delete stale work directories created by this manager (names starting with CLEANUP_PREFIX) + Delete stale work directories created by this manager (names starting with owner_prefix) whose modification time is older than the configured TTL. """ import time @@ -78,7 +85,7 @@ def _cleanup_old(self): try: if not entry.is_dir(follow_symlinks=False): continue - if not entry.name.startswith(self.CLEANUP_PREFIX): + if not entry.name.startswith(self.owner_prefix): continue try: age = now - entry.stat(follow_symlinks=False).st_mtime diff --git a/functions-python/helpers/parquet_policy.py b/functions-python/helpers/parquet_policy.py new file mode 100644 index 000000000..392af92fc --- /dev/null +++ b/functions-python/helpers/parquet_policy.py @@ -0,0 +1,86 @@ +# +# MobilityData 2026 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Which worker a Parquet build belongs on. + +Two processes decide this and they have to agree. The Operations API routes a build when +it enqueues one, and the builder re-routes it when one runs out of resources; if their +tables disagreed, a dataset would bounce between two workers forever. They used to hold +separate copies with a comment in each saying they must match, which is a comment where +a shared constant belongs. + +`sizing.py` is the mechanism and is deliberately generic. This is the Parquet builder's +policy: its bands, its budgets, its config key, and the thresholds that move a feed +between rungs. +""" + +from __future__ import annotations + +from shared.helpers.sizing import Size, Tier + +MIB = 1024**2 + +# Routing table for the build workers. The measure is the largest single uncompressed +# file in the dataset, because the builder's in-memory volume holds one at a time, so +# that file is what decides whether a build fits. Totals are the wrong signal: a feed of +# many medium files is cheaper than one with a single huge one, and the compressed size +# is wrong by a factor that runs from 4x to 13x across the catalogue. +# +# The bands come from measuring it: of 4277 feeds, the median archive is 0.2 MB and only +# 31 are above 100 MB, while about ten feeds have a single member over 1 GB and the worst +# has one of 4.8 GB. So most traffic belongs on a worker sized for a few hundred MB, and +# the large worker exists for roughly a dozen feeds. Stats as of 2026/09. +# +# `vms_budget_bytes` is the other half: what the rung provides, beside what it accepts. +# Each is the RLIMIT_AS the worker ends up with, which `limit_gcp_memory` derives as +# memory - volume - 200 MiB margin, from `local.parquet_builder_sizes` in +# infra/functions-python/main.tf. A build records what it actually used, so these are +# what that usage gets compared against when deciding a feed could run somewhere smaller. +# Change them when the terraform map changes; a worker logs its real RLIMIT_AS at +# startup, so a drift is visible rather than silent. +SIZE_TIERS = ( + # 3Gi memory, 1Gi volume + Tier(size=Size.S, max_bytes=256_000_000, vms_budget_bytes=1848 * MIB), + # 7Gi memory, 3Gi volume + Tier(size=Size.M, max_bytes=1_500_000_000, vms_budget_bytes=3896 * MIB), + # 16Gi memory, 8Gi volume + Tier(size=Size.L, max_bytes=None, vms_budget_bytes=7992 * MIB), +) + +# A feed can be pinned to a size by hand through `config_value_feed`. An override is used +# as given; the measurement is not consulted at all. +SIZE_CONFIG_NAMESPACE = "parquet_builder" +SIZE_CONFIG_KEY = "size" + +# Used only when a dataset has no per-file rows and no recorded unzipped total. GTFS +# compresses roughly 5-15x; the low end is deliberate, since overestimating the content +# of an archive routes up rather than down. +COMPRESSION_RATIO = 5 + +# How many consecutive successful builds a feed needs before its override comes down a +# rung, and how much of the smaller rung each of them has to leave unused. +# +# The two directions are deliberately asymmetric. Escalation acts on a single failure, +# because the cost of staying too small is a build that cannot finish. Coming down is +# only ever an economy, so it can afford to be slow and dull: three builds in a row, each +# leaving 40% of the smaller worker spare on both memory and disk. +# +# Both are dials. A feed whose datasets alternate between large and small is the case +# that could still move down and back up repeatedly - the disk axis usually catches it, +# because the large dataset's own build breaks the streak, but that is a property of the +# evidence rather than a guarantee. If flapping shows up in `task_execution_attempt`, +# lengthen the streak or lower the fraction. +DOWNSIZE_STREAK = 3 +DOWNSIZE_HEADROOM = 0.60 diff --git a/functions-python/helpers/requirements_dev.txt b/functions-python/helpers/requirements_dev.txt index 800a4ac11..dbea1c22d 100644 --- a/functions-python/helpers/requirements_dev.txt +++ b/functions-python/helpers/requirements_dev.txt @@ -1,4 +1,5 @@ Faker pytest~=7.4.3 urllib3-mock -requests-mock \ No newline at end of file +requests-mock +psutil diff --git a/functions-python/helpers/runtime_metrics.py b/functions-python/helpers/runtime_metrics.py index 8c048a646..9c31f7870 100644 --- a/functions-python/helpers/runtime_metrics.py +++ b/functions-python/helpers/runtime_metrics.py @@ -1,8 +1,112 @@ import functools +import logging +import resource +import sys +import threading import time import tracemalloc + import psutil -import logging + +MB = 1024**2 + + +def _max_rss_bytes() -> int: + """Peak resident set size of this process, in bytes. + + Note this is *resident*, while `limit_gcp_memory` caps *address space* via + RLIMIT_AS. Virtual is always at least RSS and for a process mapping a database + engine it is far larger, so sizing a container from RSS alone under-provisions it - + which is why `vms` is reported beside this. + + `tracemalloc` only sees blocks Python itself allocated, so for a function whose + heavy lifting happens in a C extension (DuckDB, for one) it reports the smallest + consumer and misses the one that decides the instance's memory allocation. RSS + covers both. + + This is the process high-water mark and it never resets, so on a warm instance it + spans earlier requests too. It therefore over-reports rather than under-reports, + which is the right direction for a number used to size a container. + + `ru_maxrss` is kilobytes on Linux and bytes on macOS. + """ + peak = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + return peak if sys.platform == "darwin" else peak * 1024 + + +DEFAULT_SAMPLE_INTERVAL_SECONDS = 0.25 + + +class MemorySampler: + """The high-water mark of this process's memory over a span of work. + + Neither reading available for free answers "how much did this piece of work need", + which is the number a container is sized from: + + - `ru_maxrss` is a genuine peak, but it never resets. On a warm instance it carries + whatever the previous request reached, so it over-reports an individual run. + - `memory_info().vms` resets, but it is instantaneous. Read once at the end of a job + that has already closed its database connection, it describes the quiet moment + afterwards rather than the busy one in the middle, and under-reports badly. + + Polling answers it. The peak of a Parquet build falls inside a single table + conversion, not between tables, so sampling at loop boundaries would miss it too. + + The thread is a daemon so a sampler that is never stopped cannot hold the process + open, and every read is guarded: a measurement must never fail the work it measures. + """ + + def __init__(self, interval: float = DEFAULT_SAMPLE_INTERVAL_SECONDS): + self.interval = interval + self.peak_rss_bytes = 0 + self.peak_vms_bytes = 0 + self._stop = threading.Event() + self._thread = None + + def _sample(self) -> None: + try: + info = psutil.Process().memory_info() + except Exception: + return + self.peak_rss_bytes = max(self.peak_rss_bytes, info.rss) + self.peak_vms_bytes = max(self.peak_vms_bytes, info.vms) + + def _run(self) -> None: + while not self._stop.wait(self.interval): + self._sample() + + def start(self) -> "MemorySampler": + """Begin sampling. Takes one reading immediately, so a span shorter than the + interval still reports something rather than zero.""" + self._sample() + self._stop.clear() + self._thread = threading.Thread(target=self._run, daemon=True) + self._thread.start() + return self + + def stop(self) -> dict: + """Stop sampling and return the peaks, in the shape `record_attempt` reads.""" + self._stop.set() + thread, self._thread = self._thread, None + if thread is not None: + # Bounded so a wedged sampler cannot hold up the request it measured. + thread.join(timeout=self.interval * 4) + self._sample() + return self.metrics() + + def metrics(self) -> dict: + """The peaks so far. Zero means nothing could be read, which is reported as + absent rather than as a measurement of nothing.""" + return { + "peak_rss_bytes": self.peak_rss_bytes or None, + "peak_vms_bytes": self.peak_vms_bytes or None, + } + + def __enter__(self) -> "MemorySampler": + return self.start() + + def __exit__(self, *_) -> None: + self.stop() def track_metrics(metrics=("time", "memory", "cpu")): @@ -48,7 +152,20 @@ def wrapper(*args, **kwargs): tracemalloc.stop() if metrics_message: metrics_message += ", " - metrics_message += f"memory: {current / (1024 ** 2):.2f} MB (peak: {peak / (1024 ** 2):.2f} MB)" + metrics_message += ( + f"memory: {current / MB:.2f} MB (peak: {peak / MB:.2f} MB)" + ) + # Kept beside the tracemalloc figures rather than replacing them, so + # the log line stays comparable with what is already in Cloud Logging. + try: + info = process.memory_info() + metrics_message += ( + f", rss: {info.rss / MB:.2f} MB" + f" (process peak: {_max_rss_bytes() / MB:.2f} MB)" + f", vms: {info.vms / MB:.2f} MB" + ) + except Exception as error: + logger.debug("Could not read memory info: %s", error) if "cpu" in metrics: cpu_after = process.cpu_percent(interval=None) if metrics_message: diff --git a/functions-python/helpers/sizing.py b/functions-python/helpers/sizing.py new file mode 100644 index 000000000..efe88c9f4 --- /dev/null +++ b/functions-python/helpers/sizing.py @@ -0,0 +1,573 @@ +# +# MobilityData 2026 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Routing a job to a worker sized for it. + +A function provisioned for its worst case spends most of its life running trivial work. +Measured on `pmtiles-builder-prod`, 85% of invocations finish inside 30 seconds and +account for under 7% of the compute, while every one of them is billed at the memory the +largest feed needs. Deploying the same source as two differently sized functions and +choosing between them at enqueue time is what this module exists for. + +The tier arithmetic is generic: the caller supplies the tiers and the measure. +`size_for_dataset` builds on it for the common case of a job whose weight is a GTFS +dataset, which is how every candidate function here is triggered. + +Whether the *measure* transfers is a separate question per function. The largest single +uncompressed file is right for the Parquet builder because its volume holds one at a +time; a function bounded by something else wants `choose_size` with its own measure. + +Two rules: + +- **An unknown measure routes to the largest tier.** Sizes come from database columns + that are nullable and were added without a backfill, so "I don't know" is a normal + answer, and guessing small turns a missing row into an OOM. +- **An override decides on its own.** A configured size is the answer, whichever way it + differs from the measurement. The measurement is a heuristic; someone who set the + config has looked at the feed. +""" + +from __future__ import annotations + +import errno +import logging +from dataclasses import dataclass +from datetime import datetime, timezone +from enum import Enum +from typing import TYPE_CHECKING, Iterable, Optional, Sequence + +from sqlalchemy import func + +from shared.common.config_reader import get_config_value +from shared.database_gen.sqlacodegen_models import Gtfsfile + +if TYPE_CHECKING: # pragma: no cover + from sqlalchemy.orm import Session + +# GTFS compresses roughly 5-15x. The low end is deliberate: overestimating what an +# archive expands to routes a job up a rung rather than down. +DEFAULT_COMPRESSION_RATIO = 5 + + +class Basis(Enum): + """Where a size came from, recorded so a change is legible afterwards.""" + + MEASURED = "measured" + OPERATOR = "operator" + AUTO = "auto" + + +class FailureKind(Enum): + """What ran out, when something did.""" + + RESOURCE_MEMORY = "resource_memory" + RESOURCE_DISK = "resource_disk" + OTHER = "other" + + @property + def is_resource(self) -> bool: + return self is not FailureKind.OTHER + + +# Matched on the exception's class name rather than by importing the libraries that raise +# them, so this stays usable from any function without pulling in duckdb. +_MEMORY_TYPES = {"MemoryError", "OutOfMemoryException"} +_DISK_MARKERS = ("no space left on device", "errno 28") + + +def classify_failure(exc: BaseException) -> FailureKind: + """What kind of exhaustion this failure was, if any. + + Reads the exception, not a stored message, because `str(MemoryError())` is the empty + string - CPython raises a no-args singleton - so the failure the largest worker + exists for is invisible once the message has been written to a column. + """ + name = type(exc).__name__ + message = str(exc).lower() + + if isinstance(exc, OSError) and getattr(exc, "errno", None) == errno.ENOSPC: + return FailureKind.RESOURCE_DISK + # DuckDB reports a full spill directory as its own IOException, not an OSError. + if any(marker in message for marker in _DISK_MARKERS): + return FailureKind.RESOURCE_DISK + if name in _MEMORY_TYPES or message.startswith("out of memory error"): + return FailureKind.RESOURCE_MEMORY + return FailureKind.OTHER + + +class Size(Enum): + """Worker sizes, smallest first. `value` is the routing suffix.""" + + S = "s" + M = "m" + L = "l" + + @classmethod + def parse(cls, raw) -> Optional["Size"]: + """A size from free-form config, or None when it is absent or unrecognised. + + Config values are operator-entered, so an unusable one must not be fatal - it + falls back to the measured size rather than failing the request. + """ + if raw is None: + return None + try: + return cls(str(raw).strip().lower()) + except ValueError: + logging.warning("Ignoring unrecognised size override %r", raw) + return None + + +LARGEST = Size.L + + +@dataclass(frozen=True) +class Tier: + """One row of a routing table: what this size accepts, and what it provides. + + `max_bytes` is the measure this rung takes, `None` for the catch-all, which must come + last. `vms_budget_bytes` is the address space the worker actually has - optional, + because routing does not need it; only a decision to move a job to a *different* rung + does, and that has to compare observed use against what that other rung allows. + """ + + size: Size + max_bytes: Optional[int] + vms_budget_bytes: Optional[int] = None + + +def choose_size( + measure: Optional[int], + tiers: Sequence[Tier], + override: Optional[Size] = None, +) -> Size: + """Pick a worker size for a job whose weight is `measure` bytes. + + An `override` is returned as given and the measurement is not consulted. + + `measure` is whatever bounds the job - for a Parquet build, the largest single + uncompressed file, because the volume holds one at a time. `None` means unknown, + which routes to the largest tier. + """ + if override is not None: + return override + + if measure is None: + return LARGEST + + for tier in tiers: + if tier.max_bytes is None or measure < tier.max_bytes: + return tier.size + return LARGEST + + +def escalate( + current: Optional[Size], + tiers: Sequence[Tier], + *, + attempts: int, + max_attempts: int, +) -> Optional[Size]: + """The next rung up for a job that ran out of resources, or None to stop. + + None means stop, for any of three reasons: the attempt cap is spent, the job is + already on the largest rung, or the current rung is unknown. All three are terminal, + because the alternative is a loop that re-queues forever at real cost. + """ + if attempts >= max_attempts: + return None + + ladder = [tier.size for tier in tiers] + if current is None or current not in ladder: + return None + + index = ladder.index(current) + if index + 1 >= len(ladder): + return None + return ladder[index + 1] + + +def demote(current: Optional[Size], tiers: Sequence[Tier]) -> Optional[Size]: + """The rung below `current`, or None when there is none to go to. + + The mirror of `escalate`, without its attempt cap: coming down is never a retry of + anything, so there is no loop to bound. It applies to the next job, not this one. + """ + ladder = [tier.size for tier in tiers] + if current is None or current not in ladder: + return None + + index = ladder.index(current) + if index == 0: + return None + return ladder[index - 1] + + +def tier_for(size: Size, tiers: Sequence[Tier]) -> Optional[Tier]: + """The row of the routing table for this size, or None if it has none.""" + for tier in tiers: + if tier.size is size: + return tier + return None + + +def fits_within( + tier: Tier, + *, + peak_vms_bytes: Optional[int], + largest_member_bytes: Optional[int], + headroom: float, +) -> bool: + """Whether one observed run would have fitted `tier`, with `headroom` to spare. + + Both axes, because a build can exhaust either one: address space, against what that + rung allows, and the largest file it had to hold, against what that rung's volume + accepts. `headroom` of 0.6 means each has to come in under 60% of the limit. + + An observation that is missing reads as "no". These numbers were added after the + table existed, so an older attempt simply has nothing to say, and the direction that + leaves a job where it is costs money while the other costs failed builds. + """ + if not 0 < headroom <= 1: + raise ValueError(f"headroom must be in (0, 1], got {headroom}") + + budget = tier.vms_budget_bytes + if not peak_vms_bytes or not budget: + return False + if peak_vms_bytes >= budget * headroom: + return False + + if not largest_member_bytes: + return False + # The catch-all rung accepts any file, so only the memory axis constrains it. + if tier.max_bytes is not None and largest_member_bytes >= tier.max_bytes * headroom: + return False + + return True + + +def first_known(*candidates: Optional[int]) -> Optional[int]: + """The first candidate that is a usable positive size, or None. + + Size columns are nullable and can be zero for a dataset whose files were never + recorded, so a plain `or` chain would let a 0 through as a real measurement and + route the job to the smallest worker. + """ + for candidate in candidates: + if candidate: + try: + value = int(candidate) + except (TypeError, ValueError): + continue + if value > 0: + return value + return None + + +def queue_env_var(prefix: str, size: Size) -> str: + """The env var naming the Cloud Tasks queue for this size, e.g. `..._QUEUE_L`.""" + return f"{prefix}_QUEUE_{size.value.upper()}" + + +def function_name(base: str, size: Size, environment: str) -> str: + """The deployed function name for this size, e.g. `parquet-builder-l-dev`.""" + return f"{base}-{size.value}-{environment}" + + +def describe(tiers: Iterable[Tier]) -> str: + """A routing table in one line, for the log that records why a size was chosen.""" + return ", ".join( + ( + f"{tier.size.value}<{tier.max_bytes}" + if tier.max_bytes + else f"{tier.size.value}:rest" + ) + for tier in tiers + ) + + +def largest_file_bytes(db_session: "Session", dataset) -> Optional[int]: + """The biggest single uncompressed file in the dataset, or None if unrecorded. + + One indexed aggregate on `gtfsfile`, which carries an index on `gtfs_dataset_id`. + Returns None rather than 0 for a dataset with no file rows, so the caller can tell + "no files recorded" from "files recorded, all empty". + """ + try: + return ( + db_session.query(func.max(Gtfsfile.file_size_bytes)) + .filter(Gtfsfile.gtfs_dataset_id == dataset.id) + .scalar() + ) + except Exception as error: + logging.warning("Could not measure %s: %s", dataset.stable_id, error) + return None + + +def measure_dataset( + db_session: "Session", + dataset, + compression_ratio: int = DEFAULT_COMPRESSION_RATIO, +) -> tuple[Optional[int], str]: + """How heavy a dataset is, and which rung of the fallback produced the answer. + + The size columns are nullable and were added without a backfill, so a dataset + processed before #1284 has none of them. Each rung is a worse approximation than the + one above, and all of them err upwards: `unzipped_size_bytes` is the sum rather than + the maximum, and the compressed estimate uses the low end of the ratio. + """ + largest = largest_file_bytes(db_session, dataset) + if largest: + return int(largest), "largest file" + + total = first_known(getattr(dataset, "unzipped_size_bytes", None)) + if total: + return total, "unzipped total" + + zipped = first_known(getattr(dataset, "zipped_size_bytes", None)) + if zipped: + return zipped * compression_ratio, "estimated from the archive" + + return None, "unknown" + + +# Marks a value the builder wrote for itself. Anything without it was set by a person. +SOURCE_AUTO = "auto" +SOURCE_OPERATOR = "operator" + + +@dataclass(frozen=True) +class Override: + """A configured size for a feed: what it is, who set it, and whether it may move.""" + + size: Optional[Size] = None + source: Optional[str] = None + locked: bool = False + + +def size_override( + db_session: "Session", feed, namespace: str, key: str = "size" +) -> Override: + """The size configured for this feed, or an empty `Override`. + + One row per feed holds every case. Three accepted shapes: + + - a bare `"l"` is a person's, and may be moved; + - `{"size": "l", "source": "auto"}` is the builder's own, and may be moved; + - `{"size": "l", "locked": true}` is held, and may not. + + All three decide the routing outright. `source` is only a record of who wrote the + value last. It used to decide whether the builder could overwrite the row, which + meant inferring an operator's intent from the fact of their authorship; `locked` is + that intent stated, so it is chosen rather than guessed at. + + The bare form is what a human writing SQL by hand produces, so it is the one that + needs no ceremony. + """ + try: + raw = get_config_value(namespace, key, feed_id=feed.id, db_session=db_session) + except Exception as error: + logging.warning( + "Could not read the size override for %s: %s", feed.stable_id, error + ) + return Override() + + if isinstance(raw, dict): + return Override( + size=Size.parse(raw.get("size")), + # The builder always stamps its own source, so a hand-written object that + # omits it is a person's - which is now the common case, since `locked` can + # only be expressed in the object form. + source=raw.get("source") or SOURCE_OPERATOR, + # Anything but a true boolean reads as unlocked. Config values are entered by + # hand, and a typo must not quietly freeze a feed where nothing can move it. + locked=raw.get("locked") is True, + ) + size = Size.parse(raw) + return Override(size=size, source=SOURCE_OPERATOR if size else None) + + +@dataclass(frozen=True) +class Routing: + """A sizing decision: what to run on, and where that came from.""" + + size: Size + basis: Basis + + +def size_for_dataset( + db_session: "Session", + feed, + dataset, + *, + tiers: Sequence[Tier], + namespace: str, + key: str = "size", + compression_ratio: int = DEFAULT_COMPRESSION_RATIO, +) -> Routing: + """Which worker should handle this dataset, and where that came from. + + An override decides; otherwise the measurement picks a tier. Nothing else. + + The override is absolute whoever wrote it. A person's is obvious enough. The + builder's own is absolute too, because the alternative - weighing it against the + measurement so a grown feed can overtake it - only changes the outcome when the + measurement is larger, and that case already resolves itself: the build fails and the + escalation moves it up. One wasted build is not worth a second code path. + + `source` is carried for the record, not to change the decision. + """ + override = size_override(db_session, feed, namespace, key) + if override.size is not None: + basis = Basis.AUTO if override.source == SOURCE_AUTO else Basis.OPERATOR + logging.info( + "Routing %s to %s: %s%s override on feed %s", + dataset.stable_id, + override.size.value, + "locked " if override.locked else "", + basis.value, + feed.stable_id, + ) + return Routing(size=override.size, basis=basis) + + measure, how = measure_dataset(db_session, dataset, compression_ratio) + size = choose_size(measure, tiers) + + if measure is None: + logging.warning( + "No recorded size for %s (%s); routing to %s. Run the " + "rebuild_missing_dataset_files task to record them.", + dataset.stable_id, + how, + size.value, + ) + else: + logging.info( + "Routing %s to %s: %s bytes by %s, table [%s]", + dataset.stable_id, + size.value, + measure, + how, + describe(tiers), + ) + return Routing(size=size, basis=Basis.MEASURED) + + +def set_size_override( + db_session: "Session", feed, size: Size, namespace: str, key: str = "size" +) -> None: + """Record the size this feed should use from now on, so the next run starts there. + + Written as an object so the row says the builder set it rather than a person. Both + forms decide outright; the source is kept for the record. + + Written to `config_value_feed`, which has no foreign key to `feed` and requires both + `feed_id` and a NOT NULL `feed_stable_id`, so both are supplied. `updated_at` has a + server default on insert only and is set explicitly on the update path. + """ + from sqlalchemy.dialects.postgresql import insert + + from shared.database_gen.sqlacodegen_models import ConfigValueFeed + + now = datetime.now(timezone.utc) + value = {"size": size.value, "source": SOURCE_AUTO} + statement = ( + insert(ConfigValueFeed) + .values( + feed_id=feed.id, + feed_stable_id=feed.stable_id, + namespace=namespace, + key=key, + value=value, + updated_at=now, + ) + .on_conflict_do_update( + index_elements=["feed_id", "namespace", "key"], + set_={"value": value, "updated_at": now}, + ) + ) + db_session.execute(statement) + logging.info("Feed %s pinned to the %s worker", feed.stable_id, size.value) + + +def clear_size_override( + db_session: "Session", feed, namespace: str, key: str = "size" +) -> bool: + """Remove this feed's size override. Returns whether a row went. + + A locked override stays. Everything else goes, whoever wrote it: authorship records + who set the value last, it is not a claim on the row. + """ + from shared.database_gen.sqlacodegen_models import ConfigValueFeed + + if size_override(db_session, feed, namespace, key).locked: + return False + + db_session.query(ConfigValueFeed).filter( + ConfigValueFeed.feed_id == feed.id, + ConfigValueFeed.namespace == namespace, + ConfigValueFeed.key == key, + ).delete(synchronize_session=False) + logging.info("Cleared the size override on %s", feed.stable_id) + return True + + +def record_size_override( + db_session: "Session", + feed, + dataset, + size: Size, + *, + tiers: Sequence[Tier], + namespace: str, + key: str = "size", + compression_ratio: int = DEFAULT_COMPRESSION_RATIO, +) -> bool: + """Persist `size` for this feed, unless the measurement already routes there. + + An override that agrees with the measurement changes no decision, and carries a cost: + it pins the feed, so the feed stops following its own data. A feed that shrinks keeps + paying for a worker it no longer needs, and the row has to be found and removed by + hand before it can. Storing one only where it disagrees keeps the table to the feeds + that are genuinely exceptions, and lets a feed leave the list on its own. + + A locked override is left exactly as it is. This is the only place either direction + writes, so enforcing the lock here means no caller can move a locked feed by reaching + for the write path directly. + + Returns whether a row is now stored. + """ + existing = size_override(db_session, feed, namespace, key) + if existing.locked: + logging.info( + "Not moving %s to %s: it is locked to %s", + feed.stable_id, + size.value, + existing.size.value if existing.size else "an unreadable size", + ) + return True + + measured, _ = measure_dataset(db_session, dataset, compression_ratio) + if choose_size(measured, tiers) == size: + logging.info( + "Not overriding %s to %s: it is already measured there", + feed.stable_id, + size.value, + ) + clear_size_override(db_session, feed, namespace, key) + return False + + set_size_override(db_session, feed, size, namespace, key) + return True diff --git a/functions-python/helpers/task_execution/task_execution_tracker.py b/functions-python/helpers/task_execution/task_execution_tracker.py index 95407c747..27c4295d7 100644 --- a/functions-python/helpers/task_execution/task_execution_tracker.py +++ b/functions-python/helpers/task_execution/task_execution_tracker.py @@ -29,7 +29,7 @@ class in the future. ) tracker.start_run(total_count=5000, params={"validator_endpoint": "...", "env": "staging"}) - if not tracker.is_triggered(dataset_id): + if not tracker.is_handled(dataset_id): execute_workflow(...) tracker.mark_triggered(dataset_id, execution_ref=execution.name) @@ -48,16 +48,32 @@ class in the future. from datetime import datetime, timedelta, timezone from typing import Any, Optional +from sqlalchemy import and_, func, or_, text from sqlalchemy.dialects.postgresql import insert from sqlalchemy.orm import Session -from shared.database_gen.sqlacodegen_models import TaskExecutionLog, TaskRun +from shared.database_gen.sqlacodegen_models import ( + TaskExecutionAttempt, + TaskExecutionLog, + TaskRun, +) STATUS_IN_PROGRESS = "in_progress" STATUS_COMPLETED = "completed" STATUS_FAILED = "failed" STATUS_TRIGGERED = "triggered" +# An entity in any of these is spoken for: something has dispatched it, is working on +# it, or has finished. Only an absent row or a `failed` one is free to dispatch. +_HANDLED_STATUSES = (STATUS_TRIGGERED, STATUS_IN_PROGRESS, STATUS_COMPLETED) + +# How long a claim taken by `try_acquire` stays valid without a heartbeat. A worker +# killed mid-run (OOM, timeout) cannot release its own claim, so the claim has to +# expire on its own or the entity could never be retried. Keep this comfortably +# above the worker's own timeout, so an instance GCP has not finished killing can +# never have its work started a second time underneath it. +DEFAULT_LEASE_SECONDS = 1800 + # Cap on concatenated list fields in get_summary()'s metadata_summary. The two tables # hold every entity's own metadata regardless; this only bounds the summary's size. DEFAULT_METADATA_LIST_CAP = 200 @@ -157,8 +173,11 @@ def finish_run(self, status: str = STATUS_COMPLETED) -> None: def count_already_tracked(self, entity_ids: list[str]) -> int: """ - Return how many of the given entity_ids are already tracked for this run - (status triggered or completed). Useful in dry-run to preview skips. + Return how many of the given entity_ids are already handled for this run. + Useful in dry-run to preview skips. + + Counts the same statuses `is_handled` accepts, so a dry run predicts the number + of skips the real run performs. """ if not entity_ids: return 0 @@ -167,21 +186,23 @@ def count_already_tracked(self, entity_ids: list[str]) -> int: .filter( TaskExecutionLog.task_name == self.task_name, TaskExecutionLog.run_id == self.run_id, - TaskExecutionLog.status.in_([STATUS_TRIGGERED, STATUS_COMPLETED]), + TaskExecutionLog.status.in_(_HANDLED_STATUSES), TaskExecutionLog.entity_id.in_(entity_ids), ) .count() ) - def is_triggered(self, entity_id: Optional[str]) -> bool: - """ - Return True if an execution log entry already exists for this entity - with status triggered or completed (i.e. should not be re-triggered). + def is_handled(self, entity_id: Optional[str]) -> bool: + """True when something already has this entity: triggered, running, or done. + + The dispatchable states are the two that mean nobody has it: no row at all, and + `failed`. A failed entity is meant to be retried, which is why it is not counted + here. """ query = self.db_session.query(TaskExecutionLog).filter( TaskExecutionLog.task_name == self.task_name, TaskExecutionLog.run_id == self.run_id, - TaskExecutionLog.status.in_([STATUS_TRIGGERED, STATUS_COMPLETED]), + TaskExecutionLog.status.in_(_HANDLED_STATUSES), ) if entity_id is None: query = query.filter(TaskExecutionLog.entity_id.is_(None)) @@ -265,6 +286,232 @@ def mark_failed( ) self.db_session.flush() + # ------------------------------------------------------------------ + # Exclusive claims + # ------------------------------------------------------------------ + + def try_acquire( + self, + entity_id: Optional[str], + execution_ref: Optional[str] = None, + lease_seconds: int = DEFAULT_LEASE_SECONDS, + ) -> bool: + """Atomically claim this entity for work. True iff this caller now owns it. + + `mark_triggered` is an upsert that always writes, so two callers racing on the + same entity both "succeed" — use this instead wherever the work must not run + twice at once. Cloud Tasks delivers at least once, so the worker itself has to + claim, not just the caller that enqueued it. + + The guarantee is Postgres's: a single INSERT .. ON CONFLICT DO UPDATE .. WHERE + takes the row lock, and a conflicting statement re-evaluates the predicate + against the committed row. When the predicate is false nothing is written and + nothing is returned, so exactly one concurrent caller sees a row come back. + + A claim is granted when the entity is untracked, was merely `triggered` + (enqueued, not yet started), previously `failed`, or is `in_progress` with an + expired lease. A `completed` entity is never reclaimed — call + `release_for_retry` first to redo finished work. + """ + task_run_id = self._resolve_task_run_id() + expired = TaskExecutionLog.triggered_at < func.now() - text( + "make_interval(secs => :lease_seconds)" + ).bindparams(lease_seconds=lease_seconds) + + stmt = ( + insert(TaskExecutionLog) + .values( + task_run_id=task_run_id, + task_name=self.task_name, + entity_id=entity_id, + run_id=self.run_id, + status=STATUS_IN_PROGRESS, + execution_ref=execution_ref, + error_message=None, + triggered_at=func.now(), + ) + .on_conflict_do_update( + constraint="task_execution_log_task_name_entity_id_run_id_key", + set_={ + "status": STATUS_IN_PROGRESS, + "execution_ref": execution_ref, + "error_message": None, + "triggered_at": func.now(), + "completed_at": None, + }, + # Against the existing row, never the proposed one: this is what makes + # the statement a lock rather than an upsert. + where=or_( + TaskExecutionLog.status.in_([STATUS_TRIGGERED, STATUS_FAILED]), + and_( + TaskExecutionLog.status == STATUS_IN_PROGRESS, + expired, + ), + ), + ) + .returning(TaskExecutionLog.id) + ) + acquired = self.db_session.execute(stmt).scalar_one_or_none() is not None + self.db_session.flush() + logging.info( + "TaskExecutionTracker: claim on entity=%s run=%s/%s %s", + entity_id, + self.task_name, + self.run_id, + "granted" if acquired else "refused (held elsewhere)", + ) + return acquired + + def heartbeat( + self, entity_id: Optional[str], metadata: Optional[dict[str, Any]] = None + ) -> None: + """Extend a held claim, optionally recording where the work has got to. + + Without this a job that legitimately runs longer than the lease would have its + claim stolen while it is still working. Callers should heartbeat at whatever + cadence they already report progress at, rather than adding writes for it. + """ + values: dict[Any, Any] = {"triggered_at": func.now()} + if metadata is not None: + values[TaskExecutionLog.metadata_] = metadata + self._entity_query(entity_id).filter( + TaskExecutionLog.status == STATUS_IN_PROGRESS + ).update(values, synchronize_session=False) + self.db_session.flush() + + def release_for_retry(self, entity_id: Optional[str]) -> bool: + """Make a completed entity claimable again. True if one was released. + + The deliberate route back for work that is finished but has to be redone — + a forced regeneration, say. Only touches `completed`, so it can never wrench a + claim away from a worker that currently holds one. + """ + released = ( + self._entity_query(entity_id) + .filter(TaskExecutionLog.status == STATUS_COMPLETED) + .update( + {"status": STATUS_FAILED, "completed_at": None}, + synchronize_session=False, + ) + ) + self.db_session.flush() + return bool(released) + + # ------------------------------------------------------------------ + # Attempt history + # ------------------------------------------------------------------ + + def record_attempt( + self, + entity_id: Optional[str], + *, + status: str, + started_at, + variant: Optional[str] = None, + variant_basis: Optional[str] = None, + override_at_attempt: Optional[str] = None, + escalated_to: Optional[str] = None, + failure_kind: Optional[str] = None, + error: Optional[BaseException] = None, + metrics: Optional[dict[str, Any]] = None, + metadata: Optional[dict[str, Any]] = None, + ) -> None: + """Append one attempt to the history. Never updates an existing row. + + `task_execution_log` answers where an entity stands now and rewrites itself to do + it, so it cannot answer what has happened. This can: a retry cap needs a count, + and sizing decisions need a distribution. + + The exception is taken as an object rather than a string because + `str(MemoryError())` is empty - the class name is the only thing that identifies + it. + """ + finished_at = datetime.now(timezone.utc) + metrics = metrics or {} + self.db_session.add( + TaskExecutionAttempt( + task_name=self.task_name, + entity_id=entity_id, + run_id=self.run_id, + attempt=self.attempts_since_success(entity_id) + 1, + variant=variant, + variant_basis=variant_basis, + override_at_attempt=override_at_attempt, + escalated_to=escalated_to, + status=status, + failure_kind=failure_kind, + error_type=type(error).__name__ if error is not None else None, + error_message=str(error) if error is not None else None, + started_at=started_at, + finished_at=finished_at, + duration_ms=( + int((finished_at - started_at).total_seconds() * 1000) + if started_at + else None + ), + peak_rss_bytes=metrics.get("peak_rss_bytes"), + peak_vms_bytes=metrics.get("peak_vms_bytes"), + metadata_=metadata, + ) + ) + self.db_session.flush() + + def attempts_since_success(self, entity_id: Optional[str]) -> int: + """How many times this entity has been tried since it last succeeded. + + The count an escalation is bounded by. Resets on success, so a feed that works + for months and then outgrows its worker gets the full allowance again rather than + being stuck at whatever it accumulated years ago. + """ + last_success = ( + self._attempt_query(entity_id) + .filter(TaskExecutionAttempt.status == STATUS_COMPLETED) + .order_by(TaskExecutionAttempt.finished_at.desc()) + .first() + ) + query = self._attempt_query(entity_id) + if last_success is not None: + query = query.filter( + TaskExecutionAttempt.finished_at > last_success.finished_at + ) + return query.count() + + def last_failure(self, entity_id: Optional[str]) -> Optional[TaskExecutionAttempt]: + """The most recent failed attempt, or None.""" + return ( + self._attempt_query(entity_id) + .filter(TaskExecutionAttempt.status == STATUS_FAILED) + .order_by(TaskExecutionAttempt.finished_at.desc()) + .first() + ) + + def _attempt_query(self, entity_id: Optional[str]): + query = self.db_session.query(TaskExecutionAttempt).filter( + TaskExecutionAttempt.task_name == self.task_name, + TaskExecutionAttempt.run_id == self.run_id, + ) + if entity_id is None: + return query.filter(TaskExecutionAttempt.entity_id.is_(None)) + return query.filter(TaskExecutionAttempt.entity_id == entity_id) + + def get_entity(self, entity_id: Optional[str]) -> Optional[TaskExecutionLog]: + """The tracking row for one entity, or None when it is untracked. + + `is_handled` reduces the same row to a yes/no for dispatch. Read the row itself + when the states have to be told apart - reporting `preparing` separately from + `ready` and `failed`, for instance. + """ + return self._entity_query(entity_id).one_or_none() + + def _entity_query(self, entity_id: Optional[str]): + query = self.db_session.query(TaskExecutionLog).filter( + TaskExecutionLog.task_name == self.task_name, + TaskExecutionLog.run_id == self.run_id, + ) + if entity_id is None: + return query.filter(TaskExecutionLog.entity_id.is_(None)) + return query.filter(TaskExecutionLog.entity_id == entity_id) + # ------------------------------------------------------------------ # Reporting # ------------------------------------------------------------------ diff --git a/functions-python/helpers/tests/conftest.py b/functions-python/helpers/tests/conftest.py new file mode 100644 index 000000000..ac7732390 --- /dev/null +++ b/functions-python/helpers/tests/conftest.py @@ -0,0 +1,44 @@ +# +# MobilityData 2026 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Make this package importable the way the deployed functions import it. + +Every function that consumes these helpers gets them symlinked to `src/shared/helpers`, +so production code says `from shared.helpers.sizing import ...`. This package's own +tests run with the directory itself on the path, where the same module is just `sizing`. + +`function-python-setup.sh` cannot bridge that - it refuses to link a folder into its own +descendant - so the alias is registered here, as a submodule of the real `shared` +package rather than in place of it. Without this a wrong import path passes the test +suite and fails on deploy, which is not hypothetical: it happened. +""" + +import importlib +import sys +import types +from pathlib import Path + +_HELPERS = Path(__file__).resolve().parent.parent + +try: + _shared = importlib.import_module("shared") +except ImportError: # pragma: no cover - only when src/shared has not been linked + _shared = None + +if _shared is not None and "shared.helpers" not in sys.modules: + _alias = types.ModuleType("shared.helpers") + _alias.__path__ = [str(_HELPERS)] + sys.modules["shared.helpers"] = _alias + _shared.helpers = _alias diff --git a/functions-python/helpers/tests/test_helpers.py b/functions-python/helpers/tests/test_helpers.py index 104389613..b0b16fa93 100644 --- a/functions-python/helpers/tests/test_helpers.py +++ b/functions-python/helpers/tests/test_helpers.py @@ -268,6 +268,183 @@ def test_create_http_pmtiles_builder_task( self.assertEqual(args[5], "pmtiles-queue") +class TestParquetBuilderTask(unittest.TestCase): + """The enqueue is what the Operations API records `preparing` on the strength of.""" + + ENV = { + "PARQUET_BUILDER_QUEUE_S": "parquet-queue-s", + "PARQUET_BUILDER_QUEUE_M": "parquet-queue-m", + "PARQUET_BUILDER_QUEUE_L": "parquet-queue-l", + "PROJECT_ID": "my-project", + "GCP_REGION": "northamerica-northeast1", + "ENVIRONMENT": "dev", + } + + @patch.dict(os.environ, ENV, clear=False) + @patch("shared.common.gcp_utils.create_http_task_with_name") + @patch("google.cloud.tasks_v2.CloudTasksClient") + def test_sends_the_build_request(self, mock_client_cls, mock_create_task): + from utils import create_http_parquet_builder_task + import json + + mock_client_cls.return_value = MagicMock() + + create_http_parquet_builder_task("mdb-1210", "mdb-1210-202402121801") + + _, kwargs = mock_create_task.call_args + self.assertEqual( + json.loads(kwargs["body"].decode("utf-8")), + { + "feed_stable_id": "mdb-1210", + "dataset_stable_id": "mdb-1210-202402121801", + "force": False, + "retention_days": None, + # The routing decision travels with the task so the builder can record + # why it ran where it did, without re-deriving it. + "variant_basis": None, + "override": None, + }, + ) + self.assertEqual( + kwargs["url"], + "https://northamerica-northeast1-my-project.cloudfunctions.net/" + "parquet-builder-l-dev", + ) + self.assertEqual(kwargs["queue_name"], "parquet-queue-l") + + @patch.dict(os.environ, ENV, clear=False) + @patch("shared.common.gcp_utils.create_http_task_with_name") + @patch("google.cloud.tasks_v2.CloudTasksClient") + def test_the_task_is_unnamed_so_a_retry_is_never_deduplicated_away( + self, mock_client_cls, mock_create_task + ): + """A name stays reserved after the task completes, so a retry looks like a + duplicate and is dropped.""" + from utils import create_http_parquet_builder_task + + mock_client_cls.return_value = MagicMock() + + create_http_parquet_builder_task("mdb-1210", "mdb-1210-202402121801") + + _, kwargs = mock_create_task.call_args + self.assertIsNone(kwargs["task_name"]) + self.assertTrue(kwargs["raise_on_error"]) + + @patch.dict(os.environ, ENV, clear=False) + @patch("shared.common.gcp_utils.create_http_task_with_name") + @patch("google.cloud.tasks_v2.CloudTasksClient") + def test_a_failed_enqueue_reaches_the_caller( + self, mock_client_cls, mock_create_task + ): + """Silence here would leave the dataset reporting `preparing` forever.""" + from utils import create_http_parquet_builder_task + + mock_client_cls.return_value = MagicMock() + mock_create_task.side_effect = RuntimeError("queue unavailable") + + with self.assertRaises(RuntimeError): + create_http_parquet_builder_task("mdb-1210", "mdb-1210-202402121801") + + @patch.dict( + os.environ, + {k: v for k, v in ENV.items() if k != "PARQUET_BUILDER_QUEUE_L"}, + clear=True, + ) + @patch("shared.common.gcp_utils.create_http_task_with_name") + @patch("google.cloud.tasks_v2.CloudTasksClient") + def test_no_queue_configured_is_a_logged_no_op( + self, mock_client_cls, mock_create_task + ): + """Local runs have no queue; the builder is invoked by hand there.""" + from utils import create_http_parquet_builder_task + + mock_client_cls.return_value = MagicMock() + + create_http_parquet_builder_task("mdb-1210", "mdb-1210-202402121801") + + mock_create_task.assert_not_called() + + @patch.dict(os.environ, ENV, clear=False) + @patch("shared.common.gcp_utils.create_http_task_with_name") + @patch("google.cloud.tasks_v2.CloudTasksClient") + def test_each_size_has_its_own_queue_and_target( + self, mock_client_cls, mock_create_task + ): + """Separate queues are the point: large builds must not fill the dispatch + slots in front of small ones.""" + from utils import create_http_parquet_builder_task + from sizing import Size + + mock_client_cls.return_value = MagicMock() + + for size, queue, target in ( + (Size.S, "parquet-queue-s", "parquet-builder-s-dev"), + (Size.M, "parquet-queue-m", "parquet-builder-m-dev"), + (Size.L, "parquet-queue-l", "parquet-builder-l-dev"), + ): + with self.subTest(size=size): + mock_create_task.reset_mock() + create_http_parquet_builder_task("mdb-1210", "mdb-1210-1", size=size) + + _, kwargs = mock_create_task.call_args + self.assertEqual(kwargs["queue_name"], queue) + self.assertTrue(kwargs["url"].endswith(target), kwargs["url"]) + + @patch.dict(os.environ, ENV, clear=False) + @patch("shared.common.gcp_utils.create_http_task_with_name") + @patch("google.cloud.tasks_v2.CloudTasksClient") + def test_no_size_routes_to_the_largest(self, mock_client_cls, mock_create_task): + """A caller that has not measured anything must not land on a worker that + cannot finish the job.""" + from utils import create_http_parquet_builder_task + + mock_client_cls.return_value = MagicMock() + + create_http_parquet_builder_task("mdb-1210", "mdb-1210-1") + + _, kwargs = mock_create_task.call_args + self.assertEqual(kwargs["queue_name"], "parquet-queue-l") + + +class TestCreateHttpTaskWithName(unittest.TestCase): + def _call(self, client, **kwargs): + from shared.common.gcp_utils import create_http_task_with_name + from google.cloud import tasks_v2 + + return create_http_task_with_name( + client=client, + body=b"{}", + url="https://example.test/f", + project_id="my-project", + gcp_region="northamerica-northeast1", + queue_name="a-queue", + task_name=None, + task_time=None, + http_method=tasks_v2.HttpMethod.POST, + **kwargs, + ) + + def test_errors_are_swallowed_by_default(self): + client = MagicMock() + client.create_task.side_effect = RuntimeError("queue unavailable") + + self._call(client) + + def test_raise_on_error_propagates(self): + client = MagicMock() + client.create_task.side_effect = RuntimeError("queue unavailable") + + with self.assertRaises(RuntimeError): + self._call(client, raise_on_error=True) + + def test_an_existing_task_is_never_an_error(self): + """The task this one duplicates is already queued, which is the point.""" + client = MagicMock() + client.create_task.side_effect = RuntimeError("Requested entity already exists") + + self._call(client, raise_on_error=True) + + class TestDetectEncoding(unittest.TestCase): def test_utf8_encoding(self): with tempfile.NamedTemporaryFile(delete=False, mode="w", encoding="utf-8") as f: diff --git a/functions-python/helpers/tests/test_runtime_metrics.py b/functions-python/helpers/tests/test_runtime_metrics.py new file mode 100644 index 000000000..248616916 --- /dev/null +++ b/functions-python/helpers/tests/test_runtime_metrics.py @@ -0,0 +1,139 @@ +# +# MobilityData 2026 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""The metrics line is what container sizing is read off, so it has to carry RSS.""" + +import logging +import time +import unittest +from unittest.mock import MagicMock, patch + +from runtime_metrics import MemorySampler, _max_rss_bytes, track_metrics + + +def _run(metrics=("time", "memory", "cpu"), body=lambda: 42): + logger = MagicMock(spec=logging.Logger) + + @track_metrics(metrics=metrics) + def work(logger=None): + return body() + + result = work(logger=logger) + message = logger.info.call_args.args[2] if logger.info.called else "" + return result, message + + +class TestMemoryReporting(unittest.TestCase): + def test_rss_is_reported_alongside_tracemalloc(self): + """tracemalloc cannot see a C extension's heap; RSS is what the container bills.""" + _, message = _run() + + self.assertIn("memory:", message) + self.assertIn("rss:", message) + self.assertIn("process peak:", message) + # Address space, which is what RLIMIT_AS actually caps. Sizing a container from + # RSS alone under-provisions it. + self.assertIn("vms:", message) + + def test_the_existing_fields_are_unchanged(self): + """Log history stays comparable, so nothing is renamed or dropped.""" + _, message = _run() + + self.assertIn("time:", message) + self.assertIn("peak:", message) + self.assertIn("cpu:", message) + + def test_no_memory_metric_means_no_rss(self): + _, message = _run(metrics=("time",)) + + self.assertNotIn("rss:", message) + + def test_an_unreadable_rss_does_not_break_the_line(self): + """Metrics are a courtesy; they must not fail the function they measure.""" + with patch("runtime_metrics.psutil.Process") as process_cls: + process = process_cls.return_value + process.memory_info.side_effect = RuntimeError("no /proc") + process.cpu_percent.return_value = 0.0 + + result, message = _run() + + self.assertEqual(result, 42) + self.assertIn("memory:", message) + self.assertNotIn("rss:", message) + + def test_metrics_survive_an_exception_in_the_wrapped_function(self): + def boom(): + raise ValueError("nope") + + with self.assertRaises(ValueError): + _run(body=boom) + + +class TestMaxRss(unittest.TestCase): + def test_it_returns_a_plausible_byte_count(self): + """Linux reports ru_maxrss in KiB and macOS in bytes; both must come back as bytes.""" + peak = _max_rss_bytes() + + self.assertGreater(peak, 1024 * 1024, "under 1 MB means the unit is wrong") + self.assertLess(peak, 100 * 1024**3) + + +class TestMemorySampler(unittest.TestCase): + """A number named `peak` has to be one.""" + + def test_it_reports_a_peak_rather_than_the_end_state(self): + sampler = MemorySampler(interval=0.01).start() + held = [bytearray(16 * 1024 * 1024) for _ in range(8)] + time.sleep(0.1) + at_peak = sampler.peak_vms_bytes + del held + time.sleep(0.05) + metrics = sampler.stop() + + self.assertGreaterEqual(metrics["peak_vms_bytes"], at_peak) + self.assertGreater(metrics["peak_rss_bytes"], 0) + + def test_a_span_shorter_than_the_interval_still_reports(self): + """The first reading is taken on start, not on the first tick.""" + metrics = MemorySampler(interval=60).start().stop() + + self.assertGreater(metrics["peak_vms_bytes"], 0) + + def test_an_unreadable_process_is_not_fatal(self): + with patch( + "runtime_metrics.psutil.Process", side_effect=RuntimeError("no /proc") + ): + metrics = MemorySampler(interval=0.01).start().stop() + + self.assertIsNone(metrics["peak_vms_bytes"]) + self.assertIsNone(metrics["peak_rss_bytes"]) + + def test_the_thread_does_not_outlive_the_span(self): + sampler = MemorySampler(interval=0.01).start() + thread = sampler._thread + sampler.stop() + + self.assertTrue(thread.daemon) + self.assertFalse(thread.is_alive()) + + def test_it_works_as_a_context_manager(self): + with MemorySampler(interval=0.01) as sampler: + pass + + self.assertGreater(sampler.metrics()["peak_vms_bytes"], 0) + + +if __name__ == "__main__": + unittest.main() diff --git a/functions-python/helpers/tests/test_sizing.py b/functions-python/helpers/tests/test_sizing.py new file mode 100644 index 000000000..2407bf117 --- /dev/null +++ b/functions-python/helpers/tests/test_sizing.py @@ -0,0 +1,582 @@ +# +# MobilityData 2026 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Both ways of getting this wrong are silent, so both are pinned.""" + +import unittest + +from unittest.mock import MagicMock, patch + +from sizing import ( + LARGEST, + Basis, + FailureKind, + Size, + Tier, + choose_size, + demote, + describe, + first_known, + function_name, + classify_failure, + escalate, + fits_within, + measure_dataset, + queue_env_var, + record_size_override, + size_for_dataset, + size_override, + tier_for, +) + +CONFIG_VALUE = "sizing.get_config_value" + +MB = 1024**2 +GB = 1024**3 +TIERS = ( + Tier(size=Size.S, max_bytes=256 * MB), + Tier(size=Size.M, max_bytes=2 * GB), + Tier(size=Size.L, max_bytes=None), +) + + +class TestChooseSize(unittest.TestCase): + def test_the_smallest_band_is_the_smallest_worker(self): + self.assertEqual(choose_size(1, TIERS), Size.S) + self.assertEqual(choose_size(256 * MB - 1, TIERS), Size.S) + + def test_the_middle_band(self): + self.assertEqual(choose_size(256 * MB, TIERS), Size.M) + self.assertEqual(choose_size(2 * GB - 1, TIERS), Size.M) + + def test_each_threshold_is_exclusive(self): + """A measure exactly at a bound belongs to the rung above it.""" + self.assertEqual(choose_size(256 * MB, TIERS), Size.M) + self.assertEqual(choose_size(2 * GB, TIERS), Size.L) + + def test_over_the_last_threshold_is_the_large_worker(self): + self.assertEqual(choose_size(10 * GB, TIERS), Size.L) + + def test_an_unknown_measure_routes_to_the_largest(self): + """Size columns are nullable with no backfill, so None is an ordinary answer. + + Guessing small here turns a missing row into an OOM. + """ + self.assertEqual(choose_size(None, TIERS), LARGEST) + self.assertEqual(LARGEST, Size.L) + + def test_zero_is_not_treated_as_unknown(self): + """A genuinely empty dataset is small, not unmeasured.""" + self.assertEqual(choose_size(0, TIERS), Size.S) + + +class TestOverride(unittest.TestCase): + """A configured size is the answer; the measurement is not consulted.""" + + def test_it_can_raise_the_measured_size(self): + self.assertEqual(choose_size(1, TIERS, override=Size.L), Size.L) + + def test_it_can_skip_a_rung(self): + self.assertEqual(choose_size(10 * GB, TIERS, override=Size.S), Size.S) + + def test_it_can_lower_the_measured_size(self): + """Whoever set the config has looked at the feed; the table is a heuristic.""" + self.assertEqual(choose_size(10 * GB, TIERS, override=Size.M), Size.M) + + def test_it_applies_to_an_unknown_measure(self): + self.assertEqual(choose_size(None, TIERS, override=Size.M), Size.M) + + def test_matching_the_measured_size_changes_nothing(self): + self.assertEqual(choose_size(1, TIERS, override=Size.S), Size.S) + + +class TestParse(unittest.TestCase): + def test_accepts_the_written_forms(self): + for raw in ("l", "L", " l ", "l\n"): + with self.subTest(raw=raw): + self.assertEqual(Size.parse(raw), Size.L) + + def test_accepts_every_size(self): + for raw, expected in (("s", Size.S), ("m", Size.M), ("l", Size.L)): + with self.subTest(raw=raw): + self.assertEqual(Size.parse(raw), expected) + + def test_unset_is_none(self): + self.assertIsNone(Size.parse(None)) + + def test_nonsense_is_ignored_rather_than_fatal(self): + """Config is operator-entered; a typo must not fail the request.""" + for raw in ("xl", "", "3", {}, []): + with self.subTest(raw=raw): + self.assertIsNone(Size.parse(raw)) + + +class TestFirstKnown(unittest.TestCase): + def test_returns_the_first_usable_value(self): + self.assertEqual(first_known(None, 0, 42, 99), 42) + + def test_zero_and_none_are_not_measurements(self): + self.assertIsNone(first_known(None, 0, None)) + + def test_unparsable_values_are_skipped(self): + self.assertEqual(first_known("not a number", 7), 7) + + def test_nothing_known(self): + self.assertIsNone(first_known()) + + +class TestNaming(unittest.TestCase): + def test_queue_env_var(self): + self.assertEqual( + queue_env_var("PARQUET_BUILDER", Size.S), "PARQUET_BUILDER_QUEUE_S" + ) + self.assertEqual( + queue_env_var("PARQUET_BUILDER", Size.M), "PARQUET_BUILDER_QUEUE_M" + ) + self.assertEqual( + queue_env_var("PARQUET_BUILDER", Size.L), "PARQUET_BUILDER_QUEUE_L" + ) + + def test_function_name(self): + self.assertEqual( + function_name("parquet-builder", Size.L, "dev"), "parquet-builder-l-dev" + ) + + def test_describe_is_loggable(self): + self.assertEqual(describe(TIERS), f"s<{256 * MB}, m<{2 * GB}, l:rest") + + +class TestLadder(unittest.TestCase): + def test_the_largest_is_the_last_rung(self): + self.assertEqual(LARGEST, Size.L) + self.assertEqual([s.value for s in Size], ["s", "m", "l"]) + + +def _dataset(unzipped=None, zipped=None): + dataset = MagicMock() + dataset.id = "dataset-uuid" + dataset.stable_id = "mdb-1-202401010000" + dataset.unzipped_size_bytes = unzipped + dataset.zipped_size_bytes = zipped + return dataset + + +def _session(largest=None): + session = MagicMock() + session.query.return_value.filter.return_value.scalar.return_value = largest + return session + + +class TestMeasureDataset(unittest.TestCase): + """The size columns are nullable with no backfill, so every rung gets used.""" + + def test_the_largest_file_wins(self): + self.assertEqual( + measure_dataset(_session(largest=500), _dataset(unzipped=9999)), + (500, "largest file"), + ) + + def test_it_falls_back_to_the_unzipped_total(self): + """A sum, so it overestimates - which routes up, not down.""" + self.assertEqual( + measure_dataset(_session(), _dataset(unzipped=800)), (800, "unzipped total") + ) + + def test_it_falls_back_to_an_estimate_from_the_archive(self): + measure, basis = measure_dataset(_session(), _dataset(zipped=100)) + + self.assertEqual(measure, 500) + self.assertEqual(basis, "estimated from the archive") + + def test_the_ratio_is_caller_supplied(self): + measure, _ = measure_dataset(_session(), _dataset(zipped=100), 9) + + self.assertEqual(measure, 900) + + def test_nothing_recorded(self): + self.assertEqual(measure_dataset(_session(), _dataset()), (None, "unknown")) + + def test_a_failed_query_is_not_fatal(self): + session = MagicMock() + session.query.side_effect = RuntimeError("database gone") + + self.assertEqual(measure_dataset(session, _dataset()), (None, "unknown")) + + +class TestSizeForDataset(unittest.TestCase): + def _call(self, session, dataset, override=None): + feed = MagicMock() + feed.id = "feed-uuid" + feed.stable_id = "mdb-1" + with patch(CONFIG_VALUE, return_value=override): + return size_for_dataset( + session, feed, dataset, tiers=TIERS, namespace="demo" + ).size + + def test_it_routes_on_the_measurement(self): + self.assertEqual(self._call(_session(largest=1), _dataset()), Size.S) + self.assertEqual(self._call(_session(largest=3 * GB), _dataset()), Size.L) + + def test_an_unmeasurable_dataset_goes_to_the_largest(self): + self.assertEqual(self._call(_session(), _dataset()), LARGEST) + + def test_a_pin_decides_outright(self): + self.assertEqual( + self._call(_session(largest=3 * GB), _dataset(), override="s"), Size.S + ) + + def test_a_pin_skips_the_measurement(self): + """No point costing a query for an answer that cannot change the outcome.""" + session = MagicMock() + session.query.side_effect = AssertionError("the dataset was measured") + + self.assertEqual(self._call(session, _dataset(), override="m"), Size.M) + + def test_an_unreadable_pin_falls_back_to_the_measurement(self): + feed = MagicMock() + with patch(CONFIG_VALUE, side_effect=RuntimeError("no config")): + size = size_for_dataset( + _session(largest=1), feed, _dataset(), tiers=TIERS, namespace="demo" + ).size + + self.assertEqual(size, Size.S) + + def test_the_namespace_and_key_are_the_callers(self): + feed = MagicMock() + with patch(CONFIG_VALUE, return_value=None) as config: + size_for_dataset( + _session(largest=1), + feed, + _dataset(), + tiers=TIERS, + namespace="pmtiles_builder", + key="worker", + ) + + self.assertEqual( + config.call_args_list[0].args[:2], ("pmtiles_builder", "worker") + ) + + +class TestClassifyFailure(unittest.TestCase): + """Read the exception, never a stored message.""" + + def test_a_bare_memory_error(self): + """The case that cannot be classified from text: its message is empty.""" + error = MemoryError() + + self.assertEqual(str(error), "") + self.assertEqual(classify_failure(error), FailureKind.RESOURCE_MEMORY) + + def test_duckdb_out_of_memory(self): + class OutOfMemoryException(Exception): + pass + + error = OutOfMemoryException("Out of Memory Error: failed to allocate 256 KiB") + + self.assertEqual(classify_failure(error), FailureKind.RESOURCE_MEMORY) + + def test_enospc(self): + error = OSError(28, "No space left on device", "/tmp/in-memory/stop_times.txt") + + self.assertEqual(classify_failure(error), FailureKind.RESOURCE_DISK) + + def test_duckdb_spill_filling_the_volume(self): + """DuckDB reports a full spill directory as its own error, not an OSError.""" + + class IOException(Exception): + pass + + error = IOException("IO Error: Failed to write: No space left on device") + + self.assertEqual(classify_failure(error), FailureKind.RESOURCE_DISK) + + def test_anything_else(self): + for error in (ValueError("no tables"), FileNotFoundError("archive missing")): + with self.subTest(error=error): + self.assertEqual(classify_failure(error), FailureKind.OTHER) + + def test_only_resource_kinds_are_resource(self): + self.assertTrue(FailureKind.RESOURCE_MEMORY.is_resource) + self.assertTrue(FailureKind.RESOURCE_DISK.is_resource) + self.assertFalse(FailureKind.OTHER.is_resource) + + +class TestEscalate(unittest.TestCase): + """Every None is a stop, because the alternative is a queue loop.""" + + def test_it_moves_one_rung_up(self): + self.assertEqual(escalate(Size.S, TIERS, attempts=1, max_attempts=3), Size.M) + self.assertEqual(escalate(Size.M, TIERS, attempts=1, max_attempts=3), Size.L) + + def test_the_largest_rung_is_terminal(self): + self.assertIsNone(escalate(Size.L, TIERS, attempts=1, max_attempts=3)) + + def test_the_attempt_cap_is_terminal(self): + self.assertIsNone(escalate(Size.S, TIERS, attempts=3, max_attempts=3)) + self.assertIsNone(escalate(Size.S, TIERS, attempts=9, max_attempts=3)) + + def test_an_unknown_current_rung_is_terminal(self): + """Better to stop than to guess which worker just died.""" + self.assertIsNone(escalate(None, TIERS, attempts=0, max_attempts=3)) + + +class TestConfiguredSize(unittest.TestCase): + """An override decides; otherwise the measurement does. Nothing else.""" + + def _route(self, largest, value=None): + feed = MagicMock() + with patch(CONFIG_VALUE, return_value=value): + return size_for_dataset( + _session(largest=largest), + feed, + _dataset(), + tiers=TIERS, + namespace="demo", + ) + + @staticmethod + def _auto(size): + """What the builder writes for itself, as opposed to a person's bare value.""" + return {"size": size, "source": "auto"} + + @staticmethod + def _override(value): + """The parsed override for a stored value, without routing it.""" + with patch(CONFIG_VALUE, return_value=value): + return size_override(_session(), MagicMock(), "demo") + + @staticmethod + def _locked(size): + """What an operator writes to hold a size where it is.""" + return {"size": size, "locked": True} + + def test_an_override_beats_a_larger_measurement(self): + self.assertEqual(self._route(10 * GB, "s").size, Size.S) + self.assertEqual(self._route(10 * GB, self._auto("s")).size, Size.S) + + def test_an_override_beats_a_smaller_measurement(self): + self.assertEqual(self._route(1, "l").size, Size.L) + self.assertEqual(self._route(1, self._auto("l")).size, Size.L) + + def test_the_source_is_recorded_but_does_not_change_the_outcome(self): + """Both decide outright; `basis` says who, for the record.""" + by_person = self._route(10 * GB, "s") + by_builder = self._route(10 * GB, self._auto("s")) + + self.assertEqual(by_person.size, by_builder.size) + self.assertEqual(by_person.basis, Basis.OPERATOR) + self.assertEqual(by_builder.basis, Basis.AUTO) + + def test_an_unusable_value_falls_back_to_the_measurement(self): + for value in ("xl", {"size": "xl", "source": "auto"}, {}, 3): + with self.subTest(value=value): + routing = self._route(1, value) + + self.assertEqual(routing.size, Size.S) + self.assertEqual(routing.basis, Basis.MEASURED) + + def test_a_bare_value_is_never_locked(self): + """The bare form is the one a human types; locking has to be deliberate.""" + self.assertFalse(self._override("l").locked) + + def test_an_unusable_locked_flag_reads_as_unlocked(self): + """A typo must not freeze a feed where nothing can move it again.""" + for value in ("true", 1, "yes", None): + with self.subTest(locked=value): + self.assertFalse(self._override({"size": "l", "locked": value}).locked) + + def test_locked_is_read_from_the_object_form(self): + self.assertTrue(self._override({"size": "l", "locked": True}).locked) + + def test_nothing_configured_reads_as_measured(self): + routing = self._route(1) + + self.assertEqual(routing.size, Size.S) + self.assertEqual(routing.basis, Basis.MEASURED) + + def test_a_locked_override_routes_like_any_other(self): + """The lock decides whether the value may be moved, not where it sends a build.""" + routing = self._route(1, {"size": "l", "locked": True}) + + self.assertEqual(routing.size, Size.L) + # No `source` in the object: the builder always stamps its own, so this is a + # person's. + self.assertEqual(routing.basis, Basis.OPERATOR) + + def test_an_override_skips_the_measurement_entirely(self): + feed = MagicMock() + session = MagicMock() + session.query.side_effect = AssertionError("the dataset was measured") + + with patch(CONFIG_VALUE, return_value="m"): + routing = size_for_dataset( + session, feed, _dataset(), tiers=TIERS, namespace="demo" + ) + + self.assertEqual(routing.size, Size.M) + + +class TestRecordSizeOverride(unittest.TestCase): + """An override is for feeds the measurement gets wrong, and only those.""" + + @staticmethod + def _locked(size): + return {"size": size, "locked": True} + + def _record(self, largest, size, existing=None): + feed = MagicMock() + feed.id = "feed-uuid" + feed.stable_id = "mdb-1" + session = _session(largest=largest) + with patch(CONFIG_VALUE, return_value=existing): + stored = record_size_override( + session, feed, _dataset(), size, tiers=TIERS, namespace="demo" + ) + return stored, session + + def test_a_size_the_measurement_would_not_pick_is_stored(self): + stored, session = self._record(1, Size.M) + + self.assertTrue(stored) + session.execute.assert_called_once() + + def test_a_size_it_already_picks_is_not(self): + """Storing it changes no decision and costs the feed its ability to move.""" + stored, session = self._record(1, Size.S) + + self.assertFalse(stored) + session.execute.assert_not_called() + + def test_an_earlier_override_the_measurement_caught_up_with_is_cleared(self): + stored, session = self._record( + 1, Size.S, existing={"size": "m", "source": "auto"} + ) + + self.assertFalse(stored) + session.query.return_value.filter.return_value.delete.assert_called_once() + + def test_an_unlocked_pin_is_cleared_like_any_other_row(self): + """Who wrote a value is a record, not a claim on it. Only `locked` holds a row.""" + stored, session = self._record(1, Size.S, existing="s") + + self.assertFalse(stored) + session.query.return_value.filter.return_value.delete.assert_called_once() + + def test_an_unlocked_pin_is_raised_like_any_other_row(self): + stored, session = self._record(1, Size.M, existing="s") + + self.assertTrue(stored) + session.execute.assert_called_once() + + def test_a_locked_override_is_never_raised(self): + stored, session = self._record(1, Size.M, existing=self._locked("s")) + + self.assertTrue(stored, "the locked row is still the stored override") + session.execute.assert_not_called() + + def test_a_locked_override_is_never_cleared(self): + stored, session = self._record(1, Size.S, existing=self._locked("s")) + + session.query.return_value.filter.return_value.delete.assert_not_called() + self.assertTrue(stored) + + def test_an_unmeasurable_dataset_stores_nothing_for_the_largest(self): + """Unknown already routes to `l`, so an escalation there has nothing to say.""" + stored, session = self._record(None, LARGEST) + + self.assertFalse(stored) + session.execute.assert_not_called() + + +class TestDemote(unittest.TestCase): + """The mirror of escalate, and the smallest rung is where it stops.""" + + def test_it_moves_one_rung_down(self): + self.assertEqual(demote(Size.L, TIERS), Size.M) + self.assertEqual(demote(Size.M, TIERS), Size.S) + + def test_the_smallest_rung_is_terminal(self): + self.assertIsNone(demote(Size.S, TIERS)) + + def test_an_unknown_rung_is_terminal(self): + self.assertIsNone(demote(None, TIERS)) + self.assertIsNone(demote(Size.L, (Tier(size=Size.S, max_bytes=None),))) + + +class TestTierFor(unittest.TestCase): + def test_it_finds_the_row(self): + self.assertEqual(tier_for(Size.M, TIERS).max_bytes, 2 * GB) + + def test_a_size_not_in_the_table(self): + self.assertIsNone(tier_for(Size.L, (Tier(size=Size.S, max_bytes=None),))) + + +BUDGET = 1000 +CEILING = 500 +ONE_RUNG = Tier(size=Size.S, max_bytes=CEILING, vms_budget_bytes=BUDGET) + + +class TestFitsWithin(unittest.TestCase): + """Both axes, because a build can exhaust either one.""" + + def _fits(self, vms, largest, tier=ONE_RUNG, headroom=0.6): + return fits_within( + tier, + peak_vms_bytes=vms, + largest_member_bytes=largest, + headroom=headroom, + ) + + def test_comfortably_inside_both(self): + self.assertTrue(self._fits(500, 250)) + + def test_memory_at_the_threshold_does_not_fit(self): + self.assertFalse(self._fits(600, 250)) + self.assertTrue(self._fits(599, 250)) + + def test_disk_at_the_threshold_does_not_fit(self): + """The alternating-dataset case: memory is fine, the file is not.""" + self.assertFalse(self._fits(500, 300)) + self.assertTrue(self._fits(500, 299)) + + def test_a_missing_observation_reads_as_no(self): + """These columns were added after the table, so an older attempt says nothing - + and leaving a feed where it is costs money, while the other way costs builds.""" + self.assertFalse(self._fits(None, 250)) + self.assertFalse(self._fits(500, None)) + self.assertFalse(self._fits(0, 250)) + + def test_a_rung_with_no_budget_recorded_reads_as_no(self): + self.assertFalse( + self._fits(500, 250, tier=Tier(size=Size.S, max_bytes=CEILING)) + ) + + def test_the_catch_all_rung_is_bounded_only_by_memory(self): + """It accepts any file by definition, so the disk axis cannot refuse it.""" + catch_all = Tier(size=Size.L, max_bytes=None, vms_budget_bytes=BUDGET) + + self.assertTrue(self._fits(500, 10**12, tier=catch_all)) + self.assertFalse(self._fits(900, 10**12, tier=catch_all)) + + def test_the_headroom_must_be_a_fraction(self): + for bad in (0, -0.5, 1.5): + with self.subTest(headroom=bad): + with self.assertRaises(ValueError): + self._fits(500, 250, headroom=bad) + + +if __name__ == "__main__": + unittest.main() diff --git a/functions-python/helpers/tests/test_task_execution_tracker.py b/functions-python/helpers/tests/test_task_execution_tracker.py index d869051b2..4f84072fa 100644 --- a/functions-python/helpers/tests/test_task_execution_tracker.py +++ b/functions-python/helpers/tests/test_task_execution_tracker.py @@ -16,7 +16,7 @@ import unittest import uuid -from datetime import datetime, timezone +from datetime import datetime, timedelta, timezone from unittest.mock import MagicMock from task_execution.task_execution_tracker import ( @@ -81,16 +81,21 @@ def test_start_run_resets_status_to_in_progress_on_rerun(self): self.assertIn("completed_at", stmt_compiled) -class TestTaskExecutionTrackerIsTriggered(unittest.TestCase): - def test_returns_true_when_triggered_row_exists(self): +def _status_values(criterion): + """The statuses an `in_` criterion accepts, read off the compiled SQL.""" + compiled = criterion.compile(compile_kwargs={"literal_binds": True}) + return str(compiled) + + +class TestTaskExecutionTrackerIsHandled(unittest.TestCase): + def test_returns_true_when_a_row_exists(self): tracker, session = _make_tracker() existing_row = MagicMock() session.query.return_value.filter.return_value.filter.return_value.first.return_value = ( existing_row ) - result = tracker.is_triggered("ds-123") - self.assertTrue(result) + self.assertTrue(tracker.is_handled("ds-123")) def test_returns_false_when_no_row(self): tracker, session = _make_tracker() @@ -98,8 +103,7 @@ def test_returns_false_when_no_row(self): None ) - result = tracker.is_triggered("ds-999") - self.assertFalse(result) + self.assertFalse(tracker.is_handled("ds-999")) def test_handles_none_entity_id(self): tracker, session = _make_tracker() @@ -107,8 +111,57 @@ def test_handles_none_entity_id(self): None ) - result = tracker.is_triggered(None) - self.assertFalse(result) + self.assertFalse(tracker.is_handled(None)) + + def test_in_progress_counts_as_handled(self): + """The whole point: an entity being worked on is not free to dispatch. + + It used to answer the same as an entity nobody had ever asked for, which made + `in_progress` indistinguishable from absent through this method. + """ + tracker, session = _make_tracker() + session.query.return_value.filter.return_value.filter.return_value.first.return_value = ( + None + ) + + tracker.is_handled("ds-123") + + criteria = _status_values(session.query.return_value.filter.call_args.args[2]) + for status in (STATUS_TRIGGERED, STATUS_IN_PROGRESS, STATUS_COMPLETED): + self.assertIn(status, criteria) + + def test_failed_stays_dispatchable(self): + """A failed entity is meant to be retried, so it must not read as handled.""" + tracker, session = _make_tracker() + session.query.return_value.filter.return_value.filter.return_value.first.return_value = ( + None + ) + + tracker.is_handled("ds-123") + + criteria = _status_values(session.query.return_value.filter.call_args.args[2]) + self.assertNotIn(STATUS_FAILED, criteria) + + +class TestTaskExecutionTrackerCountAlreadyTracked(unittest.TestCase): + """The dry-run preview has to agree with what the real run skips.""" + + def test_it_counts_the_same_statuses_is_handled_accepts(self): + tracker, session = _make_tracker() + session.query.return_value.filter.return_value.count.return_value = 3 + + self.assertEqual(tracker.count_already_tracked(["a", "b"]), 3) + + criteria = _status_values(session.query.return_value.filter.call_args.args[2]) + for status in (STATUS_TRIGGERED, STATUS_IN_PROGRESS, STATUS_COMPLETED): + self.assertIn(status, criteria) + self.assertNotIn(STATUS_FAILED, criteria) + + def test_an_empty_list_costs_no_query(self): + tracker, session = _make_tracker() + + self.assertEqual(tracker.count_already_tracked([]), 0) + session.query.assert_not_called() class TestTaskExecutionTrackerMarkTriggered(unittest.TestCase): @@ -331,3 +384,193 @@ def test_drops_key_seen_with_inconsistent_types(self): ) self.assertNotIn("x", result) self.assertEqual(result["total"], 9) + + +# ---------------------------------------------------------------------------- +# Exclusive claims +# +# Against a real Postgres, deliberately: the guarantee `try_acquire` provides is a +# Postgres one - a conditional upsert re-evaluated against the committed row under a +# row lock - and a mocked session would assert the shape of the statement while +# proving nothing about the exclusion it exists for. +# ---------------------------------------------------------------------------- + +import threading # noqa: E402 + +import sqlalchemy # noqa: E402 +from sqlalchemy.orm import sessionmaker # noqa: E402 + +from task_execution.task_execution_tracker import DEFAULT_LEASE_SECONDS # noqa: E402 +from shared.database_gen.sqlacodegen_models import TaskExecutionLog # noqa: E402 + +TEST_DB_URL = "postgresql://postgres:postgres@localhost:54320/MobilityDatabaseTest" +CLAIM_TASK = "test_try_acquire" + + +def _engine_or_skip(): + try: + engine = sqlalchemy.create_engine(TEST_DB_URL) + with engine.connect(): + pass + return engine + except Exception as error: # pragma: no cover - depends on the local environment + raise unittest.SkipTest(f"test database unavailable: {error}") + + +class TestTryAcquire(unittest.TestCase): + """One worker at a time, and a claim that cannot be lost forever.""" + + @classmethod + def setUpClass(cls): + cls.engine = _engine_or_skip() + cls.Session = sessionmaker(bind=cls.engine) + + def setUp(self): + self.entity = f"dataset-{uuid.uuid4()}" + self.run_id = f"v{uuid.uuid4()}" + self.session = self.Session() + self.tracker = self._tracker(self.session) + self.tracker.start_run() + self.session.commit() + + def tearDown(self): + self.session.query(TaskExecutionLog).filter( + TaskExecutionLog.task_name == CLAIM_TASK, + TaskExecutionLog.run_id == self.run_id, + ).delete(synchronize_session=False) + self.session.commit() + self.session.close() + + def _tracker(self, session): + return TaskExecutionTracker( + task_name=CLAIM_TASK, run_id=self.run_id, db_session=session + ) + + def _row(self): + return self.tracker.get_entity(self.entity) + + def _set_status(self, status, triggered_at=None): + values = {"status": status} + if triggered_at is not None: + values["triggered_at"] = triggered_at + self.session.query(TaskExecutionLog).filter( + TaskExecutionLog.task_name == CLAIM_TASK, + TaskExecutionLog.run_id == self.run_id, + TaskExecutionLog.entity_id == self.entity, + ).update(values, synchronize_session=False) + self.session.commit() + + def test_claims_an_untracked_entity(self): + self.assertTrue(self.tracker.try_acquire(self.entity)) + self.session.commit() + self.assertEqual(self._row().status, STATUS_IN_PROGRESS) + + def test_refuses_a_second_claim_while_held(self): + self.assertTrue(self.tracker.try_acquire(self.entity)) + self.session.commit() + + other = self.Session() + try: + self.assertFalse(self._tracker(other).try_acquire(self.entity)) + other.commit() + finally: + other.close() + + def test_takes_over_from_a_merely_triggered_row(self): + """The API records the enqueue; the worker that runs it must still claim.""" + self.tracker.mark_triggered(self.entity) + self.session.commit() + + self.assertTrue(self.tracker.try_acquire(self.entity)) + + def test_reclaims_a_failed_entity(self): + self.tracker.try_acquire(self.entity) + self.tracker.mark_failed(self.entity, error_message="boom") + self.session.commit() + + self.assertTrue(self.tracker.try_acquire(self.entity)) + self.session.commit() + self.assertIsNone( + self._row().error_message, "a retry starts without the old error" + ) + + def test_never_reclaims_a_completed_entity(self): + self.tracker.try_acquire(self.entity) + self.tracker.mark_completed(self.entity, metadata={"base_url": "x"}) + self.session.commit() + + self.assertFalse(self.tracker.try_acquire(self.entity)) + + def test_release_for_retry_reopens_a_completed_entity(self): + self.tracker.try_acquire(self.entity) + self.tracker.mark_completed(self.entity) + self.session.commit() + + self.assertTrue(self.tracker.release_for_retry(self.entity)) + self.assertTrue(self.tracker.try_acquire(self.entity)) + + def test_release_for_retry_cannot_steal_a_running_claim(self): + self.tracker.try_acquire(self.entity) + self.session.commit() + + self.assertFalse(self.tracker.release_for_retry(self.entity)) + self.session.commit() + self.assertEqual(self._row().status, STATUS_IN_PROGRESS) + + def test_reclaims_a_claim_whose_lease_expired(self): + """A worker killed mid-run cannot release its own claim.""" + self.tracker.try_acquire(self.entity) + self.session.commit() + self._set_status( + STATUS_IN_PROGRESS, + triggered_at=datetime.now(timezone.utc) + - timedelta(seconds=DEFAULT_LEASE_SECONDS + 60), + ) + + self.assertTrue(self.tracker.try_acquire(self.entity)) + + def test_a_heartbeat_keeps_a_long_build_from_losing_its_claim(self): + self.tracker.try_acquire(self.entity) + self.session.commit() + self._set_status( + STATUS_IN_PROGRESS, + triggered_at=datetime.now(timezone.utc) + - timedelta(seconds=DEFAULT_LEASE_SECONDS + 60), + ) + + self.tracker.heartbeat(self.entity, metadata={"phase": "convert"}) + self.session.commit() + + self.assertFalse( + self.tracker.try_acquire(self.entity), + "a heartbeat must renew the lease, not merely record progress", + ) + self.assertEqual(self._row().metadata_, {"phase": "convert"}) + + def test_exactly_one_of_two_racing_workers_wins(self): + """The case the whole mechanism exists for.""" + results = [] + barrier = threading.Barrier(2) + + def claim(): + session = self.Session() + try: + barrier.wait(timeout=10) + won = self._tracker(session).try_acquire(self.entity) + session.commit() + results.append(won) + except Exception as error: # pragma: no cover - surfaced via the assert + results.append(error) + finally: + session.close() + + threads = [threading.Thread(target=claim) for _ in range(2)] + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=20) + + self.assertEqual( + results.count(True), 1, f"exactly one worker may win, got {results}" + ) + self.assertEqual(results.count(False), 1, f"got {results}") diff --git a/functions-python/helpers/utils.py b/functions-python/helpers/utils.py index 259cf06dd..8d9f53dc3 100644 --- a/functions-python/helpers/utils.py +++ b/functions-python/helpers/utils.py @@ -21,7 +21,10 @@ import urllib3.exceptions from datetime import date, datetime, timezone from logging import Logger -from typing import Optional +from typing import Optional, TYPE_CHECKING + +if TYPE_CHECKING: # pragma: no cover + from shared.helpers.sizing import Size import requests import urllib3 @@ -535,6 +538,95 @@ def create_http_pmtiles_builder_task( ) +PARQUET_BUILDER_BASE = "parquet-builder" +PARQUET_QUEUE_PREFIX = "PARQUET_BUILDER" + + +def create_http_parquet_builder_task( + feed_stable_id: str, + dataset_stable_id: str, + force: bool = False, + retention_days: Optional[int] = None, + size: Optional["Size"] = None, + variant_basis: Optional[str] = None, + override: Optional[str] = None, +) -> None: + """ + Create a task to render a dataset as Parquet, on the worker sized for it. + + Raises if the task could not be created: the caller marks the dataset as queued on + the strength of this call. + + `size` selects both the queue and the target function. Each size has its own queue + so a run of large builds cannot fill the dispatch slots in front of the small ones. + Defaults to the largest, so a caller that has not measured anything is never routed + to a worker that cannot finish the job. + + The task is unnamed. A name after the dataset stays reserved for about an hour after + the task completes, so a retry after a failed conversion would be dropped as a + duplicate. Deduplication is the builder's database claim. + """ + from google.cloud import tasks_v2 + from google.protobuf import timestamp_pb2 + from shared.common.gcp_utils import create_http_task_with_name + from shared.helpers.sizing import LARGEST, function_name, queue_env_var + import json + + size = size or LARGEST + + client = tasks_v2.CloudTasksClient() + body = json.dumps( + { + "feed_stable_id": feed_stable_id, + "dataset_stable_id": dataset_stable_id, + "force": force, + # Omitted when the caller did not ask, so the builder applies its default. + "retention_days": retention_days, + # Why this size was chosen, carried so the builder can record the decision + # alongside the attempt rather than having to re-derive it. + "variant_basis": variant_basis, + "override": override, + } + ).encode() + queue_env = queue_env_var(PARQUET_QUEUE_PREFIX, size) + queue_name = os.getenv(queue_env) + project_id = os.getenv("PROJECT_ID") + gcp_region = os.getenv("GCP_REGION") + gcp_env = os.getenv("ENVIRONMENT") + + if not queue_name: + logging.warning( + "%s is not set; skipping parquet build task for %s", + queue_env, + dataset_stable_id, + ) + return + + proto_time = timestamp_pb2.Timestamp() + proto_time.GetCurrentTime() + target = function_name(PARQUET_BUILDER_BASE, size, gcp_env) + + logging.info( + "Enqueuing parquet build for %s on the %s worker (%s)", + dataset_stable_id, + size.value, + target, + ) + create_http_task_with_name( + client=client, + body=body, + url=f"https://{gcp_region}-{project_id}.cloudfunctions.net/{target}", + project_id=project_id, + gcp_region=gcp_region, + queue_name=queue_name, + task_name=None, + task_time=proto_time, + http_method=tasks_v2.HttpMethod.POST, + timeout_s=1800, + raise_on_error=True, + ) + + def create_http_gtfs_datasets_comparer_task( feed_stable_id: str, base_dataset_stable_id: str, diff --git a/functions-python/helpers/validation_report/validation_report_update.py b/functions-python/helpers/validation_report/validation_report_update.py index f5cf87872..06d0aa99e 100644 --- a/functions-python/helpers/validation_report/validation_report_update.py +++ b/functions-python/helpers/validation_report/validation_report_update.py @@ -57,8 +57,9 @@ def execute_workflows( :param bypass_db_update: Whether to bypass the database update :param reports_bucket_name: The name of the bucket where the reports are stored :param tracker: Optional TaskExecutionTracker for idempotent execution tracking. - When provided, datasets already in triggered/completed state are skipped - and newly triggered datasets are recorded. + When provided, datasets something already has in hand - triggered, + running or completed - are skipped, and newly triggered datasets are + recorded. A failed dataset is not skipped: it is meant to be retried. :return: List of dataset stable ids for which the workflow was executed """ project_id = f"mobility-feeds-{env}" @@ -69,7 +70,7 @@ def execute_workflows( count = 0 logging.info(f"Executing workflow for {len(latest_datasets)} datasets") for feed_id, dataset_id in latest_datasets: - if tracker and tracker.is_triggered(dataset_id): + if tracker and tracker.is_handled(dataset_id): logging.info(f"Skipping already triggered dataset {feed_id}/{dataset_id}") continue try: diff --git a/functions-python/operations_api/.openapi-generator/FILES b/functions-python/operations_api/.openapi-generator/FILES index e317b1b7f..c30bf8a8e 100644 --- a/functions-python/operations_api/.openapi-generator/FILES +++ b/functions-python/operations_api/.openapi-generator/FILES @@ -5,6 +5,8 @@ src/feeds_gen/apis/licenses_api.py src/feeds_gen/apis/licenses_api_base.py src/feeds_gen/apis/operations_api.py src/feeds_gen/apis/operations_api_base.py +src/feeds_gen/apis/parquet_api.py +src/feeds_gen/apis/parquet_api_base.py src/feeds_gen/apis/users_api.py src/feeds_gen/apis/users_api_base.py src/feeds_gen/main.py @@ -22,6 +24,8 @@ src/feeds_gen/models/early_access_program_report.py src/feeds_gen/models/early_access_report_row.py src/feeds_gen/models/early_access_report_summary.py src/feeds_gen/models/early_access_report_summary_by_source.py +src/feeds_gen/models/execution_attempt.py +src/feeds_gen/models/execution_attempts_response.py src/feeds_gen/models/external_id.py src/feeds_gen/models/extra_models.py src/feeds_gen/models/feature_flag_assignment.py @@ -71,6 +75,8 @@ src/feeds_gen/models/operation_feed.py src/feeds_gen/models/operation_gtfs_feed.py src/feeds_gen/models/operation_gtfs_rt_feed.py src/feeds_gen/models/operation_user_profile.py +src/feeds_gen/models/parquet_dataset_state.py +src/feeds_gen/models/parquet_generate_request.py src/feeds_gen/models/propagate_license_affected_feed.py src/feeds_gen/models/propagate_license_request.py src/feeds_gen/models/propagate_license_response.py diff --git a/functions-python/operations_api/README.md b/functions-python/operations_api/README.md index ea1f4cec8..4a6ea9363 100644 --- a/functions-python/operations_api/README.md +++ b/functions-python/operations_api/README.md @@ -4,6 +4,13 @@ The Operations API OpenAPI schema lives at `../../docs/OperationsAPI.yaml`. > Note: generated server stubs are created from the schema. Do not edit generated code under `src/feeds_gen/`; put implementation under `src/feeds_operations/impl/`. +## Consumers + +The API is internal. Its known client is the operations web app, in a separate repo. +For the Parquet browse flow it drives - the call sequence, authentication, polling and +what it reads from the storage bucket - see +[docs/parquet-feed-browsing.md](../../docs/parquet-feed-browsing.md). + # Function configuration The function is configured using the following environment variables: - `FEEDS_DATABASE_URL`: The URL of the feeds database. @@ -27,6 +34,14 @@ docker compose --env-file ./config/.env.local up -d liquibase-test ``` ./scripts/api-operations-update-schema.sh ``` +- Generate a dataset's Parquet locally, with no GCP access, to work on the + `/v1/operations/{gtfs_feeds,gtfs_datasets}/{id}/parquet` endpoints or the viewer that + consumes them: +``` +./scripts/parquet-generate-local.sh [--serve] +``` + See [parquet_builder](../parquet_builder/README.md) for those endpoints and the full + local testing guide. ## Development process diff --git a/functions-python/operations_api/src/feeds_operations/impl/feeds_operations_impl.py b/functions-python/operations_api/src/feeds_operations/impl/feeds_operations_impl.py index c4520249d..31c462abf 100644 --- a/functions-python/operations_api/src/feeds_operations/impl/feeds_operations_impl.py +++ b/functions-python/operations_api/src/feeds_operations/impl/feeds_operations_impl.py @@ -29,6 +29,8 @@ from starlette.responses import Response from feeds_gen.models.data_type import DataType +from feeds_gen.models.execution_attempt import ExecutionAttempt +from feeds_gen.models.execution_attempts_response import ExecutionAttemptsResponse from feeds_gen.models.get_feeds200_response import GetFeeds200Response from feeds_gen.models.gtfs_feed_availability_response import ( GtfsFeedAvailabilityResponse, @@ -57,10 +59,14 @@ from shared.database_gen.sqlacodegen_models import ( Gtfsfeed, GtfsFeedAvailabilityCheck, + Gtfsdataset, t_feedsearch, Feed, Gtfsrealtimefeed, + TaskExecutionAttempt, ) +from shared.common.config_reader import get_config_value +from shared.helpers.parquet_policy import SIZE_CONFIG_KEY, SIZE_CONFIG_NAMESPACE from shared.common.license_utils import assign_license_by_url, propagate_license_by_url from shared.common.gcp_utils import create_web_revalidation_task from shared.db_models.gtfs_feed_availability_check_impl import ( @@ -136,6 +142,64 @@ def _strip_derived_fields(dumped: dict) -> dict: return dumped +def _current_size_override(db_session, feed_id) -> tuple[Optional[str], bool]: + """The size configured for this feed and whether it is held, if any. + + An override decides which worker runs the build, so reporting it answers the question + this page exists for. Who set it is in `variant_basis` on each attempt. + + The lock matters more than it looks: a locked feed never escalates, so a run of + failed attempts with no `escalated_to` has no other explanation visible from here. + """ + value = get_config_value( + SIZE_CONFIG_NAMESPACE, SIZE_CONFIG_KEY, feed_id=feed_id, db_session=db_session + ) + if isinstance(value, dict): + return value.get("size"), value.get("locked") is True + return (value if isinstance(value, str) else None), False + + +def _as_datetime(value) -> Optional[datetime]: + """A query timestamp, from whatever the router handed over. + + The generated router types these as `str`, so an ISO string is the usual case; an + unparsable one is rejected rather than silently ignored, since a filter that quietly + does nothing is worse than an error. + """ + if not value: + return None + if isinstance(value, datetime): + return value + try: + return datetime.fromisoformat(str(value).replace("Z", "+00:00")) + except ValueError: + raise HTTPException( + status_code=400, detail=f"Invalid date-time: {value}" + ) from None + + +def _execution_attempt(row) -> ExecutionAttempt: + """One stored attempt, as the API reports it.""" + return ExecutionAttempt( + task_name=row.task_name, + dataset_stable_id=row.entity_id, + attempt=row.attempt, + variant=row.variant, + variant_basis=row.variant_basis, + override_at_attempt=row.override_at_attempt, + escalated_to=row.escalated_to, + status=row.status, + failure_kind=row.failure_kind, + error_type=row.error_type, + error_message=row.error_message, + started_at=row.started_at, + finished_at=row.finished_at, + duration_ms=row.duration_ms, + peak_rss_bytes=row.peak_rss_bytes, + peak_vms_bytes=row.peak_vms_bytes, + ) + + class OperationsApiImpl(BaseOperationsApi): """Implementation of the operations API.""" @@ -299,6 +363,75 @@ def get_gtfs_feed_availability( checks=[GtfsFeedAvailabilityCheckImpl.from_orm(c) for c in checks], ) + @with_db_session + def get_gtfs_feed_execution_attempts( + self, + id: Annotated[ + StrictStr, Field(description="The feed ID of the requested feed.") + ], + task_name: Optional[str] = None, + status: Optional[str] = None, + failure_kind: Optional[str] = None, + var_from: Optional[str] = None, + to: Optional[str] = None, + limit: Optional[int] = 100, + offset: Optional[int] = 0, + db_session: Session = None, + ) -> ExecutionAttemptsResponse: + """What background tasks have tried to do with this feed, newest first. + + Attempts are recorded against a dataset, so this spans every dataset of the feed: + an admin looking at a feed wants its whole history, not one version's. + """ + gtfs_feed = ( + db_session.query(Gtfsfeed).filter(Gtfsfeed.stable_id == id).one_or_none() + ) + if gtfs_feed is None: + raise HTTPException(status_code=404, detail="GTFS feed not found") + + # The dataset stable ids are the entity ids the attempts were recorded under. + dataset_ids = [ + row[0] + for row in db_session.query(Gtfsdataset.stable_id) + .filter(Gtfsdataset.feed_id == gtfs_feed.id) + .all() + ] + + query = db_session.query(TaskExecutionAttempt).filter( + TaskExecutionAttempt.entity_id.in_(dataset_ids or [""]) + ) + if task_name: + query = query.filter(TaskExecutionAttempt.task_name == task_name) + if status: + query = query.filter(TaskExecutionAttempt.status == status) + if failure_kind: + query = query.filter(TaskExecutionAttempt.failure_kind == failure_kind) + # On started_at, not finished_at: a long build should be found by when it was + # triggered, which is what someone looking for "what ran on Tuesday" means. + start = _as_datetime(var_from) + if start is not None: + query = query.filter(TaskExecutionAttempt.started_at >= start) + end = _as_datetime(to) + if end is not None: + query = query.filter(TaskExecutionAttempt.started_at <= end) + + total = query.count() + attempts = ( + query.order_by(TaskExecutionAttempt.finished_at.desc()) + .offset(offset) + .limit(limit) + .all() + ) + + size_override, size_locked = _current_size_override(db_session, gtfs_feed.id) + return ExecutionAttemptsResponse( + feed_stable_id=id, + current_size_override=size_override, + current_size_locked=size_locked, + total=total, + attempts=[_execution_attempt(row) for row in attempts], + ) + @with_db_session def get_gtfs_rt_feed( self, diff --git a/functions-python/operations_api/src/feeds_operations/impl/parquet_api_impl.py b/functions-python/operations_api/src/feeds_operations/impl/parquet_api_impl.py new file mode 100644 index 000000000..e31150745 --- /dev/null +++ b/functions-python/operations_api/src/feeds_operations/impl/parquet_api_impl.py @@ -0,0 +1,341 @@ +# +# MobilityData 2026 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Reporting and starting the Parquet rendering of a GTFS dataset. + +The shape of the responses here is not free: a browser viewer polls the GET roughly +twice a second and renders every state itself, so the fields are named as it names its +own load reports and `absent` is an ordinary answer rather than a 404. A 404 means the +feed or dataset does not exist, which is a different thing entirely and the interface +has to be able to tell them apart. + +The GET never starts work. Starting is the POST, which a client calls once. +""" + +import logging +from datetime import datetime, timezone +from typing import Optional + +from fastapi import HTTPException +from pydantic import StrictStr +from sqlalchemy.orm import Session + +from feeds_gen.apis.parquet_api_base import BaseParquetApi +from feeds_gen.models.parquet_dataset_state import ParquetDatasetState +from feeds_gen.models.parquet_generate_request import ParquetGenerateRequest +from shared.database.database import with_db_session +from shared.database_gen.sqlacodegen_models import Gtfsdataset, Gtfsfeed +from shared.helpers.parquet_policy import ( + COMPRESSION_RATIO, + SIZE_CONFIG_KEY, + SIZE_CONFIG_NAMESPACE, + SIZE_TIERS, +) +from shared.helpers.sizing import Basis, Routing, size_for_dataset +from shared.helpers.task_execution.task_execution_tracker import ( + STATUS_COMPLETED, + STATUS_FAILED, + TaskExecutionTracker, +) +from shared.helpers.utils import create_http_parquet_builder_task + +# Must match functions-python/parquet_builder/src/converter.py. Bumping the converter +# invalidates previously written artifacts by moving them to a different run. +TASK_NAME = "parquet_generation" +PARQUET_CONVERTER_VERSION = "2" + +STATUS_ABSENT = "absent" +STATUS_PREPARING = "preparing" +STATUS_READY = "ready" +STATUS_FAILED_STATE = "failed" + + +class ParquetApiImpl(BaseParquetApi): + """Implementation of the Parquet API.""" + + # ------------------------------------------------------------------ + # Generated entry points + # ------------------------------------------------------------------ + + def get_gtfs_feed_parquet(self, id: StrictStr) -> ParquetDatasetState: + return self.handle_status(feed_stable_id=id) + + def get_gtfs_dataset_parquet(self, id: StrictStr) -> ParquetDatasetState: + return self.handle_status(dataset_stable_id=id) + + def generate_gtfs_feed_parquet( + self, + id: StrictStr, + parquet_generate_request: Optional[ParquetGenerateRequest] = None, + ) -> ParquetDatasetState: + return self.handle_generate( + feed_stable_id=id, + force=_force(parquet_generate_request), + retention_days=_retention_days(parquet_generate_request), + ) + + def generate_gtfs_dataset_parquet( + self, + id: StrictStr, + parquet_generate_request: Optional[ParquetGenerateRequest] = None, + ) -> ParquetDatasetState: + return self.handle_generate( + dataset_stable_id=id, + force=_force(parquet_generate_request), + retention_days=_retention_days(parquet_generate_request), + ) + + # ------------------------------------------------------------------ + # Behaviour + # ------------------------------------------------------------------ + + @with_db_session + def handle_status( + self, + feed_stable_id: Optional[str] = None, + dataset_stable_id: Optional[str] = None, + db_session: Session = None, + ) -> ParquetDatasetState: + feed, dataset = _resolve(db_session, feed_stable_id, dataset_stable_id) + return _state_of(db_session, feed, dataset) + + @with_db_session + def handle_generate( + self, + feed_stable_id: Optional[str] = None, + dataset_stable_id: Optional[str] = None, + force: bool = False, + retention_days: Optional[int] = None, + db_session: Session = None, + ) -> ParquetDatasetState: + feed, dataset = _resolve(db_session, feed_stable_id, dataset_stable_id) + state = _state_of(db_session, feed, dataset) + + # A build already under way is reported, not duplicated - the viewer calls this + # once on first `absent`, but nothing stops two operators clicking at once, and + # `force` must not be able to interrupt work in flight either. + if state.status == STATUS_PREPARING: + logging.info( + "Parquet build already in progress for %s; not enqueuing another", + dataset.stable_id, + ) + return state + + if state.status == STATUS_READY and not force: + return state + + # Written before the dispatch, so the next poll reads `preparing` rather than + # `absent`. It has to come first: `mark_triggered` is an unconditional upsert, + # and Cloud Tasks can deliver before this request finishes, so writing it + # afterwards could push a row the worker already moved to `in_progress` back to + # the claimable `triggered`. A status marker, not a lock - the worker still + # claims the dataset through `try_acquire`. + tracker = TaskExecutionTracker( + task_name=TASK_NAME, + run_id=PARQUET_CONVERTER_VERSION, + db_session=db_session, + ) + tracker.start_run(params={"converter_version": PARQUET_CONVERTER_VERSION}) + tracker.mark_triggered( + dataset.stable_id, + metadata={"phase": "start", "done": 0, "total": 0, "detail": ""}, + ) + db_session.commit() + + try: + create_http_parquet_builder_task( + feed.stable_id, + dataset.stable_id, + force=force, + retention_days=retention_days, + **_enqueue_sizing(db_session, feed, dataset), + ) + except Exception as error: + logging.error( + "Failed to enqueue Parquet build for %s: %s", dataset.stable_id, error + ) + # The marker describes work that will never arrive. Left as `triggered` it + # reports `preparing`, the one state this endpoint refuses to re-trigger. + try: + tracker.mark_failed( + dataset.stable_id, + error_message=f"Could not start the conversion: {error}", + ) + db_session.commit() + except Exception: + logging.exception( + "Could not clear the trigger marker for %s", dataset.stable_id + ) + db_session.rollback() + raise HTTPException( + status_code=500, detail=f"Could not start the conversion: {error}" + ) + + return ParquetDatasetState( + status=STATUS_PREPARING, + feed_stable_id=feed.stable_id, + dataset_stable_id=dataset.stable_id, + phase="start", + done=0, + total=0, + detail="", + ) + + +def _force(request: Optional[ParquetGenerateRequest]) -> bool: + return bool(request.force) if request and request.force is not None else False + + +def _retention_days(request: Optional[ParquetGenerateRequest]) -> Optional[int]: + """None when the caller did not ask for one. + + Passed through rather than defaulted here: the schema declares the bounds but no + default, so the builder stays the single place that decides how long a set lives. + """ + return request.retention_days if request else None + + +def _resolve( + db_session: Session, + feed_stable_id: Optional[str], + dataset_stable_id: Optional[str], +): + """Find the feed and dataset a request addresses, or 404. + + A feed resolves to its latest dataset, so a caller holding only a feed id never has + to know about datasets. + """ + if dataset_stable_id: + dataset = ( + db_session.query(Gtfsdataset) + .filter(Gtfsdataset.stable_id == dataset_stable_id) + .one_or_none() + ) + if dataset is None: + raise HTTPException(status_code=404, detail="GTFS dataset not found") + return dataset.feed, dataset + + feed = ( + db_session.query(Gtfsfeed) + .filter(Gtfsfeed.stable_id == feed_stable_id) + .one_or_none() + ) + if feed is None: + raise HTTPException(status_code=404, detail="GTFS feed not found") + if feed.latest_dataset is None: + raise HTTPException(status_code=404, detail="GTFS feed has no dataset yet") + return feed, feed.latest_dataset + + +def _enqueue_sizing(db_session: Session, feed, dataset) -> dict: + """The sizing arguments for the enqueue, flattened.""" + routing = _size_for(db_session, feed, dataset) + return { + "size": routing.size, + "variant_basis": routing.basis.value, + "override": routing.size.value if routing.basis is not Basis.MEASURED else None, + } + + +def _size_for(db_session: Session, feed, dataset) -> Routing: + """Which worker should build this dataset, and why. + + A thin wrapper over the shared helper: everything here is the Parquet builder's own + policy - its bands, its config namespace - and none of the mechanism. + """ + return size_for_dataset( + db_session, + feed, + dataset, + tiers=SIZE_TIERS, + namespace=SIZE_CONFIG_NAMESPACE, + key=SIZE_CONFIG_KEY, + compression_ratio=COMPRESSION_RATIO, + ) + + +def _is_expired(metadata: dict) -> bool: + """True once the set has passed the expiry its build recorded. + + The objects are deleted by the bucket's lifecycle rule, on GCS's own schedule and + with no promptness guarantee, so the row cannot be trusted to disappear with them. + Reading the expiry here is what makes the dataset report `absent` on the date and + rebuild on the next request. Rows written without an expiry never expire. + """ + raw = metadata.get("expires_at") + if not raw: + return False + try: + expires_at = datetime.fromisoformat(raw) + except (TypeError, ValueError): + logging.warning("Ignoring unparsable expires_at %r", raw) + return False + if expires_at.tzinfo is None: + expires_at = expires_at.replace(tzinfo=timezone.utc) + return expires_at <= datetime.now(timezone.utc) + + +def _state_of(db_session: Session, feed, dataset) -> ParquetDatasetState: + """Turn the tracking row into the state a viewer can act on.""" + tracker = TaskExecutionTracker( + task_name=TASK_NAME, + run_id=PARQUET_CONVERTER_VERSION, + db_session=db_session, + ) + # The row itself, not `is_handled`: that reduces it to a yes/no for dispatch, and + # this has to tell `preparing` from `ready` and `failed`. + row = tracker.get_entity(dataset.stable_id) + + base = { + "feed_stable_id": feed.stable_id, + "dataset_stable_id": dataset.stable_id, + } + + if row is None: + return ParquetDatasetState(status=STATUS_ABSENT, **base) + + metadata = row.metadata_ or {} + + if row.status == STATUS_COMPLETED: + if _is_expired(metadata): + # The files are gone, or about to be. `absent` is also what lets the next + # request rebuild: `handle_generate` re-triggers it, which `try_acquire` + # accepts, where a `completed` row would have been refused. + return ParquetDatasetState(status=STATUS_ABSENT, **base) + # No table list: a reader handed one skips `manifest.json`, and the manifest + # is where the sizes and counts live. The builder still records the tables in + # the tracking row for diagnostics; they are simply not served. + return ParquetDatasetState( + status=STATUS_READY, + base_url=metadata.get("base_url"), + generated_at=row.completed_at, + **base, + ) + + if row.status == STATUS_FAILED: + return ParquetDatasetState( + status=STATUS_FAILED_STATE, + message=row.error_message or "The conversion failed.", + **base, + ) + + return ParquetDatasetState( + status=STATUS_PREPARING, + phase=metadata.get("phase", "start"), + done=metadata.get("done", 0), + total=metadata.get("total", 0), + detail=metadata.get("detail", ""), + **base, + ) diff --git a/functions-python/operations_api/src/main.py b/functions-python/operations_api/src/main.py index 3ce4ab925..d6cc38a0b 100644 --- a/functions-python/operations_api/src/main.py +++ b/functions-python/operations_api/src/main.py @@ -23,6 +23,7 @@ from feeds_gen.apis.early_access_api import router as EarlyAccessApiRouter from feeds_gen.apis.licenses_api import router as LicenseApiRouter from feeds_gen.apis.users_api import router as UsersApiRouter +from feeds_gen.apis.parquet_api import router as ParquetApiRouter import functions_framework import asyncio @@ -58,6 +59,7 @@ app.include_router(EarlyAccessApiRouter) app.include_router(LicenseApiRouter) app.include_router(UsersApiRouter) +app.include_router(ParquetApiRouter) def build_scope_from_wsgi(request: Request) -> dict: diff --git a/functions-python/operations_api/tests/feeds_operations/impl/test_parquet_api_impl.py b/functions-python/operations_api/tests/feeds_operations/impl/test_parquet_api_impl.py new file mode 100644 index 000000000..14976a35d --- /dev/null +++ b/functions-python/operations_api/tests/feeds_operations/impl/test_parquet_api_impl.py @@ -0,0 +1,515 @@ +# +# MobilityData 2026 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""The four states, and the rule that a build is never started twice. + +Mock-style: a `db_session` passed explicitly bypasses the `with_db_session` decorator, +so these run without a database. What matters here is the mapping from a tracking row +to the state a viewer renders, which is pure logic. +""" + +import unittest +from datetime import datetime, timedelta, timezone +from unittest.mock import MagicMock, patch + +from fastapi import HTTPException + +from shared.helpers import sizing +from feeds_operations.impl import parquet_api_impl +from feeds_operations.impl.parquet_api_impl import ParquetApiImpl + +# The override lookup now lives in the shared sizing helper. +CONFIG_VALUE = "shared.helpers.sizing.get_config_value" +GB_ = 1024**3 +FEED = "mdb-1210" +DATASET = "mdb-1210-202402121801" +BASE_URL = f"https://files.example.org/{FEED}/{DATASET}/parquet" +GENERATED_AT = datetime(2026, 9, 16, 14, 22, 31, tzinfo=timezone.utc) + + +def _entities(): + feed, dataset = MagicMock(), MagicMock() + feed.stable_id = FEED + dataset.stable_id = DATASET + dataset.feed = feed + feed.latest_dataset = dataset + return feed, dataset + + +def _row(status, metadata=None, error_message=None, completed_at=None): + row = MagicMock() + row.status = status + row.metadata_ = metadata + row.error_message = error_message + row.completed_at = completed_at + return row + + +class ParquetStateTestCase(unittest.TestCase): + """Shared plumbing: a resolved feed/dataset and a stubbed tracker.""" + + def setUp(self): + self.feed, self.dataset = _entities() + self.tracker = MagicMock() + self.tracker.get_entity.return_value = None + + self._resolve_patch = patch.object( + parquet_api_impl, "_resolve", return_value=(self.feed, self.dataset) + ) + self._tracker_patch = patch.object( + parquet_api_impl, "TaskExecutionTracker", return_value=self.tracker + ) + self._enqueue_patch = patch.object( + parquet_api_impl, "create_http_parquet_builder_task" + ) + # Routing has its own tests; here it is pinned so the flow assertions do not + # depend on what a MagicMock session reports as a file size. + self._size_patch = patch.object( + parquet_api_impl, + "_size_for", + return_value=parquet_api_impl.Routing( + size=sizing.Size.L, + basis=sizing.Basis.MEASURED, + ), + ) + self._size_patch.start() + self.addCleanup(self._size_patch.stop) + self._resolve_patch.start() + self._tracker_patch.start() + self.enqueue = self._enqueue_patch.start() + self.addCleanup(self._resolve_patch.stop) + self.addCleanup(self._tracker_patch.stop) + self.addCleanup(self._enqueue_patch.stop) + + self.api = ParquetApiImpl() + + def status(self): + return self.api.handle_status(dataset_stable_id=DATASET, db_session=MagicMock()) + + def generate(self, force=False): + return self.api.handle_generate( + dataset_stable_id=DATASET, force=force, db_session=MagicMock() + ) + + +class TestStatus(ParquetStateTestCase): + def test_untracked_dataset_is_absent_not_an_error(self): + """`absent` is the ordinary first answer; a 404 here would read as a failure.""" + state = self.status() + + self.assertEqual(state.status, "absent") + self.assertEqual(state.feed_stable_id, FEED) + self.assertEqual(state.dataset_stable_id, DATASET) + self.assertIsNone(state.base_url) + + def test_in_progress_reports_the_phase_it_reached(self): + self.tracker.get_entity.return_value = _row( + "in_progress", + metadata={ + "phase": "convert", + "done": 12, + "total": 32, + "detail": "stop_times", + }, + ) + + state = self.status() + + self.assertEqual(state.status, "preparing") + self.assertEqual( + (state.phase, state.done, state.total, state.detail), + ("convert", 12, 32, "stop_times"), + ) + + def test_a_claim_with_no_reading_yet_still_reports_preparing(self): + self.tracker.get_entity.return_value = _row("in_progress", metadata=None) + + state = self.status() + + self.assertEqual(state.status, "preparing") + self.assertEqual(state.phase, "start") + + def test_completed_reports_where_the_files_are(self): + self.tracker.get_entity.return_value = _row( + "completed", + metadata={ + "base_url": BASE_URL, + "tables": [{"name": "stops", "rows": 4821, "parquet_bytes": 148213}], + }, + completed_at=GENERATED_AT, + ) + + state = self.status() + + self.assertEqual(state.status, "ready") + self.assertEqual(state.base_url, BASE_URL) + self.assertEqual(state.generated_at, GENERATED_AT) + + def test_the_table_list_is_not_served(self): + """A reader handed one skips manifest.json, and with it every size it carries. + + The builder still records the tables in the tracking row for diagnostics; this + asserts they stay there rather than reaching the response. + """ + self.tracker.get_entity.return_value = _row( + "completed", + metadata={"base_url": BASE_URL, "tables": [{"name": "stops"}]}, + completed_at=GENERATED_AT, + ) + + state = self.status() + + self.assertFalse( + hasattr(state, "tables"), + "ParquetDatasetState must not carry a table list", + ) + + def test_failed_reports_the_reason_verbatim(self): + """The viewer shows this to an operator, so a generic string would waste it.""" + self.tracker.get_entity.return_value = _row( + "failed", error_message="Conversion ran out of memory" + ) + + state = self.status() + + self.assertEqual(state.status, "failed") + self.assertEqual(state.message, "Conversion ran out of memory") + + def test_failed_without_a_reason_still_says_something(self): + self.tracker.get_entity.return_value = _row("failed", error_message=None) + + self.assertEqual(self.status().message, "The conversion failed.") + + def test_status_never_enqueues(self): + """A client polls this twice a second.""" + self.status() + + self.enqueue.assert_not_called() + + +class TestExpiry(ParquetStateTestCase): + """The files go on the bucket's schedule; the row is what dates the set.""" + + def _ready(self, expires_at): + metadata = {"base_url": f"https://files.test/{FEED}/{DATASET}/parquet"} + if expires_at is not None: + metadata["expires_at"] = expires_at + self.tracker.get_entity.return_value = _row("completed", metadata=metadata) + return self.status() + + def test_a_past_expiry_reads_absent(self): + """Not `ready`: the objects are gone, or about to be.""" + past = (datetime.now(timezone.utc) - timedelta(days=1)).isoformat() + + state = self._ready(past) + + self.assertEqual(state.status, "absent") + self.assertIsNone(state.base_url) + + def test_a_future_expiry_still_reads_ready(self): + future = (datetime.now(timezone.utc) + timedelta(days=1)).isoformat() + + state = self._ready(future) + + self.assertEqual(state.status, "ready") + + def test_a_row_without_an_expiry_never_expires(self): + """Rows written before the field existed keep working.""" + self.assertEqual(self._ready(None).status, "ready") + + def test_an_unparsable_expiry_is_ignored_rather_than_fatal(self): + self.assertEqual(self._ready("not a date").status, "ready") + + def test_a_naive_timestamp_is_read_as_utc(self): + """Everything written is tz-aware, but a naive value must not crash the read.""" + past = (datetime.now(timezone.utc) - timedelta(days=1)).replace(tzinfo=None) + + self.assertEqual(self._ready(past.isoformat()).status, "absent") + + def test_expiry_never_starts_work(self): + self._ready((datetime.now(timezone.utc) - timedelta(days=1)).isoformat()) + + self.enqueue.assert_not_called() + + +class TestGenerate(ParquetStateTestCase): + def test_absent_enqueues_and_reports_preparing_immediately(self): + state = self.generate() + + self.enqueue.assert_called_once_with( + FEED, + DATASET, + force=False, + retention_days=None, + size=sizing.Size.L, + variant_basis="measured", + override=None, + ) + # Not `absent` again: a client polling at 500ms would otherwise ask twice. + self.assertEqual(state.status, "preparing") + self.assertEqual(state.phase, "start") + + def test_the_enqueue_is_recorded_so_the_next_poll_sees_it(self): + """Otherwise the gap before the worker starts reads as `absent` all over again.""" + self.generate() + + self.tracker.mark_triggered.assert_called_once() + self.assertEqual(self.tracker.mark_triggered.call_args.args[0], DATASET) + + def test_the_marker_is_written_before_the_task_is_dispatched(self): + """An unconditional upsert cannot follow a dispatch the worker may already have + picked up.""" + order = [] + self.tracker.mark_triggered.side_effect = lambda *a, **k: order.append("mark") + self.enqueue.side_effect = lambda *a, **k: order.append("enqueue") + + self.generate() + + self.assertEqual(order, ["mark", "enqueue"]) + + def test_a_failed_enqueue_leaves_the_dataset_retriable(self): + """`preparing` is the one state this endpoint refuses to re-trigger.""" + self.enqueue.side_effect = RuntimeError("queue unreachable") + + with self.assertRaises(HTTPException): + self.generate() + + self.tracker.mark_failed.assert_called_once() + self.assertEqual(self.tracker.mark_failed.call_args.args[0], DATASET) + self.assertIn( + "queue unreachable", + self.tracker.mark_failed.call_args.kwargs["error_message"], + ) + + def test_a_compensation_that_also_fails_still_surfaces_the_500(self): + self.enqueue.side_effect = RuntimeError("queue unreachable") + self.tracker.mark_failed.side_effect = RuntimeError("database gone too") + + with self.assertRaises(HTTPException) as caught: + self.generate() + + self.assertEqual(caught.exception.status_code, 500) + + def test_a_running_build_is_reported_not_duplicated(self): + self.tracker.get_entity.return_value = _row( + "in_progress", + metadata={"phase": "convert", "done": 1, "total": 3, "detail": "stops"}, + ) + + state = self.generate() + + self.enqueue.assert_not_called() + self.assertEqual(state.status, "preparing") + self.assertEqual(state.phase, "convert") + + def test_force_cannot_interrupt_a_running_build(self): + self.tracker.get_entity.return_value = _row("in_progress", metadata={}) + + state = self.generate(force=True) + + self.enqueue.assert_not_called() + self.assertEqual(state.status, "preparing") + + def test_ready_is_returned_without_rebuilding(self): + self.tracker.get_entity.return_value = _row( + "completed", metadata={"base_url": BASE_URL} + ) + + state = self.generate() + + self.enqueue.assert_not_called() + self.assertEqual(state.status, "ready") + + def test_force_rebuilds_a_ready_dataset(self): + self.tracker.get_entity.return_value = _row( + "completed", metadata={"base_url": BASE_URL} + ) + + self.generate(force=True) + + self.enqueue.assert_called_once_with( + FEED, + DATASET, + force=True, + retention_days=None, + size=sizing.Size.L, + variant_basis="measured", + override=None, + ) + + def test_a_failed_dataset_is_retried(self): + self.tracker.get_entity.return_value = _row("failed", error_message="boom") + + state = self.generate() + + self.enqueue.assert_called_once_with( + FEED, + DATASET, + force=False, + retention_days=None, + size=sizing.Size.L, + variant_basis="measured", + override=None, + ) + self.assertEqual(state.status, "preparing") + + def test_an_enqueue_failure_is_a_500_not_a_silent_success(self): + self.enqueue.side_effect = RuntimeError("queue unreachable") + + with self.assertRaises(HTTPException) as caught: + self.generate() + + self.assertEqual(caught.exception.status_code, 500) + + +class TestSizeRouting(unittest.TestCase): + """This impl owns the bands and the namespace; the mechanism is tested in helpers.""" + + GB = 1024**3 + + def _size_for(self, largest): + session = MagicMock() + session.query.return_value.filter.return_value.scalar.return_value = largest + dataset = MagicMock() + dataset.id = "dataset-uuid" + dataset.stable_id = DATASET + dataset.unzipped_size_bytes = None + dataset.zipped_size_bytes = None + feed = MagicMock() + feed.id = "feed-uuid" + feed.stable_id = FEED + with patch(CONFIG_VALUE, return_value=None): + return parquet_api_impl._size_for(session, feed, dataset).size + + def test_the_bands(self): + for largest, expected in ( + (10 * 1024**2, sizing.Size.S), + (255_999_999, sizing.Size.S), + (256_000_000, sizing.Size.M), + (1_499_999_999, sizing.Size.M), + (1_500_000_000, sizing.Size.L), + (4 * GB_, sizing.Size.L), + ): + with self.subTest(largest=largest): + self.assertEqual(self._size_for(largest), expected) + + def test_an_unmeasurable_dataset_goes_to_the_largest(self): + """Datasets predating #1284 have none of the size columns populated.""" + self.assertEqual(self._size_for(None), sizing.Size.L) + + def test_it_pins_against_its_own_config_namespace(self): + session = MagicMock() + session.query.return_value.filter.return_value.scalar.return_value = 1 + dataset, feed = MagicMock(), MagicMock() + with patch(CONFIG_VALUE, return_value="l") as config: + parquet_api_impl._size_for(session, feed, dataset) + + self.assertEqual(config.call_args_list[0].args[0], "parquet_builder") + self.assertEqual(config.call_args_list[0].args[1], "size") + + +class TestResolution(unittest.TestCase): + """404 means the feed or dataset does not exist - never that it is unconverted.""" + + def test_unknown_dataset_is_404(self): + session = MagicMock() + session.query.return_value.filter.return_value.one_or_none.return_value = None + + with self.assertRaises(HTTPException) as caught: + parquet_api_impl._resolve(session, None, "nope") + + self.assertEqual(caught.exception.status_code, 404) + + def test_unknown_feed_is_404(self): + session = MagicMock() + session.query.return_value.filter.return_value.one_or_none.return_value = None + + with self.assertRaises(HTTPException) as caught: + parquet_api_impl._resolve(session, "nope", None) + + self.assertEqual(caught.exception.status_code, 404) + + def test_a_feed_with_no_dataset_is_404(self): + feed = MagicMock() + feed.stable_id = FEED + feed.latest_dataset = None + session = MagicMock() + session.query.return_value.filter.return_value.one_or_none.return_value = feed + + with self.assertRaises(HTTPException) as caught: + parquet_api_impl._resolve(session, FEED, None) + + self.assertEqual(caught.exception.status_code, 404) + self.assertIn("no dataset", caught.exception.detail) + + def test_a_feed_resolves_to_its_latest_dataset(self): + feed, dataset = _entities() + session = MagicMock() + session.query.return_value.filter.return_value.one_or_none.return_value = feed + + resolved_feed, resolved_dataset = parquet_api_impl._resolve(session, FEED, None) + + self.assertIs(resolved_feed, feed) + self.assertIs(resolved_dataset, dataset) + + +class TestRetentionPassthrough(ParquetStateTestCase): + """The API declares the bounds; the builder owns the default.""" + + def test_an_omitted_value_is_passed_through_as_none(self): + self.generate() + + self.assertIsNone(self.enqueue.call_args.kwargs["retention_days"]) + + def test_a_callers_value_reaches_the_builder(self): + self.api.handle_generate( + dataset_stable_id=DATASET, retention_days=7, db_session=MagicMock() + ) + + self.assertEqual(self.enqueue.call_args.kwargs["retention_days"], 7) + + +class TestRetentionUnwrapping(unittest.TestCase): + def test_absent_body_yields_none(self): + self.assertIsNone(parquet_api_impl._retention_days(None)) + + def test_absent_field_yields_none(self): + request = MagicMock() + request.retention_days = None + self.assertIsNone(parquet_api_impl._retention_days(request)) + + def test_a_value_is_returned_unchanged(self): + request = MagicMock() + request.retention_days = 45 + self.assertEqual(parquet_api_impl._retention_days(request), 45) + + +class TestForceFlag(unittest.TestCase): + def test_absent_body_is_not_a_force(self): + self.assertFalse(parquet_api_impl._force(None)) + + def test_explicit_force(self): + request = MagicMock() + request.force = True + self.assertTrue(parquet_api_impl._force(request)) + + def test_null_force_is_not_a_force(self): + request = MagicMock() + request.force = None + self.assertFalse(parquet_api_impl._force(request)) + + +if __name__ == "__main__": + unittest.main() diff --git a/functions-python/parquet_builder/.coveragerc b/functions-python/parquet_builder/.coveragerc new file mode 100644 index 000000000..4ab57a633 --- /dev/null +++ b/functions-python/parquet_builder/.coveragerc @@ -0,0 +1,11 @@ +[run] +omit = + */test*/* + */helpers/* + */database_gen/* + */shared/* + +[report] +exclude_lines = + pragma: no cover + if __name__ == .__main__.: diff --git a/functions-python/parquet_builder/README.md b/functions-python/parquet_builder/README.md new file mode 100644 index 000000000..2d25109a6 --- /dev/null +++ b/functions-python/parquet_builder/README.md @@ -0,0 +1,540 @@ +# Parquet Builder + +Renders a GTFS dataset as Parquet: one file per GTFS table, plus a `manifest.json` +describing the set. The result is published publicly beside the dataset so a browser +can query the feed in place over HTTP range requests, without downloading or unpacking +it. + +The Operations API decides *whether* a build is needed and reports progress; this +function does the work. See [Operations API endpoints](#operations-api-endpoints) +below. + +## Usage + +The function receives the following request: + +``` +{ + "feed_stable_id": str, – stable_id of the GTFS feed + "dataset_stable_id": str, – stable_id of the dataset to convert + "force": bool (optional) – rebuild even if a finished build exists (default: false) +} +``` + +Example: + +```json +{ + "feed_stable_id": "mdb-1210", + "dataset_stable_id": "mdb-1210-202402121801" +} +``` + +The function verifies that the dataset stable id starts with the feed stable id. + +### Output + +``` +gs://///parquet/ + agency.parquet + routes.parquet + stops.parquet + ... + manifest.json +``` + +Objects are made public. They have to be: the reader builds each file's URL from the +base URL's origin and path only, discarding any query string, so a signed URL cannot +survive the round trip. + +**The publish order is part of the contract.** Tables go up first, `manifest.json` +last, and stale objects are pruned only afterwards. The manifest is what a reader +holding just the bucket URL uses to learn which tables exist, so publishing it earlier +advertises files that have not arrived - the reader asks for every table and finds only +the handful uploaded so far. Nothing is deleted before the new set is up either: a +rebuild that cleared the prefix first left an already-published dataset unreadable for +the length of the upload, and a reader that has been told the dataset is ready has +stopped polling by then and never finds out. Both orderings are pinned by +`TestPublishIsNotObservablyPartial` in `tests/test_main.py`. + +One case is not covered by this and cannot be, short of publishing to a throwaway +prefix and copying: on a **first** build there is no manifest and no `ready` status yet, +so a client that goes straight to the bucket and probes for table names can still catch +a partial set. Clients driven by the Operations API or by the manifest never do. + +`manifest.json` - the dataset's description of itself, and the only place a reader +learns what it holds: + +```json +{ + "version": 2, + "generated_at": "2026-09-17T22:41:03+00:00", + "converter_version": "2", + "source": { "kind": "zip", "bytes": 4821334 }, + "totals": { "uncompressed_bytes": 18422910, "stored_bytes": 903411 }, + "tables": [ + { "name": "stops", "file": "stops.parquet", "rows": 4821, "columns": 12, + "bytes": 481223, "compressed_bytes": 92210, "parquet_bytes": 41880 } + ] +} +``` + +**`bytes` is the source size, not the Parquet size.** In version 1 it meant the +opposite, which is why the version was bumped rather than the field added to: a reader +has to know which it is holding. `compressed_bytes` is what the file weighed inside the +archive and is null for a feed converted from a folder; `parquet_bytes` is the converted +size. The sizes are captured during conversion because they cannot be recovered +afterwards - the CSVs are deleted and the archive is gone. + +The Operations API deliberately does **not** repeat this table list. The reader skips +`manifest.json` entirely when it is handed a list, and the manifest is where the sizes +and counts live, so serving one would suppress the load report it draws from them. + +### Response + +The function returns HTTP 200 even when a build fails, and records the reason in the +database instead. Cloud Tasks retries non-2xx, and nothing here fails transiently in a +way a retry would fix - a corrupt archive is still corrupt on the second delivery. + +```json +{ "status": "success", "dataset": "mdb-1210-202402121801", "base_url": "...", "tables": ["agency", "stops"] } +{ "status": "skipped", "reason": "already in progress", "dataset": "mdb-1210-202402121801" } +{ "status": "error", "error": "Failed to build Parquet for dataset ..." } +``` + +## What the conversion guarantees + +Two properties are load-bearing for the reader and break *silently* rather than loudly +if they change, so they are asserted in `tests/test_converter.py`: + +- **Every column is text.** The reader filters with `ILIKE` and `= ''`; a typed column + still renders but stops matching. `ALL_VARCHAR` on the way in is what guarantees it + on the way out. +- **An empty CSV field becomes NULL**, not an empty string, because the reader rewrites + `= ''` into `IS NULL OR = ''` on that basis. + +Table names are the GTFS file stem (`stops.txt` → `stops`), plus `locations` flattened +from `locations.geojson`. A file with no header row is skipped; a header with no data +rows is kept, since an empty `frequencies.txt` is a legitimate part of a feed. + +`PARQUET_CONVERTER_VERSION` in `src/converter.py` is the run id of the tracking rows, +and must match the constant of the same name in the Operations API implementation. +Bump it when the output changes in a way that makes previously written files wrong: a +bump invalidates every dataset's artifacts rather than serving files a newer reader no +longer matches. It went to `"2"` with manifest v2, so datasets built before that report +as `absent` and rebuild on first request. + +## Retention + +Each build stamps every object it publishes with a `customTime` of +`now + retention_days` (default 30, 1..60 per request), and the datasets bucket carries +one lifecycle rule - `daysSinceCustomTime: 0`, Delete - that removes an object once its +own `customTime` has passed. One rule, a different date per file, no scheduled job. + +The same timestamp goes into the tracking row as `expires_at`, and the Operations API +reports an expired row as `absent`. That is deliberate rather than redundant: lifecycle +deletion is asynchronous and GCS gives no promptness guarantee, so the row cannot be +relied on to vanish with the files. Reporting `absent` from the row is what makes the +dataset rebuild on the next request, and it happens on the date rather than whenever the +bucket gets round to the delete. The window in between - `absent` reported while the +files are still there - is harmless; a rebuild overwrites them. + +**The builder must stay the only writer of `customTime` in this bucket.** The lifecycle +rule is bucket-wide and matched solely by the presence of that field (the condition is +never satisfied for an object without it, and `matchesPrefix` cannot express +`*/parquet/`). Anything else that starts setting `customTime` there becomes deletable by +this rule. + +Nothing deletes the `task_execution_log` rows; they are small, and one row per converter +version per dataset already accumulates by design. + +## Concurrency + +A dataset is never converted twice at once. Cloud Tasks delivers at least once, so the +claim is taken by this function rather than by whatever enqueued it: +`TaskExecutionTracker.try_acquire` does a conditional upsert on `task_execution_log`, +and a refused claim makes the function return `skipped` without doing any work. + +The claim carries a 30 minute lease, renewed by each progress write. It is deliberately +longer than the function's own 1680s timeout, so an instance GCP has not finished +killing cannot have its work started underneath it. A build killed by OOM or timeout +leaves its claim to expire; one that fails normally releases it immediately. + +## Worker sizes + +The builder is deployed three times from one source zip, as +`parquet-builder-{s,m,l}-`, each with its own Cloud Tasks queue. The Operations API +picks one at enqueue time. + +| Size | Memory | CPU | Volume | DuckDB | Largest uncompressed file | Share of feeds | +|---|---|---|---|---|---|---| +| `s` | 3Gi | 1 | 1Gi | 512MB | < 256 MB | ~98% | +| `m` | 7Gi | 2 | 3Gi | 1GB | < 1.5 GB | ~1.8% | +| `l` | 16Gi | 4 | 8Gi | 2GB | anything larger, or unknown | ~0.2% | + +A feed pinned through config bypasses the table entirely; see below. + +The bands are measured rather than guessed. Across 4277 feeds the median archive is +0.2 MB and only 31 exceed 100 MB, while about ten hold a single member over 1 GB and the +largest holds one of 4.8 GB. The catalogue is not a spectrum: it is a great many trivial +feeds, a thin band of medium ones, and roughly a dozen large ones. + +Two rules set the numbers for each rung: + +- **the volume** must hold the largest member plus its Parquet output, and is carved out + of the total rather than added to it; +- **the process budget** left over must cover the process's *address space*, which is + what `RLIMIT_AS` caps - not its resident size. Measured across ten feeds: + + | rung | budget | DuckDB cap | peak RSS | peak VMS | VMS of budget | + |---|---|---|---|---|---| + | `s` | 1848 MiB | 512MB | 443 MB | 816 MB | 44% | + | `m` | 3896 MiB | 1GB | 1056 MB | 2297 MB | 59% | + | `l` | 7992 MiB | 2GB | 2496 MB | 5571 MB | 70% | + + **VMS runs about 2.2-2.8x the DuckDB cap, and 2-3x peak RSS.** Baseline VMS is ~550 MB + before any work. Size a rung from `vms`; sizing from `rss` under-provisions by about + half, which is how `l` ended up at a 3896 MiB budget and killed mdb-2014 with + `MemoryError`, and how `s` ended up at 824 MiB and died on a 41.7 MB feed with a + DuckDB `OutOfMemoryException`. The `s` and `m` figures above are post-fix; both were + raised a rung after those runs. + +**The measure is the largest single file, not the archive or the feed total.** The volume +holds one source at a time, so that file is what decides whether a build fits. A total +would misjudge a feed of many medium files, and the compressed size misjudges almost +everything: mdb-2014 is a 1.08 GiB archive containing a 4.07 GiB `stop_times.txt`. + +It comes from `max(gtfsfile.file_size_bytes)` for the dataset. Those rows are absent for +datasets processed before they existed, so the measure falls back to +`unzipped_size_bytes`, then to `zipped_size_bytes` times a conservative ratio, then to +nothing. **An unknown size routes to `l`**, never `m` - guessing small turns a missing row +into an OOM. The log line says which rung answered; if it says `unknown`, run the +`rebuild_missing_dataset_files` task to record them. + +Separate queues are the point rather than a side effect. `max_concurrent_dispatches` caps +memory across concurrent instances, so a single queue has to be sized for the heaviest job +and a run of large builds leaves the short ones waiting behind them. + +### Pinning a feed by hand + +One key, `size`, under namespace `parquet_builder`, registered by +`liquibase/changes/feat_parquet_builder_size.sql` - no new table, just the `config_key` +row that `config_value_feed`'s foreign key needs. + +One row per feed. **An override decides which worker runs the build**, in place of the +size the dataset would otherwise measure into. Three accepted forms: + +| value | set by | may the builder move it? | +|---|---|---| +| `"l"` | a person | yes | +| `{"size": "l", "source": "auto"}` | the builder, after a build ran out of resources | yes | +| `{"size": "l", "locked": true}` | a person, to hold the size | no | + +All three decide the routing outright. `source` is carried for the record - it shows up +as `variant_basis` on each attempt - and changes nothing. It used to decide whether the +builder could overwrite a row, which meant inferring an operator's intent from the fact +of their authorship; `locked` is that intent stated instead. + +**Without `locked`, a pin is where the feed starts, not where it stays.** The healing +will raise it after a build runs out of resources and lower it after several that +finished with room to spare, so a value set by hand can be gone a week later. That is +usually what you want; when it is not, lock it. + +The bare form is the one to type by hand: + +```sql +INSERT INTO config_value_feed (feed_id, feed_stable_id, namespace, key, value) +SELECT id, stable_id, 'parquet_builder', 'size', '"l"'::jsonb + FROM feed WHERE stable_id = 'mdb-2014' +ON CONFLICT (feed_id, namespace, key) DO UPDATE SET value = EXCLUDED.value; +``` + +To hold it there instead, `'{"size":"l","locked":true}'::jsonb`. Only a literal `true` +locks; anything else reads as unlocked, so a typo cannot freeze a feed where nothing can +move it again. + +**A lock is a real lock.** A feed locked below what it needs fails and stays failed: the +build runs out of resources, nothing escalates, nothing is re-queued, and the dataset is +not produced until a person changes the value. Each attempt records the failure with +`escalated_to` null, and `current_size_locked` on +`GET /v1/operations/gtfs_feeds/{id}/execution_attempts` says why - without it a run of +failed attempts with no escalation has no visible explanation. + +With no per-feed row, every feed is routed purely by measurement. The key carries no +`default_value` on purpose: setting one would move the whole catalogue at once, which +belongs in the routing table in `parquet_api_impl.py` instead. + +**An override is absolute.** The measurement is not consulted at all - not even computed +- so an override can send a feed either way. The flip side is that setting one below what +a feed needs will fail it with `ENOSPC` or `MemoryError`, so check the build after +changing one. + +Weighing the builder's own override against the measurement instead, so a grown feed +could overtake it, was considered and dropped: it changes the outcome only when the +measurement is larger, and that case already resolves itself - the build fails and the +escalation moves it up. One wasted build is not worth a second code path. + +**The builder only writes a row where it disagrees with the measurement.** Escalating a +feed whose dataset already measures into the target size stores nothing, and clears an +earlier `auto` row that the measurement has since caught up with - the dataset is on its +way to that worker through the task either way. So the table holds the feeds that are +genuinely exceptions, and a feed leaves it as soon as its own data says the same thing. + +A locked override is never cleared automatically, and never raised either; everything +else the builder may move. + +An unrecognised value is logged and ignored rather than failing the request. + +There is no endpoint or UI for `config_value_feed`, so this is SQL for now. + +### Coming back down + +An `auto` override is reviewed on the success path, in `_downsize_after_success`, and +lowered a rung when the feed's recent builds say it can be. Nothing is re-queued: the +build in hand has already succeeded, so the new size applies to the feed's next one. + +The two directions are deliberately asymmetric. Escalation acts on a single failure, +because being too small costs a build that cannot finish. Coming down is only ever an +economy, so it waits for `DOWNSIZE_STREAK` builds in a row that each left +`DOWNSIZE_HEADROOM` of the smaller worker unused - 3 and 60% today, both in +`helpers/parquet_policy.py`. + +Each of those builds is checked on both axes, because either one can end a build: + +| axis | evidence | compared against | +|---|---|---| +| memory | peak address space during the build | that rung's `RLIMIT_AS` budget | +| disk | largest uncompressed member actually opened | that rung's band ceiling | + +Both numbers come from the build itself, recorded on the attempt row - `peak_vms_bytes`, +and `largest_member_bytes` in its `metadata`. The measurement in the database is +deliberately not consulted: an override exists precisely because that measurement was +wrong about this feed, so reading it again to justify undoing the override would be +circular. + +Four things stop a review short, before any history is read: the feed has no override +(it already routes on its measurement, which is as low as it goes), the override is +locked, the worker is already the smallest, or the override disagrees with the worker +that just ran. An attempt recorded before this evidence existed has no +`largest_member_bytes` and reads as unknown, which blocks the streak rather than +permitting it - leaving a feed too large costs money, the other way costs builds. + +The case to watch is a feed whose datasets alternate between large and small: it could +move down, fail on the next large one, escalate, and repeat. The disk axis usually +catches it, because the large dataset's own build breaks the streak - but that is a +property of the evidence, not a guarantee. `task_execution_attempt` shows it if it +happens; the fix is a longer streak or a smaller fraction. + +### Where the policy lives + +The routing *mechanism* is `functions-python/helpers/sizing.py`, shared so other +functions can adopt it without copying: `choose_size` for the tier arithmetic, +`size_for_dataset` for the whole decision, `escalate` and `demote` for the ladder, and +`fits_within` for the two-axis comparison. + +The Parquet builder's *policy* is `functions-python/helpers/parquet_policy.py` - its +bands and their budgets, its config namespace and key, its compression ratio, and the +two downsize thresholds. It is shared rather than per-function because two processes +decide the same thing and have to agree: the Operations API routes a build when it +enqueues one, and the builder re-routes it when one fails. A disagreement would bounce a +dataset between two workers. + +One thing to look at before a second function adopts it: the measure. The largest single +uncompressed file is right here because the volume holds one at a time, but a function +bounded by something else wants `choose_size` with a measure of its own. + +## Memory + +The function's allocation is **split**, not shared: the in-memory volume is carved out of +the total rather than added to it. + +Per variant, taking `l` as the example: + +``` +cgroup limit (the variant's "memory") 16384 MiB + - in-memory volume at PARQUET_TMPDIR - 8192 MiB + - MEMORY_MARGIN_MB - 200 MiB + = RLIMIT_AS set on the Python process 7992 MiB +``` + +The same arithmetic gives `m` 3896 MiB and `s` 1848 MiB. The small volumes are viable +only because the archive is streamed rather than written to them. + +All three rungs live in `local.parquet_builder_sizes` in `infra/functions-python/main.tf`. +Memory and volume are set together there on purpose: the volume is subtracted from the +total, so splitting them across files invites a rung whose process budget is accidentally +negative. + +`limit_gcp_memory` (`shared/common/gcp_memory_utils.py`) does that subtraction at import, +before anything allocates, and sets `RLIMIT_AS`. The point is that an overshoot raises a +catchable `MemoryError` with a traceback instead of the kernel killing the container +silently. It reads the volume's *declared* size, so it is subtracted whether or not a +byte is written to it. Both numbers are logged on every cold start: + +``` +Process memory limit: 16384.00 MiB, total tmpfs size: 8192.00 MiB, available: 8192.00 MiB +RLIMIT_AS set to 7992.00 MiB +``` + +If `total tmpfs size` reads `0.00 MiB`, the volume did not get mounted and the process is +running with the full cgroup limit as its budget. That is a broken deploy, not a safe one. + +Three consumers, bounded differently: + +| Consumer | Bound | Overshoot | +|---|---|---| +| The workdir under `PARQUET_TMPDIR` | 8Gi, by the volume's own `size-limit` | `ENOSPC` | +| DuckDB | `PARQUET_DUCKDB_MEMORY_LIMIT`, 2GB | spills, see below | +| The Python process | `RLIMIT_AS` | `MemoryError` | + +**The tmpfs is the real limit, not the total.** DuckDB's spill directory is inside the +workdir, so spilling does not release memory from the container - it moves bytes out of +DuckDB's budget and into the tmpfs. A feed whose largest single CSV plus its Parquet plus +the spill exceeds the volume fails with `ENOSPC` no matter how much total memory the +function has. Raising the total without raising the volume does not help that case. + +Size the volume from the **largest single uncompressed file** in the feed, not from the +archive or the feed total. The archive path needs more again, because the `.zip` stays +resident for the whole build: + +``` +both paths: largest file + its Parquet + DuckDB spill +``` + +The worked example is mdb-2014, whose 1.08 GiB archive holds a 4.07 GiB `stop_times.txt` +and a 2.04 GiB `shapes.txt`. It once needed 1.08 + 4.07 = 5.15 GiB before conversion +started, and failed with `ENOSPC` on a 4Gi volume; now that the archive is streamed the +same feed needs 4.07 GiB plus its output. Read those numbers off any archive without +downloading it using +`zip_member_sizes_from_file` with `TailReader` over a ranged read of the last 1 MiB, which +is what the builder itself does to record compressed sizes. + +What keeps usage low is that the conversion streams: one source file and one Parquet output +exist at a time, each deleted before the next begins, so peak tracks the largest single +table rather than the whole feed. This holds on both paths - the archive is read over the +network through `Blob.open("rb")`, a seekable reader that `zipfile` drives directly, so it +is never written to the volume. Members are converted in archive order rather than table +order, because a backward seek discards the reader's buffer and refetches; the manifest is +sorted by name afterwards so it does not depend on which path produced it. + +Sizing from observed averages is a trap here: feed sizes span orders of magnitude, so a +sample that happens to exclude the largest feeds will suggest a volume that cannot build +them at all. The process budget is the part that measurement does settle - the largest +observed `process peak` is 2329 MB on `l`, but that is resident, not address space, and +`RLIMIT_AS` caps the latter - read `vms` for sizing. Note that the +`Function metrics` log +line reports `tracemalloc` for `memory:`, which sees Python allocations only - not DuckDB's +C++ heap and not tmpfs pages. Use the `rss` figure on that same line, or Cloud Monitoring's +`run.googleapis.com/container/memory/utilizations`, when sizing. There is no telemetry from +qa or prod yet, so the largest feeds in the catalogue may not have been converted. + +# GCP environment variables + +- `DATASETS_BUCKET_NAME`: bucket where datasets are stored, including the environment + suffix (`-dev`, `-qa`, `-prod`). The function fails if this is not defined. +- `PUBLIC_HOSTED_DATASETS_URL`: public URL prefix the artifacts are served from; used + to build the `base_url` reported back to callers. +- `FEEDS_DATABASE_URL` (secret): used for the claim and for progress reporting. +- `PARQUET_TMPDIR`: the in-memory volume everything large is written to + (default `/tmp/in-memory`). `limit_gcp_memory` reads its size to compute the process + budget, so this must point at the mounted volume. +- `PARQUET_DUCKDB_MEMORY_LIMIT`: DuckDB's own budget (default `2GB`). Set explicitly + because DuckDB otherwise sizes itself from the host's RAM rather than the cgroup, and + so would spill far too late to help. +- `MEMORY_MARGIN_MB`: margin subtracted before `RLIMIT_AS` is set (default `200`). + +# Local testing + +## Generating Parquet without GCP + +`scripts/parquet-generate-local.sh` runs the same conversion this function runs - +importing `converter.convert_to_parquet` rather than reimplementing it - against a feed +on disk or a public URL. No bucket, database, task queue or credentials involved. + +```bash +# A feed id: downloads its current archive over HTTPS +scripts/parquet-generate-local.sh mdb-1210 + +# A specific dataset, from a non-production environment +scripts/parquet-generate-local.sh mdb-1210-202402121801 --env dev + +# A feed already on disk, as an archive or an unpacked folder +scripts/parquet-generate-local.sh ./gtfs.zip +scripts/parquet-generate-local.sh ./extracted/ +``` + +Output goes to `.dist/parquet/` unless `--out` says otherwise. The virtualenv +and `duckdb` are provisioned on first run; nothing else is needed. + +## Serving it to a browser + +```bash +scripts/parquet-generate-local.sh mdb-1210 --serve # http://localhost:8090 +``` + +`--serve` exists because the obvious alternative does not work: `python -m http.server` +ignores `Range` and returns whole files, so a reader asking for a Parquet footer gets +the entire table. The built-in server answers ranges (including the suffix ranges a +footer read uses) and sends the CORS headers a cross-origin worker needs. + +## With the operations web app + +The operations web app serves `public/datasets/` as a static dataset, so generating +straight into it is enough to browse a real feed: + +```bash +scripts/parquet-generate-local.sh mdb-1210 \ + --out /public/datasets/mdb-1210 + +cd && yarn dev +# then open /feeds/gtfs/mdb-1210/browse +``` + +## Exercising the whole path, including the Operations API + +Only needed when the endpoints themselves are what is being tested. Cloud Tasks does +not dispatch locally - the `PARQUET_BUILDER_QUEUE_M`/`_L` vars are unset, so the enqueue +is a logged no-op - which means the builder is invoked by hand in place of the queue. + +```bash +# 1. Database and Operations API (http://localhost:8081) +docker-compose --env-file ./config/.env.local up -d +scripts/api-operations-start.sh + +# 2. The builder, against a real bucket +gcloud auth application-default login +scripts/function-python-setup.sh --function_name parquet_builder +DATASETS_BUCKET_NAME=mobilitydata-datasets-dev \ +PUBLIC_HOSTED_DATASETS_URL=https://dev-files.mobilitydatabase.org \ + scripts/function-python-run.sh --function_name parquet_builder # http://localhost:8080 + +# 3. Stand in for Cloud Tasks +curl -X POST localhost:8080 -H 'Content-Type: application/json' \ + -d '{"feed_stable_id":"mdb-1210","dataset_stable_id":"mdb-1210-202402121801"}' + +# 4. Watch the API report absent -> preparing -> ready +curl localhost:8081/v1/operations/gtfs_datasets/mdb-1210-202402121801/parquet +``` + +To skip GCS entirely at step 2, run the builder against a local feed with +`scripts/parquet-generate-local.sh` and point the viewer at the output instead; only +the database-backed status reporting needs the function itself. + +## Unit tests + +```bash +scripts/api-tests.sh --folder functions-python/parquet_builder +``` + +The suite needs no network and no GCP: Cloud Storage is faked, and the archives the +tests convert are real zips built in-process. + +# Operations API endpoints + +The endpoints a client drives to request and follow a build, the states they report, and +what a client fetches once a set is ready, are documented in +[docs/parquet-feed-browsing.md](../../docs/parquet-feed-browsing.md) - including the +authentication the spec gets wrong and the CORS requirement the browser reader depends on. + +Implemented in `functions-python/operations_api/src/feeds_operations/impl/parquet_api_impl.py`, +specified in `docs/OperationsAPI.yaml`. diff --git a/functions-python/parquet_builder/function_config.json b/functions-python/parquet_builder/function_config.json new file mode 100644 index 000000000..b6f218444 --- /dev/null +++ b/functions-python/parquet_builder/function_config.json @@ -0,0 +1,43 @@ +{ + "name": "parquet-builder", + "description": "The Parquet Builder function renders a GTFS dataset as one Parquet file per table", + "entry_point": "build_parquet_handler", + "timeout": 1680, + "memory": "16Gi", + "trigger_http": true, + "include_folders": [ + "helpers" + ], + "include_api_folders": [ + "database_gen", + "database", + "common" + ], + "environment_variables": [ + { + "key": "DATASETS_BUCKET_NAME" + }, + { + "key": "PUBLIC_HOSTED_DATASETS_URL" + }, + { + "key": "LOGGING_LEVEL" + }, + { + "key": "PARQUET_TMPDIR" + }, + { + "key": "PARQUET_DUCKDB_MEMORY_LIMIT" + } + ], + "secret_environment_variables": [ + { + "key": "FEEDS_DATABASE_URL" + } + ], + "ingress_settings": "ALL", + "max_instance_request_concurrency": 1, + "max_instance_count": 20, + "min_instance_count": 0, + "available_cpu": 4 +} diff --git a/functions-python/parquet_builder/requirements.txt b/functions-python/parquet_builder/requirements.txt new file mode 100644 index 000000000..fa062ab97 --- /dev/null +++ b/functions-python/parquet_builder/requirements.txt @@ -0,0 +1,26 @@ +# Common packages +functions-framework==3.* +google-cloud-logging +psycopg2-binary==2.9.6 +urllib3~=2.6.3 +attrs~=23.1.0 +certifi~=2025.8.3 + +# SQL Alchemy and Geo Alchemy +SQLAlchemy==2.0.23 +geoalchemy2==0.14.7 + +# Google specific packages for this function +flask +google-cloud-storage +google-cloud-tasks + +# The conversion itself. Reads the GTFS CSVs and writes Parquet; the only +# non-infrastructure dependency this function has. +duckdb~=1.5 + +# Runtime metrics +psutil + +# Configuration +python-dotenv==1.2.2 diff --git a/functions-python/parquet_builder/requirements_dev.txt b/functions-python/parquet_builder/requirements_dev.txt new file mode 100644 index 000000000..41e75d402 --- /dev/null +++ b/functions-python/parquet_builder/requirements_dev.txt @@ -0,0 +1,4 @@ +Faker +pytest~=7.4.3 +requests-mock +psutil diff --git a/functions-python/parquet_builder/src/converter.py b/functions-python/parquet_builder/src/converter.py new file mode 100644 index 000000000..97f09f793 --- /dev/null +++ b/functions-python/parquet_builder/src/converter.py @@ -0,0 +1,527 @@ +# +# MobilityData 2026 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Rewrite an extracted GTFS feed as one Parquet file per table. + +This mirrors what `gtfs-garage --export` produces, deliberately rather than by +importing it. Two reasons, the second being the one that would survive the first: + + * The distribution carries a web server. `gtfs-garage` is on PyPI, but it declares + FastAPI, uvicorn and python-multipart as unconditional dependencies, so a function + that never starts a server would ship one. Its `tests/test_layering.py` keeps + `gtfs_garage.core` framework-free and names this repository as the consumer, so + this is a packaging gap rather than a design one - it would take a `server` extra + over there to close it. + * The shapes differ. `Feed.export_parquet` registers the whole feed into one DuckDB + database and then exports every table; `main._convert_and_publish` fetches, + converts, uploads and deletes one table at a time, so peak memory tracks the + largest single table rather than the feed, which is what fits a large one inside + the function's tmpfs. Importing would not supply that loop, only the SQL in it. + +What is actually being reproduced is small - a view per file and a COPY per table - +and it is pinned here by tests that assert the properties the viewer depends on. + +Two of those properties are load-bearing and neither is obvious: + + * Every column is text. The reader compares with ILIKE and with `= ''`, so a typed + column silently breaks filtering rather than failing loudly. `ALL_VARCHAR` on the + way in is what guarantees it on the way out. + * An empty CSV field becomes NULL, not an empty string, because the reader rewrites + `= ''` into `IS NULL OR = ''` on that basis. + +Bump PARQUET_CONVERTER_VERSION when the output changes in a way that makes previously +written files wrong. It is the run id of the tracking rows, so a bump invalidates every +dataset's artifacts instead of serving files a newer reader no longer matches. +""" + +from __future__ import annotations + +import io +import json +import logging +import zipfile +from dataclasses import dataclass, field +from datetime import datetime, timezone +from pathlib import Path +from typing import Callable, Optional + +PARQUET_CONVERTER_VERSION = "2" + +# The one GTFS file that is not a CSV, and the table it becomes. +LOCATIONS_GEOJSON = "locations.geojson" +LOCATIONS_TABLE = "locations" + +# Written beside the Parquet: the dataset's description of itself. The reader takes +# its table list from here, and its load report - how big each file was before +# conversion and after it - which is why the sizes below are worth carrying. +MANIFEST = "manifest.json" +MANIFEST_VERSION = 2 + +# phase, done, total, detail +ProgressFn = Callable[[str, int, int, str], None] + +PHASE_CONVERT = "convert" +PHASE_EXTRACT = "extract" + + +@dataclass +class SourceFacts: + """Where the feed came from, for the manifest header. + + Defaults describe a directory of GTFS files, which is what a caller that simply + points at an extracted feed has. A zip source fills all three in. + """ + + kind: str = "folder" + bytes: Optional[int] = None + # Keyed by file name. Empty for a folder: a member's compressed size exists only + # in a zip's central directory and is unrecoverable once extracted. + compressed_sizes: dict = field(default_factory=dict) + + +@dataclass +class ConvertedTable: + name: str + file: str + rows: int + columns: int + # What the source file weighed, and what it weighed inside the archive it arrived + # in. Neither survives the conversion - the CSVs are deleted and the archive is + # gone - so they are recorded here or lost. + bytes: Optional[int] + compressed_bytes: Optional[int] + parquet_bytes: int + + def as_manifest_entry(self) -> dict: + return { + "name": self.name, + "file": self.file, + "rows": self.rows, + "columns": self.columns, + "bytes": self.bytes, + "compressed_bytes": self.compressed_bytes, + "parquet_bytes": self.parquet_bytes, + } + + +def extract_feed( + archive: Path, + destination: Path, + on_progress: Optional[ProgressFn] = None, +) -> Path: + """Unpack a GTFS archive into `destination`, flattening any wrapping directory. + + Producers differ on whether the files sit at the root of the zip or inside a + folder, and the conversion looks for *.txt in one place. Flattening here rather + than searching there keeps that difference from spreading. + + Shared with the local generation script deliberately, so a feed prepared on a + laptop is unpacked exactly as one prepared in the cloud function. + """ + destination.mkdir(parents=True, exist_ok=True) + with zipfile.ZipFile(archive) as zf: + members = [m for m in zf.infolist() if not m.is_dir()] + for index, member in enumerate(members, start=1): + name = Path(member.filename).name + if not name: + continue + if on_progress: + on_progress(PHASE_EXTRACT, index, len(members), name) + with zf.open(member) as source, open(destination / name, "wb") as target: + target.write(source.read()) + return destination + + +def zip_member_sizes_from_file(fileobj) -> dict[str, int]: + """Each member's compressed size, keyed by file name, from a seekable archive. + + Only the central directory is read, so a file-like that serves just the tail of a + remote archive is enough - see `TailReader`. The number exists nowhere else: once + the files are extracted, what they weighed inside the zip is gone, and the load + report has a column for it. + """ + try: + with zipfile.ZipFile(fileobj) as zf: + return { + Path(info.filename).name: info.compress_size + for info in zf.infolist() + if not info.is_dir() + } + except (OSError, zipfile.BadZipFile): + return {} + + +def zip_member_sizes(archive: Path) -> dict[str, int]: + """`zip_member_sizes_from_file` for an archive on disk.""" + try: + with open(archive, "rb") as handle: + return zip_member_sizes_from_file(handle) + except OSError: + return {} + + +class TailReader(io.RawIOBase): + """A read-only file over the last `len(tail)` bytes of a larger file. + + Lets `zipfile` read a remote archive's central directory without fetching the + archive. The directory lives at the end, but its records carry offsets into the + *whole* file, so a plain buffer of the tail would have every offset wrong. This + reports the full size and maps absolute positions onto the tail, which is all + `zipfile` needs to enumerate members - it only seeks earlier when asked to read an + entry's data, which we never do. + """ + + def __init__(self, tail: bytes, size: int): + self._tail = tail + self._size = size + self._start = size - len(tail) + self._pos = 0 + + def readable(self) -> bool: + return True + + def seekable(self) -> bool: + return True + + def tell(self) -> int: + return self._pos + + def seek(self, offset: int, whence: int = io.SEEK_SET) -> int: + if whence == io.SEEK_SET: + self._pos = offset + elif whence == io.SEEK_CUR: + self._pos += offset + else: + self._pos = self._size + offset + return self._pos + + def read(self, size: int = -1) -> bytes: + if self._pos < self._start: + # Only reachable if the directory did not fit in the tail; the caller + # retries with more rather than getting silently truncated records. + raise ValueError("read before the start of the fetched tail") + start = self._pos - self._start + end = len(self._tail) if size is None or size < 0 else start + size + chunk = self._tail[start:end] + self._pos += len(chunk) + return chunk + + def readall(self) -> bytes: + return self.read(-1) + + +def _quote(path: Path) -> str: + """Escape a path for embedding in a DuckDB string literal.""" + return str(path).replace("'", "''") + + +def _ident(name: str) -> str: + """Quote an identifier for DuckDB, doubling any embedded quote. + + Table names come from archive member filenames, which can contain anything a + filesystem accepts - a double quote included. + """ + return '"' + name.replace('"', '""') + '"' + + +def _has_header(path: Path) -> bool: + """True when the file opens with something that could be a header row. + + Worth checking up front because an empty file is not an error to DuckDB: it + invents a single `column0`, returns no rows, and would be published as a table + that is not one. A header with no data rows is an entirely different thing and is + kept - an empty frequencies.txt is a legitimate part of a feed. + + Only the first line is read, so this stays cheap on a multi-gigabyte stop_times. + """ + try: + with path.open("r", encoding="utf-8-sig", errors="replace") as handle: + return bool(handle.readline().strip()) + except OSError: + return False + + +def open_connection( + temp_dir: Optional[Path] = None, memory_limit: Optional[str] = None +): + """A DuckDB connection told where to spill and how much it may hold. + + Both settings matter in a container. DuckDB's default memory heuristic reads the + host's RAM rather than the cgroup, so left alone it spills far too late and the + process dies before it ever tries; and its default spill target is the working + directory. Naming both turns an overshoot into slower work instead of a lost run. + """ + import duckdb + + con = duckdb.connect(database=":memory:") + if temp_dir is not None: + temp_dir.mkdir(parents=True, exist_ok=True) + con.execute(f"SET temp_directory='{_quote(temp_dir)}'") + if memory_limit: + con.execute(f"SET memory_limit='{memory_limit}'") + return con + + +def table_name_for(source: Path) -> Optional[str]: + """The table a source file becomes, or None if it is not one. + + Producers ship all sorts of things in a feed - PDFs, `__MACOSX` entries, licences - + and everything in the archive is kept, so the filter lives here. + """ + # AppleDouble resource forks (`__MACOSX/._stops.txt`) flatten to a name that looks + # exactly like a table once the folder is stripped, and every archive member is + # published to `extracted/`, so this is routine rather than hypothetical. No GTFS + # file begins with a dot. + if source.name.startswith("."): + return None + if source.name == LOCATIONS_GEOJSON: + return LOCATIONS_TABLE + if source.suffix == ".txt": + return source.stem + return None + + +def register_table(con, table: str, source: Path, logger: logging.Logger) -> bool: + """Create the view for one source file. False when the file cannot be used. + + A file DuckDB cannot parse is skipped rather than fatal: most GTFS files are + optional, and one unreadable extra must not cost the whole feed. + """ + if table == LOCATIONS_TABLE: + return _register_locations(con, source, logger) + + if not _has_header(source): + logger.warning("Skipping %s: no header row", source.name) + return False + try: + con.execute(f""" + CREATE VIEW {_ident(table)} AS + SELECT * FROM read_csv( + '{_quote(source)}', + ALL_VARCHAR = TRUE, + IGNORE_ERRORS = TRUE, + NULLSTR = '' + ) + """) + except Exception as error: + logger.warning("Skipping unreadable file %s: %s", source.name, error) + return False + return True + + +def _register_locations(con, source: Path, logger: logging.Logger) -> bool: + """Flatten locations.geojson to one row per Feature, all columns text. + + `features` is read as opaque JSON rather than letting DuckDB infer its shape, so a + feed mixing Polygon and MultiPolygon cannot change the columns this produces. + """ + try: + con.execute(f""" + CREATE VIEW {_ident(LOCATIONS_TABLE)} AS + SELECT + CAST(json_extract_string(feature, '$.id') AS VARCHAR) AS id, + CAST(json_extract_string(feature, '$.properties.stop_name') AS VARCHAR) + AS stop_name, + CAST(json_extract_string(feature, '$.properties.stop_desc') AS VARCHAR) + AS stop_desc, + CAST(json_extract_string(feature, '$.geometry.type') AS VARCHAR) + AS geometry_type, + CAST(json_extract(feature, '$.geometry') AS VARCHAR) AS geometry + FROM ( + SELECT unnest(features) AS feature + FROM read_json( + '{_quote(source)}', + columns = {{type: 'VARCHAR', features: 'JSON[]'}} + ) + ) + """) + # A view is lazy, so one over malformed JSON is created happily and only fails + # later, mid-conversion. Reading a row forces the parse while it can still be + # handled as "this feed has no zones" rather than as a failed build. + con.execute(f"SELECT * FROM {_ident(LOCATIONS_TABLE)} LIMIT 1").fetchall() + except Exception as error: + logger.warning("Ignoring unusable %s: %s", LOCATIONS_GEOJSON, error) + con.execute(f"DROP VIEW IF EXISTS {_ident(LOCATIONS_TABLE)}") + return False + return True + + +def convert_table( + con, + table: str, + source: Path, + destination: Path, + compressed_bytes: Optional[int] = None, +) -> ConvertedTable: + """Write one registered table as Parquet and describe what was written. + + Drops the view afterwards so the source file can be deleted immediately - holding + it open would pin a file the caller is about to remove to reclaim memory. + """ + destination.mkdir(parents=True, exist_ok=True) + target = destination / f"{table}.parquet" + con.execute( + f"""COPY (SELECT * FROM {_ident(table)}) TO '{_quote(target)}' """ + f"""(FORMAT PARQUET, COMPRESSION ZSTD)""" + ) + # Counted off the file rather than the source view: Parquet carries its row count + # in the footer, so this reads a few bytes instead of reparsing what was just + # written. + rows = con.execute( + f"SELECT count(*) FROM read_parquet('{_quote(target)}')" + ).fetchone()[0] + columns = con.execute(f"DESCRIBE {_ident(table)}").fetchall() + size = _size_of(source) + con.execute(f"DROP VIEW IF EXISTS {_ident(table)}") + return ConvertedTable( + name=table, + file=target.name, + rows=int(rows), + columns=len(columns), + bytes=size, + compressed_bytes=compressed_bytes, + parquet_bytes=target.stat().st_size, + ) + + +def convert_to_parquet( + data_dir: Path, + destination: Path, + on_progress: Optional[ProgressFn] = None, + logger: Optional[logging.Logger] = None, + source: Optional[SourceFacts] = None, + on_table: Optional[Callable[[ConvertedTable, Path], None]] = None, + delete_sources: bool = False, +) -> list[ConvertedTable]: + """Convert every GTFS table under `data_dir` into `destination`. + + Tables are converted one at a time and the view is dropped after each, so at most + one source is ever open. `on_table` is called with each finished file, which lets a + caller publish and delete it before the next begins; `delete_sources` removes each + input once converted. Both default off, because a caller pointed at a directory it + does not own must not have its files deleted. + + `source` carries facts about where the feed came from that cannot be recovered from + `data_dir` - the archive's own size, and each member's compressed size. + """ + logger = logger or logging.getLogger(__name__) + destination.mkdir(parents=True, exist_ok=True) + + candidates = {} + for path in sorted(data_dir.iterdir()): + table = table_name_for(path) + if table is not None: + candidates[table] = path + if not candidates: + raise ValueError(f"No readable GTFS files found in {data_dir}") + + source = source or SourceFacts(bytes=_folder_bytes(candidates.values())) + con = open_connection() + try: + return _convert_each( + con, + candidates, + destination, + source, + on_progress, + on_table, + delete_sources, + logger, + ) + finally: + con.close() + + +def _convert_each( + con, + candidates: dict, + destination: Path, + source: SourceFacts, + on_progress, + on_table, + delete_sources: bool, + logger: logging.Logger, +) -> list[ConvertedTable]: + converted: list[ConvertedTable] = [] + tables = sorted(candidates) + for index, table in enumerate(tables, start=1): + origin = candidates[table] + if on_progress: + on_progress(PHASE_CONVERT, index, len(tables), table) + if not register_table(con, table, origin, logger): + continue + entry = convert_table( + con, table, origin, destination, source.compressed_sizes.get(origin.name) + ) + converted.append(entry) + if on_table: + on_table(entry, destination / entry.file) + if delete_sources: + origin.unlink(missing_ok=True) + + if not converted: + raise ValueError("No GTFS tables could be converted") + + write_manifest(destination, converted, source) + return converted + + +def _size_of(path: Path) -> Optional[int]: + try: + return path.stat().st_size + except OSError: + return None + + +def _folder_bytes(sources) -> int: + return sum(_size_of(path) or 0 for path in sources) + + +def write_manifest( + destination: Path, tables: list[ConvertedTable], source: SourceFacts +) -> Path: + """Describe the dataset for a reader that will never see what it came from. + + Version 2 added the sizes the load report shows. It is not a superset of version 1: + there, `bytes` meant the Parquet size, where here it means the source's and the + converted size is `parquet_bytes`. A reader has to know which it is holding, which + is what the version is for - do not add fields under the old number. + """ + sized = [table for table in tables if table.bytes is not None] + manifest = destination / MANIFEST + manifest.write_text( + json.dumps( + { + "version": MANIFEST_VERSION, + "generated_at": datetime.now(timezone.utc).isoformat( + timespec="seconds" + ), + # Our own converter, not gtfs-garage's version: this writes the same + # artifact but is not that package, and claiming its version would be + # a lie a reader might act on. + "converter_version": PARQUET_CONVERTER_VERSION, + "source": {"kind": source.kind, "bytes": source.bytes}, + "totals": { + "uncompressed_bytes": sum(table.bytes for table in sized), + "stored_bytes": sum(table.parquet_bytes for table in tables), + }, + "tables": [table.as_manifest_entry() for table in tables], + }, + indent=2, + ) + + "\n", + encoding="utf-8", + ) + return manifest diff --git a/functions-python/parquet_builder/src/main.py b/functions-python/parquet_builder/src/main.py new file mode 100644 index 000000000..e6577d8ce --- /dev/null +++ b/functions-python/parquet_builder/src/main.py @@ -0,0 +1,1039 @@ +# +# MobilityData 2026 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Render a GTFS dataset as Parquet and publish it. + +Invoked by Cloud Tasks with `{feed_stable_id, dataset_stable_id}`. Downloads the +dataset archive from the datasets bucket, converts every table, and uploads the result +to `//parquet/` so a browser can query it in place over range requests. + +Two things about this function are deliberate and easy to undo by accident: + + * It claims the dataset in the database before doing any work, and gives up quietly + if the claim is refused. Cloud Tasks delivers at least once, and a conversion is + expensive enough that running it twice concurrently matters. + * It returns HTTP 200 even when the build fails. Cloud Tasks retries non-2xx, and + nothing here fails transiently in a way a retry would fix - a corrupt archive is + still corrupt the second time. Failures are recorded in the tracking row, which is + what the API reports, rather than signalled by the status code. +""" + +import json +import logging +import os +import shutil +import sys +import zipfile +from datetime import datetime, timedelta, timezone +from pathlib import Path +from typing import Optional + +import flask +import functions_framework +from google.cloud import storage +from sqlalchemy.orm import Session + +from converter import ( + MANIFEST, + PARQUET_CONVERTER_VERSION, + ConvertedTable, + SourceFacts, + TailReader, + convert_table, + open_connection, + register_table, + table_name_for, + write_manifest, + zip_member_sizes_from_file, +) +from progress import ( + PHASE_CONVERT, + PHASE_DONE, + PHASE_START, + PHASE_SUMMARISE, + PHASE_UPLOAD, + ThrottledProgress, +) +from shared.common.gcp_memory_utils import limit_gcp_memory +from shared.database.database import with_db_session +from shared.database_gen.sqlacodegen_models import Gtfsdataset, TaskExecutionAttempt +from shared.helpers.ephemeral_workdir import EphemeralOrDebugWorkdir +from shared.helpers.logger import get_logger, init_logger +from shared.helpers.runtime_metrics import track_metrics +from shared.helpers.runtime_metrics import MemorySampler +from shared.helpers.parquet_policy import ( + DOWNSIZE_HEADROOM, + DOWNSIZE_STREAK, + SIZE_CONFIG_NAMESPACE, + SIZE_TIERS, +) +from shared.helpers.sizing import ( + Basis, + Size, + classify_failure, + demote, + escalate, + fits_within, + function_name, + record_size_override, + size_override, + tier_for, +) +from shared.helpers.task_execution.task_execution_tracker import TaskExecutionTracker +from shared.helpers.utils import create_http_parquet_builder_task + +init_logger() + +TASK_NAME = "parquet_generation" +BUILDER_BASE = "parquet-builder" +# Two escalations is the whole ladder: x -> m -> l. A third attempt would only repeat the +# largest worker, so the cap is the ladder's length rather than a tuning knob. +MAX_ATTEMPTS = 3 +PARQUET_PREFIX = "parquet" +WORKDIR_PREFIX = "parquet_" + +# Everything large goes here. It is a declared in-memory volume in deployed +# environments, so its size is already subtracted from the process budget below. +TMPDIR = os.getenv("PARQUET_TMPDIR", "/tmp/in-memory") + +# DuckDB's own default reads the host's RAM rather than the cgroup, so left alone it +# spills far too late to help. Sized well under what the limiter leaves us. +DUCKDB_MEMORY_LIMIT = os.getenv("PARQUET_DUCKDB_MEMORY_LIMIT", "2GB") + +# How much of an archive's tail to fetch when recovering compressed sizes. A GTFS zip +# has tens of members, so its central directory is a few kilobytes; 1 MiB is slack. +CENTRAL_DIRECTORY_TAIL = 1024 * 1024 +# What the archive reader buffers per range request. Buffered in process memory, not on +# the volume, so this trades requests against RLIMIT_AS rather than against the tmpfs. +# Well under the library's 40 MiB default, which would be charged to the heap in full. +ARCHIVE_CHUNK_SIZE = 8 * 1024 * 1024 + +# How long a generated set lives when the caller does not say. The single source of +# truth: the API deliberately declares bounds but no default, so an omitted value +# arrives here as None rather than as someone else's idea of 30. +DEFAULT_RETENTION_DAYS = 30 +MAX_RETENTION_DAYS = 60 + +# Must run before anything allocates: turns an overshoot into a catchable MemoryError +# instead of the kernel killing the container with no traceback and no response. +MEMORY_BUDGET = limit_gcp_memory(TMPDIR) +MIB = 1024**2 + + +def _retention_days(value) -> int: + """How long this set should live, defaulted and bounded. + + The API validates the range too, but this is not a duplicated check: the builder is + reachable from Cloud Tasks and from the CLI below, neither of which goes through the + schema, and a bad value here would otherwise be written into the tracking row and + honoured by the sweep. + """ + if value is None: + return DEFAULT_RETENTION_DAYS + try: + days = int(value) + except (TypeError, ValueError): + logging.warning("Ignoring unusable retention_days %r", value) + return DEFAULT_RETENTION_DAYS + if days < 1 or days > MAX_RETENTION_DAYS: + logging.warning( + "retention_days %s is outside 1..%s; using %s", + days, + MAX_RETENTION_DAYS, + DEFAULT_RETENTION_DAYS, + ) + return DEFAULT_RETENTION_DAYS + return days + + +def _budget() -> str: + """The worker's resource budget, for the log line that reports it per build.""" + budget = MEMORY_BUDGET + rlimit = getattr(budget, "rlimit_as_bytes", None) + volume = getattr(budget, "volume_bytes", None) + return ( + ( + f"{rlimit / MIB:.0f} MiB of address space" + if rlimit + else "no address-space limit" + ) + + ( + f", a {volume / MIB:.0f} MiB volume at {TMPDIR}" + if volume + else f", an unmeasured volume at {TMPDIR}" + ) + + f", and a DuckDB cap of {DUCKDB_MEMORY_LIMIT}" + ) + + +def _attempt_metadata(tables=None) -> dict: + """What this build had to work with, beside what it used. + + An attempt that records only usage cannot answer whether that usage was comfortable, + and reconstructing the budget afterwards from the deployment means trusting the + deployment has not moved since. `largest_member_bytes` is the measure the routing + table is written in, taken from the files the build actually opened rather than from + the database columns - those columns are exactly what an override exists to correct. + """ + budget = MEMORY_BUDGET + sizes = [table.bytes for table in (tables or []) if table.bytes] + return { + "rlimit_as_bytes": getattr(budget, "rlimit_as_bytes", None), + "volume_bytes": getattr(budget, "volume_bytes", None), + "largest_member_bytes": max(sizes) if sizes else None, + } + + +def _dataset_row(db_session, dataset_stable_id: str): + """The dataset, for its feed. None when it cannot be read.""" + try: + return ( + db_session.query(Gtfsdataset) + .filter(Gtfsdataset.stable_id == dataset_stable_id) + .one_or_none() + ) + except Exception: + return None + + +def _own_variant() -> Optional[Size]: + """Which worker this is, read from the name Cloud Run gave the service. + + `K_SERVICE` is exactly what `sizing.function_name` builds - `parquet-builder-m-dev` - + so the running size is recoverable without an env var per variant, which keeps the + pattern free of per-function wiring. + """ + service = os.getenv("K_SERVICE") or "" + for size in Size: + if service == function_name(BUILDER_BASE, size, os.getenv("ENVIRONMENT", "")): + return size + return None + + +def _escalate_after_failure( + tracker, + db_session, + dataset, + dataset_stable_id: str, + feed_stable_id: str, + variant: Optional[Size], + error: BaseException, + retention_days: int, + logger, +) -> Optional[Size]: + """Move the feed up a rung and re-queue, when the failure warrants it. + + Returns the size it escalated to, or None when it declined. Declining is the common + case and has three reasons, all terminal: the failure was not a resource one, the + build was already on the largest worker, or the attempt cap is spent. Each is a stop, + because the alternative is a queue loop that costs money. + """ + kind = classify_failure(error) + if not kind.is_resource: + return None + + attempts = tracker.attempts_since_success(dataset_stable_id) + target = escalate(variant, SIZE_TIERS, attempts=attempts, max_attempts=MAX_ATTEMPTS) + if target is None: + logger.warning( + "Not escalating %s: %s on %s after %s attempt(s)", + dataset_stable_id, + kind.value, + variant.value if variant else "an unknown worker", + attempts, + ) + return None + + feed = getattr(dataset, "feed", None) + if feed is None: + logger.warning( + "Cannot escalate %s: its feed is not resolvable", dataset_stable_id + ) + return None + + # A lock means what it says: this build is not retried anywhere, and the dataset is + # not produced until a person changes the value. Checked before anything is written + # or enqueued, and reported loudly, because nothing else will raise it. + locked = size_override(db_session, feed, SIZE_CONFIG_NAMESPACE) + if locked.locked: + logger.warning( + "Not escalating %s: feed %s is locked to %s, so %s will not be retried", + dataset_stable_id, + feed.stable_id, + locked.size.value if locked.size else "an unreadable size", + kind.value, + ) + return None + + # Only when it disagrees with the measurement: an override that merely restates the + # measured tier pins the feed to it for good, and this dataset is already on its way + # to `target` through the task below either way. + record_size_override( + db_session, + feed, + dataset, + target, + tiers=SIZE_TIERS, + namespace=SIZE_CONFIG_NAMESPACE, + ) + db_session.commit() + logger.info( + "Escalating %s from %s to %s after %s", + dataset_stable_id, + variant.value if variant else "unknown", + target.value, + kind.value, + ) + create_http_parquet_builder_task( + feed_stable_id, + dataset_stable_id, + force=True, + retention_days=retention_days, + size=target, + variant_basis=Basis.AUTO.value, + override=target.value, + ) + return target + + +def _recent_attempts(db_session, feed, limit: int) -> list: + """This feed's most recent attempts, newest first, across all of its datasets. + + Attempts are recorded against a dataset, but a size override belongs to the feed, so + the question "has this feed been comfortable lately" spans its datasets. `run_id` is + pinned to the converter version deliberately: a new converter is a different + workload, and evidence gathered under the old one should not carry over. + """ + dataset_ids = [ + row[0] + for row in db_session.query(Gtfsdataset.stable_id) + .filter(Gtfsdataset.feed_id == feed.id) + .all() + ] + if not dataset_ids: + return [] + return ( + db_session.query(TaskExecutionAttempt) + .filter( + TaskExecutionAttempt.task_name == TASK_NAME, + TaskExecutionAttempt.run_id == PARQUET_CONVERTER_VERSION, + TaskExecutionAttempt.entity_id.in_(dataset_ids), + ) + .order_by(TaskExecutionAttempt.finished_at.desc()) + .limit(limit) + .all() + ) + + +def _downsize_after_success( + db_session, dataset_stable_id: str, variant: Optional[Size], logger +) -> Optional[Size]: + """Lower this feed's override one rung when its recent builds say it can be. + + The mirror of `_escalate_after_failure`, and deliberately not its equal. Escalation + acts on a single failure, because being too small costs a build that cannot finish. + Being too large only costs money, so coming down waits for `DOWNSIZE_STREAK` builds + in a row that each left `DOWNSIZE_HEADROOM` of the smaller worker unused, on both + memory and disk. + + Nothing is re-queued. This build has already succeeded; the decision applies to the + feed's next one. + + Every failure here is swallowed: an economy must never cost a build that worked. + """ + try: + if variant is None: + return None + + target = demote(variant, SIZE_TIERS) + if target is None: + return None + + dataset = _dataset_row(db_session, dataset_stable_id) + feed = getattr(dataset, "feed", None) if dataset is not None else None + if feed is None: + return None + + current = size_override(db_session, feed, SIZE_CONFIG_NAMESPACE) + # No override means the feed already routes on its measurement, which is as low + # as it goes. An override that disagrees with the worker that just ran is a state + # this cannot reason about. A locked one is refused by `record_size_override`, + # but returning here keeps a locked feed out of the history query entirely. + if current.size is None or current.locked or current.size is not variant: + return None + + tier = tier_for(target, SIZE_TIERS) + if tier is None: + return None + + attempts = _recent_attempts(db_session, feed, DOWNSIZE_STREAK) + if len(attempts) < DOWNSIZE_STREAK: + return None + + for attempt in attempts: + metadata = attempt.metadata_ or {} + if ( + attempt.status != "completed" + or attempt.variant != variant.value + or not fits_within( + tier, + peak_vms_bytes=attempt.peak_vms_bytes, + largest_member_bytes=metadata.get("largest_member_bytes"), + headroom=DOWNSIZE_HEADROOM, + ) + ): + return None + + record_size_override( + db_session, + feed, + dataset, + target, + tiers=SIZE_TIERS, + namespace=SIZE_CONFIG_NAMESPACE, + ) + db_session.commit() + logger.info( + "Lowering %s from %s to %s: %s builds in a row under %.0f%% of %s", + feed.stable_id, + variant.value, + target.value, + DOWNSIZE_STREAK, + DOWNSIZE_HEADROOM * 100, + target.value, + ) + return target + except Exception: + logger.exception("Could not review the worker size for %s", dataset_stable_id) + db_session.rollback() + return None + + +@functions_framework.http +def build_parquet_handler(request: flask.Request) -> dict: + """Entrypoint for building the Parquet rendering of a GTFS dataset.""" + payload = request.get_json(silent=True) or {} + feed_stable_id = payload.get("feed_stable_id") + dataset_stable_id = payload.get("dataset_stable_id") + force = bool(payload.get("force", False)) + retention_days = _retention_days(payload.get("retention_days")) + # Why the Operations API routed this here. Carried for the attempt record only; the + # worker that actually ran it comes from K_SERVICE, which cannot be wrong. + variant_basis = payload.get("variant_basis") + override = payload.get("override") + + if not (feed_stable_id and dataset_stable_id): + return { + "status": "error", + "error": "Both feed_stable_id and dataset_stable_id must be defined.", + } + + if not dataset_stable_id.startswith(feed_stable_id): + return { + "status": "error", + "error": ( + f"feed_stable_id={feed_stable_id} is not a prefix of " + f"dataset_stable_id={dataset_stable_id}" + ), + } + + bucket_name = os.getenv("DATASETS_BUCKET_NAME") + if not bucket_name: + return { + "status": "error", + "error": "DATASETS_BUCKET_NAME environment variable is not defined.", + } + + try: + return build_parquet( + feed_stable_id=feed_stable_id, + dataset_stable_id=dataset_stable_id, + bucket_name=bucket_name, + force=force, + retention_days=retention_days, + variant_basis=variant_basis, + override=override, + ) + except Exception as error: + # Deliberately a 200: see the module docstring. + logging.exception("Failed to build Parquet for dataset %s", dataset_stable_id) + return { + "status": "error", + "error": f"Failed to build Parquet for dataset {dataset_stable_id}: {error}", + } + + +@with_db_session +@track_metrics(metrics=("time", "memory", "cpu")) +def build_parquet( + feed_stable_id: str, + dataset_stable_id: str, + bucket_name: str, + force: bool = False, + retention_days: int = DEFAULT_RETENTION_DAYS, + variant_basis: str = None, + override: str = None, + db_session: Session = None, +) -> dict: + """Claim the dataset, convert it, publish it, and record what was written.""" + logger = get_logger(build_parquet.__name__, dataset_stable_id) + started_at = datetime.now(timezone.utc) + variant = _own_variant() + tracker = TaskExecutionTracker( + task_name=TASK_NAME, + run_id=PARQUET_CONVERTER_VERSION, + db_session=db_session, + ) + tracker.start_run(params={"converter_version": PARQUET_CONVERTER_VERSION}) + + if force: + # Only ever moves a finished row back to claimable; it cannot take a claim away + # from a build that is currently running. + tracker.release_for_retry(dataset_stable_id) + + if not tracker.try_acquire( + dataset_stable_id, execution_ref=os.getenv("K_REVISION") + ): + logger.info("Dataset %s is already being built elsewhere", dataset_stable_id) + db_session.commit() + return { + "status": "skipped", + "reason": "already in progress", + "dataset": dataset_stable_id, + } + + # Started only once the claim is held, so a request that turns out to be a duplicate + # does not leave a sampler thread behind on a warm instance. + sampler = MemorySampler().start() + + # `limit_gcp_memory` logs the same figures, but it runs at module import - during the + # cold start, before any request exists - so those lines carry no trace and no + # dataset, and pairing them with a build means joining on the instance id across + # whatever else that instance has served since. Restating them here costs one line + # per build and puts the budget in the same trace as the failure it explains. + logger.info("Worker %s has %s", variant.value if variant else "unknown", _budget()) + + # Committed on its own: `ThrottledProgress` swallows failures from the callback + # below, so the claim cannot depend on the first progress write to reach the + # database. + db_session.commit() + + def publish(state: dict) -> None: + # Renews the claim as well as recording the reading, so a long build holds its + # lock by reporting rather than by a separate keepalive. + try: + tracker.heartbeat(dataset_stable_id, metadata=state) + db_session.commit() + except Exception: + # The caller logs and carries on, so leave the session usable. + db_session.rollback() + raise + + progress = ThrottledProgress(publish=publish, logger=logger) + progress.flush(phase=PHASE_START, done=0, total=0, detail="") + + try: + with EphemeralOrDebugWorkdir( + owner_prefix=WORKDIR_PREFIX, + dir=TMPDIR, + prefix=f"{dataset_stable_id}_", + ) as workdir_name: + workdir = Path(workdir_name) + bucket = storage.Client().get_bucket(bucket_name) + + plan = _plan_sources( + bucket, + feed_stable_id, + dataset_stable_id, + workdir, + progress, + logger, + db_session, + ) + # One timestamp for both halves of the expiry: the objects carry it as + # `customTime` for the bucket's lifecycle rule, and the row carries it so + # the API can stop advertising the set on the same date. + expires_at = datetime.now(timezone.utc) + timedelta(days=retention_days) + + tables, base_url = _convert_and_publish( + bucket=bucket, + feed_stable_id=feed_stable_id, + dataset_stable_id=dataset_stable_id, + workdir=workdir, + plan=plan, + progress=progress, + logger=logger, + expires_at=expires_at, + ) + + progress.flush(phase=PHASE_SUMMARISE, done=0, total=0, detail="") + metadata = { + "phase": PHASE_DONE, + "done": len(tables), + "total": len(tables), + "detail": "", + "base_url": base_url, + "retention_days": retention_days, + "expires_at": expires_at.isoformat(), + "tables": [table.as_manifest_entry() for table in tables], + } + tracker.mark_completed(dataset_stable_id, metadata=metadata) + tracker.record_attempt( + dataset_stable_id, + status="completed", + started_at=started_at, + variant=variant.value if variant else None, + variant_basis=variant_basis, + override_at_attempt=override, + metrics=sampler.stop(), + metadata=_attempt_metadata(tables), + ) + db_session.commit() + + # After the commit, so this build counts towards its own streak. + _downsize_after_success(db_session, dataset_stable_id, variant, logger) + + logger.info( + "Built %s Parquet tables for %s", len(tables), dataset_stable_id + ) + return { + "status": "success", + "dataset": dataset_stable_id, + "base_url": base_url, + "tables": [table.name for table in tables], + } + except Exception as error: + logger.exception("Parquet build failed for %s", dataset_stable_id) + # The failure may have come from a commit, which leaves the session refusing + # every statement until it is rolled back - `mark_failed` included. + db_session.rollback() + # Releases the claim as well as recording why, so the dataset can be retried + # without waiting out the lease. + try: + tracker.mark_failed(dataset_stable_id, error_message=str(error)) + db_session.commit() + except Exception: + logger.exception( + "Could not record the failure for %s; the claim will expire with its " + "lease", + dataset_stable_id, + ) + db_session.rollback() + + # Recording the attempt and escalating are best-effort: the build has already + # failed, and losing the original exception to a bookkeeping error would hide + # the thing worth reading. + escalated_to = None + try: + dataset = _dataset_row(db_session, dataset_stable_id) + escalated_to = _escalate_after_failure( + tracker, + db_session, + dataset, + dataset_stable_id, + feed_stable_id, + variant, + error, + retention_days, + logger, + ) + tracker.record_attempt( + dataset_stable_id, + status="failed", + started_at=started_at, + variant=variant.value if variant else None, + variant_basis=variant_basis, + override_at_attempt=override, + escalated_to=escalated_to.value if escalated_to else None, + failure_kind=classify_failure(error).value, + error=error, + metrics=sampler.stop(), + metadata=_attempt_metadata(), + ) + db_session.commit() + except Exception: + logger.exception("Could not record the attempt for %s", dataset_stable_id) + db_session.rollback() + raise + + +class SourcePlan: + """Where this build's CSVs come from, resolved once up front. + + `sources` is deliberately a list of (table, fetch) rather than files on disk: the + whole point is that only one source exists locally at a time, so materialising is + deferred to the moment each table is converted. + """ + + def __init__( + self, + kind: str, + source_bytes, + compressed_sizes: dict, + sources: list, + closer=None, + ): + self.facts = SourceFacts( + kind=kind, bytes=source_bytes, compressed_sizes=compressed_sizes + ) + self.sources = sources + # The archive plan reads members from a remote handle that has to outlive the + # plan, so releasing it is the plan's job rather than the planner's. + self._closer = closer + + def close(self) -> None: + if self._closer is not None: + self._closer() + self._closer = None + + +def _plan_sources( + bucket, + feed_stable_id: str, + dataset_stable_id: str, + workdir: Path, + progress, + logger, + db_session: Session, +) -> SourcePlan: + """Prefer the files `batch_process_dataset` already extracted; fall back to the zip. + + Reading `extracted/` avoids ever holding the archive and the whole feed at once. + The dataset's `Gtfsfile` rows are the index - they name every extracted file and + its size without a bucket listing - and their absence is the signal that this + dataset predates extraction (or that it failed), in which case the archive is the + only source there is. + """ + dataset = ( + db_session.query(Gtfsdataset) + .filter(Gtfsdataset.stable_id == dataset_stable_id) + .one_or_none() + ) + files = list(dataset.gtfsfiles) if dataset else [] + + if files: + plan = _plan_from_extracted( + bucket, feed_stable_id, dataset_stable_id, dataset, files, workdir, logger + ) + if plan is not None: + return plan + # The index named files the bucket does not have. Fall through to the archive, + # which is the authoritative copy. + + logger.info( + "No extracted files recorded for %s; falling back to the archive", + dataset_stable_id, + ) + return _plan_from_archive( + bucket, feed_stable_id, dataset_stable_id, workdir, logger + ) + + +def _plan_from_extracted( + bucket, feed_stable_id, dataset_stable_id, dataset, files, workdir, logger +) -> Optional[SourcePlan]: + """Plan from `extracted/`, or None when the bucket does not match the index. + + The `Gtfsfile` rows are an index, not a guarantee: `rebuild_missing_dataset_files` + exists because a dataset can be recorded with files the bucket no longer holds. A + missing object used to surface as a `NotFound` from inside the conversion, which + lost the whole build to one absent file, so the set is checked up front instead. + + Returning None rather than converting the subset is deliberate. The manifest + asserts it describes the dataset, so quietly dropping `feed_info` or `stop_times` + would publish a feed that looks complete and is not. + """ + prefix = f"{feed_stable_id}/{dataset_stable_id}/extracted" + local_dir = workdir / "extracted" + local_dir.mkdir(parents=True, exist_ok=True) + + # One listing rather than an existence check per file. + present = {blob.name for blob in bucket.list_blobs(prefix=prefix + "/")} + missing = [ + record.file_name + for record in files + if f"{prefix}/{record.file_name}" not in present + ] + if missing: + logger.warning( + "%s records %s extracted file(s) the bucket does not have (%s); using the " + "archive instead. Run the rebuild_missing_dataset_files task to repair the " + "extracted copies.", + dataset_stable_id, + len(missing), + ", ".join(sorted(missing)[:5]), + ) + return None + + sources = [] + for record in files: + # `file_name` is the archive-relative path, so a feed wrapped in a folder is + # recorded as "feed/stops.txt". The table is named after the basename. + name = Path(record.file_name).name + table = table_name_for(Path(name)) + if table is None: + continue + blob_path = f"{prefix}/{record.file_name}" + sources.append( + (table, _fetch_blob(bucket, blob_path, local_dir / name, logger)) + ) + + archive_blob = bucket.blob( + f"{feed_stable_id}/{dataset_stable_id}/{dataset_stable_id}.zip" + ) + return SourcePlan( + kind="zip", + source_bytes=_archive_size(dataset, archive_blob, logger), + compressed_sizes=_compressed_sizes_from_blob(archive_blob, logger), + sources=sorted(sources, key=lambda item: item[0]), + ) + + +def _plan_from_archive( + bucket, feed_stable_id, dataset_stable_id, workdir, logger +) -> SourcePlan: + """Fall back to the archive, read over the network rather than downloaded. + + The archive is never written to the workdir. `Blob.open("rb")` is a seekable + reader, which is all `zipfile` needs, so members are pulled with ranged requests as + the converter asks for them. Downloading it instead used to cost its full size on + the in-memory volume for the whole build, on top of the member being converted - + for a 1 GiB archive holding a 4 GiB table, enough to exhaust the volume before + conversion started. + """ + blob_path = f"{feed_stable_id}/{dataset_stable_id}/{dataset_stable_id}.zip" + archive_blob = bucket.blob(blob_path) + if not archive_blob.exists(): + raise FileNotFoundError( + f"Dataset archive not found at gs://{bucket.name}/{blob_path}" + ) + + archive_blob.reload() + handle = archive_blob.open("rb", chunk_size=ARCHIVE_CHUNK_SIZE) + try: + zf = zipfile.ZipFile(handle) + members = [member for member in zf.infolist() if not member.is_dir()] + except Exception: + handle.close() + raise + + data_dir = workdir / "extracted" + data_dir.mkdir(parents=True, exist_ok=True) + + sources = [] + # In archive order, not table order: a backward seek throws the reader's buffer + # away, so converting alphabetically would refetch most of the file. + for member in sorted(members, key=lambda m: m.header_offset): + # Flattened: producers differ on whether the files sit at the archive root or + # inside a folder, and the table is named after the basename either way. + name = Path(member.filename).name + table = table_name_for(Path(name)) + if table is None: + continue + sources.append((table, _extract_member(zf, member.filename, data_dir / name))) + + def close(): + zf.close() + handle.close() + + return SourcePlan( + kind="zip", + # The central directory is already in hand, so the compressed sizes cost + # nothing here - no second ranged read of the tail. + source_bytes=int(archive_blob.size or 0), + compressed_sizes={Path(m.filename).name: m.compress_size for m in members}, + sources=sources, + closer=close, + ) + + +def _extract_member(zf: zipfile.ZipFile, member_name: str, target: Path): + """A callable that unpacks one member when the converter is ready for it. + + Reads from the archive handle the plan holds open rather than reopening it, so a + remote archive is not re-read per member. + """ + + def fetch() -> Path: + target.parent.mkdir(parents=True, exist_ok=True) + with zf.open(member_name) as src, open(target, "wb") as dst: + shutil.copyfileobj(src, dst) + return target + + return fetch + + +def _fetch_blob(bucket, blob_path: str, target: Path, logger): + """A callable that downloads one object when the converter is ready for it.""" + + def fetch() -> Path: + blob = bucket.blob(blob_path) + target.parent.mkdir(parents=True, exist_ok=True) + blob.download_to_filename(str(target)) + return target + + return fetch + + +def _archive_size(dataset, archive_blob, logger): + """The archive's size, without downloading it.""" + recorded = getattr(dataset, "zipped_size_bytes", None) + if recorded: + return int(recorded) + try: + archive_blob.reload() + return int(archive_blob.size or 0) + except Exception as error: + logger.warning("Could not read the archive's size: %s", error) + return None + + +def _compressed_sizes_from_blob(archive_blob, logger) -> dict: + """Per-member compressed sizes, read from the archive's tail rather than all of it. + + These exist only in the zip's central directory, so reading `extracted/` would + otherwise lose them and the load report's Zipped column would be blank for every + dataset built this way. The directory sits at the end of the file, so a ranged read + of the last megabyte is enough. + """ + try: + archive_blob.reload() + size = int(archive_blob.size or 0) + if not size: + return {} + start = max(0, size - CENTRAL_DIRECTORY_TAIL) + tail = archive_blob.download_as_bytes(start=start, end=size - 1) + return zip_member_sizes_from_file(TailReader(tail, size)) + except Exception as error: + logger.warning("Could not read compressed sizes from the archive: %s", error) + return {} + + +def _convert_and_publish( + bucket, + feed_stable_id, + dataset_stable_id, + workdir, + plan, + progress, + logger, + expires_at, +): + """Convert and publish one table at a time, holding no more than one of each. + + Each table is fetched, converted, uploaded and deleted before the next begins, so + peak memory tracks the largest single table rather than the whole feed. The + manifest still goes up last - it is what tells a reader the set is complete. + """ + dest_prefix = f"{feed_stable_id}/{dataset_stable_id}/{PARQUET_PREFIX}" + existing = {blob.name for blob in bucket.list_blobs(prefix=dest_prefix + "/")} + written = set() + + out_dir = workdir / PARQUET_PREFIX + out_dir.mkdir(parents=True, exist_ok=True) + con = open_connection(temp_dir=workdir / "duckdb", memory_limit=DUCKDB_MEMORY_LIMIT) + tables: list[ConvertedTable] = [] + try: + total = len(plan.sources) + for index, (table, fetch) in enumerate(plan.sources, start=1): + progress(PHASE_CONVERT, index, total, table) + source = fetch() + try: + if not register_table(con, table, source, logger): + continue + entry = convert_table( + con, + table, + source, + out_dir, + plan.facts.compressed_sizes.get(source.name), + ) + finally: + source.unlink(missing_ok=True) + + parquet = out_dir / entry.file + written.add(_publish(bucket, dest_prefix, parquet, logger, expires_at)) + parquet.unlink(missing_ok=True) + tables.append(entry) + finally: + con.close() + plan.close() + + # Converted in archive order for the reader's sake; reported by name so the manifest + # does not depend on which path produced it. + tables.sort(key=lambda entry: entry.name) + + if not tables: + raise ValueError(f"No GTFS tables could be converted for {dataset_stable_id}") + + progress.flush(phase=PHASE_CONVERT, done=len(tables), total=len(tables)) + + progress(PHASE_UPLOAD, 1, 1, MANIFEST) + manifest = write_manifest(out_dir, tables, plan.facts) + written.add(_publish(bucket, dest_prefix, manifest, logger, expires_at)) + manifest.unlink(missing_ok=True) + progress.flush(phase=PHASE_UPLOAD, done=1, total=1, detail=MANIFEST) + + for name in sorted(existing - written): + bucket.blob(name).delete() + logger.info("Removed stale object %s", name) + + public_base = os.getenv("PUBLIC_HOSTED_DATASETS_URL", "").rstrip("/") + return tables, f"{public_base}/{dest_prefix}" + + +def _uniform_access(bucket) -> bool: + """True when the bucket grants access by policy, so per-object ACLs do not apply.""" + try: + return bool(bucket.iam_configuration.uniform_bucket_level_access_enabled) + except Exception: + return False + + +def _publish(bucket, dest_prefix: str, path: Path, logger, expires_at) -> str: + blob = bucket.blob(f"{dest_prefix}/{path.name}") + # The date the bucket's lifecycle rule deletes this object. Set on the upload, so + # it costs no extra request; GCS ignores it until a rule names customTime. + blob.custom_time = expires_at + blob.upload_from_filename(str(path)) + try: + blob.make_public() + except Exception as error: + # Only harmless under uniform bucket-level access, where ACLs do not apply and + # the objects are public by policy. Elsewhere the file is published but + # unreadable, and the build would report `ready` on a base_url that 403s. + if not _uniform_access(bucket): + raise + logger.debug( + "Skipping ACL on %s: bucket uses uniform access (%s)", blob.name, error + ) + return blob.name + + +def main(): # pragma: no cover + if len(sys.argv) < 2: + print("Usage: python src/main.py ") + sys.exit(1) + + dataset_stable_id = sys.argv[1] + feed_stable_id = dataset_stable_id.rsplit("-", 1)[0] + payload = {"feed_stable_id": feed_stable_id, "dataset_stable_id": dataset_stable_id} + + with flask.Flask(__name__).test_request_context(json=payload): + print(json.dumps(build_parquet_handler(flask.request), indent=2)) + + +if __name__ == "__main__": # pragma: no cover + main() diff --git a/functions-python/parquet_builder/src/progress.py b/functions-python/parquet_builder/src/progress.py new file mode 100644 index 000000000..4574d5eec --- /dev/null +++ b/functions-python/parquet_builder/src/progress.py @@ -0,0 +1,125 @@ +# +# MobilityData 2026 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Publishing how far a build has got, without writing to the database per event. + +The viewer polls roughly twice a second, and a multi-gigabyte feed emits an event per +zip member - thousands of them. Forwarding each one would mean thousands of +transactions competing with ordinary API traffic to serve a reading nobody can see +change that fast. + +So the writes are throttled to what a reader can actually perceive, with one exception +that matters: a phase change is always written immediately. Phases are what the viewer +turns into words ("Unzipping", "Converting"), and dropping one because it landed inside +a quiet window would leave the interface describing the wrong step. + +Each write also renews the build's claim on the dataset, so a long conversion keeps its +lock by reporting progress rather than by a separate keepalive. +""" + +from __future__ import annotations + +import logging +import time +from typing import Any, Callable, Optional + +# Phases, matching what the viewer already knows how to word. `upload` is publishing to +# the bucket; there is no phase for the local write, which happens under `convert`. +PHASE_START = "start" +PHASE_DOWNLOAD = "download" +PHASE_EXTRACT = "extract" +PHASE_CONVERT = "convert" +PHASE_UPLOAD = "upload" +PHASE_SUMMARISE = "summarise" +PHASE_DONE = "done" + +DEFAULT_MIN_INTERVAL_S = 1.0 + + +class ThrottledProgress: + """Collects progress events and publishes the ones worth publishing. + + `publish` is handed the full state each time rather than a delta, so a reader that + missed writes is never behind - it simply sees fewer intermediate readings. + """ + + def __init__( + self, + publish: Callable[[dict[str, Any]], None], + min_interval_s: float = DEFAULT_MIN_INTERVAL_S, + clock: Callable[[], float] = time.monotonic, + logger: Optional[logging.Logger] = None, + ): + self._publish = publish + self._min_interval_s = min_interval_s + self._clock = clock + self._logger = logger or logging.getLogger(__name__) + self._last_phase: Optional[str] = None + self._last_write_at: Optional[float] = None + self.last_state: Optional[dict[str, Any]] = None + + def __call__( + self, phase: str, done: int = 0, total: int = 0, detail: str = "" + ) -> None: + """A `ProgressFn`, so it can be passed straight to the converter.""" + state = { + "phase": phase, + "done": int(done), + "total": int(total), + "detail": detail, + } + self.last_state = state + + now = self._clock() + phase_changed = phase != self._last_phase + due = ( + self._last_write_at is None + or (now - self._last_write_at) >= self._min_interval_s + ) + if not (phase_changed or due): + return + + self._last_phase = phase + self._last_write_at = now + self._write(state) + + def flush(self, **overrides: Any) -> None: + """Publish the current reading regardless of the interval. + + For the end of a phase whose final count would otherwise be dropped, and for + the terminal states, where the last thing written is what a reader is left + looking at. + """ + state = { + **( + self.last_state + or {"phase": PHASE_START, "done": 0, "total": 0, "detail": ""} + ), + **overrides, + } + self.last_state = state + self._last_phase = state.get("phase") + self._last_write_at = self._clock() + self._write(state) + + def _write(self, state: dict[str, Any]) -> None: + try: + self._publish(state) + except Exception as error: + # Progress is a courtesy; a build must not fail because a reading could not + # be recorded. The claim is renewed by these same writes, so a sustained + # outage costs the lease - which is the correct outcome, since a worker + # that cannot reach the database cannot report failure either. + self._logger.warning("Could not publish progress %s: %s", state, error) diff --git a/functions-python/parquet_builder/src/scripts/generate_local.py b/functions-python/parquet_builder/src/scripts/generate_local.py new file mode 100644 index 000000000..e3878beee --- /dev/null +++ b/functions-python/parquet_builder/src/scripts/generate_local.py @@ -0,0 +1,340 @@ +# +# MobilityData 2026 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Produce a dataset's Parquet locally, with no GCP involved. + +The cloud function needs a bucket, a database, a Cloud Tasks queue and credentials. +None of that is interesting when the thing being worked on is the viewer, so this does +the same conversion against a feed on disk or a public URL and leaves the result +somewhere a browser can read it. + +The conversion itself is imported, not reimplemented - `converter.convert_to_parquet`, +exactly what the function calls - so what comes out here is what would come out of a +real build. If the two ever disagree, that is a bug in one of them, not a difference +between local and deployed. + +`--serve` exists because the obvious way to serve the output does not work: the reader +queries Parquet over HTTP range requests, and `python -m http.server` ignores Range and +returns whole files. The server here answers ranges and sends the CORS headers a +cross-origin worker needs. +""" + +from __future__ import annotations + +import argparse +import functools +import http.server +import os +import re +import shutil +import socketserver +import sys +import tempfile +import urllib.error +import urllib.request +from pathlib import Path +from typing import Optional + +from converter import ( + MANIFEST, + SourceFacts, + convert_to_parquet, + extract_feed, + zip_member_sizes, +) + +# Datasets are public over HTTPS, which is what lets this avoid credentials entirely. +DEFAULT_HOST = "https://files.mobilitydatabase.org" +ENV_HOSTS = { + "prod": "https://files.mobilitydatabase.org", + "qa": "https://qa-files.mobilitydatabase.org", + "dev": "https://dev-files.mobilitydatabase.org", +} + +# A dataset id is its feed id with a timestamp appended: mdb-1210-202402121801. +DATASET_ID = re.compile(r"^(?P.+)-(?P\d{8,})$") + + +def _log(message: str) -> None: + print(message, file=sys.stderr, flush=True) + + +def _progress(phase: str, done: int, total: int, detail: str) -> None: + total_text = f"/{total}" if total else "" + _log(f" {phase:9} {done}{total_text} {detail}") + + +def archive_url(stable_id: str, host: str) -> str: + """Where a feed or dataset id's archive lives. + + A dataset id addresses one specific archive; a feed id addresses whatever is + current, which is the more useful thing to type. + """ + match = DATASET_ID.match(stable_id) + if match: + feed = match.group("feed") + return f"{host}/{feed}/{stable_id}/{stable_id}.zip" + return f"{host}/{stable_id}/latest.zip" + + +def download(url: str, target: Path) -> Path: + _log(f" download {url}") + try: + with urllib.request.urlopen(url) as response, open(target, "wb") as out: + shutil.copyfileobj(response, out) + except urllib.error.HTTPError as error: + raise SystemExit( + f"Could not download {url}: {error.code} {error.reason}\n" + "Check the id, or pass --env dev/qa if the feed is not in production." + ) + except urllib.error.URLError as error: + raise SystemExit(f"Could not reach {url}: {error.reason}") + return target + + +def _from_archive(archive: Path, workdir: Path) -> tuple[Path, SourceFacts]: + """Unpack, and record what the archive can tell us before it is discarded. + + The compressed sizes have to be read here: they exist only in the zip's directory, + so a manifest written without them reports a dash in the Zipped column for every + file - which is how a local build ends up describing itself differently from a real + one, for no reason a reader could see. + """ + facts = SourceFacts( + kind="zip", + bytes=archive.stat().st_size, + compressed_sizes=zip_member_sizes(archive), + ) + return extract_feed(archive, workdir / "extracted", on_progress=_progress), facts + + +def resolve_source( + source: str, workdir: Path, host: str +) -> tuple[Path, Optional[SourceFacts]]: + """Get the feed onto disk as a directory of GTFS files, however it was named. + + Returns the directory and, for an archive, the facts only the archive carries. + None means a plain folder, which the converter describes for itself. + """ + candidate = Path(source).expanduser() + + if candidate.is_dir(): + _log(f" source {candidate} (directory)") + return candidate, None + + if candidate.is_file(): + _log(f" source {candidate} (archive)") + return _from_archive(candidate, workdir) + + if candidate.suffix == ".zip" or "/" in source or "\\" in source: + # Looks like a path the user expected to exist, rather than an id. + raise SystemExit(f"No such file or directory: {source}") + + archive = download(archive_url(source, host), workdir / f"{source}.zip") + return _from_archive(archive, workdir) + + +def generate(source: str, out_dir: Path, host: str, keep: bool) -> list: + out_dir = out_dir.expanduser().resolve() + if out_dir.exists() and not keep: + # Cleared so a regeneration cannot leave a table behind that the feed no longer + # has - the same reason the cloud function clears its prefix before uploading. + shutil.rmtree(out_dir) + + with tempfile.TemporaryDirectory(prefix="parquet-local-") as tmp: + workdir = Path(tmp) + data_dir, facts = resolve_source(source, workdir, host) + tables = convert_to_parquet( + data_dir, out_dir, on_progress=_progress, source=facts + ) + + total = sum(t.bytes for t in tables) + _log("") + _log( + f"Wrote {len(tables)} tables and a manifest to {out_dir} ({total / 1e6:.1f} MB)" + ) + for table in tables: + _log(f" {table.name:24} {table.rows:>9,} rows {table.bytes / 1e3:>8.1f} kB") + return tables + + +class RangeRequestHandler(http.server.SimpleHTTPRequestHandler): + """A static handler that answers Range requests and allows cross-origin reads. + + Both are required by the reader and neither is provided by the stdlib handler: it + ignores Range entirely, so a query that asks for a few kilobytes of a footer is + answered with the whole file, and a worker on another origin is refused outright. + """ + + def end_headers(self): + self.send_header("Access-Control-Allow-Origin", "*") + self.send_header("Access-Control-Allow-Headers", "range") + self.send_header( + "Access-Control-Expose-Headers", + "content-range, content-length, accept-ranges, etag", + ) + self.send_header("Accept-Ranges", "bytes") + self.send_header("Cache-Control", "no-store") + super().end_headers() + + def do_OPTIONS(self): # noqa: N802 - name fixed by the base class + self.send_response(204) + self.end_headers() + + def do_GET(self): # noqa: N802 - name fixed by the base class + requested = self.headers.get("Range") + if not requested: + return super().do_GET() + + match = re.match(r"bytes=(\d*)-(\d*)", requested.strip()) + if not match: + return super().do_GET() + + path = self.translate_path(self.path) + try: + size = os.path.getsize(path) + except OSError: + return self.send_error(404, "File not found") + + start_text, end_text = match.groups() + if start_text: + start = int(start_text) + end = int(end_text) if end_text else size - 1 + else: + # A suffix range - "the last N bytes" - which is how a Parquet footer is + # read, so getting this wrong breaks every query. + length = int(end_text or 0) + start, end = max(0, size - length), size - 1 + end = min(end, size - 1) + + if start > end or start >= size: + self.send_response(416) + self.send_header("Content-Range", f"bytes */{size}") + self.end_headers() + return + + self.send_response(206) + self.send_header("Content-Type", self.guess_type(path)) + self.send_header("Content-Range", f"bytes {start}-{end}/{size}") + self.send_header("Content-Length", str(end - start + 1)) + self.end_headers() + with open(path, "rb") as handle: + handle.seek(start) + self.wfile.write(handle.read(end - start + 1)) + + def log_message(self, fmt, *args): + _log(" " + fmt % args) + + +def serve(directory: Path, port: int) -> None: + handler = functools.partial(RangeRequestHandler, directory=str(directory)) + + class Server(socketserver.ThreadingTCPServer): + allow_reuse_address = True + daemon_threads = True + + with Server(("", port), handler) as httpd: + _log("") + _log(f"Serving {directory} at http://localhost:{port} (ranges and CORS on)") + _log(f" manifest: http://localhost:{port}/{MANIFEST}") + _log(" Ctrl-C to stop") + try: + httpd.serve_forever() + except KeyboardInterrupt: + _log("stopped") + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + prog="parquet-generate-local.sh", + description="Convert a GTFS feed to Parquet locally, with no GCP access.", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog="""examples: + # A feed id: downloads its current archive over HTTPS + parquet-generate-local.sh mdb-1210 + + # A specific dataset + parquet-generate-local.sh mdb-1210-202402121801 + + # A feed already on disk, as a zip or an unpacked folder + parquet-generate-local.sh ./gtfs.zip + parquet-generate-local.sh ./extracted/ + + # Straight into the operations web app, which serves public/datasets/ + parquet-generate-local.sh mdb-1210 --out ../ops-web/public/datasets/mdb-1210 + + # Or serve it directly, with the range support the reader needs + parquet-generate-local.sh mdb-1210 --serve +""", + ) + parser.add_argument( + "source", help="a feed or dataset stable id, a .zip path, or a directory" + ) + parser.add_argument( + "--out", + type=Path, + help="where to write (default: .dist/parquet/)", + ) + parser.add_argument( + "--env", + choices=sorted(ENV_HOSTS), + default="prod", + help="which environment to download a stable id from (default: prod)", + ) + parser.add_argument( + "--serve", + nargs="?", + type=int, + const=8090, + metavar="PORT", + help="serve the result on PORT (default 8090) with ranges and CORS", + ) + parser.add_argument( + "--keep", + action="store_true", + help="add to an existing output directory instead of clearing it first", + ) + return parser + + +def main(argv=None) -> int: + args = build_parser().parse_args(argv) + + name = Path(args.source).name.removesuffix(".zip") or "feed" + out_dir = args.out or Path(".dist/parquet") / name + + try: + tables = generate( + source=args.source, + out_dir=out_dir, + host=ENV_HOSTS[args.env], + keep=args.keep, + ) + except ValueError as error: + # A feed that cannot be converted is a normal thing to type by mistake - a + # wrong folder, an archive that is not GTFS - so it gets a sentence rather + # than a traceback. + _log(f"error: {error}") + return 1 + if not tables: + return 1 + + if args.serve is not None: + serve(out_dir.expanduser().resolve(), args.serve) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/functions-python/parquet_builder/tests/test_converter.py b/functions-python/parquet_builder/tests/test_converter.py new file mode 100644 index 000000000..ef9a61852 --- /dev/null +++ b/functions-python/parquet_builder/tests/test_converter.py @@ -0,0 +1,375 @@ +# +# MobilityData 2026 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""The properties a Parquet reader depends on, pinned. + +These exist because the conversion is reproduced here rather than imported from +gtfs-garage, so nothing else would notice if the output drifted from what its reader +expects. Two of them fail silently rather than loudly if broken, which is the whole +reason they are asserted: a typed column breaks filtering but still renders, and an +empty string where a NULL belongs makes "is empty" quietly miss rows. +""" + +import json +from pathlib import Path + +import pytest + +from converter import ( + MANIFEST, + SourceFacts, + _ident, + convert_to_parquet, + extract_feed, + zip_member_sizes, +) + +duckdb = pytest.importorskip("duckdb") + + +AGENCY = "agency_id,agency_name,agency_url,agency_timezone\n1,Test Transit,https://example.org,UTC\n" +STOPS = ( + "stop_id,stop_name,stop_lat,stop_lon,parent_station\n" + "S1,First,45.5,-73.6,\n" + "S2,Second,45.6,-73.7,S1\n" +) +ROUTES = "route_id,agency_id,route_short_name,route_type\nR1,1,1,3\n" + +LOCATIONS = json.dumps( + { + "type": "FeatureCollection", + "features": [ + { + "id": "zone-a", + "properties": {"stop_name": "Zone A", "stop_desc": "on demand"}, + "geometry": { + "type": "Polygon", + "coordinates": [[[0, 0], [1, 0], [1, 1], [0, 0]]], + }, + }, + { + "id": "zone-b", + "properties": {"stop_name": "Zone B"}, + "geometry": { + "type": "MultiPolygon", + "coordinates": [[[[0, 0], [2, 0], [2, 2], [0, 0]]]], + }, + }, + ], + } +) + + +@pytest.fixture +def feed(tmp_path) -> Path: + data = tmp_path / "extracted" + data.mkdir() + (data / "agency.txt").write_text(AGENCY) + (data / "stops.txt").write_text(STOPS) + (data / "routes.txt").write_text(ROUTES) + return data + + +def _describe(path: Path) -> dict: + con = duckdb.connect(database=":memory:") + try: + rows = con.execute(f"DESCRIBE SELECT * FROM read_parquet('{path}')").fetchall() + return {row[0]: row[1] for row in rows} + finally: + con.close() + + +def test_one_parquet_per_table_named_after_the_file_stem(feed, tmp_path): + out = tmp_path / "parquet" + tables = convert_to_parquet(feed, out) + + assert [t.name for t in tables] == ["agency", "routes", "stops"] + assert sorted(p.name for p in out.glob("*.parquet")) == [ + "agency.parquet", + "routes.parquet", + "stops.parquet", + ] + + +def test_a_table_name_cannot_escape_its_identifier(feed, tmp_path): + """Table names come from archive member filenames, which a producer controls.""" + hostile = 'stops" AS SELECT 1; DROP TABLE agency; --' + (feed / f"{hostile}.txt").write_text(AGENCY) + + out = tmp_path / "parquet" + tables = convert_to_parquet(feed, out) + + assert hostile in [t.name for t in tables] + assert "agency" in [t.name for t in tables] + assert (out / f"{hostile}.parquet").exists() + + +def test_ident_doubles_embedded_quotes(): + assert _ident("stops") == '"stops"' + assert _ident('a"b') == '"a""b"' + + +def test_every_column_is_text(feed, tmp_path): + """The reader filters with ILIKE and `= ''`. A typed column breaks that silently.""" + out = tmp_path / "parquet" + convert_to_parquet(feed, out) + + for parquet in out.glob("*.parquet"): + types = set(_describe(parquet).values()) + assert types == {"VARCHAR"}, f"{parquet.name} has non-text columns: {types}" + + +def test_empty_fields_become_null_not_empty_string(feed, tmp_path): + """The reader rewrites `= ''` into `IS NULL OR = ''` on this basis.""" + out = tmp_path / "parquet" + convert_to_parquet(feed, out) + + con = duckdb.connect(database=":memory:") + try: + nulls = con.execute( + f"SELECT count(*) FROM read_parquet('{out / 'stops.parquet'}') " + "WHERE parent_station IS NULL" + ).fetchone()[0] + empties = con.execute( + f"SELECT count(*) FROM read_parquet('{out / 'stops.parquet'}') " + "WHERE parent_station = ''" + ).fetchone()[0] + finally: + con.close() + + assert nulls == 1 + assert empties == 0 + + +def test_an_unreadable_file_is_skipped_not_fatal(feed, tmp_path): + """Most GTFS files are optional; one bad extra must not cost the whole feed.""" + (feed / "broken.txt").write_text("") + out = tmp_path / "parquet" + + tables = convert_to_parquet(feed, out) + + assert "agency" in [t.name for t in tables] + assert not (out / "broken.parquet").exists() + + +def test_a_header_only_file_is_still_a_table(feed, tmp_path): + """An optional file present but empty of rows is a real, empty GTFS table.""" + (feed / "frequencies.txt").write_text("trip_id,start_time,end_time,headway_secs\n") + out = tmp_path / "parquet" + + tables = {t.name: t for t in convert_to_parquet(feed, out)} + + assert "frequencies" in tables + assert tables["frequencies"].rows == 0 + assert list(_describe(out / "frequencies.parquet")) == [ + "trip_id", + "start_time", + "end_time", + "headway_secs", + ] + + +def test_apple_double_resource_forks_are_not_tables(feed, tmp_path): + """`__MACOSX/._stops.txt` flattens to something that looks like a table.""" + (feed / "._stops.txt").write_text(STOPS) + (feed / ".DS_Store").write_text("junk") + out = tmp_path / "parquet" + + tables = convert_to_parquet(feed, out) + + assert [t.name for t in tables] == ["agency", "routes", "stops"] + assert not (out / "._stops.parquet").exists() + + +def test_locations_geojson_becomes_a_table_with_stable_columns(feed, tmp_path): + """Polygon and MultiPolygon in one feed must not change the column set.""" + (feed / "locations.geojson").write_text(LOCATIONS) + out = tmp_path / "parquet" + + tables = convert_to_parquet(feed, out) + + assert "locations" in [t.name for t in tables] + columns = _describe(out / "locations.parquet") + assert list(columns) == [ + "id", + "stop_name", + "stop_desc", + "geometry_type", + "geometry", + ] + assert set(columns.values()) == {"VARCHAR"} + + con = duckdb.connect(database=":memory:") + try: + rows = con.execute( + f"SELECT id, geometry_type FROM read_parquet('{out / 'locations.parquet'}') ORDER BY id" + ).fetchall() + finally: + con.close() + assert rows == [("zone-a", "Polygon"), ("zone-b", "MultiPolygon")] + + +def test_malformed_locations_geojson_is_ignored(feed, tmp_path): + """A broken optional file loses the zones, not the feed.""" + (feed / "locations.geojson").write_text("{not json at all") + out = tmp_path / "parquet" + + tables = convert_to_parquet(feed, out) + + assert "locations" not in [t.name for t in tables] + assert "stops" in [t.name for t in tables] + + +def test_manifest_describes_what_was_written(feed, tmp_path): + out = tmp_path / "parquet" + tables = convert_to_parquet(feed, out) + + manifest = json.loads((out / MANIFEST).read_text()) + + assert manifest["version"] == 2 + assert [t["name"] for t in manifest["tables"]] == [t.name for t in tables] + stops = next(t for t in manifest["tables"] if t["name"] == "stops") + assert stops["file"] == "stops.parquet" + assert stops["rows"] == 2 + assert stops["columns"] == 5 + + +def test_bytes_is_the_source_size_and_parquet_bytes_the_converted_one(feed, tmp_path): + """The v1 -> v2 trap: `bytes` used to mean the Parquet size and now means the + source's. Asserting both, and that they differ, is what catches a swap.""" + out = tmp_path / "parquet" + convert_to_parquet(feed, out) + + manifest = json.loads((out / MANIFEST).read_text()) + stops = next(t for t in manifest["tables"] if t["name"] == "stops") + + assert stops["bytes"] == (feed / "stops.txt").stat().st_size + assert stops["parquet_bytes"] == (out / "stops.parquet").stat().st_size + assert stops["bytes"] != stops["parquet_bytes"] + + +def test_the_totals_add_up(feed, tmp_path): + out = tmp_path / "parquet" + convert_to_parquet(feed, out) + + manifest = json.loads((out / MANIFEST).read_text()) + entries = manifest["tables"] + + assert manifest["totals"]["stored_bytes"] == sum( + t["parquet_bytes"] for t in entries + ) + assert manifest["totals"]["uncompressed_bytes"] == sum(t["bytes"] for t in entries) + + +def test_a_folder_source_reports_no_compressed_sizes(feed, tmp_path): + """A member's compressed size exists only inside an archive.""" + out = tmp_path / "parquet" + convert_to_parquet(feed, out) + + manifest = json.loads((out / MANIFEST).read_text()) + + assert manifest["source"]["kind"] == "folder" + assert manifest["source"]["bytes"] > 0 + assert all(t["compressed_bytes"] is None for t in manifest["tables"]) + + +def test_a_zip_source_records_what_each_file_weighed_inside_it(tmp_path): + """The Zipped column can never come back if this is not captured up front.""" + import zipfile + + archive = tmp_path / "feed.zip" + with zipfile.ZipFile(archive, "w", zipfile.ZIP_DEFLATED) as zf: + # Repetitive content so compression is unambiguously smaller. + zf.writestr("stops.txt", "stop_id,stop_name\n" + "S1,First\n" * 500) + zf.writestr("agency.txt", AGENCY) + + data_dir = extract_feed(archive, tmp_path / "extracted") + out = tmp_path / "parquet" + convert_to_parquet( + data_dir, + out, + source=SourceFacts( + kind="zip", + bytes=archive.stat().st_size, + compressed_sizes=zip_member_sizes(archive), + ), + ) + + manifest = json.loads((out / MANIFEST).read_text()) + assert manifest["source"]["kind"] == "zip" + assert manifest["source"]["bytes"] == archive.stat().st_size + + stops = next(t for t in manifest["tables"] if t["name"] == "stops") + assert stops["compressed_bytes"] > 0 + assert stops["compressed_bytes"] < stops["bytes"], "a zipped file should be smaller" + + +def test_zip_member_sizes_is_empty_for_a_non_archive(tmp_path): + not_a_zip = tmp_path / "feed.zip" + not_a_zip.write_text("definitely not a zip") + + assert zip_member_sizes(not_a_zip) == {} + + +def test_row_counts_are_reported(feed, tmp_path): + tables = {t.name: t for t in convert_to_parquet(feed, tmp_path / "parquet")} + + assert tables["stops"].rows == 2 + assert tables["agency"].rows == 1 + + +def test_progress_is_reported_once_per_table(feed, tmp_path): + seen = [] + convert_to_parquet( + feed, tmp_path / "parquet", on_progress=lambda *args: seen.append(args) + ) + + assert [s[0] for s in seen] == ["convert"] * 3 + assert [s[1] for s in seen] == [1, 2, 3] + assert all(s[2] == 3 for s in seen) + assert [s[3] for s in seen] == ["agency", "routes", "stops"] + + +def test_a_feed_with_nothing_readable_is_an_error(tmp_path): + """Every table failing means a wrong location, not a feed with no files.""" + empty = tmp_path / "extracted" + empty.mkdir() + + with pytest.raises(ValueError, match="No readable GTFS files"): + convert_to_parquet(empty, tmp_path / "parquet") + + +def test_output_reopens_as_a_parquet_feed(feed, tmp_path): + """The round trip the viewer performs: open the directory, get the same tables.""" + out = tmp_path / "parquet" + convert_to_parquet(feed, out) + + con = duckdb.connect(database=":memory:") + try: + for parquet in out.glob("*.parquet"): + table = parquet.stem + con.execute( + f"""CREATE VIEW "{table}" AS SELECT * FROM read_parquet('{parquet}')""" + ) + names = {r[0] for r in con.execute("SHOW TABLES").fetchall()} + assert names == {"agency", "routes", "stops"} + assert con.execute( + 'SELECT stop_name FROM "stops" ORDER BY stop_id' + ).fetchall() == [ + ("First",), + ("Second",), + ] + finally: + con.close() diff --git a/functions-python/parquet_builder/tests/test_generate_local.py b/functions-python/parquet_builder/tests/test_generate_local.py new file mode 100644 index 000000000..cbfc8667f --- /dev/null +++ b/functions-python/parquet_builder/tests/test_generate_local.py @@ -0,0 +1,295 @@ +# +# MobilityData 2026 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""The local generator: which archive an id means, and the server's range handling. + +The range tests are the point of the file. The reader queries Parquet by asking for +byte ranges, and a server that quietly answers with the whole file looks like it works +until a query returns nothing useful - so the suffix range that reads a footer is +asserted explicitly. +""" + +import json +import threading +import unittest +import urllib.error +import urllib.request +import zipfile +from functools import partial +from http.server import HTTPServer +from pathlib import Path +from tempfile import TemporaryDirectory + +import pytest + +from scripts.generate_local import ( + ENV_HOSTS, + RangeRequestHandler, + archive_url, + build_parser, + generate, +) + +pytest.importorskip("duckdb") + +AGENCY = "agency_id,agency_name,agency_url,agency_timezone\n1,T,https://e.org,UTC\n" +STOPS = "stop_id,stop_name\nS1,First\n" + + +class TestArchiveUrl(unittest.TestCase): + def test_a_dataset_id_addresses_that_exact_archive(self): + self.assertEqual( + archive_url("mdb-1210-202402121801", ENV_HOSTS["prod"]), + "https://files.mobilitydatabase.org/mdb-1210/mdb-1210-202402121801/mdb-1210-202402121801.zip", + ) + + def test_a_feed_id_addresses_whatever_is_current(self): + self.assertEqual( + archive_url("mdb-1210", ENV_HOSTS["prod"]), + "https://files.mobilitydatabase.org/mdb-1210/latest.zip", + ) + + def test_the_environment_selects_the_host(self): + self.assertTrue( + archive_url("mdb-1", ENV_HOSTS["dev"]).startswith( + "https://dev-files.mobilitydatabase.org/" + ) + ) + + def test_a_feed_id_containing_digits_is_not_mistaken_for_a_dataset(self): + """`mdb-1210` ends in digits too; only a timestamp-length run means a dataset.""" + self.assertTrue( + archive_url("mdb-1210", ENV_HOSTS["prod"]).endswith("/latest.zip") + ) + + +class TestGenerate(unittest.TestCase): + def _feed_zip(self, path: Path, nested: bool = False) -> Path: + prefix = "feed/" if nested else "" + with zipfile.ZipFile(path, "w") as zf: + zf.writestr(f"{prefix}agency.txt", AGENCY) + zf.writestr(f"{prefix}stops.txt", STOPS) + return path + + def test_converts_a_local_archive(self): + with TemporaryDirectory() as tmp: + tmp = Path(tmp) + archive = self._feed_zip(tmp / "feed.zip") + out = tmp / "out" + + tables = generate(str(archive), out, ENV_HOSTS["prod"], keep=False) + + self.assertEqual([t.name for t in tables], ["agency", "stops"]) + self.assertTrue((out / "manifest.json").exists()) + + def test_an_archive_source_is_described_as_one(self): + """A local build must describe itself exactly as a real one does. + + Regression: the script unpacked the zip and then converted the resulting folder + without passing the archive's facts, so every manifest it wrote claimed + `kind: folder` and lost the zipped sizes - a difference between local and + deployed output with no visible cause. + """ + with TemporaryDirectory() as tmp: + tmp = Path(tmp) + archive = tmp / "feed.zip" + with zipfile.ZipFile(archive, "w", zipfile.ZIP_DEFLATED) as zf: + zf.writestr("stops.txt", "stop_id,stop_name\n" + "S1,First\n" * 500) + out = tmp / "out" + + generate(str(archive), out, ENV_HOSTS["prod"], keep=False) + + manifest = json.loads((out / "manifest.json").read_text()) + assert manifest["source"]["kind"] == "zip" + assert manifest["source"]["bytes"] == archive.stat().st_size + stops = manifest["tables"][0] + assert stops["compressed_bytes"] < stops["bytes"] + + def test_a_directory_source_is_described_as_one(self): + with TemporaryDirectory() as tmp: + tmp = Path(tmp) + data = tmp / "extracted" + data.mkdir() + (data / "agency.txt").write_text(AGENCY) + out = tmp / "out" + + generate(str(data), out, ENV_HOSTS["prod"], keep=False) + + manifest = json.loads((out / "manifest.json").read_text()) + assert manifest["source"]["kind"] == "folder" + assert manifest["tables"][0]["compressed_bytes"] is None + + def test_converts_a_directory(self): + with TemporaryDirectory() as tmp: + tmp = Path(tmp) + data = tmp / "extracted" + data.mkdir() + (data / "agency.txt").write_text(AGENCY) + out = tmp / "out" + + tables = generate(str(data), out, ENV_HOSTS["prod"], keep=False) + + self.assertEqual([t.name for t in tables], ["agency"]) + + def test_a_nested_archive_is_flattened(self): + with TemporaryDirectory() as tmp: + tmp = Path(tmp) + archive = self._feed_zip(tmp / "feed.zip", nested=True) + + tables = generate(str(archive), tmp / "out", ENV_HOSTS["prod"], keep=False) + + self.assertEqual([t.name for t in tables], ["agency", "stops"]) + + def test_regenerating_clears_a_table_the_feed_no_longer_has(self): + with TemporaryDirectory() as tmp: + tmp = Path(tmp) + out = tmp / "out" + out.mkdir() + (out / "gone.parquet").write_text("stale") + + generate( + str(self._feed_zip(tmp / "feed.zip")), + out, + ENV_HOSTS["prod"], + keep=False, + ) + + self.assertFalse((out / "gone.parquet").exists()) + + def test_keep_leaves_an_existing_directory_alone(self): + with TemporaryDirectory() as tmp: + tmp = Path(tmp) + out = tmp / "out" + out.mkdir() + (out / "keep.parquet").write_text("mine") + + generate( + str(self._feed_zip(tmp / "feed.zip")), out, ENV_HOSTS["prod"], keep=True + ) + + self.assertTrue((out / "keep.parquet").exists()) + + def test_a_missing_path_is_reported_not_treated_as_an_id(self): + with TemporaryDirectory() as tmp: + with self.assertRaises(SystemExit) as caught: + generate( + "./no/such/feed.zip", + Path(tmp) / "out", + ENV_HOSTS["prod"], + keep=False, + ) + self.assertIn("No such file", str(caught.exception)) + + +class TestMainExit(unittest.TestCase): + def test_an_unconvertible_source_exits_1_without_a_traceback(self): + """A wrong folder is a normal typo, not a crash.""" + from scripts.generate_local import main + + with TemporaryDirectory() as tmp: + empty = Path(tmp) / "empty" + empty.mkdir() + self.assertEqual(main([str(empty), "--out", str(Path(tmp) / "out")]), 1) + + +class TestRangeServer(unittest.TestCase): + """What `python -m http.server` would get wrong.""" + + @classmethod + def setUpClass(cls): + cls._tmp = TemporaryDirectory() + cls.root = Path(cls._tmp.name) + cls.payload = bytes(range(256)) * 4 # 1024 bytes, every value distinct by index + (cls.root / "data.parquet").write_bytes(cls.payload) + (cls.root / "manifest.json").write_text( + json.dumps({"version": 1, "tables": []}) + ) + + handler = partial(RangeRequestHandler, directory=str(cls.root)) + cls.server = HTTPServer(("localhost", 0), handler) + cls.port = cls.server.server_address[1] + cls.thread = threading.Thread(target=cls.server.serve_forever, daemon=True) + cls.thread.start() + + @classmethod + def tearDownClass(cls): + cls.server.shutdown() + cls.server.server_close() + cls.thread.join(timeout=5) + cls._tmp.cleanup() + + def _get(self, path, headers=None): + request = urllib.request.Request( + f"http://localhost:{self.port}/{path}", headers=headers or {} + ) + return urllib.request.urlopen(request) + + def test_serves_a_whole_file_without_a_range(self): + response = self._get("data.parquet") + self.assertEqual(response.status, 200) + self.assertEqual(response.read(), self.payload) + + def test_advertises_range_support_and_allows_cross_origin_reads(self): + response = self._get("manifest.json") + self.assertEqual(response.headers["Accept-Ranges"], "bytes") + self.assertEqual(response.headers["Access-Control-Allow-Origin"], "*") + self.assertIn( + "content-range", response.headers["Access-Control-Expose-Headers"] + ) + + def test_answers_a_byte_range_with_exactly_those_bytes(self): + response = self._get("data.parquet", {"Range": "bytes=10-19"}) + self.assertEqual(response.status, 206) + self.assertEqual(response.headers["Content-Range"], "bytes 10-19/1024") + self.assertEqual(response.read(), self.payload[10:20]) + + def test_an_open_ended_range_runs_to_the_end(self): + response = self._get("data.parquet", {"Range": "bytes=1000-"}) + self.assertEqual(response.status, 206) + self.assertEqual(response.read(), self.payload[1000:]) + + def test_a_suffix_range_reads_the_tail(self): + """How a Parquet footer is fetched; getting it wrong breaks every query.""" + response = self._get("data.parquet", {"Range": "bytes=-8"}) + self.assertEqual(response.status, 206) + self.assertEqual(response.headers["Content-Range"], "bytes 1016-1023/1024") + self.assertEqual(response.read(), self.payload[-8:]) + + def test_a_range_past_the_end_is_clamped(self): + response = self._get("data.parquet", {"Range": "bytes=1020-9999"}) + self.assertEqual(response.status, 206) + self.assertEqual(response.read(), self.payload[1020:]) + + def test_an_unsatisfiable_range_is_refused(self): + with self.assertRaises(urllib.error.HTTPError) as caught: + self._get("data.parquet", {"Range": "bytes=5000-6000"}) + self.assertEqual(caught.exception.code, 416) + + +class TestParser(unittest.TestCase): + def test_serve_defaults_to_a_port_without_one_given(self): + self.assertEqual(build_parser().parse_args(["mdb-1", "--serve"]).serve, 8090) + + def test_serve_is_off_unless_asked_for(self): + self.assertIsNone(build_parser().parse_args(["mdb-1"]).serve) + + def test_serve_accepts_an_explicit_port(self): + self.assertEqual( + build_parser().parse_args(["mdb-1", "--serve", "9000"]).serve, 9000 + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/functions-python/parquet_builder/tests/test_main.py b/functions-python/parquet_builder/tests/test_main.py new file mode 100644 index 000000000..944982be6 --- /dev/null +++ b/functions-python/parquet_builder/tests/test_main.py @@ -0,0 +1,1109 @@ +# +# MobilityData 2026 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""The orchestration: claiming, converting, publishing, and failing safely. + +Only Google Cloud Storage is faked. The archive really is a zip, and it really is +extracted and converted, so the wiring between the phases is exercised rather than +asserted about. +""" + +import io +import json +import tempfile +import unittest +from datetime import datetime, timedelta, timezone +from types import SimpleNamespace +import zipfile +from pathlib import Path +from unittest.mock import MagicMock, patch + +import flask +import pytest +from shared.helpers.sizing import Override + +import main + +pytest.importorskip("duckdb") + +FEED = "mdb-1210" +DATASET = "mdb-1210-202402121801" +BUCKET = "test-datasets" +PUBLIC = "https://files.example.org" + +AGENCY = "agency_id,agency_name,agency_url,agency_timezone\n1,T,https://e.org,UTC\n" +STOPS = "stop_id,stop_name\nS1,First\n" + + +def _zip_bytes(nested: bool = False, reverse: bool = False) -> bytes: + """`reverse` writes stops before agency, so archive order and table order differ.""" + buffer = io.BytesIO() + prefix = "feed/" if nested else "" + members = [(f"{prefix}agency.txt", AGENCY), (f"{prefix}stops.txt", STOPS)] + if reverse: + members.reverse() + with zipfile.ZipFile(buffer, "w") as zf: + for name, body in members: + zf.writestr(name, body) + return buffer.getvalue() + + +class FakeBlob: + def __init__(self, name, store, payload=None): + self.name = name + self._store = store + self._payload = payload + self.size = len(payload) if payload else 0 + self.public = False + self.custom_time = None + + def exists(self): + return self._payload is not None + + def reload(self): + pass + + def download_to_filename(self, path): + self._store.downloaded.append(self.name) + Path(path).write_bytes(self._payload) + + def download_as_bytes(self, start=None, end=None): + """Serve a byte range, as GCS does for the archive's central directory.""" + self._store.ranged.append((self.name, start, end)) + data = self._payload or b"" + return data[start : (end + 1 if end is not None else None)] + + def upload_from_filename(self, path): + self._store.uploaded[self.name] = Path(path).read_bytes() + self._store.custom_times[self.name] = self.custom_time + self._store.events.append(("upload", self.name)) + + def open(self, mode="rb", chunk_size=None): + """Stand-in for `BlobReader`: seekable, and never writes to the workdir.""" + assert mode == "rb" + if self._payload is None: + raise FileNotFoundError(self.name) + self._store.opened.append(self.name) + return io.BytesIO(self._payload) + + def make_public(self): + if self._store.acl_error is not None: + raise self._store.acl_error + self.public = True + self._store.made_public.add(self.name) + + def delete(self): + self._store.deleted.append(self.name) + self._store.events.append(("delete", self.name)) + + +class FakeBucket: + def __init__(self, archive: bytes, existing=()): + self.name = BUCKET + self.uploaded = {} + # customTime as it stood at upload, per object. + self.custom_times = {} + self.deleted = [] + self.made_public = set() + self.events = [] + self.downloaded = [] + # Blobs read through open("rb") rather than downloaded to the workdir. + self.opened = [] + self.ranged = [] + # Blobs under //extracted/, as batch_process_dataset leaves them. + self.extracted = {} + self._archive = archive + self._existing = list(existing) + # Raised by make_public when set, standing in for a bucket that refuses ACLs. + self.acl_error = None + self.iam_configuration = SimpleNamespace( + uniform_bucket_level_access_enabled=False + ) + + def blob(self, name): + archive_path = f"{FEED}/{DATASET}/{DATASET}.zip" + if name == archive_path: + return FakeBlob(name, self, self._archive) + return FakeBlob(name, self, self.extracted.get(name)) + + def list_blobs(self, prefix): + names = set(self._existing) | set(self.extracted) + return [FakeBlob(n, self) for n in sorted(names) if n.startswith(prefix)] + + +class BuildTestCase(unittest.TestCase): + def setUp(self): + self.bucket = FakeBucket(_zip_bytes()) + self.tracker = MagicMock() + self.tracker.try_acquire.return_value = True + self.session = MagicMock() + # Default: no extracted files recorded, so the archive path is taken. + self.session.query.return_value.filter.return_value.one_or_none.return_value = ( + None + ) + + client = MagicMock() + client.get_bucket.return_value = self.bucket + self._patches = [ + patch.object(main.storage, "Client", return_value=client), + patch.object(main, "TaskExecutionTracker", return_value=self.tracker), + patch.dict( + main.os.environ, + {"DATASETS_BUCKET_NAME": BUCKET, "PUBLIC_HOSTED_DATASETS_URL": PUBLIC}, + ), + ] + for p in self._patches: + p.start() + self.addCleanup(p.stop) + + def build(self, **kwargs): + return main.build_parquet( + feed_stable_id=FEED, + dataset_stable_id=DATASET, + bucket_name=BUCKET, + db_session=self.session, + **kwargs, + ) + + +class TestValidation(unittest.TestCase): + def _call(self, payload): + with flask.Flask(__name__).test_request_context(json=payload): + return main.build_parquet_handler(flask.request) + + def test_missing_identifiers_are_rejected(self): + self.assertIn("error", self._call({"feed_stable_id": FEED})) + self.assertIn("error", self._call({})) + + def test_a_dataset_belonging_to_another_feed_is_rejected(self): + result = self._call({"feed_stable_id": "mdb-999", "dataset_stable_id": DATASET}) + self.assertIn("not a prefix", result["error"]) + + def test_a_missing_bucket_setting_is_reported(self): + with patch.dict(main.os.environ, {}, clear=True): + result = self._call({"feed_stable_id": FEED, "dataset_stable_id": DATASET}) + self.assertIn("DATASETS_BUCKET_NAME", result["error"]) + + def test_a_build_failure_returns_200_so_cloud_tasks_does_not_retry(self): + """A corrupt archive is still corrupt on the second delivery.""" + with patch.dict( + main.os.environ, {"DATASETS_BUCKET_NAME": BUCKET} + ), patch.object(main, "build_parquet", side_effect=RuntimeError("boom")): + result = self._call({"feed_stable_id": FEED, "dataset_stable_id": DATASET}) + # A dict, not a raise: functions-framework would turn a raise into a 500. + self.assertEqual(result["status"], "error") + self.assertIn("boom", result["error"]) + + +class TestClaiming(BuildTestCase): + def test_a_refused_claim_does_no_work(self): + self.tracker.try_acquire.return_value = False + + result = self.build() + + self.assertEqual(result["status"], "skipped") + self.assertEqual(self.bucket.uploaded, {}, "nothing may be published") + self.tracker.mark_completed.assert_not_called() + + def test_force_reopens_a_finished_dataset_before_claiming(self): + self.build(force=True) + + self.tracker.release_for_retry.assert_called_once_with(DATASET) + + def test_without_force_a_finished_dataset_is_not_reopened(self): + self.build() + + self.tracker.release_for_retry.assert_not_called() + + +class TestSuccessfulBuild(BuildTestCase): + def test_publishes_a_parquet_per_table_plus_a_manifest(self): + result = self.build() + + self.assertEqual(result["status"], "success") + self.assertEqual( + sorted(self.bucket.uploaded), + [ + f"{FEED}/{DATASET}/parquet/agency.parquet", + f"{FEED}/{DATASET}/parquet/manifest.json", + f"{FEED}/{DATASET}/parquet/stops.parquet", + ], + ) + + def test_every_published_object_is_public(self): + """The reader discards query strings, so signed URLs cannot work.""" + self.build() + + self.assertEqual(self.bucket.made_public, set(self.bucket.uploaded)) + + def test_the_manifest_lists_the_tables(self): + self.build() + + manifest = json.loads( + self.bucket.uploaded[f"{FEED}/{DATASET}/parquet/manifest.json"] + ) + self.assertEqual([t["name"] for t in manifest["tables"]], ["agency", "stops"]) + + def test_records_where_the_files_are(self): + self.build() + + _, kwargs = self.tracker.mark_completed.call_args + metadata = kwargs["metadata"] + self.assertEqual(metadata["base_url"], f"{PUBLIC}/{FEED}/{DATASET}/parquet") + self.assertEqual([t["name"] for t in metadata["tables"]], ["agency", "stops"]) + self.assertEqual(metadata["phase"], "done") + + def test_reports_progress_through_every_phase(self): + self.build() + + phases = [ + call.kwargs["metadata"]["phase"] + for call in self.tracker.heartbeat.call_args_list + ] + # Neither `extract` nor `download` appears any more. Unzipping is per-member + # inside `convert`, and the archive is read over the network as those members + # are asked for, so neither is a stage the build passes through. + for expected in ("start", "convert", "upload", "summarise"): + self.assertIn(expected, phases, f"no {expected} reading was published") + + def test_no_download_phase_is_reported(self): + """Nothing is downloaded, so reporting the phase would describe absent work.""" + self.build() + + phases = [ + call.kwargs["metadata"]["phase"] + for call in self.tracker.heartbeat.call_args_list + ] + self.assertNotIn("download", phases) + + def test_a_stale_previous_build_is_cleared_first(self): + """A rebuild with fewer tables must not leave an orphan behind.""" + self.bucket._existing = [f"{FEED}/{DATASET}/parquet/gone.parquet"] + + self.build() + + self.assertIn(f"{FEED}/{DATASET}/parquet/gone.parquet", self.bucket.deleted) + + def test_a_feed_wrapped_in_a_directory_still_converts(self): + """Some producers nest the files; the converter looks in one place.""" + self.bucket._archive = _zip_bytes(nested=True) + + result = self.build() + + self.assertEqual(sorted(result["tables"]), ["agency", "stops"]) + + +class TestStreamsOneFileAtATime(BuildTestCase): + """Peak memory is the point of the per-file flow, so it is asserted directly. + + `/tmp` in Cloud Functions gen2 is RAM-backed tmpfs, so a file left on disk costs the + same as a file held in memory. Converting the whole feed and then uploading it held + the archive, every CSV and every Parquet at once; this keeps at most one of each. + """ + + def _residency(self): + """Watch the workdir while the build runs, recording what coexists.""" + seen = [] + real_convert = main.convert_table + + def spy(con, table, source, destination, compressed_bytes=None): + root = source.parent.parent + live = [p for p in root.rglob("*") if p.is_file()] + seen.append( + { + "csv": [p.name for p in live if p.suffix in (".txt", ".geojson")], + "parquet": [p.name for p in live if p.suffix == ".parquet"], + "zip": [p.name for p in live if p.suffix == ".zip"], + } + ) + return real_convert(con, table, source, destination, compressed_bytes) + + with patch.object(main, "convert_table", side_effect=spy): + self.build() + return seen + + def test_only_one_source_file_is_ever_resident(self): + snapshots = self._residency() + + self.assertTrue(snapshots, "no tables were converted") + worst = max(len(s["csv"]) for s in snapshots) + self.assertEqual( + worst, 1, f"more than one CSV was on disk at once: {snapshots}" + ) + + def test_parquet_files_do_not_accumulate(self): + snapshots = self._residency() + + worst = max(len(s["parquet"]) for s in snapshots) + self.assertLessEqual( + worst, + 1, + f"converted files accumulated instead of being published: {snapshots}", + ) + + def test_the_archive_is_kept_but_never_unpacked_wholesale(self): + """The deliberate trade on the fallback path. + + Holding the compressed archive and extracting one member at a time costs far + less than extracting every CSV up front, so the zip staying resident is correct + - what must not happen is the CSVs piling up beside it. + """ + snapshots = self._residency() + + self.assertTrue(all(len(s["csv"]) <= 1 for s in snapshots), snapshots) + + def test_everything_is_cleaned_up_afterwards(self): + self.build() + + leftovers = [p for p in Path(main.TMPDIR).glob("parquet_*") if p.is_dir()] + self.assertEqual(leftovers, [], f"workdirs left behind: {leftovers}") + + +class TestPublishIsNotObservablyPartial(BuildTestCase): + """A reader must never find a dataset advertising tables that are not there yet. + + Reported from the operations web app: a feed loaded with "No tables found", then + with one table, then with seven of eleven, and only became correct once every + object had landed. The manifest was being uploaded in the middle of the + alphabetical sequence, and a rebuild cleared the whole prefix before replacing it. + """ + + def _events(self): + prefix = f"{FEED}/{DATASET}/parquet/" + return [(kind, name[len(prefix) :]) for kind, name in self.bucket.events] + + def test_the_manifest_is_published_last(self): + """It is the readiness marker for anyone reading the bucket directly.""" + self.build() + + uploads = [name for kind, name in self._events() if kind == "upload"] + self.assertEqual( + uploads[-1], + "manifest.json", + f"manifest must be the final upload, got order {uploads}", + ) + + def test_nothing_is_deleted_before_the_new_set_is_up(self): + """Clearing first leaves an already-published dataset unreadable meanwhile.""" + self.bucket._existing = [f"{FEED}/{DATASET}/parquet/gone.parquet"] + + self.build() + + kinds = [kind for kind, _ in self._events()] + self.assertIn("delete", kinds) + self.assertLess( + max(i for i, k in enumerate(kinds) if k == "upload"), + min(i for i, k in enumerate(kinds) if k == "delete"), + "every upload must precede every delete", + ) + + def test_a_rebuild_never_removes_a_table_it_is_replacing(self): + """The live objects stay readable, overwritten in place, for the whole build.""" + self.bucket._existing = [f"{FEED}/{DATASET}/parquet/agency.parquet"] + + self.build() + + self.assertNotIn( + f"{FEED}/{DATASET}/parquet/agency.parquet", + self.bucket.deleted, + "a table present in both the old and new set must never be deleted", + ) + + def test_ready_is_recorded_only_after_every_object_has_landed(self): + recorded = [] + self.tracker.mark_completed.side_effect = lambda *a, **k: recorded.append( + len(self.bucket.events) + ) + + self.build() + + self.assertEqual( + recorded, [len(self.bucket.events)], "ready was published mid-publish" + ) + + +class TestReadsPreExtractedFiles(BuildTestCase): + """The fast path: use what `batch_process_dataset` already unpacked. + + Both paths must produce the same artifact - a dataset should not describe itself + differently depending on which route the builder happened to take. + """ + + def _with_extracted(self, names, dataset_id=DATASET): + """Point the DB at Gtfsfile rows and seed the matching blobs.""" + records = [] + for name in names: + record = MagicMock() + record.file_name = name + records.append(record) + body = AGENCY if name.endswith("agency.txt") else STOPS + self.bucket.extracted[f"{FEED}/{dataset_id}/extracted/{name}"] = ( + body.encode() + ) + + dataset = MagicMock() + dataset.gtfsfiles = records + dataset.zipped_size_bytes = 4321 + self.session.query.return_value.filter.return_value.one_or_none.return_value = ( + dataset + ) + return dataset + + def test_converts_from_extracted_without_downloading_the_archive(self): + self._with_extracted(["agency.txt", "stops.txt"]) + + result = self.build() + + self.assertEqual(sorted(result["tables"]), ["agency", "stops"]) + self.assertNotIn( + f"{FEED}/{DATASET}/{DATASET}.zip", + self.bucket.opened, + "the archive was read despite extracted files being available", + ) + + def test_a_recorded_file_the_bucket_lacks_falls_back_to_the_archive(self): + """The Gtfsfile rows are an index, not a guarantee. + + `rebuild_missing_dataset_files` exists because a dataset can be recorded with + files the bucket no longer holds. This used to raise NotFound from inside the + conversion and lose the whole build. + """ + dataset = self._with_extracted(["agency.txt"]) + phantom = MagicMock() + phantom.file_name = "feed_info.txt" + dataset.gtfsfiles = list(dataset.gtfsfiles) + [phantom] + + result = self.build() + + self.assertEqual(result["status"], "success") + self.assertIn( + f"{FEED}/{DATASET}/{DATASET}.zip", + self.bucket.opened, + "the archive was not used despite the extracted set being incomplete", + ) + # The archive holds both tables, so nothing is lost by taking that route. + self.assertEqual(sorted(result["tables"]), ["agency", "stops"]) + + def test_an_incomplete_extracted_set_is_never_converted_on_its_own(self): + """Publishing the subset would assert a partial feed is the whole feed.""" + dataset = self._with_extracted(["agency.txt"]) + phantom = MagicMock() + phantom.file_name = "stop_times.txt" + dataset.gtfsfiles = list(dataset.gtfsfiles) + [phantom] + # No archive to fall back to either. + self.bucket._archive = None + self.bucket.blob = lambda name: FakeBlob(name, self.bucket, None) + + with self.assertRaises(FileNotFoundError): + self.build() + + self.assertEqual(self.bucket.uploaded, {}, "a partial set was published") + + def test_archive_relative_paths_are_flattened(self): + """`extracted/` preserves the archive's own layout, folders and all.""" + self._with_extracted(["feed/agency.txt", "feed/stops.txt"]) + + result = self.build() + + self.assertEqual(sorted(result["tables"]), ["agency", "stops"]) + + def test_non_gtfs_members_are_ignored(self): + self._with_extracted( + ["agency.txt", "stops.txt", "__MACOSX/._stops.txt", "licence.pdf"] + ) + + result = self.build() + + self.assertEqual(sorted(result["tables"]), ["agency", "stops"]) + + def test_compressed_sizes_come_from_the_archive_tail(self): + """Read by range, not by downloading the archive - the numbers exist nowhere else.""" + self._with_extracted(["agency.txt", "stops.txt"]) + + self.build() + + manifest = json.loads( + self.bucket.uploaded[f"{FEED}/{DATASET}/parquet/manifest.json"] + ) + self.assertEqual(manifest["source"]["kind"], "zip") + self.assertTrue( + any(t["compressed_bytes"] for t in manifest["tables"]), + f"no compressed sizes recovered: {manifest['tables']}", + ) + self.assertTrue(self.bucket.ranged, "the tail was never fetched") + + def test_the_recorded_archive_size_is_used_when_present(self): + self._with_extracted(["agency.txt"]) + + self.build() + + manifest = json.loads( + self.bucket.uploaded[f"{FEED}/{DATASET}/parquet/manifest.json"] + ) + self.assertEqual(manifest["source"]["bytes"], 4321) + + def test_falls_back_to_the_archive_when_nothing_was_extracted(self): + """Datasets processed before extraction existed still have to build.""" + dataset = MagicMock() + dataset.gtfsfiles = [] + self.session.query.return_value.filter.return_value.one_or_none.return_value = ( + dataset + ) + + result = self.build() + + self.assertEqual(sorted(result["tables"]), ["agency", "stops"]) + self.assertIn(f"{FEED}/{DATASET}/{DATASET}.zip", self.bucket.opened) + + +class TestArchiveIsStreamed(BuildTestCase): + """The archive is read over the network, never written to the volume. + + Downloading it cost its full size on the in-memory volume for the whole build, on + top of whichever member was being converted. A 1 GiB archive holding a 4 GiB table + exhausted the volume before conversion started. + """ + + def _archive_only(self): + """No Gtfsfile rows, so the build takes the archive path.""" + dataset = MagicMock() + dataset.gtfsfiles = [] + self.session.query.return_value.filter.return_value.one_or_none.return_value = ( + dataset + ) + + def test_the_archive_is_never_written_to_the_workdir(self): + self._archive_only() + + self.build() + + self.assertIn(f"{FEED}/{DATASET}/{DATASET}.zip", self.bucket.opened) + self.assertEqual( + self.bucket.downloaded, [], "the archive was written to the volume" + ) + + def test_members_are_read_in_archive_order_not_table_order(self): + """A backward seek throws the reader's buffer away and refetches. + + The fixture writes stops before agency, so archive order and alphabetical order + disagree and a plan built in the wrong one is visible. + """ + self.bucket._archive = _zip_bytes(reverse=True) + self._archive_only() + + plan = main._plan_from_archive( + self.bucket, FEED, DATASET, Path(tempfile.mkdtemp()), MagicMock() + ) + try: + self.assertEqual([table for table, _ in plan.sources], ["stops", "agency"]) + finally: + plan.close() + + def test_tables_are_still_reported_by_name(self): + """Conversion order is an IO detail; the manifest must not depend on it.""" + self.bucket._archive = _zip_bytes(reverse=True) + self._archive_only() + + result = self.build() + + self.assertEqual(result["tables"], ["agency", "stops"]) + manifest = json.loads( + self.bucket.uploaded[f"{FEED}/{DATASET}/parquet/manifest.json"] + ) + self.assertEqual([t["name"] for t in manifest["tables"]], ["agency", "stops"]) + + def test_compressed_sizes_come_from_the_directory_already_read(self): + """The central directory is in hand, so the tail is not fetched a second time.""" + self._archive_only() + + self.build() + + manifest = json.loads( + self.bucket.uploaded[f"{FEED}/{DATASET}/parquet/manifest.json"] + ) + for table in manifest["tables"]: + self.assertIsNotNone( + table["compressed_bytes"], f"{table['name']} lost its zipped size" + ) + self.assertEqual(self.bucket.ranged, [], "the archive tail was re-read") + + def test_the_handle_is_released_when_the_build_fails(self): + self._archive_only() + + with patch.object(main, "convert_table", side_effect=RuntimeError("no memory")): + with self.assertRaises(RuntimeError): + self.build() + + # Nothing to assert on GCS here; the point is that the failure path runs + # plan.close() rather than leaking the reader, which the build would otherwise + # hold until the instance is recycled. + self.tracker.mark_failed.assert_called_once() + + +class TestEscalationOnFailure(BuildTestCase): + """A runaway here re-queues forever at real cost, so every stop is pinned.""" + + def setUp(self): + super().setUp() + self.enqueue = patch.object(main, "create_http_parquet_builder_task").start() + self.override = patch.object(main, "record_size_override").start() + self.stored = patch.object( + main, "size_override", return_value=Override() + ).start() + self.addCleanup(patch.stopall) + self.tracker.attempts_since_success.return_value = 1 + # A resolvable feed, as the escalation needs for the config write. + dataset = MagicMock() + dataset.feed = MagicMock(id="feed-uuid", stable_id=FEED) + self.session.query.return_value.filter.return_value.one_or_none.return_value = ( + dataset + ) + + def _fail_with(self, error, service="parquet-builder-s-dev"): + with patch.dict(main.os.environ, {"K_SERVICE": service, "ENVIRONMENT": "dev"}): + with patch.object(main, "convert_table", side_effect=error): + with self.assertRaises(type(error)): + self.build() + + def test_a_memory_failure_escalates_one_rung(self): + self._fail_with(MemoryError()) + + self.override.assert_called_once() + self.assertEqual(self.override.call_args.args[3], main.Size.M) + self.enqueue.assert_called_once() + self.assertEqual(self.enqueue.call_args.kwargs["size"], main.Size.M) + + def test_a_disk_failure_escalates_too(self): + self._fail_with(OSError(28, "No space left on device")) + + self.enqueue.assert_called_once() + + def test_an_ordinary_failure_never_escalates(self): + """Retrying a corrupt archive on a bigger machine is just a second failure.""" + self._fail_with(ValueError("no tables could be converted")) + + self.override.assert_not_called() + self.enqueue.assert_not_called() + + def test_a_locked_feed_does_not_escalate(self): + """A lock means what it says: nothing is retried and nothing is built, which is + why it is logged and reported rather than quietly worked around.""" + self.stored.return_value = Override(size=main.Size.S, locked=True) + + self._fail_with(OSError(28, "No space left on device")) + + self.override.assert_not_called() + self.enqueue.assert_not_called() + + def test_a_locked_feed_still_records_the_failure(self): + self.stored.return_value = Override(size=main.Size.S, locked=True) + + self._fail_with(MemoryError()) + + kwargs = self.tracker.record_attempt.call_args.kwargs + self.assertEqual(kwargs["status"], "failed") + self.assertIsNone(kwargs["escalated_to"]) + + def test_the_largest_worker_does_not_escalate(self): + """Nothing above `l` to escalate to, so this is where the loop stops.""" + self._fail_with(MemoryError(), service="parquet-builder-l-dev") + + self.override.assert_not_called() + self.enqueue.assert_not_called() + + def test_the_attempt_cap_stops_the_loop(self): + self.tracker.attempts_since_success.return_value = main.MAX_ATTEMPTS + + self._fail_with(MemoryError()) + + self.enqueue.assert_not_called() + + def test_an_unknown_worker_does_not_escalate(self): + """If K_SERVICE does not name a known size, guessing which rung died is worse + than stopping.""" + self._fail_with(MemoryError(), service="something-else") + + self.enqueue.assert_not_called() + + def test_the_attempt_is_recorded_with_the_decision(self): + self._fail_with(MemoryError()) + + kwargs = self.tracker.record_attempt.call_args.kwargs + self.assertEqual(kwargs["status"], "failed") + self.assertEqual(kwargs["variant"], "s") + self.assertEqual(kwargs["failure_kind"], "resource_memory") + self.assertEqual(kwargs["escalated_to"], "m") + self.assertIsInstance(kwargs["error"], MemoryError) + + def test_a_declined_escalation_records_no_target(self): + """The difference between "chose not to escalate" and "never ran".""" + self._fail_with(ValueError("no tables")) + + kwargs = self.tracker.record_attempt.call_args.kwargs + self.assertIsNone(kwargs["escalated_to"]) + self.assertEqual(kwargs["failure_kind"], "other") + + def test_a_success_is_recorded_too(self): + with patch.dict( + main.os.environ, + {"K_SERVICE": "parquet-builder-m-dev", "ENVIRONMENT": "dev"}, + ): + self.build() + + kwargs = self.tracker.record_attempt.call_args.kwargs + self.assertEqual(kwargs["status"], "completed") + self.assertEqual(kwargs["variant"], "m") + self.enqueue.assert_not_called() + + +class TestBothPathsAgree(BuildTestCase): + """A dataset must not describe itself differently depending on the route taken.""" + + def _manifest(self): + return json.loads( + self.bucket.uploaded[f"{FEED}/{DATASET}/parquet/manifest.json"] + ) + + def _use_extracted(self): + records = [] + for name, body in (("agency.txt", AGENCY), ("stops.txt", STOPS)): + record = MagicMock() + record.file_name = name + records.append(record) + self.bucket.extracted[f"{FEED}/{DATASET}/extracted/{name}"] = body.encode() + dataset = MagicMock() + dataset.gtfsfiles = records + dataset.zipped_size_bytes = None # force both paths to size the archive alike + self.session.query.return_value.filter.return_value.one_or_none.return_value = ( + dataset + ) + + def test_extracted_and_archive_paths_produce_the_same_manifest(self): + self.build() # no Gtfsfile rows -> archive path + from_archive = self._manifest() + + self.setUp() + self._use_extracted() + self.build() + from_extracted = self._manifest() + + volatile = {"generated_at"} + self.assertEqual( + {k: v for k, v in from_archive.items() if k not in volatile}, + {k: v for k, v in from_extracted.items() if k not in volatile}, + "the two source paths disagree about the same dataset", + ) + + +class TestRetention(BuildTestCase): + """The builder owns the default, so a caller that says nothing still gets one.""" + + def _recorded(self): + _, kwargs = self.tracker.mark_completed.call_args + return kwargs["metadata"]["retention_days"] + + def test_an_omitted_value_uses_the_builders_default(self): + self.build() + + self.assertEqual(self._recorded(), main.DEFAULT_RETENTION_DAYS) + + def test_a_callers_value_is_honoured(self): + self.build(retention_days=7) + + self.assertEqual(self._recorded(), 7) + + def test_out_of_range_values_fall_back_rather_than_being_stored(self): + """The API bounds this too, but Cloud Tasks and the CLI do not go through it, + and a bad value here would be written to the row and obeyed by the sweep.""" + for bad in (0, -1, main.MAX_RETENTION_DAYS + 1, "abc", 3.7e9): + with self.subTest(bad=bad): + self.assertEqual(main._retention_days(bad), main.DEFAULT_RETENTION_DAYS) + + def test_the_boundaries_are_accepted(self): + self.assertEqual(main._retention_days(1), 1) + self.assertEqual( + main._retention_days(main.MAX_RETENTION_DAYS), main.MAX_RETENTION_DAYS + ) + + def test_none_means_the_default(self): + self.assertEqual(main._retention_days(None), main.DEFAULT_RETENTION_DAYS) + + def test_every_object_carries_the_expiry_as_custom_time(self): + """The bucket's lifecycle rule deletes on customTime, so an object without one + would never expire.""" + before = datetime.now(timezone.utc) + self.build(retention_days=7) + + self.assertEqual( + set(self.bucket.custom_times), + set(self.bucket.uploaded), + "an object was published without an expiry", + ) + for name, stamped in self.bucket.custom_times.items(): + with self.subTest(name=name): + self.assertIsNotNone(stamped) + delta = stamped - before + self.assertGreater(delta, timedelta(days=7) - timedelta(minutes=1)) + self.assertLess(delta, timedelta(days=7) + timedelta(minutes=1)) + + def test_the_manifest_expires_with_the_tables(self): + """A manifest outliving its tables would advertise files that are gone.""" + self.build() + + stamps = set(self.bucket.custom_times.values()) + self.assertEqual(len(stamps), 1, f"the set does not expire together: {stamps}") + + def test_a_shorter_retention_expires_sooner(self): + self.build(retention_days=7) + short = min(self.bucket.custom_times.values()) + + self.setUp() + self.build(retention_days=30) + long = min(self.bucket.custom_times.values()) + + self.assertLess(short, long) + + def test_the_row_records_the_same_instant_as_the_objects(self): + """Two sources of truth for one date would drift; the API reads the row.""" + self.build() + + recorded = datetime.fromisoformat( + self.tracker.mark_completed.call_args.kwargs["metadata"]["expires_at"] + ) + self.assertEqual(recorded, next(iter(self.bucket.custom_times.values()))) + + +class TestClaimDurability(BuildTestCase): + """The claim is what stops a second worker; it cannot ride on a best-effort write.""" + + def setUp(self): + super().setUp() + self.events = [] + self.session.commit.side_effect = lambda: self.events.append("commit") + self.tracker.heartbeat.side_effect = lambda *a, **k: self.events.append( + "heartbeat" + ) + + def test_the_claim_is_committed_before_any_progress_write(self): + self.build() + + self.assertEqual( + self.events[0], + "commit", + f"the claim was not committed first: {self.events[:3]}", + ) + self.assertLess(self.events.index("commit"), self.events.index("heartbeat")) + + def test_a_progress_write_that_fails_leaves_the_session_usable(self): + """Progress is a courtesy, so the failure is swallowed - but not the rollback.""" + self.tracker.heartbeat.side_effect = RuntimeError("connection reset") + + result = self.build() + + self.assertEqual(result["status"], "success") + self.assertTrue(self.session.rollback.called, "the session was left poisoned") + self.tracker.mark_completed.assert_called_once() + + +class TestPublicReadability(BuildTestCase): + def test_an_acl_failure_fails_the_build(self): + """A published but unreadable file must not be reported as `ready`.""" + self.bucket.acl_error = RuntimeError("403 forbidden") + + with self.assertRaises(RuntimeError): + self.build() + + self.tracker.mark_completed.assert_not_called() + self.tracker.mark_failed.assert_called_once() + + def test_uniform_bucket_level_access_needs_no_acl(self): + """There the objects are public by policy, and per-object ACLs are rejected.""" + self.bucket.acl_error = RuntimeError("cannot use ACL API") + self.bucket.iam_configuration.uniform_bucket_level_access_enabled = True + + result = self.build() + + self.assertEqual(result["status"], "success") + + +class TestFailure(BuildTestCase): + def test_a_missing_archive_is_recorded_and_released(self): + self.bucket._archive = None + self.bucket.blob = lambda name: FakeBlob(name, self.bucket, None) + + with self.assertRaises(FileNotFoundError): + self.build() + + # Recorded as failed rather than left holding the claim until the lease runs + # out, so the dataset can be retried at once. + self.tracker.mark_failed.assert_called_once() + self.assertIn( + "not found", self.tracker.mark_failed.call_args.kwargs["error_message"] + ) + + def test_a_conversion_failure_releases_the_claim(self): + with patch.object(main, "convert_table", side_effect=RuntimeError("no memory")): + with self.assertRaises(RuntimeError): + self.build() + + self.tracker.mark_failed.assert_called_once() + self.tracker.mark_completed.assert_not_called() + self.assertEqual(self.bucket.uploaded, {}) + + def test_a_failure_in_the_commit_itself_is_still_recorded(self): + """Recording a failure must not need a session the failure has just broken.""" + completed = [] + self.tracker.mark_completed.side_effect = lambda *a, **k: completed.append(1) + + def commit(): + if completed: + raise RuntimeError("could not commit") + + self.session.commit.side_effect = commit + + with self.assertRaises(RuntimeError): + self.build() + + self.assertTrue(self.session.rollback.called) + self.tracker.mark_failed.assert_called_once() + + +MIB = 1024**2 + + +def _override(size, locked=False, source="auto"): + """A stored override, as `_downsize_after_success` reads it.""" + return Override(size=size, source=source, locked=locked) + + +def _attempt( + status="completed", variant="m", vms=500 * MIB, largest=50_000_000, metadata=True +): + """One row of a feed's recent history, as the downsize rule reads it.""" + row = MagicMock() + row.status = status + row.variant = variant + row.peak_vms_bytes = vms + row.metadata_ = {"largest_member_bytes": largest} if metadata else None + return row + + +class TestDownsizeAfterSuccess(BuildTestCase): + """Coming down is only ever an economy, so it waits for a dull amount of evidence.""" + + def setUp(self): + super().setUp() + self.record = patch.object(main, "record_size_override").start() + self.override = patch.object( + main, "size_override", return_value=_override(main.Size.M) + ).start() + self.attempts = patch.object( + main, "_recent_attempts", return_value=[_attempt() for _ in range(3)] + ).start() + self.dataset = MagicMock() + self.dataset.feed = MagicMock(id="feed-uuid", stable_id=FEED) + patch.object(main, "_dataset_row", return_value=self.dataset).start() + self.addCleanup(patch.stopall) + + def _review(self, variant=main.Size.M): + return main._downsize_after_success(self.session, DATASET, variant, MagicMock()) + + def test_a_clean_streak_lowers_one_rung(self): + self.assertEqual(self._review(), main.Size.S) + + self.assertEqual(self.record.call_args.args[3], main.Size.S) + + def test_nothing_is_re_queued(self): + """The build already succeeded; the decision is for the feed's next one.""" + with patch.object(main, "create_http_parquet_builder_task") as enqueue: + self._review() + + enqueue.assert_not_called() + + def test_a_streak_one_attempt_short_does_not(self): + self.attempts.return_value = [_attempt() for _ in range(2)] + + self.assertIsNone(self._review()) + self.record.assert_not_called() + + def test_one_attempt_over_the_memory_threshold_does_not(self): + self.attempts.return_value = [_attempt(), _attempt(vms=1500 * MIB), _attempt()] + + self.assertIsNone(self._review()) + + def test_one_attempt_over_the_disk_threshold_does_not(self): + """A feed whose datasets alternate: the large one breaks its own streak.""" + self.attempts.return_value = [ + _attempt(), + _attempt(largest=200_000_000), + _attempt(), + ] + + self.assertIsNone(self._review()) + + def test_a_failure_in_the_streak_does_not(self): + self.attempts.return_value = [_attempt(), _attempt(status="failed"), _attempt()] + + self.assertIsNone(self._review()) + + def test_an_attempt_on_another_worker_does_not(self): + self.attempts.return_value = [_attempt(), _attempt(variant="l"), _attempt()] + + self.assertIsNone(self._review()) + + def test_an_attempt_recorded_before_the_evidence_existed_does_not(self): + self.attempts.return_value = [_attempt(), _attempt(metadata=False), _attempt()] + + self.assertIsNone(self._review()) + + def test_an_unlocked_pin_is_lowered_like_any_other_override(self): + """Who wrote the value is a record, not a claim on it.""" + self.override.return_value = _override(main.Size.M, source="operator") + + self.assertEqual(self._review(), main.Size.S) + + def test_a_locked_override_is_never_lowered(self): + self.override.return_value = _override(main.Size.M, locked=True) + + self.assertIsNone(self._review()) + self.record.assert_not_called() + + def test_a_feed_with_no_override_is_left_alone(self): + """It already routes on its measurement, which is as low as it goes.""" + self.override.return_value = Override() + + self.assertIsNone(self._review()) + + def test_an_override_that_disagrees_with_the_worker_is_left_alone(self): + self.override.return_value = _override(main.Size.L) + + self.assertIsNone(self._review()) + + def test_the_smallest_worker_is_terminal(self): + self.assertIsNone(self._review(variant=main.Size.S)) + self.record.assert_not_called() + + def test_an_unknown_worker_does_nothing(self): + self.assertIsNone(self._review(variant=None)) + + def test_a_failure_in_the_review_never_fails_the_build(self): + self.attempts.side_effect = RuntimeError("the database went away") + + self.assertIsNone(self._review()) + self.session.rollback.assert_called() + + +if __name__ == "__main__": + unittest.main() diff --git a/functions-python/parquet_builder/tests/test_progress.py b/functions-python/parquet_builder/tests/test_progress.py new file mode 100644 index 000000000..a4cffc91a --- /dev/null +++ b/functions-python/parquet_builder/tests/test_progress.py @@ -0,0 +1,132 @@ +# +# MobilityData 2026 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""What the throttle must and must not drop.""" + +import unittest + +from progress import PHASE_CONVERT, PHASE_DOWNLOAD, PHASE_EXTRACT, ThrottledProgress + + +class FakeClock: + def __init__(self): + self.now = 0.0 + + def __call__(self): + return self.now + + def advance(self, seconds): + self.now += seconds + + +def _recorder(): + written = [] + return written, written.append + + +class TestThrottledProgress(unittest.TestCase): + def test_writes_the_first_reading(self): + written, publish = _recorder() + progress = ThrottledProgress(publish, clock=FakeClock()) + + progress(PHASE_EXTRACT, 1, 10, "stops.txt") + + self.assertEqual(len(written), 1) + self.assertEqual(written[0]["phase"], PHASE_EXTRACT) + + def test_coalesces_events_inside_one_phase(self): + written, publish = _recorder() + clock = FakeClock() + progress = ThrottledProgress(publish, min_interval_s=1.0, clock=clock) + + # A large feed emits thousands of these; the reader polls twice a second. + for i in range(1, 501): + progress(PHASE_EXTRACT, i, 500, f"file{i}.txt") + + self.assertEqual( + len(written), 1, "500 events inside one second must be one write" + ) + + def test_writes_again_once_the_interval_passes(self): + written, publish = _recorder() + clock = FakeClock() + progress = ThrottledProgress(publish, min_interval_s=1.0, clock=clock) + + progress(PHASE_EXTRACT, 1, 10, "a.txt") + clock.advance(1.5) + progress(PHASE_EXTRACT, 2, 10, "b.txt") + + self.assertEqual([w["detail"] for w in written], ["a.txt", "b.txt"]) + + def test_never_drops_a_phase_change(self): + """The rule that matters: phases are what the viewer turns into words.""" + written, publish = _recorder() + clock = FakeClock() + progress = ThrottledProgress(publish, min_interval_s=1000.0, clock=clock) + + progress(PHASE_DOWNLOAD, 0, 0, "feed.zip") + progress(PHASE_EXTRACT, 1, 10, "stops.txt") + progress(PHASE_CONVERT, 1, 5, "stops") + + self.assertEqual( + [w["phase"] for w in written], + [PHASE_DOWNLOAD, PHASE_EXTRACT, PHASE_CONVERT], + "a phase change must be written even deep inside a quiet window", + ) + + def test_flush_writes_regardless_of_the_interval(self): + written, publish = _recorder() + clock = FakeClock() + progress = ThrottledProgress(publish, min_interval_s=1000.0, clock=clock) + + progress(PHASE_CONVERT, 1, 32, "agency") + progress(PHASE_CONVERT, 32, 32, "trips") # throttled away + progress.flush() + + self.assertEqual(len(written), 2) + self.assertEqual(written[-1]["done"], 32, "the final count must survive") + + def test_flush_overrides_are_applied(self): + written, publish = _recorder() + progress = ThrottledProgress(publish, clock=FakeClock()) + + progress.flush(phase=PHASE_DOWNLOAD, done=10, total=100, detail="feed.zip") + + self.assertEqual( + written[-1], + {"phase": PHASE_DOWNLOAD, "done": 10, "total": 100, "detail": "feed.zip"}, + ) + + def test_a_failing_publish_does_not_break_the_build(self): + """Progress is a courtesy. Losing a reading must not lose the conversion.""" + + def explode(_state): + raise RuntimeError("database went away") + + progress = ThrottledProgress(explode, clock=FakeClock()) + progress(PHASE_CONVERT, 1, 5, "stops") # must not raise + + def test_ints_are_coerced(self): + written, publish = _recorder() + progress = ThrottledProgress(publish, clock=FakeClock()) + + progress(PHASE_DOWNLOAD, 1.0, 2.0, "feed.zip") + + self.assertIsInstance(written[0]["done"], int) + self.assertIsInstance(written[0]["total"], int) + + +if __name__ == "__main__": + unittest.main() diff --git a/functions-python/pmtiles_builder/src/main.py b/functions-python/pmtiles_builder/src/main.py index da0f37bd0..f917e55fb 100644 --- a/functions-python/pmtiles_builder/src/main.py +++ b/functions-python/pmtiles_builder/src/main.py @@ -41,7 +41,7 @@ from shared.helpers.runtime_metrics import track_metrics from shared.database.database import with_db_session from shared.common.gcp_utils import create_web_revalidation_task -from ephemeral_workdir import EphemeralOrDebugWorkdir +from shared.helpers.ephemeral_workdir import EphemeralOrDebugWorkdir import flask import functions_framework @@ -91,7 +91,9 @@ def build_pmtiles_handler(request: flask.Request) -> dict: workdir_root = os.getenv("WORKDIR_ROOT", "/tmp/in-memory") # Use combined context manager that also cleans old directories with EphemeralOrDebugWorkdir( - dir=workdir_root, prefix=f"{dataset_stable_id}_" + owner_prefix="pmtiles_", + dir=workdir_root, + prefix=f"{dataset_stable_id}_", ) as workdir: result: dict[str, object] = { "params": { diff --git a/functions-python/pmtiles_builder/tests/test_workdir.py b/functions-python/pmtiles_builder/tests/test_workdir.py index f12ae00fb..c83ae9b0d 100644 --- a/functions-python/pmtiles_builder/tests/test_workdir.py +++ b/functions-python/pmtiles_builder/tests/test_workdir.py @@ -3,7 +3,7 @@ import time from pathlib import Path -from src.main import EphemeralOrDebugWorkdir +from shared.helpers.ephemeral_workdir import EphemeralOrDebugWorkdir def _backdate(path: Path, seconds_ago: int): @@ -30,7 +30,7 @@ def test_old_prefixed_dir_deleted_and_root_auto_removed(monkeypatch): fresh_prefixed.mkdir() _backdate(fresh_prefixed, 1) - with EphemeralOrDebugWorkdir(dir=root) as new_dir: + with EphemeralOrDebugWorkdir(owner_prefix="pmtiles_", dir=root) as new_dir: assert not stale_dir.exists() assert keep_other.exists() assert fresh_prefixed.exists() @@ -59,7 +59,7 @@ def test_fresh_prefixed_dir_retained(monkeypatch, tmp_path): recent_dir.mkdir() _backdate(recent_dir, 10) # Younger than TTL - with EphemeralOrDebugWorkdir(dir=str(root)): + with EphemeralOrDebugWorkdir(owner_prefix="pmtiles_", dir=str(root)): assert recent_dir.exists() # Recent directory still present @@ -88,7 +88,7 @@ def test_debug_workdir_preserved(monkeypatch, tmp_path): stale_dir.mkdir() _backdate(stale_dir, 10) - with EphemeralOrDebugWorkdir(dir=str(root)) as returned: + with EphemeralOrDebugWorkdir(owner_prefix="pmtiles_", dir=str(root)) as returned: # returned path should be the debug dir assert returned == str(debug_dir) assert Path(returned).exists() @@ -112,7 +112,7 @@ def test_debug_mode_skips_cleanup(monkeypatch, tmp_path): old_prefixed.mkdir() _backdate(old_prefixed, 1000) - with EphemeralOrDebugWorkdir(dir=str(root)): + with EphemeralOrDebugWorkdir(owner_prefix="pmtiles_", dir=str(root)): # cleanup should be skipped, so old_prefixed still present assert old_prefixed.exists() # still present after exit diff --git a/functions-python/tasks_executor/src/main.py b/functions-python/tasks_executor/src/main.py index 9ad8bfb5f..ca0993f43 100644 --- a/functions-python/tasks_executor/src/main.py +++ b/functions-python/tasks_executor/src/main.py @@ -68,6 +68,7 @@ from tasks.users.reconcile_announcements_from_brevo import ( reconcile_announcements_from_brevo_handler, ) +from tasks.task_execution_stats import task_execution_stats_handler from tasks.users.purge_early_access_invites import purge_early_access_invites_handler from tasks.notifications.dispatch_batch import notifications_dispatch_batch_handler from tasks.notifications.dispatch_worker import ( @@ -234,6 +235,18 @@ ), "handler": reconcile_announcements_from_brevo_handler, }, + "task_execution_stats": { + "description": ( + "Aggregates task_execution_attempt for one task over a window: attempts, " + "completions and failures by worker variant, peak address space and RSS per " + "variant, duration percentiles, failure counts by kind, and the entities " + "that needed more than one attempt. Read-only. Exists because worker sizing " + "has to come from a distribution, not from whichever build someone last " + "looked at; it reports, it does not retune. Parameters: task_name " + "(required), window_days (default 30), limit (default 20)." + ), + "handler": task_execution_stats_handler, + }, "purge_early_access_invites": { "description": ( "Deletes early_access_invited_email rows past their program's " diff --git a/functions-python/tasks_executor/src/tasks/task_execution_stats.py b/functions-python/tasks_executor/src/tasks/task_execution_stats.py new file mode 100644 index 000000000..a39a6c895 --- /dev/null +++ b/functions-python/tasks_executor/src/tasks/task_execution_stats.py @@ -0,0 +1,185 @@ +# +# MobilityData 2026 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Task: task_execution_stats + +What a task's attempts actually look like, aggregated from `task_execution_attempt`. + +This exists because worker sizing was set twice from a single observed build and was +wrong both times. The numbers that matter are distributions: peak address space per +worker against what that worker allows, how often each one runs out, and whether a run of +feeds is escalating off the same band. One build cannot show any of that. + +Read-only, and deliberately a report rather than an autotuner. Memory and volume live in +terraform; this says what they should be, a person decides. + +Unlike `get_summary`, which pulls every row for a run into Python and counts there, this +aggregates in the database - the table it reads grows without bound by design. +""" + +from __future__ import annotations + +import logging +from datetime import datetime, timedelta, timezone + +from sqlalchemy import text +from sqlalchemy.orm import Session + +from shared.database.database import with_db_session + +logger = logging.getLogger(__name__) + +DEFAULT_WINDOW_DAYS = 30 + +# Percentiles rather than a mean: durations here are bimodal, a few very long builds +# among many trivial ones, and a mean describes neither. +_BY_VARIANT_SQL = text(""" + SELECT variant, + count(*) AS attempts, + count(*) FILTER (WHERE status = 'completed') AS completed, + count(*) FILTER (WHERE status = 'failed') AS failed, + count(*) FILTER (WHERE escalated_to IS NOT NULL) AS escalated, + count(*) FILTER (WHERE variant_basis = 'floor') AS on_floor, + max(peak_vms_bytes) AS peak_vms_bytes, + max(peak_rss_bytes) AS peak_rss_bytes, + percentile_disc(0.5) WITHIN GROUP (ORDER BY duration_ms) AS p50_ms, + percentile_disc(0.95) WITHIN GROUP (ORDER BY duration_ms) AS p95_ms, + max(duration_ms) AS max_ms + FROM task_execution_attempt + WHERE task_name = :task_name + AND finished_at >= :since + GROUP BY variant + ORDER BY variant + """) + +_BY_FAILURE_SQL = text(""" + SELECT failure_kind, error_type, count(*) AS attempts + FROM task_execution_attempt + WHERE task_name = :task_name + AND finished_at >= :since + AND status = 'failed' + GROUP BY failure_kind, error_type + ORDER BY count(*) DESC + """) + +# Entities that needed more than one go. The retry cap means this is bounded, so a long +# list is a sign the bands are wrong rather than that one feed is unusual. +_REPEAT_OFFENDERS_SQL = text(""" + SELECT entity_id, + count(*) AS attempts, + max(escalated_to) AS escalated_to + FROM task_execution_attempt + WHERE task_name = :task_name + AND finished_at >= :since + GROUP BY entity_id + HAVING count(*) > 1 + ORDER BY count(*) DESC + LIMIT :limit + """) + + +def task_execution_stats_handler(payload: dict) -> dict: + """Entry point. + + Payload: + task_name (str): required, e.g. "parquet_generation". + window_days (int): how far back to look. Default 30. + limit (int): how many repeat offenders to list. Default 20. + """ + task_name = payload.get("task_name") + if not task_name: + raise ValueError("task_name is required") + return task_execution_stats( + task_name=task_name, + window_days=int(payload.get("window_days", DEFAULT_WINDOW_DAYS)), + limit=int(payload.get("limit", 20)), + ) + + +@with_db_session +def task_execution_stats( + task_name: str, + window_days: int = DEFAULT_WINDOW_DAYS, + limit: int = 20, + db_session: Session | None = None, +) -> dict: + """Aggregate one task's attempts over a window.""" + since = datetime.now(timezone.utc) - timedelta(days=window_days) + params = {"task_name": task_name, "since": since} + + by_variant = [ + { + "variant": row.variant, + "attempts": row.attempts, + "completed": row.completed, + "failed": row.failed, + "escalated": row.escalated, + "on_floor": row.on_floor, + "peak_vms_mib": _mib(row.peak_vms_bytes), + "peak_rss_mib": _mib(row.peak_rss_bytes), + "p50_ms": row.p50_ms, + "p95_ms": row.p95_ms, + "max_ms": row.max_ms, + } + for row in db_session.execute(_BY_VARIANT_SQL, params).all() + ] + + failures = [ + { + "failure_kind": row.failure_kind, + "error_type": row.error_type, + "attempts": row.attempts, + } + for row in db_session.execute(_BY_FAILURE_SQL, params).all() + ] + + repeats = [ + { + "entity_id": row.entity_id, + "attempts": row.attempts, + "escalated_to": row.escalated_to, + } + for row in db_session.execute( + _REPEAT_OFFENDERS_SQL, {**params, "limit": limit} + ).all() + ] + + totals = { + "attempts": sum(v["attempts"] for v in by_variant), + "completed": sum(v["completed"] for v in by_variant), + "failed": sum(v["failed"] for v in by_variant), + "escalated": sum(v["escalated"] for v in by_variant), + } + logger.info( + "%s over %s days: %s attempts, %s failed, %s escalated", + task_name, + window_days, + totals["attempts"], + totals["failed"], + totals["escalated"], + ) + + return { + "task_name": task_name, + "window_days": window_days, + "since": since.isoformat(), + "totals": totals, + "by_variant": by_variant, + "failures": failures, + "repeat_entities": repeats, + } + + +def _mib(value) -> float | None: + """Bytes as MiB, which is the unit the worker limits are expressed in.""" + return round(value / (1024**2), 1) if value else None diff --git a/infra/batch/main.tf b/infra/batch/main.tf index ca671d18a..c40930d7c 100644 --- a/infra/batch/main.tf +++ b/infra/batch/main.tf @@ -84,9 +84,25 @@ resource "google_storage_bucket" "datasets_bucket" { soft_delete_policy { retention_duration_seconds = local.retention_duration_seconds } + # Deletes objects on the date their own `customTime` names, which is how generated + # Parquet expires - the builder stamps each file with its retention date on upload. + # The condition is never satisfied for an object with no customTime, so this rule + # reaches nothing else in the bucket. The Parquet builder must stay the only writer + # of that field here: anything else that sets it becomes deletable by this rule. + lifecycle_rule { + action { + type = "Delete" + } + condition { + days_since_custom_time = 0 + } + } cors { origin = ["*"] - method = ["GET"] + # HEAD as well as GET: a browser reading Parquet over range requests probes with + # HEAD first, and GCS matches CORS methods literally rather than treating HEAD as + # a kind of GET. + method = ["GET", "HEAD"] response_header = ["*"] } } diff --git a/infra/functions-python/main.tf b/infra/functions-python/main.tf index 74fd046fe..956aca2b2 100644 --- a/infra/functions-python/main.tf +++ b/infra/functions-python/main.tf @@ -71,6 +71,54 @@ locals { function_gtfs_file_data_extractor_config = jsondecode(file("${path.module}/../../functions-python/gtfs_file_data_extractor/function_config.json")) function_gtfs_file_data_extractor_zip = "${path.module}/../../functions-python/gtfs_file_data_extractor/.dist/gtfs_file_data_extractor.zip" + + function_parquet_builder_config = jsondecode(file("${path.module}/../../functions-python/parquet_builder/function_config.json")) + # Parquet builds are routed to one of these by the Operations API, on the largest + # single uncompressed file in the dataset. The bands are measured, not guessed: of + # 4277 feeds the median archive is 0.2 MB and only 31 exceed 100 MB, while about ten + # have a single member over 1 GB and the worst has one of 4.8 GB. + # + # Two rules set the numbers: + # - `volume` must hold the largest member plus its Parquet output. It is carved OUT + # of `memory`, not added to it: limit_gcp_memory subtracts it to set RLIMIT_AS. + # - the process budget left over (memory - volume - 200MiB) tracks `duckdb_memory` + # plus ~330 MB, because RSS stops growing once DuckDB hits its cap. It does not + # scale with the feed, which is why the volume is what changes between rungs. + parquet_builder_sizes = { + s = { + # 3Gi, for a 1848 MiB budget. At 2Gi the budget was 824 MiB and DuckDB died with + # OutOfMemoryException on a 41.7 MB feed: measured VMS is ~816 MB on a 0.7 MB + # feed, so there was nothing left. Baseline VMS is ~550 MB before any work. + memory = "3Gi" + cpu = "1" + volume = "1Gi" + duckdb_memory = "512MB" + max_instances = 60 + max_dispatches = 40 + } + m = { + # 7Gi, for a 3896 MiB budget. At 6Gi it worked, but peak VMS reached 2297 MB of a + # 2872 MiB budget (80%) on a 1332 MB member, and the band admits up to 1.5 GB. + memory = "7Gi" + cpu = "2" + volume = "3Gi" + duckdb_memory = "1GB" + max_instances = 40 + max_dispatches = 20 + } + l = { + # 16Gi, for a 7992 MiB budget. Measured peak VMS is 5571 MB (70%); peak RSS is + # only 2496 MB, which is why sizing this from RSS put it at 3896 MiB and killed + # mdb-2014 with MemoryError. VMS runs ~2.2-2.8x the DuckDB cap - size from that. + memory = "16Gi" + cpu = "4" + volume = "8Gi" + duckdb_memory = "2GB" + max_instances = local.function_parquet_builder_config.max_instance_count + max_dispatches = 5 + } + } + function_parquet_builder_zip = "${path.module}/../../functions-python/parquet_builder/.dist/parquet_builder.zip" } locals { @@ -84,7 +132,8 @@ locals { local.function_tasks_executor_config.secret_environment_variables, local.function_pmtiles_builder_config.secret_environment_variables, local.function_gtfs_datasets_comparer_config.secret_environment_variables, - local.function_gtfs_file_data_extractor_config.secret_environment_variables + local.function_gtfs_file_data_extractor_config.secret_environment_variables, + local.function_parquet_builder_config.secret_environment_variables ) # Remove duplicates by key, keeping the first occurrence @@ -251,6 +300,13 @@ resource "google_storage_bucket_object" "pmtiles_builder_zip" { source = local.function_pmtiles_builder_zip } +# 18. Parquet Builder +resource "google_storage_bucket_object" "parquet_builder_zip" { + bucket = google_storage_bucket.functions_bucket.name + name = "parquet-builder-${substr(filebase64sha256(local.function_parquet_builder_zip), 0, 10)}.zip" + source = local.function_parquet_builder_zip +} + # 16. GTFS Change Tracker resource "google_storage_bucket_object" "gtfs_datasets_comparer_zip" { bucket = google_storage_bucket.functions_bucket.name @@ -868,6 +924,11 @@ resource "google_cloudfunctions2_function" "operations_api" { GOOGLE_CLIENT_ID = var.operations_oauth2_client_id DATASET_PROCESSING_TOPIC_NAME = "datasets-batch-topic-${var.environment}" WEB_REVALIDATION_QUEUE = google_cloud_tasks_queue.web_revalidation_task_queue.name + PARQUET_BUILDER_QUEUE_S = google_cloud_tasks_queue.parquet_builder_task_queue["s"].name + PARQUET_BUILDER_QUEUE_M = google_cloud_tasks_queue.parquet_builder_task_queue["m"].name + PARQUET_BUILDER_QUEUE_L = google_cloud_tasks_queue.parquet_builder_task_queue["l"].name + DATASETS_BUCKET_NAME = "${var.datasets_bucket_name}-${var.environment}" + PUBLIC_HOSTED_DATASETS_URL = local.public_hosted_datasets_url } available_memory = local.function_operations_api_config.memory timeout_seconds = local.function_operations_api_config.timeout @@ -1238,6 +1299,27 @@ resource "google_cloudfunctions2_function_iam_member" "pmtiles_builder_invoker" member = "serviceAccount:${google_service_account.functions_service_account.email}" } +# Grant execution permission to the service account to the parquet_builder function. +# The operations-api function enqueues the Cloud Task, and the task is dispatched with +# an OIDC token for this same service account. +resource "google_cloudfunctions2_function_iam_member" "parquet_builder_invoker" { + for_each = local.parquet_builder_sizes + project = var.project_id + location = var.gcp_region + cloud_function = google_cloudfunctions2_function.parquet_builder[each.key].name + role = "roles/cloudfunctions.invoker" + member = "serviceAccount:${google_service_account.functions_service_account.email}" +} + +resource "google_cloud_run_service_iam_member" "parquet_builder_cloud_run_invoker" { + for_each = local.parquet_builder_sizes + project = var.project_id + location = var.gcp_region + service = google_cloudfunctions2_function.parquet_builder[each.key].name + role = "roles/run.invoker" + member = "serviceAccount:${google_service_account.functions_service_account.email}" +} + # Grant execution permission to batchfunctions service account to the gtfs_datasets_comparer function resource "google_cloudfunctions2_function_iam_member" "gtfs_datasets_comparer_invoker_batch_sa" { project = var.project_id @@ -1402,6 +1484,31 @@ resource "google_cloud_tasks_queue" "pmtiles_builder_task_queue" { } } +# One task queue per worker size. Separate queues are the point rather than a side +# effect: with a single queue a run of large builds fills every dispatch slot and the +# short ones wait behind them, and the concurrency cap has to be set for the heaviest +# job. max_concurrent_dispatches caps memory across concurrent instances, not throughput, +# so the cheap worker can run far more of them. +resource "google_cloud_tasks_queue" "parquet_builder_task_queue" { + for_each = local.parquet_builder_sizes + project = var.project_id + location = var.gcp_region + name = "parquet-builder-queue-${each.key}-${var.environment}-${local.deployment_timestamp}" + + rate_limits { + max_concurrent_dispatches = each.value.max_dispatches + max_dispatches_per_second = 1 + } + + retry_config { + # One attempt. A conversion that fails does so deterministically - a corrupt + # archive is still corrupt - and the builder records the reason and returns 200 + # rather than letting Cloud Tasks retry it. Retrying would also fight the + # database claim the builder takes for the duration. + max_attempts = 1 + } +} + # Task queue to invoke gtfs_datasets_comparer function for backfill changelog tasks resource "google_cloud_tasks_queue" "gtfs_datasets_comparer_backfill_task_queue" { project = var.project_id @@ -1549,6 +1656,115 @@ resource "google_cloudfunctions2_function" "pmtiles_builder" { } +# 18. functions/parquet_builder cloud function +resource "google_cloudfunctions2_function" "parquet_builder" { + for_each = local.parquet_builder_sizes + name = "${local.function_parquet_builder_config.name}-${each.key}-${var.environment}" + project = var.project_id + description = local.function_parquet_builder_config.description + location = var.gcp_region + depends_on = [google_secret_manager_secret_iam_member.secret_iam_member] + + build_config { + runtime = var.python_runtime + entry_point = local.function_parquet_builder_config.entry_point + source { + storage_source { + bucket = google_storage_bucket.functions_bucket.name + object = google_storage_bucket_object.parquet_builder_zip.name + } + } + } + service_config { + environment_variables = { + ENVIRONMENT = var.environment + PROJECT_ID = var.project_id + GCP_REGION = var.gcp_region + SERVICE_ACCOUNT_EMAIL = google_service_account.functions_service_account.email + DATASETS_BUCKET_NAME = "${var.datasets_bucket_name}-${var.environment}" + PUBLIC_HOSTED_DATASETS_URL = local.public_hosted_datasets_url + # Everything large is written here, on the in-memory volume mounted below. + # limit_gcp_memory subtracts the volume's size from the process budget, so a + # conversion that overshoots raises MemoryError instead of being SIGKILLed. + PARQUET_TMPDIR = "/tmp/in-memory" + # Well under what the limiter leaves us. DuckDB's own default reads the host's + # RAM rather than the cgroup, so unset it spills far too late to help. + PARQUET_DUCKDB_MEMORY_LIMIT = each.value.duckdb_memory + # The builder re-queues its own work one rung up when it runs out of resources, so + # it needs the queues the Operations API uses. Without these the enqueue is a + # logged no-op and the escalation would silently do nothing. + PARQUET_BUILDER_QUEUE_S = google_cloud_tasks_queue.parquet_builder_task_queue["s"].name + PARQUET_BUILDER_QUEUE_M = google_cloud_tasks_queue.parquet_builder_task_queue["m"].name + PARQUET_BUILDER_QUEUE_L = google_cloud_tasks_queue.parquet_builder_task_queue["l"].name + } + available_memory = each.value.memory + timeout_seconds = local.function_parquet_builder_config.timeout + available_cpu = each.value.cpu + max_instance_request_concurrency = local.function_parquet_builder_config.max_instance_request_concurrency + max_instance_count = each.value.max_instances + min_instance_count = local.function_parquet_builder_config.min_instance_count + service_account_email = google_service_account.functions_service_account.email + ingress_settings = "ALLOW_ALL" + vpc_connector = data.google_vpc_access_connector.vpc_connector.id + vpc_connector_egress_settings = "PRIVATE_RANGES_ONLY" + + dynamic "secret_environment_variables" { + for_each = local.function_parquet_builder_config.secret_environment_variables + content { + key = secret_environment_variables.value["key"] + project_id = var.project_id + secret = lookup(secret_environment_variables.value, "secret", "${upper(var.environment)}_${secret_environment_variables.value["key"]}") + version = "latest" + } + } + } +} + + +# google_cloudfunctions2_function does not expose volume mounts in its schema, so the +# in-memory volume is attached to the underlying Cloud Run service after deploy. +# triggers_replace includes `source_zip` because a new deployment replaces the Cloud +# Run revision without this volume, and `size` because the idempotency check below +# matches volume names only. +resource "terraform_data" "parquet_builder_volume_mount" { + for_each = local.parquet_builder_sizes + triggers_replace = { + function_name = google_cloudfunctions2_function.parquet_builder[each.key].name + region = var.gcp_region + project = var.project_id + size = each.value.volume + source_zip = google_storage_bucket_object.parquet_builder_zip.name + } + + provisioner "local-exec" { + command = <<-EOT + MOUNTS=$(gcloud run services describe ${google_cloudfunctions2_function.parquet_builder[each.key].name} \ + --project ${var.project_id} \ + --region ${var.gcp_region} \ + --format='value(spec.template.spec.volumes[].name)' 2>/dev/null) + + if echo "$MOUNTS" | grep -q "in-memory"; then + echo "In-memory volume already mounted; removing it so the size limit is re-applied." + gcloud run services update ${google_cloudfunctions2_function.parquet_builder[each.key].name} \ + --project ${var.project_id} \ + --region ${var.gcp_region} \ + --remove-volume-mount volume=in-memory \ + --remove-volume in-memory \ + --quiet + fi + + gcloud run services update ${google_cloudfunctions2_function.parquet_builder[each.key].name} \ + --project ${var.project_id} \ + --region ${var.gcp_region} \ + --add-volume name=in-memory,type=in-memory,size-limit=${each.value.volume} \ + --add-volume-mount volume=in-memory,mount-path=/tmp/in-memory \ + --quiet + EOT + } + + depends_on = [google_cloudfunctions2_function.parquet_builder] +} + # 16. functions/gtfs_datasets_comparer cloud function resource "google_cloudfunctions2_function" "gtfs_datasets_comparer" { name = "${local.function_gtfs_datasets_comparer_config.name}-${var.environment}" diff --git a/infra/functions-python/vars.tf b/infra/functions-python/vars.tf index 106359bac..e77e0d2d6 100644 --- a/infra/functions-python/vars.tf +++ b/infra/functions-python/vars.tf @@ -30,158 +30,158 @@ variable "environment" { } variable "python_runtime" { - type = string + type = string description = "Python runtime version" - default = "python311" + default = "python311" } variable "datasets_bucket_name" { type = string description = "Name of the bucket where the datasets are stored" - default = "mobilitydata-datasets" + default = "mobilitydata-datasets" } variable "public_hosted_datasets_dns" { - type = string + type = string description = "Public hosted DNS for datasets" - default = "files.mobilitydatabase.org" + default = "files.mobilitydatabase.org" } variable "validator_endpoint" { - type = string + type = string description = "URL of the web validator" - default = "https://stg-gtfs-validator-web-mbzoxaljzq-ue.a.run.app" + default = "https://stg-gtfs-validator-web-mbzoxaljzq-ue.a.run.app" } variable "gbfs_bucket_name" { - type = string - description = "Name of the bucket where the GBFS feeds are stored" - default = "mobilitydata-gbfs-snapshots" + type = string + description = "Name of the bucket where the GBFS feeds are stored" + default = "mobilitydata-gbfs-snapshots" } variable "gbfs_scheduler_schedule" { - type = string - description = "Schedule for the GBFS validator daily job (version extraction and validation, Mon–Sat only)" - default = "0 0 * * 1-6" # At 00:00 Mon–Sat (Sunday is handled by the weekly geolocation job) + type = string + description = "Schedule for the GBFS validator daily job (version extraction and validation, Mon–Sat only)" + default = "0 0 * * 1-6" # At 00:00 Mon–Sat (Sunday is handled by the weekly geolocation job) } variable "gbfs_geolocation_scheduler_schedule" { - type = string - description = "Schedule for the GBFS geolocation extraction weekly job" - default = "0 0 * * 0" # At 00:00 every Sunday + type = string + description = "Schedule for the GBFS geolocation extraction weekly job" + default = "0 0 * * 0" # At 00:00 every Sunday } variable "jbda_scheduler_schedule" { - type = string - description = "Schedule for the JBDA scheduler job" - default = "0 0 3 * *" # At 00:00 on the 3rd day of every month + type = string + description = "Schedule for the JBDA scheduler job" + default = "0 0 3 * *" # At 00:00 on the 3rd day of every month } variable "odpt_scheduler_schedule" { - type = string - description = "Schedule for the ODPT scheduler job" - default = "0 0 5 * *" # At 00:00 on the 5th day of every month + type = string + description = "Schedule for the ODPT scheduler job" + default = "0 0 5 * *" # At 00:00 on the 5th day of every month } variable "gtfs_feed_availability_check_schedule" { - type = string - description = "Cron schedule for the GTFS feed availability check job" - default = "0 2 * * *" # Daily at 02:00 UTC + type = string + description = "Cron schedule for the GTFS feed availability check job" + default = "0 2 * * *" # Daily at 02:00 UTC } variable "notification_dispatch_daily_schedule" { - type = string - description = "Cron schedule for the daily notification dispatcher job" - default = "0 8 * * *" # Daily at 08:00 UTC + type = string + description = "Cron schedule for the daily notification dispatcher job" + default = "0 8 * * *" # Daily at 08:00 UTC } variable "reconcile_announcements_schedule" { - type = string - description = "Cron schedule for the Brevo announcements reconciliation job" - default = "0 3 * * *" # Daily at 03:00 UTC + type = string + description = "Cron schedule for the Brevo announcements reconciliation job" + default = "0 3 * * *" # Daily at 03:00 UTC } variable "seal_orchestrator_schedule" { - type = string - description = "Cron schedule for the nightly Seal of Reliability orchestrator job" - # Daily at 09:00 UTC. Deliberately after process-validation-report has finished writing the - # night's reports to the database - it works through them at about 80/hour and does not finish - # until roughly 08:30. - default = "0 9 * * *" + type = string + description = "Cron schedule for the nightly Seal of Reliability orchestrator job" + # Daily at 09:00 UTC. Deliberately after process-validation-report has finished writing the + # night's reports to the database - it works through them at about 80/hour and does not finish + # until roughly 08:30. + default = "0 9 * * *" } variable "purge_early_access_invites_schedule" { - type = string - description = "Cron schedule for the nightly early access invited-email purge job" - default = "0 5 * * *" # Daily at 05:00 UTC + type = string + description = "Cron schedule for the nightly early access invited-email purge job" + default = "0 5 * * *" # Daily at 05:00 UTC } variable "notification_dispatch_weekly_weekday" { - type = number - description = "Weekday the weekly digest is sent by the daily dispatcher (Monday=0 .. Sunday=6)" - default = 0 # Monday + type = number + description = "Weekday the weekly digest is sent by the daily dispatcher (Monday=0 .. Sunday=6)" + default = 0 # Monday } variable "tdg_api_token" { - type = string - description = "TDG API key" + type = string + description = "TDG API key" } variable "web_app_revalidate_url" { - type = string + type = string description = "URL of the website revalidation endpoint for cache invalidation" - default = "" + default = "" } variable "web_app_revalidate_secret" { - type = string + type = string description = "Secret token used to authenticate requests to the website revalidation endpoint" - sensitive = true - default = "" + sensitive = true + default = "" } variable "transitland_scraping_schedule" { - type = string - description = "Schedule Transitland scraping job" - default = "0 15 3 * *" # Runs at 00:00 JST on the 3rd day of every month + type = string + description = "Schedule Transitland scraping job" + default = "0 15 3 * *" # Runs at 00:00 JST on the 3rd day of every month } variable "transitland_api_key" { - type = string - description = "Transitland API key" + type = string + description = "Transitland API key" } variable "operations_oauth2_client_id" { - type = string + type = string description = "value of the OAuth2 client id for the Operations API" } variable "export_csv_schedule" { - type = string - description = "Schedule the export_csv function" - default = "0 4 * * 2,5" # At 4am every Tuesday and Friday. + type = string + description = "Schedule the export_csv function" + default = "0 4 * * 2,5" # At 4am every Tuesday and Friday. } variable "update_feed_status_schedule" { - type = string - description = "Schedule the update_feed_status function" - default = "0 4 * * *" # At 4am every day. + type = string + description = "Schedule the update_feed_status function" + default = "0 4 * * *" # At 4am every day. } variable "generate_sitemap_schedule" { - type = string - description = "Schedule the mobilitydatabase.org sitemap generation task" - default = "0 8 * * *" # At 08:00 UTC every day. + type = string + description = "Schedule the mobilitydatabase.org sitemap generation task" + default = "0 8 * * *" # At 08:00 UTC every day. } variable "brevo_api_announcements_list_id" { - type = string - description = "Brevo list ID for API announcements" - default = "" + type = string + description = "Brevo list ID for API announcements" + default = "" } variable "gtfs_datasets_comparer_in_memory_size" { - type = string - description = "Size limit for the gtfs_datasets_comparer in-memory tmpfs volume" - default = "3Gi" + type = string + description = "Size limit for the gtfs_datasets_comparer in-memory tmpfs volume" + default = "3Gi" } diff --git a/liquibase/changelog.xml b/liquibase/changelog.xml index 336fa821d..66740b0db 100644 --- a/liquibase/changelog.xml +++ b/liquibase/changelog.xml @@ -144,6 +144,10 @@ + + + + diff --git a/liquibase/changes/feat_parquet_builder_size.sql b/liquibase/changes/feat_parquet_builder_size.sql new file mode 100644 index 000000000..325365c98 --- /dev/null +++ b/liquibase/changes/feat_parquet_builder_size.sql @@ -0,0 +1,18 @@ +-- Register the per-feed Parquet worker size key in the existing config_key catalogue. +-- No new table: config_key and config_value_feed are generic. This row only satisfies the +-- FK config_value_feed(namespace, key) -> config_key, without which no value can be +-- written for a feed at all. +-- +-- The value is "s", "m" or "l". It may also be an object carrying the same size plus the +-- source that set it, which is how the builder records a size it chose itself, and a +-- "locked" flag an operator sets to hold the size where it is. +-- +-- No default_value on purpose: one would move the whole catalogue at once, which belongs +-- in the routing table in the code. With no per-feed row, a feed is routed by measurement. +INSERT INTO config_key (namespace, key, description) +VALUES ( + 'parquet_builder', + 'size', + 'Parquet build worker size for this feed: "s", "m" or "l", or an object of the same size with the source that set it and an optional "locked" flag.' +) +ON CONFLICT (namespace, key) DO NOTHING; diff --git a/liquibase/changes/feat_task_execution_attempt.sql b/liquibase/changes/feat_task_execution_attempt.sql new file mode 100644 index 000000000..533fb7740 --- /dev/null +++ b/liquibase/changes/feat_task_execution_attempt.sql @@ -0,0 +1,49 @@ +-- Append-only history of individual task attempts. +-- +-- task_execution_log holds one row per (task_name, entity_id, run_id) and rewrites it in +-- place: try_acquire nulls error_message and completed_at and resets triggered_at, and +-- metadata is overwritten whole by every heartbeat. That row answers "where does this +-- entity stand now" and cannot answer "what has happened to it", which is what sizing +-- decisions and a retry cap both need. +-- +-- One row per attempt, never updated. Deliberately not keyed to task_execution_log: a row +-- there may be rewritten or, for a run that is cleaned up, removed, and the history should +-- outlive it. +CREATE TABLE task_execution_attempt ( + id BIGSERIAL PRIMARY KEY, + task_name VARCHAR NOT NULL, + entity_id VARCHAR, + run_id VARCHAR NOT NULL, + -- 1-based, counted since the entity last succeeded. + attempt INTEGER NOT NULL, + -- Which worker ran it, and why. `variant` alone says what ran; the other three are + -- what make a change legible rather than inferred from neighbouring rows. + variant VARCHAR, + variant_basis VARCHAR, + override_at_attempt VARCHAR, + escalated_to VARCHAR, + status VARCHAR NOT NULL, + -- Null on success. Resource kinds are the ones an escalation may act on. + failure_kind VARCHAR, + -- The exception class, separately from its message: str(MemoryError()) is the empty + -- string, so the failure the largest worker exists for cannot be recognised from the + -- message alone. + error_type VARCHAR, + error_message TEXT, + started_at TIMESTAMPTZ NOT NULL, + finished_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + duration_ms BIGINT, + peak_rss_bytes BIGINT, + -- Address space, which is what RLIMIT_AS caps. Sizing from RSS under-provisions by + -- roughly half. + peak_vms_bytes BIGINT, + metadata JSONB +); + +-- "What has happened to this entity", for the retry cap and the per-feed history. +CREATE INDEX idx_task_execution_attempt_entity + ON task_execution_attempt (task_name, entity_id, run_id, finished_at DESC); + +-- "How is this task doing lately", for the stats report. +CREATE INDEX idx_task_execution_attempt_task_time + ON task_execution_attempt (task_name, finished_at DESC); diff --git a/scripts/parquet-generate-local.sh b/scripts/parquet-generate-local.sh new file mode 100755 index 000000000..d047596b3 --- /dev/null +++ b/scripts/parquet-generate-local.sh @@ -0,0 +1,63 @@ +#!/bin/bash +# +# +# MobilityData 2026 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# + +# Converts a GTFS feed to Parquet on this machine, so the viewer can be worked on +# without a bucket, a database, a task queue or any GCP credentials. The conversion is +# the parquet_builder function's own, imported rather than copied, so the output is +# what a real build would produce. +# +# Only duckdb is needed; it is installed into a throwaway virtualenv on first run. +# +# Usage: +# parquet-generate-local.sh [options] +# +# Options are passed straight through; run with --help for the full list. +# +# Examples: +# parquet-generate-local.sh mdb-1210 +# parquet-generate-local.sh mdb-1210-202402121801 --env dev +# parquet-generate-local.sh ./gtfs.zip --out ../ops-web/public/datasets/mdb-1210 +# parquet-generate-local.sh mdb-1210 --serve + +set -euo pipefail + +SCRIPT_PATH="$(cd "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +FUNCTION_PATH="$SCRIPT_PATH/../functions-python/parquet_builder" +VENV="$FUNCTION_PATH/venv" + +if [ $# -eq 0 ]; then + echo "Usage: parquet-generate-local.sh [--out DIR] [--serve [PORT]]" >&2 + echo "Run with --help for all options." >&2 + exit 1 +fi + +# The function's own venv when it exists (tests create it), otherwise a minimal one. +# Only duckdb is imported by this path - none of the cloud dependencies are touched. +if [ ! -x "$VENV/bin/python" ]; then + echo "INFO: creating virtualenv for parquet_builder" + python3 -m venv "$VENV" +fi + +if ! "$VENV/bin/python" -c "import duckdb" >/dev/null 2>&1; then + echo "INFO: installing duckdb" + "$VENV/bin/python" -m pip install --disable-pip-version-check -q duckdb +fi + +PYTHONPATH="$FUNCTION_PATH/src" exec "$VENV/bin/python" \ + "$FUNCTION_PATH/src/scripts/generate_local.py" "$@"