cancel_batch.py

import argparse
import sys
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor, as_completed

import scaleapi
from scaleapi.exceptions import ScaleException, ScaleUnauthorized

# Script that takes in an array of batch names (split by comma) and
# applies a bulk action to cancel all tasks in each batch.
# By default, this script makes 50 concurrent API calls.

# Example:
# python cancel_batch.py --api_key "SCALE_API_KEY"
# --batches "batch1,batch2" --clear "True"

# Change this for update concurrency
MAX_WORKERS = 50


def cancel_batch(client, batch_name, clear_unique_id):
 print(f"\nProcessing Batch: {batch_name}")
 try:
 batch = client.get_batch(batch_name)
 except ScaleException:
 print(f"-ERROR: Batch {batch_name} not found.")
 return

 task_count = client.get_tasks_count(
 project_name=batch.project, batch_name=batch.name
 )
 print(f"-Batch {batch.name} contains {task_count} tasks.")

 summary_metrics = defaultdict(lambda: 0)
 task_in_progress = 0
 processes = []

 tasks = client.get_tasks(project_name=batch.project, batch_name=batch.name)

 with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
 for task in tasks:
 task_in_progress += 1
 if task_in_progress % 1000 == 0:
 print(f"-Processing Task # {task_in_progress}")
 processes.append(
 executor.submit(
 cancel_task_with_response, client, task, clear_unique_id
 )
 )

 for process in as_completed(processes):
 result = process.result()
 summary_metrics[result["status"]] += 1

 for k, v in summary_metrics.items():
 print(f"--{k}: {v} tasks")


def cancel_task_with_response(client: scaleapi.ScaleClient, task, clear_unique_ud):
 task_status = task.as_dict()["status"]
 if task_status in ["completed", "canceled"]:
 return {"task": task.id, "status": task_status}

 try:
 task = client.cancel_task(task.id, clear_unique_ud)
 return {"task": task.id, "status": task.as_dict()["status"]}
 except ScaleException:
 return {"task": task.id, "status": "Can not cancel"}
 except Exception as err:
 print(err)
 return {"task": task.id, "status": "Errored"}


def get_args():
 ap = argparse.ArgumentParser()
 ap.add_argument("--api_key", required=True, help="Please provide Scale API Key")
 ap.add_argument(
 "--batches", required=True, help="Please enter batch names separated by a comma"
 )
 ap.add_argument(
 "--clear",
 type=bool,
 help="Set to True if you want to remove unique_id upon cancel",
 )
 return ap.parse_args()


def main():
 args = get_args()
 clear_unique_id = args.clear or False

 client = scaleapi.ScaleClient(args.api_key)

 # Testing API Key
 try:
 client.projects()
 except ScaleUnauthorized as err:
 print(err.message)
 sys.exit(1)

 batch_list = args.batches.split(",")
 batches = [word.strip() for word in batch_list]

 for batch_name in batches:
 cancel_batch(client, batch_name, clear_unique_id)


if __name__ == "__main__":
 main()