Change how NMPI session quota is requested to get default applied

SpiNNakerManchester · Oct 23, 2024 · 4d527e7 · 4d527e7
1 parent 7ad9ea9
commit 4d527e7
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 13 deletions.
diff --git a/...aker-allocserv/src/main/java/uk/ac/manchester/spinnaker/alloc/allocator/QuotaManager.java b/...aker-allocserv/src/main/java/uk/ac/manchester/spinnaker/alloc/allocator/QuotaManager.java
@@ -53,6 +53,7 @@
 import uk.ac.manchester.spinnaker.alloc.nmpi.ResourceUsage;
 import uk.ac.manchester.spinnaker.alloc.nmpi.SessionRequest;
 import uk.ac.manchester.spinnaker.alloc.nmpi.SessionResourceUpdate;
+import uk.ac.manchester.spinnaker.alloc.nmpi.SessionResponse;
 
 /**
  * Manages user quotas.
@@ -371,7 +372,7 @@ private QuotaInfo parseQuotaData(List<Project> projects) {
 		return new QuotaInfo(total, units);
 	}
 
-	final Optional<String> mayCreateNMPISession(String collab) {
+	final Optional<String> checkQuota(String collab) {
 		// Read collab from NMPI; fail if not there
 		var projects = nmpi.getProjects(STATUS_ACCEPTED, collab);
 		var info = parseQuotaData(projects);
@@ -390,10 +391,13 @@ final Optional<String> mayCreateNMPISession(String collab) {
 		return Optional.empty();
 	}
 
-	void associateNMPISession(int jobId, String user, String collab,
-			String quotaUnits) {
-		var sessionId = nmpi.createSession(collab, user);
+	final SessionResponse createSession(String collab, String user) {
+		var session = nmpi.createSession(collab, user);
+		checkQuota(collab);
+		return session;
+	}
 
+	void associateNMPISession(int jobId, int sessionId, String quotaUnits) {
 		// Associate NMPI session with Job in the database
 		try (var c = getConnection();
 				var setSession = c.update(SET_JOB_SESSION)) {
@@ -485,7 +489,7 @@ final Optional<NMPIJobQuotaDetails> mayUseNMPIJob(String user,
 		}
 
 		// This is now a collab so check there instead
-		var quotaUnits = mayCreateNMPISession(job.getCollab());
+		var quotaUnits = checkQuota(job.getCollab());
 		if (quotaUnits.isEmpty()) {
 			return Optional.empty();
 		}
@@ -702,12 +706,12 @@ List<Project> getProjects(String status, String collab) {
 		return proxy.getProjects(apiKey, status, collab);
 	}
 
-	int createSession(String collab, String user) {
+	SessionResponse createSession(String collab, String user) {
 		var request = new SessionRequest();
 		request.setCollab(collab);
 		request.setHardwarePlatform(platform);
 		request.setUserId(user);
-		return proxy.createSession(apiKey, request).getId();
+		return proxy.createSession(apiKey, request);
 	}
 
 	void setSessionStatusAndResources(int sessionId, String status,

diff --git a/SpiNNaker-allocserv/src/main/java/uk/ac/manchester/spinnaker/alloc/allocator/Spalloc.java b/SpiNNaker-allocserv/src/main/java/uk/ac/manchester/spinnaker/alloc/allocator/Spalloc.java
@@ -530,10 +530,8 @@ public Optional<Job> createJobInCollabSession(String owner,
 			String nmpiCollab, CreateDescriptor descriptor,
 			String machineName, List<String> tags, Duration keepaliveInterval,
 			byte[] originalRequest) {
-		var quotaUnits = quotaManager.mayCreateNMPISession(nmpiCollab);
-		if (quotaUnits.isEmpty()) {
-			return Optional.empty();
-		}
+		var session = quotaManager.createSession(nmpiCollab, owner);
+		var quotaUnits = session.getResourceUsage().getUnits();
 
 		// Use the Collab name as the group, as it should exist
 		var job = execute(conn -> createJobInGroup(
@@ -545,8 +543,8 @@ public Optional<Job> createJobInCollabSession(String owner,
 			return job;
 		}
 
-		quotaManager.associateNMPISession(
-				job.get().getId(), owner, nmpiCollab, quotaUnits.get());
+		quotaManager.associateNMPISession(job.get().getId(), session.getId(),
+				quotaUnits);
 
 		// Return the job created
 		return job;

diff --git a/SpiNNaker-allocserv/src/main/java/uk/ac/manchester/spinnaker/alloc/nmpi/SessionResponse.java b/SpiNNaker-allocserv/src/main/java/uk/ac/manchester/spinnaker/alloc/nmpi/SessionResponse.java
@@ -25,6 +25,9 @@ public class SessionResponse {
 	/** The ID of the session. */
 	private Integer id;
 
+	/** A count of how much resource has been used by the job. */
+	private ResourceUsage resourceUsage;
+
 	/**
 	 * Get the ID of the session.
 	 *
@@ -44,6 +47,25 @@ public void setId(final Integer id) {
 		this.id = id;
 	}
 
+	/**
+	 * Get the count of how much resource has been used by the job.
+	 *
+	 * @return the resourceUsage
+	 */
+	public ResourceUsage getResourceUsage() {
+		return resourceUsage;
+	}
+
+	/**
+	 * Sets the resourceUsage.
+	 *
+	 * @param resourceUsage
+	 *            the resourceUsage to set
+	 */
+	public void setResourceUsage(final ResourceUsage resourceUsage) {
+		this.resourceUsage = resourceUsage;
+	}
+
 	/**
 	 * Used for JSON serialisation; ignores other properties we don't care
 	 * about.