record costs for all gpu tasks

daras_ai_v2/gpu_server.py

@@ -2,6 +2,7 @@
 import datetime
 import os
 import typing
+from time import time
 
 import requests
 from furl import furl
@@ -138,6 +139,7 @@ def get_celery():
         _app = Celery()
         _app.conf.broker_url = settings.GPU_CELERY_BROKER_URL
         _app.conf.result_backend = settings.GPU_CELERY_RESULT_BACKEND
+        _app.conf.result_extended = True
 
     return _app
 
@@ -149,8 +151,16 @@ def call_celery_task(
     inputs: dict,
     queue_prefix: str = "gooey-gpu",
 ):
+    from usage_costs.cost_utils import record_cost_auto
+    from usage_costs.models import ModelSku
+
     queue = os.path.join(queue_prefix, pipeline["model_id"].strip()).strip("/")
     result = get_celery().send_task(
         task_name, kwargs=dict(pipeline=pipeline, inputs=inputs), queue=queue
     )
-    return result.get(disable_sync_subtasks=False)
+    s = time()
+    ret = result.get(disable_sync_subtasks=False)
+    record_cost_auto(
+        model=queue, sku=ModelSku.gpu_ms, quantity=int((time() - s) * 1000)
+    )
+    return ret

scripts/init_self_hosted_pricing.py

@@ -0,0 +1,2 @@
+def run():
+    pass

usage_costs/admin.py

@@ -1,5 +1,3 @@
-from decimal import Decimal
-
 from django.contrib import admin
 
 from bots.admin_links import open_in_new_tab, change_obj_url
@@ -21,6 +19,7 @@ class UsageCostAdmin(admin.ModelAdmin, CostQtyMixin):
         "display_dollar_amount",
         "view_pricing",
         "view_saved_run",
+        "view_parent_published_run",
         "notes",
         "created_at",
     ]
@@ -41,6 +40,11 @@ class UsageCostAdmin(admin.ModelAdmin, CostQtyMixin):
     def display_dollar_amount(self, obj):
         return f"${obj.dollar_amount.normalize()}"
 
+    @admin.display(description="Published Run")
+    def view_parent_published_run(self, obj):
+        pr = obj.saved_run.parent_published_run()
+        return pr and change_obj_url(pr)
+
     @admin.display(description="Saved Run", ordering="saved_run")
     def view_saved_run(self, obj):
         return change_obj_url(

usage_costs/migrations/0002_alter_modelpricing_category_and_more.py

@@ -0,0 +1,38 @@
+# Generated by Django 4.2.7 on 2024-02-07 15:44
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('usage_costs', '0001_initial'),
+    ]
+
+    operations = [
+        migrations.AlterField(
+            model_name='modelpricing',
+            name='category',
+            field=models.IntegerField(choices=[(1, 'LLM'), (2, 'Self-Hosted')], help_text='The category of the model. Only used for Display purposes.'),
+        ),
+        migrations.AlterField(
+            model_name='modelpricing',
+            name='model_name',
+            field=models.CharField(choices=[('gpt_4_vision', 'GPT-4 Vision (openai)'), ('gpt_4_turbo', 'GPT-4 Turbo (openai)'), ('gpt_4', 'GPT-4 (openai)'), ('gpt_4_32k', 'GPT-4 32K (openai)'), ('gpt_3_5_turbo', 'ChatGPT (openai)'), ('gpt_3_5_turbo_16k', 'ChatGPT 16k (openai)'), ('llama2_70b_chat', 'Llama 2 (Meta AI)'), ('palm2_chat', 'PaLM 2 Chat (Google)'), ('palm2_text', 'PaLM 2 Text (Google)'), ('text_davinci_003', 'GPT-3.5 Davinci-3 (openai)'), ('text_davinci_002', 'GPT-3.5 Davinci-2 (openai)'), ('text_curie_001', 'Curie (openai)'), ('text_babbage_001', 'Babbage (openai)'), ('text_ada_001', 'Ada (openai)'), ('code_davinci_002', 'Codex [Deprecated] (openai)'), ('protogen_2_2', 'Protogen V2.2 (darkstorm2150)'), ('epicdream', 'epiCDream (epinikion)')], help_text='The name of the model. Only used for Display purposes.', max_length=255),
+        ),
+        migrations.AlterField(
+            model_name='modelpricing',
+            name='provider',
+            field=models.IntegerField(choices=[(1, 'OpenAI'), (2, 'Google'), (3, 'TogetherAI'), (4, 'Azure OpenAI'), (5, 'Azure Kubernetes Service')], help_text='The provider of the model. Only used for Display purposes.'),
+        ),
+        migrations.AlterField(
+            model_name='modelpricing',
+            name='sku',
+            field=models.IntegerField(choices=[(1, 'LLM Prompt'), (2, 'LLM Completion'), (3, 'GPU Milliseconds')], help_text="The model's SKU. Model ID + SKU should be unique together."),
+        ),
+        migrations.AlterField(
+            model_name='modelpricing',
+            name='unit_quantity',
+            field=models.PositiveIntegerField(default=1, help_text='The quantity of the unit. (e.g. 1000 tokens)'),
+        ),
+    ]

usage_costs/models.py

@@ -40,6 +40,7 @@ def __str__(self):
 
 class ModelCategory(models.IntegerChoices):
     LLM = 1, "LLM"
+    SELF_HOSTED = 2, "Self-Hosted"
 
 
 class ModelProvider(models.IntegerChoices):
@@ -48,17 +49,24 @@ class ModelProvider(models.IntegerChoices):
     together_ai = 3, "TogetherAI"
     azure_openai = 4, "Azure OpenAI"
 
+    aks = 5, "Azure Kubernetes Service"
+
 
 def get_model_choices():
     from daras_ai_v2.language_model import LargeLanguageModels
+    from recipes.DeforumSD import AnimationModels
 
-    return [(api.name, api.value) for api in LargeLanguageModels]
+    return [(api.name, api.value) for api in LargeLanguageModels] + [
+        (model.name, model.label) for model in AnimationModels
+    ]
 
 
 class ModelSku(models.IntegerChoices):
     llm_prompt = 1, "LLM Prompt"
     llm_completion = 2, "LLM Completion"
 
+    gpu_ms = 3, "GPU Milliseconds"
+
 
 class ModelPricing(models.Model):
     model_id = models.TextField(
@@ -75,7 +83,7 @@ class ModelPricing(models.Model):
         help_text="The cost per unit.",
     )
     unit_quantity = models.PositiveIntegerField(
-        help_text="The quantity of the unit. (e.g. 1000 tokens)"
+        help_text="The quantity of the unit. (e.g. 1000 tokens)", default=1
     )
 
     category = models.IntegerField(