dynamo-llm 1.0.2

Dynamo LLM Library
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84

// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

use super::*;

use anyhow::Result;
use dynamo_runtime::{component::Component, pipeline::PushRouter, protocols::annotated::Annotated};
use futures::StreamExt;
use serde::de::DeserializeOwned;
use serde_json::Value;

pub struct ControlClient {
    client: PushRouter<ControlMessage, Annotated<Value>>,
    instance_id: i64,
}

impl ControlClient {
    pub async fn new(component: Component, instance_id: i64) -> Result<Self> {
        let client = component.endpoint("controller").client().await?;
        client.wait_for_instances().await?;
        let client =
            PushRouter::<ControlMessage, Annotated<Value>>::from_client(client, Default::default())
                .await?;

        Ok(Self {
            client,
            instance_id,
        })
    }

    #[tracing::instrument(level = "debug", skip(self), ret)]
    pub async fn status(&self, cache_level: CacheLevel) -> Result<BlockPoolStatus> {
        self.execute::<BlockPoolStatus>(ControlMessage::Status(cache_level))
            .await
    }

    #[tracing::instrument(level = "debug", skip(self), ret)]
    pub async fn reset_pool(&self, cache_level: CacheLevel) -> Result<()> {
        self.execute::<()>(ControlMessage::ResetPool(cache_level))
            .await
    }

    #[tracing::instrument(level = "debug", skip(self), ret)]
    pub async fn reset_blocks(
        &self,
        cache_level: CacheLevel,
        sequence_hashes: Vec<SequenceHash>,
    ) -> Result<ResetBlocksResponse> {
        self.execute::<ResetBlocksResponse>(ControlMessage::ResetBlocks(ResetRequest {
            cache_level,
            sequence_hashes,
        }))
        .await
    }

    #[tracing::instrument(level = "debug", skip(self), ret)]
    pub async fn reset_all_pools(&self) -> Result<()> {
        self.execute::<()>(ControlMessage::ResetAll).await
    }

    async fn execute<T: DeserializeOwned>(&self, message: ControlMessage) -> Result<T> {
        let mut stream = self
            .client
            .direct(message.into(), self.instance_id as u64)
            .await?;
        let resp = stream
            .next()
            .await
            .ok_or(anyhow::anyhow!("Failed to get a response from controller"))?;
        match resp.into_result() {
            Ok(data) => match data {
                Some(value) => {
                    let result: T = serde_json::from_value(value)?;
                    Ok(result)
                }
                None => {
                    let result: T = serde_json::from_value(Value::Null)?;
                    Ok(result)
                }
            },
            Err(e) => Err(e)?,
        }
    }
}