bug fixes for site invalidation mode #994

bug fixes for site invalidation mode

Author: juanpablosalas

DMOps/file_invalidation_tool/src/container_invalidation.py

@@ -6,6 +6,7 @@
 from CMSSpark.spark_utils import get_spark_session
 from hadoop_queries import get_df_rse_locks, get_df_rse_replicas, get_df_contents
 from pyspark.sql.window import Window
+from file_invalidation import includes_rse_safe
 
 
 @click.command()
@@ -15,7 +16,6 @@
               help='RSE to look at')
 @click.option('--mode', required=False, type=click.Choice(['rucio','spark']), default='rucio', help='List generation mode')
 def invalidate_containers(filename,rse, mode):
-    #TODO: Check rse option
 
     if mode=='rucio':
         #Start Rucio Client
@@ -77,7 +77,7 @@ def invalidate_containers(filename,rse, mode):
                     df_rules = pd.concat([df_rules,df_rules_i])
 
         if rse is not None:
-            df_rules['includes_rse'] = df_rules['rse_expression'].apply(lambda exp: {'rse':rse} in list(rucio_client.list_rses(rse_expression=exp)))
+            df_rules['includes_rse'] = df_rules['rse_expression'].apply(lambda exp: includes_rse_safe(rucio_client,exp, rse))
             df_rules = df_rules.loc[df_rules.includes_rse]
 
         df_rules.columns = df_rules.columns.str.upper()

DMOps/file_invalidation_tool/src/dataset_invalidation.py

@@ -19,6 +19,7 @@
 from pyspark.sql.functions import col, collect_list, concat_ws
 from hadoop_queries import get_df_rse_locks, get_df_rse_replicas, get_df_contents, get_df_dataset_level_rules
 from pyspark.sql.window import Window
+from file_invalidation import includes_rse_safe
 
 @click.command()
 @click.option('--filename', required=True, default=None, type=str,
@@ -65,7 +66,7 @@ def invalidate_datasets(filename,rse, mode):
             df_rules = pd.DataFrame(columns=['rse','rule_id'])
 
         if rse is not None:
-            df_rules['includes_rse'] = df_rules['rse_expression'].apply(lambda exp: {'rse':rse} in list(rucio_client.list_rses(rse_expression=exp)))
+            df_rules['includes_rse'] = df_rules['rse_expression'].apply(lambda exp: includes_rse_safe(rucio_client,exp, rse))
             df_rules = df_rules.loc[df_rules.includes_rse]
 
         df_rules.columns = df_rules.columns.str.upper()

DMOps/file_invalidation_tool/src/file_invalidation.py

@@ -19,6 +19,23 @@
 from pyspark.sql.window import Window
 import pandas as pd
 from rucio.client import Client
+from rucio.common.exception import InvalidRSEExpression, DataIdentifierNotFound
+
+def includes_rse_safe(client,exp, rse):
+    try:
+        return {'rse': rse} in list(
+            client.list_rses(rse_expression=exp)
+        )
+    except InvalidRSEExpression as e:
+        return False
+    except Exception as e:
+        return False
+    
+def list_rules_safe(client,d):
+    try:
+        return client.list_associated_rules_for_file(scope=d["scope"], name=d["name"])
+    except DataIdentifierNotFound:
+        return []
 
 
 @click.command()
@@ -59,16 +76,18 @@ def invalidate_files(filename, rse, mode):
         df_delete.drop_duplicates().to_csv('/input/rucio_replicas_inv.csv',index=False)
 
         #Replicas to the rules
-        df_rules = pd.DataFrame(columns=['subscription_id', 'rse_expression', 'source_replica_expression', 'ignore_account_limit', 'created_at', 'account', 'copies', 'activity', 'priority', 'updated_at', 'scope', 'expires_at', 'grouping', 'name', 'weight', 'notification', 'comments', 'did_type', 'locked', 'stuck_at', 'child_rule_id', 'state', 'locks_ok_cnt', 'purge_replicas', 'eol_at', 'id', 'error', 'locks_replicating_cnt', 'ignore_availability', 'split_container', 'locks_stuck_cnt', 'meta', 'bytes'])
-        rules = [list(rucio_client.list_associated_rules_for_file(scope=d['scope'],name=d['name'])) for d in dict_delete]
-        for r in rules:
-            if df_rules.shape[0]==0 and len(r)>0:
-                df_rules = pd.DataFrame(r)
-            else:
-                df_rules = pd.concat([df_rules,pd.DataFrame(r)],axis=0)
+        rules = [rule for d in dict_delete for rule in list_rules_safe(rucio_client,d)]
+        
+        df_rules = pd.DataFrame(
+                    rules,
+                    columns=['subscription_id', 'rse_expression', 'source_replica_expression','ignore_account_limit', 'created_at', 'account', 'copies', 'activity',
+                        'priority', 'updated_at', 'scope', 'expires_at', 'grouping', 'name','weight', 'notification', 'comments', 'did_type', 'locked', 'stuck_at',
+                        'child_rule_id', 'state', 'locks_ok_cnt', 'purge_replicas', 'eol_at','id', 'error', 'locks_replicating_cnt', 'ignore_availability',
+                        'split_container', 'locks_stuck_cnt', 'meta', 'bytes']
+                    )
 
         if rse is not None:
-            df_rules['includes_rse'] = df_rules['rse_expression'].apply(lambda exp: {'rse':rse} in list(rucio_client.list_rses(rse_expression=exp)))
+            df_rules['includes_rse'] = df_rules['rse_expression'].apply(lambda exp: includes_rse_safe(rucio_client,exp, rse))
             df_rules = df_rules.loc[df_rules.includes_rse]
 
         #Rules protecting the replicas at File level