11use crate :: uniffi:: {
22 error:: { OrcaError , Result , selector} ,
3- model:: pod:: { PodJob , PodResult } ,
43 orchestrator:: agent:: { Agent , AgentClient } ,
5- store:: ModelID ,
64} ;
7- use chrono:: Utc ;
5+ use chrono:: { DateTime , Utc } ;
86use futures_util:: future:: FutureExt as _;
97use regex:: Regex ;
108use serde:: { Deserialize , Serialize } ;
@@ -21,15 +19,20 @@ use tokio::{
2119use tokio_util:: task:: TaskTracker ;
2220
2321#[ expect( clippy:: expect_used, reason = "Valid static regex" ) ]
24- static RE_PODJOB_ACTION : LazyLock < Regex > = LazyLock :: new ( || {
22+ static RE_AGENT_KEY_EXPR : LazyLock < Regex > = LazyLock :: new ( || {
2523 Regex :: new (
2624 r"(?x)
2725 ^
28- group\/(?<group>[a-z_\-]+)\/
29- (?<action>request|reservation|success|failure)\/
30- pod_job\/(?<pod_job_hash>[0-9a-f]+)\/
31- host\/(?<host>[a-z_]+)\/
32- timestamp\/(?<timestamp>.*?)
26+ group/
27+ (?<group>[a-z_\-]+)/
28+ (?<action>request|success|failure)/
29+ (?<model_type>[a-z_]+)/
30+ (?<ref>[0-9a-f]+)/
31+ .*?
32+ host/
33+ (?<host>[a-z_]+)/
34+ timestamp/
35+ (?<timestamp>.*?)
3336 $
3437 " ,
3538 )
@@ -40,33 +43,14 @@ static RE_PODJOB_ACTION: LazyLock<Regex> = LazyLock::new(|| {
4043 dead_code,
4144 reason = "Need to be able to initialize to pass metadata as input."
4245) ]
43- #[ derive( Debug , Clone ) ]
46+ #[ derive( Debug ) ]
4447pub struct EventMetadata {
45- group : String ,
46- host : String ,
47- subgroup : String ,
48- }
49-
50- #[ expect(
51- dead_code,
52- reason = "Need to be able to initialize to pass metadata as input."
53- ) ]
54- #[ derive( Debug , Clone ) ]
55- pub enum EventPayload {
56- Request ( PodJob ) ,
57- Reservation ( ModelID ) ,
58- Success ( PodResult ) ,
59- Failure ( PodResult ) ,
60- }
61-
62- #[ expect(
63- dead_code,
64- reason = "Need to be able to initialize to pass metadata as input."
65- ) ]
66- #[ derive( Debug , Clone ) ]
67- pub struct Event {
68- metadata : EventMetadata ,
69- payload : EventPayload ,
48+ pub group : String ,
49+ pub action : String ,
50+ pub model_type : String ,
51+ pub r#ref : String ,
52+ pub host : String ,
53+ pub timestamp : DateTime < Utc > ,
7054}
7155
7256impl AgentClient {
@@ -95,7 +79,6 @@ impl AgentClient {
9579///
9680/// Will fail if there is an issue sending the message.
9781pub ( crate ) async fn log ( & self , message : & str ) -> Result < ( ) > {
98- println ! ( "{message}" ) ;
9982 self . publish ( "log" , message) . await
10083 }
10184}
@@ -106,23 +89,20 @@ impl AgentClient {
10689 reason = "`result::Result<(), SendError<_>>` is the only uncaptured result since it would mean we can't transmit results over mpsc."
10790) ]
10891pub async fn start_service <
109- EventClassifierF , // function to classify the event payload e.g. EventPayload::{Request | Reservation | ..}
110- RequestF , // function to run on requests
111- RequestI , // input to the function for requests
112- RequestR , // output to the function for requests
113- ResponseF , // function to run on completing a request i.e. response
114- ResponseI , // input to the function for responses
115- ResponseR , // output to the function for responses
92+ RequestF , // function to run on requests
93+ RequestI , // input to the function for requests
94+ RequestR , // output to the function for requests
95+ ResponseF , // function to run on completing a request i.e. response
96+ ResponseI , // input to the function for responses
97+ ResponseR , // output to the function for responses
11698> (
11799 agent : Arc < Agent > ,
118100 request_key_expr : String ,
119101 namespace_lookup : HashMap < String , PathBuf > ,
120- event_classifier : EventClassifierF ,
121102 request_task : RequestF ,
122103 response_task : ResponseF ,
123104) -> Result < ( ) >
124105where
125- EventClassifierF : Fn ( & RequestI ) -> EventPayload + Send + ' static ,
126106 RequestI : for < ' serde > Deserialize < ' serde > + Send + ' static ,
127107 RequestF : FnOnce ( Arc < Agent > , HashMap < String , PathBuf > , EventMetadata , RequestI ) -> RequestR
128108 + Clone
@@ -156,15 +136,17 @@ where
156136 while let Ok ( sample) = subscriber. recv_async ( ) . await {
157137 if let ( Ok ( input) , Some ( metadata) ) = (
158138 serde_json:: from_slice :: < RequestI > ( & sample. payload ( ) . to_bytes ( ) ) ,
159- RE_PODJOB_ACTION . captures ( sample. key_expr ( ) . as_str ( ) ) ,
139+ RE_AGENT_KEY_EXPR . captures ( sample. key_expr ( ) . as_str ( ) ) ,
160140 ) {
161141 let inner_response_tx = response_tx. clone ( ) ;
162142 let event_metadata = EventMetadata {
163143 group : metadata[ "group" ] . to_string ( ) ,
144+ action : metadata[ "action" ] . to_string ( ) ,
145+ model_type : metadata[ "model_type" ] . to_string ( ) ,
146+ r#ref : metadata[ "ref" ] . to_string ( ) ,
164147 host : metadata[ "host" ] . to_string ( ) ,
165- subgroup : metadata[ "pod_job_hash" ] . to_string ( ) ,
148+ timestamp : DateTime :: parse_from_rfc3339 ( & metadata[ "timestamp" ] ) ? . into ( ) ,
166149 } ;
167- let _event_payload = event_classifier ( & input) ;
168150 tasks. spawn ( {
169151 let inner_request_task = request_task. clone ( ) ;
170152 let inner_inner_agent = Arc :: clone ( & inner_agent) ;
@@ -190,8 +172,8 @@ where
190172 }
191173 } ) ;
192174 services. spawn ( async move {
193- while let Some ( content ) = response_rx. recv ( ) . await {
194- response_task ( Arc :: clone ( & agent. client ) , content ?) . await ?;
175+ while let Some ( response ) = response_rx. recv ( ) . await {
176+ response_task ( Arc :: clone ( & agent. client ) , response ?) . await ?;
195177 }
196178 Ok ( ( ) )
197179 } ) ;
0 commit comments