Skip to content
This repository has been archived by the owner on Aug 2, 2023. It is now read-only.

Commit

Permalink
Add domain_name, group_id, user_uuid fields in kernels (#148)
Browse files Browse the repository at this point in the history
Migration:
* 'default' group in the 'default' domain will be created.
* Those fields are automatically filled in migration with values of
(domain_name = 'default', group_id=<gid of the default group>,
 user_uuid = user's uuid whose access_key is used to run a container>)
* Every user is automatically associated with the 'default' group.
adrysn committed May 28, 2019
1 parent cee2f9d commit 9701c88
Showing 2 changed files with 99 additions and 3 deletions.
Original file line number Diff line number Diff line change
@@ -0,0 +1,94 @@
"""add domain, group, user fields to kernels
Revision ID: c5e4e764f9e3
Revises: 6f1c1b83870a
Create Date: 2019-05-28 10:22:56.904061
"""
from alembic import op
import sqlalchemy as sa
import ai.backend.manager.models.base # noqa


# revision identifiers, used by Alembic.
revision = 'c5e4e764f9e3'
down_revision = '6f1c1b83870a'
branch_labels = None
depends_on = None


def upgrade():
# ### commands auto generated by Alembic - please adjust! ###
op.add_column('kernels', sa.Column('domain_name', sa.String(length=64), nullable=True))
op.add_column('kernels', sa.Column('group_id', ai.backend.manager.models.base.GUID(), nullable=True))
op.add_column('kernels', sa.Column('user_uuid', ai.backend.manager.models.base.GUID(), nullable=True))
op.create_foreign_key(op.f('fk_kernels_group_id_groups'), 'kernels', 'groups', ['group_id'], ['id'])
op.create_foreign_key(op.f('fk_kernels_user_uuid_users'), 'kernels', 'users', ['user_uuid'], ['uuid'])
op.create_foreign_key(op.f('fk_kernels_domain_name_domains'), 'kernels', 'domains', ['domain_name'], ['name'])
# ### end Alembic commands ###

import textwrap
from ai.backend.manager.models import kernels, keypairs, groups, users, association_groups_users

# Create default group in the default domain.
# Assumption: "default" domain must exist
connection = op.get_bind()
query = (sa.insert(groups)
.values(name='default', description='Default group', is_active=True,
domain_name='default'))
query = textwrap.dedent('''\
INSERT INTO groups (name, description, is_active, domain_name)
VALUES ('default', 'Default group', True, 'default')
ON CONFLICT (name, domain_name) DO NOTHING
RETURNING id;
''')
result = connection.execute(query).fetchone()
gid = result.id if hasattr(result, 'id') else None
if gid is None: # group already exists
query = textwrap.dedent('''\
SELECT id FROM groups where name = 'default' and domain_name = 'default';
''')
gid = connection.execute(query).fetchone().id

# Fill in kernels' domain_name, group_id, and user_uuid.
query = sa.select([kernels]).select_from(kernels)
all_kernels = connection.execute(query).fetchall()
for kernel in all_kernels:
# Get kernel's keypair (access_key).
query = (sa.select([keypairs.c.user]).select_from(keypairs)
.where(keypairs.c.access_key == kernel.access_key))
kp = connection.execute(query).fetchone()
# Update kernel information.
query = '''\
UPDATE kernels SET domain_name = 'default', group_id = '%s', user_uuid = '%s'
WHERE id = '%s';
''' % (gid, kp.user, kernel.id)
connection.execute(query)

# Associate every users with the default group.
# NOTE: this operation is not undoable unless you drop groups table.
query = sa.select([users.c.uuid]).select_from(users)
all_users = connection.execute(query).fetchall()
for user in all_users:
query = '''\
INSERT INTO association_groups_users (group_id, user_id)
VALUES ('%s', '%s')
ON CONFLICT (group_id, user_id) DO NOTHING;
''' % (gid, user.uuid,)
connection.execute(query)

# Make kernel's new fields non-nullable.
op.alter_column('kernels', column_name='domain_name', nullable=False)
op.alter_column('kernels', column_name='group_id', nullable=False)
op.alter_column('kernels', column_name='user_uuid', nullable=False)


def downgrade():
# ### commands auto generated by Alembic - please adjust! ###
op.drop_constraint(op.f('fk_kernels_domain_name_domains'), 'kernels', type_='foreignkey')
op.drop_constraint(op.f('fk_kernels_user_uuid_users'), 'kernels', type_='foreignkey')
op.drop_constraint(op.f('fk_kernels_group_id_groups'), 'kernels', type_='foreignkey')
op.drop_column('kernels', 'user_uuid')
op.drop_column('kernels', 'group_id')
op.drop_column('kernels', 'domain_name')
# ### end Alembic commands ###
8 changes: 5 additions & 3 deletions src/ai/backend/manager/models/kernel.py
Original file line number Diff line number Diff line change
@@ -9,7 +9,7 @@
from ai.backend.common.types import BinarySize
from .base import (
metadata, zero_if_none,
BigInt, IDColumn, EnumType,
BigInt, GUID, IDColumn, EnumType,
ResourceSlotColumn,
Item, PaginatedList,
)
@@ -46,8 +46,10 @@ class KernelStatus(enum.Enum):
sa.Column('role', sa.String(length=16), nullable=False, default='master'),
sa.Column('agent', sa.String(length=64), sa.ForeignKey('agents.id')),
sa.Column('agent_addr', sa.String(length=128), nullable=False),
sa.Column('access_key', sa.String(length=20),
sa.ForeignKey('keypairs.access_key')),
sa.Column('domain_name', sa.String(length=64), sa.ForeignKey('domains.name'), nullable=False),
sa.Column('group_id', GUID, sa.ForeignKey('groups.id'), nullable=False),
sa.Column('user_uuid', GUID, sa.ForeignKey('users.uuid'), nullable=False),
sa.Column('access_key', sa.String(length=20), sa.ForeignKey('keypairs.access_key')),
sa.Column('image', sa.String(length=512)),
sa.Column('registry', sa.String(length=512)),
sa.Column('tag', sa.String(length=64), nullable=True),

0 comments on commit 9701c88

Please sign in to comment.