diff --git a/backend/src/ref_backend/api/routes/diagnostics.py b/backend/src/ref_backend/api/routes/diagnostics.py index b5ede0a3..d2abc33f 100644 --- a/backend/src/ref_backend/api/routes/diagnostics.py +++ b/backend/src/ref_backend/api/routes/diagnostics.py @@ -1,7 +1,7 @@ from typing import Literal, NamedTuple from fastapi import APIRouter, HTTPException, Query, Request -from sqlalchemy.orm import joinedload, selectinload +from sqlalchemy.orm import joinedload from starlette.responses import StreamingResponse from climate_ref import models @@ -60,15 +60,8 @@ def _get_diagnostic( return diagnostic -def _listed_diagnostics( - app_context: AppContextDep, *, with_group_ids: bool = False -) -> list[models.Diagnostic]: +def _listed_diagnostics(app_context: AppContextDep) -> list[models.Diagnostic]: query = app_context.session.query(models.Diagnostic).options(joinedload(models.Diagnostic.provider)) - if with_group_ids: - # One query for every diagnostic's group ids, rather than one per diagnostic. - query = query.options( - selectinload(models.Diagnostic.execution_groups).load_only(models.ExecutionGroup.id) - ) if app_context.settings.DIAGNOSTIC_PROVIDERS: query = query.join(models.Provider).filter( models.Provider.slug.in_(app_context.settings.DIAGNOSTIC_PROVIDERS) @@ -85,12 +78,10 @@ class _Scope(NamedTuple): filters: DiagnosticFilter -def _scoped_diagnostics( - app_context: AppContextDep, mip_era: MipEra | None, source_id: str | None, *, with_group_ids: bool = False -) -> _Scope: +def _scoped_diagnostics(app_context: AppContextDep, mip_era: MipEra | None, source_id: str | None) -> _Scope: """Get the listed diagnostics, counting only the groups that draw on `mip_era` and ran `source_id`""" return _Scope( - _listed_diagnostics(app_context, with_group_ids=with_group_ids), + _listed_diagnostics(app_context), DiagnosticFilter( include_superseded=False, include_retracted=False, @@ -156,7 +147,7 @@ def _list( Counts cover the promoted version of each diagnostic and leave out retracted executions. """ - scope = _scoped_diagnostics(app_context, mip_era, source_id, with_group_ids=True) + scope = _scoped_diagnostics(app_context, mip_era, source_id) value_types = _value_types(app_context, scope) return Collection( data=[ @@ -182,7 +173,7 @@ def catalog( Checking for metric values is most of the cost of the full listing. Fetch those from `/diagnostics/value-flags`. """ - scope = _scoped_diagnostics(app_context, mip_era, source_id, with_group_ids=True) + scope = _scoped_diagnostics(app_context, mip_era, source_id) return Collection(data=_catalog_entries(app_context, scope)) diff --git a/backend/src/ref_backend/api/routes/executions.py b/backend/src/ref_backend/api/routes/executions.py index f570a920..df6d197c 100644 --- a/backend/src/ref_backend/api/routes/executions.py +++ b/backend/src/ref_backend/api/routes/executions.py @@ -1,9 +1,4 @@ import mimetypes -import os -import tarfile -import tempfile -from collections.abc import Generator -from pathlib import Path from typing import Literal, get_args from fastapi import APIRouter, HTTPException, Query, Request @@ -18,7 +13,7 @@ from climate_ref_core.pycmec.metric import CMECMetric from ref_backend.api.deps import AppContext, AppContextDep from ref_backend.core.execution_groups import load_execution_groups -from ref_backend.core.file_handling import file_iterator, resolve_artifact, tail_lines +from ref_backend.core.file_handling import file_iterator, resolve_artifact, tail_lines, tar_gz_iterator from ref_backend.core.metric_values import ( MetricValueType, parse_id_list, @@ -83,7 +78,7 @@ def group_counts(mip_era: MipEra | None = None) -> GroupCounts: @router.get("/") def list_recent_execution_groups( # noqa: PLR0913, PLR0917 app_context: AppContextDep, - limit: int = Query(10, ge=0), + limit: int = Query(10, ge=0, le=1000), offset: int = Query(0, ge=0), diagnostic_name_contains: str | None = None, provider_name_contains: str | None = None, @@ -243,7 +238,7 @@ def execution_logs( logger.warning(f"Log file not found: {file_path}") raise HTTPException(status_code=404, detail="Log file not found") - headers = {"Content-Disposition": f"attachment; filename=execution_result_{execution_id}.log"} + headers = {"Content-Disposition": f"attachment; filename=execution_result_{execution.id}.log"} if lines is not None: tail = "".join(f"{line}\n" for line in tail_lines(file_path, lines)) return Response(tail, media_type=mime_type, headers=headers) @@ -367,28 +362,8 @@ def execution_archive( if not result_path.exists(): raise HTTPException(status_code=404, detail="Execution output not found") - # This is an arbitrary value as a placeholder - # No experimentation has been done to find the best chunk size - CHUNK_SIZE = 128 * 1024 # 128 KB chunks - - def generate_archive() -> Generator[bytes]: - with tempfile.NamedTemporaryFile(delete=True) as temp_tar: - # Open the tar file in write mode with gzip compression - with tarfile.open(temp_tar.name, mode="w:gz") as tar: - for root, _, files in os.walk(result_path): - for file in files: - file_path = Path(root) / file - arcname = str(file_path.relative_to(result_path)) - tar.add(file_path, arcname=arcname) - - # Read and stream the tar file in chunks - with open(temp_tar.name, "rb") as f: - while chunk := f.read(CHUNK_SIZE): - yield chunk - # The temp file will be deleted automatically when closed - return StreamingResponse( - generate_archive(), + tar_gz_iterator(result_path), media_type="application/x-gzip", headers={"Content-Disposition": f"attachment; filename=execution_{execution.id}.tar.gz"}, ) diff --git a/backend/src/ref_backend/core/file_handling.py b/backend/src/ref_backend/core/file_handling.py index f6547b4c..ba482b16 100644 --- a/backend/src/ref_backend/core/file_handling.py +++ b/backend/src/ref_backend/core/file_handling.py @@ -1,4 +1,7 @@ import os +import stat +import tarfile +import zlib from collections.abc import Callable, Generator from pathlib import Path @@ -71,3 +74,54 @@ def file_iterator(file_path: str, chunk_size: int = 64 * 1024) -> Generator[byte with open(file_path, "rb") as file: while chunk := file.read(chunk_size): yield chunk + + +def _tar_entries(root: Path, chunk_size: int) -> Generator[bytes]: + """Yield an uncompressed tar of the regular files below `root`, one chunk at a time.""" + for directory, _, files in os.walk(root): + for file in sorted(files): + file_path = Path(directory) / file + file_stat = file_path.lstat() + if not stat.S_ISREG(file_stat.st_mode): + continue + + info = tarfile.TarInfo(str(file_path.relative_to(root))) + info.size = file_stat.st_size + info.mtime = int(file_stat.st_mtime) + info.mode = stat.S_IMODE(file_stat.st_mode) + yield info.tobuf(tarfile.PAX_FORMAT) + + remaining = info.size + with open(file_path, "rb") as f: + while remaining and (chunk := f.read(min(chunk_size, remaining))): + remaining -= len(chunk) + yield chunk + # A file that shrank while being read is zero-filled, so the header stays true. + padding = remaining + (-info.size % tarfile.BLOCKSIZE) + while padding: + chunk_length = min(chunk_size, padding) + padding -= chunk_length + yield bytes(chunk_length) + + yield bytes(2 * tarfile.BLOCKSIZE) + + +def tar_gz_iterator(root: Path, chunk_size: int = 64 * 1024) -> Generator[bytes]: + """ + Stream a gzipped tar of the regular files below a directory + + Memory use is bounded by `chunk_size`, whatever the size of the files. + Symlinks and other special files are left out. + + Parameters + ---------- + root + Directory to archive, which becomes the root of the archive + chunk_size + Size of each chunk to read from the files + """ + compressor = zlib.compressobj(wbits=zlib.MAX_WBITS | 16) + for data in _tar_entries(root, chunk_size): + if compressed := compressor.compress(data): + yield compressed + yield compressor.flush() diff --git a/backend/src/ref_backend/models/diagnostics.py b/backend/src/ref_backend/models/diagnostics.py index a44db7c1..73d33f27 100644 --- a/backend/src/ref_backend/models/diagnostics.py +++ b/backend/src/ref_backend/models/diagnostics.py @@ -64,10 +64,6 @@ class DiagnosticCatalogEntry(BaseModel): """ Description of the diagnostic """ - execution_groups: list[int] - """ - List of IDs for the provider executions associated with this provider - """ groups: GroupCounts """ Execution groups at the promoted version, by the outcome of their latest execution @@ -109,9 +105,8 @@ def _ensure_metadata_cache(app_context: "AppContext") -> Mapping[str, Diagnostic return load_diagnostic_metadata_cached(app_context.settings.diagnostic_metadata_path_resolved) @staticmethod - def _build_group_by_summary(diagnostic: models.Diagnostic, app_context: "AppContext") -> list[GroupBy]: + def _build_group_by_summary(concrete_diagnostic: "Diagnostic | None") -> list[GroupBy]: """Extract and build group_by information from diagnostic data requirements.""" - concrete_diagnostic = _concrete_diagnostic(diagnostic, app_context) if concrete_diagnostic is None: return [] data_requirements = sorted( @@ -147,7 +142,7 @@ def _get_aft_link(diagnostic: models.Diagnostic) -> "AFTDiagnosticDetail | None" if aft_id: return get_aft_diagnostic_by_id(aft_id) else: - logger.warning(f"No AFT found for diagnostic {diagnostic.provider.slug}/{diagnostic.slug}") + logger.debug(f"No AFT found for diagnostic {diagnostic.provider.slug}/{diagnostic.slug}") return None @staticmethod @@ -183,10 +178,9 @@ def build_with_stats( ) -> "DiagnosticCatalogEntry": """Build a DiagnosticCatalogEntry with pre-computed statistics to avoid N+1 queries.""" metadata_cache = DiagnosticCatalogEntry._ensure_metadata_cache(app_context) - group_by_summary = DiagnosticCatalogEntry._build_group_by_summary(diagnostic, app_context) - aft = DiagnosticCatalogEntry._get_aft_link(diagnostic) - concrete_diagnostic = _concrete_diagnostic(diagnostic, app_context) + group_by_summary = DiagnosticCatalogEntry._build_group_by_summary(concrete_diagnostic) + aft = DiagnosticCatalogEntry._get_aft_link(diagnostic) description = (concrete_diagnostic.__doc__ or "") if concrete_diagnostic else "" # Build the base diagnostic summary @@ -196,7 +190,6 @@ def build_with_stats( slug=diagnostic.slug, name=diagnostic.name, description=description, - execution_groups=[e.id for e in diagnostic.execution_groups], groups=groups, promoted_version=diagnostic.promoted_version, group_by=group_by_summary, diff --git a/backend/tests/test_api/test_routes/test_executions.py b/backend/tests/test_api/test_routes/test_executions.py index 175b353e..20bd1309 100644 --- a/backend/tests/test_api/test_routes/test_executions.py +++ b/backend/tests/test_api/test_routes/test_executions.py @@ -1,3 +1,6 @@ +import io +import tarfile + import pytest from fastapi.testclient import TestClient @@ -498,6 +501,8 @@ def test_execution_logs(client: TestClient, settings, monkeypatch, tmp_path): assert r.status_code == 200 assert r.text == log_file.read_text() + latest_id = client.get(f"{settings.API_V1_STR}/executions/{group_id}/execution").json()["id"] + assert r.headers["content-disposition"] == f"attachment; filename=execution_result_{latest_id}.log" def test_execution_logs_lines(client: TestClient, settings, monkeypatch, tmp_path): @@ -657,3 +662,24 @@ def test_execution_value_facets_of_an_unknown_group(client: TestClient, settings r = client.get(f"{settings.API_V1_STR}/executions/999999/values/facets", params={"value_type": "scalar"}) assert r.status_code == 404 + + +def test_execution_list_limit_capped(client: TestClient, settings) -> None: + r = client.get(f"{settings.API_V1_STR}/executions/?limit=1001") + + assert r.status_code == 422 + + +def test_execution_archive(client: TestClient, settings, monkeypatch, tmp_path) -> None: + group_id = get_execution_group_id(client, settings) + (tmp_path / "plots").mkdir() + (tmp_path / "diagnostic.json").write_text("{}") + (tmp_path / "plots" / "map.png").write_bytes(b"png" * 1000) + monkeypatch.setattr(executions_routes, "resolve_artifact", lambda *_: tmp_path) + + r = client.get(f"{settings.API_V1_STR}/executions/{group_id}/archive") + + assert r.status_code == 200 + with tarfile.open(fileobj=io.BytesIO(r.content), mode="r:gz") as tar: + assert sorted(tar.getnames()) == ["diagnostic.json", "plots/map.png"] + assert tar.extractfile("plots/map.png").read() == b"png" * 1000 diff --git a/backend/tests/test_core/test_file_handling.py b/backend/tests/test_core/test_file_handling.py index 27067e46..bc7b911d 100644 --- a/backend/tests/test_core/test_file_handling.py +++ b/backend/tests/test_core/test_file_handling.py @@ -1,6 +1,11 @@ -"""Tests for chunked file reading used by the log-download endpoint.""" +"""Tests for the chunked file reading behind the log and archive downloads.""" -from ref_backend.core.file_handling import file_iterator, tail_lines +import io +import os +import tarfile +import tracemalloc + +from ref_backend.core.file_handling import _tar_entries, file_iterator, tail_lines, tar_gz_iterator class TestFileIterator: @@ -69,3 +74,63 @@ def test_invalid_utf8_is_replaced(self, tmp_path): file_path.write_bytes(b"ok\n\xff bad\n") assert tail_lines(file_path, 5) == ["ok", "\ufffd bad"] + + +def _read_archive(chunks) -> dict[str, bytes]: + with tarfile.open(fileobj=io.BytesIO(b"".join(chunks)), mode="r:gz") as tar: + return {m.name: tar.extractfile(m).read() for m in tar.getmembers()} + + +class TestTarGzIterator: + """Test streaming a directory as a tar.gz.""" + + def test_round_trips_nested_files(self, tmp_path): + long_name = "d" * 150 + (tmp_path / "plots" / long_name).mkdir(parents=True) + (tmp_path / "diagnostic.json").write_text("{}") + (tmp_path / "empty.txt").write_bytes(b"") + (tmp_path / "plots" / long_name / "map.png").write_bytes(os.urandom(3000)) + + archive = _read_archive(tar_gz_iterator(tmp_path, chunk_size=1024)) + + assert archive == { + "diagnostic.json": b"{}", + "empty.txt": b"", + f"plots/{long_name}/map.png": (tmp_path / "plots" / long_name / "map.png").read_bytes(), + } + + def test_empty_directory(self, tmp_path): + assert _read_archive(tar_gz_iterator(tmp_path)) == {} + + def test_leaves_out_symlinks(self, tmp_path): + (tmp_path / "real.txt").write_text("real") + (tmp_path / "link.txt").symlink_to(tmp_path / "real.txt") + + assert _read_archive(tar_gz_iterator(tmp_path)) == {"real.txt": b"real"} + + def test_file_that_shrinks_is_zero_filled(self, tmp_path): + """The header records the size first, so a file that shrinks mid-read is padded to match.""" + path = tmp_path / "log.txt" + path.write_bytes(b"x" * 5000) + entries = _tar_entries(tmp_path, chunk_size=1024) + next(entries) # the header, which recorded 5000 bytes + path.write_bytes(b"y" * 100) + + rest = list(entries) + + assert max(len(chunk) for chunk in rest) <= 1024 + assert b"".join(rest)[:5000] == b"y" * 100 + bytes(4900) + + def test_memory_does_not_grow_with_file_size(self, tmp_path): + # Random bytes do not compress, so holding the file would show up as 4 MB. + (tmp_path / "big.nc").write_bytes(os.urandom(4 * 1024 * 1024)) + + tracemalloc.start() + try: + for _ in tar_gz_iterator(tmp_path, chunk_size=16 * 1024): + pass + _, peak = tracemalloc.get_traced_memory() + finally: + tracemalloc.stop() + + assert peak < 1024 * 1024 diff --git a/changelog/172.breaking.md b/changelog/172.breaking.md new file mode 100644 index 00000000..351a5dce --- /dev/null +++ b/changelog/172.breaking.md @@ -0,0 +1,2 @@ +Removed the unused `execution_groups` list of ids from diagnostic summaries in the API. +`GET /executions/` now rejects a `limit` above 1000 with a 422 error, so use `limit` and `offset` to page through more groups. diff --git a/changelog/172.fix.md b/changelog/172.fix.md new file mode 100644 index 00000000..154ee27c --- /dev/null +++ b/changelog/172.fix.md @@ -0,0 +1 @@ +Log downloads are now named after the execution, rather than `execution_result_None.log`. diff --git a/changelog/172.improvement.md b/changelog/172.improvement.md new file mode 100644 index 00000000..2ec24431 --- /dev/null +++ b/changelog/172.improvement.md @@ -0,0 +1,2 @@ +Execution archive downloads now start straight away and use little server memory, however large the outputs. +Symbolic links in an execution's output folder are no longer included in the archive. diff --git a/frontend/src/client/types.gen.ts b/frontend/src/client/types.gen.ts index fab10203..d1074d13 100644 --- a/frontend/src/client/types.gen.ts +++ b/frontend/src/client/types.gen.ts @@ -657,10 +657,6 @@ export type DiagnosticCatalogEntry = { * Description */ description: string; - /** - * Execution Groups - */ - execution_groups: Array; groups: GroupCounts; /** * Promoted Version @@ -737,10 +733,6 @@ export type DiagnosticSummary = { * Description */ description: string; - /** - * Execution Groups - */ - execution_groups: Array; groups: GroupCounts; /** * Promoted Version diff --git a/frontend/src/lib/resourceUsageCsv.test.ts b/frontend/src/lib/resourceUsageCsv.test.ts index 23a11181..62f34088 100644 --- a/frontend/src/lib/resourceUsageCsv.test.ts +++ b/frontend/src/lib/resourceUsageCsv.test.ts @@ -5,7 +5,6 @@ import { resourceUsageCsv } from "./resourceUsageCsv"; const base = { id: 1, description: "", - execution_groups: [], has_metric_values: false, has_scalar_values: false, has_series_values: false,