diff --git a/cyborg/api/controllers/v1/accelerators.py b/cyborg/api/controllers/v1/accelerators.py index fddceff0..95a4c2c6 100644 --- a/cyborg/api/controllers/v1/accelerators.py +++ b/cyborg/api/controllers/v1/accelerators.py @@ -230,9 +230,8 @@ class AcceleratorsController(AcceleratorsControllerBase): @expose.expose(None, types.uuid, status_code=http_client.NO_CONTENT) def delete(self, uuid): """Delete an accelerator. - :param uuid: UUID of an accelerator. """ - obj_acc = self._resource or self._get_resource(uuid) context = pecan.request.context + obj_acc = self._resource or self._get_resource(uuid) pecan.request.conductor_api.accelerator_delete(context, obj_acc) diff --git a/cyborg/api/controllers/v1/deployables.py b/cyborg/api/controllers/v1/deployables.py index 641c0203..c7c2ca69 100644 --- a/cyborg/api/controllers/v1/deployables.py +++ b/cyborg/api/controllers/v1/deployables.py @@ -27,6 +27,7 @@ from cyborg.api import expose from cyborg.common import exception from cyborg.common import policy from cyborg import objects +from cyborg.quota import QUOTAS class Deployable(base.APIBase): @@ -226,13 +227,29 @@ class DeployablesController(base.CyborgController): :param patch: a json PATCH document to apply to this deployable. """ context = pecan.request.context + reservations = None + obj_dep = objects.Deployable.get(context, uuid) try: + # TODO(xinran): need more discussion on quota's granularity. + # Now we count by board. + for p in patch: + if p["path"] == "/instance_uuid" and p["op"] == "replace": + if not p["value"]: + obj_dep["assignable"] = True + reserve_opts = {obj_dep["board"]: -1} + else: + obj_dep["assignable"] = False + reserve_opts = {obj_dep["board"]: 1} + reservations = QUOTAS.reserve(context, reserve_opts) api_dep = Deployable( **api_utils.apply_jsonpatch(obj_dep.as_dict(), patch)) except api_utils.JSONPATCH_EXCEPTIONS as e: + QUOTAS.rollback(context, reservations, project_id=None) raise exception.PatchError(patch=patch, reason=e) + QUOTAS.commit(context, reservations) + # Update only the fields that have changed for field in objects.Deployable.fields: try: diff --git a/cyborg/common/exception.py b/cyborg/common/exception.py index a7492d9d..effa5794 100644 --- a/cyborg/common/exception.py +++ b/cyborg/common/exception.py @@ -289,3 +289,19 @@ class InventoryInUse(InvalidInventory): # cyborg.services.client.report._RE_INV_IN_USE regex. msg_fmt = _("Inventory for '%(resource_classes)s' on " "resource provider '%(resource_provider)s' in use.") + + +class QuotaNotFound(NotFound): + message = _("Quota could not be found") + + +class QuotaUsageNotFound(QuotaNotFound): + message = _("Quota usage for project %(project_id)s could not be found.") + + +class QuotaResourceUnknown(QuotaNotFound): + message = _("Unknown quota resources %(unknown)s.") + + +class InvalidReservationExpiration(Invalid): + message = _("Invalid reservation expiration %(expire)s.") diff --git a/cyborg/db/__init__.py b/cyborg/db/__init__.py index e69de29b..d18e7586 100644 --- a/cyborg/db/__init__.py +++ b/cyborg/db/__init__.py @@ -0,0 +1,20 @@ +# Copyright 2010 United States Government as represented by the +# Administrator of the National Aeronautics and Space Administration. +# All Rights Reserved. +# +# 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 +# +# http://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. +""" +DB abstraction for Cyborg +""" + +from cyborg.db.api import * # noqa diff --git a/cyborg/db/api.py b/cyborg/db/api.py index 320c5285..048bb18f 100644 --- a/cyborg/db/api.py +++ b/cyborg/db/api.py @@ -123,3 +123,13 @@ class Connection(object): @abc.abstractmethod def attribute_delete(self, context, uuid): """Delete an attribute.""" + + @abc.abstractmethod + def quota_reserve(self, context, resources, deltas, expire, + until_refresh, max_age, project_id=None, + is_allocated_reserve=False): + """Check quotas and create appropriate reservations.""" + + @abc.abstractmethod + def reservation_commit(self, context, reservations, project_id=None): + """Check quotas and create appropriate reservations.""" diff --git a/cyborg/db/sqlalchemy/alembic/versions/d6f033d8fa5b_add_quota_related_tables.py b/cyborg/db/sqlalchemy/alembic/versions/d6f033d8fa5b_add_quota_related_tables.py new file mode 100644 index 00000000..36ad93ff --- /dev/null +++ b/cyborg/db/sqlalchemy/alembic/versions/d6f033d8fa5b_add_quota_related_tables.py @@ -0,0 +1,70 @@ +# 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 +# +# http://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. + +"""add-quota-related-tables + +Revision ID: d6f033d8fa5b +Revises: f50980397351 +Create Date: 2018-04-28 03:07:06.857245 + +""" + +# revision identifiers, used by Alembic. +revision = 'd6f033d8fa5b' +down_revision = 'f50980397351' + +from alembic import op +import sqlalchemy as sa + + +def upgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.create_table( + 'quota_usages', + sa.Column('created_at', sa.DateTime(), nullable=True), + sa.Column('updated_at', sa.DateTime(), nullable=True), + sa.Column('id', sa.Integer(), nullable=False), + sa.Column('project_id', sa.String(length=255), nullable=True), + sa.Column('user_id', sa.String(length=255), nullable=True), + sa.Column('resource', sa.String(length=255), nullable=False), + sa.Column('in_use', sa.Integer(), nullable=False), + sa.Column('reserved', sa.Integer(), nullable=False), + sa.Column('until_refresh', sa.Integer(), nullable=True), + sa.PrimaryKeyConstraint('id') + ) + op.create_index('ix_quota_usages_project_id', 'quota_usages', + ['project_id'], unique=False) + op.create_index('ix_quota_usages_user_id', 'quota_usages', ['user_id'], + unique=False) + + op.create_table( + 'reservations', + sa.Column('created_at', sa.DateTime(), nullable=True), + sa.Column('updated_at', sa.DateTime(), nullable=True), + sa.Column('id', sa.Integer(), nullable=False), + sa.Column('uuid', sa.String(length=36), nullable=False), + sa.Column('usage_id', sa.Integer(), nullable=False), + sa.Column('project_id', sa.String(length=255), nullable=True), + sa.Column('user_id', sa.String(length=255), nullable=True), + sa.Column('resource', sa.String(length=255), nullable=True), + sa.Column('delta', sa.Integer(), nullable=False), + sa.Column('expire', sa.DateTime(), nullable=True), + sa.ForeignKeyConstraint(['usage_id'], ['quota_usages.id'], ), + sa.PrimaryKeyConstraint('id') + ) + op.create_index('ix_reservations_project_id', 'reservations', + ['project_id'], unique=False) + op.create_index('ix_reservations_user_id', 'reservations', ['user_id'], + unique=False) + op.create_index('reservations_uuid_idx', 'reservations', ['uuid'], + unique=False) + # ### end Alembic commands ### diff --git a/cyborg/db/sqlalchemy/api.py b/cyborg/db/sqlalchemy/api.py index a346bac2..b45988bc 100644 --- a/cyborg/db/sqlalchemy/api.py +++ b/cyborg/db/sqlalchemy/api.py @@ -17,6 +17,7 @@ import threading import copy +import uuid from oslo_db import api as oslo_db_api from oslo_db import exception as db_exc @@ -24,8 +25,12 @@ from oslo_db.sqlalchemy import enginefacade from oslo_db.sqlalchemy import utils as sqlalchemyutils from oslo_log import log from oslo_utils import strutils +from oslo_utils import timeutils from oslo_utils import uuidutils +from sqlalchemy.orm import load_only from sqlalchemy.orm.exc import NoResultFound +from sqlalchemy.sql import func + from cyborg.common import exception from cyborg.common.i18n import _ @@ -37,6 +42,8 @@ from sqlalchemy import and_ _CONTEXT = threading.local() LOG = log.getLogger(__name__) +main_context_manager = enginefacade.transaction_context() + def get_backend(): """The backend is this module itself.""" @@ -51,6 +58,11 @@ def _session_for_write(): return enginefacade.writer.using(_CONTEXT) +def get_session(use_slave=False, **kwargs): + return main_context_manager._factory.get_legacy_facade().get_session( + use_slave=use_slave, **kwargs) + + def model_query(context, model, *args, **kwargs): """Query helper for simpler session usage. @@ -498,6 +510,185 @@ class Connection(api.Connection): if count != 1: raise exception.AttributeNotFound(uuid=uuid) + def _get_quota_usages(self, context, project_id, resources=None): + # Broken out for testability + query = model_query(context, models.QuotaUsage,).filter_by( + project_id=project_id) + if resources: + query = query.filter(models.QuotaUsage.resource.in_( + list(resources))) + rows = query.order_by(models.QuotaUsage.id.asc()). \ + with_for_update().all() + return {row.resource: row for row in rows} + + def _quota_usage_create(self, project_id, resource, until_refresh, + in_use, reserved, session=None): + + quota_usage_ref = models.QuotaUsage() + quota_usage_ref.project_id = project_id + quota_usage_ref.resource = resource + quota_usage_ref.in_use = in_use + quota_usage_ref.reserved = reserved + quota_usage_ref.until_refresh = until_refresh + quota_usage_ref.save(session=session) + + return quota_usage_ref + + def _reservation_create(self, uuid, usage, project_id, resource, delta, + expire, session=None): + usage_id = usage['id'] if usage else None + reservation_ref = models.Reservation() + reservation_ref.uuid = uuid + reservation_ref.usage_id = usage_id + reservation_ref.project_id = project_id + reservation_ref.resource = resource + reservation_ref.delta = delta + reservation_ref.expire = expire + reservation_ref.save(session=session) + return reservation_ref + + def _get_reservation_resources(self, context, reservation_ids): + """Return the relevant resources by reservations.""" + + reservations = model_query(context, models.Reservation). \ + options(load_only('resource')). \ + filter(models.Reservation.uuid.in_(reservation_ids)). \ + all() + return {r.resource for r in reservations} + + def _quota_reservations(self, session, context, reservations): + """Return the relevant reservations.""" + + # Get the listed reservations + return model_query(context, models.Reservation). \ + filter(models.Reservation.uuid.in_(reservations)). \ + with_lockmode('update'). \ + all() + + def quota_reserve(self, context, resources, deltas, expire, + until_refresh, max_age, project_id=None, + is_allocated_reserve=False): + """ Create reservation record in DB according to params""" + with _session_for_write() as session: + if project_id is None: + project_id = context.project_id + usages = self._get_quota_usages(context, project_id, + resources=deltas.keys()) + work = set(deltas.keys()) + while work: + resource = work.pop() + + # Do we need to refresh the usage? + refresh = False + # create quota usage in DB if there is no record of this type + # of resource + if resource not in usages: + usages[resource] = self._quota_usage_create(project_id, + resource, + until_refresh + or None, + in_use=0, + reserved=0, + session=session + ) + refresh = True + elif usages[resource].in_use < 0: + # Negative in_use count indicates a desync, so try to + # heal from that... + refresh = True + elif usages[resource].until_refresh is not None: + usages[resource].until_refresh -= 1 + if usages[resource].until_refresh <= 0: + refresh = True + elif max_age and usages[resource].updated_at is not None and ( + (timeutils.utcnow() - + usages[resource].updated_at).total_seconds() >= + max_age): + refresh = True + + # refresh the usage + if refresh: + # Grab the sync routine + updates = self._sync_acc_res(context, resource, project_id) + for res, in_use in updates.items(): + # Make sure we have a destination for the usage! + if res not in usages: + usages[res] = self._quota_usage_create( + project_id, + res, + until_refresh or None, + in_use=0, + reserved=0, + session=session + ) + + # Update the usage + usages[res].in_use = in_use + usages[res].until_refresh = until_refresh or None + + # Because more than one resource may be refreshed + # by the call to the sync routine, and we don't + # want to double-sync, we make sure all refreshed + # resources are dropped from the work set. + work.discard(res) + + # NOTE(Vek): We make the assumption that the sync + # routine actually refreshes the + # resources that it is the sync routine + # for. We don't check, because this is + # a best-effort mechanism. + + unders = [r for r, delta in deltas.items() + if delta < 0 and delta + usages[r].in_use < 0] + reservations = [] + for resource, delta in deltas.items(): + usage = usages[resource] + reservation = self._reservation_create( + str(uuid.uuid4()), usage, project_id, resource, + delta, expire, session=session) + reservations.append(reservation.uuid) + usages[resource].reserved += delta + session.flush() + if unders: + LOG.warning("Change will make usage less than 0 for the following " + "resources: %s", unders) + return reservations + + def _sync_acc_res(self, context, resource, project_id): + """Quota sync funciton""" + res_in_use = self._accelerator_data_get_for_project(context, resource, + project_id) + return {resource: res_in_use} + + def _accelerator_data_get_for_project(self, context, resource, project_id): + """Return the number of resource which is being used by a project""" + query = model_query(context, models.Accelerator).\ + filter_by(project_id=project_id).filter_by(device_type=resource) + + return query.count() + + def _dict_with_usage_id(self, usages): + return {row.id: row for row in usages.values()} + + def reservation_commit(self, context, reservations, project_id=None): + """Commit quota reservation to quota usage table""" + with _session_for_write() as session: + quota_usage = self._get_quota_usages( + context, project_id, + resources=self._get_reservation_resources(context, + reservations)) + usages = self._dict_with_usage_id(quota_usage) + + for reservation in self._quota_reservations(session, context, + reservations): + + usage = usages[reservation.usage_id] + if reservation.delta >= 0: + usage.reserved -= reservation.delta + usage.in_use += reservation.delta + session.flush() + reservation.delete(session=session) + def process_sort_params(self, sort_keys, sort_dirs, default_keys=['created_at', 'id'], default_dir='asc'): diff --git a/cyborg/db/sqlalchemy/models.py b/cyborg/db/sqlalchemy/models.py index 842bdeb9..733dcb9a 100644 --- a/cyborg/db/sqlalchemy/models.py +++ b/cyborg/db/sqlalchemy/models.py @@ -17,11 +17,14 @@ from oslo_db import options as db_options from oslo_db.sqlalchemy import models +from oslo_utils import timeutils import six.moves.urllib.parse as urlparse from sqlalchemy.ext.declarative import declarative_base from sqlalchemy import Column, String, Integer, Boolean, ForeignKey, Index from sqlalchemy import Text from sqlalchemy import schema +from sqlalchemy import DateTime +from sqlalchemy import orm from cyborg.common import paths from cyborg.conf import CONF @@ -48,6 +51,18 @@ class CyborgBase(models.TimestampMixin, models.ModelBase): d[c.name] = self[c.name] return d + @staticmethod + def delete_values(): + return {'deleted': True, + 'deleted_at': timeutils.utcnow()} + + def delete(self, session): + """Delete this object.""" + updated_values = self.delete_values() + self.update(updated_values) + self.save(session=session) + return updated_values + Base = declarative_base(cls=CyborgBase) @@ -124,3 +139,54 @@ class Attribute(Base): nullable=False) key = Column(Text, nullable=False) value = Column(Text, nullable=False) + + +class QuotaUsage(Base): + """Represents the current usage for a given resource.""" + + __tablename__ = 'quota_usages' + __table_args__ = ( + Index('ix_quota_usages_project_id', 'project_id'), + Index('ix_quota_usages_user_id', 'user_id'), + ) + id = Column(Integer, primary_key=True) + + project_id = Column(String(255)) + user_id = Column(String(255)) + resource = Column(String(255), nullable=False) + + in_use = Column(Integer, nullable=False) + reserved = Column(Integer, nullable=False) + + @property + def total(self): + return self.in_use + self.reserved + + until_refresh = Column(Integer) + + +class Reservation(Base): + """Represents a resource reservation for quotas.""" + + __tablename__ = 'reservations' + __table_args__ = ( + Index('ix_reservations_project_id', 'project_id'), + Index('reservations_uuid_idx', 'uuid'), + Index('ix_reservations_user_id', 'user_id'), + ) + id = Column(Integer, primary_key=True, nullable=False) + uuid = Column(String(36), nullable=False) + + usage_id = Column(Integer, ForeignKey('quota_usages.id'), nullable=False) + + project_id = Column(String(255)) + user_id = Column(String(255)) + resource = Column(String(255)) + + delta = Column(Integer, nullable=False) + expire = Column(DateTime) + + usage = orm.relationship( + "QuotaUsage", + foreign_keys=usage_id, + primaryjoin=usage_id == QuotaUsage.id) diff --git a/cyborg/quota.py b/cyborg/quota.py new file mode 100644 index 00000000..9a2ed260 --- /dev/null +++ b/cyborg/quota.py @@ -0,0 +1,190 @@ +# Copyright 2018 Intel, Inc. +# +# 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 +# +# http://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. + +import datetime +from oslo_config import cfg +from oslo_log import log as logging +from oslo_utils import timeutils +import six + +from cyborg import db as db_api +from cyborg.common import exception + +LOG = logging.getLogger(__name__) + +quota_opts = [ + cfg.IntOpt('reservation_expire', + default=86400, + help='Number of seconds until a reservation expires'), + cfg.IntOpt('until_refresh', + default=0, + help='Count of reservations until usage is refreshed'), + cfg.StrOpt('quota_driver', + default="cyborg.quota.DbQuotaDriver", + help='Default driver to use for quota checks'), + cfg.IntOpt('quota_fpgas', + default=10, + help='Total amount of fpga allowed per project'), + cfg.IntOpt('quota_gpus', + default=10, + help='Total amount of storage allowed per project'), + cfg.IntOpt('max_age', + default=0, + help='Number of seconds between subsequent usage refreshes') + ] + +CONF = cfg.CONF +CONF.register_opts(quota_opts) + + +class QuotaEngine(object): + """Represent the set of recognized quotas.""" + + def __init__(self, quota_driver_class=None): + """Initialize a Quota object.""" + + self._resources = {} + self._driver = DbQuotaDriver() + + def register_resource(self, resource): + """Register a resource.""" + self._resources[resource.name] = resource + + def register_resources(self, resources): + """Register a list of resources.""" + for resource in resources: + self.register_resource(resource) + + def reserve(self, context, deltas, expire=None, project_id=None): + """Check quotas and reserve resources. + + For counting quotas--those quotas for which there is a usage + synchronization function--this method checks quotas against + current usage and the desired deltas. The deltas are given as + keyword arguments, and current usage and other reservations + are factored into the quota check. + + This method will raise a QuotaResourceUnknown exception if a + given resource is unknown or if it does not have a usage + synchronization function. + + If any of the proposed values is over the defined quota, an + OverQuota exception will be raised with the sorted list of the + resources which are too high. Otherwise, the method returns a + list of reservation UUIDs which were created. + + :param context: The request context, for access checks. + :param expire: An optional parameter specifying an expiration + time for the reservations. If it is a simple + number, it is interpreted as a number of + seconds and added to the current time; if it is + a datetime.timedelta object, it will also be + added to the current time. A datetime.datetime + object will be interpreted as the absolute + expiration time. If None is specified, the + default expiration time set by + --default-reservation-expire will be used (this + value will be treated as a number of seconds). + :param project_id: Specify the project_id if current context + is admin and admin wants to impact on + common user's tenant. + """ + if not project_id: + project_id = context.project_id + reservations = self._driver.reserve(context, self._resources, deltas, + expire=expire, + project_id=project_id) + + LOG.debug("Created reservations %s", reservations) + + return reservations + + def commit(self, context, reservations, project_id=None): + """Commit reservations. + + :param context: The request context, for access checks. + :param reservations: A list of the reservation UUIDs, as + returned by the reserve() method. + :param project_id: Specify the project_id if current context + is admin and admin wants to impact on + common user's tenant. + """ + project_id = context.project_id + try: + self._driver.commit(context, reservations, project_id=project_id) + except Exception: + # NOTE(Vek): Ignoring exceptions here is safe, because the + # usage resynchronization and the reservation expiration + # mechanisms will resolve the issue. The exception is + # logged, however, because this is less than optimal. + LOG.exception("Failed to commit reservations %s", reservations) + + def rollback(self, context, reservations, project_id=None): + pass + + +class DbQuotaDriver(object): + """Driver to perform check to enforcement of quotas. + + Also allows to obtain quota information. + The default driver utilizes the local database. + """ + dbapi = db_api.get_instance() + + def reserve(self, context, resources, deltas, expire=None, + project_id=None): + # Set up the reservation expiration + if expire is None: + expire = CONF.reservation_expire + if isinstance(expire, six.integer_types): + expire = datetime.timedelta(seconds=expire) + if isinstance(expire, datetime.timedelta): + expire = timeutils.utcnow() + expire + if not isinstance(expire, datetime.datetime): + raise exception.InvalidReservationExpiration(expire=expire) + + # If project_id is None, then we use the project_id in context + if project_id is None: + project_id = context.project_id + + return self._reserve(context, resources, deltas, expire, + project_id) + + def _reserve(self, context, resources, deltas, expire, project_id): + return self.dbapi.quota_reserve(context, resources, deltas, expire, + CONF.until_refresh, CONF.max_age, + project_id=project_id) + + def commit(self, context, reservations, project_id=None): + """Commit reservations. + + :param context: The request context, for access checks. + :param reservations: A list of the reservation UUIDs, as + returned by the reserve() method. + :param project_id: Specify the project_id if current context + is admin and admin wants to impact on + common user's tenant. + """ + + try: + self.dbapi.reservation_commit(context, reservations, + project_id=project_id) + except Exception: + # NOTE(Vek): Ignoring exceptions here is safe, because the + # usage resynchronization and the reservation expiration + # mechanisms will resolve the issue. The exception is + # logged, however, because this is less than optimal. + LOG.exception("Failed to commit reservations %s", reservations) + +QUOTAS = QuotaEngine() diff --git a/cyborg/tests/unit/db/test_db_api.py b/cyborg/tests/unit/db/test_db_api.py new file mode 100644 index 00000000..1d1e69c8 --- /dev/null +++ b/cyborg/tests/unit/db/test_db_api.py @@ -0,0 +1,129 @@ +# Copyright 2018 Intel, Inc. +# +# 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 +# +# http://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. + +"""Unit tests for the DB api.""" + +import datetime +from cyborg.tests.unit.db import base +from cyborg.db import api as dbapi +from cyborg.db.sqlalchemy import api as sqlalchemyapi + + +def _quota_reserve(context, project_id): + """Create sample QuotaUsage and Reservation objects. + + There is no method db.quota_usage_create(), so we have to use + db.quota_reserve() for creating QuotaUsage objects. + + Returns reservations uuids. + + """ + sqlalchemy_api = sqlalchemyapi.get_backend() + resources = {} + deltas = {} + for i, resource in enumerate(('fpga', 'gpu')): + deltas[resource] = i + 1 + return sqlalchemy_api.quota_reserve( + context, resources, deltas, + datetime.datetime.utcnow(), datetime.datetime.utcnow(), + datetime.timedelta(days=1), project_id + ) + + +class DBAPIQuotaUsageTestCase(base.DbTestCase): + + """Tests for db.api.quota_usage_* methods.""" + + def test_quota_reserve(self): + sqlalchemy_api = sqlalchemyapi.get_backend() + reservations = _quota_reserve(self.context, 'project1') + self.assertEqual(2, len(reservations)) + quota_usages = sqlalchemy_api._get_quota_usages(self.context, + 'project1') + result = {'project_id': "project1"} + for k, v in quota_usages.items(): + result[v.resource] = dict(in_use=v.in_use, reserved=v.reserved) + + self.assertEqual({'project_id': 'project1', + 'gpu': {'reserved': 2, 'in_use': 0}, + 'fpga': {'reserved': 1, 'in_use': 0}}, + result) + + def test__get_quota_usages(self): + _quota_reserve(self.context, 'project1') + sqlalchemy_api = sqlalchemyapi.get_backend() + quota_usages = sqlalchemy_api._get_quota_usages(self.context, + 'project1') + + self.assertEqual(['fpga', 'gpu'], + sorted(quota_usages.keys())) + + def test__get_quota_usages_with_resources(self): + _quota_reserve(self.context, 'project1') + sqlalchemy_api = sqlalchemyapi.get_backend() + quota_usage = sqlalchemy_api._get_quota_usages( + self.context, 'project1', resources=['gpu']) + + self.assertEqual(['gpu'], list(quota_usage.keys())) + + +class DBAPIReservationTestCase(base.DbTestCase): + + """Tests for db.api.reservation_* methods.""" + + def setUp(self): + super(DBAPIReservationTestCase, self).setUp() + self.values = { + 'uuid': 'sample-uuid', + 'project_id': 'project1', + 'resource': 'resource', + 'delta': 42, + 'expire': (datetime.datetime.utcnow() + + datetime.timedelta(days=1)), + 'usage': {'id': 1} + } + + def test__get_reservation_resources(self): + sqlalchemy_api = sqlalchemyapi.get_backend() + reservations = _quota_reserve(self.context, 'project1') + expected = ['fpga', 'gpu'] + resources = sqlalchemy_api._get_reservation_resources( + self.context, reservations) + self.assertEqual(expected, sorted(resources)) + + def test_reservation_commit(self): + db_api = dbapi.get_instance() + reservations = _quota_reserve(self.context, 'project1') + expected = {'project_id': 'project1', + 'fpga': {'reserved': 1, 'in_use': 0}, + 'gpu': {'reserved': 2, 'in_use': 0}, + } + quota_usages = db_api._get_quota_usages(self.context, 'project1') + result = {'project_id': "project1"} + for k, v in quota_usages.items(): + result[v.resource] = dict(in_use=v.in_use, reserved=v.reserved) + + self.assertEqual(expected, result) + + db_api.reservation_commit(self.context, reservations, 'project1') + expected = {'project_id': 'project1', + 'fpga': {'reserved': 0, 'in_use': 1}, + 'gpu': {'reserved': 0, 'in_use': 2}, + } + quota_usages = db_api._get_quota_usages(self.context, 'project1') + result = {'project_id': "project1"} + for k, v in quota_usages.items(): + result[v.resource] = dict(in_use=v.in_use, + reserved=v.reserved) + self.assertEqual(expected, result)