CTORM load tester PR-7456

.github/workflows/lint.yml

@@ -14,9 +14,9 @@ jobs:
           python-version: 3.11
       - uses: TrueBrain/actions-flake8@v2
         with:
-          flake8_version: 6.0.0
+          flake8_version: 7.3.0
           max_line_length: 120
-          plugins: flake8-isort==6.0.0 flake8-quotes==3.3.2
+          plugins: flake8-isort==7.0.0, flake8-quotes==3.4.0
 
   ruff-format:
     runs-on: ubuntu-latest
@@ -25,5 +25,5 @@ jobs:
       - uses: actions/checkout@v4
       - uses: astral-sh/ruff-action@v3
         with:
-          version: "~=0.13.3"
+          version: "~=0.15.13"
           args: format --check --diff --output-format=github

.github/workflows/test.yml

@@ -8,17 +8,23 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ['3.9', '3.10', '3.11', '3.12']
+        python-version: [ '3.9', '3.10', '3.11', '3.12' ]
       fail-fast: false
 
     steps:
-    - uses: actions/checkout@v4
-    - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v5
-      with:
-        python-version: ${{ matrix.python-version }}
-    - name: Install dependencies
-      run: |
-        python -m pip install --upgrade pip
-        python -m pip install poetry "tox<4.0.0" tox-poetry tox-gh-actions
-    - run: make tests
+      - uses: actions/checkout@v4
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          python -m pip install poetry "tox<4.0.0" tox-poetry tox-gh-actions
+      - name: Run tests
+        run: |
+          make tests
+          if [ "${{ matrix.python-version }}" = "3.12" ]; then
+            # ctorm supports py312 only. 
+            make ctorm-test
+          fi

.gitignore

@@ -5,5 +5,7 @@
 poetry.lock
 .tox
 
+**.tfstate.backup
+
 # Auto generated by Make
 tox.ini

Makefile

@@ -13,3 +13,7 @@ TOX_INIS := $(PACKAGES:%=%/tox.ini)
 .PHONY: tests
 tests: $(TOX_INIS)
 	ss=0; $(foreach package,$(PACKAGES),cd $(package) && tox || ss=1;cd ..;) exit $$ss
+
+.PHONY: ctorm-test
+ctorm-test:
+	cd ctorm && tox -c pyproject.toml

ctorm/.python-version

@@ -0,0 +1 @@
+3.12

ctorm/README.md

@@ -0,0 +1,15 @@
+# CTORM
+
+Cumulus Throughput... uh... ORM?
+
+Load tester tool for Cumulus.
+
+Operates as x steps:
+
+* cumulus-db-md-extract
+    * Operates in the prod account.
+    * Gets the metadata from the cumulus database and writes it to a bunch of .jsonl files in a S3 bucket in the
+      operational account.
+* granule-md-db-loader
+    * Operates in the account in which the load test will operate.
+    * Reads the .jsonl files from the S3 bucket and loads them into the CTORM granules dynamo db.

ctorm/ctorm.cfg.example

@@ -0,0 +1,35 @@
+[ctorm]
+
+calc_md5 = true
+granule_goal = 200
+granules_sqs_queue_url = "https://sqs.us-west-2.amazonaws.com/123456789012/ctorm-dev-granules"
+
+[[ ctorm.pipelines ]]
+bucketname = "nisar-bucket-0"
+ingest_keypair_name = ""
+ingest_queue_name = "sds-n-cumulus-test-nisar-workflow-queue"
+ingest_sqs_queue_url = ""
+prepare_keypair_name = "NISAR_P"
+provider = "FOO"
+share = 50
+ummg_prefix = "UMMG/FOO_COLLECTION/"
+
+[[ctorm.pipelines]]
+bucketname = "nisar-bucket-2"
+ingest_keypair_name = ""
+ingest_queue_name = "sds-n-cumulus-test-nisar-workflow-queue"
+ingest_sqs_queue_url = ""
+prepare_keypair_name = "NISAR"
+provider = "FOO"
+share = 45
+ummg_prefix = "UMMG/BAR_COLLECTION/"
+
+[[ctorm.pipelines]]
+bucketname = "opera-bucket-0"
+ingest_keypair_name = ""
+ingest_queue_name = "asf-cumulus-test-opera-workflow-queue"
+ingest_sqs_queue_url = ""
+prepare_keypair_name = "OPERA"
+provider = "FOO"
+share = 5
+ummg_prefix = "UMMG/BAZ_COLLECTION/"

ctorm/ctorm/__main__.py

@@ -0,0 +1,3 @@
+from ctorm.main import main
+
+main()

ctorm/ctorm/cnm.py

@@ -0,0 +1,131 @@
+import re
+import uuid
+from collections.abc import Callable
+from datetime import datetime, timezone
+from logging import getLogger
+from pathlib import Path
+
+from ctorm.config import (
+    MD5_CHECKSUM_PATTERN,
+    CtormPipeline,
+    CtormPreparedFile,
+    CtormPreparedGranule,
+)
+
+log = getLogger(__name__)
+
+DATA_TYPE_MAP = {
+    ".context.json": "metadata",
+    ".dataset.json": "metadata",
+    ".h5": "data",
+    ".iso.xml": "metadata",
+    ".jpg": "browse",
+    ".log": "metadata",
+    ".md5": "metadata",
+    ".met.json": "metadata",
+    ".nc": "data",
+    ".pdf": "qa",
+    ".png": "browse",
+    ".qa.h5": "qa",
+    ".rc.yaml": "metadata",
+    ".tif": "data",
+    ".xml": "data",
+}
+
+type CtormCnmSGeneratorType = Callable[[CtormPreparedGranule, str], dict]  # noqa
+
+
+class CtormCnmSGenerator:
+    """input dict like:
+    ```
+    {
+      "bm": {
+        "sds-n-cumulus-prod-nisar-products": "B1",
+        "sds-n-cumulus-prod-nisar-jpl-private-data": "B2"
+      },
+      "g": "NISAR_L0_RRST_VC08_20250821T101036_20250821T101041_P00408_J_001",
+      "c": "NISAR_L0A_RRST_BETA_V1",
+      "cv": "1",
+      "f": [
+        {
+          "f": "s3://$B1/NISAR_L0A_RRST_BETA_V1/$G/$G.bin",
+          "m": "e989430f4c5bfa04eeb26e56f4cd8375",
+          "s": 1073829760
+        },
+        {
+          "f": "s3://$B1/NISAR_L0A_RRST_BETA_V1/$G/$G.rc.yaml",
+          "s": 120597,
+          "m": "11bcaec780f0d246f3996b3f08680410"
+        },
+        {
+          "f": "s3://$B1/NISAR_L0A_RRST_BETA_V1/$G/$G.bin.qa",
+          "s": 1134,
+        },
+      ]
+    },
+    ```
+    """
+
+    def _assemble_file_dict(
+        self, file_dict: CtormPreparedFile, g_name: str, bucket_map: dict
+    ):
+        file_s3uri = file_dict["f"].replace("$G", g_name)
+        for bucket, replacetoken in bucket_map.items():
+            file_s3uri = file_s3uri.replace(f"${replacetoken}", bucket)
+        filename = Path(file_s3uri).name
+        return {
+            "name": filename,
+            "type": self._get_type(filename),
+            "uri": file_s3uri,
+            "size": file_dict["s"],
+            "checksum": self._get_checksum(file_dict),
+            "checksumType": "md5",
+        }
+
+    def __call__(
+        self,
+        granule_dict: CtormPreparedGranule,
+        provider: str,
+    ) -> dict:
+        cnm_s = {
+            "identifier": str(uuid.uuid4()),
+            "collection": granule_dict["c"],
+            "version": "1.3",
+            "submissionTime": datetime.now(tz=timezone.utc).strftime(
+                "%Y-%m-%dT%H:%M:%S.%fZ"
+            ),
+            "product": {
+                "name": granule_dict["g"],
+                "dataVersion": granule_dict["cv"],
+                "files": [
+                    self._assemble_file_dict(
+                        file, granule_dict["g"], granule_dict["bm"]
+                    )
+                    for file in granule_dict["f"]
+                ],
+            },
+            "provider": provider,
+        }
+
+        return cnm_s
+
+    def _get_type(self, filename: str):
+
+        suffixes = Path(filename).suffixes
+        while suffixes:
+            data_type = DATA_TYPE_MAP.get("".join(suffixes))
+
+            if data_type:
+                return data_type
+
+            suffixes = suffixes[1:]
+
+        log.debug("suffix from %s is not in map. Using 'data'", filename)
+        return "data"
+
+    def _get_checksum(self, file_dict: CtormPreparedFile):
+        m = MD5_CHECKSUM_PATTERN.match(file_dict.get("m", "not md5"))
+        if m:
+            return m.group(1)
+
+        return "00000000000000000000000000000000"

ctorm/ctorm/cnm_sender.py

@@ -0,0 +1,29 @@
+from logging import getLogger
+from typing import List
+
+from ctorm.cnm import CtormCnmSGenerator
+from ctorm.config import CtormConfig, CtormPreparedGranule
+
+log = getLogger(__name__)
+
+
+class CnmSender:
+    def __init__(
+        self,
+        cfg: CtormConfig,
+        granules: List[CtormPreparedGranule],
+        provider: str = "TODO: PROVIDER",
+    ):
+        self.cfg = cfg
+        self.granules = granules
+        self.provider = provider
+        self.cnm_s_generator = CtormCnmSGenerator()
+
+    def send_all(self) -> bool:
+        for granule in self.granules:
+            cnms = self.cnm_s_generator(granule, self.provider)
+            log.debug("Sending %s", cnms)
+            # TODO: actually send
+
+        # TODO: determine successfulllness
+        return True

ctorm/ctorm/commands/__init__.py

@@ -0,0 +1,27 @@
+import argparse
+import logging
+
+from ctorm.config import CtormConfig
+from ctorm.prepare import CtormPrepare
+
+log = logging.getLogger(__name__)
+
+
+def cmd_prepare(args: argparse.Namespace):
+    cfg = CtormConfig.from_file(
+        cfg_file=args.cfg_file,
+    )
+    prepare = CtormPrepare(cfg)
+    prepare.prepare()
+
+
+def cmd_cnm_sender(args: argparse.Namespace):
+    # TODO: something like:
+    #
+    # from ctorm.cnm_sender import CnmSender
+    #
+    # cfg = CtormConfig.from_file(cfg_file=args.cfg_file)
+    # cnm_sender = CnmSender(cfg)
+    # cnm_sender.run(max_messages=args.max_messages)
+    #
+    raise NotImplementedError("cnm_sender command is not implemented yet")

ctorm/ctorm/config.py

@@ -0,0 +1,80 @@
+import re
+import tomllib
+from dataclasses import dataclass, fields
+from typing import NotRequired, TypedDict
+
+from boto3.session import Session
+from mypy_boto3_s3 import S3Client
+
+try:
+    from typing import Self
+except ImportError:
+    Self = "CtormConfig"
+
+DEFAULT = "default"
+AWS_REGION = "us-west-2"  # We will never not want us-west-2
+MD5_CHECKSUM_PATTERN = re.compile(r"^([\da-f]{32})$")
+
+
+class CtormPreparedFile(TypedDict):
+    f: str
+    m: NotRequired[str]
+    s: int
+
+
+class CtormPreparedGranule(TypedDict):
+    bm: dict[str, str]
+    g: str
+    c: str
+    cv: str
+    f: list[CtormPreparedFile]
+
+
+@dataclass
+class CtormPipeline:
+    bucketname: str
+
+    # Looks for a keypair in the form of `AWS_ACCESS_KEY_ID_[keypair_name]`
+    # This should have permissions to send SQS to the maturity we're sending CNM to
+    ingest_keypair_name: str
+    ingest_queue_name: str
+    ingest_sqs_queue_url: str
+
+    # Looks for a keypair in the form of `AWS_ACCESS_KEY_ID_[keypair_name]`
+    # This should have permissions to read from prod bucket where data resides
+    prepare_keypair_name: str
+    share: int
+
+    # These are used at runtime
+    next_cont_token: str = None
+    s3_client: S3Client = None
+    session: Session = None
+    ummg_prefix: str = "UMMG/"
+    provider: str = "TODO"
+
+
+@dataclass
+class CtormConfig:
+    pipelines: list
+    granules_sqs_queue_url: str
+
+    calc_md5: bool = True
+
+    # This is the number of granules that will be prepared for test.
+    granule_goal: int = 2000000
+
+    @classmethod
+    def from_file(
+        cls,
+        cfg_file: str,
+    ) -> Self:
+        with open(cfg_file, "rb") as f:
+            cfg = tomllib.load(f)
+        valid_fields = {f.name for f in fields(cls)}
+        # Filter the input dictionary
+        kwargs = {k: v for k, v in cfg["ctorm"].items() if k in valid_fields}
+        # create CtormPipeline for each bucket in cfg['pipelines']
+        kwargs["pipelines"] = [CtormPipeline(**b) for b in kwargs["pipelines"]]
+        obj = cls(**kwargs)
+
+        return obj