Fixed model cleanup on cluster

CHANGELOG.md

@@ -19,6 +19,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Add create events waiting on db apply.
 - Refactor secrets loading method.
 - Add db.load in db wait
+- Add model component cleanup
 
 #### New Features & Functionality
 

superduper/backends/base/cluster.py

@@ -91,3 +91,13 @@ def initialize(self, with_compute: bool = False):
         self.vector_search.initialize()
         self.crontab.initialize()
         self.cdc.initialize()
+
+    def drop_component(self, uuid: str):
+        """Drop component and its services rom the cluster.
+
+        :param uuid: Component uuid.
+        """
+        try:
+            del self.cache[uuid]
+        except KeyError:
+            pass

superduper/backends/base/compute.py

@@ -110,3 +110,9 @@ def db(self, value: 'Datalayer'):
         :param value: ``Datalayer`` instance.
         """
         self._db = value
+
+    def drop_component(self, uuid: str):
+        """Drop the component from compute.
+
+        :param uuid: Component uuid.
+        """

superduper/base/datalayer.py

@@ -509,6 +509,7 @@ def remove(
 
             for v in sorted(versions_in_use):
                 self.metadata.hide_component_version(type_id, identifier, v)
+
         else:
             logging.warn('aborting.')
 

superduper/components/model.py

@@ -400,6 +400,13 @@ def __post_init__(self, db, example):
         if not self.identifier:
             raise Exception('_Predictor identifier must be non-empty')
 
+    def cleanup(self, db: "Datalayer") -> None:
+        """Clean up when the model is deleted.
+
+        :param db: Data layer instance to process.
+        """
+        db.cluster.compute.drop_component(self.uuid)
+
     @property
     def inputs(self) -> Inputs:
         """Instance of `Inputs` to represent model params."""

superduper/ext/llm/prompter.py

@@ -67,9 +67,9 @@ class RetrievalPrompt(QueryModel):
     prompt_introduction: str = PROMPT_INTRODUCTION
     join: str = "\n---\n"
 
-    def __post_init__(self, db):
+    def __post_init__(self, db, example):
         assert 'prompt' in self.select.variables
-        return super().__post_init__(db)
+        return super().__post_init__(db, example)
 
     @property
     def inputs(self):

superduper/rest/base.py

@@ -1,6 +1,7 @@
 import sys
 import threading
 import time
+import os
 import typing as t
 from functools import cached_property
 from traceback import format_exc
@@ -198,13 +199,11 @@ def _add_templates(self, db):
 
             existing = db.show('template')
             for t in self.templates:
-                if t in existing:
+                if t in existing or t is None:
                     logging.info(f'Found existing template: {t}')
                     continue
                 logging.info(f'Applying template: {t}')
 
-                import os
-
                 if t is None:
                     continue
 

superduper/templates/__init__.py

@@ -30,6 +30,8 @@ def ls():
 def __getattr__(name: str):
     import re
 
+    breakpoint()
+
     if not re.match('.*[0-9]+\.[0-9]+\.[0-9]+.*', name):
         assert name in TEMPLATES, f'{name} not in supported templates {TEMPLATES}'
         file = TEMPLATES[name].split('/')[-1]