Add support for JSON and SSE responses to invoke endpoint (#1643)

* Add support for JSON and SSE responses to invoke endpoint

* Update API endpoint for storing and getting parameters and return type

* Using application/json as default accept header

* Using the new API in fe scaling test

* backward compatibility for invoke endpoint for inkwell

* Add output uri paths to task event

* remove return

Author: diptanu

.vscode/settings.json

@@ -2,5 +2,19 @@
     "python.analysis.extraPaths": [
         "./indexify/src",
         "./tensorlake/src"
+    ],
+    "cSpell.words": [
+        "cpus",
+        "nanos",
+        "readahead",
+        "retriable",
+        "rocksdb",
+        "sched",
+        "schedulable",
+        "strs",
+        "tensorlake",
+        "tombstoned",
+        "Upserted",
+        "utoipa"
     ]
 }

indexify/tests/cli/test_server_fe_scaling.py

@@ -80,7 +80,7 @@ def test_server_uses_the_same_function_executor_if_fe_task_queue_doesnt_overflow
 
         invocation_ids: List[str] = []
         for _ in range(_FE_ALLOCATIONS_QUEUE_SIZE):
-            invocation_id = graph.run(block_until_done=False, sleep_secs=0.01)
+            invocation_id = graph.call(sleep_secs=0.01)
             invocation_ids.append(invocation_id)
 
         fe_ids: Set[str] = set()

server/src/data_model/mod.rs

@@ -339,6 +339,10 @@ pub struct ComputeFn {
     #[serde(default)]
     pub retry_policy: FunctionRetryPolicy,
     pub cache_key: Option<CacheKey>,
+    #[serde(default)]
+    pub parameters: Vec<ParameterMetadata>,
+    #[serde(default)]
+    pub return_type: Option<serde_json::Value>,
 }
 
 impl ComputeFn {
@@ -414,6 +418,14 @@ pub struct RuntimeInformation {
     pub sdk_version: String,
 }
 
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+pub struct ParameterMetadata {
+    pub name: String,
+    pub description: Option<String>,
+    pub required: bool,
+    pub data_type: serde_json::Value,
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
 pub struct ComputeGraph {
     pub namespace: String,
@@ -1048,7 +1060,7 @@ pub enum TaskFailureReason {
     InternalError,
     // Clear function code failure typically by raising an exception from the function code.
     FunctionError,
-    // Function code run time exceeded its cofigured timeout.
+    // Function code run time exceeded its configured timeout.
     FunctionTimeout,
     // Function code raised InvocationError to mark the invocation as permanently failed.
     InvocationError,

server/src/http_objects.rs

@@ -424,6 +424,10 @@ pub struct ComputeFn {
     pub retry_policy: NodeRetryPolicy,
     #[serde(rename = "cache_key")]
     pub cache_key: Option<CacheKey>,
+    #[serde(default)]
+    pub parameters: Vec<ParameterMetadata>,
+    #[serde(default)]
+    pub return_type: Option<serde_json::Value>,
 }
 
 impl From<ComputeFn> for data_model::ComputeFn {
@@ -442,6 +446,8 @@ impl From<ComputeFn> for data_model::ComputeFn {
             resources: val.resources.into(),
             retry_policy: val.retry_policy.into(),
             cache_key: val.cache_key.and_then(|v| Some(v.into())),
+            parameters: val.parameters.into_iter().map(|p| p.into()).collect(),
+            return_type: val.return_type,
         }
     }
 }
@@ -461,6 +467,8 @@ impl From<data_model::ComputeFn> for ComputeFn {
             resources: c.resources.into(),
             retry_policy: c.retry_policy.into(),
             cache_key: c.cache_key.and_then(|v| Some(v.into())),
+            parameters: c.parameters.into_iter().map(|p| p.into()).collect(),
+            return_type: c.return_type,
         }
     }
 }
@@ -512,6 +520,35 @@ impl From<data_model::RuntimeInformation> for RuntimeInformation {
     }
 }
 
+#[derive(Debug, Serialize, Deserialize, ToSchema, Clone)]
+pub struct ParameterMetadata {
+    pub name: String,
+    pub description: Option<String>,
+    pub required: bool,
+    pub data_type: serde_json::Value,
+}
+
+impl From<data_model::ParameterMetadata> for ParameterMetadata {
+    fn from(parameter: data_model::ParameterMetadata) -> Self {
+        Self {
+            name: parameter.name,
+            description: parameter.description,
+            required: parameter.required,
+            data_type: parameter.data_type,
+        }
+    }
+}
+
+impl From<ParameterMetadata> for data_model::ParameterMetadata {
+    fn from(parameter: ParameterMetadata) -> Self {
+        Self {
+            name: parameter.name,
+            description: parameter.description,
+            required: parameter.required,
+            data_type: parameter.data_type,
+        }
+    }
+}
 #[derive(Debug, Serialize, Deserialize, ToSchema)]
 pub struct ComputeGraph {
     pub name: String,
@@ -1136,11 +1173,6 @@ pub struct ExecutorsAllocationsResponse {
     pub executors: Vec<ExecutorAllocations>,
 }
 
-#[derive(Debug, Serialize, Deserialize, ToSchema)]
-pub struct RequestQueryParams {
-    pub block_until_finish: Option<bool>,
-}
-
 #[cfg(test)]
 mod tests {
     use crate::http_objects::ComputeFn;

server/src/routes/compute_graphs.rs

@@ -203,10 +203,10 @@ pub async fn list_compute_graphs(
 /// Get a compute graph definition
 #[utoipa::path(
     get,
-    path = "/namespaces/{namespace}/compute_graphs/{compute_graph}",
+    path = "/v1/namespaces/{namespace}/compute-graphs/{compute_graph}",
     tag = "operations",
     responses(
-        (status = 200, description = "compute graph definition", body = http_objects::ComputeGraph),
+        (status = 200, description = "compute graph definition", body = http_objects_v1::ComputeGraph),
         (status = INTERNAL_SERVER_ERROR, description = "internal server error")
     ),
 )]

server/src/routes/download.rs

@@ -13,46 +13,6 @@ use crate::{
     http_objects::{IndexifyAPIError, RequestError},
 };
 
-pub async fn download_invocation_payload(
-    Path((namespace, compute_graph, invocation_id)): Path<(String, String, String)>,
-    State(state): State<RouteState>,
-) -> Result<Response<Body>, IndexifyAPIError> {
-    let output = state
-        .indexify_state
-        .reader()
-        .invocation_payload(&namespace, &compute_graph, &invocation_id)
-        .map_err(|e| {
-            IndexifyAPIError::internal_error(anyhow!(
-                "failed to download invocation payload: {}",
-                e
-            ))
-        })?;
-    let storage_reader = state
-        .blob_storage
-        .get(&output.payload.path)
-        .await
-        .map_err(IndexifyAPIError::internal_error)?;
-
-    if output.encoding == "application/json" {
-        let json_bytes = storage_reader
-            .map_ok(|chunk| chunk.to_vec())
-            .try_concat()
-            .await
-            .map_err(|e| IndexifyAPIError::internal_error(anyhow!("Failed to read JSON: {}", e)))?;
-
-        return Response::builder()
-            .header("Content-Type", output.encoding)
-            .body(Body::from(json_bytes))
-            .map_err(|e| IndexifyAPIError::internal_error_str(&e.to_string()));
-    }
-
-    Response::builder()
-        .header("Content-Type", output.encoding)
-        .header("Content-Length", output.payload.size.to_string())
-        .body(Body::from_stream(storage_reader))
-        .map_err(|e| IndexifyAPIError::internal_error_str(&e.to_string()))
-}
-
 pub async fn download_invocation_error(
     invocation_error: Option<GraphInvocationError>,
     blob_storage: &BlobStorage,