Merge pull request #959 from ericvaandering/revert_ext_md

Revert "CMSDM-315: Include additional metadata in the tape colocation plugin"

Author: ericvaandering

src/policy/CMSRucioPolicy/algorithms/tape_colocation.py

@@ -2,29 +2,24 @@
 Tape colocation algorithm for placement of CMS data on tape
 '''
 
-import json
-import logging
 from typing import Any, Optional, Union
-
 from rucio.common.exception import DataIdentifierNotFound
-from rucio.common.types import InternalScope
-from rucio.core.did import get_did, list_content, list_parent_dids
-from rucio.db.sqla.constants import DIDType
 from rucio.transfertool.fts3_plugins import FTS3TapeMetadataPlugin
+from rucio.core.did import list_parent_dids, get_did
+from rucio.db.sqla.constants import DIDType
+from rucio.common.types import InternalScope
+import logging 
+import json
 
 logger = logging.getLogger(__name__)
 
-class CMSTapeColocation(FTS3TapeMetadataPlugin):
-    policy_algorithm = "tape_metadata"
+class CMSTapeColocation(FTS3TapeMetadataPlugin): 
+    policy_algorithm = "tape_colocation"
 
-    # Logic for tape colocation
     allowed_types = ['data', 'hidata', 'mc', 'himc', 'relval', 'hirelval']
     parking_name = "parking"
     raw_name = "raw"
     hiraw_name = "hiraw"
-    
-    # Schema version as of June 26, 2025
-    schema_version = 1
 
     def __init__(self, policy: Union[str, None] = None) -> None:
 
@@ -37,8 +32,8 @@ def _module_init_(cls) -> None:
         logger.info("Registered plugin %s", cls.policy_algorithm)
         cls.register(
             cls.policy_algorithm, 
-            func=lambda x: cls.tape_metadata(x)
-        )
+            func=lambda x: cls.cms_colocation(x)
+        ) 
 
     @staticmethod
     def _encode(name: Any) -> Optional[str]: 
@@ -49,33 +44,33 @@ def _encode(name: Any) -> Optional[str]:
             return None
 
     @staticmethod
-    def parent_container(scope, name): 
+    def parent_container(name): 
         # Custom logic for CMS
         # If dataset - look for the parent container
         # If file - look for the parent dataset and then the parent container
-        scope = InternalScope(scope)
-        try:
+        scope = InternalScope("cms")
+        try: 
             is_file = get_did(scope=scope, name=name)['type'] == DIDType.FILE
         except DataIdentifierNotFound: 
             logger.warning("DID not found for %s:%s", scope, name)
             return None
-        try:
+        try: 
             if is_file:
                 parent_dataset = [parent
-                    for parent
-                    in list_parent_dids(scope=scope, name=name)
+                    for parent 
+                    in list_parent_dids(scope=scope, name=name) 
                 ][0]
                 containers = [
                     parent['name']
-                    for parent
-                    in list_parent_dids(scope=scope, name=parent_dataset['name'])
+                    for parent 
+                    in list_parent_dids(scope=scope, name=parent_dataset['name']) 
                     if parent['type'] == DIDType.CONTAINER
                 ]
             else: 
                 containers = [
                     parent['name']
-                    for parent
-                    in list_parent_dids(scope=scope, name=name)
+                    for parent 
+                    in list_parent_dids(scope=scope, name=name) 
                     if parent['type']==DIDType.CONTAINER
                 ]
             container = CMSTapeColocation._encode(containers[0])
@@ -87,161 +82,104 @@ def parent_container(scope, name):
             logger.debug("No parent container found for %s:%s", scope, name)
 
     @staticmethod
-    def _is_raw(name):
+    def _is_raw(name): 
         # Raw always contains "RAW" in the name
         return any(i=="RAW" for i in name.split('/'))
 
     @staticmethod
     def _is_parking(name):
         # Parking is denoted by having ParkingXXXX in the lfn
-        try:
+        try: 
             return any(n.startswith("Parking") for n in name.split('/'))
         except IndexError:
             return False
 
     @staticmethod
-    def data_type(name):
+    def data_type(name): 
         data_type = name.removeprefix('/store/').split('/')[0] # First index that isn't `store`
+        
         # Custom logic: Use parking or raw over "data", use hiraw if heavy ion and raw 
-        if data_type not in CMSTapeColocation.allowed_types:
+        if data_type not in CMSTapeColocation.allowed_types: 
             data_type = "n/a"
-        elif CMSTapeColocation._is_parking(name):
+        elif CMSTapeColocation._is_parking(name): 
             data_type = CMSTapeColocation.parking_name
         elif CMSTapeColocation._is_raw(name):
-            if data_type.startswith("hi"):
+            if data_type.startswith("hi"): 
                 data_type = CMSTapeColocation.hiraw_name
-            else:
+            else: 
                 data_type = CMSTapeColocation.raw_name
 
         return data_type
 
     @staticmethod
-    def data_tier(name):
-        try:
+    def data_tier(name): 
+        try: 
             tier = name.removeprefix('/store/').split('/')[3]
             if CMSTapeColocation._encode(tier) is None:
                 logger.debug("Could not encode data tier for %s", name)
             return tier
-        except IndexError:
+        except IndexError: 
             logger.debug("Could not determine data tier for %s", name)
 
     @staticmethod
-    def era(name):
-        try:
+    def era(name): 
+        try: 
             era = name.removeprefix('/store/').split('/')[1]
             if CMSTapeColocation._encode(era) is None:
                 logger.debug("Could not encode era for %s", name)
             return era
-        except IndexError:
+        except IndexError: 
             logger.debug("Could not determine era for %s", name)
 
-    @staticmethod
-    def _get_container_stats(scope, name):
-        size = 0
-        length = 0
-        scope  = InternalScope(scope)
-        try:
-            contents = list_content(scope, name)
-            for item in contents:
-                if item['type'] == DIDType.FILE:
-                    size += item.get('bytes', 0)
-                    length += 1
-                elif item['type'] in [DIDType.DATASET, DIDType.CONTAINER]:
-                    # Recursively get size of nested datasets/containers
-                    sub_length, sub_size = CMSTapeColocation._get_container_stats(item['name'])
-                    length += sub_length
-                    size += sub_size
-        except DataIdentifierNotFound:
-            logger.warning("DID not found for container %s:%s", scope, name)
-
-        return length, size
-
     @classmethod
-    def tape_metadata(cls, hints):
+    def cms_colocation(cls, hints):
         """
         https://github.com/dmwm/CMSRucio/issues/753
         https://github.com/dmwm/CMSRucio/issues/323
-        https://its.cern.ch/jira/browse/CMSDM-315
+
+        Level 0
+        Data/MC/HIData/HiMC (from /store/(data/mc/hi/data/himc) plus RAW and HIRAW, and Parking.
+        
+        Level 1
+        Data tier - either in the LFN or the end of the parent container
+
+        Level 2
+        Era (which for MC is the Campaign)
+
+        Level 3 
+        Parent Container (parent container of dataset if file)
+        
         
-        Tape Colocation: 
-            Level 0
-            Data/MC/HIData/HiMC (from /store/(data/mc/hi/data/himc) plus RAW and HIRAW, and Parking.
-            
-            Level 1
-            Data tier - either in the LFN or the end of the parent container
-
-            Level 2
-            Era (which for MC is the Campaign)
-
-            Level 3
-            Parent Container (parent container of dataset if file)
-            
-            
-            Examples:
-                * Parking data:
-                    /store/data/Run2024C/ParkingVBF5/RAW/v1/000/380/115/00000/b4c0513e-f732-42b1-858d-572c86ce4b97.root
-                    -->  {'0': 'parking', '1': 'RAW', '2': 'Run2024C', '3': '/ParkingVBF5/Run2024C-v1/RAW'}
-                * Raw:
-                    /store/hidata/HIRun2024B/HIPhysicsRawPrime5/RAW/v1/000/388/624/00000/fa0795b5-633b-461c-bc21-02d40a118dd2.root
-                    -->  {'0': 'hiraw', '1': 'RAW', '2': 'HIRun2024B', '3': '/HIPhysicsRawPrime5/HIRun2024B-v1/RAW'}
-
-        File Metadata: 
-            As given by Rucio hints:
-            * Size
-            * MD5
-            * Adler32
-
-        Additional Hints:
-            * Activity (default to "default" if not given)
-            * Level 3 length and size (if parent container exists)
-                - Length - total number of files in the parent container, including nested datasets/containers
-                - Size - total size of all files in the parent container, including nested datasets/containers
-                        
+        Examples: 
+            * Parking data: 
+                /store/data/Run2024C/ParkingVBF5/RAW/v1/000/380/115/00000/b4c0513e-f732-42b1-858d-572c86ce4b97.root
+                -->  {'0': 'parking', '1': 'RAW', '2': 'Run2024C', '3': '/ParkingVBF5/Run2024C-v1/RAW'}
+            * Raw: 
+                /store/hidata/HIRun2024B/HIPhysicsRawPrime5/RAW/v1/000/388/624/00000/fa0795b5-633b-461c-bc21-02d40a118dd2.root
+                -->  {'0': 'hiraw', '1': 'RAW', '2': 'HIRun2024B', '3': '/HIPhysicsRawPrime5/HIRun2024B-v1/RAW'}
+
         """
 
         lfn = hints['name']
         data_type = cls.data_type(lfn)
+        
         colocation = {
             "0": data_type,
         }
 
         if data_type != "n/a":
             tier = cls.data_tier(lfn)
             era = cls.era(lfn)
-            parent = cls.parent_container(hints['scope'], hints['name'])
-            if tier is not None:
+            parent = cls.parent_container(hints['name'])
+            if tier is not None: 
                 colocation['1'] = tier
-            if era is not None:
+            if era is not None: 
                 colocation['2'] = era
-            if parent is not None:
+            if parent is not None: 
                 colocation['3'] = parent
-        else:
-            parent = None
+        else: 
             logger.debug("Could not determine data type for %s", lfn)
 
+        # TODO Speak with FTS3 Team about these headers
         logger.debug("Setting colocation hints %s", colocation)
-
-        additional_hints = {
-            "activity": hints.get("activity", "default"),
-        }
-        if parent is not None:
-            length, size = CMSTapeColocation._get_container_stats(hints['scope'], parent)
-            additional_hints['3'] = {
-                "length": length, #The number of files in the parent container
-                "size": size,   #The total size of the parent container
-            }
-        logger.debug("Setting additional hints %s", additional_hints)
-
-        metadata = {
-            "size": hints['metadata'].get('filesize', 0),  # File size
-            "md5": hints['metadata'].get("md5"),  # MD5 checksum
-            "adler32": hints['metadata'].get("adler32"), # Adler32 checksum
-        }
-        logger.debug("File metadata: %s", metadata)
-
-        return {
-            "collocation_hints": colocation,
-            "additional_hints": additional_hints,
-            "file_metadata": metadata,
-            "schema_version": cls.schema_version
-        }
+        return {"collocation_hints": colocation}