<commit_before>package org.apache.hadoop.yarn.server.nodemanager;

import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.yarn.api.records.ContainerId;
import org.apache.hadoop.yarn.conf.YarnConfiguration;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;

public class CoresManagerImpl implements CoresManager {
	
	 private static final Log LOG = LogFactory
		      .getLog(CoresManagerImpl.class);
	
	private Set<Integer> totalCores = new HashSet<Integer>();
	
	private Set<Integer> unUsedCores = new HashSet<Integer>();
	
	//should be initialized at start
	private Map<Integer,Set<ContainerId>> coresToContainer = new HashMap<Integer,Set<ContainerId>>();
    
	private Map<ContainerId, Set<Integer>> containerToCores = new HashMap<ContainerId, Set<Integer>>();

	@Override
	public void init(Configuration conf) {
        //we get cores fisrt		
		int virtualCores =
		        conf.getInt(
		            YarnConfiguration.NM_VCORES, YarnConfiguration.DEFAULT_NM_VCORES);

         //we initialize the total cores and unused cores
		for(int i=0;i<virtualCores;i++){
			totalCores.add(i);
			unUsedCores.add(i);
			Set<ContainerId> cntIdSet = new HashSet<ContainerId>();
			coresToContainer.put(i, cntIdSet);
		}
	}

	
	private Set<Integer> getAvailableCores(int num) {
		Set<Integer> returnedResults = new HashSet<Integer>();
		int index = 0;
		assert(num <= totalCores.size());
		if(unUsedCores.size() > 0){
			for(Integer core : unUsedCores){
				returnedResults.add(core);
				index++;
				if(index >= num){
					break;
				}
			}
			
			for(Integer core : returnedResults){
				unUsedCores.remove(core);
			}
			
		}
		while(index < num){
	       Integer value = 0;
		   int     min = Integer.MAX_VALUE;
		for(Map.Entry<Integer, Set<ContainerId>> entry: coresToContainer.entrySet()){
			//find min core each time
		     if(returnedResults.contains(entry.getKey())){
		    	 continue;
		     }
		     
		     if(entry.getValue().size() < min){
		    	 value = entry.getKey();
		    	 min   = entry.getValue().size();
		     }
		  }
		returnedResults.add(value);
		}
		
		return returnedResults;
	}
	
	@Override
	public Set<Integer> allocateCores(ContainerId cntId, int num){
		LogOverlapWarning();
		Set<Integer> returnedResults = this.getAvailableCores(num);
		this.allcoateCoresforContainer(returnedResults, cntId);
		
		return returnedResults;
	}
	
	private void allcoateCoresforContainer(Set<Integer> cores,ContainerId cntId){
		for(Integer core : cores){
		    coresToContainer.get(core).add(cntId);
		}
		
		if(containerToCores.get(cntId) == null){
		  //first allocated
		  containerToCores.put(cntId, cores);	
		}else{
		  //newly allocated cores
		   containerToCores.get(cntId).addAll(cores);
		}
		LOG.info("allocate cpuset "+cores);
	}

	@Override
	public void releaseCores(ContainerId cntId) {
		LogOverlapWarning();
		Set<Integer> cores= new HashSet<Integer>();
		cores.addAll(containerToCores.get(cntId));
		this.releaseCoresforContainer(cntId, cores);
	}
	
	private void releaseCoresforContainer(ContainerId cntId, Set<Integer> cores){
		for(Integer core : cores){
			coresToContainer.get(core).remove(cntId);
			if(coresToContainer.get(core).size() == 0){
				unUsedCores.add(core);
			}
		}
		
		for(Integer core : cores){
			//remove core one by one
			containerToCores.get(cntId).remove(core);
		}
		
		//if there are no cores on this map, remove the entry
		if(containerToCores.get(cntId).size() == 0){
		    containerToCores.remove(cntId);
		}
		
		LOG.info("release cpuset "+cores);
	}
	
  @Override
  public Set<Integer> resetCores(ContainerId cntId, int num) {
		Set<Integer> cores = this.containerToCores.get(cntId);
		Set<Integer> returnedCores = new HashSet<Integer>();
		
	   //for a partially preempted container, its cores are null
	  if(cores != null){
	   if(num < cores.size()){<commit_msg>loop bug 1<commit_after>package org.apache.hadoop.yarn.server.nodemanager;

import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.yarn.api.records.ContainerId;
import org.apache.hadoop.yarn.conf.YarnConfiguration;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;

public class CoresManagerImpl implements CoresManager {
	
	 private static final Log LOG = LogFactory
		      .getLog(CoresManagerImpl.class);
	
	private Set<Integer> totalCores = new HashSet<Integer>();
	
	private Set<Integer> unUsedCores = new HashSet<Integer>();
	
	//should be initialized at start
	private Map<Integer,Set<ContainerId>> coresToContainer = new HashMap<Integer,Set<ContainerId>>();
    
	private Map<ContainerId, Set<Integer>> containerToCores = new HashMap<ContainerId, Set<Integer>>();

	@Override
	public void init(Configuration conf) {
        //we get cores fisrt		
		int virtualCores =
		        conf.getInt(
		            YarnConfiguration.NM_VCORES, YarnConfiguration.DEFAULT_NM_VCORES);

         //we initialize the total cores and unused cores
		for(int i=0;i<virtualCores;i++){
			totalCores.add(i);
			unUsedCores.add(i);
			Set<ContainerId> cntIdSet = new HashSet<ContainerId>();
			coresToContainer.put(i, cntIdSet);
		}
	}

	
	private Set<Integer> getAvailableCores(int num) {
		Set<Integer> returnedResults = new HashSet<Integer>();
		int index = 0;
		assert(num <= totalCores.size());
		
	    if(unUsedCores.size() > 0){
			for(Integer core : unUsedCores){
				returnedResults.add(core);
				index++;
				if(index >= num){
					break;
				}
			}
		}
			
		while(index < num){
	       Integer value = 0;
		   int     min = Integer.MAX_VALUE;
			   
		   for(Map.Entry<Integer, Set<ContainerId>> entry: coresToContainer.entrySet()){
			//find min core each time
		     if(returnedResults.contains(entry.getKey())){
		    	 continue;
		     }
		     
		     if(entry.getValue().size() < min){
		    	 value = entry.getKey();
		    	 min   = entry.getValue().size();
		     }
		  }
		returnedResults.add(value);
		index++;
		}
		
		return returnedResults;
	}
	
	@Override
	public Set<Integer> allocateCores(ContainerId cntId, int num){
		LogOverlapWarning();
		Set<Integer> returnedResults = this.getAvailableCores(num);
		this.allcoateCoresforContainer(returnedResults, cntId);
		
		return returnedResults;
	}
	
	private synchronized void allcoateCoresforContainer(Set<Integer> cores,ContainerId cntId){
		
		for(Integer core : cores){
			 unUsedCores.remove(core);
		 }
		 
		for(Integer core : cores){
		     coresToContainer.get(core).add(cntId);
		}

		if(containerToCores.get(cntId) == null){
		  //first allocated
		  containerToCores.put(cntId, cores);	
		}else{
		  //newly allocated cores
		   containerToCores.get(cntId).addAll(cores);
		}
		LOG.info("allocate cpuset "+cores);
	}

	@Override
	public void releaseCores(ContainerId cntId) {
		LogOverlapWarning();
		Set<Integer> cores= new HashSet<Integer>();
		cores.addAll(containerToCores.get(cntId));
		this.releaseCoresforContainer(cntId, cores);
	}
	
	private synchronized void releaseCoresforContainer(ContainerId cntId, Set<Integer> cores){
		for(Integer core : cores){
			if(coresToContainer.get(core).size() == 0){
				unUsedCores.add(core);
			}
		}
		
		for(Integer core : cores){
			coresToContainer.get(core).remove(cntId);
		}
		for(Integer core : cores){
			//remove core one by one
			containerToCores.get(cntId).remove(core);
		}
		
		//if there are no cores on this map, remove the entry
		if(containerToCores.get(cntId).size() == 0){
		    containerToCores.remove(cntId);
		}
		
		LOG.info("release cpuset "+cores);
	}
	
  @Override
  public Set<Integer> resetCores(ContainerId cntId, int num) {
		Set<Integer> cores = this.containerToCores.get(cntId);
		Set<Integer> returnedCores = new HashSet<Integer>();
		
	   //for a partially preempted container, its cores are null
	  if(cores != null){
	   if(num < cores.size()){
{'fraction_non_alphanumeric': 0.124194348041646, 'fraction_numerical': 0.0016113039167079822, 'mean_word_length': 3.7576650943396226, 'pattern_counts': {'":': 0, '<': 71, '<?xml version=': 0, '>': 65, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5050069', 'n_tokens_mistral': 2983, 'n_tokens_neox': 2573, 'n_words': 537}