feat: include an optional list of tags when deploying a compute graph (…

…#1083) A remote graph can now include a map of tags that will be displayed on the Remote Graph overview.
tensorlakeai · Dec 6, 2024 · 855b4c9 · 855b4c9
1 parent d1bf656
commit 855b4c9
Show file tree

Hide file tree

Showing 9 changed files with 149 additions and 54 deletions.
diff --git a/python-sdk/indexify/functions_sdk/graph.py b/python-sdk/indexify/functions_sdk/graph.py
@@ -65,14 +65,19 @@ def is_pydantic_model_from_annotation(type_annotation):
 
 class Graph:
     def __init__(
-        self, name: str, start_node: IndexifyFunction, description: Optional[str] = None
+        self,
+        name: str,
+        start_node: IndexifyFunction,
+        description: Optional[str] = None,
+        tags: Dict[str, str] = {},
     ):
         self.name = name
         self.description = description
         self.nodes: Dict[str, Union[IndexifyFunction, IndexifyRouter]] = {}
         self.routers: Dict[str, List[str]] = defaultdict(list)
         self.edges: Dict[str, List[str]] = defaultdict(list)
         self.accumulator_zero_values: Dict[str, Any] = {}
+        self.tags = tags
 
         self.add_node(start_node)
         if issubclass(start_node, IndexifyRouter):
@@ -204,12 +209,15 @@ def definition(self) -> ComputeGraphMetadata:
                     )
                 )
 
+        print("TAGS", self.tags)
+
         return ComputeGraphMetadata(
             name=self.name,
             description=self.description or "",
             start_node=NodeMetadata(compute_fn=start_node),
             nodes=metadata_nodes,
             edges=metadata_edges,
+            tags=self.tags,
             runtime_information=RuntimeInformation(
                 major_version=sys.version_info.major,
                 minor_version=sys.version_info.minor,

diff --git a/python-sdk/indexify/functions_sdk/graph_definition.py b/python-sdk/indexify/functions_sdk/graph_definition.py
@@ -44,6 +44,7 @@ class ComputeGraphMetadata(BaseModel):
     name: str
     description: str
     start_node: NodeMetadata
+    tags: Dict[str, str] = {}
     nodes: Dict[str, NodeMetadata]
     edges: Dict[str, List[str]]
     accumulator_zero_values: Dict[str, bytes] = {}

diff --git a/server/data_model/src/lib.rs b/server/data_model/src/lib.rs
@@ -345,6 +345,8 @@ pub struct ComputeGraph {
     pub name: String,
     pub description: String,
     pub version: GraphVersion, // Version incremented with code update
+    #[serde(default)]
+    pub tags: HashMap<String, String>,
     pub code: ComputeGraphCode,
     pub created_at: u64,
     pub start_fn: Node,
@@ -376,6 +378,7 @@ impl ComputeGraph {
 
         self.description = update.description;
         self.runtime_information = update.runtime_information;
+        self.tags = update.tags;
 
         if self.code.sha256_hash != update.code.sha256_hash ||
             self.edges != update.edges ||
@@ -1053,6 +1056,7 @@ mod tests {
             namespace: TEST_NAMESPACE.to_string(),
             name: "graph1".to_string(),
             description: "description1".to_string(),
+            tags: HashMap::new(),
             nodes: HashMap::from([
                 ("fn_a".to_string(), Node::Compute(fn_a.clone())),
                 ("fn_b".to_string(), Node::Compute(fn_b.clone())),
@@ -1083,6 +1087,7 @@ mod tests {
             namespace: TEST_NAMESPACE.to_string(),
             name: "graph1".to_string(),
             description: "description2".to_string(),
+            tags: HashMap::from([("tag1".to_string(), "val1".to_string())]),
             nodes: HashMap::from([
                 ("fn_a".to_string(), Node::Compute(fn_a.clone())),
                 ("fn_b".to_string(), Node::Compute(fn_b.clone())),
@@ -1117,6 +1122,7 @@ mod tests {
         assert_eq!(graph.code.sha256_hash, "hash_code2", "update code");
         assert_eq!(graph.start_fn.name(), "fn_a", "update start_fn");
         assert_eq!(graph.version, GraphVersion(2), "update version");
+        assert!(graph.tags.contains_key("tag1"), "update tags");
         assert_eq!(
             graph.runtime_information.minor_version, 12,
             "update runtime_information"
@@ -1144,6 +1150,7 @@ mod tests {
                 namespace: String::new(),
                 name: String::new(),
                 description: String::new(),
+                tags: HashMap::new(),
                 version: GraphVersion::default(),
                 code: ComputeGraphCode {
                     path: String::new(),

diff --git a/server/data_model/src/test_objects.rs b/server/data_model/src/test_objects.rs
@@ -167,6 +167,10 @@ pub mod tests {
         ComputeGraph {
             namespace: TEST_NAMESPACE.to_string(),
             name: "graph_A".to_string(),
+            tags: HashMap::from([
+                ("tag1".to_string(), "val1".to_string()),
+                ("tag2".to_string(), "val2".to_string()),
+            ]),
             nodes: HashMap::from([
                 ("fn_b".to_string(), Node::Compute(fn_b)),
                 ("fn_c".to_string(), Node::Compute(fn_c)),
@@ -221,6 +225,10 @@ pub mod tests {
         ComputeGraph {
             namespace: TEST_NAMESPACE.to_string(),
             name: "graph_B".to_string(),
+            tags: HashMap::from([
+                ("tag1".to_string(), "val1".to_string()),
+                ("tag2".to_string(), "val2".to_string()),
+            ]),
             nodes: HashMap::from([
                 ("fn_b".to_string(), Node::Compute(fn_b)),
                 ("fn_c".to_string(), Node::Compute(fn_c)),
@@ -252,6 +260,10 @@ pub mod tests {
         ComputeGraph {
             namespace: TEST_NAMESPACE.to_string(),
             name: "graph_R".to_string(),
+            tags: HashMap::from([
+                ("tag1".to_string(), "val1".to_string()),
+                ("tag2".to_string(), "val2".to_string()),
+            ]),
             nodes: HashMap::from([
                 ("fn_a".to_string(), Node::Compute(fn_a.clone())),
                 ("fn_b".to_string(), Node::Compute(fn_b)),

diff --git a/server/src/http_objects.rs b/server/src/http_objects.rs
@@ -286,6 +286,8 @@ pub struct ComputeGraph {
     pub start_node: Node,
     #[serde(skip_deserializing)]
     pub version: u32,
+    #[serde(default)]
+    pub tags: Option<HashMap<String, String>>,
     pub nodes: HashMap<String, Node>,
     pub edges: HashMap<String, Vec<String>>,
     #[serde(default = "get_epoch_time_in_ms")]
@@ -313,6 +315,7 @@ impl ComputeGraph {
             namespace: self.namespace,
             description: self.description,
             start_fn,
+            tags: self.tags.unwrap_or_default(),
             version: Default::default(),
             code: ComputeGraphCode {
                 sha256_hash: sha256_hash.to_string(),
@@ -344,6 +347,7 @@ impl From<data_model::ComputeGraph> for ComputeGraph {
             namespace: compute_graph.namespace,
             description: compute_graph.description,
             start_node: start_fn,
+            tags: Some(compute_graph.tags),
             version: compute_graph.version.0,
             nodes,
             edges: compute_graph.edges,

diff --git a/server/src/scheduler.rs b/server/src/scheduler.rs
@@ -609,6 +609,7 @@ mod tests {
             ComputeGraph {
                 namespace: TEST_NAMESPACE.to_string(),
                 name: "graph_R".to_string(),
+                tags: HashMap::new(),
                 nodes: HashMap::from([
                     ("fn_gen".to_string(), Node::Compute(fn_gen.clone())),
                     ("fn_map".to_string(), Node::Compute(fn_map)),
@@ -856,6 +857,7 @@ mod tests {
             ComputeGraph {
                 namespace: TEST_NAMESPACE.to_string(),
                 name: "graph_R".to_string(),
+                tags: HashMap::new(),
                 nodes: HashMap::from([
                     ("fn_gen".to_string(), Node::Compute(fn_gen.clone())),
                     ("fn_map".to_string(), Node::Compute(fn_map)),
@@ -1168,6 +1170,7 @@ mod tests {
             ComputeGraph {
                 namespace: TEST_NAMESPACE.to_string(),
                 name: "graph_R".to_string(),
+                tags: HashMap::new(),
                 nodes: HashMap::from([
                     ("fn_gen".to_string(), Node::Compute(fn_gen.clone())),
                     ("fn_map".to_string(), Node::Compute(fn_map)),
@@ -1447,6 +1450,7 @@ mod tests {
             ComputeGraph {
                 namespace: TEST_NAMESPACE.to_string(),
                 name: "graph_R".to_string(),
+                tags: HashMap::new(),
                 nodes: HashMap::from([
                     ("fn_gen".to_string(), Node::Compute(fn_gen.clone())),
                     ("fn_map".to_string(), Node::Compute(fn_map)),
@@ -1826,6 +1830,7 @@ mod tests {
             ComputeGraph {
                 namespace: TEST_NAMESPACE.to_string(),
                 name: "graph_R".to_string(),
+                tags: HashMap::new(),
                 nodes: HashMap::from([
                     ("fn_gen".to_string(), Node::Compute(fn_gen.clone())),
                     ("fn_map".to_string(), Node::Compute(fn_map)),

diff --git a/server/ui/src/components/cards/ComputeGraphsCard.tsx b/server/ui/src/components/cards/ComputeGraphsCard.tsx
@@ -1,5 +1,14 @@
 import { useState } from 'react'
-import { Alert, Card, CardContent, Grid, IconButton, Typography } from '@mui/material'
+import {
+  Alert,
+  Card,
+  CardContent,
+  Chip,
+  Grid,
+  IconButton,
+  ListItem,
+  Typography,
+} from '@mui/material'
 import { Box, Stack } from '@mui/system'
 import { Cpu, InfoCircle } from 'iconsax-react'
 import DeleteIcon from '@mui/icons-material/Delete'
@@ -15,14 +24,22 @@ interface ComputeGraphsCardProps {
   namespace: string
 }
 
-export function ComputeGraphsCard({ client, computeGraphs, namespace }: ComputeGraphsCardProps) {
-  const [localGraphs, setLocalGraphs] = useState<ComputeGraph[]>(computeGraphs.compute_graphs || [])
+export function ComputeGraphsCard({
+  client,
+  computeGraphs,
+  namespace,
+}: ComputeGraphsCardProps) {
+  const [localGraphs, setLocalGraphs] = useState<ComputeGraph[]>(
+    computeGraphs.compute_graphs || []
+  )
   const [error, setError] = useState<string | null>(null)
 
   async function handleDeleteGraph(graphName: string) {
     try {
       await client.deleteComputeGraph(graphName)
-      setLocalGraphs(prevGraphs => prevGraphs.filter(graph => graph.name !== graphName))
+      setLocalGraphs((prevGraphs) =>
+        prevGraphs.filter((graph) => graph.name !== graphName)
+      )
     } catch (err) {
       console.error('Error deleting compute graph:', err)
       setError('Failed to delete compute graph. Please try again.')
@@ -32,16 +49,22 @@ export function ComputeGraphsCard({ client, computeGraphs, namespace }: ComputeG
   function renderGraphCard(graph: ComputeGraph) {
     return (
       <Grid item xs={12} sm={6} md={4} key={graph.name} mb={2}>
-        <Card sx={{ minWidth: 275, height: '100%', boxShadow: '0px 0px 2px 0px rgba(51, 132, 252, 0.5) inset' }}>
+        <Card
+          sx={{
+            minWidth: 275,
+            height: '100%',
+            boxShadow: '0px 0px 2px 0px rgba(51, 132, 252, 0.5) inset',
+          }}
+        >
           <CardContent>
             <Box display="flex" justifyContent="space-between">
               <Link to={`/${namespace}/compute-graphs/${graph.name}`}>
                 <TruncatedText text={graph.name} maxLength={20} />
               </Link>
               <Box display="flex" flexDirection="row">
                 <CopyText text={graph.name} />
-                <IconButton 
-                  onClick={() => handleDeleteGraph(graph.name)} 
+                <IconButton
+                  onClick={() => handleDeleteGraph(graph.name)}
                   aria-label="delete compute graph"
                 >
                   <DeleteIcon color="error" />
@@ -57,20 +80,53 @@ export function ComputeGraphsCard({ client, computeGraphs, namespace }: ComputeG
             <Typography variant="subtitle2" color="text.secondary">
               Number of Nodes: {Object.keys(graph.nodes || {}).length}
             </Typography>
+            {graph.created_at && (
+              <Typography variant="subtitle2" color="text.secondary">
+                Created At: {new Date(graph.created_at).toLocaleString()}
+              </Typography>
+            )}
+            <Typography variant="subtitle2" color="text.secondary">
+              Tags:
+              {Object.keys(graph.tags || {}).length > 0 && (
+                <Box display="flex" flexWrap="wrap" mt={1}>
+                  {Object.entries(graph.tags).map(([key, value]) => (
+                    <ListItem key={key}>
+                      <Chip
+                        key={key}
+                        label={`${key}: ${value}`}
+                        color="primary"
+                        size="small"
+                        sx={{ m: 0.5 }}
+                      />
+                    </ListItem>
+                  ))}
+                </Box>
+              )}
+            </Typography>
           </CardContent>
         </Card>
       </Grid>
     )
   }
 
   function renderContent() {
-    if (error) 
-      return <Alert variant="outlined" severity="error" sx={{ my: 2 }}>{error}</Alert>
+    if (error)
+      return (
+        <Alert variant="outlined" severity="error" sx={{ my: 2 }}>
+          {error}
+        </Alert>
+      )
 
     if (!localGraphs.length)
-      return <Alert variant="outlined" severity="info" sx={{ my: 2 }}>No Graphs Found</Alert>
+      return (
+        <Alert variant="outlined" severity="info" sx={{ my: 2 }}>
+          No Graphs Found
+        </Alert>
+      )
 
-    const sortedGraphs = [...localGraphs].sort((a, b) => a.name.localeCompare(b.name))
+    const sortedGraphs = [...localGraphs].sort((a, b) =>
+      a.name.localeCompare(b.name)
+    )
 
     return (
       <Box sx={{ width: '100%', overflow: 'auto', borderRadius: '5px' }} mt={2}>