manage compression many indexes at a time

Author: sagehen03

bioindex/lib/aws.py

@@ -22,11 +22,21 @@
 
 
 def get_bgzip_job_status(job_id: str):
-    batch_client = boto3.client('batch')
-    job_response = batch_client.describe_jobs(jobs=[job_id])
-    if len(job_response['jobs']) > 0:
-        return job_response['jobs'][0]['status']
-    return None
+    """
+    Get the status of a batch job. Returns the status string or None if job not found.
+    Valid statuses: SUBMITTED, PENDING, RUNNABLE, STARTING, RUNNING, SUCCEEDED, FAILED
+    """
+    try:
+        batch_client = boto3.client('batch')
+        job_response = batch_client.describe_jobs(jobs=[job_id])
+        if len(job_response['jobs']) > 0:
+            return job_response['jobs'][0]['status']
+        else:
+            print(f"Warning: Job {job_id} not found in AWS Batch")
+            return None
+    except Exception as e:
+        print(f"Error getting status for job {job_id}: {str(e)}")
+        return None
 
 
 def start_batch_job(s3_bucket: str, index_name: str, s3_path: str, job_definition: str, additional_parameters: dict = None):

bioindex/lib/index.py

@@ -268,9 +268,6 @@ def batch_run_function(self, config, obj):
 
         return start_and_wait_for_indexer_job(obj['Key'], self.name, self.schema.arity, config.s3_bucket, config.rds_secret,
                                               config.bio_schema, obj['Size'])
-        """
-        Index the objects using a lambda function.
-        """
 
     def lambda_run_function(self, config, obj):
         logging.info(f'Processing {relative_key(obj["Key"], self.s3_prefix)}...')

bioindex/main.py

@@ -182,8 +182,26 @@ def validate_job_type(ctx, param, value):
 
 
 def convert_cli_arg_to_list(value):
-    if value:
-        return [item.strip() for item in value]
+    if not value:
+        return None
+
+    # For multiple=True options, Click returns a tuple of values
+    if isinstance(value, tuple):
+        result = []
+        for v in value:
+            # Each value could be comma-separated
+            if v and ',' in v:
+                result.extend([item.strip() for item in v.split(',')])
+            elif v:
+                result.append(v.strip())
+        return result
+    # For single string values
+    elif isinstance(value, str) and ',' in value:
+        return [item.strip() for item in value.split(',')]
+    # For single value
+    elif value:
+        return [value.strip()]
+
     return None
 
 
@@ -194,25 +212,75 @@ def convert_cli_arg_to_list(value):
               help="List of index names to exclude from processing. E.g. --exclude=name1,name2")
 @click.option('--job-type', type=str, required=True, callback=validate_job_type,
               help="Type of job to perform. Must be one of COMPRESS, DECOMPRESS, DELETE_JSON")
+@click.option('--wait-for-completion/--no-wait', is_flag=True, default=True,
+              help="Wait for all AWS batch jobs to complete before exiting")
 @click.pass_obj
-def cli_bulk_compression_management(cfg, include, exclude, job_type):
+def cli_bulk_compression_management(cfg, include, exclude, job_type, wait_for_completion):
     engine = migrate.migrate(cfg)
     indexes = index.Index.list_indexes(engine, False)
 
+    # Debug information
+    console.print(f"Raw include: {include}")
+    console.print(f"Raw exclude: {exclude}")
+
     inclusion_list = convert_cli_arg_to_list(include)
     exclusion_list = convert_cli_arg_to_list(exclude)
 
+    console.print(f"Processed inclusion list: {inclusion_list}")
+    console.print(f"Processed exclusion list: {exclusion_list}")
+    console.print(f"Wait for completion: {wait_for_completion}")
+
+    job_ids = []
+    console.print(f"Starting {job_type.value} jobs for selected indexes...")
+
     with concurrent.futures.ThreadPoolExecutor(max_workers=3) as executor:
         futures = []
         for i in indexes:
             if (inclusion_list is None or i.name in inclusion_list) and (
                 exclusion_list is None or i.name not in exclusion_list):
                 futures.append(
-                    executor.submit(check_index_and_launch_job, cfg, i.name, i.s3_prefix, job_type))
+                    executor.submit(check_index_and_launch_job, cfg, i.name, i.s3_prefix, job_type, return_job_id=True))
 
-        # wait for all futures to complete
+        # collect all job IDs
         for future in concurrent.futures.as_completed(futures):
-            future.result()
+            try:
+                job_id = future.result()
+                if job_id:
+                    job_ids.append(job_id)
+                    console.print(f"Successfully queued job {job_id}")
+            except Exception as e:
+                console.print(f"Error launching job: {str(e)}")
+
+    if job_ids and wait_for_completion:
+        console.print(f"Launched {len(job_ids)} jobs. Waiting for completion...")
+        console.print(f"Job IDs to monitor: {job_ids}")
+
+        initial_wait = 90
+        check_interval = 30
+        console.print(f"Waiting initial {initial_wait} seconds before checking status...")
+        time.sleep(initial_wait)
+
+        # Check all jobs until they're complete
+        iteration = 1
+        while job_ids:
+            console.print(f"Status check iteration {iteration}, {len(job_ids)} jobs remaining")
+            for job_id in list(job_ids):  # Create a copy to safely modify during iteration
+                try:
+                    status = aws.get_bgzip_job_status(job_id)
+                    console.print(f'{job_type} job, id: {job_id} status: {status}')
+                    if status in ['SUCCEEDED', 'FAILED']:
+                        console.print(f'{job_type} job {job_id} finished with status: {status}')
+                        job_ids.remove(job_id)
+                except Exception as e:
+                    console.print(f"Error checking status for job {job_id}: {str(e)}")
+
+            if job_ids:  # Only sleep if there are still jobs running
+                console.print(f"Waiting {check_interval} seconds before next status check...")
+                time.sleep(check_interval)
+
+            iteration += 1
+
+        console.print("All jobs completed.")
 
 
 @click.command(name='decompress')
@@ -239,29 +307,82 @@ def is_index_prefix_valid(cfg, idx: str, prefix: str):
     return len(selected_index) == 1
 
 
-def check_index_and_launch_job(cfg, index_name, prefix, job_type, additional_parameters=None):
-    if is_index_prefix_valid(cfg, index_name, prefix):
-        start_and_monitor_aws_batch_job(cfg.s3_bucket, job_type, index_name, prefix, additional_parameters=additional_parameters)
-    else:
-        console.print(f'Could not find unique index with name {index_name} and prefix {prefix}, quitting')
+def check_index_and_launch_job(cfg, index_name, prefix, job_type, additional_parameters=None, return_job_id=False):
+    try:
+        if is_index_prefix_valid(cfg, index_name, prefix):
+            try:
+                job_id = aws.start_batch_job(cfg.s3_bucket, index_name, prefix, job_type.value, additional_parameters=additional_parameters)
+                console.print(f'{job_type} started with id {job_id}')
+
+                if return_job_id:
+                    return job_id
+                else:
+                    # For backward compatibility with existing command usage
+                    start_and_monitor_aws_batch_job(cfg.s3_bucket, job_type, index_name, prefix,
+                                                  additional_parameters=additional_parameters, job_id=job_id)
+                    return None
+            except Exception as e:
+                console.print(f'Error starting job for index {index_name} with prefix {prefix}: {str(e)}')
+                return None
+        else:
+            console.print(f'Could not find unique index with name {index_name} and prefix {prefix}, skipping')
+            return None
+    except Exception as e:
+        console.print(f'Unexpected error for index {index_name}: {str(e)}')
+        return None
 
 
 def start_and_monitor_aws_batch_job(s3_bucket: str, job_type: BgzipJobType, index_name: str, s3_path: str,
                                     initial_wait: int = 90, check_interval: int = 30,
-                                    additional_parameters: dict = None):
-    job_id = aws.start_batch_job(s3_bucket, index_name, s3_path, job_type.value, additional_parameters=additional_parameters)
-    console.print(f'{job_type} started with id {job_id}')
-    time.sleep(initial_wait)
-    while True:
-        status = aws.get_bgzip_job_status(job_id)
-        console.print(f'{job_type} job, id: {job_id} status: {status}')
-        if status == 'SUCCEEDED':
-            console.print(f'{job_type} job {job_id} succeeded')
-            break
-        elif status == 'FAILED':
-            console.print(f'{job_type} job {job_id} failed')
-            break
-        time.sleep(check_interval)
+                                    additional_parameters: dict = None, job_id: str = None):
+    try:
+        # Start a job if job_id is not provided
+        if job_id is None:
+            job_id = aws.start_batch_job(s3_bucket, index_name, s3_path, job_type.value, additional_parameters=additional_parameters)
+            console.print(f'{job_type} started with id {job_id}')
+            console.print(f'Waiting {initial_wait} seconds before checking status...')
+            time.sleep(initial_wait)
+
+        # Monitor the job until completion or max retries
+        max_retries = 5
+        retry_count = 0
+
+        while True:
+            try:
+                status = aws.get_bgzip_job_status(job_id)
+
+                if status is None:
+                    retry_count += 1
+                    if retry_count >= max_retries:
+                        console.print(f'Error: Could not get status for {job_type} job {job_id} after {max_retries} attempts, giving up')
+                        break
+                    console.print(f'Warning: Could not get status for {job_type} job {job_id}, retrying ({retry_count}/{max_retries})...')
+                    time.sleep(check_interval)
+                    continue
+
+                retry_count = 0  # Reset retry count on successful status check
+                console.print(f'{job_type} job, id: {job_id} status: {status}')
+
+                if status == 'SUCCEEDED':
+                    console.print(f'{job_type} job {job_id} succeeded')
+                    break
+                elif status == 'FAILED':
+                    console.print(f'{job_type} job {job_id} failed')
+                    break
+
+                console.print(f'Waiting {check_interval} seconds before next status check...')
+                time.sleep(check_interval)
+
+            except Exception as e:
+                console.print(f'Error checking job status: {str(e)}')
+                retry_count += 1
+                if retry_count >= max_retries:
+                    console.print(f'Error: Failed to check job status after {max_retries} attempts, giving up')
+                    break
+                time.sleep(check_interval)
+
+    except Exception as e:
+        console.print(f'Error in start_and_monitor_aws_batch_job: {str(e)}')
 
 
 @click.command(name='update-compressed-status')