Greatly simplify the template saving logic

Author: blythed

CHANGELOG.md

@@ -18,6 +18,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Use databackend to perform metadata duties
 - Add `db.create` and `db.insert` instead of `auto_schema`
 - Merkel-tree implementation replacing random `.uuid` with deterministic implementation
+- Simplify the `Template` class
 
 #### New Features & Functionality
 

superduper/__init__.py

@@ -21,6 +21,7 @@
 from .base.schema import Schema
 from .components.application import Application
 from .components.component import Component
+from .components.cron_job import CronJob, FunctionCronJob
 from .components.dataset import Dataset
 from .components.listener import Listener
 from .components.metric import Metric
@@ -35,7 +36,7 @@
 from .components.plugin import Plugin
 from .components.streamlit import Streamlit
 from .components.table import Table
-from .components.template import QueryTemplate, Template
+from .components.template import Template
 from .components.vector_index import VectorIndex
 
 REQUIRES = [
@@ -54,6 +55,8 @@
     'QueryModel',
     'Validation',
     'Model',
+    'CronJob',
+    'FunctionCronJob',
     'Trainer',
     'model',
     'Listener',
@@ -65,7 +68,6 @@
     'Table',
     'Application',
     'Template',
-    'QueryTemplate',
     'Application',
     'Component',
     'pickle_serializer',

superduper/backends/base/backends.py

@@ -1,11 +1,99 @@
 import typing as t
 from abc import ABC, abstractmethod
+from collections import defaultdict
 
 if t.TYPE_CHECKING:
     from superduper.base.datalayer import Datalayer
     from superduper.components.component import Component
 
 
+class Bookkeeping(ABC):
+    """Mixin class for tracking components and associated tools."""
+
+    def __init__(self):
+        self.component_uuid_mapping = defaultdict(set)
+        self.uuid_component_mapping = {}
+        self.tool_uuid_mapping = defaultdict(set)
+        self.uuid_tool_mapping = {}
+        self.tools = {}
+
+    def build_tool(self, component: 'Component'):
+        """Build a tool from a component.
+
+        :param component: Component to build tool from.
+        """
+        pass
+
+    def get_tool(self, uuid: str):
+        """Get the tool from a uuid.
+
+        :param uuid: UUID of the tool.
+        """
+        tool_id = self.uuid_tool_mapping[uuid]
+        return self.tools[tool_id]
+
+    def put_component(self, component: 'Component', **kwargs):
+        """Put a component to the backend.
+
+        :param component: Component to put.
+        :param kwargs: kwargs dictionary.
+        """
+        tool = self.build_tool(component)
+        tool.db = self.db
+        self.component_uuid_mapping[(component.component, component.identifier)].add(
+            component.uuid
+        )
+        self.uuid_component_mapping[component.uuid] = (
+            component.component,
+            component.identifier,
+        )
+        self.uuid_tool_mapping[component.uuid] = tool.identifier
+        self.tool_uuid_mapping[tool.identifier].add(component.uuid)
+        self.tools[tool.identifier] = tool
+        tool.initialize(**kwargs)
+
+    def drop_component(self, component: str, identifier: str):
+        """Drop the component from backend.
+
+        :param component: Component name.
+        :param identifier: Component identifier.
+        """
+        uuids = self.component_uuid_mapping[(component, identifier)]
+        tool_ids = []
+        for uuid in uuids:
+            del self.uuid_component_mapping[uuid]
+            tool_id = self.uuid_tool_mapping[uuid]
+            tool_ids.append(tool_id)
+            del self.uuid_tool_mapping[uuid]
+            self.tool_uuid_mapping[tool_id].remove(uuid)
+            if not self.tool_uuid_mapping[tool_id]:
+                self.tools[tool_id].drop()
+                del self.tools[tool_id]
+        del self.component_uuid_mapping[(component, identifier)]
+
+    def drop(self):
+        """Drop the backend."""
+        for tool in self.tools.values():
+            tool.drop()
+        self.component_uuid_mapping = defaultdict(set)
+        self.uuid_component_mapping = {}
+        self.tool_uuid_mapping = defaultdict(set)
+        self.uuid_tool_mapping = {}
+        self.tools = {}
+
+    def list_components(self):
+        """List components, and identifiers deployed."""
+        return list(self.component_uuid_mapping.keys())
+
+    def list_tools(self):
+        """List tools deployed."""
+        return list(self.tools.keys())
+
+    def list_uuids(self):
+        """List uuids deployed."""
+        return list(self.uuid_component_mapping.keys())
+
+
 class BaseBackend(ABC):
     """Base backend class for cluster client."""
 
@@ -34,34 +122,19 @@ def initialize(self):
         """To be called on program start."""
         pass
 
-    def put_component(self, component: 'Component', **kwargs):
+    @abstractmethod
+    def put_component(self, component: 'Component'):
         """Add a component to the deployment.
 
         :param component: ``Component`` to put.
-        :param kwargs: kwargs dictionary.
         """
-        # This is to make sure that we only have 1 version
-        # of each component implemented at any given time
-        # TODO: get identifier in string component argument.
-        identifier = ''
-        if isinstance(component, str):
-            uuid = component
-        else:
-            uuid = component.uuid
-            identifier = component.identifier
-
-        if uuid in self.list_uuids():
-            return
-        if identifier in self.list_components():
-            del self[component.identifier]
-
-        self._put(component, **kwargs)
 
     @abstractmethod
     def drop_component(self, component: str, identifier: str):
         """Drop the component from backend.
 
-        :param identifier: Component identifier
+        :param component: Component name.
+        :param identifier: Component identifier.
         """
 
     @property

superduper/backends/base/cluster.py

@@ -1,8 +1,8 @@
 # TODO rename queue + crontab to scheduler
 import dataclasses as dc
 import time
-from abc import ABC, abstractmethod
 import typing as t
+from abc import ABC, abstractmethod
 
 from superduper.backends.base.cache import Cache
 from superduper.backends.base.cdc import CDCBackend
@@ -16,16 +16,15 @@
 class Cluster(ABC):
     """Cluster object for managing the backend.
 
-    :param compute: The compute backend.
     :param cache: The cache backend.
-    :param queue: The queue backend.
+    :param scheduler: The scheduler backend.
     :param vector_search: The vector search backend.
     :param cdc: The change data capture backend.
     :param crontab: The crontab backend.
     """
 
-    cache: Cache                
-    scheduler: BaseScheduler 
+    cache: Cache
+    scheduler: BaseScheduler
     vector_search: VectorSearchBackend
     cdc: CDCBackend
     crontab: CrontabBackend
@@ -87,11 +86,7 @@ def load_custom_plugins(self):
                 plugin = self.db.load('Plugin', plugin)
 
     def initialize(self):
-        """Initialize the cluster.
-
-        :param with_compute: Boolean to init
-                             compute.
-        """
+        """Initialize the cluster."""
         from superduper import logging
 
         start = time.time()

superduper/backends/base/compute.py

@@ -4,9 +4,9 @@
 from superduper.backends.base.backends import BaseBackend
 
 if t.TYPE_CHECKING:
-    from superduper.components.component import Component
     from superduper.base.datalayer import Datalayer
     from superduper.base.event import Job
+    from superduper.components.component import Component
 
 
 class ComputeBackend(BaseBackend):
@@ -23,14 +23,10 @@ class ComputeBackend(BaseBackend):
     def release_futures(self, context: str):
         """Release futures from backend.
 
-        :param context: Context of futures to release.
+        :param context: Futures context to release.
         """
         pass
 
-    def get_local_client(self):
-        """Returns a local version of self."""
-        pass
-
     @abstractmethod
     def submit(self, job: 'Job') -> t.Any:
         """
@@ -45,14 +41,15 @@ def disconnect(self) -> None:
         """Disconnect the client."""
         pass
 
+    @abstractmethod
     def initialize(self):
         """Connect to address."""
 
+    @abstractmethod
     def put_component(self, component: 'Component'):
         """Create handler on component declare.
 
-        :param args: *args for `create_handler`
-        :param kwargs: *kwargs for `create_handler`
+        :param component: Component to put.
         """
 
     @property
@@ -68,8 +65,10 @@ def db(self, value: 'Datalayer'):
         """
         self._db = value
 
-    def drop_component(self, uuid: str):
+    @abstractmethod
+    def drop_component(self, component: str, identifier: str):
         """Drop the component from compute.
 
-        :param uuid: Component uuid.
+        :param component: Component name.
+        :param identifier: Component identifier.
         """

superduper/backends/base/crontab.py

@@ -3,4 +3,3 @@
 
 class CrontabBackend(BaseBackend):
     """Base class for crontab backends."""
-