Coverage for src / quber / files / output.py: 53%

1"""Write extraction artifacts to a local directory or an ``s3://`` prefix. 

2 

3`output_sink` is a context manager that hands back a local directory to write 

4into. When the destination is an ``s3://`` URI, everything written there is 

5uploaded to the bucket on exit; local destinations are written in place and 

6never copied. 

7 

8The write-local-then-upload shape exists because the artifact writers (pymupdf 

9for the annotated PDF, the HTML renderer, ``json.dump``) all need a real local 

10filesystem path -- none can stream to S3 directly. Uploads go through 

11cloudpathlib's default client, which resolves AWS credentials from the boto3 

12chain (env / ``~/.aws`` / IAM role), so no keys are passed explicitly. 

13""" 

14 

15from __future__ import annotations 

16 

17import json 

18import tempfile 

19from contextlib import contextmanager 

20from pathlib import Path 

21from typing import Any, Generator 

22 

23from cloudpathlib import AnyPath, S3Path 

24from loguru import logger 

25 

26 

27def is_s3(output_dir: str) -> bool: 

28 """True when ``output_dir`` is an ``s3://`` URI rather than a local path.""" 

29 return isinstance(AnyPath(output_dir), S3Path) 

30 

31 

32@contextmanager 

33def output_sink(output_dir: str) -> Generator[Path, None, None]: 

34 """Yield a local directory to write artifacts into. 

35 

36 For a local ``output_dir`` the directory is created and yielded directly, 

37 so writes land in their final location. For an ``s3://`` ``output_dir`` a 

38 temporary directory is yielded; on clean exit every file written under it 

39 is uploaded to the bucket under the same relative path, then the temporary 

40 directory is removed. An exception inside the ``with`` block skips the 

41 upload and propagates, so a failed run never leaves partial artifacts in 

42 the bucket. 

43 """ 

44 dest = AnyPath(output_dir) 

45 if isinstance(dest, S3Path): 

46 with tempfile.TemporaryDirectory(prefix="quber-out-") as tmp: 

47 work = Path(tmp) 

48 yield work 

49 uploaded = 0 

50 for f in sorted(work.rglob("*")): 

51 if not f.is_file(): 

52 continue 

53 target = dest / f.relative_to(work).as_posix() 

54 target.upload_from(f, force_overwrite_to_cloud=True) 

55 logger.info("Uploaded {} -> {}", f.name, target) 

56 uploaded += 1 

57 logger.info("Uploaded {} artifact(s) to {}", uploaded, dest) 

58 else: 

59 local = Path(output_dir) 

60 local.mkdir(parents=True, exist_ok=True) 

61 yield local 

62 

63 

64def write_completion_marker(output_dir: str, filename: str, payload: dict[str, Any]) -> None: 

65 """Write a job's completion marker as its own, final upload. 

66 

67 The marker's existence tells downstream (the join step that launches 

68 fusion) that this job's artifacts are fully in place, so it must be the 

69 last object to land. It goes through its own `output_sink` after the 

70 artifact sink has closed: a single sink uploads files in sorted order, 

71 and a marker name can sort before the artifacts it vouches for. 

72 """ 

73 with output_sink(output_dir) as out: 

74 (out / filename).write_text(json.dumps(payload, indent=2, default=str), encoding="utf-8")