Fix scoring scripts

scripts/comparison_utils.py

@@ -9,8 +9,7 @@
 import numpy as np
 from conversion.data_formats import MarkImpressionType, MarkStriationType, MarkType
 from conversion.surface_comparison.models import ComparisonParams
-
-from scripts.conversion_utils import parse_db_scratch
+from conversion_utils import parse_db_scratch
 
 logger = logging.getLogger(__name__)
 _MARK_TYPE_FOLDER_MAP: list[tuple[str, MarkType]] = sorted(

scripts/convert_scores.py

@@ -33,17 +33,16 @@
 from typing import Any
 
 import requests
-from conversion.data_formats import MarkImpressionType
-
-from scripts.comparison_utils import (
+from comparison_utils import (
     ComparisonEntry,
     _build_body,
     _save_result,
     find_all_mark_types,
     generate_pairs,
 )
-from scripts.conversion_utils import ConversionConfig, run_parallel
-from scripts.csv_pairs import (
+from conversion.data_formats import MarkImpressionType
+from conversion_utils import ConversionConfig, run_parallel
+from csv_pairs import (
     DONE_STATUSES,
     CsvTask,
     ScoreWriter,
@@ -53,7 +52,7 @@
     find_result_file,
     read_pairs_csv,
 )
-from scripts.http_utils import _cleanup_vault, _post_with_retry, download_urls
+from http_utils import _cleanup_vault, _post_with_retry, download_urls
 
 logging.basicConfig(level=logging.WARNING, format="%(levelname)s: %(message)s")
 logger = logging.getLogger(__name__)

scripts/csv_pairs.py

@@ -14,19 +14,20 @@
 
 from __future__ import annotations
 
+import contextlib
 import csv
 import logging
 import os
 import threading
+import uuid
 from collections import defaultdict, deque
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Any
 
+from comparison_utils import ComparisonEntry, infer_mark_type
 from conversion.data_formats import MarkImpressionType, MarkType
-
-from scripts.comparison_utils import ComparisonEntry, infer_mark_type
-from scripts.conversion_utils import ConversionConfig
+from conversion_utils import ConversionConfig
 
 logger = logging.getLogger(__name__)
 
@@ -103,8 +104,10 @@ def extract_metrics(result: dict[str, Any] | None, mark_type: MarkType) -> dict[
         return {}
 
     if isinstance(mark_type, MarkImpressionType):
+        cells = result.get("cells") if isinstance(result, dict) else None
+        total_cells = len(cells) if isinstance(cells, list) else None
         metrics = {
-            "total_cells": comparison_results.get(TOTAL_CELLS_KEY),
+            "total_cells": total_cells,
             "matching_cells": comparison_results.get(MATCHING_CELLS_KEY),
         }
     else:
@@ -313,8 +316,10 @@ def __init__(
         out_dir.mkdir(parents=True, exist_ok=True)
         for mark_type in mark_types:
             if resume:
-                self.values[mark_type] = self._read_previous(mark_type)
-            self._write(mark_type)
+                with self._lock:
+                    self.values[mark_type] = self._read_previous(mark_type)
+            with self._lock:
+                self._write(mark_type)
         logger.info("Writing scored copies to %s", ", ".join(p.name for p in self.paths.values()))
 
     def _read_previous(self, mark_type: MarkType) -> dict[int, dict[str, Any]]:
@@ -375,18 +380,23 @@ def flush(self) -> None:
                 self._pending[mark_type] = 0
 
     def _write(self, mark_type: MarkType) -> None:
-        """Rewrite one file. The caller must hold the lock."""
+        """Rewrite one file atomically. The caller must hold the lock."""
         path = self.paths[mark_type]
         columns = self.columns[mark_type]
         values = self.values[mark_type]
-        tmp = path.with_name(path.name + ".tmp")
-        with tmp.open("w", newline="", encoding="utf-8") as fh:
-            writer = csv.writer(fh, delimiter=self.delimiter)
-            if self.header is not None:
-                writer.writerow([*self.header, *columns])
-            for row in self.rows:
-                scores = values.get(row.index, {})
-                writer.writerow([*row.fields, *("" if scores.get(c) is None else scores[c] for c in columns)])
-            fh.flush()
-            os.fsync(fh.fileno())
-        os.replace(tmp, path)
+        tmp = path.with_name(f".{path.name}.{uuid.uuid4().hex}.tmp")
+        try:
+            with tmp.open("w", newline="", encoding="utf-8") as fh:
+                writer = csv.writer(fh, delimiter=self.delimiter)
+                if self.header is not None:
+                    writer.writerow([*self.header, *columns])
+                for row in self.rows:
+                    scores = values.get(row.index, {})
+                    writer.writerow([*row.fields, *("" if scores.get(c) is None else scores[c] for c in columns)])
+                fh.flush()
+                os.fsync(fh.fileno())
+            os.replace(tmp, path)
+        except BaseException:
+            with contextlib.suppress(OSError):
+                tmp.unlink()
+            raise