factory.rs 3.44 KB
Newer Older
1
2
3
//! Factory for creating router instances

use super::{pd_router::PDRouter, router::Router, RouterTrait};
4
use crate::config::{PolicyConfig, RoutingMode};
5
use crate::policies::PolicyFactory;
6
7
use crate::server::AppContext;
use std::sync::Arc;
8
9
10
11
12

/// Factory for creating router instances based on configuration
pub struct RouterFactory;

impl RouterFactory {
13
    /// Create a router instance from application context
14
    pub async fn create_router(ctx: &Arc<AppContext>) -> Result<Box<dyn RouterTrait>, String> {
15
        match &ctx.router_config.mode {
16
            RoutingMode::Regular { worker_urls } => {
17
                Self::create_regular_router(worker_urls, &ctx.router_config.policy, ctx).await
18
19
20
21
            }
            RoutingMode::PrefillDecode {
                prefill_urls,
                decode_urls,
22
23
                prefill_policy,
                decode_policy,
24
25
26
27
28
29
30
31
32
33
34
            } => {
                Self::create_pd_router(
                    prefill_urls,
                    decode_urls,
                    prefill_policy.as_ref(),
                    decode_policy.as_ref(),
                    &ctx.router_config.policy,
                    ctx,
                )
                .await
            }
35
36
37
38
        }
    }

    /// Create a regular router with injected policy
39
    async fn create_regular_router(
40
41
        worker_urls: &[String],
        policy_config: &PolicyConfig,
42
        ctx: &Arc<AppContext>,
43
44
45
46
    ) -> Result<Box<dyn RouterTrait>, String> {
        // Create policy
        let policy = PolicyFactory::create_from_config(policy_config);

47
        // Create regular router with injected policy and client
48
49
50
        let router = Router::new(
            worker_urls.to_vec(),
            policy,
51
52
53
54
55
            ctx.client.clone(),
            ctx.router_config.worker_startup_timeout_secs,
            ctx.router_config.worker_startup_check_interval_secs,
            ctx.router_config.dp_aware,
            ctx.router_config.api_key.clone(),
56
            ctx.router_config.retry.clone(),
57
            ctx.router_config.circuit_breaker.clone(),
58
            ctx.router_config.health_check.clone(),
59
60
        )
        .await?;
61
62
63
64
65

        Ok(Box::new(router))
    }

    /// Create a PD router with injected policy
66
    async fn create_pd_router(
67
68
        prefill_urls: &[(String, Option<u16>)],
        decode_urls: &[String],
69
70
71
        prefill_policy_config: Option<&PolicyConfig>,
        decode_policy_config: Option<&PolicyConfig>,
        main_policy_config: &PolicyConfig,
72
        ctx: &Arc<AppContext>,
73
    ) -> Result<Box<dyn RouterTrait>, String> {
74
75
76
77
78
        // Create policies - use specific policies if provided, otherwise fall back to main policy
        let prefill_policy =
            PolicyFactory::create_from_config(prefill_policy_config.unwrap_or(main_policy_config));
        let decode_policy =
            PolicyFactory::create_from_config(decode_policy_config.unwrap_or(main_policy_config));
79

80
        // Create PD router with separate policies and client
81
82
83
        let router = PDRouter::new(
            prefill_urls.to_vec(),
            decode_urls.to_vec(),
84
85
            prefill_policy,
            decode_policy,
86
87
88
            ctx.client.clone(),
            ctx.router_config.worker_startup_timeout_secs,
            ctx.router_config.worker_startup_check_interval_secs,
89
            ctx.router_config.retry.clone(),
90
            ctx.router_config.circuit_breaker.clone(),
91
            ctx.router_config.health_check.clone(),
92
93
        )
        .await?;
94
95
96
97

        Ok(Box::new(router))
    }
}