Fix/cloud oom

Dockerfile

@@ -20,10 +20,10 @@ FROM base AS app
 WORKDIR /app
 
 # Copy everything needed for the application install at once
-COPY requirements.txt constraints.txt dist/*.whl ./
+COPY requirements.txt dist/*.whl ./
 
 # Install app
-RUN uv pip install --system --compile-bytecode -c constraints.txt -r requirements.txt *.whl
+RUN uv pip install --system --compile-bytecode -r requirements.txt *.whl
 
 # --------------------------------------------------------------------
 # --- Test Target ---
@@ -36,7 +36,7 @@ COPY pyproject.toml README.md ./
 
 # Install dev dependencies
 ENV SETUPTOOLS_SCM_PRETEND_VERSION=0.0.0
-RUN uv pip install --system --group dev -c constraints.txt
+RUN uv pip install --system --group dev
 
 # Run tests
 RUN pytest tests/ -v

Makefile

@@ -1,7 +1,11 @@
 # Define all abstract commands here
-.PHONY: test test-image init 
+.PHONY: test test-image init lint
 
-# Rune tests
+# Lint
+lint:
+	uvx ruff check --fix && uvx ruff format
+
+# Run tests
 test:
 	uv run pytest tests/ -v
 

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "data-index"
 dynamic = ["version"]
-description = "Add your description here"
+description = "AODN Data Indexer for the `imos-data` public bucket."
 readme = "README.md"
 authors = [
     { name = "Tom Galindo", email = "[email protected]" }
@@ -30,6 +30,12 @@ dependencies = [
     "xarray>=2026.4.0",
 ]
 
+[tool.uv]
+constraint-dependencies = [
+    "prefect==3.6.28",
+    "prefect-aws==0.7.7",
+]
+
 [build-system]
 requires = ["hatchling", "hatch-vcs"]
 build-backend = "hatchling.build"

src/data_index/defaults/cluster.py

@@ -9,12 +9,12 @@
 from data_index.file_fetcher import S3Fetcher, S5CMDFetcher, ThresholdFileFetcher
 from data_index.iceberg_config import (
     IcebergTableConfig,
-    IcebergTableScanConfig,
     S3TablesCatalogConfig,
 )
-from data_index.inventory_source import LiveS3InventorySource, ParquetInventorySource
-from data_index.inventory_source.live_s3_facility_subset import (
-    LiveS3InventorySourceFacilitySubset,
+from data_index.inventory_source import (
+    LiveS3InventorySource,
+    ParquetInventorySource,
+    S3TableInventorySource,
 )
 from data_index.metadata_extractor import NetCDFExtractor, UnstructuedNetCDFExtractor
 from data_index.structured_metadata import StructuredMetadata
@@ -41,30 +41,18 @@
 # --- Live Inventory Source config
 _s3_metadata_catalog_config = S3TablesCatalogConfig(
     region=REGION,
-    arn="arn:aws:s3tables:ap-southeast-2:104044260116:bucket/aws-s3",
+    arn="arn:aws:s3tables:ap-southeast-2:704910415367:bucket/imos-data-inventory",
 )
 
 _inventory_table_config = IcebergTableConfig(
     catalog_config=_s3_metadata_catalog_config,
-    namespace="b_imos-data",
-    table_name="inventory",
+    namespace="inventory",
+    table_name="live",
 )
 
-_inventory_table_scan_config = IcebergTableScanConfig(
-    row_filter="key LIKE 'IMOS/%'",
-)
-
-_live_inventory_source = LiveS3InventorySourceFacilitySubset(
+_live_inventory_source = S3TableInventorySource(
     table_config=_inventory_table_config,
-    table_scan_config=_inventory_table_scan_config,
-    path=pathlib.Path(".extract/s3_metadata"),
-    skip_if_exists=True,
-    subset_per_facility=2_000,
-)
-
-# --- Static Inventory Source config ---
-_static_inventory_source = ParquetInventorySource(
-    source="s3://aodn-dataflow-dev/thomas.galindo/processing/stored/s3_metadata/"
+    subset_per_facility=10_000,
 )
 
 # --- Partitioner config ---

src/data_index/inventory_source/__init__.py

@@ -1,4 +1,5 @@
 from .live_s3 import LiveS3InventorySource
 from .parquet import ParquetInventorySource
+from .s3_table import S3TableInventorySource
 
-__all__ = ["LiveS3InventorySource", "ParquetInventorySource"]
+__all__ = ["LiveS3InventorySource", "ParquetInventorySource", "S3TableInventorySource"]

src/data_index/inventory_source/s3_table.py

@@ -0,0 +1,83 @@
+from __future__ import annotations
+
+import polars
+import pydantic
+
+from data_index.iceberg_config.iceberg_table_config import IcebergTableConfig
+from data_index.iceberg_config.table_scan_config import IcebergTableScanConfig
+from data_index.inventory_source import LiveS3InventorySource
+
+
+class S3TableInventorySource(LiveS3InventorySource):
+    """InventorySource that runs on a pre-live-evaluated S3 Metadata Table."""
+
+    subset_per_facility: int = pydantic.Field(default=10_000, ge=1)
+    table_config: IcebergTableConfig
+    table_scan_config: IcebergTableScanConfig = pydantic.Field(
+        default_factory=lambda: IcebergTableScanConfig(
+            selected_fields=("bucket", "key", "size", "facility"),
+        )
+    )
+    subset_per_facility: int | None = pydantic.Field(default=None)
+
+    @staticmethod
+    def _s3_uri_column() -> polars.Expr:
+        return polars.concat_str(
+            polars.lit("s3://"),
+            polars.col("bucket"),
+            polars.lit("/"),
+            polars.col("key"),
+        ).alias("s3_uri")
+
+    def _full_inventory(self) -> polars.DataFrame:
+        table = self.table_config.load()
+        df = table.scan(selected_fields=("bucket", "key", "size")).to_polars()
+        return df.select(
+            self._s3_uri_column(),
+            polars.col("size"),
+        )
+
+    def _subset_per_facility(self) -> polars.DataFrame:
+        table = self.table_config.load()
+
+        facilities_df = table.scan(selected_fields=("facility",)).to_polars()
+
+        # Find facilities, defined as second part of path
+        facilities = facilities_df["facility"].drop_nulls().unique()
+
+        # Takes samples of size `subset_per_facility` per facility
+        sampled_slices: list[polars.DataFrame] = []
+        for facility in facilities:
+            df = table.scan(
+                selected_fields=("bucket", "key", "size"),
+                row_filter=f"facility = '{facility}'",
+            ).to_polars()
+            if df.is_empty():
+                continue
+
+            sample_n = min(self.subset_per_facility, df.height)
+            sampled_slices.append(
+                df.sample(n=sample_n, with_replacement=False, shuffle=True)
+            )
+
+        # No samples case
+        if not sampled_slices:
+            return polars.DataFrame(
+                schema={"s3_uri": polars.String, "size": polars.Int64}
+            )
+
+        # Concat and adjust to s3_uri
+        return polars.concat(sampled_slices).select(
+            self._s3_uri_column(),
+            polars.col("size"),
+        )
+
+    def inventory(self) -> polars.DataFrame:
+        """
+        Return a subset from each facility
+        """
+
+        if self.subset_per_facility:
+            return self._subset_per_facility()
+        else:
+            return self._full_inventory()

src/data_index/orchestrate.py

@@ -82,13 +82,13 @@ def orchestrate(
     if metadata_factory is None:
         metadata_factory = DiskCachedUnstructuredMetadata
 
+    logger.info(f"Provisioning inventory: `{inventory_source}`")
+    inventory = inventory_source.inventory()
+
     logger.info(f"Provisioning sinks: `{structured_sink}`, `{unstructured_sink}`")
     structured_sink.provision()
     unstructured_sink.provision()
 
-    logger.info(f"Provisioning inventory: `{inventory_source}`")
-    inventory = inventory_source.inventory()
-
     logger.info(f"Batch workers: `{partitioner}, `{fetcher}`, `{extractor}`")
     logger.info(f"Dispatching ({len(inventory)} files total)")
     futures = [