Fixed model cleanup on cluster

CHANGELOG.md

@@ -13,6 +13,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Deprecate vanilla `DataType`
 - Remove `_Encodable` from project 
+- 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

@@ -541,6 +541,7 @@ def remove(
 
             for v in sorted(versions_in_use):
                 self.metadata.hide_component_version(type_id, identifier, v)
+
         else:
             logging.warn('aborting.')
 
@@ -664,10 +665,7 @@ def _remove_component_version(
                 type_id, identifier, version=version, allow_hidden=force
             )
             component.cleanup(self)
-            try:
-                del self.cluster.cache[component.uuid]
-            except KeyError:
-                pass
+            self.cluster.drop_component(component.uuid)
 
             self._delete_artifacts(r['uuid'], info)
             self.metadata.delete_component_version(type_id, identifier, version=version)

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

@@ -198,14 +198,16 @@ 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 os.path.exists(t):
                     from superduper import Template
+
                     t = Template.read(t)
                 else:
                     t = templates.get(t)

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]