def test_gpu_accounting_move_quota_from_one_vc_to_another(self): vc_info = { "A": { "P40": 20 }, "B": { "P40": 20 }, } # previous A has quota of 30, and A used them all, later admin moved # 10 to B vc_usage = { "A": { "P40": 30 }, "B": { "P40": 5 }, } cluster_total = {"P40": 40} cluster_available = {"P40": 5} cluster_unschedulable = {} result = quota.calculate_vc_gpu_counts(cluster_total, cluster_available, cluster_unschedulable, vc_info, vc_usage) vc_total, vc_used, vc_available, vc_unschedulable = result self.assertEqual(vc_info, vc_total) self.assertEqual(vc_usage, vc_used) target_vc_available = { "A": { "P40": 0 }, "B": { "P40": 5 }, } self.assertEqual(target_vc_available, vc_available) target_vc_unschedulable = { "A": { "P40": 0 }, "B": { "P40": 10 }, } self.assertEqual(target_vc_unschedulable, vc_unschedulable)
def test_gpu_accounting_real_case2(self): vc_info = { "quantus": { "P40": 150 }, "relevance2": { "P40": 234 }, "relevance2-inf": { "P40": 40 }, } vc_usage = { "quantus": { "P40": 125 }, "relevance2": { "P40": 231 }, "relevance2-inf": { "P40": 0 }, } cluster_total = {"P40": 424} cluster_available = {"P40": 68} cluster_unschedulable = {"P40": 1} result = quota.calculate_vc_gpu_counts(cluster_total, cluster_available, cluster_unschedulable, vc_info, vc_usage) vc_total, vc_used, vc_available, vc_unschedulable = result self.assertEqual(vc_info, vc_total) self.assertEqual(vc_usage, vc_used) target_vc_available = { "quantus": { "P40": 25 }, "relevance2": { "P40": 2 }, "relevance2-inf": { "P40": 40 }, } self.assertEqual(target_vc_available, vc_available) target_vc_unschedulable = { "quantus": { "P40": 0 }, "relevance2": { "P40": 1 }, "relevance2-inf": { "P40": 0 }, } self.assertEqual(target_vc_unschedulable, vc_unschedulable)
def test_gpu_accounting_idle_gpus_become_unscheduable(self): vc_info = { "A": { "P40": 40 }, "B": { "P40": 40 }, "C": { "P40": 40 }, } vc_usage = { "A": { "P40": 40 }, "B": { "P40": 31 }, "C": { "P40": 0 }, } cluster_total = {"P40": 120} cluster_available = {"P40": 29} cluster_unschedulable = {"P40": 20} result = quota.calculate_vc_gpu_counts(cluster_total, cluster_available, cluster_unschedulable, vc_info, vc_usage) vc_total, vc_used, vc_available, vc_unschedulable = result self.assertEqual(vc_info, vc_total) self.assertEqual(vc_usage, vc_used) target_vc_available = { "A": { "P40": 0 }, "B": { "P40": 1 }, "C": { "P40": 27 }, } self.assertEqual(target_vc_available, vc_available) target_vc_unschedulable = { "A": { "P40": 0 }, "B": { "P40": 8 }, "C": { "P40": 13 }, } self.assertEqual(target_vc_unschedulable, vc_unschedulable)
def test_gpu_accounting_real_case(self): vc_info = { "platform": { "P40": 48 }, "relevance": { "P40": 200 }, "quantus": { "P40": 100 }, "AU": { "P40": 20 }, } vc_usage = { "platform": { "P40": 57 }, "relevance": { "P40": 164 }, "quantus": { "P40": 93 }, "AU": { "P40": 0 }, } cluster_total = {"P40": 368} cluster_available = {"P40": 54} cluster_unschedulable = {} result = quota.calculate_vc_gpu_counts(cluster_total, cluster_available, cluster_unschedulable, vc_info, vc_usage) vc_total, vc_used, vc_available, vc_unschedulable = result self.assertEqual(vc_info, vc_total) self.assertEqual(vc_usage, vc_used) target_vc_available = { "platform": { "P40": 0 }, "relevance": { "P40": 30 }, "quantus": { "P40": 6 }, "AU": { "P40": 17 }, } self.assertEqual(target_vc_available, vc_available) target_vc_unschedulable = { "platform": { "P40": 0 }, "relevance": { "P40": 6 }, "quantus": { "P40": 1 }, "AU": { "P40": 3 }, } self.assertEqual(target_vc_unschedulable, vc_unschedulable)
def GetVC(userName, vcName): ret = None data_handler = DataHandler() cluster_status, _ = data_handler.GetClusterStatus() cluster_total = cluster_status["gpu_capacity"] cluster_available = cluster_status["gpu_avaliable"] cluster_reserved = cluster_status["gpu_reserved"] user_status = collections.defaultdict(lambda: ResourceInfo()) user_status_preemptable = collections.defaultdict(lambda: ResourceInfo()) vc_list = getClusterVCs() vc_info = {} vc_usage = collections.defaultdict( lambda: collections.defaultdict(lambda: 0)) vc_preemptable_usage = collections.defaultdict( lambda: collections.defaultdict(lambda: 0)) for vc in vc_list: vc_info[vc["vcName"]] = json.loads(vc["quota"]) active_job_list = data_handler.GetActiveJobList() for job in active_job_list: jobParam = json.loads(base64.b64decode(job["jobParams"])) if "gpuType" in jobParam: if not jobParam["preemptionAllowed"]: vc_usage[job["vcName"]][jobParam["gpuType"]] += GetJobTotalGpu( jobParam) else: vc_preemptable_usage[job["vcName"]][ jobParam["gpuType"]] += GetJobTotalGpu(jobParam) result = quota.calculate_vc_gpu_counts(cluster_total, cluster_available, cluster_reserved, vc_info, vc_usage) vc_total, vc_used, vc_available, vc_unschedulable = result for vc in vc_list: if vc["vcName"] == vcName and AuthorizationManager.HasAccess( userName, ResourceType.VC, vcName, Permission.User): num_active_jobs = 0 for job in active_job_list: if job["vcName"] == vcName and job["jobStatus"] == "running": num_active_jobs += 1 username = job["userName"] jobParam = json.loads(base64.b64decode(job["jobParams"])) if "gpuType" in jobParam: if not jobParam["preemptionAllowed"]: if username not in user_status: user_status[username] = ResourceInfo() user_status[username].Add( ResourceInfo({ jobParam["gpuType"]: GetJobTotalGpu(jobParam) })) else: if username not in user_status_preemptable: user_status_preemptable[ username] = ResourceInfo() user_status_preemptable[username].Add( ResourceInfo({ jobParam["gpuType"]: GetJobTotalGpu(jobParam) })) vc["gpu_capacity"] = vc_total[vcName] vc["gpu_used"] = vc_used[vcName] vc["gpu_preemptable_used"] = vc_preemptable_usage[vcName] vc["gpu_unschedulable"] = vc_unschedulable[vcName] vc["gpu_avaliable"] = vc_available[vcName] vc["AvaliableJobNum"] = num_active_jobs vc["node_status"] = cluster_status["node_status"] vc["user_status"] = [] for user_name, user_gpu in user_status.iteritems(): # TODO: job_manager.getAlias should be put in a util file user_name = user_name.split("@")[0].strip() vc["user_status"].append({ "userName": user_name, "userGPU": user_gpu.ToSerializable() }) vc["user_status_preemptable"] = [] for user_name, user_gpu in user_status_preemptable.iteritems(): user_name = user_name.split("@")[0].strip() vc["user_status_preemptable"].append({ "userName": user_name, "userGPU": user_gpu.ToSerializable() }) try: gpu_idle_url = config["gpu_reporter"] + '/gpu_idle' gpu_idle_params = {"vc": vcName} gpu_idle_response = requests.get(gpu_idle_url, params=gpu_idle_params) gpu_idle_json = gpu_idle_response.json() vc["gpu_idle"] = gpu_idle_json except Exception: logger.exception("Failed to fetch gpu_idle from gpu-exporter") ret = vc break return ret
def TakeJobActions(data_handler, redis_conn, launcher, jobs): vc_list = data_handler.ListVCs() cluster_status, _ = data_handler.GetClusterStatus() cluster_total = cluster_status["gpu_capacity"] cluster_available = cluster_status["gpu_avaliable"] cluster_reserved = cluster_status["gpu_reserved"] vc_info = {} vc_usage = collections.defaultdict( lambda: collections.defaultdict(lambda: 0)) for vc in vc_list: vc_info[vc["vcName"]] = json.loads(vc["quota"]) active_job_list = data_handler.GetActiveJobList() for job in active_job_list: jobParam = json.loads(base64.b64decode(job["jobParams"])) if "gpuType" in jobParam: vc_usage[job["vcName"]][jobParam["gpuType"]] += GetJobTotalGpu( jobParam) result = quota.calculate_vc_gpu_counts(cluster_total, cluster_available, cluster_reserved, vc_info, vc_usage) vc_total, vc_used, vc_available, vc_unschedulable = result cluster_gpu_capacity = cluster_status["gpu_capacity"] cluster_gpu_unschedulable = cluster_status["gpu_unschedulable"] global_total = ResourceInfo(cluster_gpu_capacity) global_unschedulable = ResourceInfo(cluster_gpu_unschedulable) vc_resources = {} globalResInfo = ResourceInfo.Difference(global_total, global_unschedulable) priority_dict = get_priority_dict() logging.info("Job priority dict: {}".format(priority_dict)) for vc in vc_list: vc_name = vc["vcName"] vc_schedulable = {} for gpu_type, total in vc_total[vc_name].items(): vc_schedulable[ gpu_type] = total - vc_unschedulable[vc_name][gpu_type] vc_resources[vc_name] = ResourceInfo(vc_schedulable) jobsInfo = [] for job in jobs: if job["jobStatus"] in ["queued", "scheduling", "running"]: singleJobInfo = {} singleJobInfo["job"] = job job_params = json.loads(base64.b64decode(job["jobParams"])) singleJobInfo["preemptionAllowed"] = job_params[ "preemptionAllowed"] singleJobInfo["jobId"] = job_params["jobId"] jobGpuType = "any" if "gpuType" in job_params: jobGpuType = job_params["gpuType"] singleJobInfo["globalResInfo"] = ResourceInfo( {jobGpuType: GetJobTotalGpu(job_params)}) # Job lists will be sorted based on and in the order of below # 1. non-preemptible precedes preemptible # 2. running precedes scheduling, precedes queued # 3. larger priority value precedes lower priority value # 4. early job time precedes later job time # Non-Preemptible jobs first preemptible = 1 if singleJobInfo["preemptionAllowed"] else 0 # Job status job_status = 0 if job["jobStatus"] == "scheduling": job_status = 1 elif job["jobStatus"] == "queued": job_status = 2 # Priority value reverse_priority = get_job_priority(priority_dict, singleJobInfo["jobId"]) priority = 999999 - reverse_priority # Job time job_time = str(job["jobTime"]) singleJobInfo["sortKey"] = "{}_{}_{:06d}_{}".format( preemptible, job_status, priority, job_time) singleJobInfo["allowed"] = False jobsInfo.append(singleJobInfo) jobsInfo.sort(key=lambda x: x["sortKey"]) logging.info("TakeJobActions : local resources : %s" % (vc_resources)) logging.info("TakeJobActions : global resources : %s" % (globalResInfo.CategoryToCountMap)) for sji in jobsInfo: logging.info("TakeJobActions : job : %s : %s : %s" % (sji["jobId"], sji["globalResInfo"].CategoryToCountMap, sji["sortKey"])) vc_name = sji["job"]["vcName"] vc_resource = vc_resources[vc_name] if (not sji["preemptionAllowed"]) and (vc_resource.CanSatisfy( sji["globalResInfo"])): vc_resource.Subtract(sji["globalResInfo"]) globalResInfo.Subtract(sji["globalResInfo"]) sji["allowed"] = True logging.info( "TakeJobActions : local assignment : %s : %s" % (sji["jobId"], sji["globalResInfo"].CategoryToCountMap)) for sji in jobsInfo: if sji["preemptionAllowed"] and (sji["allowed"] is False): if globalResInfo.CanSatisfy(sji["globalResInfo"]): logging.info( "TakeJobActions : job : %s : %s" % (sji["jobId"], sji["globalResInfo"].CategoryToCountMap)) # Strict FIFO policy not required for global (bonus) tokens since these jobs are anyway pre-emptible. globalResInfo.Subtract(sji["globalResInfo"]) sji["allowed"] = True logging.info( "TakeJobActions : global assignment : %s : %s" % (sji["jobId"], sji["globalResInfo"].CategoryToCountMap)) logging.info("TakeJobActions : global resources : %s" % (globalResInfo.CategoryToCountMap)) for sji in jobsInfo: try: if sji["job"]["jobStatus"] == "queued" and (sji["allowed"] is True): launcher.submit_job(sji["job"]) update_job_state_latency(redis_conn, sji["jobId"], "scheduling") logging.info("TakeJobActions : submitting job : %s : %s" % (sji["jobId"], sji["sortKey"])) elif sji["preemptionAllowed"] and ( sji["job"]["jobStatus"] == "scheduling" or sji["job"]["jobStatus"] == "running") and (sji["allowed"] is False): launcher.kill_job(sji["job"]["jobId"], "queued") logging.info("TakeJobActions : pre-empting job : %s : %s" % (sji["jobId"], sji["sortKey"])) except Exception as e: logging.error("Process job failed {}".format(sji["job"]), exc_info=True) logging.info("TakeJobActions : job desired actions taken")