Refactor producer/consumer logic, add to operations (#52)

# Description

- Refactors the way we're attaching producers and consumers
- Adds input/output tensors to operation list

backend/models.py

@@ -60,7 +60,6 @@
     PrimaryKeyConstraint("operation_id", "output_index", "tensor_id"),
 )
 
-
 stack_traces = Table(
     "stack_traces",
     db.metadata,
@@ -69,7 +68,6 @@
     PrimaryKeyConstraint("operation_id", "stack_trace")
 )
 
-
 buffers = Table(
     "buffers",
     db.metadata,
@@ -104,14 +102,20 @@
     Column("cb_limit", Integer),
 )
 
+
 class Device(db.Model):
     __table__ = devices
 
 
-
 class Tensor(db.Model):
     __table__ = tensors
 
+    def producers(self):
+        return [c.operation_id for c in OutputTensor.query.filter_by(tensor_id=self.tensor_id)]
+
+    def consumers(self):
+        return [c.operation_id for c in InputTensor.query.filter_by(tensor_id=self.tensor_id)]
+
 
 class Buffer(db.Model):
     __table__ = buffers
@@ -126,11 +130,13 @@ class InputTensor(db.Model):
 class StackTrace(db.Model):
     __table__ = stack_traces
 
+
 class OutputTensor(db.Model):
     __table__ = output_tensors
     tensor = db.relationship("Tensor", backref="output")
 
 
+
 class Operation(db.Model):
     __table__ = operations
     arguments = db.relationship("OperationArgument", backref="operation")

backend/schemas.py

@@ -1,5 +1,3 @@
-from typing import Optional
-
 from marshmallow import fields, validates
 
 from backend.extensions import ma
@@ -19,15 +17,20 @@
 class TensorSchema(ma.SQLAlchemyAutoSchema):
     class Meta:
         model = Tensor
+
     shape = ma.auto_field()
     address = ma.auto_field()
+    consumers = fields.List(fields.Integer, default=[])
+    producers = fields.List(fields.Integer, default=[])
 
 
 class StackTraceSchema(ma.SQLAlchemyAutoSchema):
     class Meta:
         model = StackTrace
+
     stack_trace = ma.Function(lambda obj: obj.stack_trace or "")
 
+
 class InputOutputSchema(object):
     # TODO - We can probably create a model to avoid backrefs
     shape = ma.Function(lambda obj: obj.tensor.shape)
@@ -38,8 +41,8 @@ class InputOutputSchema(object):
     buffer_type = ma.Function(lambda obj: obj.tensor.buffer_type)
     dtype = ma.Function(lambda obj: obj.tensor.dtype)
     tensor_id = ma.Function(lambda obj: obj.tensor.tensor_id)
-    consumers = fields.List(fields.Integer, default=[])
-    producers = fields.List(fields.Integer, default=[])
+    producers = ma.Function(lambda obj: obj.tensor.producers())
+    consumers = ma.Function(lambda obj: obj.tensor.consumers())
     operation_id = ma.auto_field()
 
 

backend/views.py

@@ -1,16 +1,14 @@
-from http import HTTPStatus
 import json
-from pathlib import Path
 import shutil
+from http import HTTPStatus
+from pathlib import Path
 
 from flask import Blueprint, Response, current_app, request
 
 from backend.models import (
     Device,
     Operation,
     Buffer,
-    InputTensor,
-    OutputTensor,
     Tensor,
     StackTrace,
 )
@@ -47,23 +45,11 @@ def operation_list():
         many=True,
         exclude=[
             "buffers",
-            "input_tensors",
-            "output_tensors",
             "operation_id",
         ],
     ).dump(operations)
 
 
-def attach_producers_consumers(t: Tensor):
-    t.consumers = [
-        c.operation_id for c in InputTensor.query.filter_by(tensor_id=t.tensor_id)
-    ]
-    t.producers = [
-        c.operation_id for c in OutputTensor.query.filter_by(tensor_id=t.tensor_id)
-    ]
-    return t
-
-
 @api.route("/operations/<operation_id>", methods=["GET"])
 def operation_detail(operation_id):
     operation = Operation.query.get(operation_id)
@@ -80,10 +66,10 @@ def operation_detail(operation_id):
     stack_trace_dump = StackTraceSchema().dump(stack_trace, many=False)
     stack_trace_value = stack_trace_dump.get("stack_trace")
     input_tensors = InputTensorSchema().dump(
-        map(attach_producers_consumers, operation.input_tensors), many=True
+        operation.input_tensors, many=True
     )
     output_tensors = OutputTensorSchema().dump(
-        map(attach_producers_consumers, operation.output_tensors), many=True
+        operation.output_tensors, many=True
     )
 
     return dict(
@@ -125,7 +111,7 @@ def get_config():
 
 @api.route("/tensors", methods=["GET"])
 def get_tensors():
-    tensors = map(attach_producers_consumers, Tensor.query.all())
+    tensors = Tensor.query.all()
     return TensorSchema().dump(tensors, many=True)
 
 
@@ -134,7 +120,7 @@ def get_tensor(tensor_id):
     tensor = Tensor.query.get(tensor_id)
     if not tensor:
         return Response(status=HTTPStatus.NOT_FOUND)
-    return TensorSchema().dump(attach_producers_consumers(tensor))
+    return TensorSchema().dump(tensor)
 
 
 @api.route(
@@ -146,7 +132,6 @@ def get_tensor(tensor_id):
 def create_upload_files():
     """
     Copies the folder upload into the active data directory
-    :param files:
     :return:
     """
     files = request.files.getlist("files")