Merge pull request #463 from Datura-ai/feature/pods-module

Feature/pods module

apps/server/controllers/pod.py

@@ -0,0 +1,62 @@
+from fastapi_sqlalchemy import db
+from fastapi import APIRouter, Depends, HTTPException, status
+from utils.auth import authenticate
+from typings.auth import UserAccount
+from models.pod import PodModel
+from typings.pod import PodOutput, PodInput, CreatePodOutput
+
+router = APIRouter()
+
+
+@router.get(
+    "",
+    response_model=list[PodOutput],
+    status_code=200
+)
+def get_pods(auth: UserAccount = Depends(authenticate)):
+    try:
+        data = PodModel.get_pods(
+            db=db,
+            account=auth.account
+        )
+
+        data_to_list = [PodOutput.from_orm(pod) for pod in data]
+
+        return data_to_list
+    except Exception as e:
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail=str(e)
+        )
+
+
+@router.post(
+    "",
+    response_model=CreatePodOutput,
+    status_code=201
+)
+def create_pod(input: PodInput, auth: UserAccount = Depends(authenticate)):
+    try:
+        pod = PodModel.create_pod(
+            db=db,
+            pod=input,
+            user=auth.user,
+            account=auth.account
+        )
+
+        if not pod:
+            raise HTTPException(
+                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+                detail="Pod not created"
+            )
+
+        return {
+                "success": True,
+                "message": "Pod created successfully"
+            }
+
+    except Exception as e:
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail=str(e)
+        )

apps/server/controllers/resource.py

@@ -0,0 +1,77 @@
+from fastapi_sqlalchemy import db
+from fastapi import APIRouter, Depends, HTTPException, status
+from utils.auth import authenticate
+from typings.auth import UserAccount
+from models.resource import ResourceModel
+from typings.resource import ResourceOutput, ResourceInput
+
+router = APIRouter()
+
+
+@router.get(
+    "",
+    response_model=list[ResourceOutput],
+    status_code=200
+)
+def get_resources():
+    """_summary_
+
+    Raises:
+        HTTPException: _description_
+
+    Returns:
+        _type_: _description_
+    """
+    try:
+        data = ResourceModel.get_resources(db=db)
+
+        data_to_list = [ResourceOutput.from_orm(resource) for resource in data]
+
+        return data_to_list
+    except Exception as e:
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail=str(e)
+        )
+
+
+@router.post(
+    "",
+    # response_model=list[PodOutput],
+    status_code=201
+)
+def create_resource(input: ResourceInput):
+    """_summary_
+
+    Args:
+        input (ResourceInput): _description_
+
+    Raises:
+        HTTPException: _description_
+        HTTPException: _description_
+
+    Returns:
+        _type_: _description_
+    """
+    try:
+        resource = ResourceModel.create_resource(
+            db=db,
+            resource=input,
+        )
+
+        if not resource:
+            raise HTTPException(
+                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+                detail="Pod not created"
+            )
+
+        return {
+                "success": True,
+                "message": "Resource created successfully"
+            }
+
+    except Exception as e:
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail=str(e)
+        )

apps/server/controllers/template.py

@@ -0,0 +1,68 @@
+from fastapi_sqlalchemy import db
+from fastapi import APIRouter, Depends, HTTPException, status
+from utils.auth import authenticate
+from typings.auth import UserAccount
+from models.template import TemplateModel
+from typings.template import (
+    TemplateOutput,
+    TemplateInput,
+    CreateTemplateOutput
+)
+
+router = APIRouter()
+
+
+@router.get(
+    "",
+    response_model=list[TemplateOutput],
+    status_code=200
+)
+def get_templates(auth: UserAccount = Depends(authenticate)):
+    try:
+        data = TemplateModel.get_templates(
+            db=db,
+            account=auth.account
+        )
+
+        data_to_list = [TemplateOutput.from_orm(template) for template in data]
+
+        return data_to_list
+    except Exception as e:
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail=str(e)
+        )
+
+
+@router.post(
+    "",
+    response_model=CreateTemplateOutput,
+    status_code=201
+)
+def create_template(
+        input: TemplateInput,
+        auth: UserAccount = Depends(authenticate)):
+    try:
+        pod = TemplateModel.create_template(
+            db=db,
+            template=input,
+            user=auth.user,
+            account=auth.account
+        )
+
+        if not pod:
+            raise HTTPException(
+                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+                detail="Pod not created"
+            )
+
+        return {
+                "success": True,
+                "message": "Template created successfully"
+            }
+
+    except Exception as e:
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail=str(e)
+        )

apps/server/exceptions.py

@@ -158,5 +158,9 @@ class UserAccessNotFoundException(DatasourceException):
     pass
 
 
+class TemplateNotFoundException(DatasourceException):
+    pass
+
+
 class UserAccountAccessException(AppBaseException):
     pass

apps/server/main.py

@@ -30,6 +30,9 @@
 from controllers.voice import router as voice_router
 from controllers.workspace import router as workspace_router
 from controllers.user_account_access import router as user_account_access_router
+from controllers.pod import router as pod_router
+from controllers.resource import router as resource_router
+from controllers.template import router as template_router
 from models import *  # noqa: F403
 from models.db import Base, engine  # noqa: F401
 from resolvers.account import AccountMutation, AccountQuery
@@ -135,6 +138,9 @@ def jwt_exception_handler(request: Request, exc: AuthJWTException):
 app.include_router(fine_tuning_router, prefix="/fine-tuning", include_in_schema=False)
 app.include_router(voice_router, prefix="/voice", include_in_schema=True)
 app.include_router(user_account_access_router, prefix="/user-account-access", include_in_schema=False)
+app.include_router(pod_router, prefix="/pod", include_in_schema=False)
+app.include_router(resource_router, prefix="/resource", include_in_schema=False)
+app.include_router(template_router, prefix="/template", include_in_schema=False)
 
 
 @app.get("/", include_in_schema=False)

apps/server/migrations/versions/041177c95b30_add_podmodel_resourcemodel.py

@@ -0,0 +1,86 @@
+"""add PodModel ResourceModel
+
+Revision ID: 041177c95b30
+Revises: f1d5bc37bceb
+Create Date: 2024-06-06 18:28:45.199706
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision: str = '041177c95b30'
+down_revision: Union[str, None] = 'f1d5bc37bceb'
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table('resource',
+    sa.Column('id', sa.UUID(), nullable=False),
+    sa.Column('name', sa.String(length=255), nullable=False),
+    sa.Column('display_name', sa.String(length=255), nullable=False),
+    sa.Column('type', sa.Enum('cpu', 'gpu', name='resourcetypeenum'), nullable=False),
+    sa.Column('category', sa.String(), nullable=True),
+    sa.Column('ram', sa.Numeric(precision=5, scale=2), nullable=False),
+    sa.Column('secure_price', sa.Numeric(precision=5, scale=2), nullable=True),
+    sa.Column('one_month_price', sa.Numeric(precision=5, scale=2), nullable=True),
+    sa.Column('three_month_price', sa.Numeric(precision=5, scale=2), nullable=True),
+    sa.Column('six_month_price', sa.Numeric(precision=5, scale=2), nullable=True),
+    sa.Column('max_gpu', sa.Numeric(precision=5, scale=2), nullable=True),
+    sa.Column('lowest_price', postgresql.JSONB(astext_type=sa.Text()), nullable=False),
+    sa.Column('status', sa.Enum('low', 'high', 'unavailable', name='resourcestatusenum'), nullable=True),
+    sa.Column('disc_type', sa.Enum('ssd', 'nvme', 'unavailable', name='resourcedisctypeenum'), nullable=True),
+    sa.Column('cloud_type', sa.Enum('secure_cloud', 'community_cloud', 'unavailable', name='resourcecloudtypeenum'), nullable=True),
+    sa.Column('region', sa.String(length=255), nullable=True),
+    sa.Column('cuda_version', sa.Numeric(precision=5, scale=2), nullable=True),
+    sa.Column('is_deleted', sa.Boolean(), nullable=True),
+    sa.Column('created_on', sa.DateTime(timezone=True), nullable=False),
+    sa.Column('updated_on', sa.DateTime(timezone=True), nullable=False),
+    sa.PrimaryKeyConstraint('id')
+    )
+    op.create_index(op.f('ix_resource_is_deleted'), 'resource', ['is_deleted'], unique=False)
+    op.create_table('pod',
+    sa.Column('id', sa.UUID(), nullable=False),
+    sa.Column('pod_name', sa.String(), nullable=True),
+    sa.Column('price', sa.Numeric(precision=5, scale=2), nullable=True),
+    sa.Column('status', sa.Enum('running', 'stopped', name='podstatusenum'), nullable=True),
+    sa.Column('provider', sa.String(), nullable=True),
+    sa.Column('category', sa.String(), nullable=True),
+    sa.Column('type', sa.Enum('cpu', 'gpu', name='podtypeenum'), nullable=True),
+    sa.Column('resource', sa.UUID(), nullable=True),
+    sa.Column('gpu_count', sa.Numeric(precision=5, scale=2), nullable=True),
+    sa.Column('isinstance_pricing', postgresql.JSONB(astext_type=sa.Text()), nullable=False),
+    sa.Column('is_deleted', sa.Boolean(), nullable=True),
+    sa.Column('account_id', sa.UUID(), nullable=True),
+    sa.Column('created_by', sa.UUID(), nullable=True),
+    sa.Column('modified_by', sa.UUID(), nullable=True),
+    sa.Column('created_on', sa.DateTime(timezone=True), nullable=False),
+    sa.Column('updated_on', sa.DateTime(timezone=True), nullable=False),
+    sa.ForeignKeyConstraint(['account_id'], ['account.id'], ondelete='CASCADE'),
+    sa.ForeignKeyConstraint(['created_by'], ['user.id'], name='fk_created_by', ondelete='CASCADE'),
+    sa.ForeignKeyConstraint(['modified_by'], ['user.id'], name='fk_modified_by', ondelete='CASCADE'),
+    sa.ForeignKeyConstraint(['resource'], ['resource.id'], ondelete='CASCADE'),
+    sa.PrimaryKeyConstraint('id')
+    )
+    op.create_index(op.f('ix_pod_created_by'), 'pod', ['created_by'], unique=False)
+    op.create_index(op.f('ix_pod_is_deleted'), 'pod', ['is_deleted'], unique=False)
+    op.create_index(op.f('ix_pod_modified_by'), 'pod', ['modified_by'], unique=False)
+    op.create_index(op.f('ix_pod_resource'), 'pod', ['resource'], unique=False)
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_index(op.f('ix_pod_resource'), table_name='pod')
+    op.drop_index(op.f('ix_pod_modified_by'), table_name='pod')
+    op.drop_index(op.f('ix_pod_is_deleted'), table_name='pod')
+    op.drop_index(op.f('ix_pod_created_by'), table_name='pod')
+    op.drop_table('pod')
+    op.drop_index(op.f('ix_resource_is_deleted'), table_name='resource')
+    op.drop_table('resource')
+    # ### end Alembic commands ###

apps/server/migrations/versions/8ee8b2ab331f_template_pod_relation.py

@@ -0,0 +1,34 @@
+"""Template Pod relation
+
+Revision ID: 8ee8b2ab331f
+Revises: aa2dd6264fd3
+Create Date: 2024-06-07 10:07:32.210936
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision: str = '8ee8b2ab331f'
+down_revision: Union[str, None] = 'aa2dd6264fd3'
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column('pod', sa.Column('template', sa.UUID(), nullable=True))
+    op.create_index(op.f('ix_pod_template'), 'pod', ['template'], unique=False)
+    op.create_foreign_key(None, 'pod', 'template', ['template'], ['id'], ondelete='CASCADE')
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_constraint(None, 'pod', type_='foreignkey')
+    op.drop_index(op.f('ix_pod_template'), table_name='pod')
+    op.drop_column('pod', 'template')
+    # ### end Alembic commands ###