asterai 1.0.0-alpha.16

CLI for asterai - the portable AI compute platform
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
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132

use crate::command::component::ComponentArgs;
use crate::config::ARTIFACTS_DIR;
use crate::registry::RegistryClient;
use crate::version_resolver::ComponentRef;
use asterai_runtime::resource::Resource;
use asterai_runtime::resource::metadata::ResourceKind;
use eyre::{Context, OptionExt, bail};
use std::fs;
use std::path::{Path, PathBuf};
use std::str::FromStr;

#[derive(Debug)]
pub(super) struct PullArgs {
    /// Component reference (version optional, will resolve to latest if omitted).
    component_ref: ComponentRef,
    /// Output directory (defaults to local resource storage).
    output: Option<String>,
}

impl PullArgs {
    pub fn parse(mut args: impl Iterator<Item = String>) -> eyre::Result<Self> {
        let mut component_ref: Option<ComponentRef> = None;
        let mut output: Option<String> = None;
        while let Some(arg) = args.next() {
            match arg.as_str() {
                "-o" | "--output" => {
                    output = Some(args.next().ok_or_eyre("missing value for output flag")?);
                }
                other => {
                    if other.starts_with('-') {
                        bail!("unknown flag: {}", other);
                    }
                    if component_ref.is_some() {
                        bail!("unexpected argument: {}", other);
                    }
                    component_ref = Some(ComponentRef::parse(other).wrap_err(
                        "invalid component reference \
                         (expected namespace:name or namespace:name@version)",
                    )?);
                }
            }
        }
        let component_ref = component_ref.ok_or_eyre(
            "missing component reference \
            (e.g., namespace:component or namespace:component@version)",
        )?;
        Ok(Self {
            component_ref,
            output,
        })
    }

    async fn execute(&self, api_url: &str, registry_url: &str) -> eyre::Result<()> {
        // Resolve version if not specified.
        let resolved = self.component_ref.resolve(api_url, registry_url).await?;
        let component = Resource::from_str(&resolved)?;
        let namespace = component.namespace();
        let name = component.name();
        let version = component.version().to_string();
        let repo_name = format!("{}/{}", namespace, name);
        let tag = &version;
        println!("pulling {}@{}", repo_name, tag);
        let client = reqwest::Client::new();
        let registry = RegistryClient::new(&client, api_url, registry_url);
        let token = registry.get_token(None, &repo_name).await?;
        println!("fetching manifest...");
        let (manifest, manifest_digest) = registry.fetch_manifest(&repo_name, tag, &token).await?;
        let output_dir = self.determine_output_dir(namespace, name, &version)?;
        // Download layers.
        for (i, layer) in manifest.layers.iter().enumerate() {
            println!("downloading layer {} ({})...", i + 1, &layer.digest[..19]);
            let blob_bytes = registry
                .download_blob(&repo_name, &layer.digest, &token)
                .await?;
            let filename = match i {
                0 => "component.wasm",
                _ => "package.wasm",
            };
            let file_path = output_dir.join(filename);
            fs::write(&file_path, &blob_bytes)
                .wrap_err_with(|| format!("failed to write {}", filename))?;
            println!("  saved to {}", file_path.display());
        }
        // Fetch WIT package via referrers API.
        if let Ok(Some(wit_bytes)) = registry
            .fetch_wit_referrer(&repo_name, &manifest_digest, &token)
            .await
        {
            let file_path = output_dir.join("package.wasm");
            fs::write(&file_path, &wit_bytes).wrap_err("failed to write package.wasm")?;
            println!("  saved to {}", file_path.display());
        }
        self.write_metadata(&output_dir, &repo_name, tag)?;
        println!("pulled component saved to {}", output_dir.display());
        Ok(())
    }

    fn determine_output_dir(
        &self,
        namespace: &str,
        name: &str,
        version: &str,
    ) -> eyre::Result<PathBuf> {
        let output_dir = match &self.output {
            Some(dir) => PathBuf::from(dir),
            None => ARTIFACTS_DIR
                .join(namespace)
                .join(format!("{}@{}", name, version)),
        };
        fs::create_dir_all(&output_dir).wrap_err("failed to create output directory")?;
        Ok(output_dir)
    }

    fn write_metadata(&self, output_dir: &Path, repo_name: &str, tag: &str) -> eyre::Result<()> {
        let metadata = serde_json::json!({
            "kind": ResourceKind::Component.to_string(),
            "pulled_from": format!("{}@{}", repo_name, tag),
        });
        let metadata_path = output_dir.join("metadata.json");
        fs::write(&metadata_path, serde_json::to_string_pretty(&metadata)?)
            .wrap_err("failed to write metadata")?;
        Ok(())
    }
}

impl ComponentArgs {
    pub async fn pull(&self) -> eyre::Result<()> {
        let args = self.pull_args.as_ref().ok_or_eyre("no pull args")?;
        args.execute(&self.api_endpoint, &self.registry_endpoint)
            .await
    }
}