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

define_api_id!(0x75c9_e1a5_cd99_4b85, "ml-v5");

pub use super::ml_v1::FutureHandle;
pub use super::ml_v1::TrainingHandle;

use crate::ml_v3::InferenceHandle;
use crate::FFIResult;

#[allow(clippy::too_many_arguments)]
#[ark_api_macros::ark_bindgen(imports = "ark-ml-v5")]
mod ml {
    use super::*;

    #[derive(Clone, Debug, Copy, Eq, PartialEq)]
    #[repr(u32)]
    #[non_exhaustive]
    pub enum HardwareType {
        CPU = 0,
        GPU = 1,
        LargeGPU = 2,
    }

    extern "C" {
        /// Starts training.
        ///
        /// The promise outputs a `TrainingHandle`.
        pub fn start_training(
            hive_url: &str,
            hive_port: u32,
            experiment_name: &str,
            configuration: &str,
            namespace: &str,
            module_id: &str,
            duration: u32,
            worker_count: u32,
            hardware_type: HardwareType,
        ) -> FFIResult<FutureHandle>;

        /// Retrieves the configuration.
        pub fn raw_experiment_config(context: TrainingHandle) -> FFIResult<FutureHandle>;

        /// Submits observations previously pushed using `push_inference_observation`.
        ///
        /// If an observation was terminal or interrupted then the actions for that observation will be empty.
        ///
        /// Returns a serialized byte-slice in NE ordering. The serialization format is:
        ///
        /// u8  : version
        /// u8  : response-count
        /// u16 : action-count
        /// repeated :
        ///   u64 : id
        ///   u8  : has-data   -- the behavior of returning empty data for finished episodes is deprecated but changing right now is not worth it
        ///   if has-data != 0 :
        ///     repeated f32 : actions
        pub fn submit_inference_observations(context: InferenceHandle) -> FFIResult<Vec<u8>>;

        /// Submits observations previously pushed using `push_training_observation`.
        ///
        /// If an observation was terminal or interrupted then the actions for that observation will be empty.
        ///
        /// Returns a serialized byte-slice following the same serialization format as [`submit_inference_observations`].
        pub fn submit_training_observations(context: TrainingHandle) -> FFIResult<Vec<u8>>;

        /// Loads the provided brain data into the snapshot.
        ///
        /// `snapshot_format` has to match one of the `SnapshotFormat` members.
        pub fn start_inference(
            snapshot_data: &[u8],
            snapshot_format: u32,
        ) -> FFIResult<InferenceHandle>;
    }
}

pub use ml::*;