6233274834

Committed 19 Sep 2023 08:46AM UTC coverage: 36.409% (-0.6%) from 36.982%

Build # 6233274834

Build Type

Pull #658

github

Committed by

dkfellows

Commit Message

Merge branch 'master' into java-17

Pull Request Pull Request #658: Update Java version to 17

Run Details

1656 of 1656 new or added lines in 260 files covered. (100.0%)

8373 of 22997 relevant lines covered (36.41%)

0.36 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

28.03

/SpiNNaker-allocserv/src/main/java/uk/ac/manchester/spinnaker/alloc/allocator/QuotaManager.java

/*
 * Copyright (c) 2021 The University of Manchester
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package uk.ac.manchester.spinnaker.alloc.allocator;

import static java.util.Objects.isNull;
import static org.slf4j.LoggerFactory.getLogger;
import static uk.ac.manchester.spinnaker.alloc.db.Row.integer;
import static uk.ac.manchester.spinnaker.alloc.db.Utils.isBusy;
import static uk.ac.manchester.spinnaker.alloc.nmpi.ResourceUsage.BOARD_SECONDS;
import static uk.ac.manchester.spinnaker.alloc.nmpi.ResourceUsage.CORE_HOURS;
import static uk.ac.manchester.spinnaker.alloc.model.GroupRecord.GroupType.COLLABRATORY;
import static uk.ac.manchester.spinnaker.alloc.security.TrustLevel.USER;

import java.util.List;
import java.util.Optional;

import javax.annotation.PostConstruct;
import javax.ws.rs.BadRequestException;

import org.slf4j.Logger;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.dao.DataAccessException;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Service;

import com.google.errorprone.annotations.RestrictedApi;

import uk.ac.manchester.spinnaker.alloc.ForTestingOnly;
import uk.ac.manchester.spinnaker.alloc.ServiceMasterControl;
import uk.ac.manchester.spinnaker.alloc.SpallocProperties.QuotaProperties;
import uk.ac.manchester.spinnaker.alloc.db.DatabaseAPI.Connection;
import uk.ac.manchester.spinnaker.alloc.db.DatabaseAPI.Query;
import uk.ac.manchester.spinnaker.alloc.db.DatabaseAPI.Update;
import uk.ac.manchester.spinnaker.alloc.db.DatabaseAwareBean;
import uk.ac.manchester.spinnaker.alloc.db.Row;
import uk.ac.manchester.spinnaker.alloc.nmpi.Job;
import uk.ac.manchester.spinnaker.alloc.nmpi.JobResourceUpdate;
import uk.ac.manchester.spinnaker.alloc.nmpi.NMPIv3API;
import uk.ac.manchester.spinnaker.alloc.nmpi.Project;
import uk.ac.manchester.spinnaker.alloc.nmpi.ResourceUsage;
import uk.ac.manchester.spinnaker.alloc.nmpi.SessionRequest;
import uk.ac.manchester.spinnaker.alloc.nmpi.SessionResourceUpdate;

/**
 * Manages user quotas.
 *
 * @author Donal Fellows
 */
@Service
public class QuotaManager extends DatabaseAwareBean {
        private static final Logger log = getLogger(QuotaManager.class);

        /**
         * The status of the quote to request from the NMPI service.
         */
        private static final String STATUS_ACCEPTED = "accepted";

        /**
         * An approximation of cores per board.
         */
        private static final int APPROX_CORES_PER_BOARD = 48 * 16;

        /**
         * The number of seconds per hour.
         */
        private static final int SECONDS_PER_HOUR = 60 * 60;

        @Autowired
        private ServiceMasterControl control;

        @Autowired
        private QuotaProperties quotaProps;

        /** The wrapped NMPI proxy, bound to the API key. */
        private NMPI nmpi;

        /**
         * Make the NMPI access interface.
         */
        @PostConstruct
        private void createProxy() {
                var nmpiUrl = quotaProps.getNMPIUrl();
                if (!nmpiUrl.isEmpty()) {
                        nmpi = new NMPI(quotaProps);
                }
        }

        /**
         * Can the user (in a specific group) create another job at this point? If
         * not, they're currently out of resources.
         *
         * @param groupId
         *            What group will the job be accounted against.
         * @return True if they can make a job. False if they can't.
         */
        public boolean mayCreateJob(int groupId) {
                try (var sql = new CreateCheckSQL()) {
                        return sql.transactionRead(() -> sql.mayCreateJob(groupId));
                }
        }

        private final class CreateCheckSQL extends AbstractSQL {
                // TODO These should be combined (but one is an aggregate so...)
                private final Query getQuota = conn.query(GET_GROUP_QUOTA);

                private final Query getCurrentUsage = conn.query(GET_CURRENT_USAGE);

                @Override
                public void close() {
                        getQuota.close();
                        getCurrentUsage.close();
                        super.close();
                }

                private boolean mayCreateJob(int groupId) {
                        return getQuota.call1(result -> {
                                var quota = result.getInteger("quota");
                                log.debug("Group {} has quota {}", groupId, quota);
                                if (isNull(quota)) {
                                        return true;
                                }
                                // Quota is defined; check if current usage exceeds it
                                int usage = getCurrentUsage.call1(
                                                integer("current_usage"), groupId).orElse(0);
                                log.debug("Group {} has usage {}", groupId, usage);
                                // If board-seconds are left, we're good to go
                                return (quota > usage);
                        }, groupId).orElse(true);
                }
        }

        /**
         * Has the execution of a job remained within its group's resource
         * allocation at this point?
         *
         * @param jobId
         *            What job is consuming resources?
         * @return True if the job can continue to run. False if it can't.
         */
        public boolean mayLetJobContinue(int jobId) {
                try (var sql = new ContinueCheckSQL()) {
                        return sql.transactionRead(() -> sql.mayLetJobContinue(jobId));
                }
        }

        /**
         * Has the execution of a job exceeded its group's resource allocation at
         * this point?
         *
         * @param jobId
         *            What job is consuming resources?
         * @return False if the job should be killed. True otherwise.
         */
        public boolean shouldKillJob(int jobId) {
                return !mayLetJobContinue(jobId);
        }

        private final class ContinueCheckSQL extends AbstractSQL {
                private final Query getUsageAndQuota =
                                conn.query(GET_JOB_USAGE_AND_QUOTA);

                @Override
                public void close() {
                        getUsageAndQuota.close();
                        super.close();
                }

                private boolean mayLetJobContinue(int jobId) {
                        return getUsageAndQuota.call1(
                                        // If we have an entry, check if usage <= quota
                                        row -> row.getInt("quota_used") <= row.getInt("quota"),
                                        jobId)

                                        // Otherwise, we'll just allow it
                                        .orElse(true);
                }
        }

        /**
         * Adjust a group's quota.
         *
         * @param groupId
         *            Which group's quota to change
         * @param delta
         *            Amount to change by, in board-seconds
         * @return Information about what group's quota was adjusted and what it has
         *         become.
         */
        public Optional<AdjustedQuota> addQuota(int groupId, int delta) {
                try (var sql = new AdjustQuotaSQL()) {
                        return sql.transaction(() -> sql.adjustQuota(groupId, delta));
                }
        }

        /**
         * Describes the result of the {@link QuotaManager#addQuota(int,int)}
         * operation.
         *
         * @author Donal Fellows
         * @param name
         *            The name of the group.
         * @param quota
         *            The new quota of the group.
         */
        public record AdjustedQuota(String name, Long quota) {
                private AdjustedQuota(Row row) {
                        this(row.getString("group_name"), row.getLong("quota"));
                }
        }

        private final class AdjustQuotaSQL extends AbstractSQL {
                private final Update adjustQuota = conn.update(ADJUST_QUOTA);

                private final Query getQuota = conn.query(GET_GROUP_QUOTA);

                @Override
                public void close() {
                        adjustQuota.close();
                        getQuota.close();
                        super.close();
                }

                private Optional<AdjustedQuota> adjustQuota(int groupId, int delta) {
                        if (adjustQuota.call(delta, groupId) == 0) {
                                return Optional.empty();
                        }
                        return getQuota.call1(AdjustedQuota::new, groupId);
                }
        }

        /**
         * Consolidates usage from finished jobs onto quotas. Runs hourly.
         */
        @Scheduled(cron = "#{quotaProperties.consolidationSchedule}")
        public void consolidateQuotas() {
                if (control.isPaused()) {
                        return;
                }
                // Split off for testability
                try (var c = getConnection()) {
                        doConsolidate(c);
                } catch (DataAccessException e) {
                        if (isBusy(e)) {
                                log.info("database is busy; "
                                                + "will try job quota consolidation processing later");
                                return;
                        }
                        throw e;
                }
        }

        private void doConsolidate(Connection c) {
                try (var sql = new ConsolidateSQL(c)) {
                        sql.transaction(sql::consolidate);
                }
        }

        private final class ConsolidateSQL extends AbstractSQL {
                private final Query getConsoldationTargets =
                                conn.query(GET_CONSOLIDATION_TARGETS);

                private final Update decrementQuota = conn.update(DECREMENT_QUOTA);

                private final Update markConsolidated = conn.update(MARK_CONSOLIDATED);

                ConsolidateSQL(Connection c) {
                        super(c);
                }

                @Override
                public void close() {
                        getConsoldationTargets.close();
                        decrementQuota.close();
                        markConsolidated.close();
                        super.close();
                }

                // Result is arbitrary and ignored
                private Void consolidate() {
                        for (var target : getConsoldationTargets.call(Target::new)) {
                                decrementQuota.call(target.quotaUsed, target.groupId);
                                markConsolidated.call(target.jobId);
                        }
                        return null;
                }

                private class Target {
                        final Object quotaUsed;

                        final int groupId;

                        final int jobId;

                        Target(Row row) {
                                quotaUsed = row.getObject("quota_used");
                                groupId = row.getInt("group_id");
                                jobId = row.getInt("job_id");
                        }
                }
        }

        /**
         * @param quota
         *            The size of quota remaining, in board-seconds.
         * @param units
         *            The units that the quota was measured in on the NMPI.
         */
        private record QuotaInfo(long quota, String units) {
        }

        private QuotaInfo parseQuotaData(List<Project> projects) {
                String units = null;
                long total = 0;
                for (var project : projects) {
                        for (var quota : project.getQuotas()) {
                                if (!quota.getPlatform().equals(quotaProps.getNMPIPlaform())) {
                                        continue;
                                }
                                if (units == null) {
                                        units = quota.getUnits();
                                } else if (units != quota.getUnits()) {
                                        throw new RuntimeException("Quotas have multiple units!");
                                }

                                switch (quota.getUnits()) {
                                case BOARD_SECONDS:
                                        total += (long) (quota.getLimit() - quota.getUsage());
                                        break;
                                case CORE_HOURS:
                                        total += toBoardSeconds(quota.getLimit())
                                                        - toBoardSeconds(quota.getUsage());
                                        break;
                                default:
                                        throw new RuntimeException(
                                                        "Unknown Quota units: " + quota.getUnits());
                                }
                        }
                }
                return new QuotaInfo(total, units);
        }

        final Optional<String> mayCreateNMPISession(String collab) {
                // Read collab from NMPI; fail if not there
                var projects = nmpi.getProjects(STATUS_ACCEPTED, collab);
                var info = parseQuotaData(projects);

                log.debug("Setting quota of collab {} to {}", collab, info.quota());

                // Update quota in group for collab from NMPI
                try (var c = getConnection();
                                var setQuota = c.update(SET_COLLAB_QUOTA)) {
                        c.transaction(() -> setQuota.call(info.quota(), collab));
                }

                if (info.quota() > 0) {
                        return Optional.of(info.units());
                }
                return Optional.empty();
        }

        void associateNMPISession(int jobId, String user, String collab,
                        String quotaUnits) {
                var sessionId = nmpi.createSession(collab, user);

                // Associate NMPI session with Job in the database
                try (var c = getConnection();
                                var setSession = c.update(SET_JOB_SESSION)) {
                        c.transaction(
                                        () -> setSession.call(jobId, sessionId, quotaUnits));
                }
        }

        private final class InflateUser extends AbstractSQL {
                private final Query getUserByName =
                                conn.query(GET_USER_DETAILS_BY_NAME);

                private final Query getGroupByName = conn.query(GET_GROUP_BY_NAME);

                private final Update createUser = conn.update(CREATE_USER);

                private final Update createGroup =
                                conn.update(CREATE_GROUP_IF_NOT_EXISTS);

                private final Update addUserToGroup = conn.update(ADD_USER_TO_GROUP);

                @Override
                public void close() {
                        getUserByName.close();
                        getGroupByName.close();
                        createUser.close();
                        createGroup.close();
                        addUserToGroup.close();
                        super.close();
                }

                private Optional<Integer> getUserByName(String user) {
                        return getUserByName.call1(r -> r.getInt("user_id"), user);
                }

                private Optional<Integer> getGroupByName(String group) {
                        return getGroupByName.call1(r -> r.getInt("group_id"), group);
                }

                private boolean createUser(String user) {
                        return createUser.call(user, null, USER, false, null) > 0;
                }

                private boolean createGroup(String group) {
                        return createGroup.call(group, 0.0, COLLABRATORY) > 0;
                }

                private boolean addUserToGroup(int userId, Optional<Integer> groupId) {
                        return addUserToGroup.call(userId, groupId) > 0;
                }

                /**
                 * Construct the user and group records if necessary and associate them.
                 *
                 * @param user
                 *            The user name.
                 * @param job
                 *            The NMPI job details containing the collabratory (group)
                 *            name.
                 */
                void inflateUser(String user, Job job) {
                        var userId = getUserByName(user).or(() -> {
                                /*
                                 * The user has never logged in directly, so we have to create
                                 * them.
                                 */
                                if (!createUser(user)) {
                                        log.warn("failed to make user: {}", user);
                                }
                                return getUserByName(user);
                        });

                        createGroup(job.getCollab());
                        var groupId = getGroupByName(job.getCollab());
                        addUserToGroup(userId.orElseThrow(() -> new IllegalStateException(
                                        "failed to find or create user")), groupId);
                }
        }

        final Optional<NMPIJobQuotaDetails> mayUseNMPIJob(String user,
                        int nmpiJobId) {
                // Read job from NMPI to get collab ID
                var job = nmpi.getJob(nmpiJobId);

                // If it is possible to run this job, we need to associate the user
                // with it because only special users can run jobs like this.
                try (var sql = new InflateUser()) {
                        sql.transaction(() -> sql.inflateUser(user, job));
                }

                // This is now a collab so check there instead
                var quotaUnits = mayCreateNMPISession(job.getCollab());
                if (quotaUnits.isEmpty()) {
                        return Optional.empty();
                }

                return Optional.of(
                                new NMPIJobQuotaDetails(job.getCollab(), quotaUnits.get()));
        }

        /**
         * @param collabId
         *            The collaboratory ID.
         * @param quotaUnits
         *            The units of the Quota.
         */
        record NMPIJobQuotaDetails(String collabId, String quotaUnits) {
        }

        void associateNMPIJob(int jobId, int nmpiJobId, String quotaUnits) {
                // Associate NMPI Job with job in database
                try (var c = getConnection();
                                var setNMPIJob = c.update(SET_JOB_NMPI_JOB)) {
                        c.transaction(() -> setNMPIJob.call(jobId, nmpiJobId, quotaUnits));
                }
        }

        /** Results of database queries. */
        private record FinishInfo(Optional<Long> quota, Optional<Session> session,
                        Optional<NMPIJob> job) {
        }

        private FinishInfo getFinishingInfo(int jobId) {
                try (var c = getConnection();
                                var getSession = c.query(GET_JOB_SESSION);
                                var getNMPIJob = c.query(GET_JOB_NMPI_JOB);
                                var getUsage = c.query(GET_JOB_USAGE_AND_QUOTA)) {
                        // Get the quota used
                        return c.transaction(false,
                                        () -> new FinishInfo(
                                                        getUsage.call1(r -> r.getLong("quota_used"), jobId),
                                                        getSession.call1(Session::new, jobId),
                                                        getNMPIJob.call1(NMPIJob::new, jobId)));
                }
        }

        final void finishJob(int jobId) {
                // Get the information about the job from the DB
                var info = getFinishingInfo(jobId);

                // From here on, we don't touch the DB but we do touch the network

                if (!info.quota().isPresent()) {
                        // No quota? No update!
                        return;
                }

                // If job has associated session, update quota in session
                info.session().ifPresent(session -> {
                        try {
                                nmpi.setSessionStatusAndResources(session.id(), "finished",
                                                getResourceUsage(info.quota().get(),
                                                                session.quotaUnits()));
                        } catch (BadRequestException e) {
                                log.error(e.getResponse().readEntity(String.class));
                                throw e;
                        }
                });

                // If job has associated NMPI job, update quota on NMPI job
                info.job().ifPresent(nmpiJob -> {
                        try {
                                nmpi.setJobResources(nmpiJob.id(), getResourceUsage(
                                                info.quota().get(), nmpiJob.quotaUnits()));
                        } catch (BadRequestException e) {
                                log.error(e.getResponse().readEntity(String.class));
                                throw e;
                        }
                });
        }

        private record Session(int id, String quotaUnits) {
                private Session(Row r) {
                        this(r.getInt("session_id"), r.getString("quota_units"));
                }
        }

        private record NMPIJob(int id, String quotaUnits) {
                private NMPIJob(Row r) {
                        this(r.getInt("nmpi_job_id"), r.getString("quota_units"));
                }
        }

        private static ResourceUsage getResourceUsage(long boardSeconds,
                        String units) {
                var resourceUsage = new ResourceUsage();
                resourceUsage.setUnits(units);
                if (units.equals(BOARD_SECONDS)) {
                        resourceUsage.setValue(boardSeconds);
                } else if (units.equals(CORE_HOURS)) {
                        resourceUsage.setValue(toCoreHours(boardSeconds));
                } else {
                        throw new IllegalArgumentException("Unknown units " + units);
                }
                return resourceUsage;
        }

        /**
         * Convert board-seconds to core-hours (approximately).
         *
         * @param boardSeconds
         *            The number of board-seconds to convert.
         * @return The number of board-hours, which may have fractional values.
         */
        private static double toCoreHours(long boardSeconds) {
                return ((double) (boardSeconds * APPROX_CORES_PER_BOARD))
                                / SECONDS_PER_HOUR;
        }

        /**
         * Convert core-hours to board-seconds (approximately).
         *
         * @param coreHours
         *            The number of core-hours to convert.
         * @return The integer number of board seconds.
         */
        private static long toBoardSeconds(double coreHours) {
                return (long) ((coreHours * SECONDS_PER_HOUR) / APPROX_CORES_PER_BOARD);
        }

        /**
         * Operations for testing only.
         *
         * @hidden
         */
        @ForTestingOnly
        interface TestAPI {
                /**
                 * Consolidates usage from finished jobs onto quotas.
                 *
                 * @param c
                 *            How to talk to the DB.
                 */
                void doConsolidate(Connection c);
        }

        /**
         * @return The test interface.
         * @deprecated This interface is just for testing.
         * @hidden
         */
        @ForTestingOnly
        @RestrictedApi(explanation = "just for testing", link = "index.html",
                        allowedOnPath = ".*/src/test/java/.*")
        @Deprecated
        TestAPI getTestAPI() {
                ForTestingOnly.Utils.checkForTestClassOnStack();
                return new TestAPI() {
                        @Override
                        public void doConsolidate(Connection c) {
                                QuotaManager.this.doConsolidate(c);
                        }
                };
        }
}

/**
 * Wrapper round the proxy and its API key. Does a bit of wrapping and
 * unwrapping of arguments and results.
 *
 * @author Donal Fellows
 */
final class NMPI {
        private final NMPIv3API proxy;

        private final String apiKey;

        private final String platform;

        NMPI(QuotaProperties quotaProps) {
                proxy = NMPIv3API.createClient(quotaProps.getNMPIUrl());
                apiKey = quotaProps.getNMPIApiKey();
                platform = quotaProps.getNMPIPlaform();
        }

        Job getJob(int jobId) {
                return proxy.getJob(apiKey, jobId);
        }

        void setJobResources(int jobId, ResourceUsage resources) {
                var wrapper = new JobResourceUpdate();
                wrapper.setResourceUsage(resources);
                proxy.setJobResources(apiKey, jobId, wrapper);
        }

        List<Project> getProjects(String status, String collab) {
                return proxy.getProjects(apiKey, status, collab);
        }

        int createSession(String collab, String user) {
                var request = new SessionRequest();
                request.setCollab(collab);
                request.setHardwarePlatform(platform);
                request.setUserId(user);
                return proxy.createSession(apiKey, request).getId();
        }

        void setSessionStatusAndResources(int sessionId, String status,
                        ResourceUsage resources) {
                var wrapper = new SessionResourceUpdate();
                wrapper.setStatus(status);
                wrapper.setResourceUsage(resources);
                proxy.setSessionStatusAndResources(apiKey, sessionId, wrapper);
        }
}

1	/*
2	* Copyright (c) 2021 The University of Manchester
3	*
4	* Licensed under the Apache License, Version 2.0 (the "License");
5	* you may not use this file except in compliance with the License.
6	* You may obtain a copy of the License at
7	*
8	* https://www.apache.org/licenses/LICENSE-2.0
9	*
10	* Unless required by applicable law or agreed to in writing, software
11	* distributed under the License is distributed on an "AS IS" BASIS,
12	* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13	* See the License for the specific language governing permissions and
14	* limitations under the License.
15	*/
16	package uk.ac.manchester.spinnaker.alloc.allocator;
17
18	import static java.util.Objects.isNull;
19	import static org.slf4j.LoggerFactory.getLogger;
20	import static uk.ac.manchester.spinnaker.alloc.db.Row.integer;
21	import static uk.ac.manchester.spinnaker.alloc.db.Utils.isBusy;
22	import static uk.ac.manchester.spinnaker.alloc.nmpi.ResourceUsage.BOARD_SECONDS;
23	import static uk.ac.manchester.spinnaker.alloc.nmpi.ResourceUsage.CORE_HOURS;
24	import static uk.ac.manchester.spinnaker.alloc.model.GroupRecord.GroupType.COLLABRATORY;
25	import static uk.ac.manchester.spinnaker.alloc.security.TrustLevel.USER;
26
27	import java.util.List;
28	import java.util.Optional;
29
30	import javax.annotation.PostConstruct;
31	import javax.ws.rs.BadRequestException;
32
33	import org.slf4j.Logger;
34	import org.springframework.beans.factory.annotation.Autowired;
35	import org.springframework.dao.DataAccessException;
36	import org.springframework.scheduling.annotation.Scheduled;
37	import org.springframework.stereotype.Service;
38
39	import com.google.errorprone.annotations.RestrictedApi;
40
41	import uk.ac.manchester.spinnaker.alloc.ForTestingOnly;
42	import uk.ac.manchester.spinnaker.alloc.ServiceMasterControl;
43	import uk.ac.manchester.spinnaker.alloc.SpallocProperties.QuotaProperties;
44	import uk.ac.manchester.spinnaker.alloc.db.DatabaseAPI.Connection;
45	import uk.ac.manchester.spinnaker.alloc.db.DatabaseAPI.Query;
46	import uk.ac.manchester.spinnaker.alloc.db.DatabaseAPI.Update;
47	import uk.ac.manchester.spinnaker.alloc.db.DatabaseAwareBean;
48	import uk.ac.manchester.spinnaker.alloc.db.Row;
49	import uk.ac.manchester.spinnaker.alloc.nmpi.Job;
50	import uk.ac.manchester.spinnaker.alloc.nmpi.JobResourceUpdate;
51	import uk.ac.manchester.spinnaker.alloc.nmpi.NMPIv3API;
52	import uk.ac.manchester.spinnaker.alloc.nmpi.Project;
53	import uk.ac.manchester.spinnaker.alloc.nmpi.ResourceUsage;
54	import uk.ac.manchester.spinnaker.alloc.nmpi.SessionRequest;
55	import uk.ac.manchester.spinnaker.alloc.nmpi.SessionResourceUpdate;
56
57	/**
58	* Manages user quotas.
59	*
60	* @author Donal Fellows
61	*/
62	@Service
63	public class QuotaManager extends DatabaseAwareBean {	1✔
64	private static final Logger log = getLogger(QuotaManager.class);	1✔
65
66	/**
67	* The status of the quote to request from the NMPI service.
68	*/
69	private static final String STATUS_ACCEPTED = "accepted";
70
71	/**
72	* An approximation of cores per board.
73	*/
74	private static final int APPROX_CORES_PER_BOARD = 48 * 16;
75
76	/**
77	* The number of seconds per hour.
78	*/
79	private static final int SECONDS_PER_HOUR = 60 * 60;
80
81	@Autowired
82	private ServiceMasterControl control;
83
84	@Autowired
85	private QuotaProperties quotaProps;
86
87	/** The wrapped NMPI proxy, bound to the API key. */
88	private NMPI nmpi;
89
90	/**
91	* Make the NMPI access interface.
92	*/
93	@PostConstruct
94	private void createProxy() {
95	var nmpiUrl = quotaProps.getNMPIUrl();	1✔
96	if (!nmpiUrl.isEmpty()) {	1✔
97	nmpi = new NMPI(quotaProps);	×
98	}
99	}	1✔
100
101	/**
102	* Can the user (in a specific group) create another job at this point? If
103	* not, they're currently out of resources.
104	*
105	* @param groupId
106	* What group will the job be accounted against.
107	* @return True if they can make a job. False if they can't.
108	*/
109	public boolean mayCreateJob(int groupId) {
110	try (var sql = new CreateCheckSQL()) {	1✔
111	return sql.transactionRead(() -> sql.mayCreateJob(groupId));	1✔
112	}
113	}
114
115	private final class CreateCheckSQL extends AbstractSQL {	1✔
116	// TODO These should be combined (but one is an aggregate so...)
117	private final Query getQuota = conn.query(GET_GROUP_QUOTA);	1✔
118
119	private final Query getCurrentUsage = conn.query(GET_CURRENT_USAGE);	1✔
120
121	@Override
122	public void close() {
123	getQuota.close();	1✔
124	getCurrentUsage.close();	1✔
125	super.close();	1✔
126	}	1✔
127
128	private boolean mayCreateJob(int groupId) {
129	return getQuota.call1(result -> {	1✔
130	var quota = result.getInteger("quota");	1✔
131	log.debug("Group {} has quota {}", groupId, quota);	1✔
132	if (isNull(quota)) {	1✔
133	return true;	×
134	}
135	// Quota is defined; check if current usage exceeds it
136	int usage = getCurrentUsage.call1(	1✔
137	integer("current_usage"), groupId).orElse(0);	1✔
138	log.debug("Group {} has usage {}", groupId, usage);	1✔
139	// If board-seconds are left, we're good to go
140	return (quota > usage);	1✔
141	}, groupId).orElse(true);	1✔
142	}
143	}
144
145	/**
146	* Has the execution of a job remained within its group's resource
147	* allocation at this point?
148	*
149	* @param jobId
150	* What job is consuming resources?
151	* @return True if the job can continue to run. False if it can't.
152	*/
153	public boolean mayLetJobContinue(int jobId) {
154	try (var sql = new ContinueCheckSQL()) {	×
155	return sql.transactionRead(() -> sql.mayLetJobContinue(jobId));	×
156	}
157	}
158
159	/**
160	* Has the execution of a job exceeded its group's resource allocation at
161	* this point?
162	*
163	* @param jobId
164	* What job is consuming resources?
165	* @return False if the job should be killed. True otherwise.
166	*/
167	public boolean shouldKillJob(int jobId) {
168	return !mayLetJobContinue(jobId);	×
169	}
170
171	private final class ContinueCheckSQL extends AbstractSQL {	×
172	private final Query getUsageAndQuota =	×
173	conn.query(GET_JOB_USAGE_AND_QUOTA);	×
174
175	@Override
176	public void close() {
177	getUsageAndQuota.close();	×
178	super.close();	×
179	}	×
180
181	private boolean mayLetJobContinue(int jobId) {
182	return getUsageAndQuota.call1(	×
183	// If we have an entry, check if usage <= quota
184	row -> row.getInt("quota_used") <= row.getInt("quota"),	×
185	jobId)	×
186
187	// Otherwise, we'll just allow it
188	.orElse(true);	×
189	}
190	}
191
192	/**
193	* Adjust a group's quota.
194	*
195	* @param groupId
196	* Which group's quota to change
197	* @param delta
198	* Amount to change by, in board-seconds
199	* @return Information about what group's quota was adjusted and what it has
200	* become.
201	*/
202	public Optional<AdjustedQuota> addQuota(int groupId, int delta) {
203	try (var sql = new AdjustQuotaSQL()) {	×
204	return sql.transaction(() -> sql.adjustQuota(groupId, delta));	×
205	}
206	}
207
208	/**
209	* Describes the result of the {@link QuotaManager#addQuota(int,int)}
210	* operation.
211	*
212	* @author Donal Fellows
213	* @param name
214	* The name of the group.
215	* @param quota
216	* The new quota of the group.
217	*/
218	public record AdjustedQuota(String name, Long quota) {	×
219	private AdjustedQuota(Row row) {
220	this(row.getString("group_name"), row.getLong("quota"));	×
221	}	×
222	}
223
224	private final class AdjustQuotaSQL extends AbstractSQL {	×
225	private final Update adjustQuota = conn.update(ADJUST_QUOTA);	×
226
227	private final Query getQuota = conn.query(GET_GROUP_QUOTA);	×
228
229	@Override
230	public void close() {
231	adjustQuota.close();	×
232	getQuota.close();	×
233	super.close();	×
234	}	×
235
236	private Optional<AdjustedQuota> adjustQuota(int groupId, int delta) {
237	if (adjustQuota.call(delta, groupId) == 0) {	×
238	return Optional.empty();	×
239	}
240	return getQuota.call1(AdjustedQuota::new, groupId);	×
241	}
242	}
243
244	/**
245	* Consolidates usage from finished jobs onto quotas. Runs hourly.
246	*/
247	@Scheduled(cron = "#{quotaProperties.consolidationSchedule}")
248	public void consolidateQuotas() {
249	if (control.isPaused()) {	×
250	return;	×
251	}
252	// Split off for testability
253	try (var c = getConnection()) {	×
254	doConsolidate(c);	×
255	} catch (DataAccessException e) {	×
256	if (isBusy(e)) {	×
257	log.info("database is busy; "	×
258	+ "will try job quota consolidation processing later");
259	return;	×
260	}
261	throw e;	×
262	}	×
263	}	×
264
265	private void doConsolidate(Connection c) {
266	try (var sql = new ConsolidateSQL(c)) {	1✔
267	sql.transaction(sql::consolidate);	1✔
268	}
269	}	1✔
270
271	private final class ConsolidateSQL extends AbstractSQL {
272	private final Query getConsoldationTargets =	1✔
273	conn.query(GET_CONSOLIDATION_TARGETS);	1✔
274
275	private final Update decrementQuota = conn.update(DECREMENT_QUOTA);	1✔
276
277	private final Update markConsolidated = conn.update(MARK_CONSOLIDATED);	1✔
278
279	ConsolidateSQL(Connection c) {	1✔
280	super(c);	1✔
281	}	1✔
282
283	@Override
284	public void close() {
285	getConsoldationTargets.close();	1✔
286	decrementQuota.close();	1✔
287	markConsolidated.close();	1✔
288	super.close();	1✔
289	}	1✔
290
291	// Result is arbitrary and ignored
292	private Void consolidate() {
293	for (var target : getConsoldationTargets.call(Target::new)) {	1✔
294	decrementQuota.call(target.quotaUsed, target.groupId);	1✔
295	markConsolidated.call(target.jobId);	1✔
296	}	1✔
297	return null;	1✔
298	}
299
300	private class Target {
301	final Object quotaUsed;
302
303	final int groupId;
304
305	final int jobId;
306
307	Target(Row row) {	1✔
308	quotaUsed = row.getObject("quota_used");	1✔
309	groupId = row.getInt("group_id");	1✔
310	jobId = row.getInt("job_id");	1✔
311	}	1✔
312	}
313	}
314
315	/**
316	* @param quota
317	* The size of quota remaining, in board-seconds.
318	* @param units
319	* The units that the quota was measured in on the NMPI.
320	*/
321	private record QuotaInfo(long quota, String units) {	×
322	}
323
324	private QuotaInfo parseQuotaData(List<Project> projects) {
325	String units = null;	×
326	long total = 0;	×
327	for (var project : projects) {	×
328	for (var quota : project.getQuotas()) {	×
329	if (!quota.getPlatform().equals(quotaProps.getNMPIPlaform())) {	×
330	continue;	×
331	}
332	if (units == null) {	×
333	units = quota.getUnits();	×
334	} else if (units != quota.getUnits()) {	×
335	throw new RuntimeException("Quotas have multiple units!");	×
336	}
337
338	switch (quota.getUnits()) {	×
339	case BOARD_SECONDS:
340	total += (long) (quota.getLimit() - quota.getUsage());	×
341	break;	×
342	case CORE_HOURS:
343	total += toBoardSeconds(quota.getLimit())	×
344	- toBoardSeconds(quota.getUsage());	×
345	break;	×
346	default:
347	throw new RuntimeException(	×
348	"Unknown Quota units: " + quota.getUnits());	×
349	}
350	}	×
351	}	×
352	return new QuotaInfo(total, units);	×
353	}
354
355	final Optional<String> mayCreateNMPISession(String collab) {
356	// Read collab from NMPI; fail if not there
357	var projects = nmpi.getProjects(STATUS_ACCEPTED, collab);	×
358	var info = parseQuotaData(projects);	×
359
360	log.debug("Setting quota of collab {} to {}", collab, info.quota());	×
361
362	// Update quota in group for collab from NMPI
363	try (var c = getConnection();	×
364	var setQuota = c.update(SET_COLLAB_QUOTA)) {	×
365	c.transaction(() -> setQuota.call(info.quota(), collab));	×
366	}
367
368	if (info.quota() > 0) {	×
369	return Optional.of(info.units());	×
370	}
371	return Optional.empty();	×
372	}
373
374	void associateNMPISession(int jobId, String user, String collab,
375	String quotaUnits) {
376	var sessionId = nmpi.createSession(collab, user);	×
377
378	// Associate NMPI session with Job in the database
379	try (var c = getConnection();	×
380	var setSession = c.update(SET_JOB_SESSION)) {	×
381	c.transaction(	×
382	() -> setSession.call(jobId, sessionId, quotaUnits));	×
383	}
384	}	×
385
386	private final class InflateUser extends AbstractSQL {	×
387	private final Query getUserByName =	×
388	conn.query(GET_USER_DETAILS_BY_NAME);	×
389
390	private final Query getGroupByName = conn.query(GET_GROUP_BY_NAME);	×
391
392	private final Update createUser = conn.update(CREATE_USER);	×
393
394	private final Update createGroup =	×
395	conn.update(CREATE_GROUP_IF_NOT_EXISTS);	×
396
397	private final Update addUserToGroup = conn.update(ADD_USER_TO_GROUP);	×
398
399	@Override
400	public void close() {
401	getUserByName.close();	×
402	getGroupByName.close();	×
403	createUser.close();	×
404	createGroup.close();	×
405	addUserToGroup.close();	×
406	super.close();	×
407	}	×
408
409	private Optional<Integer> getUserByName(String user) {
410	return getUserByName.call1(r -> r.getInt("user_id"), user);	×
411	}
412
413	private Optional<Integer> getGroupByName(String group) {
414	return getGroupByName.call1(r -> r.getInt("group_id"), group);	×
415	}
416
417	private boolean createUser(String user) {
418	return createUser.call(user, null, USER, false, null) > 0;	×
419	}
420
421	private boolean createGroup(String group) {
422	return createGroup.call(group, 0.0, COLLABRATORY) > 0;	×
423	}
424
425	private boolean addUserToGroup(int userId, Optional<Integer> groupId) {
426	return addUserToGroup.call(userId, groupId) > 0;	×
427	}
428
429	/**
430	* Construct the user and group records if necessary and associate them.
431	*
432	* @param user
433	* The user name.
434	* @param job
435	* The NMPI job details containing the collabratory (group)
436	* name.
437	*/
438	void inflateUser(String user, Job job) {
439	var userId = getUserByName(user).or(() -> {	×
440	/*
441	* The user has never logged in directly, so we have to create
442	* them.
443	*/
444	if (!createUser(user)) {	×
445	log.warn("failed to make user: {}", user);	×
446	}
447	return getUserByName(user);	×
448	});
449
450	createGroup(job.getCollab());	×
451	var groupId = getGroupByName(job.getCollab());	×
452	addUserToGroup(userId.orElseThrow(() -> new IllegalStateException(	×
453	"failed to find or create user")), groupId);
454	}	×
455	}
456
457	final Optional<NMPIJobQuotaDetails> mayUseNMPIJob(String user,
458	int nmpiJobId) {
459	// Read job from NMPI to get collab ID
460	var job = nmpi.getJob(nmpiJobId);	×
461
462	// If it is possible to run this job, we need to associate the user
463	// with it because only special users can run jobs like this.
464	try (var sql = new InflateUser()) {	×
465	sql.transaction(() -> sql.inflateUser(user, job));	×
466	}
467
468	// This is now a collab so check there instead
469	var quotaUnits = mayCreateNMPISession(job.getCollab());	×
470	if (quotaUnits.isEmpty()) {	×
471	return Optional.empty();	×
472	}
473
474	return Optional.of(	×
475	new NMPIJobQuotaDetails(job.getCollab(), quotaUnits.get()));	×
476	}
477
478	/**
479	* @param collabId
480	* The collaboratory ID.
481	* @param quotaUnits
482	* The units of the Quota.
483	*/
484	record NMPIJobQuotaDetails(String collabId, String quotaUnits) {	×
485	}
486
487	void associateNMPIJob(int jobId, int nmpiJobId, String quotaUnits) {
488	// Associate NMPI Job with job in database
489	try (var c = getConnection();	×
490	var setNMPIJob = c.update(SET_JOB_NMPI_JOB)) {	×
491	c.transaction(() -> setNMPIJob.call(jobId, nmpiJobId, quotaUnits));	×
492	}
493	}	×
494
495	/** Results of database queries. */
496	private record FinishInfo(Optional<Long> quota, Optional<Session> session,	1✔
497	Optional<NMPIJob> job) {
498	}
499
500	private FinishInfo getFinishingInfo(int jobId) {
501	try (var c = getConnection();	1✔
502	var getSession = c.query(GET_JOB_SESSION);	1✔
503	var getNMPIJob = c.query(GET_JOB_NMPI_JOB);	1✔
504	var getUsage = c.query(GET_JOB_USAGE_AND_QUOTA)) {	1✔
505	// Get the quota used
506	return c.transaction(false,	1✔
507	() -> new FinishInfo(	1✔
508	getUsage.call1(r -> r.getLong("quota_used"), jobId),	1✔
509	getSession.call1(Session::new, jobId),	1✔
510	getNMPIJob.call1(NMPIJob::new, jobId)));	1✔
511	}
512	}
513
514	final void finishJob(int jobId) {
515	// Get the information about the job from the DB
516	var info = getFinishingInfo(jobId);	1✔
517
518	// From here on, we don't touch the DB but we do touch the network
519
520	if (!info.quota().isPresent()) {	1✔
521	// No quota? No update!
522	return;	×
523	}
524
525	// If job has associated session, update quota in session
526	info.session().ifPresent(session -> {	1✔
527	try {
528	nmpi.setSessionStatusAndResources(session.id(), "finished",	×
529	getResourceUsage(info.quota().get(),	×
530	session.quotaUnits()));	×
531	} catch (BadRequestException e) {	×
532	log.error(e.getResponse().readEntity(String.class));	×
533	throw e;	×
534	}	×
535	});	×
536
537	// If job has associated NMPI job, update quota on NMPI job
538	info.job().ifPresent(nmpiJob -> {	1✔
539	try {
540	nmpi.setJobResources(nmpiJob.id(), getResourceUsage(	×
541	info.quota().get(), nmpiJob.quotaUnits()));	×
542	} catch (BadRequestException e) {	×
543	log.error(e.getResponse().readEntity(String.class));	×
544	throw e;	×
545	}	×
546	});	×
547	}	1✔
548
549	private record Session(int id, String quotaUnits) {	×
550	private Session(Row r) {
551	this(r.getInt("session_id"), r.getString("quota_units"));	×
552	}	×
553	}
554
555	private record NMPIJob(int id, String quotaUnits) {	×
556	private NMPIJob(Row r) {
557	this(r.getInt("nmpi_job_id"), r.getString("quota_units"));	×
558	}	×
559	}
560
561	private static ResourceUsage getResourceUsage(long boardSeconds,
562	String units) {
563	var resourceUsage = new ResourceUsage();	×
564	resourceUsage.setUnits(units);	×
565	if (units.equals(BOARD_SECONDS)) {	×
566	resourceUsage.setValue(boardSeconds);	×
567	} else if (units.equals(CORE_HOURS)) {	×
568	resourceUsage.setValue(toCoreHours(boardSeconds));	×
569	} else {
570	throw new IllegalArgumentException("Unknown units " + units);	×
571	}
572	return resourceUsage;	×
573	}
574
575	/**
576	* Convert board-seconds to core-hours (approximately).
577	*
578	* @param boardSeconds
579	* The number of board-seconds to convert.
580	* @return The number of board-hours, which may have fractional values.
581	*/
582	private static double toCoreHours(long boardSeconds) {
583	return ((double) (boardSeconds * APPROX_CORES_PER_BOARD))	×
584	/ SECONDS_PER_HOUR;
585	}
586
587	/**
588	* Convert core-hours to board-seconds (approximately).
589	*
590	* @param coreHours
591	* The number of core-hours to convert.
592	* @return The integer number of board seconds.
593	*/
594	private static long toBoardSeconds(double coreHours) {
595	return (long) ((coreHours * SECONDS_PER_HOUR) / APPROX_CORES_PER_BOARD);	×
596	}
597
598	/**
599	* Operations for testing only.
600	*
601	* @hidden
602	*/
603	@ForTestingOnly
604	interface TestAPI {
605	/**
606	* Consolidates usage from finished jobs onto quotas.
607	*
608	* @param c
609	* How to talk to the DB.
610	*/
611	void doConsolidate(Connection c);
612	}
613
614	/**
615	* @return The test interface.
616	* @deprecated This interface is just for testing.
617	* @hidden
618	*/
619	@ForTestingOnly
620	@RestrictedApi(explanation = "just for testing", link = "index.html",
621	allowedOnPath = "./src/test/java/.")
622	@Deprecated
623	TestAPI getTestAPI() {
624	ForTestingOnly.Utils.checkForTestClassOnStack();	1✔
625	return new TestAPI() {	1✔
626	@Override
627	public void doConsolidate(Connection c) {
628	QuotaManager.this.doConsolidate(c);	1✔
629	}	1✔
630	};
631	}
632	}
633
634	/**
635	* Wrapper round the proxy and its API key. Does a bit of wrapping and
636	* unwrapping of arguments and results.
637	*
638	* @author Donal Fellows
639	*/
640	final class NMPI {
641	private final NMPIv3API proxy;
642
643	private final String apiKey;
644
645	private final String platform;
646
647	NMPI(QuotaProperties quotaProps) {	×
648	proxy = NMPIv3API.createClient(quotaProps.getNMPIUrl());	×
649	apiKey = quotaProps.getNMPIApiKey();	×
650	platform = quotaProps.getNMPIPlaform();	×
651	}	×
652
653	Job getJob(int jobId) {
654	return proxy.getJob(apiKey, jobId);	×
655	}
656
657	void setJobResources(int jobId, ResourceUsage resources) {
658	var wrapper = new JobResourceUpdate();	×
659	wrapper.setResourceUsage(resources);	×
660	proxy.setJobResources(apiKey, jobId, wrapper);	×
661	}	×
662
663	List<Project> getProjects(String status, String collab) {
664	return proxy.getProjects(apiKey, status, collab);	×
665	}
666
667	int createSession(String collab, String user) {
668	var request = new SessionRequest();	×
669	request.setCollab(collab);	×
670	request.setHardwarePlatform(platform);	×
671	request.setUserId(user);	×
672	return proxy.createSession(apiKey, request).getId();	×
673	}
674
675	void setSessionStatusAndResources(int sessionId, String status,
676	ResourceUsage resources) {
677	var wrapper = new SessionResourceUpdate();	×
678	wrapper.setStatus(status);	×
679	wrapper.setResourceUsage(resources);	×
680	proxy.setSessionStatusAndResources(apiKey, sessionId, wrapper);	×
681	}	×
682	}

SpiNNakerManchester / JavaSpiNNaker / 6233274834

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous